Feature/data augment

README.md

@@ -184,11 +184,6 @@ poetry run coverage run -m pytest -s && poetry run coverage report -m
 ```
 
 ## Mkdocs Documentation:
-You need to install the following packages:
-```shell
-pip install mkdocs
-pip install mkdocstrings
-```
 To run the mkdocs documentation, you can run the following lines below:
 ```shell
 cd src

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "plastic-origins"
-version = "2.3.0a0"
+version = "2.4.0a0"
 
 description = "A package containing methods commonly used to make inferences"
 repository = "https://github.com/surfriderfoundationeurope/surfnet"
@@ -39,6 +39,9 @@ azure-identity = "^1.12.0"
 azure-keyvault-secrets = "^4.6.0"
 python-dotenv = "^0.21.0"
 tensorflow = "^2.12.0"
+mkdocstrings = "^0.22.0"
+mkdocstrings-python = "^1.5.1"
+mkdocs-material = "^9.2.3"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.1"

src/plasticorigins/training/data/DA_for_GenData.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+import argparse
+from argparse import Namespace
+import os
+
+
+#############
+# COPY AND RENAME OLD TRAIN AND VAL TXT
+# FILL NEW TRAIN TXT by excuting :
+# python src/plasticorigins/training/data/DA_for_GenData.py --data-dir /datadrive/data/data_20062022 --artificial-data /datadrive/data/artificial_data
+# RENAME WITH _DA SUFFIX AFTER
+#############
+
+
+def main(args: Namespace) -> None:
+    """Main Function to write new images paths for training.
+
+    Args:
+        args (argparse): list of arguments to build dataset for label mapping and training
+    """
+
+    data_dir = Path(args.data_dir)
+
+    # use data augmentation for artificial data only if original data have been processed
+    artificial_data_dir = Path(args.artificial_data)
+    artificial_train_files = [Path(path).as_posix() for path in os.listdir(artificial_data_dir / "images")]
+
+    # concatenate original images and artificial data
+    with open(data_dir / "train.txt", "w") as f:
+        for path in artificial_train_files:
+            f.write(path + "\n")
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Build dataset")
+    parser.add_argument("--data-dir", type=str, help="path to main data folder")
+    parser.add_argument("--artificial-data", type=str, help="path to artificial data folder")
+
+    args = parser.parse_args()
+    main(args)

src/plasticorigins/training/data/data_processing.py

@@ -58,7 +58,8 @@
 import numpy as np
 from numpy import ndarray, array
 from sklearn.model_selection import train_test_split
-import tensorflow as tf
+import torchvision.transforms as transforms
+from torchvision.utils import save_image
 
 import pandas as pd
 from pandas import DataFrame
@@ -281,6 +282,7 @@ def flip_left_right_annotations(bboxes: np.ndarray) -> np.ndarray:
         bboxes (np.ndarray): array of bounding boxes after flip
     """
 
+    bboxes = bboxes.copy()
     bboxes[:, 0] = 1 - bboxes[:, 0]
     return bboxes
 
@@ -296,7 +298,7 @@ def flip_up_down_annotations(bboxes: np.ndarray) -> np.ndarray:
         bboxes (np.ndarray): array of bounding boxes after flip
     """
 
-    bboxes = flip_left_right_annotations(bboxes)
+    bboxes = bboxes.copy()
     bboxes[:, 1] = 1 - bboxes[:, 1]
     return bboxes
 
@@ -312,8 +314,28 @@ def rot90_annotations(bboxes: np.ndarray) -> np.ndarray:
         bboxes (np.ndarray): array of bounding boxes after rotation
     """
 
-    bboxes[:, [0, 1]] = 1 - bboxes[:, [1, 0]]
-    bboxes[:, [2, 3]] = bboxes[:, [3, 2]]
+    bboxes = bboxes.copy()
+    bboxes[:, [0,1]] = 1 - bboxes[:, [1,0]]
+    bboxes[:, 1] = 1 - bboxes[:, 1]
+    bboxes[:, [2,3]] = bboxes[:, [3,2]]
+    return bboxes
+
+
+def rot270_annotations(bboxes: np.ndarray) -> np.ndarray:
+
+    """Applying 270 degrees rotation transformation to annotations in numpy array format.
+
+    Args:
+        bboxes (np.ndarray): array of bounding boxes (x and y positions with height and width) for each object in the current image
+
+    Returns:
+        bboxes (np.ndarray): array of bounding boxes after rotation
+    """
+
+    bboxes = bboxes.copy()
+    bboxes[:, [0,1]] = 1 - bboxes[:, [1,0]]
+    bboxes[:, 0] = 1 - bboxes[:, 0]
+    bboxes[:, [2,3]] = bboxes[:, [3,2]]
     return bboxes
 
 
@@ -498,103 +520,84 @@ def data_augmentation_for_yolo_data(
     images_dir = Path("images")
     labels_dir = Path("labels")
     used_images = train_files.copy()
-    used_imgs_ids = set(
-        [image_name.split("/")[-1].split(".")[0] for image_name in used_images]
-    )
+    used_imgs_ids = set([image_name.split("/")[-1].split(".")[0] for image_name in used_images])
+
+    # to transform PIL image to Tensor
+    trans_to_tensor = transforms.ToTensor()
 
     for image_id in tqdm(used_imgs_ids):
 
-        image = Image.open(data_dir / images_dir / f"{image_id}.jpg")
+        img_path = data_dir / images_dir / f"{image_id}.jpg"
+        image = cv2.imread(img_path.as_posix())  # in BGR
         src = data_dir / labels_dir / f"{image_id}.txt"
         labels, bboxes = transform_anns_to_array(src)
 
-        # Add new image with flip left right
-        flipped = tf.image.flip_left_right(image)
-        flipped_img_name = data_dir / images_dir / f"{image_id}_flip_left_right.jpg"
+        # Add new image with 90 degrees rotation
+        # rotate the image by 90 degree clockwise
+        img_cw_90 = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
+        new_img_path = data_dir / images_dir / f"{image_id}_rot90.jpg"
+        cv2.imwrite(new_img_path.as_posix(), img_cw_90)
+        used_images.append(new_img_path.as_posix())
+        bboxes_rot90 = rot90_annotations(bboxes)
+        transform_anns_to_str(labels, bboxes_rot90, data_dir / labels_dir / f"{image_id}_rot90.txt")
+
+        # Add new image with 270 degrees rotation
+        # rotate the image by 90 degree counter clockwise
+        img_cw_270 = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
+        new_img_path = data_dir / images_dir / f"{image_id}_rot270.jpg"
+        cv2.imwrite(new_img_path.as_posix(), img_cw_270)
+        used_images.append(new_img_path.as_posix())
+        bboxes_rot270 = rot270_annotations(bboxes)
+        transform_anns_to_str(labels, bboxes_rot270, data_dir / labels_dir / f"{image_id}_rot270.txt")
+
+        image = cv2.imread(img_path.as_posix())[:,:,::-1]  # BGR to RGB
+
+        # Add new image with horizontal flip with a given probability p
+        Horizontal_Flipping_Transformation = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.RandomHorizontalFlip(p=1),
+        ])
+        Flipping_Img = Horizontal_Flipping_Transformation(image)
+        flipped_img_name = data_dir / images_dir / f"{image_id}_horiz_flip.jpg"
         used_images.append(flipped_img_name.as_posix())
-        tf.keras.utils.save_img(flipped_img_name, flipped)
+        save_image(trans_to_tensor(Flipping_Img), flipped_img_name)
         bboxes_flipped = flip_left_right_annotations(bboxes)
-        transform_anns_to_str(
-            labels,
-            bboxes_flipped,
-            data_dir / labels_dir / f"{image_id}_flip_left_right.txt",
-        )
-
-        # Add new image with flip up down
-        flipped = tf.image.flip_up_down(image)
-        flipped_img_name = data_dir / images_dir / f"{image_id}_flip_up_down.jpg"
+        transform_anns_to_str(labels, bboxes_flipped, data_dir / labels_dir / f"{image_id}_horiz_flip.txt")
+
+        # Add new image with vertical flip with a given probability p
+        Vertical_Flipping_Transformation = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.RandomVerticalFlip(p=1)
+        ])
+        Flipping_Img = Vertical_Flipping_Transformation(image)
+        flipped_img_name = data_dir / images_dir / f"{image_id}_vert_flip.jpg"
         used_images.append(flipped_img_name.as_posix())
-        tf.keras.utils.save_img(flipped_img_name, flipped)
+        save_image(trans_to_tensor(Flipping_Img), flipped_img_name)
         bboxes_flipped = flip_up_down_annotations(bboxes)
-        transform_anns_to_str(
-            labels,
-            bboxes_flipped,
-            data_dir / labels_dir / f"{image_id}_flip_up_down.txt",
-        )
-
-        # Add new image with gray color
-        grayscaled = tf.image.rgb_to_grayscale(image)
-        gray_img_name = data_dir / images_dir / f"{image_id}_gray.jpg"
-        used_images.append(gray_img_name.as_posix())
-        tf.keras.utils.save_img(gray_img_name, grayscaled)
-        dest = data_dir / labels_dir / f"{image_id}_gray.txt"
-        shutil.copy2(src, dest)
-
-        # Add new image with rotation 90
-        rotated = tf.image.rot90(image)
-        rot90_img_name = data_dir / images_dir / f"{image_id}_rot90.jpg"
-        used_images.append(rot90_img_name.as_posix())
-        tf.keras.utils.save_img(rot90_img_name, rotated)
-        bboxes_rot90 = rot90_annotations(bboxes)
-        transform_anns_to_str(
-            labels, bboxes_rot90, data_dir / labels_dir / f"{image_id}_rot90.txt"
-        )
-
-        # Add new image with saturation
-        saturated = tf.image.adjust_saturation(image, 3)
-        saturated_img_name = data_dir / images_dir / f"{image_id}_saturated.jpg"
-        used_images.append(saturated_img_name.as_posix())
-        tf.keras.utils.save_img(saturated_img_name, saturated)
-        dest = data_dir / labels_dir / f"{image_id}_saturated.txt"
+        transform_anns_to_str(labels, bboxes_flipped, data_dir / labels_dir / f"{image_id}_vert_flip.txt")
+
+        # Add new image with contrast and high brightness
+        Color_Transformation = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.ColorJitter(brightness=0.8, contrast=0.5,saturation=0.2, hue=0.4)
+        ])
+        Transformed_Img = Color_Transformation(image)
+        transformed_img_name = data_dir / images_dir / f"{image_id}_high_bright_contrast.jpg"
+        used_images.append(transformed_img_name.as_posix())
+        save_image(trans_to_tensor(Transformed_Img), transformed_img_name)
+        dest = data_dir / labels_dir / f"{image_id}_high_bright_contrast.txt"
         shutil.copy2(src, dest)
 
-        # Add new images with random contrast
-        for i in range(2):
-            seed = (i, 0)  # tuple of size (2,)
-
-            stateless_random_brightness = tf.image.stateless_random_brightness(
-                image, max_delta=0.95, seed=seed
-            )
-            bright_img_name = (
-                data_dir / images_dir / f"{image_id}_random_bright_{seed[0]}.jpg"
-            )
-            used_images.append(bright_img_name.as_posix())
-            tf.keras.utils.save_img(bright_img_name, stateless_random_brightness)
-            dest = data_dir / labels_dir / f"{image_id}_random_bright_{seed[0]}.txt"
-            shutil.copy2(src, dest)
-
-            stateless_random_contrast = tf.image.stateless_random_contrast(
-                image, lower=0.1, upper=0.9, seed=seed
-            )
-            contrast_img_name = (
-                data_dir / images_dir / f"{image_id}_random_contrast_{seed[0]}.jpg"
-            )
-            used_images.append(contrast_img_name.as_posix())
-            tf.keras.utils.save_img(contrast_img_name, stateless_random_contrast)
-            dest = data_dir / labels_dir / f"{image_id}_random_contrast_{seed[0]}.txt"
-            shutil.copy2(src, dest)
-
     return used_images
 
 
 def get_train_valid(
-    data_dir: str, list_files: List[str], split: float = 0.85
+    list_files: List[str], split: float = 0.85
 ) -> Tuple[List[str], List[str]]:
 
     """Split data into train and validation partitions with data augmentation
 
     Args:
-        data_dir (WindowsPath): path of the root data directory. It should contain a folder with all useful data for images and annotations.
         list_files (List[Any,type[str]]): list of image files to split into train and test partitions
         split (float, optional): train_size between 0 and 1. Set as default to 0.85.
 
@@ -605,7 +608,6 @@ def get_train_valid(
 
     train_files, val_files = train_test_split(list_files, train_size=split)
     train_files = list(set(train_files))
-    train_files = data_augmentation_for_yolo_data(data_dir, train_files)
     val_files = list(set(val_files))
 
     return train_files, val_files

src/plasticorigins/training/data/make_dataset2.py

@@ -7,6 +7,7 @@
     get_annotations_from_files,
     get_annotations_from_db,
     find_img_ids_to_exclude,
+    data_augmentation_for_yolo_data,
 )
 from plasticorigins.training.data.data_processing import (
     generate_yolo_files,
@@ -16,6 +17,7 @@
 from pathlib import Path
 import argparse
 from argparse import Namespace
+import os
 
 
 def main(args: Namespace) -> None:
@@ -63,17 +65,39 @@ def main(args: Namespace) -> None:
     print(
         f"found {cpos} valid annotations with images and {cneg} unmatched annotations"
     )
+    if args.split:
+        train_files, val_files = get_train_valid(yolo_filelist, args.split)
+
+    # use data augmentation for artificial data only if original data have been processed
+    if args.artificial_data and args.data_augmentation:
+
+        artificial_data_dir = Path(args.artificial_data)
+        artificial_data_list = [Path(path).as_posix() for path in os.listdir(artificial_data_dir / "images")]
+        artificial_train_files, artificial_val_files = get_train_valid(artificial_data_list, args.split)
+        artificial_train_files = data_augmentation_for_yolo_data(artificial_data_dir, artificial_train_files)
+        # concatenate original images and artificial data
+        with open(data_dir / "train.txt", "w") as f:
+            for path in artificial_train_files:
+                f.write(path + "\n")
+        with open(data_dir / "val.txt", "w") as f:
+            for path in artificial_val_files:
+                f.write(path + "\n")
 
-    train_files, val_files = get_train_valid(data_dir, yolo_filelist, args.split)
+    else:
+        train_files, val_files = get_train_valid(yolo_filelist)
+
+        if args.data_augmentation:
+            train_files = data_augmentation_for_yolo_data(data_dir, train_files)
 
-    generate_yolo_files(data_dir, train_files, val_files, args.nb_classes)
+        generate_yolo_files(data_dir, train_files, val_files, args.nb_classes)
 
 
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Build dataset")
     parser.add_argument("--data-dir", type=str, help="path to main data folder")
     parser.add_argument("--images-dir", type=str, help="path to image folder")
+    parser.add_argument("--artificial-data", type=str, help="path to artificial data folder")
     parser.add_argument("--bboxes-filename", type=str, default="")
     parser.add_argument("--images-filename", type=str, default="")
     parser.add_argument("--user", type=str, help="username for connection to DB")
@@ -93,6 +117,7 @@ def main(args: Namespace) -> None:
     )
     parser.add_argument("--nb-classes", type=int, default=10)
     parser.add_argument("--split", type=float, default=0.85)
+    parser.add_argument("--data-augmentation", type=int, default=0)
     parser.add_argument("--limit-data", type=int, default=0)
     parser.add_argument(
         "--exclude-img-folder",