carvekit/pipelines/postprocessing.py

"""
Source url: https://github.com/OPHoperHPO/image-background-remove-tool
Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
License: Apache License 2.0
"""
from carvekit.ml.wrap.fba_matting import FBAMatting
from typing import Union, List
from PIL import Image
from pathlib import Path
from carvekit.trimap.cv_gen import CV2TrimapGenerator
from carvekit.trimap.generator import TrimapGenerator
from carvekit.utils.mask_utils import apply_mask
from carvekit.utils.pool_utils import thread_pool_processing
from carvekit.utils.image_utils import load_image, convert_image

__all__ = ["MattingMethod"]


class MattingMethod:
    """
    Improving the edges of the object mask using neural networks for matting and algorithms for creating trimap.
    Neural network for matting performs accurate object edge detection by using a special map called trimap,
    with unknown area that we scan for boundary, already known general object area and the background."""

    def __init__(
        self,
        matting_module: Union[FBAMatting],
        trimap_generator: Union[TrimapGenerator, CV2TrimapGenerator],
        device="cpu",
    ):
        """
        Initializes Matting Method class.

        Args:
            matting_module: Initialized matting neural network class
            trimap_generator: Initialized trimap generator class
            device: Processing device used for applying mask to image
        """
        self.device = device
        self.matting_module = matting_module
        self.trimap_generator = trimap_generator

    def __call__(
        self,
        images: List[Union[str, Path, Image.Image]],
        masks: List[Union[str, Path, Image.Image]],
    ):
        """
        Passes data through apply_mask function

        Args:
            images: list of images
            masks: list pf masks

        Returns:
            list of images
        """
        if len(images) != len(masks):
            raise ValueError("Images and Masks lists should have same length!")
        images = thread_pool_processing(lambda x: convert_image(load_image(x)), images)
        masks = thread_pool_processing(
            lambda x: convert_image(load_image(x), mode="L"), masks
        )
        trimaps = thread_pool_processing(
            lambda x: self.trimap_generator(original_image=images[x], mask=masks[x]),
            range(len(images)),
        )
        alpha = self.matting_module(images=images, trimaps=trimaps)
        return list(
            map(
                lambda x: apply_mask(
                    image=images[x], mask=alpha[x], device=self.device
                ),
                range(len(images)),
            )
        )