FEAT: Add final data loaders for the whole dataset

data_preparation.py

@@ -1,5 +1,4 @@
 # --- Standard Python Library ---
-import os
 import random
 
 # --- Data Handling & Analysis ---
@@ -17,7 +16,7 @@ from torchvision import transforms
 from torch.utils.data import DataLoader
 
 # --- Experiment Tracking ---
-from clearml import Task, Logger
+from clearml import Task, Dataset
 
 
 # Setting up the SEED to be able to repeat experiments
@@ -170,37 +169,8 @@ print("✅ Checkpoint: Plot with classes distributions is created and saved")
 
 # --------------- Data Splits ------------
 def get_prototype_loaders(batch_size=32):
-        
-    # Standard ImageNet mean and std 
-    # These values are used to normalize the tensors
-    IMAGENET_MEAN = [0.485, 0.456, 0.406]
-    IMAGENET_STD = [0.229, 0.224, 0.225]
-
-    # Defining pipeline to ensure that images are consistently formatted (for Val/Test)
-    normalisation_pipeline = transforms.Compose([
-        # Convert PIL Image to a PyTorch Tensor
-        # This also scales pixel values from [0, 255] to [0.0, 1.0]
-        transforms.ToTensor(), 
-
-        # Normalise the Tensor; Standartises pixel values
-        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
-    ])
-    print("✅ Checkpoint: Transform pipeline created")
-
-    # Augmentation pipeline (to change some parameters of the pictures to create "new" ones)
-    augmentation_pipeline = transforms.Compose([
-        # Randomly changing some parameters of pictures to enrich dataset
-        transforms.RandomRotation(degrees=30),
-        transforms.ColorJitter(brightness=0.2, saturation=0.2),
-        transforms.GaussianBlur(kernel_size=3),
-
-        # Convert to Tensor and Normalise
-        transforms.ToTensor(), 
-        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
-    ])
-
-    print("✅ Checkpoint: Augmentation pipeline created")
-
+    # Calling function to define pipelines
+    normalisation_pipeline, augmentation_pipeline = get_transform_pipelines()
 
     # -- Split the prototype dataset -- 
     # This returns a dictionary: {'train': 70%, 'test': 30%}
@@ -235,6 +205,71 @@ def get_prototype_loaders(batch_size=32):
     return proto_train_loader, proto_val_loader, proto_test_loader
 
 
+
+def get_final_loaders(batch_size=32):
+    # Calling function to define pipelines
+    normalisation_pipeline, augmentation_pipeline = get_transform_pipelines()
+
+    # -- Split the prototype dataset -- 
+    # This returns a dictionary: {'train': 70%, 'test': 30%}
+    split_1_dict = data_plants.train_test_split(test_size=0.3, seed=SEED)
+
+    # Assign the 70% part to final train split
+    train_split = split_1_dict['train']
+
+    # Assign the 30% part to a temporary var
+    temp_split = split_1_dict['test']
+
+    # Split 30% into 2 15%
+    # This returns a dictionary: {'train': 50%, 'test': 50%}
+    split_2_dict = temp_split.train_test_split(test_size=0.5, seed=SEED)
+
+    val_split = split_2_dict['train']
+    test_split = split_2_dict['test']
+
+    print("✅ Checkpoint: Dataset splitted")
+
+    # -- Putting splits through pipelines --
+    train_split.set_transform(augmentation_pipeline)
+    val_split.set_transform(normalisation_pipeline)
+    test_split.set_transform(normalisation_pipeline)
+
+    # -- Creating the prototype dataloaders -- 
+    train_loader = DataLoader(dataset = train_split, batch_size = batch_size, shuffle = True )
+    val_loader = DataLoader(dataset = val_split, batch_size = batch_size, shuffle = False )
+    test_loader = DataLoader(dataset = test_split, batch_size = batch_size, shuffle = False )
+
+    print("✅ Checkpoint: DataLoaders are set")
+    return train_loader, val_loader, test_loader
+
+
+def get_transform_pipelines():
+    """
+    Defines and returns the normalization and augmentation pipelines.
+    """
+    IMAGENET_MEAN = [0.485, 0.456, 0.406]
+    IMAGENET_STD = [0.229, 0.224, 0.225]
+
+    # Defining pipeline for Val/Test
+    normalisation_pipeline = transforms.Compose([
+        transforms.ToTensor(), 
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+    ])
+
+    # Augmentation pipeline for Train
+    augmentation_pipeline = transforms.Compose([
+        transforms.RandomRotation(degrees=30),
+        transforms.ColorJitter(brightness=0.2, saturation=0.2),
+        transforms.GaussianBlur(kernel_size=3),
+        transforms.ToTensor(), 
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+    ])
+    
+    print("✅ Checkpoint: Transform pipelines created")
+    
+    # Return both pipelines
+    return normalisation_pipeline, augmentation_pipeline
+
 # ----------------------------------------------------------------------
 if __name__ == "__main__":