libcom 0.1.8__py3-none-any.whl

libcom/naive_composition/__init__.py

@@ -0,0 +1,3 @@
+from .generate_composite_image import get_composite_image
+
+__all__ = ['get_composite_image']

libcom/naive_composition/generate_composite_image.py

@@ -0,0 +1,116 @@
+import numpy as np
+import os
+import cv2
+from libcom.utils.process_image import *
+
+
+def get_composite_image(foreground_image, foreground_mask, background_image, bbox, option="none"):
+    """
+    Generate composite image through copy-and-paste.
+
+    Args:
+        foreground_image (str | numpy.ndarray): The path to foreground image or the background image in ndarray form.
+        foreground_mask (str | numpy.ndarray): Mask of foreground image which indicates the foreground object region in the foreground image.
+        background_image (str | numpy.ndarray): The path to background image or the background image in ndarray form. 
+        bbox (list): The bounding box which indicates the foreground's location in the background. [x1, y1, x2, y2]. 
+        option (str): 'none', 'gaussian', or 'poisson'. Image blending method. default: None.
+
+    Returns:
+        composite_image (numpy.ndarray): Generated composite image with the same resolution as input background image.
+        composite_mask (numpy.ndarray): Generated composite mask with the same resolution as composite image.
+    
+    Examples:
+        >>> from libcom import get_composite_image
+        >>> from libcom.utils.process_image import make_image_grid, draw_bbox_on_image
+        >>> import cv2
+        >>> test_dir = 'source/'
+        >>> img_list = ['1.jpg', '8.jpg']
+        >>> bbox_list = [[1000, 895, 1480, 1355], [1170, 944, 2331, 3069]] 
+        >>> for i,img_name in enumerate(img_list):
+        >>>     bg_img  = test_dir + 'background/' + img_name
+        >>>     bbox    = bbox_list[i] # x1,y1,x2,y2
+        >>>     fg_img  = test_dir + 'foreground/' + img_name
+        >>>     fg_mask = test_dir + 'foreground_mask/' + img_name.replace('.jpg', '.png')
+        >>>     # generate composite images by naive methods
+        >>>     comp_img1, comp_mask1 = get_composite_image(fg_img, fg_mask, bg_img, bbox, 'none')
+        >>>     comp_img2, comp_mask2 = get_composite_image(fg_img, fg_mask, bg_img, bbox, 'gaussian')
+        >>>     comp_img3, comp_mask3 = get_composite_image(fg_img, fg_mask, bg_img, bbox, 'poisson')
+        >>>     vis_list = [bg_img, fg_img, comp_img1, comp_mask1, comp_img2, comp_mask2, comp_img3, comp_mask3]
+        >>>     # visualization results
+        >>>     grid_img  = make_image_grid(vis_list, cols=4)
+        >>>     cv2.imwrite(f'../docs/_static/image/generatecomposite_result{i+1}.jpg', grid_img)
+    
+    Expected result:
+
+    .. image:: _static/image/generatecomposite_result1.jpg
+        :scale: 38 %
+        
+    .. image:: _static/image/generatecomposite_result2.jpg
+        :scale: 38 %
+            
+    """
+    choices = ['none', 'gaussian', 'poisson']
+    assert option in choices, f'{option} must be one of {choices}'
+    fg_img   = read_image_opencv(foreground_image)
+    fg_mask  = read_mask_opencv(foreground_mask)
+    bg_img   = read_image_opencv(background_image)
+    if option == 'none':
+        comp_img,comp_mask = simple_composite_image(bg_img, fg_img, fg_mask, bbox)
+    elif option == 'gaussian':
+        comp_img,comp_mask = gaussian_composite_image(bg_img, fg_img, fg_mask, bbox)
+    else: # option == 'poisson'
+        comp_img,comp_mask = poisson_composite_image(bg_img, fg_img, fg_mask, bbox)
+    return comp_img, comp_mask
+
+
+def crop_and_resize_foreground(fg_img, fg_mask, bbox):
+    if fg_mask.shape[:2] != fg_img.shape[:2]:
+        fg_mask = cv2.resize(fg_mask, (fg_img.shape[1], fg_img.shape[0]))
+    fg_bbox = convert_mask_to_bbox(fg_mask)
+    fg_region = fg_img[fg_bbox[1] : fg_bbox[3], fg_bbox[0] : fg_bbox[2]]
+    x1, y1, x2, y2 = bbox
+    fg_region = cv2.resize(fg_region, (x2-x1, y2-y1), cv2.INTER_CUBIC)
+    fg_mask   = fg_mask[fg_bbox[1] : fg_bbox[3], fg_bbox[0] : fg_bbox[2]]
+    fg_mask   = cv2.resize(fg_mask, (x2-x1, y2-y1))
+    fg_mask   = np.where(fg_mask > 127, 255, 0).astype(fg_img.dtype)
+    return fg_region, fg_mask 
+
+
+def simple_composite_image(bg_img, fg_img, fg_mask, bbox):
+    fg_region, fg_mask = crop_and_resize_foreground(fg_img, fg_mask, bbox)
+    x1, y1, x2, y2 = bbox
+    comp_mask = np.zeros((bg_img.shape[0], bg_img.shape[1]), dtype=np.uint8) 
+    comp_mask[y1:y2, x1:x2] = fg_mask
+    comp_img = bg_img.copy()
+    comp_img[y1:y2, x1:x2]  = np.where(fg_mask[:,:,np.newaxis]> 127, fg_region, comp_img[y1:y2, x1:x2])
+    return comp_img, comp_mask
+    
+    
+def gaussian_composite_image(bg_img, fg_img, fg_mask, bbox, kernel_size=15):
+    if fg_mask.shape[:2] != fg_img.shape[:2]:
+        fg_mask = cv2.resize(fg_mask, (fg_img.shape[1], fg_img.shape[0]))
+    fg_mask = cv2.GaussianBlur(255-fg_mask, (kernel_size, kernel_size), kernel_size/3.)
+    fg_mask = 255 - fg_mask
+    fg_bbox = convert_mask_to_bbox(fg_mask)
+    fg_region = fg_img[fg_bbox[1] : fg_bbox[3], fg_bbox[0] : fg_bbox[2]]
+    x1, y1, x2, y2 = bbox
+    fg_region = cv2.resize(fg_region, (x2-x1, y2-y1), cv2.INTER_CUBIC)
+    fg_mask   = fg_mask[fg_bbox[1] : fg_bbox[3], fg_bbox[0] : fg_bbox[2]]
+    fg_mask   = cv2.resize(fg_mask, (x2-x1, y2-y1))
+    norm_mask = (fg_mask.astype(np.float32) / 255)[:,:,np.newaxis]
+    
+    comp_mask = np.zeros((bg_img.shape[0], bg_img.shape[1]), dtype=np.uint8) 
+    comp_mask[y1:y2, x1:x2] = fg_mask
+    comp_img = bg_img.copy()
+    comp_img[y1:y2, x1:x2] = (fg_region * norm_mask + comp_img[y1:y2, x1:x2] * (1-norm_mask)).astype(comp_img.dtype) 
+    return comp_img, comp_mask
+    
+
+def poisson_composite_image(bg_img, fg_img, fg_mask, bbox, clone_method=cv2.NORMAL_CLONE):
+    fg_region, fg_mask = crop_and_resize_foreground(fg_img, fg_mask, bbox)
+    x1, y1, x2, y2 = bbox
+    center    = (x1+x2)//2, (y1+y2)//2
+    comp_img  = cv2.seamlessClone(fg_region, bg_img, fg_mask[:,:,np.newaxis], center, clone_method)
+    comp_mask = np.zeros((bg_img.shape[0], bg_img.shape[1]), dtype=np.uint8) 
+    comp_mask[y1:y2, x1:x2] = fg_mask
+    return comp_img, comp_mask

libcom/opa_score/__init__.py

@@ -0,0 +1,3 @@
+from .opa_score_prediction import OPAScoreModel
+
+__all__ = ['OPAScoreModel']

libcom/opa_score/opa_score_prediction.py

@@ -0,0 +1,101 @@
+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.opa_score.source import ObjectPlaceNet
+import torch 
+import os
+import torchvision.transforms as transforms
+
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['SimOPA'] 
+
+
+class OPAScoreModel:
+    """
+    OPA 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 OPAScoreModel
+        >>> from libcom import get_composite_image
+        >>> from libcom.utils.process_image import make_image_grid
+        >>> import cv2
+        >>> net = OPAScoreModel(device=0, model_type='SimOPA')
+        >>> test_dir  = './source'
+        >>> bg_img    = 'source/background/17.jpg'
+        >>> fg_img    = 'source/foreground/17.jpg'
+        >>> fg_mask   = 'source/foreground_mask/17.png'
+        >>> bbox_list = [[475, 697, 1275, 1401], [475, 300, 1275, 1004]]
+        >>> comp1, comp_mask1 = get_composite_image(fg_img, fg_mask, bg_img, bbox_list[0])
+        >>> comp2, comp_mask2 = get_composite_image(fg_img, fg_mask, bg_img, bbox_list[1])
+        >>> score1 = net(comp1, comp_mask1)
+        >>> score2 = net(comp2, comp_mask2)
+        >>> grid_img  = make_image_grid([comp1, comp_mask1, comp2, comp_mask2], text_list=[f'opa_score:{score1:.2f}', 'composite-mask', f'opa_score:{score2:.2f}', 'composite-mask'])
+        >>> cv2.imwrite('../docs/_static/image/opascore_result1.jpg', grid_img)
+
+    Expected result:
+
+    .. image:: _static/image/opascore_result1.jpg
+        :scale: 38 %
+
+            
+    """
+    def __init__(self, device=0, model_type='SimOPA', **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', 'SimOPA.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 = ObjectPlaceNet(False)
+        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.transformer = transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+        ])
+    
+    def inputs_preprocess(self, composite_image, composite_mask):
+        img  = read_image_pil(composite_image)
+        img  = self.transformer(img)
+        mask = read_mask_pil(composite_mask)
+        mask = self.transformer(mask)
+        cat_img = torch.cat([img, mask], dim=0)
+        cat_img = cat_img.unsqueeze(0).to(self.device)
+        return cat_img
+    
+    def outputs_postprocess(self, outputs):
+        score   = torch.softmax(outputs, dim=-1)[0, 1].cpu().item()
+        return score
+    
+    @torch.no_grad()
+    def __call__(self, composite_image, composite_mask):
+        """
+        Predicting the object placement assessment (opa) score for the given composite image, which evaluates the rationality of foreground object placement.
+
+        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:
+            opa_score (float): Predicted opa score ranges from 0 to 1, where a larger score indicates more reasonable placement.
+
+        """
+        inputs    = self.inputs_preprocess(composite_image, composite_mask)
+        outputs   = self.model(inputs)
+        preds     = self.outputs_postprocess(outputs)
+        return preds

libcom/opa_score/source/__init__.py

@@ -0,0 +1,3 @@
+import os,sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from object_place_net import ObjectPlaceNet

libcom/opa_score/source/object_place_net.py

@@ -0,0 +1,55 @@
+import os,sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from oblect_place_config import opt
+from resnet_4ch import resnet
+
+class ObjectPlaceNet(nn.Module):
+    def __init__(self, backbone_pretrained=True):
+        super(ObjectPlaceNet, self).__init__()
+        
+        ## Backbone, only resnet
+        resnet_layers = int(opt.backbone.split('resnet')[-1])
+        backbone = resnet(resnet_layers,
+                          backbone_pretrained,
+                          os.path.join(opt.pretrained_model_path, opt.backbone+'.pth'))
+
+        # drop pool layer and fc layer
+        features = list(backbone.children())[:-2]
+        backbone = nn.Sequential(*features)
+        self.backbone = backbone
+
+        ## global predict
+        self.global_feature_dim = 512 if opt.backbone in ['resnet18', 'resnet34'] else 2048
+   
+        self.avgpool3x3 = nn.AdaptiveAvgPool2d(3)
+        self.avgpool1x1 = nn.AdaptiveAvgPool2d(1)
+
+        self.prediction_head = nn.Linear(self.global_feature_dim, opt.class_num, bias=False)
+
+    def forward(self, img_cat):
+        '''  img_cat:b,4,256,256  '''
+        global_feature = None
+        if opt.without_mask:
+            img_cat = img_cat[:,0:3] 
+        feature_map = self.backbone(img_cat)  # b,512,8,8 (resnet layer4 output shape: b,c,8,8, if resnet18, c=512)
+        global_feature = self.avgpool1x1(feature_map)  # b,512,1,1
+        global_feature = global_feature.flatten(1) # b,512
+
+        prediction = self.prediction_head(global_feature) 
+
+        return prediction
+
+if __name__ == '__main__':
+    device = torch.device('cuda:0')
+    b = 4
+    img_cat = torch.randn(b, 4, 256, 256).to(device)
+    model = ObjectPlaceNet(backbone_pretrained=False).to(device)
+    local_pre = model(img_cat)
+    print(local_pre)
+
+
+

libcom/opa_score/source/oblect_place_config.py

@@ -0,0 +1,66 @@
+import os
+
+class Config(object):
+    # path
+    dataset_path = './dataset'
+    
+    train_data_path = os.path.join(dataset_path, 'train_set.csv')
+    test_data_path = os.path.join(dataset_path, 'test_set.csv')
+
+    img_path = os.path.join('./')
+    mask_path = os.path.join('.')
+ 
+    pretrained_model_path = './pretrained_models'
+
+    ## 
+    class_num = 2 
+    img_size = 256
+    batch_size = 128
+    num_workers = 8
+    global_feature_size = 8
+
+    # train
+    base_lr = 1e-4
+    lr_milestones = [10, 16]
+    lr_gamma = 0.1
+    epochs = 25
+    eval_freq = 1
+    save_freq = 5
+    display_freq = 20
+
+    # Config
+    backbone = 'resnet18'
+    gpu_id = 1
+
+    without_mask = False
+
+    # Save path 
+    prefix = backbone
+    if without_mask:
+        prefix += '+without_mask'
+
+    exp_root = os.path.join(os.getcwd(), './experiments/simopa/')
+    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)
+
+    checkpoint_dir = os.path.join(exp_path, 'checkpoints')
+    log_dir = os.path.join(exp_path, 'logs')
+
+    def create_path(self):
+        print('Create experiments directory: ', self.exp_path)
+        os.makedirs(self.exp_path)
+        os.makedirs(self.checkpoint_dir)
+        os.makedirs(self.log_dir)
+
+opt = Config()
+
+if __name__ == "__main__":
+    print(opt.exp_root)

libcom/opa_score/source/resnet_4ch.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+import os
+from oblect_place_config import opt
+
+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):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, 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):
+        residual = 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:
+            residual = self.downsample(x)
+
+        out += residual
+        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 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = 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:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        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)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i 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)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+def resnet(layers, pretrained=False, pretrained_weight=None, **kwargs):
+    model = None
+    if layers == 18:
+        model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    elif layers == 34:
+        model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    elif layers == 50:
+        model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    elif layers == 101:
+        model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    elif layers == 152:
+        model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    else:
+        raise Exception('Unsupport resnet layers numbers: ', layers)
+    if opt.without_mask:
+        if pretrained:
+            model.load_state_dict(torch.load(pretrained_weight))
+            print('loaded pretrained resnet{} from {}'.format(layers, pretrained_weight))
+    else:
+        model.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        if pretrained:
+            pretrained_state_dict = torch.load(pretrained_weight)
+            conv1 = pretrained_state_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]
+            pretrained_state_dict['conv1.weight'] = torch.cat((conv1, new), dim=1)
+            model.load_state_dict(pretrained_state_dict)
+            print('loaded pretrained resnet{} from {}'.format(layers, pretrained_weight))
+    return model
+
+def resnet_for_depth(layers, pretrained=False, pretrained_weight=None, **kwargs):
+    model = None
+    if layers == 18:
+        model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    elif layers == 34:
+        model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    elif layers == 50:
+        model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    elif layers == 101:
+        model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    elif layers == 152:
+        model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    else:
+        raise Exception('Unsupport resnet layers numbers: ', layers)
+
+    model.conv1 = nn.Conv2d(2, 64, kernel_size=7, stride=2, padding=3, bias=False)
+    if pretrained:
+        pretrained_state_dict = torch.load(pretrained_weight)
+        conv1 = pretrained_state_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]
+        pretrained_state_dict['conv1.weight'] = torch.cat((new, new), dim=1)
+        model.load_state_dict(pretrained_state_dict)
+        print('loaded pretrained resnet{} from {}'.format(layers, pretrained_weight))
+    return model
+
+if __name__ == '__main__':
+    model = resnet(101, pretrained=False)
+    input = torch.randn(2,4,256,256)
+    print('Input shape {}, output shape {}'.format(input.shape, model(input).shape))

libcom/os_insert/__init__.py

@@ -0,0 +1,3 @@
+from .os_insert import OSInsertModel
+
+__all__ = ['OSInsertModel']

libcom/os_insert/diffusers_osinsert/__init__.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import importlib
+from contextlib import contextmanager
+import sys
+import types
+from pathlib import Path
+
+
+def _load_patched_pipeline_source() -> str:
+    src_path = Path(__file__).resolve().parent / "_patched_pipeline_flux_fill.py"
+    return src_path.read_text(encoding="utf-8")
+
+
+_ORIG_CALL = None
+_PATCHED_CALL = None
+_PATCH_ENABLED = False
+
+
+def _load_patched_call():
+    global _PATCHED_CALL
+    if _PATCHED_CALL is not None:
+        return _PATCHED_CALL
+
+    patched_mod_name = "diffusers_osinsert._patched_pipeline_flux_fill"
+    patched_mod = types.ModuleType(patched_mod_name)
+    patched_mod.__file__ = str(Path(__file__).resolve().parent / "_patched_pipeline_flux_fill.py")
+    # Make relative imports inside the patched file resolve against diffusers.*
+    patched_mod.__package__ = "diffusers.pipelines.flux"
+
+    code = _load_patched_pipeline_source()
+    exec(compile(code, patched_mod.__file__, "exec"), patched_mod.__dict__)
+    sys.modules[patched_mod_name] = patched_mod
+
+    _PATCHED_CALL = patched_mod.FluxFillPipeline.__call__
+    return _PATCHED_CALL
+
+
+def enable_patch() -> None:
+    """Enable the OSInsert FluxFillPipeline patch for the current Python process."""
+    global _ORIG_CALL, _PATCH_ENABLED
+    if _PATCH_ENABLED:
+        return
+
+    upstream_mod = importlib.import_module("diffusers.pipelines.flux.pipeline_flux_fill")
+    if _ORIG_CALL is None:
+        _ORIG_CALL = upstream_mod.FluxFillPipeline.__call__
+
+    upstream_mod.FluxFillPipeline.__call__ = _load_patched_call()
+    _PATCH_ENABLED = True
+
+
+def disable_patch() -> None:
+    """Disable the OSInsert FluxFillPipeline patch and restore upstream behavior."""
+    global _PATCH_ENABLED
+    if not _PATCH_ENABLED:
+        return
+
+    upstream_mod = importlib.import_module("diffusers.pipelines.flux.pipeline_flux_fill")
+    if _ORIG_CALL is not None:
+        upstream_mod.FluxFillPipeline.__call__ = _ORIG_CALL
+
+    _PATCH_ENABLED = False
+
+
+@contextmanager
+def patch_context():
+    """Temporarily enable the patch within a `with` block and restore afterwards."""
+    enable_patch()
+    try:
+        yield
+    finally:
+        disable_patch()
+
+
+# Re-export upstream classes for convenience.
+from diffusers.pipelines.flux.pipeline_flux_fill import FluxFillPipeline  # noqa: E402
+from diffusers import FluxPriorReduxPipeline  # noqa: E402