libcom 0.1.4__tar.gz → 0.1.5.post1__tar.gz

{libcom-0.1.4 → libcom-0.1.5.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libcom
-Version: 0.1.4
+Version: 0.1.5.post1
 Summary: Image Composition Toolbox
 Home-page: https://github.com/bcmi/libcom
 Author: BCMI Lab
@@ -29,11 +29,15 @@ Requires-Dist: omegaconf
 Requires-Dist: opencv_python
 Requires-Dist: opencv_python_headless
 Requires-Dist: pytorch_lightning
-Requires-Dist: transformers
+Requires-Dist: transformers==4.44.2
 Requires-Dist: tqdm
 Requires-Dist: open_clip_torch
 Requires-Dist: taming-transformers-rom1504
 Requires-Dist: scikit-learn
+Requires-Dist: segment_anything
+Requires-Dist: clip
+Requires-Dist: pydantic==2.6.1
+Requires-Dist: lpips==0.1.4
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier
@@ -59,12 +63,11 @@ Dynamic: summary
 [![PyPI](https://img.shields.io/pypi/v/libcom)](https://pypi.org/project/libcom)
 [![Downloads](https://static.pepy.tech/badge/libcom)](https://pepy.tech/project/libcom)
 [![Hits](https://hits.sh/github.com/bcmi/libcom.svg?label=views)](https://hits.sh/github.com/bcmi/libcom/)
-[![Static Badge](https://img.shields.io/badge/Online_Workbench-yellow)](http://libcom.ustcnewly.com/)
+[![Static Badge](https://img.shields.io/badge/Online_Workbench-yellow)](https://libcom.ustcnewly.com/)
 [![Static Badge](https://img.shields.io/badge/Dataset&Code_Resources-green)](https://github.com/bcmi/Awesome-Object-Insertion)
 [![Static Badge](https://img.shields.io/badge/survey-arxiv%3A2106.14490-red)](https://arxiv.org/pdf/2106.14490.pdf)
 [![Static Badge](https://img.shields.io/github/stars/bcmi/libcom.svg?style=social)](https://github.com/bcmi/libcom/stargazers)
 
-
 ---
 
 ## Introduction
@@ -74,21 +77,23 @@ Dynamic: summary
 <img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/image_composition_task.gif" width="600" />
 </div>
 
-**_libcom_ covers a diversity of related tasks in the field of image composition**, including image blending, standard/painterly image harmonization, shadow generation, object placement, generative composition, quality evaluation, *etc*. For each task, we integrate one or two selected methods considering both efficiency and effectiveness. The selected methods will be continuously updated upon the emergence of better methods. 
-
-**The ultimate goal of this library is solving all the problems related to image composition with simple `import libcom`.**
+**_libcom_ covers a diversity of related tasks in the field of image composition**, including image blending, standard/painterly image harmonization, shadow generation, reflection generation, object placement, generative composition, quality evaluation, *etc*. For each task, we integrate one or two selected methods considering both efficiency and effectiveness. The selected methods will be continuously updated upon the emergence of better methods. 
 
+**The ultimate goal of this library is solving all the problems related to image composition with simple `import libcom`.** 
 Welcome to follow WeChat public account ["Newly AIGCer"](https://www.ustcnewly.com/blog.html) or Zhihu Column ["Newly CVer"](https://www.zhihu.com/column/c_1333918224900206592) to get the latest information about image composition! 
 
+We compare the performance of Libcom and Banana Pro across various image composition tasks and write a brief summary in [English](https://zhuanlan.zhihu.com/p/1992987450860118091)/[Chinese](https://zhuanlan.zhihu.com/p/1976625829741082597). 
+
 ## Online Workbench
 
-We have built the [image composition workbench](http://libcom.ustcnewly.com/) based on our Libcom toolbox. The interface is as follows:
+We have built the [image composition workbench](https://libcom.ustcnewly.com/) based on our Libcom toolbox. The interface is as follows:
 
 <div align="center">
 </br>
 <img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/online_workbench.jpg" width="1000" />
 </div>
 
+[![]](https://github.com/user-attachments/assets/87416ec5-2461-42cb-9f2d-5030b1e1b5ec)
 
 ## Main Functions
 
@@ -96,16 +101,15 @@ We have built the [image composition workbench](http://libcom.ustcnewly.com/) ba
 - **OPAScoreModel** [[OPA]](https://github.com/bcmi/Object-Placement-Assessment-Dataset-OPA) evaluates the rationality of foreground object placement in a composite image.
 - **FOPAHeatMapModel** [[FOPA]](https://github.com/bcmi/FOPA-Fast-Object-Placement-Assessment) can predict the rationality scores for all locations/scales given a background-foreground pair, and output the composite image with optimal location/scale.  
 - **color_transfer** adjusts the foreground color to match the background using traditional color transfer method.
-- **ImageHarmonizationModel** [[PCTNet]](https://github.com/rakutentech/PCT-Net-Image-Harmonization) adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
+- **ImageHarmonizationModel** [[PCTNet]](https://github.com/rakutentech/PCT-Net-Image-Harmonization) [[LBM]](https://github.com/gojasper/LBM)  adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
 - **PainterlyHarmonizationModel** [[PHDNet]](https://github.com/bcmi/PHDNet-Painterly-Image-Harmonization)  [[PHDiffusion]](https://github.com/bcmi/PHDiffusion-Painterly-Image-Harmonization) adjusts the foreground style to be compatible with the background given artistic background and photorealistic foreground.
 - **HarmonyScoreModel** [[BargainNet]](https://github.com/bcmi/BargainNet-Image-Harmonization) evaluates the harmony level between foreground and background in a composite image.
 - **InharmoniousLocalizationModel** [[MadisNet]](https://github.com/bcmi/MadisNet-Inharmonious-Region-Localization) localizes the inharmonious region in a synthetic image.
-- **FOSScoreModel** [[DiscoFOS]](https://github.com/bcmi/Foreground-Object-Search-Dataset-FOSD) evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics. Due to limited training data, the generalization ability of this model is poor. 
+- **FOSScoreModel** [[DiscoFOS]](https://github.com/bcmi/Foreground-Object-Search-Dataset-FOSD) evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics. 
 - **ShadowGenerationModel** [[GPSDiffusion]](https://github.com/bcmi/GPSDiffusion-Object-Shadow-Generation) generates plausible shadow for the inserted object in a composite image. 
-- **ReflectionGenerationModel** generates plausible reflection for the inserted object in a composite image. 
+- **ReflectionGenerationModel** [[RGDiffusion]](https://github.com/bcmi/Object-Reflection-Generation-Dataset-DEROBA) generates plausible reflection for the inserted object in a composite image. 
 - **KontextBlendingHarmonizationModel** [[FluxKontext]](https://github.com/black-forest-labs/flux) is a generative image composition model which inserts foreground into the specified bounding box in the background. The pose and view of foreground stay unchanged. The "blending" mode does not adjust the foreground illumination, while the "harmonization" mode adjusts the foreground illumination to make the composite image harmonious. 
-- **InsertAnythingModel** [[InsertAnything]](https://github.com/song-wensong/insert-anything) is a generative image composition model which inserts foreground into the specified bounding box in the background. The model has reasonable ability to adjust the pose and view of foreground according to the background. 
-
+- **OSInsertModel** [[OSInsert]](https://github.com/bcmi/OSInsert-Image-Composition) is a generative image composition model which inserts foreground into the specified bounding box in the background. The "conservative" mode has weak ability to adjust the pose and view of foreground according to the background. The "aggressive" mode has stronger ability to adjust the pose and view of foreground according to the background. 
 
 **For the detailed method descriptions, code examples, visualization results, and performance comments, please refer to our [[documents]](https://libcom.readthedocs.io/en/latest/).** If the model performance is not satisfactory, you can finetune the pretrained model on your own dataset using the source repository and replace the original model. 
 
@@ -120,7 +124,6 @@ Please refer to [[Installation]](https://github.com/bcmi/libcom/blob/main/docs/g
 - Institution: [Brain-like Computing and Machine Intelligence (BCMI) Lab](https://bcmi.sjtu.edu.cn/).
 - Project Initiator & Team Manager: [Li Niu](https://www.ustcnewly.com/index.html). 
 - Architect & Lead Developer: [Bo Zhang](https://bo-zhang-cs.github.io/), [Yujie Zhou](https://github.com/YujieOuO).
-- Documentation Manager: [Jiacheng Sui](https://github.com/charlessjc).
 - Module Developers: [Jiacheng Sui](https://github.com/charlessjc), [Bingjie Gao](https://github.com/WhynotHAHA), [Lingxiao Lu](https://github.com/pokaaa), [Xinhao Tao](https://github.com/taoxinhao13), [Junyan Cao](https://github.com/cjy-4), [Haonan Zhao](https://github.com/nononononno), [Jiaxuan Chen](https://github.com/csdahunzi), [Junqi You](https://github.com/dhmbb2).
 
 ## License
@@ -145,3 +148,6 @@ If you use our toolbox, please cite our survey paper using the following BibTeX
 
 
 
+
+
+

{libcom-0.1.4 → libcom-0.1.5.post1}/README.md

@@ -11,12 +11,11 @@
 [![PyPI](https://img.shields.io/pypi/v/libcom)](https://pypi.org/project/libcom)
 [![Downloads](https://static.pepy.tech/badge/libcom)](https://pepy.tech/project/libcom)
 [![Hits](https://hits.sh/github.com/bcmi/libcom.svg?label=views)](https://hits.sh/github.com/bcmi/libcom/)
-[![Static Badge](https://img.shields.io/badge/Online_Workbench-yellow)](http://libcom.ustcnewly.com/)
+[![Static Badge](https://img.shields.io/badge/Online_Workbench-yellow)](https://libcom.ustcnewly.com/)
 [![Static Badge](https://img.shields.io/badge/Dataset&Code_Resources-green)](https://github.com/bcmi/Awesome-Object-Insertion)
 [![Static Badge](https://img.shields.io/badge/survey-arxiv%3A2106.14490-red)](https://arxiv.org/pdf/2106.14490.pdf)
 [![Static Badge](https://img.shields.io/github/stars/bcmi/libcom.svg?style=social)](https://github.com/bcmi/libcom/stargazers)
 
-
 ---
 
 ## Introduction
@@ -26,21 +25,23 @@
 <img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/image_composition_task.gif" width="600" />
 </div>
 
-**_libcom_ covers a diversity of related tasks in the field of image composition**, including image blending, standard/painterly image harmonization, shadow generation, object placement, generative composition, quality evaluation, *etc*. For each task, we integrate one or two selected methods considering both efficiency and effectiveness. The selected methods will be continuously updated upon the emergence of better methods. 
-
-**The ultimate goal of this library is solving all the problems related to image composition with simple `import libcom`.**
+**_libcom_ covers a diversity of related tasks in the field of image composition**, including image blending, standard/painterly image harmonization, shadow generation, reflection generation, object placement, generative composition, quality evaluation, *etc*. For each task, we integrate one or two selected methods considering both efficiency and effectiveness. The selected methods will be continuously updated upon the emergence of better methods. 
 
+**The ultimate goal of this library is solving all the problems related to image composition with simple `import libcom`.** 
 Welcome to follow WeChat public account ["Newly AIGCer"](https://www.ustcnewly.com/blog.html) or Zhihu Column ["Newly CVer"](https://www.zhihu.com/column/c_1333918224900206592) to get the latest information about image composition! 
 
+We compare the performance of Libcom and Banana Pro across various image composition tasks and write a brief summary in [English](https://zhuanlan.zhihu.com/p/1992987450860118091)/[Chinese](https://zhuanlan.zhihu.com/p/1976625829741082597). 
+
 ## Online Workbench
 
-We have built the [image composition workbench](http://libcom.ustcnewly.com/) based on our Libcom toolbox. The interface is as follows:
+We have built the [image composition workbench](https://libcom.ustcnewly.com/) based on our Libcom toolbox. The interface is as follows:
 
 <div align="center">
 </br>
 <img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/online_workbench.jpg" width="1000" />
 </div>
 
+[![]](https://github.com/user-attachments/assets/87416ec5-2461-42cb-9f2d-5030b1e1b5ec)
 
 ## Main Functions
 
@@ -48,16 +49,15 @@ We have built the [image composition workbench](http://libcom.ustcnewly.com/) ba
 - **OPAScoreModel** [[OPA]](https://github.com/bcmi/Object-Placement-Assessment-Dataset-OPA) evaluates the rationality of foreground object placement in a composite image.
 - **FOPAHeatMapModel** [[FOPA]](https://github.com/bcmi/FOPA-Fast-Object-Placement-Assessment) can predict the rationality scores for all locations/scales given a background-foreground pair, and output the composite image with optimal location/scale.  
 - **color_transfer** adjusts the foreground color to match the background using traditional color transfer method.
-- **ImageHarmonizationModel** [[PCTNet]](https://github.com/rakutentech/PCT-Net-Image-Harmonization) adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
+- **ImageHarmonizationModel** [[PCTNet]](https://github.com/rakutentech/PCT-Net-Image-Harmonization) [[LBM]](https://github.com/gojasper/LBM)  adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
 - **PainterlyHarmonizationModel** [[PHDNet]](https://github.com/bcmi/PHDNet-Painterly-Image-Harmonization)  [[PHDiffusion]](https://github.com/bcmi/PHDiffusion-Painterly-Image-Harmonization) adjusts the foreground style to be compatible with the background given artistic background and photorealistic foreground.
 - **HarmonyScoreModel** [[BargainNet]](https://github.com/bcmi/BargainNet-Image-Harmonization) evaluates the harmony level between foreground and background in a composite image.
 - **InharmoniousLocalizationModel** [[MadisNet]](https://github.com/bcmi/MadisNet-Inharmonious-Region-Localization) localizes the inharmonious region in a synthetic image.
-- **FOSScoreModel** [[DiscoFOS]](https://github.com/bcmi/Foreground-Object-Search-Dataset-FOSD) evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics. Due to limited training data, the generalization ability of this model is poor. 
+- **FOSScoreModel** [[DiscoFOS]](https://github.com/bcmi/Foreground-Object-Search-Dataset-FOSD) evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics. 
 - **ShadowGenerationModel** [[GPSDiffusion]](https://github.com/bcmi/GPSDiffusion-Object-Shadow-Generation) generates plausible shadow for the inserted object in a composite image. 
-- **ReflectionGenerationModel** generates plausible reflection for the inserted object in a composite image. 
+- **ReflectionGenerationModel** [[RGDiffusion]](https://github.com/bcmi/Object-Reflection-Generation-Dataset-DEROBA) generates plausible reflection for the inserted object in a composite image. 
 - **KontextBlendingHarmonizationModel** [[FluxKontext]](https://github.com/black-forest-labs/flux) is a generative image composition model which inserts foreground into the specified bounding box in the background. The pose and view of foreground stay unchanged. The "blending" mode does not adjust the foreground illumination, while the "harmonization" mode adjusts the foreground illumination to make the composite image harmonious. 
-- **InsertAnythingModel** [[InsertAnything]](https://github.com/song-wensong/insert-anything) is a generative image composition model which inserts foreground into the specified bounding box in the background. The model has reasonable ability to adjust the pose and view of foreground according to the background. 
-
+- **OSInsertModel** [[OSInsert]](https://github.com/bcmi/OSInsert-Image-Composition) is a generative image composition model which inserts foreground into the specified bounding box in the background. The "conservative" mode has weak ability to adjust the pose and view of foreground according to the background. The "aggressive" mode has stronger ability to adjust the pose and view of foreground according to the background. 
 
 **For the detailed method descriptions, code examples, visualization results, and performance comments, please refer to our [[documents]](https://libcom.readthedocs.io/en/latest/).** If the model performance is not satisfactory, you can finetune the pretrained model on your own dataset using the source repository and replace the original model. 
 
@@ -72,7 +72,6 @@ Please refer to [[Installation]](https://github.com/bcmi/libcom/blob/main/docs/g
 - Institution: [Brain-like Computing and Machine Intelligence (BCMI) Lab](https://bcmi.sjtu.edu.cn/).
 - Project Initiator & Team Manager: [Li Niu](https://www.ustcnewly.com/index.html). 
 - Architect & Lead Developer: [Bo Zhang](https://bo-zhang-cs.github.io/), [Yujie Zhou](https://github.com/YujieOuO).
-- Documentation Manager: [Jiacheng Sui](https://github.com/charlessjc).
 - Module Developers: [Jiacheng Sui](https://github.com/charlessjc), [Bingjie Gao](https://github.com/WhynotHAHA), [Lingxiao Lu](https://github.com/pokaaa), [Xinhao Tao](https://github.com/taoxinhao13), [Junyan Cao](https://github.com/cjy-4), [Haonan Zhao](https://github.com/nononononno), [Jiaxuan Chen](https://github.com/csdahunzi), [Junqi You](https://github.com/dhmbb2).
 
 ## License
@@ -97,3 +96,6 @@ If you use our toolbox, please cite our survey paper using the following BibTeX
 
 
 
+
+
+

{libcom-0.1.4 → libcom-0.1.5.post1}/libcom/__init__.py

@@ -3,20 +3,21 @@ from .naive_composition import *
 from .opa_score import *
 from .harmony_score import *
 from .inharmonious_region_localization import *
-from .image_harmonization import *
+from .image_harmonization_old import *
 from .painterly_image_harmonization import *
 from .fopa_heat_map import *
 from .fos_score import *
-from .insert_anything 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',
-    'InsertAnythingModel', 'KontextBlendingHarmonizationModel',
-    'ShadowGenerationModel', 'ReflectionGenerationModel'
+    'KontextBlendingHarmonizationModel',
+    'ShadowGenerationModel', 'ReflectionGenerationModel',
+    'OSInsertModel'
 ]

{libcom-0.1.4 → libcom-0.1.5.post1}/libcom/fos_score/fos_score_prediction.py

@@ -14,7 +14,7 @@ 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', 'FOS_E'] 
+model_set = ['FOS_D'] 
 
 class FOSScoreModel:
     """
@@ -31,7 +31,7 @@ class FOSScoreModel:
         >>> import cv2
         >>> import torch
         >>> task_name = 'fos_score_prediction'
-        >>> MODEL_TYPE = 'FOS_D' # choose from 'FOS_D', 'FOS_E'
+        >>> 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'
@@ -70,14 +70,10 @@ class FOSScoreModel:
         """
         Build pretrained model from path of weight.
         """
-        if self.model_type == "FOS_E":
-            model = StudentModel(self.cfg)
-            model.load_state_dict(torch.load(weight_path, map_location='cpu', weights_only=True), strict=True)
-            self.model = model.to(self.device).eval()
-        else:
-            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()
+
+        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
@@ -94,17 +90,14 @@ class FOSScoreModel:
         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)] = 255
-        if self.model_type == 'FOS_D':
-            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()
-            scale_comp_r = Image.fromarray(scale_comp, mode="RGB")
-            return scale_comp_r
-
-        return bg, fg
+            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
@@ -185,19 +178,11 @@ class FOSScoreModel:
             fos_score (float): Predicted compatibility score between the given background image and the given foreground image.
 
         """
-        if self.model_type == 'FOS_E':
-            background_image, foreground_image = self.inputs_preprocess(background_image, foreground_image, foreground_mask, bounding_box)
-            sample = self.prepare_input_encoders(background_image, foreground_image, bounding_box)
-            bg_im = sample['bg'].to(self.device)
-            fg_im = sample['fg'].to(self.device)
-            q_box  = sample['query_box'].to(self.device)
-            c_box  = sample['crop_box'].to(self.device)
-            output = self.model(bg_im, fg_im, q_box, c_box)[-1].item()
-        else:
-            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()
+        
+        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):

libcom-0.1.5.post1/libcom/image_harmonization/image_harmonization.py

@@ -0,0 +1,195 @@
+import torch
+import torchvision
+from libcom.utils.model_download import download_pretrained_model, download_entire_folder
+from libcom.utils.process_image import *
+from libcom.utils.environment import *
+import os
+import torchvision.transforms as transforms
+import sys
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision.transforms import ToTensor
+
+
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+src_dir = os.path.join(cur_dir, 'source/src')
+if src_dir not in sys.path:
+    sys.path.insert(0, src_dir)
+
+from libcom.image_harmonization.source.pct_net import *
+from libcom.image_harmonization.source.cdt_net import *
+
+
+from libcom.image_harmonization.source.src.lbm.inference import get_model
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+# =======================================================
+
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+
+model_set = ['PCTNet', 'LBM'] 
+
+class ImageHarmonizationModel:
+    """
+    Image harmonization model.
+
+    Args:
+        device (str | torch.device): gpu id
+        model_type (str): predefined model type, 'PCTNet' or 'LBM'
+        kwargs (dict): other parameters for building model.
+                       For LBM, you can set 'ckpt_path' here.
+
+    Examples:
+        >>> from libcom import ImageHarmonizationModel
+        >>> import cv2
+        >>> import os
+        >>> import numpy as np
+        >>> from PIL import Image
+        >>> #Use PCTNet
+        >>> PCTNet = ImageHarmonizationModel(device=0, model_type='PCTNet')
+        >>> comp_img1  = '../tests/source/composite/comp1_PCTNet.jpg'
+        >>> comp_mask1 = '../tests/source/composite_mask/mask1_PCTNet.png'
+        >>> PCT_result1 = PCTNet(comp_img1, comp_mask1)
+        >>> cv2.imwrite('../docs/_static/image/image_harmonization_PCT_result1.jpg', np.concatenate([cv2.imread(comp_img1), cv2.imread(comp_mask1), PCT_result1],axis=1))
+        
+        >>> #Use LBM
+        >>> LBM = ImageHarmonizationModel(device=0, model_type='LBM')
+        >>> comp_img  = '../tests/source/composite/1.jpg'
+        >>> comp_mask = '../tests/source/composite_mask/1.png'
+        >>> LBM_result = LBM(comp_img, comp_mask, steps=4)
+        >>> cv2.imwrite('../docs/_static/image/image_harmonization_LBM_result.jpg', np.concatenate([cv2.imread(comp_img), cv2.imread(comp_mask), LBM_result],axis=1))
+
+    Expected result:
+
+    .. image:: _static/image/image_harmonization_PCT_result1.jpg
+    .. image:: _static/image/image_harmonization_LBM_result.jpg
+
+    """
+    def __init__(self, device=0, model_type='PCTNet', **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)
+        
+        if self.model_type == 'PCTNet':
+            
+            weight_path = os.path.join(model_dir, 'pretrained_models', 'PCTNet.pth')
+            download_pretrained_model(weight_path)
+            lut_path = os.path.join(model_dir, 'pretrained_models', 'IdentityLUT33.txt')
+            download_pretrained_model(lut_path)
+            self.build_pretrained_model(weight_path)
+            self.build_data_transformer()
+            
+        elif self.model_type == 'LBM':
+            lbm_dir = os.path.join(model_dir, 'pretrained_models', 'lbm_ckpt')
+            download_entire_folder(lbm_dir)
+            self.build_pretrained_model(lbm_dir)
+
+    def build_pretrained_model(self, weight_path):
+        if self.model_type == 'LBM':
+            self.model = get_model(weight_path, torch_dtype=torch.bfloat16, device=self.device)
+            self.model.bridge_noise_sigma = 0.005
+            if self.model.sampling_noise_scheduler is None:
+                self.model.sampling_noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+                    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler"
+                )
+            self.model.eval()
+        else:
+            model = PCTNet()
+            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.transformer = transforms.Compose([
+            transforms.ToTensor(),
+        ])
+    
+    def inputs_preprocess(self, composite_image, composite_mask):
+        img = read_image_opencv(composite_image)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        mask = read_mask_opencv(composite_mask) / 255.0
+        img_lr = cv2.resize(img, (256, 256))
+        mask_lr = cv2.resize(mask, (256, 256))
+
+        #to tensor
+        img = self.transformer(img).float().to(self.device)
+        mask = self.transformer(mask).float().to(self.device)
+        img_lr = self.transformer(img_lr).float().to(self.device)
+        mask_lr = self.transformer(mask_lr).float().to(self.device)
+        return img, mask, img_lr, mask_lr
+    
+    def outputs_postprocess(self, outputs):
+        if len(outputs.shape) == 4:
+            outputs = outputs.squeeze(0)
+        outputs = (torch.clamp(255.0 * outputs.permute(1, 2, 0), 0, 255)).cpu().numpy()
+        outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
+        return outputs
+    
+    @torch.no_grad()
+    def __call__(self, composite_image, composite_mask, **kwargs):
+        """
+        Given a composite image and a foreground mask, perform harmonization on the foreground.
+
+        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.
+            **kwargs: Extra parameters for inference (e.g., steps=4, resolution=1024 for LBM).
+
+        Returns:
+            harmonized_image (np.array): The harmonized result.
+        
+        """
+        if self.model_type == 'LBM':
+            return self._inference_lbm(composite_image, composite_mask, **kwargs)
+
+        img, mask, img_lr, mask_lr = self.inputs_preprocess(composite_image, composite_mask)
+
+        outputs = self.model(img_lr, img, mask_lr, mask)
+        preds = self.outputs_postprocess(outputs)
+        return preds
+
+    @torch.no_grad()
+    def _inference_lbm(self, composite_image, composite_mask, **kwargs):
+        steps = kwargs.get('steps', self.option.get('steps', 4))
+        inference_size = kwargs.get('resolution', self.option.get('resolution', 1024))
+        latent_size = inference_size // 8
+
+        if isinstance(composite_image, str):
+            src_raw = Image.open(composite_image).convert("RGB")
+        else:
+            src_raw = Image.fromarray(cv2.cvtColor(composite_image, cv2.COLOR_BGR2RGB))
+
+        if isinstance(composite_mask, str):
+            mask_raw = Image.open(composite_mask).convert("L")
+        else:
+            if len(composite_mask.shape) == 3:
+                composite_mask = composite_mask[:, :, 0]
+            mask_raw = Image.fromarray(composite_mask).convert("L")
+
+        ori_w, ori_h = src_raw.size
+
+        src_resized = src_raw.resize((inference_size, inference_size), Image.BILINEAR)
+        src_t = (ToTensor()(src_resized).unsqueeze(0) * 2 - 1).to(self.device, dtype=torch.bfloat16)
+        batch = {"source_image_paste": src_t}
+
+        mask_latent_img = mask_raw.resize((latent_size, latent_size), Image.BILINEAR)
+        mask_t = ToTensor()(mask_latent_img).unsqueeze(0).to(self.device, dtype=torch.bfloat16)
+
+        z_source = self.model.vae.encode(batch["source_image_paste"])
+        
+        output_tensor = self.model.sample(
+            z=z_source,
+            num_steps=int(steps),
+            conditioner_inputs=batch,
+            max_samples=1,
+            mask=mask_t
+        ).clamp(-1, 1)
+
+        res_tensor = (output_tensor[0].float().cpu() + 1) / 2
+        preds = (torch.clamp(255.0 * res_tensor.permute(1, 2, 0), 0, 255)).numpy().astype(np.uint8)
+        preds = cv2.cvtColor(preds, cv2.COLOR_RGB2BGR)
+        
+        preds = cv2.resize(preds, (ori_w, ori_h))
+        
+        return preds

{libcom-0.1.4 → libcom-0.1.5.post1}/libcom/image_harmonization/source/pct_net.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from functools import partial
 
-from .functions import PCT, ViT_Harmonizer
+from libcom.image_harmonization.source.functions import PCT, ViT_Harmonizer
 
 
 class PCTNet(nn.Module):

{libcom-0.1.4 → libcom-0.1.5.post1}/libcom/kontext_blending_harmonization/kontext_blending_harmonization.py

@@ -102,7 +102,7 @@ class KontextBlendingHarmonizationModel:
         fg_resized = fg_img.resize((x2 - x1, y2 - y1))
         # 把前景贴到背景上
         bg_img.paste(fg_resized, (x1, y1, x2, y2))
-
+        bg_img = bg_img.resize((1024, 1024))
         return bg_img
         
     
@@ -145,4 +145,4 @@ class KontextBlendingHarmonizationModel:
             ).images[0]
             comp_image_list.append(cv2.cvtColor(np.array(comp_img), cv2.COLOR_RGB2BGR))
         composite_images = np.stack(comp_image_list, axis=0)
-        return composite_images
+        return composite_images

libcom-0.1.5.post1/libcom/os_insert/__init__.py

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