libcom 0.1.8__py3-none-any.whl

libcom/fos_score/fos_score_prediction.py

@@ -0,0 +1,218 @@
+import torch
+import torchvision
+from libcom.utils.model_download import download_pretrained_model
+from libcom.utils.process_image import *
+from libcom.utils.environment import *
+from libcom.fos_score.source.config import Config
+from libcom.fos_score.source.networks import StudentModel, SingleScaleD
+import torch 
+import os
+import math
+import datetime
+import torchvision.transforms as transforms
+from torchvision.utils import save_image
+from einops import rearrange
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['FOS_D'] 
+
+class FOSScoreModel:
+    """
+    Foreground object search score prediction model.
+
+    Args:
+        device (str | torch.device): gpu id
+        model_type (str): predefined model type
+        kwargs (dict): other parameters for building model
+    
+    Examples:
+        >>> from libcom.utils.process_image import make_image_grid
+        >>> from libcom import FOSScoreModel
+        >>> import cv2
+        >>> import torch
+        >>> task_name = 'fos_score_prediction'
+        >>> MODEL_TYPE = 'FOS_D'
+        >>> background = '../tests/source/background/f80eda2459853824_m09g1w_b2413ec8_11.png'
+        >>> fg_bbox    = [175, 82, 309, 310] # x1,y1,x2,y2
+        >>> foreground = '../tests/source/foreground/f80eda2459853824_m09g1w_b2413ec8_11.png'
+        >>> foreground_mask = '../tests/source/foreground_mask/f80eda2459853824_m09g1w_b2413ec8_11.png'
+        >>> composite_image = '../tests/source/composite/f80eda2459853824_m09g1w_b2413ec8_11.png'
+        >>> net = FOSScoreModel(device=0, model_type=MODEL_TYPE)
+        >>> score = net(background, foreground, fg_bbox, foreground_mask=foreground_mask)
+        >>> grid_img  = make_image_grid([background, foreground, composite_image], text_list=[f'fos_score:{score:.2f}'])
+        >>> cv2.imshow('fos_score_demo', grid_img)
+
+    Expected result:
+
+    .. image:: _static/image/fos_score_result3.jpg
+        :scale: 50 %
+    .. image:: _static/image/fos_score_result2.jpg
+        :scale: 50 %
+
+            
+    """
+    def __init__(self, device=0, model_type='FOS_D', **kwargs):
+        
+        assert model_type in model_set, f'Not implementation for {model_type}'
+        self.model_type = model_type
+        self.option = kwargs
+        self.IMAGE_NET_MEAN = [0.5, 0.5, 0.5]
+        self.IMAGE_NET_STD  = [0.5, 0.5, 0.5]
+        config_file = os.path.join(cur_dir, 'source/config/config_rfosd.yaml')
+        self.cfg    = Config(config_file)
+        weight_path = os.path.join(model_dir, 'pretrained_models', '{}.pth'.format(self.model_type))
+        download_pretrained_model(weight_path)
+        self.device = check_gpu_device(device)
+        self.build_pretrained_model(weight_path)
+        self.build_data_transformer()
+
+    def build_pretrained_model(self, weight_path):
+        """
+        Build pretrained model from path of weight.
+        """
+
+        model = SingleScaleD(False)
+        model.load_state_dict(torch.load(weight_path, map_location='cpu', weights_only=True), strict=True)
+        self.model = model.to(self.device).eval()
+        
+    def build_data_transformer(self):
+        self.image_size = self.cfg.image_size
+        self.transformer = transforms.Compose([
+            transforms.Resize(self.cfg.image_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=self.IMAGE_NET_MEAN, std=self.IMAGE_NET_STD)
+        ])
+    
+    def inputs_preprocess(self, background_image, foreground_image, foreground_mask, bbox):
+        bg  = cv2.cvtColor(read_image_opencv(background_image), cv2.COLOR_BGR2RGB)
+        bg_h, bg_w, _ = bg.shape
+        fg = cv2.cvtColor(read_image_opencv(foreground_image), cv2.COLOR_BGR2RGB)
+        if foreground_mask is not None:
+            fg_mask = cv2.cvtColor(read_image_opencv(foreground_mask), cv2.COLOR_BGR2GRAY)
+            assert fg.shape[:2] == fg_mask.shape
+            fg[np.where(fg_mask != 255)] = 128
+        x1, y1, x2, y2 = bbox
+        rs_fg = cv2.resize(fg.copy(), (x2 - x1, y2 - y1))
+        bg[y1:y2, x1:x2] = rs_fg
+        x1, y1, x2, y2 = self.get_crop_bbox(bbox, bg_h, bg_w)
+        scale_comp = bg.copy()[y1:y2, x1:x2]
+        scale_comp_r = Image.fromarray(scale_comp, mode="RGB")
+        return scale_comp_r
+    
+    def get_crop_bbox(self, bbox, bg_h, bg_w):
+        x1, y1, x2, y2 = bbox
+        ori_tar_w = x2 - x1
+        ori_tar_h = y2 - y1
+        add_w = int(ori_tar_w * (math.sqrt(2) - 1) / 2)
+        add_h = int(ori_tar_h * (math.sqrt(2) - 1) / 2)
+        new_y1 = max(0, y1 - add_h)
+        new_y2 = min(bg_h, y2 + add_h)
+        new_x1 = max(0, x1 - add_w)
+        new_x2 = min(bg_w, x2 + add_w)
+        return new_x1, new_y1, new_x2, new_y2
+    
+    def prepare_input_encoders(self, background_image, foreground_image, bounding_box):
+        background_image = fill_box_with_specified_pixel(background_image, bounding_box, self.cfg.fill_pixel)
+        bg_t = self.transformer(background_image)
+
+        bg_ori_w, bg_ori_h = background_image.size
+        ori_x1, ori_y1 = bounding_box[0:2]
+        ori_x2, ori_y2 = bounding_box[2:4]
+        query_box = torch.tensor([ori_x1 / bg_ori_w, ori_y1 / bg_ori_h, ori_x2 / bg_ori_w, ori_y2 / bg_ori_h]) * self.cfg.image_size
+        query_box = torch.round(query_box)
+        query_box = query_box.float()
+
+        ori_tar_w = ori_x2 - ori_x1
+        ori_tar_h = ori_y2 - ori_y1
+
+        add_w = int(ori_tar_w * (math.sqrt(2) - 1) / 2)
+        add_h = int(ori_tar_h * (math.sqrt(2) - 1) / 2)
+        new_y1 = max(0, ori_y1 - add_h)
+        new_y2 = min(bg_ori_h, ori_y2 + add_h)
+        new_x1 = max(0, ori_x1 - add_w)
+        new_x2 = min(bg_ori_w, ori_x2 + add_w)
+        new_box = torch.tensor([new_x1 / bg_ori_w, new_y1 / bg_ori_h, new_x2 / bg_ori_w, new_y2 / bg_ori_h]) * self.cfg.image_size
+        new_box = torch.round(new_box)
+        new_box = new_box.float()
+
+        pad_fg = padding_to_square(foreground_image.copy(), self.cfg.pad_pixel)
+        pad_fg = Image.fromarray(pad_fg)
+        fg_t = self.transformer(pad_fg)
+
+        sample = {
+            'bg': bg_t.unsqueeze(0),
+            'fg': fg_t.unsqueeze(0),
+            'query_box': query_box.unsqueeze(0),
+            'crop_box': new_box.unsqueeze(0),
+        }
+        return sample
+
+    def prepare_input_disc(self, composite_image):
+        comp = self.transformer(composite_image).unsqueeze(0)
+        return comp
+    
+    def preprocess_image(self, image):
+        image = torch.from_numpy(image).float() / 127.5 - 1 # [-1, 1]
+        image = rearrange(image, "h w c -> 1 c h w")
+        image = image.to(self.device)
+        return image
+    
+    @torch.no_grad()
+    def __call__(
+            self,
+            background_image,
+            foreground_image,
+            bounding_box,
+            foreground_mask=None
+            ):
+        """
+        Predicting the compatibility score between the given background and the given foreground.
+
+        Args:
+            background_image (str | numpy.ndarray): The path to background image or the background image in ndarray form.
+            foreground_image (str | numpy.ndarray): The path to foreground image or the background image in ndarray form.
+            bounding_box (list): The bounding box which indicates the foreground's location in the background. [x1, y1, x2, y2].
+            foreground_mask (str | numpy.ndarray): Mask of foreground image which indicates the foreground object region in the foreground image. default: None.
+        
+        Returns:
+            fos_score (float): Predicted compatibility score between the given background image and the given foreground image.
+
+        """
+        
+        composite_image = self.inputs_preprocess(background_image, foreground_image, foreground_mask, bounding_box)
+        composite_image = self.prepare_input_disc(composite_image).to(self.device)
+        _, score = self.model(composite_image)
+        output = score[-1].item()
+        return output
+    
+def fill_box_with_specified_pixel(bg_im, query_box, fill_value):
+    x1, y1 = query_box[0:2]
+    x2, y2 = query_box[2:4]
+    bg_im = np.array(bg_im)
+    bg_im[y1:y2, x1:x2] = fill_value
+    bg_im = Image.fromarray(bg_im)
+    return bg_im
+
+def padding_to_square(src_img, pad_pixel=255):
+    src_h, src_w = src_img.shape[:2]
+    if src_h == src_w:
+        return src_img
+    if src_w > src_h:
+        pad_w = 0
+        pad_h = src_w - src_h
+    else:
+        pad_w = src_h - src_w
+        pad_h = 0
+
+    pad_y1 = int(pad_h // 2)
+    pad_y2 = int(pad_h - pad_y1)
+    pad_x1 = int(pad_w // 2)
+    pad_x2 = int(pad_w - pad_x1)
+
+    if len(src_img.shape) == 3:
+        pad_im = np.pad(src_img, ((pad_y1, pad_y2), (pad_x1, pad_x2), (0,0)),
+                        'constant', constant_values=pad_pixel)
+    else:
+        pad_im = np.pad(src_img, ((pad_y1, pad_y2), (pad_x1, pad_x2)),
+                        'constant', constant_values=pad_pixel)
+    return pad_im

libcom/fos_score/source/config/__init__.py

@@ -0,0 +1 @@
+from .config import Config

libcom/fos_score/source/config/config.py

@@ -0,0 +1,68 @@
+import os
+import yaml
+
+def load_yaml_params(config_file):
+    with open(config_file, 'r') as yaml_file:
+        params = yaml.full_load(yaml_file.read())
+        return params
+
+class Config:
+
+    def __init__(self, config_file=None):
+        if config_file is None:
+            cur_dir = os.path.dirname(os.path.abspath(__file__))
+            self.config_file = os.path.join(cur_dir, 'config_sfosd.yaml')
+        else:
+            self.config_file = config_file
+        self.refresh_params()
+
+    def refresh_params(self, config_file=None):
+        if config_file is not None:
+            self.config_file = config_file
+        self.load_params_from_yaml()
+
+    def print_yaml_params(self):
+        params = load_yaml_params(self.config_file)
+        print('*'*30, os.path.basename(self.config_file), '*'*30)
+        for k, v in params.items():
+            print(k, v)
+        print('*'*30, os.path.basename(self.config_file), '*'*30)
+
+    def load_params_from_yaml(self):
+        # add parameters from yaml file
+        names = self.__dict__
+        params = load_yaml_params(self.config_file)
+        for k, v in params.items():
+            names[k] = v
+    
+    def generate_path(self):
+        prefix = '{}class-bg{}-{}'.format(self.num_classes, self.encoder_feature, self.distill_type)
+        if self.encoder_classify:
+            prefix += '-encodercls'
+        exp_root = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                                'experiments', 'student')
+        os.makedirs(exp_root, exist_ok=True)
+        exp_name = prefix
+        exp_path = os.path.join(exp_root, prefix)
+        while os.path.exists(exp_path):
+            index = os.path.basename(exp_path).split(prefix)[-1].split('repeat')[-1]
+            try:
+                index = int(index) + 1
+            except:
+                index = 1
+            exp_name = prefix + ('_repeat{}'.format(index))
+            exp_path = os.path.join(exp_root, exp_name)
+        # print('Experiment name {} \n'.format(os.path.basename(exp_path)))
+        self.exp_name = exp_name
+        self.exp_path = exp_path
+        self.checkpoint_dir = os.path.join(exp_path, 'checkpoints')
+        self.log_dir  = os.path.join(exp_path,  'logs')
+        self.code_dir = os.path.join(exp_path, 'code')
+
+    def create_path(self):
+        # self.generate_path()
+        print('Create experiment directory: ', self.exp_path)
+        os.makedirs(self.exp_path, exist_ok=True)
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+        os.makedirs(self.code_dir, exist_ok=True)
+        os.makedirs(self.log_dir, exist_ok=True)

libcom/fos_score/source/config/config.sfosd.yaml

@@ -0,0 +1,43 @@
+config_name: sfosd
+E_lr: 1.0e-05
+backbone: vgg19
+batch_size: 10
+classify_weight: 1
+display_freq: 50
+distill_type: roiconcat
+distill_weight: 1
+encoder_classify: false
+encoder_feature: gap
+eval_bs: 64
+eval_testset1: true
+eval_testset2: false
+feature_dim: 512
+fill_pixel: 128
+gpu_id: 2
+image_size: 224
+local_rank: 1
+max_epoch: 100
+num_classes: 32
+num_negative: 10
+num_positive: 1
+num_workers: 8
+pad_pixel: 255
+save_freq: 10
+padding_max: 0
+img_aug: false
+teacher_model: teacher_sfosd.pth
+teacher_path: checkpoints/teacher
+test1_metrics:
+- Recall@1
+- Recall@5
+- Recall@10
+- Recall@20
+test2_metrics:
+- mAP
+- mAP-20
+- Precision@1
+- Precision@5
+- Precision@10
+- Precision@20
+test_freq: 10
+triplet_margin: 0.1

libcom/fos_score/source/config/config_rfosd.yaml

@@ -0,0 +1,43 @@
+config_name: rfosd
+E_lr: 1.0e-05
+backbone: vgg19
+batch_size: 6
+classify_weight: 1
+display_freq: 50
+distill_type: roiconcat
+distill_weight: 1
+encoder_classify: false
+encoder_feature: gap
+eval_bs: 64
+eval_testset1: false
+eval_testset2: true
+feature_dim: 512
+fill_pixel: 128
+gpu_id: 2
+image_size: 224
+local_rank: 1
+max_epoch: 100
+num_classes: 32
+num_negative: 10
+num_positive: 5
+num_workers: 8
+pad_pixel: 255
+save_freq: 10
+padding_max: 0.3
+img_aug: true
+teacher_model: teacher_rfosd.pth
+teacher_path: checkpoints/teacher
+test1_metrics:
+- Recall@1
+- Recall@5
+- Recall@10
+- Recall@20
+test2_metrics:
+- mAP
+- mAP-20
+- Precision@1
+- Precision@5
+- Precision@10
+- Precision@20
+test_freq: 10
+triplet_margin: 0.1

libcom/fos_score/source/loss_function.py

@@ -0,0 +1,118 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+class TripletLoss(nn.Module):
+    def __init__(self, margin=0.3):
+        super(TripletLoss, self).__init__()
+        self.margin = margin
+
+    def forward(self, anchor, positive, negative, need_norm=True):
+        '''
+        :param anchor:   num_anchor, emb_dim
+        :param positive: num_anchor, num_pos, emb_dim
+        :param negative: num_anchor, num_neg, emb_dim
+        :return:
+        '''
+        if need_norm:
+            anchor   = F.normalize(anchor,   2, -1)
+            positive = F.normalize(positive, 2, -1)
+            negative = F.normalize(negative, 2, -1)
+        pos_cos  = torch.einsum('b d, b n d -> b n', anchor, positive)
+        # print('cosine similarity', anchor.shape, positive.shape, pos_cos.shape, pos_cos.min(), pos_cos.max())
+        pos_dist = 1 - pos_cos
+        neg_cos  = torch.einsum('b d, b n d -> b n', anchor, negative)
+        neg_dist = 1 - neg_cos
+        loss = F.relu(pos_dist.unsqueeze(2) - neg_dist.unsqueeze(1) + self.margin)
+        return loss.mean()
+
+class TripletHardLoss(nn.Module):
+    """Triplet loss with hard positive/negative mining.
+
+    Reference:
+        Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.
+
+    Imported from `<https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py>`_.
+
+    Args:
+        margin (float, optional): margin for triplet. Default is 0.3.
+    """
+
+    def __init__(self, margin=0.3):
+        super(TripletHardLoss, self).__init__()
+        self.margin = margin
+
+    def forward(self, anchor, positive, negative, need_norm=True):
+        '''
+        :param anchor: num_anchor, emb_dim
+        :param positive: num_anchor, num_pos, emb_dim
+        :param negative: num_anchor, num_neg, emb_dim
+        :return:
+        '''
+        if need_norm:
+            anchor   = F.normalize(anchor,   2, -1)
+            positive = F.normalize(positive, 2, -1)
+            negative = F.normalize(negative, 2, -1)
+        pos_cos = torch.einsum('b d, b n d -> b n', anchor, positive)
+        pos_dist = 1 - pos_cos
+        neg_cos = torch.einsum('b d, b n d -> b n', anchor, negative)
+        neg_dist = 1 - neg_cos
+        loss = F.relu(pos_dist.max(-1)[0] - neg_dist.min(-1)[0] + self.margin)
+        return loss.mean()
+
+# Defines the GAN loss which uses either LSGAN or the regular GAN.
+# When LSGAN is used, it is basically same as MSELoss,
+# but it abstracts away the need to create the target label tensor
+# that has the same size as the input
+class GANLoss(nn.Module):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0):
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        if use_lsgan:
+            self.loss = nn.MSELoss()
+        else:
+            self.loss = nn.BCELoss()
+
+    def get_target_tensor(self, input, target_is_real):
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(input)
+
+    def __call__(self, input, target_is_real):
+        target_tensor = self.get_target_tensor(input, target_is_real)
+        return self.loss(input, target_tensor)
+
+def kl_divergence_loss(log_var, mu):
+    '''
+    the KL-Divergence.
+    KL-Divergence = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
+    '''
+    KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
+    return KLD
+
+
+def custom_kl_loss(p_mu, p_var, q_mu, q_var):
+    p_distribution = torch.distributions.MultivariateNormal(p_mu, torch.diag_embed(p_var))
+    q_distribution = torch.distributions.MultivariateNormal(q_mu, torch.diag_embed(q_var))
+    p_q_kl = torch.distributions.kl_divergence(p_distribution, q_distribution).mean()
+    return p_q_kl
+    # b,d = p_mu.shape
+    # normal_distribution = torch.distributions.MultivariateNormal(torch.zeros(b, d),
+    #                                                              torch.diag_embed(torch.ones(b,d)).cuda())
+    # p_normal_kl = torch.distributions.kl_divergence(p_distribution, normal_distribution).mean()
+
+def mask_iou_loss(inputs, targets):
+    # flatten label and prediction tensors
+    inputs = inputs.flatten(1).float()
+    targets = targets.flatten(1).float()
+
+    # intersection is equivalent to True Positive count
+    # union is the mutually inclusive area of all labels & predictions
+    intersection = (inputs * targets).sum(1)
+    total = (inputs + targets).sum(1)
+    union = total - intersection
+    IoU = intersection / (union + 1e-12)
+    return 1 - IoU.mean()