libcom 0.1.6__tar.gz → 0.1.8__tar.gz

libcom-0.1.8/MANIFEST.in

@@ -0,0 +1,2 @@
+include requirements.txt
+include README.md

{libcom-0.1.6 → libcom-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libcom
-Version: 0.1.6
+Version: 0.1.8
 Summary: Image Composition Toolbox
 Home-page: https://github.com/bcmi/libcom
 Author: BCMI Lab
@@ -14,24 +14,21 @@ Classifier: Programming Language :: Python :: 3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: einops
-Requires-Dist: huggingface_hub
+Requires-Dist: huggingface_hub==0.35.1
 Requires-Dist: imageio
 Requires-Dist: torch==2.6.0
 Requires-Dist: torchvision==0.21.0
 Requires-Dist: diffusers==0.34.0
-Requires-Dist: mmdet==3.2.0
-Requires-Dist: mmpose==1.2.0
-Requires-Dist: peft
+Requires-Dist: peft==0.17.1
 Requires-Dist: sentencepiece
-Requires-Dist: protobuf
-Requires-Dist: modelscope
-Requires-Dist: omegaconf
-Requires-Dist: opencv_python
+Requires-Dist: protobuf==6.32.1
+Requires-Dist: modelscope==1.30.0
+Requires-Dist: omegaconf==2.3.0
 Requires-Dist: opencv_python_headless
-Requires-Dist: pytorch_lightning
+Requires-Dist: pytorch_lightning==2.5.5
 Requires-Dist: transformers==4.44.2
 Requires-Dist: tqdm
-Requires-Dist: open_clip_torch
+Requires-Dist: open_clip_torch==3.2.0
 Requires-Dist: taming-transformers-rom1504
 Requires-Dist: scikit-learn
 Requires-Dist: segment_anything
@@ -62,14 +59,19 @@ 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/)
+![visitors](https://visitor-badge.laobi.icu/badge?page_id=bcmi.libcom&left_color=2ea44f&right_color=green&left_text=visitors)
 [![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/Dataset&Code_Resources-green)](https://github.com/bcmi/Awesome-Image-Composition)
 [![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)
 
 ---
 
+Welcome to scan the WeChat QR code to join the discussion group on image composition/compositing, object insertion/addition.
+<p align="center">
+<img width="150" height="150" alt="Image" src="https://github.com/user-attachments/assets/8572a7c4-0f62-4a5f-bc06-9a1d9fcb14e3" />
+</p>
+
 ## Introduction
 **_libcom_ (the library of image composition) is an image composition toolbox.** The goal of [image composition](https://github.com/bcmi/Awesome-Image-Composition) ([object insertion](https://github.com/bcmi/Awesome-Object-Insertion)) is inserting one foreground into a background image to get a realistic composite image, by addressing the inconsistencies (appearance, geometry, and semantic inconsistency) between foreground and background. Generally speaking, image composition could be used to combine the visual elements from different images.
 <div align="center">
@@ -82,7 +84,9 @@ Dynamic: summary
 **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). 
+**Welcome to recommend your self-developed models or models you have tested by opening an issue. If a model delivers outstanding and stable performance, we will consider integrating it into libcom.**
+
+**If you specialize in model inference acceleration, we welcome both voluntary and paid collaborations. Please contact us at ustcnewly@sjtu.edu.cn.**
 
 ## Online Workbench
 
@@ -97,19 +101,19 @@ We have built the [image composition workbench](https://libcom.ustcnewly.com/) b
 
 ## Main Functions
 
-- **get_composite_image** generates composite images using naive copy-and-paste followed by image blending.
-- **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.  
+- **get_composite_image** generates composite images using naive copy-and-paste followed by image blending (alpha blending or Poisson blending).
+- **OPAScoreModel** [[OPA]](https://github.com/bcmi/Object-Placement-Assessment-Dataset-OPA) evaluates the rationality of foreground object placement (location, scale) 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, when placing the foreground with certain scale at certain location in the background. Based on the predicted rationality score map, the model can 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) [[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.
+- **HarmonyScoreModel** [[BargainNet]](https://github.com/bcmi/BargainNet-Image-Harmonization) evaluates the harmony level between foreground and background in a composite image, that is, whether the illumination of foreground is harmonious with that of background. 
 - **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. 
+- **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, that is, whether the geometric perspective of foreground matches that of background, and whether the semantic context of foreground in the background is reasonable.
 - **ShadowGenerationModel** [[GPSDiffusion]](https://github.com/bcmi/GPSDiffusion-Object-Shadow-Generation) generates plausible shadow 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. 
-- **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. 
+- **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. You can also use foreground shape mask for controllable composition. 
 
 **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. 
 

{libcom-0.1.6 → libcom-0.1.8}/README.md

@@ -10,14 +10,19 @@
 
 [![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/)
+![visitors](https://visitor-badge.laobi.icu/badge?page_id=bcmi.libcom&left_color=2ea44f&right_color=green&left_text=visitors)
 [![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/Dataset&Code_Resources-green)](https://github.com/bcmi/Awesome-Image-Composition)
 [![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)
 
 ---
 
+Welcome to scan the WeChat QR code to join the discussion group on image composition/compositing, object insertion/addition.
+<p align="center">
+<img width="150" height="150" alt="Image" src="https://github.com/user-attachments/assets/8572a7c4-0f62-4a5f-bc06-9a1d9fcb14e3" />
+</p>
+
 ## Introduction
 **_libcom_ (the library of image composition) is an image composition toolbox.** The goal of [image composition](https://github.com/bcmi/Awesome-Image-Composition) ([object insertion](https://github.com/bcmi/Awesome-Object-Insertion)) is inserting one foreground into a background image to get a realistic composite image, by addressing the inconsistencies (appearance, geometry, and semantic inconsistency) between foreground and background. Generally speaking, image composition could be used to combine the visual elements from different images.
 <div align="center">
@@ -30,7 +35,9 @@
 **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). 
+**Welcome to recommend your self-developed models or models you have tested by opening an issue. If a model delivers outstanding and stable performance, we will consider integrating it into libcom.**
+
+**If you specialize in model inference acceleration, we welcome both voluntary and paid collaborations. Please contact us at ustcnewly@sjtu.edu.cn.**
 
 ## Online Workbench
 
@@ -45,19 +52,19 @@ We have built the [image composition workbench](https://libcom.ustcnewly.com/) b
 
 ## Main Functions
 
-- **get_composite_image** generates composite images using naive copy-and-paste followed by image blending.
-- **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.  
+- **get_composite_image** generates composite images using naive copy-and-paste followed by image blending (alpha blending or Poisson blending).
+- **OPAScoreModel** [[OPA]](https://github.com/bcmi/Object-Placement-Assessment-Dataset-OPA) evaluates the rationality of foreground object placement (location, scale) 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, when placing the foreground with certain scale at certain location in the background. Based on the predicted rationality score map, the model can 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) [[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.
+- **HarmonyScoreModel** [[BargainNet]](https://github.com/bcmi/BargainNet-Image-Harmonization) evaluates the harmony level between foreground and background in a composite image, that is, whether the illumination of foreground is harmonious with that of background. 
 - **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. 
+- **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, that is, whether the geometric perspective of foreground matches that of background, and whether the semantic context of foreground in the background is reasonable.
 - **ShadowGenerationModel** [[GPSDiffusion]](https://github.com/bcmi/GPSDiffusion-Object-Shadow-Generation) generates plausible shadow 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. 
-- **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. 
+- **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. You can also use foreground shape mask for controllable composition. 
 
 **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. 
 

{libcom-0.1.6 → libcom-0.1.8}/libcom/__init__.py

@@ -3,7 +3,7 @@ from .naive_composition import *
 from .opa_score import *
 from .harmony_score import *
 from .inharmonious_region_localization import *
-from .image_harmonization_old import *
+from .image_harmonization import *
 from .painterly_image_harmonization import *
 from .fopa_heat_map import *
 from .fos_score import *

{libcom-0.1.6 → libcom-0.1.8}/libcom/image_harmonization/image_harmonization.py

@@ -18,8 +18,6 @@ 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

libcom-0.1.8/libcom/image_harmonization/source/src/lbm/config.py

@@ -0,0 +1,141 @@
+import json
+import os
+import warnings
+from dataclasses import asdict, field
+from typing import Any, Dict, Union
+
+import yaml
+from pydantic import ValidationError
+from pydantic.dataclasses import dataclass
+from yaml import safe_load
+
+
+@dataclass
+class BaseConfig:
+    """This is the BaseConfig class which defines all the useful loading and saving methods
+    of the configs"""
+
+    name: str = field(init=False)
+
+    def __post_init__(self):
+        self.name = self.__class__.__name__
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any]) -> "BaseConfig":
+        """Creates a BaseConfig instance from a dictionnary
+
+        Args:
+            config_dict (dict): The Python dictionnary containing all the parameters
+
+        Returns:
+            :class:`BaseConfig`: The created instance
+        """
+        try:
+            config = cls(**config_dict)
+        except (ValidationError, TypeError) as e:
+            raise e
+        return config
+
+    @classmethod
+    def _dict_from_json(cls, json_path: Union[str, os.PathLike]) -> Dict[str, Any]:
+        try:
+            with open(json_path) as f:
+                try:
+                    config_dict = json.load(f)
+                    return config_dict
+
+                except (TypeError, json.JSONDecodeError) as e:
+                    raise TypeError(
+                        f"File {json_path} not loadable. Maybe not json ? \n"
+                        f"Catch Exception {type(e)} with message: " + str(e)
+                    ) from e
+
+        except FileNotFoundError:
+            raise FileNotFoundError(
+                f"Config file not found. Please check path '{json_path}'"
+            )
+
+    @classmethod
+    def from_json(cls, json_path: str) -> "BaseConfig":
+        """Creates a BaseConfig instance from a JSON config file
+
+        Args:
+            json_path (str): The path to the json file containing all the parameters
+
+        Returns:
+            :class:`BaseConfig`: The created instance
+        """
+        config_dict = cls._dict_from_json(json_path)
+
+        config_name = config_dict.pop("name")
+
+        if cls.__name__ != config_name:
+            warnings.warn(
+                f"You are trying to load a "
+                f"`{ cls.__name__}` while a "
+                f"`{config_name}` is given."
+            )
+
+        return cls.from_dict(config_dict)
+
+    def to_dict(self) -> dict:
+        """Transforms object into a Python dictionnary
+
+        Returns:
+            (dict): The dictionnary containing all the parameters"""
+        return asdict(self)
+
+    def to_json_string(self):
+        """Transforms object into a JSON string
+
+        Returns:
+            (str): The JSON str containing all the parameters"""
+        return json.dumps(self.to_dict())
+
+    def save_json(self, file_path: str):
+        """Saves a ``.json`` file from the dataclass
+
+        Args:
+            file_path (str): path to the file
+        """
+        with open(os.path.join(file_path), "w", encoding="utf-8") as fp:
+            fp.write(self.to_json_string())
+
+    def save_yaml(self, file_path: str):
+        """Saves a ``.yaml`` file from the dataclass
+
+        Args:
+            file_path (str): path to the file
+        """
+        with open(os.path.join(file_path), "w", encoding="utf-8") as fp:
+            yaml.dump(self.to_dict(), fp)
+
+    @classmethod
+    def from_yaml(cls, yaml_path: str) -> "BaseConfig":
+        """Creates a BaseConfig instance from a YAML config file
+
+        Args:
+            yaml_path (str): The path to the yaml file containing all the parameters
+
+        Returns:
+            :class:`BaseConfig`: The created instance
+        """
+        with open(yaml_path, "r") as f:
+            try:
+                config_dict = safe_load(f)
+            except yaml.YAMLError as e:
+                raise yaml.YAMLError(
+                    f"File {yaml_path} not loadable. Maybe not yaml ? \n"
+                    f"Catch Exception {type(e)} with message: " + str(e)
+                ) from e
+
+        config_name = config_dict.pop("name")
+
+        if cls.__name__ != config_name:
+            warnings.warn(
+                f"You are trying to load a "
+                f"`{ cls.__name__}` while a "
+                f"`{config_name}` is given."
+            )
+
+        return cls.from_dict(config_dict)

libcom-0.1.8/libcom/image_harmonization/source/src/lbm/data/__init__.py

@@ -0,0 +1,62 @@
+"""
+This module contains a collection of data related classes and functions to train the :mod:`cr.models`.
+In a training loop a batch of data is struvtued as a dictionnary on which the modules :mod:`cr.data.datasets`
+and :mod:`cr.data.filters` allow to perform several operations.
+
+
+Examples
+########
+
+Create a DataModule to train a model
+
+.. code-block::python
+
+    from cr.data import DataModule, DataModuleConfig
+    from cr.data.filters import KeyFilter, KeyFilterConfig
+    from cr.data.mappers import KeyRenameMapper, KeyRenameMapperConfig
+
+    # Create the filters and mappers
+    filters_mappers = [
+        KeyFilter(KeyFilterConfig(keys=["image", "txt"])),
+        KeyRenameMapper(
+            KeyRenameMapperConfig(key_map={"jpg": "image", "txt": "text"})
+        )
+    ]
+
+    # Create the DataModule
+    data_module = DataModule(
+        train_config=DataModuleConfig(
+            shards_path_or_urls="your urls or paths",
+            decoder="pil",
+            shuffle_buffer_size=100,
+            per_worker_batch_size=32,
+            num_workers=4,
+        ),
+        train_filters_mappers=filters_mappers,
+        eval_config=DataModuleConfig(
+            shards_path_or_urls="your urls or paths",
+            decoder="pil",
+            shuffle_buffer_size=100,
+            per_worker_batch_size=32,
+            num_workers=4,
+        ),
+        eval_filters_mappers=filters_mappers,
+    )
+
+    # This can then be passed to a :mod:`pytorch_lightning.Trainer` to train a model
+
+
+
+
+
+The :mod:`cr.data` includes the following submodules:
+
+- :mod:`cr.data.datasets`: a collection of :mod:`pytorch_lightning.LightningDataModule` used to train the models. In particular,
+    they can used to create the dataloaders and setup the data pipelines.
+- :mod:`cr.data.filters`: a collection of filters used apply filters on a training batch of data/
+
+"""
+
+from .datasets import DataModule
+
+__all__ = ["DataModule"]

libcom-0.1.8/libcom/image_harmonization/source/src/lbm/data/datasets/__init__.py

@@ -0,0 +1,9 @@
+"""
+A collection of :mod:`pytorch_lightning.LightningDataModule` used to train the models. In particular,
+they can be used to create the dataloaders and setup the data pipelines.
+"""
+
+from .dataset import DataModule
+from .datasets_config import DataModuleConfig
+
+__all__ = ["DataModule", "DataModuleConfig"]

libcom-0.1.8/libcom/image_harmonization/source/src/lbm/data/datasets/collation_fn.py

@@ -0,0 +1,41 @@
+from typing import Dict, List, Union
+
+import numpy as np
+import torch
+
+
+def custom_collation_fn(
+    samples: List[Dict[str, Union[int, float, np.ndarray, torch.Tensor]]],
+    combine_tensors: bool = True,
+    combine_scalars: bool = True,
+) -> dict:
+    """
+    Collate function for PyTorch DataLoader.
+
+    Args:
+        samples(List[Dict[str, Union[int, float, np.ndarray, torch.Tensor]]]): List of samples.
+        combine_tensors (bool): Whether to turn lists of tensors into a single tensor.
+        combine_scalars (bool): Whether to turn lists of scalars into a single ndarray.
+    """
+    keys = set.intersection(*[set(sample.keys()) for sample in samples])
+    batched = {key: [] for key in keys}
+    for s in samples:
+        [batched[key].append(s[key]) for key in batched]
+
+    result = {}
+    for key in batched:
+        if isinstance(batched[key][0], (int, float)):
+            if combine_scalars:
+                result[key] = np.array(list(batched[key]))
+        elif isinstance(batched[key][0], torch.Tensor):
+            if combine_tensors:
+                result[key] = torch.stack(list(batched[key]))
+        elif isinstance(batched[key][0], np.ndarray):
+            if combine_tensors:
+                result[key] = np.array(list(batched[key]))
+        else:
+            result[key] = list(batched[key])
+
+    del samples
+    del batched
+    return result