libcom 0.1.8__py3-none-any.whl

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

@@ -0,0 +1,36 @@
+from typing import Any, Dict, List, Union
+
+from .base import BaseFilter
+
+
+class FilterWrapper:
+    """
+    Wrapper for multiple filters. This class allows to apply multiple filters to a batch of data.
+    The filters are applied in the order they are passed to the wrapper.
+
+    Args:
+
+        filters (List[BaseFilter]):
+            List of filters to apply to the batch of data
+    """
+
+    def __init__(
+        self,
+        filters: Union[List[BaseFilter], None] = None,
+    ):
+        self.filters = filters
+
+    def __call__(self, batch: Dict[str, Any]) -> None:
+        """
+        Forward pass through all filters
+
+        Args:
+
+            batch: batch of data
+        """
+        filter_output = True
+        for filter in self.filters:
+            filter_output = filter(batch)
+            if not filter_output:
+                return False
+        return True

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

@@ -0,0 +1,33 @@
+import logging
+
+from .base import BaseFilter
+from .filters_config import KeyFilterConfig
+
+logging.basicConfig(level=logging.INFO)
+
+
+class KeyFilter(BaseFilter):
+    """
+    This filter checks if ALL the given keys are present in the sample
+
+    Args:
+
+        config (KeyFilterConfig): configuration for the filter
+    """
+
+    def __init__(self, config: KeyFilterConfig):
+        super().__init__(config)
+        keys = config.keys
+        if isinstance(keys, str):
+            keys = [keys]
+
+        self.keys = set(keys)
+
+    def __call__(self, batch: dict) -> bool:
+        try:
+            res = self.keys.issubset(set(batch.keys()))
+            return res
+        except Exception as e:
+            if self.verbose:
+                logging.error(f"Error in KeyFilter: {e}")
+            return False

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

@@ -0,0 +1,32 @@
+from typing import List, Union
+
+from pydantic.dataclasses import dataclass
+
+from ...config import BaseConfig
+
+
+@dataclass
+class BaseFilterConfig(BaseConfig):
+    """
+    Base configuration for filters
+
+    Args:
+
+        verbose (bool):
+            If True, print debug information. Defaults to False"""
+
+    verbose: bool = False
+
+
+@dataclass
+class KeyFilterConfig(BaseFilterConfig):
+    """
+    This filter checks if the keys are present in a sample.
+
+    Args:
+
+        keys (Union[str, List[str]]):
+            Key or list of keys to check. Defaults to "txt"
+    """
+
+    keys: Union[str, List[str]] = "txt"

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

@@ -0,0 +1,19 @@
+from .base import BaseMapper
+from .mappers import KeyRenameMapper, RescaleMapper, TorchvisionMapper
+from .mappers_config import (
+    KeyRenameMapperConfig,
+    RescaleMapperConfig,
+    TorchvisionMapperConfig,
+)
+from .mappers_wrapper import MapperWrapper
+
+__all__ = [
+    "BaseMapper",
+    "KeyRenameMapper",
+    "RescaleMapper",
+    "TorchvisionMapper",
+    "KeyRenameMapperConfig",
+    "RescaleMapperConfig",
+    "TorchvisionMapperConfig",
+    "MapperWrapper",
+]

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

@@ -0,0 +1,26 @@
+from typing import Any, Dict
+
+from .mappers_config import BaseMapperConfig
+
+
+class BaseMapper:
+    """
+    Base class for the mappers used to modify the samples in the data pipeline.
+
+    Args:
+
+        config (BaseMapperConfig):
+            Configuration for the mapper.
+    """
+
+    def __init__(self, config: BaseMapperConfig):
+        self.config = config
+        self.key = config.key
+
+        if config.output_key is None:
+            self.output_key = config.key
+        else:
+            self.output_key = config.output_key
+
+    def map(self, batch: Dict[str, Any], *args, **kwargs) -> Dict[str, Any]:
+        raise NotImplementedError

libcom/image_harmonization/source/src/lbm/data/mappers/mappers.py

@@ -0,0 +1,135 @@
+from typing import Any, Dict
+
+from torchvision import transforms
+
+from .base import BaseMapper
+from .mappers_config import (
+    KeyRenameMapperConfig,
+    RescaleMapperConfig,
+    TorchvisionMapperConfig,
+)
+
+
+class KeyRenameMapper(BaseMapper):
+    """
+    Rename keys in a sample according to a key map
+
+    Args:
+
+        config (KeyRenameMapperConfig): Configuration for the mapper
+
+    Examples
+    ########
+
+    1. Rename keys in a sample according to a key map
+
+    .. code-block:: python
+
+        from cr.data.mappers import KeyRenameMapper, KeyRenameMapperConfig
+
+        config = KeyRenameMapperConfig(
+            key_map={"old_key": "new_key"}
+        )
+
+        mapper = KeyRenameMapper(config)
+
+        sample = {"old_key": 1}
+        new_sample = mapper(sample)
+        print(new_sample)  # {"new_key": 1}
+
+    2. Rename keys in a sample according to a key map and a condition key
+
+    .. code-block:: python
+
+        from cr.data.mappers import KeyRenameMapper, KeyRenameMapperConfig
+
+        config = KeyRenameMapperConfig(
+            key_map={"old_key": "new_key"},
+            condition_key="condition",
+            condition_fn=lambda x: x == 1
+        )
+
+        mapper = KeyRenameMapper(config)
+
+        sample = {"old_key": 1, "condition": 1}
+        new_sample = mapper(sample)
+        print(new_sample)  # {"new_key": 1}
+
+        sample = {"old_key": 1, "condition": 0}
+        new_sample = mapper(sample)
+        print(new_sample)  # {"old_key": 1}
+
+    ```
+    """
+
+    def __init__(self, config: KeyRenameMapperConfig):
+        super().__init__(config)
+        self.key_map = config.key_map
+        self.condition_key = config.condition_key
+        self.condition_fn = config.condition_fn
+        self.else_key_map = config.else_key_map
+
+    def __call__(self, batch: Dict[str, Any], *args, **kwrags):
+        if self.condition_key is not None:
+            condition_key = batch[self.condition_key]
+            if self.condition_fn(condition_key):
+                for old_key, new_key in self.key_map.items():
+                    if old_key in batch:
+                        batch[new_key] = batch.pop(old_key)
+
+            elif self.else_key_map is not None:
+                for old_key, new_key in self.else_key_map.items():
+                    if old_key in batch:
+                        batch[new_key] = batch.pop(old_key)
+
+        else:
+            for old_key, new_key in self.key_map.items():
+                if old_key in batch:
+                    batch[new_key] = batch.pop(old_key)
+        return batch
+
+
+class TorchvisionMapper(BaseMapper):
+    """
+    Apply torchvision transforms to a sample
+
+    Args:
+
+        config (TorchvisionMapperConfig): Configuration for the mapper
+    """
+
+    def __init__(self, config: TorchvisionMapperConfig):
+        super().__init__(config)
+        chained_transforms = []
+        for transform, kwargs in zip(config.transforms, config.transforms_kwargs):
+            transform = getattr(transforms, transform)
+            chained_transforms.append(transform(**kwargs))
+        self.transforms = transforms.Compose(chained_transforms)
+
+    def __call__(self, batch: Dict[str, Any], *args, **kwrags) -> Dict[str, Any]:
+        if self.key in batch:
+            batch[self.output_key] = self.transforms(batch[self.key])
+        return batch
+
+
+class RescaleMapper(BaseMapper):
+    """
+    Rescale a sample from [0, 1] to [-1, 1]
+
+    Args:
+
+        config (RescaleMapperConfig): Configuration for the mapper
+    """
+
+    def __init__(self, config: RescaleMapperConfig):
+        super().__init__(config)
+
+    def __call__(self, batch: Dict[str, Any], *args, **kwrags) -> Dict[str, Any]:
+        if isinstance(batch[self.key], list):
+            tmp = []
+            for i, image in enumerate(batch[self.key]):
+                tmp.append(2 * image - 1)
+            batch[self.output_key] = tmp
+        else:
+            batch[self.output_key] = 2 * batch[self.key] - 1
+        return batch

libcom/image_harmonization/source/src/lbm/data/mappers/mappers_config.py

@@ -0,0 +1,109 @@
+from typing import Any, Callable, Dict, List, Optional
+
+from pydantic.dataclasses import dataclass
+
+from ...config import BaseConfig
+
+
+@dataclass
+class BaseMapperConfig(BaseConfig):
+    """
+    Base configuration for mappers.
+
+    Args:
+
+        verbose (bool):
+            If True, print debug information. Defaults to False
+
+        key (Optional[str]):
+            Key to apply the mapper to. Defaults to None
+
+        output_key (Optional[str]):
+            Key to store the output of the mapper. Defaults to None
+    """
+
+    verbose: bool = False
+    key: Optional[str] = None
+    output_key: Optional[str] = None
+
+
+@dataclass
+class KeyRenameMapperConfig(BaseMapperConfig):
+    """
+    Rename keys in a sample according to a key map
+
+    Args:
+
+        key_map (Dict[str, str]): Dictionary with the old keys as keys and the new keys as values
+        condition_key (Optional[str]): Key to use for the condition. Defaults to None
+        condition_fn (Optional[Callable[[Any], bool]]): Function to use for the condition to be met so
+            the key map is applied. Defaults to None.
+        else_key_map (Optional[Dict[str, str]]): Dictionary with the old keys as keys and the new keys as values
+            if the condition is not met. Defaults to None *i.e.* the original key will be used.
+    """
+
+    key_map: Dict[str, str] = None
+    condition_key: Optional[str] = None
+    condition_fn: Optional[Callable[[Any], bool]] = None
+    else_key_map: Optional[Dict[str, str]] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.key_map is not None, "key_map should be provided"
+        assert all(
+            isinstance(old_key, str) and isinstance(new_key, str)
+            for old_key, new_key in self.key_map.items()
+        ), "key_map should be a dictionary with string keys and values"
+        if self.condition_key is not None:
+            assert self.condition_fn is not None, "condition_fn should be provided"
+            assert callable(self.condition_fn), "condition_fn should be callable"
+        if self.condition_fn is not None:
+            assert self.condition_key is not None, "condition_key should be provided"
+            assert isinstance(
+                self.condition_key, str
+            ), "condition_key should be a string"
+        if self.else_key_map is not None:
+            assert all(
+                isinstance(old_key, str) and isinstance(new_key, str)
+                for old_key, new_key in self.else_key_map.items()
+            ), "else_key_map should be a dictionary with string keys and values"
+
+
+@dataclass
+class TorchvisionMapperConfig(BaseMapperConfig):
+    """
+    Apply torchvision transforms to a sample
+
+    Args:
+
+        key (str): Key to apply the transforms to
+        transforms (torchvision.transforms): List of torchvision transforms to apply
+        transforms_kwargs (Dict[str, Any]): List of kwargs for the transforms
+    """
+
+    key: str = "image"
+    transforms: List[str] = None
+    transforms_kwargs: List[Dict[str, Any]] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.transforms is None:
+            self.transforms = []
+        if self.transforms_kwargs is None:
+            self.transforms_kwargs = []
+        assert len(self.transforms) == len(
+            self.transforms_kwargs
+        ), "Number of transforms and kwargs should be same"
+
+
+@dataclass
+class RescaleMapperConfig(BaseMapperConfig):
+    """
+    Rescale a sample from [0, 1] to [-1, 1]
+
+    Args:
+
+        key (str): Key to rescale
+    """
+
+    key: str = "image"

libcom/image_harmonization/source/src/lbm/data/mappers/mappers_wrapper.py

@@ -0,0 +1,31 @@
+from typing import Any, Dict, List, Union
+
+from .base import BaseMapper
+
+
+class MapperWrapper:
+    """
+    Wrapper for the mappers to allow iterating over several mappers in one go.
+
+    Args:
+
+        mappers (Union[List[BaseMapper], None]): List of mappers to apply to the batch
+    """
+
+    def __init__(
+        self,
+        mappers: Union[List[BaseMapper], None] = None,
+    ):
+        self.mappers = mappers
+
+    def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Forward pass through all mappers
+
+        Args:
+
+            batch (Dict[str, Any]): batch of data
+        """
+        for mapper in self.mappers:
+            batch = mapper(batch)
+        return batch

libcom/image_harmonization/source/src/lbm/inference/__init__.py

@@ -0,0 +1,4 @@
+from .inference import evaluate
+from .utils import get_model
+
+__all__ = ["evaluate", "get_model"]

libcom/image_harmonization/source/src/lbm/inference/inference.py

@@ -0,0 +1,70 @@
+import logging
+
+import PIL
+import torch
+from torchvision.transforms import ToPILImage, ToTensor
+
+from lbm.models.lbm import LBMModel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ASPECT_RATIOS = {
+    str(512 / 2048): (512, 2048),
+    str(1024 / 1024): (1024, 1024),
+    str(2048 / 512): (2048, 512),
+    str(896 / 1152): (896, 1152),
+    str(1152 / 896): (1152, 896),
+    str(512 / 1920): (512, 1920),
+    str(640 / 1536): (640, 1536),
+    str(768 / 1280): (768, 1280),
+    str(1280 / 768): (1280, 768),
+    str(1536 / 640): (1536, 640),
+    str(1920 / 512): (1920, 512),
+}
+
+
+@torch.no_grad()
+def evaluate(
+    model: LBMModel,
+    source_image: PIL.Image.Image,
+    num_sampling_steps: int = 1,
+):
+    """
+    Evaluate the model on an image coming from the source distribution and generate a new image from the target distribution.
+
+    Args:
+        model (LBMModel): The model to evaluate.
+        source_image (PIL.Image.Image): The source image to evaluate the model on.
+        num_sampling_steps (int): The number of sampling steps to use for the model.
+
+    Returns:
+        PIL.Image.Image: The generated image.
+    """
+
+    ori_h_bg, ori_w_bg = source_image.size
+    ar_bg = ori_h_bg / ori_w_bg
+    closest_ar_bg = min(ASPECT_RATIOS, key=lambda x: abs(float(x) - ar_bg))
+    source_dimensions = ASPECT_RATIOS[closest_ar_bg]
+
+    source_image = source_image.resize(source_dimensions)
+
+    img_pasted_tensor = ToTensor()(source_image).unsqueeze(0) * 2 - 1
+    batch = {
+        "source_image": img_pasted_tensor.cuda().to(torch.bfloat16),
+    }
+
+    z_source = model.vae.encode(batch[model.source_key])
+
+    output_image = model.sample(
+        z=z_source,
+        num_steps=num_sampling_steps,
+        conditioner_inputs=batch,
+        max_samples=1,
+    ).clamp(-1, 1)
+
+    output_image = (output_image[0].float().cpu() + 1) / 2
+    output_image = ToPILImage()(output_image)
+    output_image.resize((ori_h_bg, ori_w_bg))
+
+    return output_image