diff --git a/core/setup.py b/core/setup.py
index bba6af3..4e132a6 100644
--- a/core/setup.py
+++ b/core/setup.py
@@ -147,7 +147,7 @@ def setup_cfg(args, cfg: Optional[CfgNode] = None) -> CfgNode:
 
     assert all(
         len(i) == 2 and all(isinstance(i_i, str) for i_i in i) for i in cfg.PREPROCESS.OUTPUT
-    ), "PREPROCESS.OUTPUT must be a list of tuples with two strings"
+    ), "PREPROCESS.OUTPUT must be a list of tuples with two strings. Example [['image', 'png'], ['sem_seg', 'png']]"
 
     # Deprecation warnings
     if cfg.PREPROCESS.RESIZE.USE:
diff --git a/data/augmentations.py b/data/augmentations.py
index e85ea03..5a3f178 100644
--- a/data/augmentations.py
+++ b/data/augmentations.py
@@ -25,12 +25,23 @@
 
 
 class TimedAugmentationList(T.AugmentationList):
+    """
+    A wrapper class for the AugmentationList class that times the __call__ method.
+    """
+
     def __init_subclass__(cls) -> None:
+        """
+        Initialize the subclass by adding the timed wrapper to the __call__ method.
+        """
         super().__init_subclass__()
         cls.__call__ = cls._timed(cls.__call__)
 
     @classmethod
     def _timed(cls, func):
+        """
+        A decorator that times the function call.
+        """
+
         def wrapper(*args, **kwargs):
             start = time.perf_counter()
             result = func(*args, **kwargs)
@@ -53,7 +64,7 @@ def __init__(self, tfm_or_aug: T.Augmentation | T.Transform, prob=0.5) -> None:
         Randomly apply an augmentation to an image with a given probability.
 
         Args:
-            tfm_or_aug (Augmentation | Transform): transform or augmentation to apply
+            tfm_or_aug (Augmentation | Transform): transform or augmentation to apply.
             prob (float, optional): probability between 0.0 and 1.0 that
                 the wrapper transformation is applied. Defaults to 0.5.
         """
@@ -91,13 +102,13 @@ def __repr__(self):
 class Augmentation(T.Augmentation):
     def get_transform_aug_input(self, aug_input: T.AugInput) -> T.Transform:
         """
-        Get the transform from the input
+        Get the transform from the input.
 
         Args:
-            aug_input (T.AugInput): input to the augmentation
+            aug_input (T.AugInput): input to the augmentation.
 
         Returns:
-            T.Transform: transform
+            T.Transform: transform.
         """
         args = _get_aug_input_args(self, aug_input)
         transform = self.get_transform(*args)
@@ -105,11 +116,11 @@ def get_transform_aug_input(self, aug_input: T.AugInput) -> T.Transform:
 
     def get_output_shape(self, old_height: int, old_width: int, dpi: Optional[int] = None) -> tuple[int, int]:
         """
-        Get the output shape of the image
+        Get the output shape of the image.
 
         Args:
-            old_height (int): height of the image
-            old_width (int): width of the image
+            old_height (int): height of the image.
+            old_width (int): width of the image.
             dpi (Optional[int], optional): dpi of the image. Defaults to None.
 
         Returns:
@@ -121,10 +132,10 @@ def get_output_shape(self, old_height: int, old_width: int, dpi: Optional[int] =
 class ResizeScaling(Augmentation):
     def __init__(self, scale: float, max_size: Optional[int] = None, target_dpi: Optional[int] = None) -> None:
         """
-        Resize the image by a given scale
+        Resize the image by a given scale.
 
         Args:
-            scale (float): scale percentage
+            scale (float): scale percentage.
             max_size (Optional[int], optional): max size of the image. Defaults to None.
             target_dpi (Optional[int], optional): target dpi of the image. Defaults to None.
         """
@@ -201,7 +212,7 @@ def __init__(
         edge_length: Optional[int] = None,
     ) -> None:
         """
-        Resize image alternative using cv2 instead of PIL or Pytorch
+        Resize image alternative using cv2 instead of PIL or Pytorch.
 
         Args:
             min_size (int | Sequence[int]): The minimum length of the side.
@@ -414,7 +425,7 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 
 class RandomElastic(Augmentation):
     """
-    Apply a random elastic transformation to the image, made using random warpfield and gaussian filters
+    Apply a random elastic transformation to the image, made using random warpfield and gaussian filters.
     """
 
     def __init__(
@@ -424,7 +435,7 @@ def __init__(
         ignore_value: int = 255,
     ) -> None:
         """
-        Apply a random elastic transformation to the image, made using random warpfield and gaussian filters
+        Apply a random elastic transformation to the image, made using random warpfield and gaussian filters.
 
         Args:
             alpha (int, optional): scale factor of the warpfield (sets max value). Defaults to 0.045.
@@ -491,7 +502,7 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 # IDEA Use super class for RandomAffine, RandomTranslation, RandomRotation, RandomShear, RandomScale
 class RandomAffine(Augmentation):
     """
-    Apply a random affine transformation to the image
+    Apply a random affine transformation to the image.
     """
 
     def __init__(
@@ -504,7 +515,7 @@ def __init__(
         ignore_value: int = 255,
     ) -> None:
         """
-        Apply a random affine transformation to the image
+        Apply a random affine transformation to the image.
 
         Args:
             t_stdv (float, optional): standard deviation used for the translation. Defaults to 0.02.
@@ -529,14 +540,14 @@ def __init__(
 
     def get_random_matrix(self, height: int, width: int) -> np.ndarray:
         """
-        Get a random affine transformation matrix
+        Get a random affine transformation matrix.
 
         Args:
             height (int): image height
             width (int): image width
 
         Returns:
-            np.ndarray: random affine transformation matrix
+            np.ndarray: random affine transformation matrix.
         """
         center = np.eye(3)
         center[:2, 2:] = np.asarray([width, height])[:, None] / 2
@@ -612,12 +623,12 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 
 class RandomTranslation(Augmentation):
     """
-    Apply a random translation to the image
+    Apply a random translation to the image.
     """
 
     def __init__(self, t_stdv: float = 0.02, ignore_value: int = 255) -> None:
         """
-        Apply a random affine transformation to the image
+        Apply a random affine transformation to the image.
 
         Args:
             t_stdv (float, optional): standard deviation used for the translation. Defaults to 0.02.
@@ -629,14 +640,14 @@ def __init__(self, t_stdv: float = 0.02, ignore_value: int = 255) -> None:
 
     def get_random_matrix(self, height: int, width: int) -> np.ndarray:
         """
-        Get a random translation matrix
+        Get a random translation matrix.
 
         Args:
-            height (int): image height
-            width (int): image width
+            height (int): image height.
+            width (int): image width.
 
         Returns:
-            np.ndarray: random translation matrix
+            np.ndarray: random translation matrix.
         """
         matrix = np.eye(3)
 
@@ -672,12 +683,12 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 
 class RandomRotation(Augmentation):
     """
-    Apply a random rotation to the image
+    Apply a random rotation to the image.
     """
 
     def __init__(self, r_kappa: float = 30, ignore_value: int = 255) -> None:
         """
-        Apply a random rotation to the image
+        Apply a random rotation to the image.
 
         Args:
             r_kappa (float, optional): kappa value used for sampling the rotation. Defaults to 30.
@@ -689,14 +700,14 @@ def __init__(self, r_kappa: float = 30, ignore_value: int = 255) -> None:
 
     def get_random_matrix(self, height: int, width: int) -> np.ndarray:
         """
-        Get a random rotation matrix
+        Get a random rotation matrix.
 
         Args:
-            height (int): image height
-            width (int): image width
+            height (int): image height.
+            width (int): image width.
 
         Returns:
-            np.ndarray: random rotation matrix
+            np.ndarray: random rotation matrix.
         """
 
         center = np.eye(3)
@@ -740,12 +751,12 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 
 class RandomShear(Augmentation):
     """
-    Apply a random shearing to the image
+    Apply a random shearing to the image.
     """
 
     def __init__(self, sh_kappa: float = 20, ignore_value: int = 255) -> None:
         """
-        Apply a random shearing to the image
+        Apply a random shearing to the image.
 
         Args:
             sh_kappa (float, optional): kappa value used for sampling the shear. Defaults to 20.
@@ -757,14 +768,14 @@ def __init__(self, sh_kappa: float = 20, ignore_value: int = 255) -> None:
 
     def get_random_matrix(self, height: int, width: int) -> np.ndarray:
         """
-        Get a random shearing matrix
+        Get a random shearing matrix.
 
         Args:
-            height (int): image height
-            width (int): image width
+            height (int): image height.
+            width (int): image width.
 
         Returns:
-            np.ndarray: random shearing matrix
+            np.ndarray: random shearing matrix.
         """
         center = np.eye(3)
         center[:2, 2:] = np.asarray([width, height])[:, None] / 2
@@ -818,12 +829,12 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 
 class RandomScale(Augmentation):
     """
-    Apply a random scaling to the image
+    Apply a random scaling to the image.
     """
 
     def __init__(self, sc_stdv: float = 0.12, ignore_value: int = 255) -> None:
         """
-        Apply a random scaling to the image
+        Apply a random scaling to the image.
 
         Args:
             sc_stdv (float, optional): standard deviation used for the scale. Defaults to 0.12.
@@ -835,14 +846,14 @@ def __init__(self, sc_stdv: float = 0.12, ignore_value: int = 255) -> None:
 
     def get_random_matrix(self, height: int, width: int) -> np.ndarray:
         """
-        Get a random scaling matrix
+        Get a random scaling matrix.
 
         Args:
-            height (int): image height
-            width (int): image width
+            height (int): image height.
+            width (int): image width.
 
         Returns:
-            np.ndarray: random scaling matrix
+            np.ndarray: random scaling matrix.
         """
         center = np.eye(3)
         center[:2, 2:] = np.asarray([width, height])[:, None] / 2
@@ -884,12 +895,12 @@ def get_transform(self, image: np.ndarray) -> T.Transform:
 
 class Grayscale(Augmentation):
     """
-    Randomly convert the image to grayscale
+    Randomly convert the image to grayscale.
     """
 
     def __init__(self, image_format="RGB") -> None:
         """
-        Randomly convert the image to grayscale
+        Randomly convert the image to grayscale.
 
         Args:
             image_format (str, optional): Color formatting. Defaults to "RGB".
@@ -908,12 +919,12 @@ def get_transform(self, image) -> T.Transform:
 
 class Invert(Augmentation):
     """
-    Invert the image
+    Invert the image.
     """
 
     def __init__(self, max_value=255) -> None:
         """
-        Invert the image
+        Invert the image.
         """
         super().__init__()
         self.max_value = max_value
@@ -935,12 +946,12 @@ def get_transform(self, image) -> T.Transform:
 
 class RandomJPEGCompression(Augmentation):
     """
-    Apply JPEG compression to the image
+    Apply JPEG compression to the image.
     """
 
     def __init__(self, min_quality: int = 40, max_quality: int = 100) -> None:
         """
-        Apply JPEG compression to the image
+        Apply JPEG compression to the image.
 
         Args:
             quality_range (tuple[int, int], optional): range of the quality of the image. Defaults to (40, 100).
@@ -972,12 +983,12 @@ def get_transform(self, image) -> T.Transform:
 
 class RandomGaussianFilter(Augmentation):
     """
-    Apply random gaussian kernels
+    Apply random gaussian kernels.
     """
 
     def __init__(self, min_sigma: float = 1, max_sigma: float = 3, order: int = 0, iterations: int = 1) -> None:
         """
-        Apply random gaussian kernels
+        Apply random gaussian kernels.
 
         Args:
             min_sigma (float, optional): min Gaussian deviation. Defaults to 1.
@@ -1003,12 +1014,12 @@ def get_transform(self, image) -> T.Transform:
 
 class RandomNoise(Augmentation):
     """
-    Apply random noise to the image
+    Apply random noise to the image.
     """
 
     def __init__(self, min_noise_std: float = 10, max_noise_std: float = 32) -> None:
         """
-        Apply random noise to the image
+        Apply random noise to the image.
 
         Args:
             min_noise_std (float, optional): min noise standard deviation. Defaults to 10.
@@ -1039,7 +1050,7 @@ def get_transform(self, image) -> T.Transform:
 
 class RandomSaturation(Augmentation):
     """
-    Change the saturation of an image
+    Change the saturation of an image.
 
     Saturation intensity is uniformly sampled in (intensity_min, intensity_max).
     - intensity < 1 will reduce saturation
@@ -1049,11 +1060,11 @@ class RandomSaturation(Augmentation):
 
     def __init__(self, intensity_min: float = 0.5, intensity_max: float = 1.5, image_format="RGB") -> None:
         """
-        Change the saturation of an image
+        Change the saturation of an image.
 
         Args:
-            intensity_min (float): Minimum augmentation
-            intensity_max (float): Maximum augmentation
+            intensity_min (float): Minimum augmentation.
+            intensity_max (float): Maximum augmentation.
         """
         super().__init__()
 
@@ -1093,7 +1104,7 @@ def get_transform(self, image) -> T.Transform:
 
 class RandomContrast(Augmentation):
     """
-    Randomly transforms image contrast
+    Randomly transforms image contrast.
 
     Contrast intensity is uniformly sampled in (intensity_min, intensity_max).
     - intensity < 1 will reduce contrast
@@ -1105,11 +1116,11 @@ class RandomContrast(Augmentation):
 
     def __init__(self, intensity_min: float = 0.5, intensity_max: float = 1.5):
         """
-        Randomly transforms image contrast
+        Randomly transforms image contrast.
 
         Args:
-            intensity_min (float): Minimum augmentation
-            intensity_max (float): Maximum augmentation
+            intensity_min (float): Minimum augmentation.
+            intensity_max (float): Maximum augmentation.
         """
         super().__init__()
 
@@ -1135,7 +1146,7 @@ def get_transform(self, image):
 
 class RandomBrightness(Augmentation):
     """
-    Randomly transforms image brightness
+    Randomly transforms image brightness.
 
     Brightness intensity is uniformly sampled in (intensity_min, intensity_max).
     - intensity < 1 will reduce brightness
@@ -1150,8 +1161,8 @@ def __init__(self, intensity_min: float = 0.5, intensity_max: float = 1.5):
         Randomly transforms image brightness.
 
         Args:
-            intensity_min (float): Minimum augmentation
-            intensity_max (float): Maximum augmentation
+            intensity_min (float): Minimum augmentation.
+            intensity_max (float): Maximum augmentation.
         """
         super().__init__()
 
@@ -1179,7 +1190,7 @@ def get_transform(self, image):
 
 class RandomHue(Augmentation):
     """
-    Randomly transforms image hue
+    Randomly transforms image hue.
     """
 
     def __init__(self, hue_delta_min: float = -0.5, hue_delta_max: float = 0.5, image_format="RGB") -> None:
@@ -1209,12 +1220,12 @@ def get_transform(self, image) -> T.Transform:
 
 class AdaptiveThresholding(Augmentation):
     """
-    Apply Adaptive thresholding to the image
+    Apply Adaptive thresholding to the image.
     """
 
     def __init__(self, image_format="RGB") -> None:
         """
-        Apply Adaptive thresholding to the image
+        Apply Adaptive thresholding to the image.
 
         Args:
             image_format (str, optional): Color formatting. Defaults to "RGB".
@@ -1233,7 +1244,7 @@ def get_transform(self, image) -> T.Transform:
 
 class RandomOrientation(Augmentation):
     """
-    Apply a random orientation to the image
+    Apply a random orientation to the image.
     """
 
     def __init__(self, orientation_percentages: Optional[list[float | int]] = None) -> None:
@@ -1353,7 +1364,7 @@ class RandomCrop(T.Augmentation):
     Randomly crop a rectangle region out of an image.
     """
 
-    def __init__(self, crop_type: str, crop_size):
+    def __init__(self, crop_type: str, crop_size: tuple[float, float]):
         """
         Args:
             crop_type (str): one of "relative_range", "relative", "absolute", "absolute_range".
@@ -1400,8 +1411,10 @@ def get_transform(self, image) -> T.Transform:
 
     def get_crop_size(self, image_size) -> tuple[int, int]:
         """
+        Returns the crop size in absolute pixels.
+
         Args:
-            image_size (tuple): height, width
+            image_size (tuple): height, width.
 
         Returns:
             crop_size (tuple): height, width in absolute pixels
@@ -1441,13 +1454,13 @@ class RandomCrop_CategoryAreaConstraint(T.Augmentation):
     def __init__(
         self,
         crop_type: str,
-        crop_size,
+        crop_size: tuple[float, float],
         single_category_max_area: float = 1.0,
         ignored_category: Optional[int] = None,
     ):
         """
         Args:
-            crop_type, crop_size: same as in :class:`RandomCrop`
+            crop_type (str), crop_size (tuple[float,float]): same as in :class:`RandomCrop`
             single_category_max_area: the maximum allowed area ratio of a
                 category. Set to 1.0 to disable
             ignored_category: allow this category in the semantic segmentation
@@ -1509,16 +1522,36 @@ def get_transform(self, image, sem_seg) -> T.Transform:
 
 
 class RandomCropResize(T.Augmentation):
+    """
+    Combine a random crop and a resize operation. This is useful to prevent resizing the whole image.
+    """
+
     def __init__(
         self,
-        crop_type,
-        crop_size,
-        resize_mode,
-        scale=None,
-        target_dpi=None,
-        min_size=None,
-        max_size=None,
+        crop_type: str,
+        crop_size: tuple[float, float],
+        resize_mode: str,
+        scale: Optional[float] = None,
+        target_dpi: Optional[int] = None,
+        min_size: Optional[int] = None,
+        max_size: Optional[int] = None,
     ):
+        """
+        Combine a random crop and a resize operation. This is useful to prevent resizing the whole image.
+
+        Args:
+            crop_type (str), crop_size (tuple[float, float]) : same as in :class:`RandomCrop`
+            resize_mode (str): one of "none", "shortest_edge", "longest_edge", "scaling"
+            scale (float): scaling factor, only used when resize_mode == "scaling"
+            target_dpi (int): target dpi, only used when resize_mode == "none" or "scaling"
+            min_size (int): min size, only used when resize_mode == "shortest_edge" or "longest_edge"
+            max_size (int): max size, only used when resize_mode == "shortest_edge" or "longest_edge"
+
+        Raises:
+            ValueError: If the resize mode specified in the configuration is not recognized.
+            ValueError: If the min_size and max_size are not provided for shortest_edge or longest_edge mode.
+            ValueError: If the scale and max_size are not provided for scaling mode.
+        """
         if resize_mode == "none":
             self.resize = ResizeScaling(scale=1.0, target_dpi=target_dpi)
         if resize_mode == "shortest_edge":
diff --git a/data/dataset.py b/data/dataset.py
index d7a39a2..e1a6211 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -85,16 +85,16 @@ def dict_of_list_to_list_of_dicts(input_dict: dict[str, list[Any]]) -> list[dict
     return output_list
 
 
-def convert_to_paths(dataset_dir: Path, input_data: dict[str, list]) -> list[dict[str, Path | Any]]:
+def convert_to_paths(dataset_dir: Path, input_data: dict[str, list]) -> list[dict[str, Path]]:
     """
-    Turn expected paths into actual Path objects instead of just str, the rest stays the same
+    Turn expected paths into actual Path objects instead of just strings
 
     Args:
         dataset_dir (Path): base dataset dir
-        input_data (dict[str, list]): data dict with some of the lists representing paths
+        input_data (dict[str, list]): data dict with lists of paths
 
     Returns:
-        list[dict[str, Path | Any]]: list of dicts containing paths where applicable
+        list[dict[str, Path]]: list of dicts containing paths
     """
     converted_data = dict_of_list_to_list_of_dicts(input_data)
     converted_data = [{key: dataset_dir.joinpath(value) for key, value in item.items()} for item in converted_data]
diff --git a/data/mapper.py b/data/mapper.py
index e13509f..47d1645 100644
--- a/data/mapper.py
+++ b/data/mapper.py
@@ -47,7 +47,16 @@ def _transform_to_aug(tfm_or_aug):
         return _TransformToAug(tfm_or_aug)
 
 
-def _check_img_dtype(img):
+def _check_img_dtype(img: np.ndarray | torch.Tensor) -> None:
+    """
+    Check if the image is a numpy array or torch tensor.
+
+    Args:
+        img: The image to check.
+
+    Raises:
+        ValueError: If the image is not a numpy array or torch tensor.
+    """
     if isinstance(img, torch.Tensor):
         assert img.dtype == torch.uint8 or img.dtype == torch.float32, f"[Augmentation] Got image of type {img.dtype}!"
         assert img.dim() == 3, img.dim()
@@ -59,6 +68,10 @@ def _check_img_dtype(img):
 
 
 class AugInput(T.AugInput):
+    """
+    A class that holds input data for augmentation.
+    """
+
     def __init__(
         self,
         image: np.ndarray | torch.Tensor,
@@ -70,6 +83,18 @@ def __init__(
         default_dpi: Optional[int] = None,
         manual_dpi: Optional[int] = None,
     ):
+        """
+        Initialize the AugInput class.
+
+        Args:
+            image (np.ndarray | torch.Tensor): The image to augment.
+            boxes (Optional[np.ndarray], optional): The boxes to augment. Defaults to None.
+            sem_seg (Optional[np.ndarray  |  torch.Tensor], optional): The semantic segmentation to augment. Defaults to None.
+            dpi (Optional[int], optional): The DPI of the image. Defaults to None.
+            auto_dpi (Optional[bool], optional): Whether to automatically detect the DPI. Defaults to False.
+            default_dpi (Optional[int], optional): The default DPI to use. Defaults to None.
+            manual_dpi (Optional[int], optional): The manual DPI to use. Defaults to None.
+        """
 
         _check_img_dtype(image)
         self.image = image
@@ -106,6 +131,8 @@ def apply_augmentations(self, augmentations: list[T.Augmentation | T.Transform])
 def check_image_size(dataset_dict, image):
     """
     Raise an error if the image does not match the size specified in the dict.
+
+    If height or width is not specified in the dict, it will be filled with the image's shape.
     """
     if "width" in dataset_dict or "height" in dataset_dict:
         if isinstance(image, torch.Tensor):
@@ -137,6 +164,10 @@ def check_image_size(dataset_dict, image):
 
 
 class Mapper(DatasetMapper):
+    """
+    Base class for mappers that handle data augmentation.
+    """
+
     @configurable
     def __init__(
         self,
@@ -272,6 +303,10 @@ def __call__(self, dataset_dict):
 
 
 class SemSegMapper(Mapper):
+    """
+    Specialized mapper for semantic segmentation tasks.
+    """
+
     @configurable
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -347,6 +382,10 @@ def __call__(self, dataset_dict):
 
 
 class SemSegInstancesMapper(Mapper):
+    """
+    Specialized mapper for semantic segmentation tasks with instance annotations.
+    """
+
     @configurable
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -425,6 +464,10 @@ def __call__(self, dataset_dict):
 
 
 class BinarySegMapper(Mapper):
+    """
+    Specialized mapper for binary segmentation tasks.
+    """
+
     @configurable
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
diff --git a/data/numpy_transforms.py b/data/numpy_transforms.py
index d763fa7..fd6b3c9 100644
--- a/data/numpy_transforms.py
+++ b/data/numpy_transforms.py
@@ -10,7 +10,17 @@
 
 
 class TimedTransform(T.Transform):
+    """
+    A wrapper class to time the transforms.
+
+    Args:
+        T (T.Transform): Transform class to be timed
+    """
+
     def __init_subclass__(cls) -> None:
+        """
+        Initialize the subclass. Apply the timing to the class methods.
+        """
         super().__init_subclass__()
         cls.apply_image = cls._timed(cls.apply_image)
         cls.apply_segmentation = cls._timed(cls.apply_segmentation)
@@ -19,6 +29,13 @@ def __init_subclass__(cls) -> None:
 
     @classmethod
     def _timed(cls, func):
+        """
+        Decorator to time the function
+
+        Args:
+            func (function): function to be timed
+        """
+
         def wrapper(*args, **kwargs):
             start = time.perf_counter()
             result = func(*args, **kwargs)
@@ -371,12 +388,13 @@ def apply_image(self, img: np.ndarray) -> np.ndarray:
             np.ndarray: transformed image
         """
         img = img.astype(np.uint8)
+        border_value = [0] if len(img.shape) == 2 else [0] * img.shape[-1]
         return cv2.warpAffine(
             img,
             self.matrix[:2, :],
             (img.shape[1], img.shape[0]),
             flags=cv2.INTER_LINEAR,
-            borderValue=0,
+            borderValue=border_value,
         )
 
     def apply_coords(self, coords: np.ndarray) -> np.ndarray:
@@ -403,12 +421,13 @@ def apply_segmentation(self, segmentation: np.ndarray) -> np.ndarray:
             np.ndarray: transformed segmentation
         """
         # borderValue=0 means background, borderValue=255 means ignored
+        border_value = [255] if len(segmentation.shape) == 2 else [self.ignore_value] * segmentation.shape[-1]
         return cv2.warpAffine(
             segmentation,
             self.matrix[:2, :],
             (segmentation.shape[1], segmentation.shape[0]),
             flags=cv2.INTER_NEAREST,
-            borderValue=self.ignore_value,
+            borderValue=border_value,
         )
 
     def inverse(self) -> T.Transform:
@@ -858,6 +877,8 @@ def inverse(self) -> T.Transform:
 class JPEGCompressionTransform(T.Transform):
     def __init__(self, quality: int):
         """
+        Compress the image using JPEG compression. This is useful for simulating artifacts in low-quality images.
+
         Args:
             quality (int): JPEG compression quality. A number between 0 and 100.
         """
@@ -920,6 +941,8 @@ def __init__(
         orig_h: Optional[int] = None,
     ):
         """
+        Crop the image(s) as a transform.
+
         Args:
             x0, y0, w, h (int): crop the image(s) by img[y0:y0+h, x0:x0+w].
             orig_w, orig_h (int): optional, the original width and height
@@ -1028,6 +1051,8 @@ def __init__(
         seg_pad_value: int = 0,
     ):
         """
+        Pad the images as a transform.
+
         Args:
             x0, y0: number of padded pixels on the left and top
             x1, y1: number of padded pixels on the right and bottom
diff --git a/data/preprocess.py b/data/preprocess.py
index 009d203..d1c201d 100644
--- a/data/preprocess.py
+++ b/data/preprocess.py
@@ -63,7 +63,7 @@ def get_arguments() -> argparse.Namespace:
 
 class Preprocess:
     """
-    Used for almost all preprocessing steps to prepare datasets to be used by the training loop
+    Used for almost all preprocessing steps to prepare datasets to be used by the training loop.
     """
 
     @configurable
@@ -183,7 +183,7 @@ def get_parser(cls) -> argparse.ArgumentParser:
         Return argparser that has the arguments required for the preprocessing.
 
         Returns:
-            argparse.ArgumentParser: the argparser for preprocessing
+            argparse.ArgumentParser: the argparser for preprocessing.
         """
         parser = argparse.ArgumentParser(add_help=False)
         pre_process_args = parser.add_argument_group("preprocessing")
@@ -226,10 +226,10 @@ def set_input_paths(
         ignore_duplicates: bool = False,
     ) -> None:
         """
-        Setter of the input paths, turn string to path. And resolve full path
+        Setter of the input paths, turn string to path. And resolve full path.
 
         Args:
-            input_paths (str | Path | Sequence[str  |  Path]): path(s) from which to extract the images
+            input_paths (str | Path | Sequence[str  |  Path]): path(s) from which to extract the images.
             ignore_duplicates (bool, optional): Ignore duplicate names in the input paths. Defaults to False.
         """
         input_paths = get_file_paths(input_paths, supported_image_formats, self.disable_check)
@@ -282,19 +282,19 @@ def check_duplicates(
 
     def get_input_paths(self) -> Optional[Sequence[Path]]:
         """
-        Getter of the input paths
+        Getter of the input paths.
 
         Returns:
-            Optional[Sequence[Path]]: path(s) from which to extract the images
+            Optional[Sequence[Path]]: path(s) from which to extract the images.
         """
         return self.input_paths
 
     def set_output_dir(self, output_dir: str | Path) -> None:
         """
-        Setter of output dir, turn string to path. And resolve full path
+        Setter of output dir, turn string to path. And resolve full path.
 
         Args:
-            output_dir (str | Path): output path of the processed images
+            output_dir (str | Path): output path of the processed images.
         """
         if isinstance(output_dir, str):
             output_dir = Path(output_dir)
@@ -317,20 +317,19 @@ def get_output_dir(self) -> Optional[Path]:
     @staticmethod
     def check_paths_exists(paths: Sequence[Path]) -> None:
         """
-        Check if all paths given exist and are readable
-
+        Check if all paths given exist and are readable.
         Args:
-            paths (list[Path]): paths to be checked
+            paths (list[Path]): paths to be checked.
         """
         all(check_path_accessible(path) for path in paths)
 
     def save_array_to_path(self, array: np.ndarray | torch.Tensor, path: Path) -> None:
         """
-        Save an array to a path with a specific method
+        Save an array to a path with a specific method.
 
         Args:
-            array (np.ndarray): array to be saved
-            path (Path): path to save the array
+            array (np.ndarray): array to be saved.
+            path (Path): path to save the array.
         """
 
         method = path.suffix
diff --git a/data/preprocess_coco.py b/data/preprocess_coco.py
index dfd4989..f384f78 100644
--- a/data/preprocess_coco.py
+++ b/data/preprocess_coco.py
@@ -464,7 +464,7 @@ def process_single_file(self, image_path: Path) -> dict:
 
         return results
 
-    def run(self) -> Path:
+    def run(self):
         """
         Run preprocessing on all images currently on input paths, save to output dir
 
@@ -517,8 +517,6 @@ def run(self) -> Path:
                 )
             )
 
-        return None
-
 
 def list_of_dict_to_dict_of_list(input_list: list[dict[str, Any]]) -> dict[str, list[Any]]:
     """
diff --git a/data/test.py b/data/test.py
deleted file mode 100644
index c7f0472..0000000
--- a/data/test.py
+++ /dev/null
@@ -1,80 +0,0 @@
-_sum = sum(
-    [
-        5.587800114881247e-05,
-        6.416998076019809e-06,
-        9.657000191509724e-06,
-        0.012764667000737973,
-        0.0063344160007545725,
-        0.007565762996819103,
-        0.00273162400117144,
-        0.011208104999241186,
-        1.1949996405746788e-06,
-        0.00031858100192039274,
-        0.00024359000235563144,
-        0.0002492499988875352,
-        0.00021513100000447594,
-        0.00021272400044836104,
-        0.0002167509992432315,
-        0.0015182800016191322,
-        6.990012479946017e-07,
-        0.0019864770001731813,
-        1.0140029189642519e-06,
-        0.0026487370014365297,
-        1.1990014172624797e-06,
-        0.03411611299816286,
-        8.450006134808064e-07,
-        0.004043042001285357,
-        1.1810006981249899e-06,
-        0.004515399999945657,
-        6.890004442539066e-07,
-        0.0016385679991799407,
-        1.0760013537947088e-06,
-        0.011471959001937648,
-        9.800023690331727e-07,
-        0.002451975997246336,
-        1.2540003808680922e-06,
-    ]
-)
-print(_sum)
-print(0.9876395419996697 - _sum)
-
-_sum = sum(
-    [
-        2.969799970742315e-05,
-        1.7359998309984803e-06,
-        4.665998858399689e-06,
-        0.008956777001003502,
-        0.00044370899922796525,
-        0.19733788799931062,
-        0.024980511003377615,
-        0.03417967600034899,
-        1.0799994925037026e-06,
-        0.00010624999777064659,
-        8.179199721780606e-05,
-        9.058800060302019e-05,
-        0.0030348259970196523,
-        0.00035102599940728396,
-        0.0003587179999158252,
-        0.010983324998960597,
-        6.359987310133874e-07,
-        0.004484345998207573,
-        7.33998604118824e-07,
-        0.02561083100226824,
-        9.190007403958589e-07,
-        0.007476146998669719,
-        3.5150005714967847e-06,
-        0.01888903299914091,
-        6.330010364763439e-07,
-        0.004940410999552114,
-        5.299989425111562e-07,
-        0.0015438919981534127,
-        6.800000846851617e-07,
-        0.0031471509973926004,
-        7.419985195156187e-07,
-        0.007156359002692625,
-        1.3349999790079892e-06,
-    ]
-)
-
-print(_sum)
-print(0.49114640400148346 - _sum)
diff --git a/environment_dev.yml b/environment_dev.yml
index c5663c8..13edeb8 100644
--- a/environment_dev.yml
+++ b/environment_dev.yml
@@ -25,6 +25,7 @@ dependencies:
   - pillow
   - shapely
   - natsort
+  - ultralytics
   # - jpeg # For loading JPEG2000 images
   # - pygments #Optional for colors
   - pip: