CHORE: separate aug pipeline & parametrise aug transforms

dataPrep/data_preparation.py

@@ -6,7 +6,7 @@ import random
 import numpy as np
 import pandas as pd
 from datasets import load_dataset
-from helpers.create_dataset import load_subset_from_dataset
+from helpers.create_dataset import make_subset
 from helpers.transforms_loaders import make_dataset_loaders
 
 # --- Visualization ---
@@ -15,17 +15,28 @@ import matplotlib.pyplot as plt
 
 # --- PyTorch (Machine Learning) ---
 import torch
-from torchvision import transforms
-from torch.utils.data import DataLoader
 
 # --- Experiment Tracking ---
-from clearml import Task, Logger, Dataset
+from clearml import Task
 
 
-# Setting up the SEED to be able to repeat experiments
+# -------- Controllable parameters --------
+# Dataset parameters
 SEED = 42
+DATASET_LINK = "DScomp380/plant_village"
 DATASET_SUBSET_RATIO = 0.25
 
+# Augmentation parameters
+ROTATION = 30
+BRIGHTNESS = 0.2
+SATURATION = 0.2
+BLUR = 3
+
+# DataLoader parameters
+BATCH_SIZE = 32
+TEST_SIZE = 0.3
+
+# Setting up the SEED to be able to repeat experiments
 random.seed(SEED)
 np.random.seed(SEED)
 torch.manual_seed(SEED)
@@ -34,10 +45,15 @@ if torch.cuda.is_available():
 
 
 # ----- ClearML Setup -----
-task = Task.init(project_name= 'Small Group CW', task_name = 'data_prep')
+task = Task.init(project_name='Small Group Project', task_name='data_prep')
 task.set_random_seed(SEED)
 clearml_logger = task.get_logger()
 
+# ----- Load a subset from a given dataset & track with ClearML -----
+data_plants, prototyping_dataset, features, clearml_dataset = make_subset(
+    DATASET_LINK, DATASET_SUBSET_RATIO, clearml_logger
+)
+
 # Log subset config to ClearML
 task.connect_configuration(
     {"subset_ratio": DATASET_SUBSET_RATIO},
@@ -45,18 +61,12 @@ task.connect_configuration(
 )
 
 
-# ----- Load a subset from a given dataset & track with ClearML -----
-data_plants, prototyping_dataset, features, clearml_dataset = load_subset_from_dataset(
-    SEED, DATASET_SUBSET_RATIO, clearml_logger
-)
-
-
 # ---- Exploratory data analysis (EDA) ----
 
 # Reformatting the label feature to understand bias
 labels_list = prototyping_dataset['label']
 df_labels = pd.Series(labels_list)
-label_count = df_labels.value_counts(sort = False)
+label_count = df_labels.value_counts(sort=False)
 
 # Checking the amount of samples in each class and logging it to clearML
 
@@ -114,8 +124,15 @@ clearml_logger.report_image(
 if __name__ == "__main__":
     
     # ------------------- Dataset splits ----------------------------------
+    aug_config = {
+        'rotation': ROTATION,
+        'brightness': BRIGHTNESS,
+        'saturation': SATURATION,
+        'blur': BLUR
+    }
+
     prototype_loaders = make_dataset_loaders(
-        prototyping_dataset, seed=SEED, batch_size=32, test_size=0.3
+        prototyping_dataset, SEED, BATCH_SIZE, TEST_SIZE, aug_config
     )
 
     print("\n--- Handoff Test Successful ---")
@@ -124,7 +141,7 @@ if __name__ == "__main__":
     print(f"Prototype Test loader batches: {len(prototype_loaders['test'])}")
 
     final_loaders = make_dataset_loaders(
-        data_plants, seed=SEED, batch_size=32, test_size=0.3
+        data_plants, SEED, BATCH_SIZE, TEST_SIZE, aug_config
     )
 
     print("\n--- Handoff Test Successful ---")

dataPrep/helpers/create_dataset.py

@@ -2,19 +2,19 @@
 A collection of dataset (DS) loading and subsetting functions.
 """
 
+import os
 import random
 import numpy as np
 from datasets import load_dataset
 from clearml import Dataset
 
 
-# Load a DS from HuggingFace Link and subset - upload both to ClearML
-def load_subset_from_dataset(seed, subset_ratio, clearml_logger):
-    DATASET_LINK = "DScomp380/plant_village"
+# Load a DS from HuggingFace Link and subset the DS - upload both to ClearML
+def make_subset(dataset_link, subset_ratio, clearml_logger):
 
     # Load dataset
     try:
-        ds = load_dataset(DATASET_LINK)
+        ds = load_dataset(dataset_link)
     except Exception as e:
         raise RuntimeError(f"Error loading the dataset: {e}")
 
@@ -35,8 +35,9 @@ def load_subset_from_dataset(seed, subset_ratio, clearml_logger):
     # ---------- Register subset in ClearML ----------
     clearml_dataset = Dataset.create(
         dataset_name="Plant Village Prototype",
-        dataset_project="smallGroupProject",
-        dataset_tags=["prototype", "subset"]
+        dataset_project="Small Group Project",
+        dataset_tags=["prototype", "subset"],
+        use_current_task=True
     )
 
     # Save indices
@@ -44,6 +45,7 @@ def load_subset_from_dataset(seed, subset_ratio, clearml_logger):
     np.save(subset_path, subset_indices)
     clearml_dataset.add_files(subset_path)
     clearml_dataset.set_metadata({
+        "huggingface_dataset": dataset_link,
         "subset_ratio": subset_ratio,
         "total_samples": len(prototyping_dataset)
     })
@@ -52,4 +54,7 @@ def load_subset_from_dataset(seed, subset_ratio, clearml_logger):
     clearml_dataset.finalize()
     clearml_logger.report_text(f"Created ClearML Dataset: {clearml_dataset.id}")
 
+    # Clean up local file
+    os.remove(subset_path)
+
     return data_plants, prototyping_dataset, features, clearml_dataset

dataPrep/helpers/transforms_loaders.py

@@ -6,13 +6,12 @@ from torchvision import transforms
 from torch.utils.data import DataLoader
 
 
+# Standard ImageNet mean and std - Used to normalize the tensors
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
 
-# Defines and returns the normalization and augmentation pipelines.
-def make_transform_pipelines():
-    
-    # Standard ImageNet mean and std - Used to normalize the tensors
-    IMAGENET_MEAN = [0.485, 0.456, 0.406]
-    IMAGENET_STD = [0.229, 0.224, 0.225]
+# Defines and returns the normalization pipeline.
+def make_norm_pipeline():
 
     # Pipeline ensures image format is consistent (for Val/Test)
     normalisation = transforms.Compose([
@@ -24,28 +23,39 @@ def make_transform_pipelines():
         transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)
     ])
 
+    return normalisation
+
+# Defines and returns the augmentation (rotation, brightness, saturation, blur) pipeline.
+def make_augment_pipeline(aug_config):
+
+    rotation = aug_config['rotation']
+    brightness = aug_config['brightness']
+    saturation = aug_config['saturation']
+    blur = aug_config['blur']
+
     # Augmentation pipeline (to create "new" images by changing some parameters)
     augmentation = transforms.Compose([
         
         # Randomly changing some parameters of pictures to enrich dataset
-        transforms.RandomRotation(30),
-        transforms.ColorJitter(brightness=0.2, saturation=0.2),
-        transforms.GaussianBlur(3),
+        transforms.RandomRotation(rotation),
+        transforms.ColorJitter(brightness, saturation),
+        transforms.GaussianBlur(blur),
         transforms.ToTensor(),
         transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)
     ])
 
-    return normalisation, augmentation
+    return augmentation
 
 
 """
 Creates and returns DataLoaders (train, val, test) for a given dataset.
 Performs a 70/15/15 split
 """
-def make_dataset_loaders(dataset, seed, batch_size=32, test_size=0.3):
+def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config):
 
     # Define transformation pipelines for the dataset
-    normalisation, augmentation = make_transform_pipelines()
+    normalisation = make_norm_pipeline()
+    augmentation = make_augment_pipeline(aug_config)
 
     # 70/30 split creates train set
     split_1 = dataset.train_test_split(test_size=test_size, seed=seed)
@@ -53,7 +63,7 @@ def make_dataset_loaders(dataset, seed, batch_size=32, test_size=0.3):
     remaining_split = split_1['test']
 
     # 15/15 split on remaining data - validation and test sets
-    val_split = test_size/2
+    val_split = 0.5
     split_2 = remaining_split.train_test_split(test_size=val_split, seed=seed)
     val_split, test_split = split_2['train'], split_2['test']