libcom 0.1.8__py3-none-any.whl

libcom/__init__.py

@@ -0,0 +1,23 @@
+from .color_transfer import *
+from .naive_composition import *
+from .opa_score import *
+from .harmony_score import *
+from .inharmonious_region_localization import *
+from .image_harmonization import *
+from .painterly_image_harmonization import *
+from .fopa_heat_map import *
+from .fos_score import *
+from .kontext_blending_harmonization import *
+from .shadow_generation import *
+from .reflection_generation import *
+from .os_insert import *
+
+__all__ = [
+    'color_transfer', 'get_composite_image', 'OPAScoreModel', 
+    'HarmonyScoreModel', 'InharmoniousLocalizationModel',
+    'ImageHarmonizationModel', 'PainterlyHarmonizationModel',
+    'FOPAHeatMapModel', 'FOSScoreModel',
+    'KontextBlendingHarmonizationModel',
+    'ShadowGenerationModel', 'ReflectionGenerationModel',
+    'OSInsertModel'
+]

libcom/color_transfer/__init__.py

@@ -0,0 +1,3 @@
+from .reinhard import color_transfer
+
+__all__ = ['color_transfer']

libcom/color_transfer/reinhard.py

@@ -0,0 +1,57 @@
+import numpy as np
+import os
+import cv2
+from libcom.utils.process_image import *
+
+def color_transfer(composite_image, composite_mask):
+    """
+    Generate composite image through copy-and-paste.
+
+    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:
+        transfered image (numpy.ndarray): Transfered image with the same resolution as input image.
+    
+    Examples:
+        >>> from libcom import color_transfer
+        >>> from libcom.utils.process_image import make_image_grid
+        >>> import cv2
+        >>> comp_img1  = '../tests/source/composite/1.jpg'
+        >>> comp_mask1 = '../tests/source/composite_mask/1.png'
+        >>> trans_img1 = color_transfer(comp_img1, comp_mask1)
+        >>> comp_img2  = '../tests/source/composite/8.jpg'
+        >>> comp_mask2 = '../tests/source/composite_mask/8.png'
+        >>> trans_img2 = color_transfer(comp_img2, comp_mask2)
+        >>> # visualization results
+        >>> grid_img  = make_image_grid([comp_img1, comp_mask1, trans_img1, 
+        >>>                             comp_img2, comp_mask2, trans_img2], cols=3)
+        >>> cv2.imwrite('../docs/_static/image/colortransfer_result1.jpg', grid_img)
+
+    Expected result:
+
+    .. image:: _static/image/colortransfer_result1.jpg
+        :scale: 50 %
+        
+    """
+    comp_img  = read_image_opencv(composite_image)
+    comp_mask = read_mask_opencv(composite_mask)
+    if comp_mask.shape[:2] != comp_img.shape[:2]:
+        comp_mask = cv2.resize(comp_mask, (comp_img.shape[1], comp_img.shape[0]))
+    trans_comp = reinhard(comp_img, comp_mask)
+    return trans_comp
+        
+def reinhard(comp_img, comp_mask):
+    # reinhard is a classical color transfer algorithm
+    # paper: https://www.cs.tau.ac.il/~turkel/imagepapers/ColorTransfer.pdf
+    comp_mask = np.where(comp_mask > 127, 255, 0).astype(np.uint8)
+    comp_lab  = cv2.cvtColor(comp_img, cv2.COLOR_BGR2Lab)
+    bg_mean, bg_std = cv2.meanStdDev(comp_lab, mask=255 - comp_mask)
+    fg_mean, fg_std = cv2.meanStdDev(comp_lab, mask=comp_mask)
+    ratio = (bg_std / fg_std).reshape(-1)
+    offset = (bg_mean - fg_mean * bg_std / fg_std).reshape(-1)
+    trans_lab = cv2.convertScaleAbs(comp_lab * ratio + offset)
+    trans_img = cv2.cvtColor(trans_lab, cv2.COLOR_Lab2BGR)
+    trans_comp = np.where(comp_mask[:,:,np.newaxis] > 127, trans_img, comp_img)
+    return trans_comp

libcom/fopa_heat_map/__init__.py

@@ -0,0 +1,4 @@
+# change to your lib name
+from .fopa_heat_map import FOPAHeatMapModel
+
+__all__ = ['FOPAHeatMapModel']

libcom/fopa_heat_map/fopa_heat_map.py

@@ -0,0 +1,212 @@
+import csv
+import random
+import shutil
+import torch
+from libcom.utils.model_download import download_pretrained_model
+from libcom.utils.process_image import *
+from libcom.utils.environment import *
+import torch 
+import os
+import torchvision.transforms as transforms
+
+import numpy as np
+from PIL import Image
+import libcom.fopa_heat_map.source.network.ObPlaNet_simple as network
+from libcom.fopa_heat_map.source.data.OBdataset import make_composite_PIL
+from libcom.fopa_heat_map.source.prepare_multi_fg_scales import prepare_multi_fg_scales
+from libcom.fopa_heat_map.source.data.all_transforms import Compose, JointResize
+
+cur_dir = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['fopa'] 
+
+
+class FOPAHeatMapModel:
+    """
+    Generate a heatmap for a pair of scaled foreground and background.
+
+    Args:
+        device (str | torch.device): gpu id
+        model_type (str): predefined model type
+
+    Examples:
+        >>> test_set = get_test_list_fopa_heatmap()
+        >>> result_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results', task_name)
+        >>> if os.path.exists(result_dir):
+        >>>     shutil.rmtree(result_dir)
+        >>> os.makedirs(result_dir, exist_ok=True)
+        >>> os.makedirs(os.path.join(result_dir, 'grid'), exist_ok=True)
+        >>> print(f'begin testing {task_name}...')
+        >>> net = FOPAHeatMapModel(device=0)
+        >>> for pair in test_set[:1]:
+        >>>     fg_img, fg_mask, bg_img = pair['foreground'], pair['foreground_mask'], pair['background']
+        >>>     bboxes, heatmaps = net(fg_img, fg_mask, bg_img, cache_dir=os.path.join(result_dir, 'cache'), heatmap_dir=os.path.join(result_dir, 'heatmap'))
+        >>>     img_name  = os.path.basename(bg_img).replace('.png', '.jpg')
+        >>>     grid_img  = make_image_grid([bg_img, fg_img, heatmaps[0]])
+        >>>     res_path  = os.path.join(result_dir, 'grid', img_name)
+        >>>     cv2.imwrite(res_path, grid_img)
+        >>>     print('save result to ', res_path)
+        >>> print(f'end testing {task_name}!') 
+
+    Expected result:
+
+    .. image:: _static/image/fopa_heatmap_FOPA_result1.png
+        :scale: 72 %
+
+    .. image:: _static/image/fopa_heatmap_FOPA_result3.png
+        :scale: 52 %
+        
+    """
+    def __init__(self, device=0, model_type='fopa', **kwargs):
+        
+        assert model_type in model_set, f'Not implementation for {model_type}'
+        self.model_type = model_type
+        self.option = kwargs
+        self.device = check_gpu_device(device)
+        fopa_weight = os.path.join(model_dir, 'pretrained_models', 'FOPA.pth')
+        download_pretrained_model(fopa_weight)
+        sopa_weight = os.path.join(model_dir, 'pretrained_models', 'SOPA.pth')
+        download_pretrained_model(sopa_weight)
+        self.build_pretrained_model(sopa_weight, fopa_weight)
+        self.build_data_transformer()
+        
+    def build_pretrained_model(self, sopa_weight, fopa_weight):
+        model = getattr(network, "ObPlaNet_resnet18")(pretrained=False, weight_path=sopa_weight).to(self.device)
+        model.load_state_dict(torch.load(fopa_weight, map_location='cpu'))
+        self.model = model.to(self.device).eval()
+        
+    def build_data_transformer(self):
+        self.triple_transform = Compose([JointResize(256)])
+        self.image_transform = transforms.Compose(
+            [transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),]
+        )
+        self.mask_transform = transforms.ToTensor()
+    
+    def generate_heatmap(self, background_image, csv_path, scaled_fg_dir, scaled_mask_dir, heatmap_dir):
+        os.makedirs(heatmap_dir, exist_ok=True)
+        heatmap_list = []
+        with open(csv_path, mode='r', newline='') as csv_file:
+            csv_reader = csv.DictReader(csv_file)
+            for row in csv_reader:
+                fg_name   = '{}_{}_{}_{}.jpg'.format(row["fg_name"].split(".")[0],row["bg_name"].split(".")[0],int(row["newWidth"]),int(row["newHeight"]))
+                mask_name = '{}_{}_{}_{}.jpg'.format(row["fg_name"].split(".")[0],row["bg_name"].split(".")[0],int(row["newWidth"]),int(row["newHeight"]))
+                scale     = row['scale']
+                save_name = fg_name.split(".")[0] + '_' + str(scale) + '.jpg'
+                
+                bg_img    = read_image_pil(background_image)
+                fg_img    = read_image_pil(os.path.join(scaled_fg_dir, fg_name))
+                mask      = read_mask_pil(os.path.join(scaled_mask_dir, mask_name))
+                bg_t, fg_t, mask_t = self.triple_transform(bg_img, fg_img, mask)
+                mask_t    = self.mask_transform(mask_t).to(self.device)
+                fg_t      = self.image_transform(fg_t).to(self.device)
+                bg_t      = self.image_transform(bg_t).to(self.device)
+                
+                outs, _   = self.model(bg_t.unsqueeze(0), fg_t.unsqueeze(0), mask_t.unsqueeze(0), 'test') 
+                outs      = torch.softmax(outs, dim=1)[:,1,:,:]
+                outs      = transforms.ToPILImage()(outs)
+                out_file  = os.path.join(heatmap_dir, save_name)
+                outs.save(out_file) 
+                heatmap_list.append(out_file)
+        return heatmap_list
+    
+    def generate_bounding_box(self, foreground_image, foreground_mask, background_image, csv_path,
+                              cache_dir, heatmap_dir, fg_scale_num, composite_num_choose, composite_num):
+        icount = 0
+        with open(csv_path, mode='r', newline='') as csv_file:
+            csv_reader = csv.DictReader(csv_file)
+            for row in csv_reader:
+                scale = row['scale']
+                w = int(row['newWidth'])
+                h = int(row['newHeight'])
+                save_name = row['fg_name'].split(".")[0] + '_' + str(scale) + '.jpg'
+                fg_name = '{}_{}_{}_{}.jpg'.format(row["fg_name"].split(".")[0],row["bg_name"].split(".")[0],int(row["newWidth"]),int(row["newHeight"]))
+                save_name = fg_name.split(".")[0] + '_' + str(scale) + '.jpg'
+            
+                if icount == 0:
+                    bg_img     = read_image_pil(background_image)  
+                    fg_tocp    = read_image_pil(foreground_image)
+                    mask_tocp  = read_image_pil(foreground_mask)
+                    
+                    composite_dir = os.path.join(cache_dir, f'{fg_scale_num}scales_composite', f'{row["fg_name"].split(".")[0]}_{row["bg_name"].split(".")[0]}')
+                    os.makedirs(composite_dir, exist_ok=True)
+                    heatmap_center_list = []
+                    fg_size_list = []
+                
+                icount += 1
+                heatmap = Image.open(os.path.join(heatmap_dir, save_name))
+                heatmap = np.array(heatmap)
+                heatmap_center = np.zeros_like(heatmap, dtype=np.float64)
+                hb = int(h / bg_img.height * heatmap.shape[0] / 2)
+                wb = int(w / bg_img.width  * heatmap.shape[1] / 2)
+                heatmap_center[hb:-hb, wb:-wb] = heatmap[hb:-hb, wb:-wb]
+                heatmap_center_list.append(heatmap_center)
+                fg_size_list.append((h,w))
+                
+                if icount == fg_scale_num:
+                    icount = 0
+                    heatmap_center_stack = np.stack(heatmap_center_list)
+                    sorted_indices = np.argsort(-heatmap_center_stack, axis=None)
+                    sorted_indices = np.unravel_index(sorted_indices, heatmap_center_stack.shape)
+                    
+                    for i in range(composite_num):
+                        iscale, y_, x_ = sorted_indices[0][i], sorted_indices[1][i], sorted_indices[2][i]
+                        h, w = fg_size_list[iscale]
+                        x_   = x_/heatmap.shape[1]*bg_img.width
+                        y_   = y_/heatmap.shape[0]*bg_img.height
+                        x    = int(x_ - w / 2)
+                        y    = int(y_ - h / 2)
+
+                        composite_img, composite_msk = make_composite_PIL(fg_tocp, mask_tocp, bg_img, [x, y, w, h], return_mask=True)
+                        save_img_path = os.path.join(composite_dir, f'{row["fg_name"].split(".")[0]}_{row["bg_name"].split(".")[0]}_{x}_{y}_{w}_{h}.jpg')
+                        save_msk_path = os.path.join(composite_dir, f'{row["fg_name"].split(".")[0]}_{row["bg_name"].split(".")[0]}_{x}_{y}_{w}_{h}.png')
+                        composite_img.save(save_img_path)
+                        composite_msk.save(save_msk_path)
+
+        source_folder = composite_dir
+        composite_dir_choose = os.path.join(cache_dir, f'{fg_scale_num}scales_composite_RandomSelect_{composite_num_choose}', f'{row["fg_name"].split(".")[0]}_{row["bg_name"].split(".")[0]}')
+        os.makedirs(composite_dir_choose, exist_ok=True)
+        
+        bbox_list = []
+        image_files = [f for f in os.listdir(composite_dir) if f.endswith((".jpg"))]
+        selected_images = random.sample(image_files, composite_num_choose)
+        for image in selected_images:
+            source_path = os.path.join(source_folder, image)
+            destination_path = os.path.join(composite_dir_choose, image)
+            png_source_path = source_path.replace(".jpg", ".png")
+            destination_path_png = os.path.join(composite_dir_choose, image.replace(".jpg", ".png"))
+            shutil.copyfile(source_path, destination_path)
+            shutil.copyfile(png_source_path, destination_path_png)
+            bbox_list_parts = os.path.splitext(os.path.split(source_path)[1])[0].split('_')
+            x_box, y_box, w_box, h_box = int(bbox_list_parts[2]), int(bbox_list_parts[3]), int(bbox_list_parts[4]), int(bbox_list_parts[5])
+            last_four_numbers = [x_box, y_box, w_box, h_box]
+            bbox_list.append(last_four_numbers)
+        return bbox_list
+
+
+    @torch.no_grad()
+    def __call__(self, foreground_image, foreground_mask, background_image, cache_dir, heatmap_dir,
+                 fg_scale_num=16, composite_num_choose=3, composite_num=50):
+        """
+        Generate a heatmap for a pair of scaled foreground and background.
+        
+        Args:
+            foreground_image: foreground image path
+            foreground_mask: foreground mask path
+            background_image: background image path
+            cache_dir: folder path where scaled foreground images, scaled mask images and composite images are stored
+            heatmap_dir: folder path where heatmaps are stored
+            fg_scale_num: number of scales of scaled foreground images and mask images
+            composite_num_choose: the number of chosen composite images
+            composite_num: the number of composite images with the highest score
+
+        Returns:
+            box_list: the path of concatenated background image, foreground image and corresponding heatmap
+            heatmap_list: the path of heatmaps
+
+        """
+        scaled_fg_dir, scaled_mask_dir, csv_file = prepare_multi_fg_scales(cache_dir, foreground_image, foreground_mask, background_image, fg_scale_num)
+        heatmap_list = self.generate_heatmap(background_image, csv_file, 
+                                             scaled_fg_dir, scaled_mask_dir, heatmap_dir)
+        box_list     = self.generate_bounding_box(foreground_image, foreground_mask, background_image, csv_file,
+                                                  cache_dir, heatmap_dir, fg_scale_num, composite_num_choose, composite_num)
+        return box_list, heatmap_list

libcom/fopa_heat_map/source/backbone/ResNet.py

@@ -0,0 +1,229 @@
+# import torchvision.models as models
+# import torch.nn as nn
+# # https://pytorch.org/docs/stable/torchvision/models.html#id3
+#
+import os
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+
+model_urls = {
+    "resnet18": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
+    "resnet34": "https://download.pytorch.org/models/resnet34-333f7ec4.pth",
+    "resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
+    "resnet101": "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth",
+    "resnet152": "https://download.pytorch.org/models/resnet152-b121ed2d.pth",
+}
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, zero_init_residual=False):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)  # 6
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)  # 3
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride), nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        return x
+
+def resnet18(pretrained=False, **kwargs):
+    """Constructs a ResNet-18 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        pretrained_dict = model_zoo.load_url(model_urls["resnet18"])
+
+        model_dict = model.state_dict()
+        # 1. filter out unnecessary keys
+        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+        # 2. overwrite entries in the existing state dict
+        model_dict.update(pretrained_dict)
+        # 3. load the new state dict
+        model.load_state_dict(model_dict)
+    return model
+
+
+def Backbone_ResNet18_in3(pretrained=True, weight_path=None):
+    if pretrained:
+        print("The backbone model loads the pretrained parameters...")
+    net = pretrained_resnet18_4ch(pretrained=True, weight_path=weight_path)
+    div_2 = nn.Sequential(*list(net.children())[:3])
+    div_4 = nn.Sequential(*list(net.children())[3:5])
+    div_8 = net.layer2
+    div_16 = net.layer3
+    div_32 = net.layer4
+
+    return div_2, div_4, div_8, div_16, div_32
+
+
+def Backbone_ResNet18_in3_1(pretrained=True):
+    if pretrained:
+        print("The backbone model loads the pretrained parameters...")
+    net = resnet18(pretrained=pretrained)
+
+    model_dict = net.state_dict()
+    conv1 = model_dict['conv1.weight']
+    new = torch.zeros(64, 1, 7, 7)
+    for i, output_channel in enumerate(conv1):
+        new[i] = 0.299 * output_channel[0] + 0.587 * output_channel[1] + 0.114 * output_channel[2]
+    net.conv1 = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), padding=3, bias=False)
+    model_dict['conv1.weight'] = torch.cat((conv1, new), dim=1)
+    net.load_state_dict(model_dict)
+
+    div_1 = nn.Sequential(*list(net.children())[:1])
+    div_2 = nn.Sequential(*list(net.children())[1:3])
+    div_4 = nn.Sequential(*list(net.children())[3:5])
+    div_8 = net.layer2
+    div_16 = net.layer3
+    # div_32 = make_layer_4(BasicBlock, 448, 2, stride=2)
+    div_32 = net.layer4
+    return div_1, div_2, div_4, div_8, div_16, div_32
+
+
+
+def pretrained_resnet18_4ch(pretrained=True, weight_path=None, **kwargs):
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    model.conv1 = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), padding=3, bias=False)
+    
+    if pretrained:
+        # load the pretrained binary classification model for slow object placement assessment (SOPA)
+        # checkpoint = torch.load('./pretrained_models/SOPA.pth.tar')
+        current_directory = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+        checkpoint = torch.load(weight_path)
+        model.load_state_dict(checkpoint['state_dict'], strict=False)
+
+    return model
+
+
+