From b89c7675d5a318a0bdc3c720c3d7e69cfb910d17 Mon Sep 17 00:00:00 2001
From: Pedro <palmelapedro@gmail.com>
Date: Tue, 26 Apr 2022 12:02:28 +0200
Subject: [PATCH] added docs, fixed types

---
 darwin/torch/transforms.py | 98 +++++++++++++++++++++++++++++++++-----
 1 file changed, 86 insertions(+), 12 deletions(-)

diff --git a/darwin/torch/transforms.py b/darwin/torch/transforms.py
index 8a5d0f5f5..4eb6eeadb 100644
--- a/darwin/torch/transforms.py
+++ b/darwin/torch/transforms.py
@@ -13,6 +13,10 @@
 
 
 class Compose(transforms.Compose):
+    """
+    Composes a sequence of Transformations.
+    """
+
     def __call__(self, image: PILImage.Image, target: Optional[TargetType] = None):
         if target is None:
             return super(Compose, self).__call__(image)
@@ -22,9 +26,30 @@ def __call__(self, image: PILImage.Image, target: Optional[TargetType] = None):
 
 
 class RandomHorizontalFlip(transforms.RandomHorizontalFlip):
+    """
+    Allows for horizontal flipping of an image, randomly.
+    """
+
     def forward(
-        self, image: PILImage.Image, target: Optional[TargetType] = None
-    ) -> Union[Union[torch.Tensor, PILImage.Image], Tuple[Union[torch.Tensor, PILImage.Image], TargetType]]:
+        self, image: torch.Tensor, target: Optional[TargetType] = None
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, TargetType]]:
+        """
+        May or may not horizontally flip an image depending on a random factor.
+
+        Parameters
+        ----------
+        image : torch.Tensor
+            Image ``Tensor`` to flip.
+        target : Optional[TargetType] = None
+            The target.
+
+        Returns
+        -------
+        Union[torch.Tensor, Tuple[torch.Tensor, TargetType]]
+            Will return a single image ``Tensor`` if the flip did not happen, or a tuple of the
+            image tensor and the target type if the flip did happen.
+
+        """
         if random.random() < self.p:
             image = F.hflip(image)
             if target is None:
@@ -45,9 +70,30 @@ def forward(
 
 
 class RandomVerticalFlip(transforms.RandomVerticalFlip):
+    """
+    Allows for vertical flipping of an image, randomly.
+    """
+
     def forward(
-        self, image: PILImage.Image, target: Optional[TargetType] = None
-    ) -> Union[Union[torch.Tensor, PILImage.Image], Tuple[Union[torch.Tensor, PILImage.Image], TargetType]]:
+        self, image: torch.Tensor, target: Optional[TargetType] = None
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, TargetType]]:
+        """
+        May or may not vertically flip an image depending on a random factor.
+
+        Parameters
+        ----------
+        image : torch.Tensor
+            Image ``Tensor`` to flip.
+        target : Optional[TargetType] = None
+            The target.
+
+        Returns
+        -------
+        Union[torch.Tensor, Tuple[torch.Tensor, TargetType]]
+            Will return a single image ``Tensor`` if the flip did not happen, or a tuple of the
+            image tensor and the target type if the flip did happen.
+
+        """
         if random.random() < self.p:
             image = F.vflip(image)
             if target is None:
@@ -67,6 +113,10 @@ def forward(
 
 
 class ColorJitter(transforms.ColorJitter):
+    """
+    Jitters the colors of the given transformation.
+    """
+
     def __call__(
         self, image: PILImage.Image, target: Optional[TargetType] = None
     ) -> Union[PILImage.Image, Tuple[PILImage.Image, TargetType]]:
@@ -78,26 +128,38 @@ def __call__(
 
 
 class ToTensor(transforms.ToTensor):
+    """
+    Converts given ``PILImage`` to a ``Tensor``.
+    """
+
     def __call__(
-        self, image: torch.Tensor, target: Optional[TargetType] = None
+        self, image: PILImage.Image, target: Optional[TargetType] = None
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, TargetType]]:
-        image = F.to_tensor(image)
+        image_tensor: torch.Tensor = F.to_tensor(image)
         if target is None:
-            return image
-        return image, target
+            return image_tensor
+        return image_tensor, target
 
 
 class ToPILImage(transforms.ToPILImage):
+    """
+    Converts given ``Tensor`` to a ``PILImage``.
+    """
+
     def __call__(
-        self, image: PILImage.Image, target: Optional[TargetType] = None
+        self, image: torch.Tensor, target: Optional[TargetType] = None
     ) -> Union[PILImage.Image, Tuple[PILImage.Image, TargetType]]:
-        image = F.to_pil_image(image)
+        pil_image: PILImage.Image = F.to_pil_image(image)
         if target is None:
-            return image
-        return image, target
+            return pil_image
+        return pil_image, target
 
 
 class Normalize(transforms.Normalize):
+    """
+    Normalizes the given ``Tensor``.
+    """
+
     def __call__(
         self, tensor: torch.Tensor, target: Optional[TargetType] = None
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, TargetType]]:
@@ -109,6 +171,10 @@ def __call__(
 
 
 class ConvertPolygonsToInstanceMasks(object):
+    """
+    Converts given polygon to an ``InstanceMask``.
+    """
+
     def __call__(self, image: PILImage.Image, target: TargetType) -> Tuple[PILImage.Image, TargetType]:
         w, h = image.size
 
@@ -164,6 +230,10 @@ def __call__(self, image: PILImage.Image, target: TargetType) -> Tuple[PILImage.
 
 
 class ConvertPolygonsToSemanticMask(object):
+    """
+    Converts given polygon to an ``SemanticMask``.
+    """
+
     def __call__(self, image: PILImage.Image, target: TargetType) -> Tuple[PILImage.Image, TargetType]:
         w, h = image.size
         image_id = target["image_id"]
@@ -189,6 +259,10 @@ def __call__(self, image: PILImage.Image, target: TargetType) -> Tuple[PILImage.
 
 
 class ConvertPolygonToMask(object):
+    """
+    Converts given polygon to a ``Mask``.
+    """
+
     def __call__(self, image: PILImage.Image, annotation: Dict[str, Any]) -> Tuple[PILImage.Image, PILImage.Image]:
         w, h = image.size
         segmentations = [obj["segmentation"] for obj in annotation]