libcom 0.1.8__py3-none-any.whl

libcom/image_harmonization/source/pct_net.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+
+from libcom.image_harmonization.source.functions import PCT, ViT_Harmonizer
+
+
+class PCTNet(nn.Module):
+    
+    def __init__(
+        self,
+        backbone_type='ViT', 
+        input_normalization = {'mean': [0.0, 0.0, 0.0], 'std':[1.0, 1.0, 1.0]},
+        dim=3, transform_type='linear', affine=True,
+        clamp=True, color_space = 'RGB', use_attn = False
+    ):
+        super(PCTNet, self).__init__()
+              
+        self.color_space = color_space
+        self.use_attn = use_attn
+
+        self.dim = dim
+        self.transform_type = transform_type
+        self.affine = affine
+        self.PCT = PCT(transform_type, dim, affine, color_space, input_normalization['mean'], input_normalization['std'], clamp=clamp)
+        self.out_dim = self.PCT.get_out_dim()
+
+        input_dim = self.out_dim
+        self.backbone_type = backbone_type
+        self.encoder = ViT_Harmonizer(output_nc=input_dim)
+        self.decoder = lambda intermediates, img, mask: (intermediates, mask)
+
+        self.get_params = nn.Conv2d(input_dim, self.out_dim, kernel_size=1)
+
+
+    def forward(self, image, image_fullres=None, mask=None, mask_fullres=None, backbone_features=None):
+        self.device = image.get_device()
+  
+        image = image.unsqueeze(0)    
+        mask = mask.unsqueeze(0)   
+        image_fullres = image_fullres.unsqueeze(0)   
+        mask_fullres = mask_fullres.unsqueeze(0)   
+        x = torch.cat((image, mask), dim=1)
+        
+        intermediates = self.encoder(x, backbone_features)
+        latent, attention_map = self.decoder(intermediates, image, mask) # (N, 32, 256, 256), (N, 1, 256, 256)
+        params = self.get_params(latent)
+
+        output_lowres = self.PCT(image, params)
+
+        if self.use_attn:
+            output_lowres = output_lowres * attention_map + image * (1-attention_map)
+        else:
+            output_lowres = output_lowres * mask + image * (1 - mask)    
+
+        # Full resolution branch
+        if torch.is_tensor(image_fullres):
+            fr_imgs = [image_fullres]
+            fr_masks = [mask_fullres]
+            idx = [[n for n in range(image_fullres.shape[0])]]
+        else:
+            fr_imgs = [img_fr.unsqueeze(0).to(image.get_device()) for img_fr in image_fullres]
+            fr_masks = [mask_fr.unsqueeze(0).to(image.get_device()) for mask_fr in mask_fullres]
+            idx = [[n] for n in range(len(image_fullres))]
+
+        out_fr = []
+        param_fr = []
+        for id, fr_img, fr_mask in zip(idx, fr_imgs, fr_masks):
+            H = fr_img.size(2)
+            W = fr_img.size(3)
+            params_fullres = F.interpolate(params[id], size=(H, W), mode='bicubic', align_corners=True)
+            output_fullres = self.PCT(fr_img, params_fullres)
+            if self.use_attn:
+                attention_map_fullres = F.interpolate(attention_map[id], size=(H, W), mode='bicubic', align_corners=True)
+                output_fullres = output_fullres * attention_map_fullres + fr_img * (1-attention_map_fullres)
+            else:
+                output_fullres = output_fullres * fr_mask + fr_img * (1 - fr_mask)
+            out_fr.append(output_fullres.squeeze())
+            param_fr.append(params_fullres.squeeze())
+        
+        if len(out_fr) == 1:
+            out_fr = out_fr[0]
+            
+        return out_fr

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/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/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/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

libcom/image_harmonization/source/src/lbm/data/datasets/dataset.py

@@ -0,0 +1,243 @@
+from typing import Callable, List, Union
+
+import pytorch_lightning as pl
+import webdataset as wds
+from webdataset import DataPipeline
+
+from ..filters import BaseFilter, FilterWrapper
+from ..mappers import BaseMapper, MapperWrapper
+from .collation_fn import custom_collation_fn
+from .datasets_config import DataModuleConfig
+
+
+class DataPipeline:
+    """
+    DataPipeline class for creating a dataloader from a single configuration
+
+    Args:
+
+        config (DataModuleConfig):
+            Configuration for the dataset
+
+        filters_mappers (Union[List[Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]]):
+            List of filters and mappers for the dataset. These will be sequentially applied.
+
+        batched_filters_mappers (List[Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]]):
+            List of batched transforms for the dataset. These will be sequentially applied.
+    """
+
+    def __init__(
+        self,
+        config: DataModuleConfig,
+        filters_mappers: List[
+            Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]
+        ],
+        batched_filters_mappers: List[
+            Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]
+        ] = None,
+    ):
+        self.config = config
+        self.shards_path_or_urls = config.shards_path_or_urls
+        self.filters_mappers = filters_mappers
+        self.batched_filters_mappers = batched_filters_mappers or []
+
+        if filters_mappers is None:
+            filters_mappers = []
+
+        # set processing pipeline
+        self.processing_pipeline = [wds.decode(config.decoder, handler=config.handler)]
+        self.processing_pipeline.extend(
+            self._add_filters_mappers(
+                filters_mappers=filters_mappers,
+                handler=config.handler,
+            )
+        )
+
+    def _add_filters_mappers(
+        self,
+        filters_mappers: List[
+            Union[
+                FilterWrapper,
+                MapperWrapper,
+            ]
+        ],
+        handler: Callable = wds.warn_and_continue,
+    ) -> List[Union[FilterWrapper, MapperWrapper]]:
+        tmp_pipeline = []
+        for filter_mapper in filters_mappers:
+            if isinstance(filter_mapper, FilterWrapper) or isinstance(
+                filter_mapper, BaseFilter
+            ):
+                tmp_pipeline.append(wds.select(filter_mapper))
+            elif isinstance(filter_mapper, MapperWrapper) or isinstance(
+                filter_mapper, BaseMapper
+            ):
+                tmp_pipeline.append(wds.map(filter_mapper, handler=handler))
+            elif isinstance(filter_mapper) or isinstance(filter_mapper):
+                tmp_pipeline.append(wds.map(filter_mapper, handler=handler))
+            else:
+                raise ValueError("Unknown type of filter/mapper")
+        return tmp_pipeline
+
+    def setup(self):
+        pipeline = [wds.SimpleShardList(self.shards_path_or_urls)]
+
+        # shuffle before split by node
+        if self.config.shuffle_before_split_by_node_buffer_size is not None:
+            pipeline.append(
+                wds.shuffle(
+                    self.config.shuffle_before_split_by_node_buffer_size,
+                    handler=self.config.handler,
+                )
+            )
+        # split by node
+        pipeline.append(wds.split_by_node)
+
+        # shuffle before split by workers
+        if self.config.shuffle_before_split_by_workers_buffer_size is not None:
+            pipeline.append(
+                wds.shuffle(
+                    self.config.shuffle_before_split_by_workers_buffer_size,
+                    handler=self.config.handler,
+                )
+            )
+        # split by worker
+        pipeline.extend(
+            [
+                wds.split_by_worker,
+                wds.tarfile_to_samples(
+                    handler=self.config.handler,
+                    rename_files=self.config.rename_files_fn,
+                ),
+            ]
+        )
+
+        # shuffle before filter mappers
+        if self.config.shuffle_before_filter_mappers_buffer_size is not None:
+            pipeline.append(
+                wds.shuffle(
+                    self.config.shuffle_before_filter_mappers_buffer_size,
+                    handler=self.config.handler,
+                )
+            )
+
+        # apply filters and mappers
+        pipeline.extend(self.processing_pipeline)
+
+        # shuffle after filter mappers
+        if self.config.shuffle_after_filter_mappers_buffer_size is not None:
+            pipeline.append(
+                wds.shuffle(
+                    self.config.shuffle_after_filter_mappers_buffer_size,
+                    handler=self.config.handler,
+                ),
+            )
+
+        # batching
+        pipeline.append(
+            wds.batched(
+                self.config.per_worker_batch_size,
+                collation_fn=custom_collation_fn,
+            )
+        )
+
+        # apply batched transforms
+        pipeline.extend(
+            self._add_filters_mappers(
+                filters_mappers=self.batched_filters_mappers,
+                handler=self.config.handler,
+            )
+        )
+
+        # create the data pipeline
+        pipeline = wds.DataPipeline(*pipeline, handler=self.config.handler)
+
+        # set the pipeline
+        self.pipeline = pipeline
+
+    def dataloader(self):
+        # return the loader
+        return wds.WebLoader(
+            self.pipeline,
+            batch_size=None,
+            num_workers=self.config.num_workers,
+        )
+
+
+class DataModule(pl.LightningDataModule):
+    """
+    Main DataModule class for creating data loaders and training/evaluating models
+
+    Args:
+
+        train_config (DataModuleConfig):
+            Configuration for the training dataset
+
+        train_filters_mappers (Union[List[Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]]):
+            List of filters and mappers for the training dataset. These will be sequentially applied.
+
+        train_batched_filters_mappers (List[Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]]):
+            List of batched transforms for the training dataset. These will be sequentially applied.
+
+        eval_config (DataModuleConfig):
+            Configuration for the evaluation dataset
+
+        eval_filters_mappers (List[Union[FilterWrapper, MapperWrapper]]):
+            List of filters and mappers for the evaluation dataset.These will be sequentially applied.
+
+        eval_batched_filters_mappers (List[Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]]):
+            List of batched transforms for the evaluation dataset. These will be sequentially applied.
+    """
+
+    def __init__(
+        self,
+        train_config: DataModuleConfig,
+        train_filters_mappers: List[
+            Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]
+        ] = None,
+        train_batched_filters_mappers: List[
+            Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]
+        ] = None,
+        eval_config: DataModuleConfig = None,
+        eval_filters_mappers: List[Union[FilterWrapper, MapperWrapper]] = None,
+        eval_batched_filters_mappers: List[
+            Union[BaseMapper, BaseFilter, FilterWrapper, MapperWrapper]
+        ] = None,
+    ):
+        super().__init__()
+
+        self.train_config = train_config
+        self.train_filters_mappers = train_filters_mappers
+        self.train_batched_filters_mappers = train_batched_filters_mappers
+
+        self.eval_config = eval_config
+        self.eval_filters_mappers = eval_filters_mappers
+        self.eval_batched_filters_mappers = eval_batched_filters_mappers
+
+    def setup(self, stage=None):
+        """
+        Setup the data module and create the webdataset processing pipelines
+        """
+
+        # train pipeline
+        self.train_pipeline = DataPipeline(
+            config=self.train_config,
+            filters_mappers=self.train_filters_mappers,
+            batched_filters_mappers=self.train_batched_filters_mappers,
+        )
+        self.train_pipeline.setup()
+
+        # eval pipeline
+        if self.eval_config is not None:
+            self.eval_pipeline = DataPipeline(
+                config=self.eval_config,
+                filters_mappers=self.eval_filters_mappers,
+                batched_filters_mappers=self.eval_batched_filters_mappers,
+            )
+            self.eval_pipeline.setup()
+
+    def train_dataloader(self):
+        return self.train_pipeline.dataloader()
+
+    def val_dataloader(self):
+        return self.eval_pipeline.dataloader()

libcom/image_harmonization/source/src/lbm/data/datasets/datasets_config.py

@@ -0,0 +1,42 @@
+from typing import Callable, List, Optional, Union
+
+import webdataset as wds
+from pydantic.dataclasses import dataclass
+
+from ...config import BaseConfig
+
+
+@dataclass
+class DataModuleConfig(BaseConfig):
+    """
+    Configuration for the DataModule
+
+    Args:
+
+        shards_path_or_urls (Union[str, List[str]]): The path or url to the shards. Defaults to None.
+        per_worker_batch_size (int): The batch size for the dataset. Defaults to 16.
+        num_workers (int): The number of workers to use. Defaults to 1.
+        shuffle_before_split_by_node_buffer_size (Optional[int]): The buffer size for the shuffle before split by node. Defaults to 100.
+        shuffle_before_split_by_workers_buffer_size (Optional[int]): The buffer size for the shuffle before split by workers. Defaults to 100.
+        shuffle_before_filter_mappers_buffer_size (Optional[int]): The buffer size for the shuffle before filter mappers. Defaults to 1000.
+        shuffle_after_filter_mappers_buffer_size (Optional[int]): The buffer size for the shuffle after filter mappers. Defaults to 1000.
+        decoder (str): The decoder to use. Defaults to "pil".
+        handler (Callable): A callable to handle the warnings. Defaults to wds.warn_and_continue.
+        rename_files_fn (Optional[Callable[[str], str]]): A callable to rename the files. Defaults to None.
+    """
+
+    shards_path_or_urls: Union[str, List[str]] = None
+    per_worker_batch_size: int = 16
+    num_workers: int = 1
+    shuffle_before_split_by_node_buffer_size: Optional[int] = 100
+    shuffle_before_split_by_workers_buffer_size: Optional[int] = 100
+    shuffle_before_filter_mappers_buffer_size: Optional[int] = 1000
+    shuffle_after_filter_mappers_buffer_size: Optional[int] = 1000
+    decoder: str = "pil"
+    handler: Callable = wds.warn_and_continue
+    rename_files_fn: Optional[Callable[[str], str]] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.rename_files_fn is not None:
+            assert callable(self.rename_files_fn), "rename_files must be a callable"

libcom/image_harmonization/source/src/lbm/data/filters/__init__.py

@@ -0,0 +1,12 @@
+from .base import BaseFilter
+from .filter_wrapper import FilterWrapper
+from .filters import KeyFilter
+from .filters_config import BaseFilterConfig, KeyFilterConfig
+
+__all__ = [
+    "BaseFilter",
+    "FilterWrapper",
+    "KeyFilter",
+    "BaseFilterConfig",
+    "KeyFilterConfig",
+]

libcom/image_harmonization/source/src/lbm/data/filters/base.py

@@ -0,0 +1,21 @@
+from typing import Any, Dict
+
+from .filters_config import BaseFilterConfig
+
+
+class BaseFilter:
+    """
+    Base class for filters. This class should be subclassed to create a new filter.
+
+    Args:
+
+        config (BaseFilterConfig):
+            Configuration for the filter
+    """
+
+    def __init__(self, config: BaseFilterConfig):
+        self.verbose = config.verbose
+
+    def __call__(self, sample: Dict[str, Any]) -> bool:
+        """This function should be implemented by the subclass"""
+        raise NotImplementedError