libcom 0.1.8__py3-none-any.whl

libcom/fos_score/source/networks.py

@@ -0,0 +1,355 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models
+from einops import rearrange
+import copy
+from .loss_function import TripletLoss
+from torchvision.ops import RoIAlign, roi_align
+
+class SingleScaleD(nn.Module):
+    def __init__(self, loadweights=False):
+        super().__init__()
+        model = models.vgg19(pretrained = loadweights)
+        self.backbone = model.features
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(in_features=512, out_features=1, bias=True)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        f3 = self.backbone[:19](x)
+        f4 = self.backbone[19:28](f3)
+        f5 = self.backbone[28:](f4)
+        fv = self.gap(f5).flatten(1)
+        out = self.sigmoid(self.fc(fv)).flatten()
+        return [f5, fv], out
+
+def conv_relu_x2(in_ch, out_ch):
+    convs = nn.Sequential(
+        nn.Conv2d(in_ch, out_ch, 3, 1, 1),
+        nn.ReLU(),
+        nn.Conv2d(out_ch, out_ch, 3, 1, 1),
+        nn.ReLU())
+    return convs
+
+def binary_classifier(in_ch):
+    classifier = nn.Sequential(
+        nn.AdaptiveAvgPool2d(1),
+        nn.Conv2d(in_ch, 1, 1, 1),
+        nn.Sigmoid())
+    return classifier
+
+def get_distillation_module(distill_type, in_ch, out_ch, alignsize, downsample):
+    if distill_type == 'roiconcat':
+        return RoIConcatDistillation(in_ch, out_ch, alignsize, downsample)
+    elif distill_type == 'roicropresize':
+        return CropResizeDistillation(in_ch, out_ch, alignsize, downsample)
+    elif distill_type == 'gloablconcat':
+        return GlobalConcatDistillation(in_ch, out_ch)
+    elif distill_type == 'roivectorconcat':
+        return RoIVectorConcatDistillation(in_ch, out_ch, downsample)
+    elif distill_type in ['gapconcat', 'onlyclassify']:
+        return VectorConcatDistillation(in_ch, out_ch)
+    else:
+        return None
+
+class CropResizeDistillation(nn.Module):
+    def __init__(self, in_ch, out_ch, alignsize, downsample):
+        super().__init__()
+        # self.RoIAlign = RoIAlignAvg(alignsize, alignsize, 1.0 / 2 ** downsample)
+        self.RoIAlign = RoIAlign(alignsize, 1.0 / 2 ** downsample, -1)
+        self.convs = conv_relu_x2(in_ch, out_ch)
+        # self.scale = 1.0 / 2 ** downsample
+        self.alignsize = alignsize
+        self.classifier = binary_classifier(out_ch)
+
+    def forward(self, bg_feat, fg_feat, q_box, c_box, num_bg):
+        '''
+        :param bg_feat: b,c,h,w
+        :param fg_feat: b,c,h,w
+        :param q_box: query box, b,4
+        :param c_box: cropped bounding box, b,4
+        :return:
+        '''
+        B = bg_feat.shape[0]
+        index = torch.arange(B).view(-1, 1).to(q_box.device)
+        idx_cbox = torch.cat((index, c_box), dim=-1).contiguous()
+        # obtain the feature of cropped background region
+        roi_feat = self.RoIAlign(bg_feat, idx_cbox) # b,c,roi_size,roi_size
+        # compute the target bounding box for foreground feature
+        tar_box = torch.zeros_like(q_box)
+        crop_w  = c_box[:,2] - c_box[:,0]
+        crop_h  = c_box[:,3] - c_box[:,1]
+        tar_box[:, 0] = torch.floor((q_box[:,0] - c_box[:,0]) / crop_w * self.alignsize)
+        tar_box[:, 1] = torch.floor((q_box[:,1] - c_box[:,1]) / crop_h * self.alignsize)
+        tar_box[:, 2] = torch.ceil((q_box[:,2] - c_box[:,0])  / crop_w * self.alignsize)
+        tar_box[:, 3] = torch.ceil((q_box[:,3] - c_box[:,1])  / crop_h * self.alignsize)
+        tar_box = tar_box.int()
+        assert tar_box.min() >= 0 and tar_box.max() <= self.alignsize, \
+            f'tar box coordinates max={tar_box.max()}, min={tar_box.min()}'
+        # as the target size of foreground feature varies with background image,
+        # we separately process different backgrounds.
+
+        roi_feat = rearrange(roi_feat, '(b n) c h w -> b n c h w', b=num_bg)
+        fg_feat = rearrange(fg_feat, '(b n) c h w -> b n c h w', b=num_bg)
+        tar_box = rearrange(tar_box, '(b n) d -> b n d', b=num_bg)
+
+        feat_collect = []
+        for i in range(num_bg):
+            comp_feat = roi_feat[i].clone()
+            per_fg_feat = fg_feat[i]
+            x1,y1,x2,y2 = tar_box[i, 0]
+            rs_w, rs_h = x2 - x1, y2 - y1
+            per_fg_feat = F.interpolate(per_fg_feat, (rs_h, rs_w))
+            comp_feat[:, :, y1:y2, x1:x2] = per_fg_feat
+            feat_collect.append(comp_feat)
+        feat_collect = torch.stack(feat_collect, dim=0)
+        feat_collect = rearrange(feat_collect, 'b n c h w -> (b n) c h w')
+        out_feat = self.convs(feat_collect)
+        score = self.classifier(out_feat)
+        return out_feat, score
+
+class GlobalConcatDistillation(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.convs = conv_relu_x2(in_ch*2, out_ch)
+        self.classifier = binary_classifier(out_ch)
+
+    def forward(self, bg_feat, fg_feat, q_box, c_box, num_bg):
+        '''
+        :param bg_feat: b,c,h,w
+        :param fg_feat: b,c,h,w
+        :param q_box: query box, b,4
+        :param c_box: cropped bounding box, b,4
+        :return:
+        '''
+        concat_feat = torch.cat([bg_feat, fg_feat], dim=1)
+        out_feat = self.convs(concat_feat)
+        score = self.classifier(out_feat)
+        return out_feat, score
+
+class RoIConcatDistillation(nn.Module):
+    def __init__(self, in_ch, out_ch, alignsize, downsample):
+        super().__init__()
+        self.RoIAlign = RoIAlign(alignsize, 1.0 / 2 ** downsample, -1)
+        self.convs = conv_relu_x2(in_ch*2, out_ch)
+        self.alignsize = alignsize
+        self.classifier = binary_classifier(out_ch)
+
+    def forward(self, bg_feat, fg_feat, q_box, c_box, num_bg):
+        '''
+        :param bg_feat: b,c,h,w
+        :param fg_feat: b,c,h,w
+        :param q_box: query box, b,4
+        :param c_box: cropped bounding box, b,4
+        :return:
+        '''
+        B = bg_feat.shape[0]
+        index = torch.arange(B).view(-1, 1).to(q_box.device)
+        idx_cbox = torch.cat((index, c_box), dim=-1).contiguous()
+        # obtain the feature of cropped background region
+        roi_feat = self.RoIAlign(bg_feat, idx_cbox)  # b,c,h,w
+        fg_feat_rs = F.interpolate(fg_feat, (self.alignsize, self.alignsize))
+        concat_feat = torch.cat([roi_feat, fg_feat_rs], dim=1)
+        out_feat = self.convs(concat_feat)
+        score = self.classifier(out_feat)
+        return out_feat, score
+
+class VectorConcatDistillation(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(2 * in_ch, out_ch),
+            nn.ReLU(),
+            nn.Linear(out_ch, out_ch),
+            nn.ReLU())
+        self.gap  = nn.AdaptiveAvgPool2d(1)
+        self.classifier = binary_classifier(out_ch)
+
+    def forward(self, bg_feat, fg_feat, q_box, c_box, num_bg):
+        '''
+        :param bg_feat: b,c,h,w
+        :param fg_feat: b,c,h,w
+        :param q_box: query box, b,4
+        :param c_box: cropped bounding box, b,4
+        :return:
+        '''
+        bg_feat  = self.gap(bg_feat)
+        fg_feat  = self.gap(fg_feat)
+        bg_vec = bg_feat.flatten(1)
+        fg_vec = fg_feat.flatten(1)
+        out_vec = self.fc(torch.concat([bg_vec, fg_vec], dim=1))
+        out_feat = out_vec.unsqueeze(-1).unsqueeze(-1)
+        score = self.classifier(out_feat)
+        return out_feat, score
+
+class RoIVectorConcatDistillation(nn.Module):
+    def __init__(self, in_ch, out_ch, downsample):
+        super().__init__()
+        self.RoIAlign = RoIAlign(1, 1.0 / 2 ** downsample, -1)
+        self.fc = nn.Sequential(
+            nn.Linear(2*in_ch, out_ch),
+            nn.ReLU(),
+            nn.Linear(out_ch, out_ch),
+            nn.ReLU())
+        self.gap  = nn.AdaptiveAvgPool2d(1)
+        self.classifier = binary_classifier(out_ch)
+
+    def forward(self, bg_feat, fg_feat, q_box, c_box, num_bg):
+        '''
+        :param bg_feat: b,c,h,w
+        :param fg_feat: b,c,h,w
+        :param q_box: query box, b,4
+        :param c_box: cropped bounding box, b,4
+        :return:
+        '''
+        B = bg_feat.shape[0]
+        index = torch.arange(B).view(-1, 1).to(q_box.device)
+        idx_cbox = torch.cat((index, c_box), dim=-1).contiguous()
+        # obtain the feature of cropped background region
+        roi_feat  = self.RoIAlign(bg_feat, idx_cbox)  # b,c,h,w
+        bg_vec = roi_feat.flatten(1)
+        fg_vec = self.gap(fg_feat).flatten(1)
+        out_vec = self.fc(torch.concat([bg_vec, fg_vec], dim=1))
+        out_feat = out_vec.unsqueeze(-1).unsqueeze(-1)
+        score = self.classifier(out_feat)
+        return out_feat, score
+
+class StudentModel(nn.Module):
+    def __init__(self, cfg, loadweights=True):
+        super().__init__()
+        assert cfg.backbone == 'vgg19', cfg.backbone
+        model = models.vgg19(pretrained=loadweights)
+        backbone = model.features
+        self.base_dim = 512
+        self.bg_encoder = nn.ModuleList([
+            nn.Sequential(backbone[:23]),
+            nn.Sequential(backbone[23:30]),
+            nn.Sequential(backbone[30:])
+        ])
+        backbone_copy = copy.deepcopy(backbone)
+        self.fg_encoder = nn.ModuleList([
+            nn.Sequential(backbone_copy[:23]),
+            nn.Sequential(backbone_copy[23:30]),
+            nn.Sequential(backbone_copy[30:])
+        ])
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.out_dim = cfg.feature_dim
+        out_ch = 512
+        # branch of feature distillation
+        '''
+        encoder_feature: gap or roi
+        encoder_classify: true or false
+        distill_type: roiconcat, roicropresize, gloablconcat, roivectorconcat, gapconcat, onlyclassify, none
+        '''
+        self.encoder_feat = cfg.encoder_feature
+        self.encoder_classify = cfg.encoder_classify
+        self.distill_type = cfg.distill_type
+
+        if self.encoder_classify:
+            self.auxiliary_classifier = nn.Sequential(
+                nn.Linear(out_ch * 2, out_ch),
+                nn.ReLU(),
+                nn.Linear(out_ch, 1, bias=False),
+                nn.Sigmoid())
+
+        alignsize, downsample = 7, 5
+        self.distill_module = get_distillation_module(self.distill_type,
+                                                      out_ch, out_ch,
+                                                      alignsize, downsample)
+        self.triplet_loss = TripletLoss(cfg.triplet_margin)
+
+    def forward_bg_encoder(self, bg_im, crop_box):
+        bg_f1 = self.bg_encoder[0](bg_im)
+        bg_f2 = self.bg_encoder[1](bg_f1)
+        bg_f3 = self.bg_encoder[2](bg_f2)
+        if self.encoder_feat == 'gap':
+            bg_emb = self.gap(bg_f3).flatten(1)
+        else:
+            bg_emb = self.extract_bg_roi(bg_f3, crop_box)
+        return bg_emb
+
+    def forward_fg_encoder(self, fg_im):
+        fg_f1 = self.fg_encoder[0](fg_im)
+        fg_f2 = self.fg_encoder[1](fg_f1)
+        fg_f3 = self.fg_encoder[2](fg_f2)
+        fg_emb = self.gap(fg_f3).flatten(1)
+        return fg_emb
+
+    def extract_bg_roi(self, bg_fm, crop_box):
+        # extract background RoI feature
+        B = bg_fm.shape[0]
+        index = torch.arange(B).view(-1, 1).to(crop_box.device)
+        idx_cbox = torch.cat((index, crop_box), dim=-1).contiguous()
+        bg_roi = roi_align(bg_fm, idx_cbox, output_size=1,
+                           spatial_scale=1. / 2 ** 5)
+        bg_emb = bg_roi.flatten(1)
+        return bg_emb
+
+    def forward(self, bg_im, fg_im, query_box, crop_box, num_bg=1):
+        bg_f1 = self.bg_encoder[0](bg_im)
+        bg_f2 = self.bg_encoder[1](bg_f1)
+        bg_f3 = self.bg_encoder[2](bg_f2)
+        if self.encoder_feat == 'gap':
+            self.bg_emb = self.gap(bg_f3).flatten(1)
+        else:
+            self.bg_emb = self.extract_bg_roi(bg_f3, crop_box)
+
+        fg_f1 = self.fg_encoder[0](fg_im)
+        fg_f2 = self.fg_encoder[1](fg_f1)
+        fg_f3 = self.fg_encoder[2](fg_f2)
+        self.fg_emb = self.gap(fg_f3).flatten(1)
+
+        if self.encoder_classify:
+            self.aux_score = self.auxiliary_classifier(torch.cat([self.bg_emb, self.fg_emb], dim=1))
+
+        if self.distill_type == 'none':
+            self.distill_feat, self.score = None, torch.empty([])
+        else:
+            if num_bg == None:
+                num_bg = bg_im.shape[0]
+            self.distill_feat, self.score = self.distill_module(bg_f3, fg_f3, query_box, crop_box, num_bg)
+            self.score = self.score.flatten()
+        return self.bg_emb, self.fg_emb, self.score
+
+
+    def calcu_classify_loss(self, gt_label):
+        gt_label = gt_label.reshape(-1,1)
+        loss = 0.
+        if self.distill_type != 'none':
+            pr_score = self.score.reshape(-1,1)
+            loss += F.binary_cross_entropy(pr_score, gt_label)
+        if self.encoder_classify:
+            pr_score2 = self.aux_score.reshape(-1,1)
+            loss += F.binary_cross_entropy(pr_score2, gt_label)
+            if self.distill_type != 'none':
+                loss /= 2
+        return loss
+
+    def calcu_triplet_loss(self, gt_label, num_bg):
+        pos_index = (gt_label > 0.5)
+        neg_index = (gt_label < 0.5)
+        pos_emb, neg_emb = self.fg_emb[pos_index], self.fg_emb[neg_index]
+        bg_emb = rearrange(self.bg_emb, '(b n) c -> b n c', b=num_bg)
+        bg_emb = bg_emb[:, 0]
+        pos_emb = rearrange(pos_emb, '(b n) d -> b n d', b=num_bg)
+        neg_emb = rearrange(neg_emb, '(b n) d -> b n d', b=num_bg)
+        loss = self.triplet_loss(bg_emb, pos_emb, neg_emb)
+        return loss
+
+    def calcu_distillation_loss(self, feat_list, loss_type='l1'):
+        if self.distill_type in ['none', 'onlyclassify']:
+            return 0.
+        if loss_type.lower() == 'l1':
+            loss_func = F.l1_loss
+        elif loss_type.lower() == 'l2':
+            loss_func = F.mse_loss
+        elif loss_type.lower() == 'smoothl1':
+            loss_func = F.smooth_l1_loss
+        else:
+            raise Exception(f'undefined loss type {loss_type}')
+        gt_featmap, gt_featvec = feat_list
+        if gt_featmap.shape[2:] != self.distill_feat.shape[2:]:
+            gt_featmap = F.interpolate(gt_featmap, self.distill_feat.shape[2:])
+        kd_loss = loss_func(self.distill_feat, gt_featmap)
+        return kd_loss

libcom/harmony_score/__init__.py

@@ -0,0 +1,3 @@
+from .harmony_score_prediction import HarmonyScoreModel
+
+__all__ = ['HarmonyScoreModel']

libcom/harmony_score/harmony_score_prediction.py

@@ -0,0 +1,119 @@
+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.harmony_score.source.bargainnet import StyleEncoder
+import torch 
+import os
+import torchvision.transforms as transforms
+import math
+
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['BargainNet'] 
+
+class HarmonyScoreModel:
+    """
+    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 import HarmonyScoreModel
+        >>> from libcom.utils.process_image import make_image_grid
+        >>> import cv2
+        >>> net = HarmonyScoreModel(device=0, model_type='BargainNet')
+        >>> test_dir   = '../tests/harmony_score_prediction/'
+        >>> img_names  = ['vaulted-cellar-247391_inharm.jpg', 'ameland-5651866_harm.jpg']
+        >>> vis_list,scores = [], []
+        >>> for img_name in img_names:
+        >>>     comp_img  = test_dir + 'composite/' + img_name
+        >>>     comp_mask = test_dir + 'composite_mask/' + img_name
+        >>>     score     = net(comp_img, comp_mask)
+        >>>     vis_list += [comp_img, comp_mask]
+        >>>     scores.append(score)
+        >>> grid_img  = make_image_grid(vis_list, text_list=[f'harmony_score:{scores[0]:.2f}', 'composite-mask', f'harmony_score:{scores[1]:.2f}', 'composite-mask'])
+        >>> cv2.imwrite('../docs/_static/image/harmonyscore_result1.jpg', grid_img)
+ 
+    Expected result:
+
+    .. image:: _static/image/harmonyscore_result1.jpg
+        :scale: 38 %
+
+    """
+    def __init__(self, device=0, model_type='BargainNet', **kwargs):
+        assert model_type in model_set, f'Not implementation for {model_type}'
+        self.model_type = model_type
+        self.option = kwargs
+        weight_path = os.path.join(model_dir, 'pretrained_models', 'BargainNet.pth')
+        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):
+        model = StyleEncoder(style_dim=16)
+        model.load_state_dict(torch.load(weight_path, map_location='cpu', weights_only=True))
+        self.model = model.to(self.device).eval()
+        
+    def build_data_transformer(self):
+        self.image_size = 256
+        self.transform = transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ])
+        self.mask_transform = transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+        ])
+    
+    def inputs_preprocess(self, composite_image, composite_mask):
+        image = read_image_opencv(composite_image)
+        mask  = read_mask_opencv(composite_mask)
+        fg_mask = mask.astype(np.float32) / 255.
+        bg_mask = 1 - fg_mask
+        fg_mask = self.mask_transform(Image.fromarray(fg_mask))
+        fg_mask = fg_mask.unsqueeze(0).to(self.device)
+        bg_mask = self.mask_transform(Image.fromarray(bg_mask))
+        bg_mask = bg_mask.unsqueeze(0).to(self.device)
+
+        image   = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image   = self.transform(Image.fromarray(image))
+        image   = image.unsqueeze(0).to(self.device)
+        return image, bg_mask, fg_mask
+    
+    def outputs_postprocess(self, bg_style, fg_style):
+        eucl_dist = self.Euclidean_distance(bg_style, fg_style)
+        # convert distance to harmony level which lies in 0 and 1
+        harm_level = math.exp(-0.04212 * eucl_dist)
+        return harm_level
+    
+    def Euclidean_distance(self, vec1, vec2):
+        vec1 = vec1.cpu().numpy()
+        vec2 = vec2.cpu().numpy()
+        dist = np.sqrt(np.sum((vec1 - vec2)**2))
+        return dist
+    
+    @torch.no_grad()
+    def __call__(self, composite_image, composite_mask):
+        """
+        Predicting the compatibility score between background and foreground in the given composite image.
+
+        Args:
+            composite_image (str | numpy.ndarray): The path to composite image or the compposite image in ndarray form.
+            composite_mask (str | numpy.ndarray): Mask of composite image which indicates the foreground object region in the composite image.
+        
+        Returns:
+            harmony_score (float): Predicted harmony score within [0,1] between background region and foreground region of the given composite image. Larger harmony score implies more harmonious composite image.
+
+        """
+        im, bg_mask, fg_mask = self.inputs_preprocess(composite_image, composite_mask)
+        bg_style = self.model(im, bg_mask)
+        fg_style = self.model(im, fg_mask)
+        preds    = self.outputs_postprocess(bg_style, fg_style)
+        return preds