FIX: modify positional params in aug pipeline

dataPrep/data_preparation.py

@@ -45,10 +45,33 @@ if torch.cuda.is_available():
 
 
 # ----- ClearML Setup -----
-task = Task.init(project_name='Small Group Project', task_name='data_prep')
+task = Task.init(
+    project_name='Small Group Project',
+    task_name='Data Preparation',
+    task_type=Task.TaskTypes.data_processing
+)
 task.set_random_seed(SEED)
 clearml_logger = task.get_logger()
 
+# --- Track full configuration in ClearML ---
+task.connect({
+    "seed": SEED,
+    "dataset": {
+        "link": DATASET_LINK,
+        "subset_ratio": DATASET_SUBSET_RATIO,
+    },
+    "augmentation": {
+        "rotation": ROTATION,
+        "brightness": BRIGHTNESS,
+        "saturation": SATURATION,
+        "blur": BLUR
+    },
+    "dataloaders": {
+        "batch_size": BATCH_SIZE,
+        "test_size": TEST_SIZE
+    }
+})
+
 # ----- 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
@@ -110,12 +133,11 @@ plt.title("Class Distribution in Prototype Dataset")
 plt.xlabel("Class")
 plt.ylabel("Count")
 plt.tight_layout()
-plt.savefig("class_distribution.png")
 
-clearml_logger.report_image(
+clearml_logger.report_matplotlib_figure(
     title="EDA Class Distribution",
     series="Prototype Subset",
-    local_path="class_distribution.png",
+    figure=plt.gcf(),
     iteration=1
 )
 
@@ -140,6 +162,13 @@ if __name__ == "__main__":
     print(f"Prototype Validation loader batches: {len(prototype_loaders['val'])}")
     print(f"Prototype Test loader batches: {len(prototype_loaders['test'])}")
 
+    clearml_logger.report_text(
+        f"Prototype loaders created: "
+        f"train={len(prototype_loaders['train'])}, "
+        f"val={len(prototype_loaders['val'])}, "
+        f"test={len(prototype_loaders['test'])}"
+    )
+
     final_loaders = make_dataset_loaders(
         data_plants, SEED, BATCH_SIZE, TEST_SIZE, aug_config
     )
@@ -154,6 +183,8 @@ if __name__ == "__main__":
         {"dataset_id": clearml_dataset.id},
         name="Dataset Metadata"
     )
+    task.mark_completed()
+
     
     # Close the ClearML task 
     task.close()

dataPrep/helpers/create_dataset.py

@@ -39,6 +39,10 @@ def make_subset(dataset_link, subset_ratio, clearml_logger):
         dataset_tags=["prototype", "subset"],
         use_current_task=True
     )
+    clearml_dataset.add_tags([
+        f"subset_ratio_{subset_ratio}",
+        "hf_source"
+    ])
 
     # Save indices
     subset_path = "subset_indices.npy"

dataPrep/helpers/transforms_loaders.py

@@ -38,7 +38,7 @@ def make_augment_pipeline(aug_config):
         
         # Randomly changing some parameters of pictures to enrich dataset
         transforms.RandomRotation(rotation),
-        transforms.ColorJitter(brightness, saturation),
+        transforms.ColorJitter(brightness=brightness, saturation=saturation),
         transforms.GaussianBlur(blur),
         transforms.ToTensor(),
         transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)