Rebase conflict resolved

Author: MrFox131

carvekit/api/high.py

@@ -1,6 +1,8 @@
 """
 Source url: https://github.com/OPHoperHPO/image-background-remove-tool
+
 Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
+
 License: Apache License 2.0
 """
 import warnings
@@ -31,26 +33,26 @@ def __init__(
         Initializes High Level interface.
 
         Args:
-            object_type: Interest object type. Can be "object" or "hairs-like".
-            matting_mask_size:  The size of the input image for the matting neural network.
-            seg_mask_size: The size of the input image for the segmentation neural network.
-            batch_size_seg: Number of images processed per one segmentation neural network call.
-            batch_size_matting: Number of images processed per one matting neural network call.
-            device: Processing device
-            fp16: Use half precision. Reduce memory usage and increase speed. Experimental support
-            trimap_prob_threshold: Probability threshold at which the prob_filter and prob_as_unknown_area operations will be applied
-            trimap_dilation: The size of the offset radius from the object mask in pixels when forming an unknown area
-            trimap_erosion_iters: The number of iterations of erosion that the object's mask will be subjected to before forming an unknown area
+            object_type (str, default=object): Interest object type. Can be "object" or "hairs-like".
+            matting_mask_size (int, default=2048):  The size of the input image for the matting neural network.
+            seg_mask_size (int, default=640): The size of the input image for the segmentation neural network.
+            batch_size_seg (int, default=2): Number of images processed per one segmentation neural network call.
+            batch_size_matting (int, matting=1): Number of images processed per one matting neural network call.
+            device (Literal[cpu, cuda], default=cpu): Processing device
+            fp16 (bool, default=False): Use half precision. Reduce memory usage and increase speed.
+            .. CAUTION:: ⚠️ **Experimental support**
+            trimap_prob_threshold (int, default=231): Probability threshold at which the prob_filter and prob_as_unknown_area operations will be applied
+            trimap_dilation (int, default=30): The size of the offset radius from the object mask in pixels when forming an unknown area
+            trimap_erosion_iters (int, default=5): The number of iterations of erosion that the object's mask will be subjected to before forming an unknown area
 
-        Notes:
-            1. Changing seg_mask_size may cause an out-of-memory error if the value is too large, and it may also
-            result in reduced precision. I do not recommend changing this value. You can change matting_mask_size in
-            range from (1024 to 4096) to improve object edge refining quality, but it will cause extra large RAM and
+        .. NOTE::
+            1. Changing seg_mask_size may cause an `out-of-memory` error if the value is too large, and it may also
+            result in reduced precision. I do not recommend changing this value. You can change `matting_mask_size` in
+            range from `(1024 to 4096)` to improve object edge refining quality, but it will cause extra large RAM and
             video memory consume. Also, you can change batch size to accelerate background removal, but it also causes
             extra large video memory consume, if value is too big.
-
-            2. Changing trimap_prob_threshold, trimap_kernel_size, trimap_erosion_iters may improve object edge
-            refining quality,
+            2. Changing `trimap_prob_threshold`, `trimap_kernel_size`, `trimap_erosion_iters` may improve object edge
+            refining quality.
         """
         if object_type == "object":
             self.u2net = TracerUniversalB7(

carvekit/api/interface.py

@@ -31,10 +31,10 @@ def __init__(
         Initializes an object for interacting with pipelines and other components of the CarveKit framework.
 
         Args:
-            pre_pipe: Initialized pre-processing pipeline object
-            seg_pipe: Initialized segmentation network object
-            post_pipe: Initialized postprocessing pipeline object
-            device: The processing device that will be used to apply the masks to the images.
+            pre_pipe (Union[U2NET, BASNET, DeepLabV3, TracerUniversalB7]): Initialized pre-processing pipeline object
+            seg_pipe (Optional[Union[PreprocessingStub]]): Initialized segmentation network object
+            post_pipe (Optional[Union[MattingMethod]]): Initialized postprocessing pipeline object
+            device (Literal[cpu, cuda], default=cpu): The processing device that will be used to apply the masks to the images.
         """
         self.device = device
         self.preprocessing_pipeline = pre_pipe

carvekit/ml/files/models_loc.py

@@ -1,6 +1,8 @@
 """
 Source url: https://github.com/OPHoperHPO/image-background-remove-tool
+
 Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
+
 License: Apache License 2.0
 """
 import pathlib
@@ -12,7 +14,7 @@ def u2net_full_pretrained() -> pathlib.Path:
     """Returns u2net pretrained model location
 
     Returns:
-        pathlib.Path to model location
+        pathlib.Path: model location
     """
     return downloader("u2net.pth")
 
@@ -21,7 +23,7 @@ def basnet_pretrained() -> pathlib.Path:
     """Returns basnet pretrained model location
 
     Returns:
-        pathlib.Path to model location
+        pathlib.Path: model location
     """
     return downloader("basnet.pth")
 
@@ -30,7 +32,7 @@ def deeplab_pretrained() -> pathlib.Path:
     """Returns basnet pretrained model location
 
     Returns:
-        pathlib.Path to model location
+        pathlib.Path: model location
     """
     return downloader("deeplab.pth")
 
@@ -39,7 +41,7 @@ def fba_pretrained() -> pathlib.Path:
     """Returns basnet pretrained model location
 
     Returns:
-        pathlib.Path to model location
+        pathlib.Path: model location
     """
     return downloader("fba_matting.pth")
 
@@ -48,7 +50,7 @@ def tracer_b7_pretrained() -> pathlib.Path:
     """Returns TRACER with EfficientNet v1 b7 encoder pretrained model location
 
     Returns:
-        pathlib.Path to model location
+        pathlib.Path: model location
     """
     return downloader("tracer_b7.pth")
 
@@ -57,7 +59,7 @@ def tracer_hair_pretrained() -> pathlib.Path:
     """Returns TRACER with EfficientNet v1 b7 encoder model for hair segmentation location
 
     Returns:
-        pathlib.Path to model location
+        pathlib.Path: model location
     """
     return downloader("tracer_hair.pth")
 

carvekit/ml/wrap/basnet.py

@@ -1,6 +1,8 @@
 """
 Source url: https://github.com/OPHoperHPO/image-background-remove-tool
+
 Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
+
 License: Apache License 2.0
 """
 import pathlib
@@ -34,12 +36,11 @@ def __init__(
         Initialize the BASNET model
 
         Args:
-            device: processing device
-            input_image_size: input image size
-            batch_size: the number of images that the neural network processes in one run
-            load_pretrained: loading pretrained model
-            fp16: use fp16 precision // not supported at this moment
-
+            device (Literal[cpu, cuda], default=cpu): processing device
+            input_image_size (Union[List[int], int], default=320): input image size
+            batch_size (int, default=10): the number of images that the neural network processes in one run
+            load_pretrained (bool, default=True): loading pretrained model
+            fp16 (bool, default=True): use fp16 precision **not supported at this moment**
         """
         super(BASNET, self).__init__(n_channels=3, n_classes=1)
         self.device = device
@@ -60,10 +61,10 @@ def data_preprocessing(self, data: PIL.Image.Image) -> torch.Tensor:
         Transform input image to suitable data format for neural network
 
         Args:
-            data: input image
+            data (PIL.Image.Image): input image
 
         Returns:
-            input for neural network
+            torch.Tensor: input for neural network
 
         """
         resized = data.resize(self.input_image_size)
@@ -81,18 +82,18 @@ def data_preprocessing(self, data: PIL.Image.Image) -> torch.Tensor:
 
     @staticmethod
     def data_postprocessing(
-        data: torch.tensor, original_image: PIL.Image.Image
+        data: torch.Tensor, original_image: PIL.Image.Image
     ) -> PIL.Image.Image:
         """
         Transforms output data from neural network to suitable data
         format for using with other components of this framework.
 
         Args:
-            data: output data from neural network
-            original_image: input image which was used for predicted data
+            data (torch.Tensor): output data from neural network
+            original_image (PIL.Image.Image): input image which was used for predicted data
 
         Returns:
-            Segmentation mask as PIL Image instance
+            PIL.Image.Image: Segmentation mask as `PIL Image` instance
 
         """
         data = data.unsqueeze(0)
@@ -109,13 +110,13 @@ def __call__(
         self, images: List[Union[str, pathlib.Path, PIL.Image.Image]]
     ) -> List[PIL.Image.Image]:
         """
-        Passes input images through neural network and returns segmentation masks as PIL.Image.Image instances
+        Passes input images through neural network and returns segmentation masks as `PIL.Image.Image` instances
 
         Args:
-            images: input images
+            images (List[Union[str, pathlib.Path, PIL.Image.Image]]): input images
 
         Returns:
-            segmentation masks as for input images, as PIL.Image.Image instances
+            List[PIL.Image.Image]: segmentation masks as for input images, as `PIL.Image.Image` instances
 
         """
         collect_masks = []

carvekit/ml/wrap/deeplab_v3.py

@@ -1,6 +1,8 @@
 """
 Source url: https://github.com/OPHoperHPO/image-background-remove-tool
+
 Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
+
 License: Apache License 2.0
 """
 import pathlib
@@ -29,14 +31,14 @@ def __init__(
         fp16: bool = False,
     ):
         """
-        Initialize the DeepLabV3 model
+        Initialize the `DeepLabV3` model
 
         Args:
-            device: processing device
-            input_image_size: input image size
-            batch_size: the number of images that the neural network processes in one run
-            load_pretrained: loading pretrained model
-            fp16: use half precision
+            device (Literal[cpu, cuda], default=cpu): processing device
+            input_image_size (): input image size
+            batch_size (int, default=10): the number of images that the neural network processes in one run
+            load_pretrained (bool, default=True): loading pretrained model
+            fp16 (bool, default=False): use half precision
 
         """
         self.device = device
@@ -69,9 +71,7 @@ def to(self, device: str):
         Moves neural network to specified processing device
 
         Args:
-            device (:class:`torch.device`): the desired device.
-        Returns:
-            None
+            device (Literal[cpu, cuda]): the desired device.
 
         """
         self.network.to(device)
@@ -81,10 +81,10 @@ def data_preprocessing(self, data: PIL.Image.Image) -> torch.Tensor:
         Transform input image to suitable data format for neural network
 
         Args:
-            data: input image
+            data (PIL.Image.Image): input image
 
         Returns:
-            input for neural network
+            torch.Tensor: input for neural network
 
         """
         copy = data.copy()
@@ -93,18 +93,18 @@ def data_preprocessing(self, data: PIL.Image.Image) -> torch.Tensor:
 
     @staticmethod
     def data_postprocessing(
-        data: torch.tensor, original_image: PIL.Image.Image
+        data: torch.Tensor, original_image: PIL.Image.Image
     ) -> PIL.Image.Image:
         """
         Transforms output data from neural network to suitable data
         format for using with other components of this framework.
 
         Args:
-            data: output data from neural network
-            original_image: input image which was used for predicted data
+            data (torch.Tensor): output data from neural network
+            original_image (PIL.Image.Image): input image which was used for predicted data
 
         Returns:
-            Segmentation mask as PIL Image instance
+            PIL.Image.Image: Segmentation mask as `PIL Image` instance
 
         """
         return (
@@ -115,13 +115,13 @@ def __call__(
         self, images: List[Union[str, pathlib.Path, PIL.Image.Image]]
     ) -> List[PIL.Image.Image]:
         """
-        Passes input images though neural network and returns segmentation masks as PIL.Image.Image instances
+        Passes input images though neural network and returns segmentation masks as `PIL.Image.Image` instances
 
         Args:
-            images: input images
+            images (List[Union[str, pathlib.Path, PIL.Image.Image]]): input images
 
         Returns:
-            segmentation masks as for input images, as PIL.Image.Image instances
+            List[PIL.Image.Image]: segmentation masks as for input images, as `PIL.Image.Image` instances
 
         """
         collect_masks = []

carvekit/ml/wrap/fba_matting.py

@@ -43,12 +43,14 @@ def __init__(
         Initialize the FBAMatting model
 
         Args:
-            device: processing device
-            input_tensor_size: input image size
-            batch_size: the number of images that the neural network processes in one run
-            encoder: neural network encoder head
-            load_pretrained: loading pretrained model
-            fp16: use half precision
+            device (Literal[cpu, cuda], default=cpu): processing device
+            input_tensor_size (Union[List[int], int], default=2048): input image size
+            batch_size (int, default=2): the number of images that the neural network processes in one run
+            encoder (str, default=resnet50_GN_WS): neural network encoder head
+            .. TODO::
+                Add more encoders to documentation as Literal typehint.
+            load_pretrained (bool, default=True): loading pretrained model
+            fp16 (bool, default=False): use half precision
 
         """
         super(FBAMatting, self).__init__(encoder=encoder)
@@ -71,10 +73,10 @@ def data_preprocessing(
         Transform input image to suitable data format for neural network
 
         Args:
-            data: input image
+            data (Union[PIL.Image.Image, np.ndarray]): input image
 
         Returns:
-            input for neural network
+            Tuple[torch.FloatTensor, torch.FloatTensor]: input for neural network
 
         """
         resized = data.copy()
@@ -114,18 +116,18 @@ def data_preprocessing(
 
     @staticmethod
     def data_postprocessing(
-        data: torch.tensor, trimap: PIL.Image.Image
+        data: torch.Tensor, trimap: PIL.Image.Image
     ) -> PIL.Image.Image:
         """
         Transforms output data from neural network to suitable data
         format for using with other components of this framework.
 
         Args:
-            data: output data from neural network
-            trimap: Map with the area we need to refine
+            data (torch.Tensor): output data from neural network
+            trimap (PIL.Image.Image): Map with the area we need to refine
 
         Returns:
-            Segmentation mask as PIL Image instance
+            PIL.Image.Image: Segmentation mask
 
         """
         if trimap.mode != "L":
@@ -149,11 +151,11 @@ def __call__(
         Passes input images though neural network and returns segmentation masks as PIL.Image.Image instances
 
         Args:
-            images: input images
-            trimaps: Maps with the areas we need to refine
+            images (List[Union[str, pathlib.Path, PIL.Image.Image]]): input images
+            trimaps (List[Union[str, pathlib.Path, PIL.Image.Image]]): Maps with the areas we need to refine
 
         Returns:
-            segmentation masks as for input images, as PIL.Image.Image instances
+            List[PIL.Image.Image]: segmentation masks as for input images
 
         """