Merge pull request #10 from K23064919/develop

dataPrep/data_preparation.py

@@ -45,8 +45,9 @@ if torch.cuda.is_available():
 
 
 # ----- ClearML Setup -----
+project_name = "Small Group Project"
 task = Task.init(
-    project_name='Small Group Project',
+    project_name=f'{project_name}/Data Preparation',
     task_name='Data Preparation',
     task_type=Task.TaskTypes.data_processing
 )
@@ -74,15 +75,15 @@ task.connect({
 })
 
 # ----- 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
+data_plants, subset_dataset, features = make_subset(
+    DATASET_LINK, DATASET_SUBSET_RATIO, task
 )
 
 
 # ---- Exploratory data analysis (EDA) ----
 
 # Reformatting the label feature to understand bias
-labels_list = prototyping_dataset['label']
+labels_list = subset_dataset['label']
 df_labels = pd.Series(labels_list)
 label_count = df_labels.value_counts(sort=False)
 
@@ -111,6 +112,7 @@ clearml_logger.report_scalar(
     value=(max_count / min_count), 
     iteration=1
 )
+
 print("--- Class imbalance analysis --- ")
 print(f"Max labels in a class: {max_count}")
 print(f"Min labels in a class: {min_count}")
@@ -122,16 +124,17 @@ class_names = features['label'].names
 formatted_class_names = [" ".join(name.replace('_', ' ').split()) for name in class_names]
 label_count.index = formatted_class_names
 
+# Plotting class distribution
 plt.figure(figsize=(10,6))
 label_count.plot(kind='bar', color='skyblue')
-plt.title("Class Distribution in Prototype Dataset")
+plt.title("Class Distribution in Subset Dataset")
 plt.xlabel("Class")
 plt.ylabel("Count")
 plt.tight_layout()
 
 clearml_logger.report_matplotlib_figure(
     title="EDA Class Distribution",
-    series="Prototype Subset",
+    series="Subset Dataset",
     figure=plt.gcf(),
     iteration=1
 )
@@ -149,7 +152,7 @@ if __name__ == "__main__":
     }
 
     prototype_loaders = make_dataset_loaders(
-        prototyping_dataset, SEED, BATCH_SIZE, TEST_SIZE, aug_config
+        subset_dataset, SEED, BATCH_SIZE, TEST_SIZE, aug_config
     )
 
     print("\n--- Handoff Test Successful ---")
@@ -173,14 +176,9 @@ if __name__ == "__main__":
     print(f"Validation loader batches: {len(final_loaders['val'])}")
     print(f"Test loader batches: {len(final_loaders['test'])}")
 
-    # Record dataset info in ClearML
-    task.connect_configuration(
-        {"dataset_id": clearml_dataset.id},
-        name="Dataset Metadata"
-    )
-    task.mark_completed()
 
-    
     # Close the ClearML task 
+    task.mark_completed()
     task.close()
+
     print("\n--- Script Finished ---")

dataPrep/helpers/clearml_data.py

@@ -0,0 +1,136 @@
+import os
+import numpy as np
+
+from clearml import Task, Dataset
+from datasets import load_dataset
+from dataPrep.helpers.transforms_loaders import make_dataset_loaders
+
+
+'''
+Takes latest Data Prep ClearML task from project and reconstruct:
+- data loaders for both full and subset datasets
+- Aug settings used
+'''
+def extract_latest_data_task(project_name: str = "Small Group Project", num_workers: int = 8):
+
+  # --------- Get latest Data Preparation task from ClearML ---------
+
+  all_tasks = Task.get_tasks(
+    project_name=f'{project_name}/Data Preparation',
+    allow_archived=False,
+    task_filter={'order_by': ["-last_update"]},
+  )
+
+  if not all_tasks:
+      raise RuntimeError(f"No tasks found in project '{project_name}'")
+
+  dp_tasks = [
+    t for t in all_tasks
+    if t.task_type == Task.TaskTypes.data_processing
+    and t.completed is not None
+  ]
+
+  if not dp_tasks:
+      raise RuntimeError("No 'Data Preparation' tasks found in this project!")
+
+  # Latest Data Prep Task 
+  latest_task = dp_tasks[0]
+  DYNAMIC_TASK_ID = latest_task.id
+  DATA_PREP = Task.get_task(task_id=DYNAMIC_TASK_ID)
+
+  # Load subset indices artifact from Data Prep task
+  artifacts = DATA_PREP.artifacts
+  if "subset_indices" not in artifacts:
+      raise RuntimeError("Data Prep task did not upload 'subset_indices' artifact!")
+
+  artifact = artifacts["subset_indices"]
+  subset_indices_path = artifact.get_local_copy()
+  subset_indices = np.load(subset_indices_path)
+
+  # Load dataset metadata from Data Prep task
+  data_params = DATA_PREP.get_parameters()
+
+  subset_ratio = float(data_params['General/dataset/subset_ratio'])
+  dataset_link = data_params['General/dataset/link']
+  seed = int(data_params['General/seed'])
+  batch_size = int(data_params['General/dataloaders/batch_size'])
+  test_size = float(data_params['General/dataloaders/test_size'])
+
+  aug_config = {
+      'rotation': float(data_params['General/augmentation/rotation']),
+      'brightness': float(data_params['General/augmentation/brightness']),
+      'saturation': float(data_params['General/augmentation/saturation']),
+      'blur': float(data_params['General/augmentation/blur']),
+  }
+
+
+  # Load Full Dataset 
+  try:
+      ds = load_dataset(dataset_link)
+  except Exception as e:
+      raise RuntimeError(f"Error loading the dataset: {e}")
+
+  full_dataset = ds['train']
+
+  # Apply subset indices to full dataset - this gives you the same subset as data prep
+  subset_dataset = full_dataset.select(subset_indices)
+
+  # Get data loaders for both full and subset datasets
+  subset_loaders, full_loaders, aug_config = get_data_loaders(data_params, subset_dataset, full_dataset, num_workers=num_workers)
+  batch_size = int(data_params['General/dataloaders/batch_size'])
+  seed = int(data_params['General/seed'])
+
+
+  # Gather data prep task metadata
+  data_prep_metadata = {
+    "data_prep_task_id": DYNAMIC_TASK_ID,
+    "dataset_link": dataset_link,
+    "subset_ratio_used": subset_ratio,
+    "augmentation_used": aug_config,
+    "batch_size_used": batch_size,
+    "seed_used": seed,
+    "test_size_used": test_size
+  }
+
+  return subset_loaders, full_loaders, data_prep_metadata
+
+
+'''
+Takes a given dataset, subset, data params to create DataLoaders
+Loaders split data into train, val, test
+'''
+def get_data_loaders(data_params, subset_dataset, full_dataset, num_workers):
+
+  # Extract data parameters- these will be used in the DataLoaders
+  seed = int(data_params['General/seed'])
+  batch_size = int(data_params['General/dataloaders/batch_size'])
+  test_size = float(data_params['General/dataloaders/test_size'])
+
+  aug_config = {
+      'rotation': float(data_params['General/augmentation/rotation']),
+      'brightness': float(data_params['General/augmentation/brightness']),
+      'saturation': float(data_params['General/augmentation/saturation']),
+      'blur': float(data_params['General/augmentation/blur'])
+  }
+
+  # Create DataLoaders using the parameters from data prep
+  subset_loaders = make_dataset_loaders(
+      subset_dataset, seed, batch_size, test_size, aug_config, workers=num_workers
+  )
+
+  print("\n--- Handoff Test Successful ---")
+  print(f"Prototype Train loader batches: {len(subset_loaders['train'])}")
+  print(f"Prototype Validation loader batches: {len(subset_loaders['val'])}")
+  print(f"Prototype Test loader batches: {len(subset_loaders['test'])}")
+
+
+  full_loaders = make_dataset_loaders(
+      full_dataset, seed, batch_size, test_size, aug_config, workers=num_workers
+  )
+
+  print("\n--- Handoff Test Successful ---")
+  print(f"Train loader batches: {len(full_loaders['train'])}")
+  print(f"Validation loader batches: {len(full_loaders['val'])}")
+  print(f"Test loader batches: {len(full_loaders['test'])}")
+
+  return subset_loaders, full_loaders, aug_config

dataPrep/helpers/create_dataset.py

@@ -6,7 +6,6 @@ import os
 import random
 import numpy as np
 from datasets import load_dataset
-from clearml import Dataset
 
 
 '''
@@ -14,7 +13,7 @@ Load a DS from HuggingFace Link & randomly subset it - upload subset to ClearML
 Subset indicies are uploaded to ClearML for reproducibility
 REPRODUCE: Load full DS, then load indicies from ClearML to get same subset
 '''
-def make_subset(dataset_link, subset_ratio, clearml_logger):
+def make_subset(dataset_link, subset_ratio, clearml_task):
 
     # Load dataset
     try:
@@ -34,36 +33,16 @@ def make_subset(dataset_link, subset_ratio, clearml_logger):
     random.shuffle(indices)
     subset_indices = indices[:subset_size]
 
-    prototyping_dataset = data_plants.select(subset_indices)
-# I THINK WE NEED TO REMOVE THIS LATER
-# We dont really need to upload subset everytime (Im not sure tho)
-    # Register subset in ClearML 
-    clearml_dataset = Dataset.create(
-        dataset_name="Plant Village Prototype",
-        dataset_project="Small Group Project",
-        dataset_tags=["prototype", "subset"],
-        use_current_task=False
-    )
-    clearml_dataset.add_tags([
-        f"subset_ratio_{subset_ratio}",
-        "hf_source"
-    ])
+    subset_dataset = data_plants.select(subset_indices)
 
-    # Save indices
+    # -------- Upload the subset indices as a ClearML artifact --------
     subset_path = "subset_indices.npy"
     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)
-    })
-
-    clearml_dataset.upload()
-    clearml_dataset.finalize()
-    clearml_logger.report_text(f"Created ClearML Dataset: {clearml_dataset.id}")
 
-    # Clean up local file
-    os.remove(subset_path)
+    clearml_task.upload_artifact(
+        name="subset_indices",
+        artifact_object=subset_path
+    )
+    clearml_task.get_logger().report_text(f"Uploaded subset indices as artifact: {subset_path}")
 
-    return data_plants, prototyping_dataset, features, clearml_dataset
+    return data_plants, subset_dataset, features

dataPrep/helpers/transforms_loaders.py

@@ -47,24 +47,25 @@ def make_augment_pipeline(aug_config):
     return augmentation
 
 
+def apply_augmentation(batch, augmentation):
+    batch['image'] = [augmentation(x) for x in batch['image']]
+    return batch
+
+def apply_normalisation(batch, normalisation):
+    batch['image'] = [normalisation(x) for x in batch['image']]
+    return batch
+    
+
 """
 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, test_size, aug_config):
+def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config, workers=8):
 
     # Define transformation pipelines for the dataset
     normalisation = make_norm_pipeline()
     augmentation = make_augment_pipeline(aug_config)
 
-    def apply_augmentation(batch):
-        batch['image'] = [augmentation(x) for x in batch['image']]
-        return batch
-
-    def apply_normalisation(batch):
-        batch['image'] = [normalisation(x) for x in batch['image']]
-        return batch
-    
     # 70/30 split creates train set
     split_1 = dataset.train_test_split(test_size=test_size, seed=seed)
     train_split = split_1['train']
@@ -76,14 +77,34 @@ def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config):
     val_split, test_split = split_2['train'], split_2['test']
 
     # Put each split through pipelines
-    train_split.set_transform(apply_augmentation)
-    val_split.set_transform(apply_normalisation)
-    test_split.set_transform(apply_normalisation)
+    train_split.set_transform(lambda batch: apply_augmentation(batch, augmentation))
+    val_split.set_transform(lambda batch: apply_normalisation(batch, normalisation))
+    test_split.set_transform(lambda batch: apply_normalisation(batch, normalisation))
 
     # Create dataloader for each
-    train_loader = DataLoader(train_split, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_split, batch_size=batch_size, shuffle=False)
-    test_loader = DataLoader(test_split, batch_size=batch_size, shuffle=False)
+    train_loader = DataLoader(
+        train_split,
+        batch_size=batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=workers
+    )
+    val_loader = DataLoader(
+        val_split,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=workers
+    )
+    test_loader = DataLoader(
+        test_split,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=workers
+    )
+
+    print(f"\nWorkers used in DataLoaders: {workers}\n")
 
     dataset_loaders = {
         "train": train_loader,

requirements.txt

@@ -1,29 +1,19 @@
 
 # Core dependencies
-torch>=2.0.0
-torchvision>=0.15.0
-gradio>=4.0.0
-numpy>=1.24.0
-Pillow>=10.0.0
+torch==2.2.2
+torchvision==0.17.2
+torcheval==0.0.7
+numpy==1.26.4
+Pillow==10.3.0
+gradio==4.19.0
 
-# For model deployment and tracking
-huggingface-hub>=0.19.0
-clearml>=1.14.0
-
-# Optional: for advanced features
-datasets>=2.14.0  # For loading PlantVillage dataset from HuggingFace
-# -- Data prep requirements --
 # Data Handling & Analysis
-numpy
-pandas
-datasets
+pandas==2.2.2
+datasets==2.18.0
 
 # Visualization
-matplotlib
+matplotlib==3.8.4
 
-# PyTorch (Machine Learning)
-torch
-torchvision
-
-# Experiment Tracking
-clearml
+# For model deployment and tracking
+huggingface-hub==0.23.0
+clearml==2.0.2

testingModel/helpers/evaluation.py

@@ -0,0 +1,43 @@
+import torch
+from torch.nn import CrossEntropyLoss
+
+
+"""
+Evaluates a trained model on a dataloader that returns batches like:
+    batch["image"] -> Tensor [B, 3, 256, 256]
+    batch["label"] -> Tensor [B]
+
+Returns dict:
+    { "accuracy": float, "loss": float }
+"""
+def make_predictions(model, dataloader, device):
+
+    model.eval()
+    criterion = CrossEntropyLoss()
+
+    total_loss = 0
+    total_correct = 0
+    total_samples = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+
+            # Move tensors to device
+            images = batch["image"].to(device)
+            labels = batch["label"].to(device).long()
+
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            total_loss += loss.item() * images.size(0)
+            total_correct += (outputs.argmax(dim=1) == labels).sum().item()
+            total_samples += labels.size(0)
+
+    accuracy = total_correct / total_samples
+    avg_loss = total_loss / total_samples
+
+    return {
+        "accuracy": accuracy,
+        "loss": avg_loss,
+    }

testingModel/run_testing.py

@@ -0,0 +1,76 @@
+from clearml import Task
+from dataPrep.helpers.clearml_data import extract_latest_data_task
+
+import torch
+from models.modelOne import modelOne
+from testingModel.helpers.evaluation import make_predictions
+
+
+# -------------- Load Data --------------
+project_name = "Small Group Project"
+subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name)
+
+
+# -------- ClearML Testing Task Setup --------
+testing_task = Task.init(
+    project_name=f"{project_name}/Model Testing",
+    task_name="Model Testing",
+    task_type=Task.TaskTypes.testing,
+    reuse_last_task_id=False,
+)
+
+# Reference the data prep task used
+testing_logger = testing_task.get_logger()
+testing_task.connect(data_prep_metadata, name="data_prep_metadata_READONLY")
+
+CLEARML_TRAINING_ID = "5bac154a885b4acbaa07d8588027bb27"
+
+# Testing parameters - Modify these when experimenting
+testing_config = {
+    "model_train_id": CLEARML_TRAINING_ID,
+    "num_classes": 39,
+    "model_path": "best_model.pt",
+}
+testing_task.connect(testing_config)
+
+# Load the model weights from ClearML training task
+training_task = Task.get_task(task_id=testing_config["model_train_id"])
+model_artifact = training_task.artifacts.get("best_model")
+model_path = model_artifact.get_local_copy()
+
+# Reference training metadata
+training_hyperparams = training_task.get_parameters_as_dict()
+testing_task.connect(training_hyperparams['General'], name="training_metadata_READONLY")
+
+
+# -------- Rebuild the ML model --------
+model = modelOne()
+state_dict = torch.load(model_path, map_location="cpu") # Load to CPU first
+model.load_state_dict(state_dict)
+model.eval() # set dropout & batch norm layers to eval mode
+
+# Move model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+
+# -------------------- Test model on test set --------------------
+testing_logger.report_text("Starting evaluation on TEST SUBSET...\n")
+test_subset = subset_loaders['test']
+
+subset_results = make_predictions(model, test_subset, device)
+
+
+# Accuracy & Loss logging
+testing_logger.report_single_value(name="Test Subset Accuracy", value=subset_results["accuracy"])
+testing_logger.report_single_value(name="Test Subset Loss", value=subset_results["loss"])
+
+
+# --------- Complete -----------------
+print("\n------ Testing Complete ------")
+testing_logger.report_text(
+    f"TEST SUBSET RESULTS:\n"
+    f"Loss: {subset_results['loss']:.4f}\n"
+    f"Accuracy: {subset_results['accuracy']:.4f}\n"
+)
+testing_task.close()

trainingModel/Training.py

@@ -1,150 +0,0 @@
-import torch 
-import torch.nn as nn
-import numpy as np
-from torcheval.metrics import MulticlassAccuracy
-from torch.utils.data import DataLoader
-
-
-# fix errors in runtime
-
-
-def train_model(
-    model: nn.Module,
-    train_loader: DataLoader,
-    val_loader: DataLoader,
-    device: torch.device,
-    n_epochs: int = 4,
-    lr: float = 1e-3,
-    save_path: str = "best_model.pt",
-    flatten_input = False,
-    num_classes : int = 39,
-
-):
-    """
-    Trains the given model and returns:
-    - training_losses: numpy array of loss per batch
-    - training_accuracies: numpy array of running accuracy per batch
-    - val_accuracies: numpy array of accuracy per epoch
-    - best_accuracy: highest validation accuracy achieved
-
-    Expected batch format:
-        batch["image"] → Tensor [B, C, H, W]
-        batch["label"] → Tensor [B] with class IDs (int64)
-    Model output:
-        outputs → Tensor [B, num_classes] (logits)
-    """
-
-    
-    # Move model to device
-    model.to(device)
-
-    # Loss and optimizer
-    criterion = nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=lr ) # might add momentum 0.9 later
-
-    # Metric trackers 
-    train_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
-    val_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
-
-    # Arrays to log metrics
-    num_batches = len(train_loader)
-
-    if num_batches == 0:
-        raise RuntimeError("UH OH!!!! empty train loader")
-
-    # Store training losses and accuracies for every batch
-    # num_batches is the number of batches for every epoch
-    training_losses = np.zeros(num_batches * n_epochs)
-    training_accuracies = np.zeros(num_batches * n_epochs)
-
-    # store validation accuracy for every epoch
-    val_accuracies = np.zeros(n_epochs)
-
-    # keep track of best validation accuracy and best model
-    best_accuracy = 0.0
-
-    #----------------------
-    # training loop
-    #----------------------
-    
-    for epoch in range(n_epochs):
-        model.train()
-        train_accuracy_fn.reset()
-
-        # iterate over all the dataloader's mini-batches
-        for i, batch in enumerate(train_loader):
-
-            # move to GPU memory
-            inputs = batch["image"].to(device)
-            labels = batch["label"].to(device).long()
-
-            # flatten if not cnn REVISE LATER
-            if flatten_input:
-                inputs = inputs.view(inputs.size(0), -1)
-            
-            optimizer.zero_grad()
-
-            # Forward pass
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
-            
-            # Backward pass
-            loss.backward()
-
-            # updates the parameters
-            optimizer.step()
-            
-            # log the loss value
-            training_losses[epoch * num_batches + i] = loss.item() 
-
-            #updates the accuracy computation with new data 
-            train_accuracy_fn.update(outputs, labels)
-
-            #compute accuracy with the current data
-            training_accuracies[epoch * num_batches + i] = train_accuracy_fn.compute().item() 
-
-        print(f'Epoch {epoch + 1} training complete')
-
-        # ----------------------
-        # validation loop
-        # ----------------------
-
-        model.eval()
-        val_accuracy_fn.reset()
-
-
-        with torch.no_grad():
-            for batch in val_loader:
-                inputs = batch["image"].to(device)
-                labels = batch["label"].to(device).long()
-
-                # flatten if not cnn REVISE LATER
-                if flatten_input:
-                    inputs = inputs.view(inputs.size(0), -1)
-
-                outputs = model(inputs)
-
-                val_accuracy_fn.update(outputs, labels)
-
-        current_accuracy = val_accuracy_fn.compute().item()
-        val_accuracies[epoch] = current_accuracy
-
-        # keep track of best validation accuracy and save best model so far
-        if current_accuracy > best_accuracy:
-            best_accuracy = current_accuracy
-            torch.save(model.state_dict(), save_path)
-            print(f'Epoch {epoch + 1} (validation accuracy: {best_accuracy})')
-        print(f'Epoch {epoch + 1} validation complete')
-
-    print(f"\nTraining finished. Best val accuracy: {best_accuracy:.4f}")
-    print(f"Best model weights saved to: {save_path}")
-
-    training_metrics = {
-        "losses": training_losses,
-        "accuracies": training_accuracies,
-        "val_accuracies": val_accuracies,
-        "best_accuracy": best_accuracy,
-    }
-
-    return training_metrics
-

trainingModel/helpers/Training.py

@@ -0,0 +1,199 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torcheval.metrics import MulticlassAccuracy
+from torch.utils.data import DataLoader
+
+
+
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", DEVICE)
+
+def train_model(
+   model: nn.Module,
+   train_loader: DataLoader,
+   val_loader: DataLoader,
+   n_epochs: int = 4,
+   lr: float = 1e-3,
+   save_path: str = "best_model.pt",
+   num_classes : int = 39,
+   early_stop : int = 3,
+
+
+):
+   """
+   Trains the given model and returns:
+   - training_losses: numpy array of loss per epoch
+   - training_accuracies: numpy array of running accuracy per epoch
+   - val_accuracies: numpy array of accuracy per epoch
+   - best_accuracy: highest validation accuracy achieved
+
+
+   Expected batch format:
+       batch["image"] → Tensor [B, C, H, W]
+       batch["label"] → Tensor [B] with class IDs (int64)
+   Model output:
+       outputs → Tensor [B, num_classes] (logits)
+   """
+
+
+   # Move model to device
+   model.to(DEVICE)
+
+
+   # Loss and optimizer
+   criterion = nn.CrossEntropyLoss()
+   optimizer = torch.optim.Adam(model.parameters(), lr=lr ) # might add momentum 0.9 later
+
+
+   # Metric trackers
+   train_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
+   val_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
+
+
+   # Arrays to log metrics
+   num_batches = len(train_loader)
+
+
+   if num_batches == 0:
+       raise RuntimeError("UH OH!!!! empty train loader")
+
+
+   # Store training losses and accuracies for every epoch
+   training_losses = np.zeros(n_epochs)
+   training_accuracies = np.zeros(n_epochs)
+
+
+   # store validation accuracy for every epoch
+   val_accuracies = np.zeros(n_epochs)
+
+
+   # keep track of best validation accuracy and best model
+   best_accuracy = 0.0
+
+
+   # keep track of accuracy improvement
+   improv_counter = 0
+
+
+   #----------------------
+   # training loop
+   #----------------------
+  
+   for epoch in range(n_epochs):
+       model.train()
+       train_accuracy_fn.reset()
+
+
+       training_loss = 0.0
+
+
+       # iterate over all the dataloader's mini-batches
+       for i, batch in enumerate(train_loader):
+
+
+           # move to GPU memory
+           inputs = batch["image"].to(DEVICE)
+           labels = batch["label"].to(DEVICE).long()
+
+
+
+
+           optimizer.zero_grad()
+
+
+           # Forward pass
+           outputs = model(inputs)
+           loss = criterion(outputs, labels)
+          
+           # Backward pass
+           loss.backward()
+
+
+           # updates the parameters
+           optimizer.step()
+          
+           # log the loss value for epoch
+           training_loss += loss.item()
+
+
+           #updates the accuracy computation with new data
+           train_accuracy_fn.update(outputs, labels)
+
+
+       # compute epoch-level training metrics
+       training_losses[epoch] = training_loss / num_batches
+       training_accuracies[epoch] = train_accuracy_fn.compute().item()
+
+
+       print(f'Epoch {epoch + 1} training complete. Training Accuracy: {training_accuracies[epoch]:.4f}')
+
+
+       # ----------------------
+       # validation loop
+       # ----------------------
+
+
+       model.eval()
+       val_accuracy_fn.reset()
+
+
+
+
+       with torch.no_grad():
+           for batch in val_loader:
+               inputs = batch["image"].to(DEVICE)
+               labels = batch["label"].to(DEVICE).long()
+
+
+               outputs = model(inputs)
+
+
+               val_accuracy_fn.update(outputs, labels)
+
+
+       current_accuracy = val_accuracy_fn.compute().item()
+       val_accuracies[epoch] = current_accuracy
+
+
+       # keep track of best validation accuracy and save best model so far
+       if current_accuracy > best_accuracy:
+           best_accuracy = current_accuracy
+           torch.save(model.state_dict(), save_path)
+           improv_counter = 0  #Resets coounter if accuracy improves
+           print(f'Epoch {epoch + 1} (validation accuracy: {best_accuracy})')
+
+
+       else:
+           improv_counter +=1
+           print(f'No improvement for {improv_counter} epoch')
+
+
+           if improv_counter >= early_stop:
+               print (f"Early stopping at epoch {epoch +1}")
+               break
+
+
+
+
+       print(f'Epoch {epoch + 1} validation complete')
+
+
+   print(f"\nTraining finished. Best val accuracy: {best_accuracy:.4f}")
+   print(f"Best model weights saved to: {save_path}")
+
+   
+   training_metrics = {
+       "losses": training_losses,
+       "accuracies": training_accuracies,
+       "val_accuracies": val_accuracies,
+       "best_accuracy": best_accuracy
+
+   }
+
+   return training_metrics
+
+
+
+

trainingModel/run_training.py

@@ -1,124 +1,37 @@
 import os
-import numpy as np
-from clearml import Task, Dataset
-from datasets import load_dataset
+from clearml import Task
+from dataPrep.helpers.clearml_data import extract_latest_data_task
 
-# Latest Data Prep Task 
-all_tasks = Task.get_tasks(project_name="Small Group Project")
-if not all_tasks:
-    raise RuntimeError("No tasks found in project 'Small Group Project'")
+import torch
+from models.modelOne import modelOne
+from trainingModel.helpers.Training import train_model
 
-dp_tasks = [t for t in all_tasks if t.name == "Data Preparation"]
-if not dp_tasks:
-    raise RuntimeError("No 'Data Preparation' tasks found in this project!")
 
-<<<<<<< HEAD
 # -------------- Load Data --------------
-
-all_tasks = Task.get_tasks(project_name="Small Group Project")
-if not all_tasks:
-    raise RuntimeError("No tasks found in project 'Small Group Project'")
-
-dp_tasks = [t for t in all_tasks if t.name == "Data Preparation"]
-if not dp_tasks:
-    raise RuntimeError("No 'Data Preparation' tasks found in this project!")
-
-# Latest Data Prep Task 
-latest_task = max(dp_tasks, key=lambda t: t.id)
-DYNAMIC_TASK_ID = latest_task.id
-DATA_PREP = Task.get_task(task_id=DYNAMIC_TASK_ID)
-
-=======
-latest_task = max(dp_tasks, key=lambda t: t.id)
-DYNAMIC_TASK_ID = latest_task.id
-DATA_PREP = Task.get_task(task_id=DYNAMIC_TASK_ID)
-
->>>>>>> 20050ad82ebca27a376e15837a7abf79fca23e98
-# Dataset ID 
-config_objects = DATA_PREP.get_configuration_objects()
-raw_meta = config_objects["Dataset Metadata"]
-dataset_id = raw_meta.split("=")[1].strip().replace('"', "")
-
-# Load ClearML Dataset
-subset_clearml = Dataset.get(dataset_id=dataset_id)
-local_folder = subset_clearml.get_local_copy()
-
-<<<<<<< HEAD
-subset_indices = np.load(os.path.join(local_folder, "subset_indices.npy"))
-=======
-subset_indices_path = os.path.join(local_folder, "subset_indices.npy")
-subset_indices = np.load(subset_indices_path)
->>>>>>> 20050ad82ebca27a376e15837a7abf79fca23e98
-
-# Load Dataset Parameters 
-data_params = DATA_PREP.get_parameters()
-dataset_link = data_params['General/dataset/link']
-
-# Load Full Dataset 
-try:
-    ds = load_dataset(dataset_link)
-except Exception as e:
-    raise RuntimeError(f"Error loading the dataset: {e}")
-
-full_dataset = ds['train']
-
-
-
-# Apply subset indices to full dataset - this gives you the same subset as data prep
-subset_dataset = full_dataset.select(subset_indices)
-
-
-# Extract parameters from data prep task - these will create the DataLoaders
-seed = int(data_params['General/seed'])
-batch_size = int(data_params['General/dataloaders/batch_size'])
-test_size = float(data_params['General/dataloaders/test_size'])
-
-aug_config = {
-    'rotation': float(data_params['General/augmentation/rotation']),
-    'brightness': float(data_params['General/augmentation/brightness']),
-    'saturation': float(data_params['General/augmentation/saturation']),
-    'blur': float(data_params['General/augmentation/blur'])
-}
-
-# Create DataLoaders using the parameters from data prep
-subset_loaders = make_dataset_loaders(
-    subset_dataset, seed, batch_size, test_size, aug_config
-)
-
-print("\n--- Handoff Test Successful ---")
-print(f"Prototype Train loader batches: {len(subset_loaders['train'])}")
-print(f"Prototype Validation loader batches: {len(subset_loaders['val'])}")
-print(f"Prototype Test loader batches: {len(subset_loaders['test'])}")
-
-
-full_loaders = make_dataset_loaders(
-    full_dataset, seed, batch_size, test_size, aug_config
-)
-
-print("\n--- Handoff Test Successful ---")
-print(f"Train loader batches: {len(full_loaders['train'])}")
-print(f"Validation loader batches: {len(full_loaders['val'])}")
-print(f"Test loader batches: {len(full_loaders['test'])}")
-# -------------- DATA PREP ENDS --------------
+NUM_WORKERS = 0
+project_name = "Small Group Project"
+subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name, num_workers=NUM_WORKERS)
 
 
 # -------- ClearML Training Task Setup --------
 training_task = Task.init(
-    project_name="Small Group Project",
+    project_name=f"{project_name}/Model Training",
     task_name="Model Training",
     reuse_last_task_id=False,
 )
 
+# Detail the data prep task used
 training_logger = training_task.get_logger()
-training_task.connect({"data_prep_task_used": DYNAMIC_TASK_ID})
+training_task.connect(data_prep_metadata, name="data_prep_metadata_READONLY")
 
 # Training parameters - Modify these to experiment
 training_config = {
     "num_classes": 39,
     "n_epochs": 1,
     "learning_rate": 1e-3,
-    "batch_size": batch_size,
+    "optimizer": "adam",
     "save_path": "best_model.pt",
+    "num_workers": NUM_WORKERS
 }
 training_task.connect(training_config)
 
@@ -126,48 +39,45 @@ training_task.connect(training_config)
 # -------- Build the ML model --------
 model = modelOne(noOfClasses=training_config["num_classes"])
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
+# Print device info
+print(f"\n**Using device: {device}**\n") 
+if device.type == 'cuda':
+    print(f"GPU Name: {torch.cuda.get_device_name(0)}")
 
 # ------- Train the model (on subset for now) -------
  
-<<<<<<< HEAD
 print("\n--- Starting Model Training on Subset ---")
 training_metrics = train_model(
-=======
-#When calling this function, the model should be trained on the given dataset
-
-print("\n--- Starting Model Training on Subset ---")
-train_model(
->>>>>>> 20050ad82ebca27a376e15837a7abf79fca23e98
     model=model,
     train_loader=subset_loaders['train'],
     val_loader=subset_loaders['val'],
-    device=device,
     n_epochs=training_config["n_epochs"],
     lr=training_config["learning_rate"],
+    num_classes=training_config["num_classes"],
     save_path=training_config["save_path"],
+    early_stop=3,
 )
-<<<<<<< HEAD
 
 
 # ----------- Log metrics to ClearML -----------
-# Per-batch training losses and accuracies
-for i, loss in enumerate(training_metrics["losses"]):
-    training_logger.report_scalar("train", "loss_per_batch", value=loss, iteration=i)
+# Per-epoch training losses and accuracies
+for epoch, loss in enumerate(training_metrics["losses"]):
+    training_logger.report_scalar("training epoch loss", "loss", value=loss, iteration=epoch)
 
-for i, acc in enumerate(training_metrics["accuracies"]):
-    training_logger.report_scalar("train", "accuracy_per_batch", value=acc, iteration=i)
+for epoch, acc in enumerate(training_metrics["accuracies"]):
+    training_logger.report_scalar("training epoch accuracy", "accuracy", value=acc, iteration=epoch)
 
-# Per-epoch validation accuracy
+# Per-epoch validation accuracies
 for epoch, acc in enumerate(training_metrics["val_accuracies"]):
-    training_logger.report_scalar("validation", "accuracy_per_epoch", value=acc, iteration=epoch)
+    training_logger.report_scalar("validation epoch accuracy", "accuracy", value=acc, iteration=epoch)
 
+# Best validation accuracy
 training_logger.report_single_value("best_val_accuracy", training_metrics["best_accuracy"])
 
 # Upload best model as artifact
 training_task.upload_artifact("best_model", training_config["save_path"])
 
 print("\nTraining complete.")
-training_task.close()
-=======
->>>>>>> 20050ad82ebca27a376e15837a7abf79fca23e98
+training_task.close()

ui/app.py

@@ -14,6 +14,10 @@ sys.path.append(str(Path(__file__).parent))
 sys.path.append(str(Path(__file__).parent.parent))
 
 from model_loader import ModelLoader
+import utils
+from utils import *
+import config
+from config import *
 
 
 class PlantDiseaseApp:
@@ -22,60 +26,98 @@ class PlantDiseaseApp:
         self.current_modelName = "CNN from Scratch"
         self.model = self.model_loader.loadModel(self.current_modelName)
         self.flagged_predictions = []
+        self.class_names = utils.get_class_names()
 
     def predict(self, image, modelName, confidence_threshold):
+        """
+        Predict plant disease from a single image.
+
+        Args:
+            image: PIL Image or numpy array from Gradio upload
+            modelName: Name of the model to use
+            confidence_threshold: float (0-100), only show predictions above this confidence
+
+        Returns:
+            display_predictions: dict, class_name -> probability
+            result_text: str, formatted top prediction info
+            raw_predictions: str, JSON-formatted top predictions
+        """
         if image is None:
             return None, "Please upload an image", ""
 
         try:
+            # Load model if needed
             if modelName != self.current_modelName:
-                self.model = self.model_loader.loadModel(modelName)
+                self.model, self.class_names = self.model_loader.loadModel(modelName)
                 self.current_modelName = modelName
 
             # Preprocess image
-            tensor = preprocess_image(image)
-            tensor = tensor.to(self.model_loader.device)
+            tensor = preprocess_image(image).to(self.model_loader.device)
 
-            # Get prediction
+            # Model inference
             with torch.no_grad():
                 logits = self.model(tensor)
 
-            # Postprocess
-            top_predictions, all_predictions = postprocess_predictions(
-                logits, config.CLASS_NAMES, config.TOP_K_PREDICTIONS
-            )
+            # Convert logits to probabilities
+            probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy()[0]
+
+            # Map to class names
+            predictions = {name: float(prob) for name, prob in zip(self.class_names, probs)}
 
             # Filter by confidence threshold
-            filtered_predictions = {
-                k: v for k, v in top_predictions.items() if v >= confidence_threshold / 100
-            }
+            filtered_predictions = {k: v for k, v in predictions.items() if v >= confidence_threshold / 100.0}
 
-            # Get top prediction info
+            # Top prediction info
             if filtered_predictions:
                 top_class = max(filtered_predictions.items(), key=lambda x: x[1])[0]
                 top_prob = filtered_predictions[top_class]
                 disease_info = get_disease_info(top_class)
 
                 result_text = f"""
-                **Top Prediction:** {disease_info['formatted_name']}
-                **Confidence:** {top_prob*100:.2f}%
-                **Plant:** {disease_info['plant']}
-                **Status:** {'Healthy' if disease_info['is_healthy'] else 'Disease Detected'}
-                """
+                    **Top Prediction:** {disease_info['formatted_name']}
+                    **Confidence:** {top_prob*100:.2f}%
+                    **Plant:** {disease_info['plant']}
+                    **Status:** {'Healthy' if disease_info['is_healthy'] else 'Disease Detected'}
+                    """
             else:
                 result_text = "No predictions above confidence threshold"
 
             # Format for Gradio Label component
-            display_predictions = {
-                format_class_name(k): v for k, v in filtered_predictions.items()
-            }
+            display_predictions = {format_class_name(k): v for k, v in filtered_predictions.items()}
 
-            return display_predictions, result_text, json.dumps(filtered_predictions, indent=2)
+            # Raw JSON output
+            import json
+            raw_predictions = json.dumps(filtered_predictions, indent=2)
+
+            return display_predictions, result_text, raw_predictions
 
         except Exception as e:
             return None, f"Error during prediction: {str(e)}", ""
 
 
+    def flag_prediction(self, image, result_info, feedback_text):
+        if image is None:
+            return "No image uploaded."
+
+        if not feedback_text.strip():
+            return "Please enter feedback before submitting."
+
+        try:
+            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+            entry = {
+                "timestamp": timestamp,
+                "feedback": feedback_text,
+                "model": self.current_modelName,
+                "result_info": result_info
+            }
+
+            self.flagged_predictions.append(entry)
+
+            return "Thanks! Your feedback has been recorded."
+
+        except Exception as e:
+            return f"Error saving feedback: {str(e)}"
 
 def create_interface():
     app = PlantDiseaseApp()
@@ -176,29 +218,6 @@ def create_interface():
                     outputs=flag_output
                 )
 
-            with gr.Tab("Example Images"):
-                gr.Markdown("### Try these example plant images")
-                gr.Markdown("Click on an example below to load it into the predictor")
-
-                example_images = app.get_example_images()
-
-                if example_images:
-                    examples = gr.Examples(
-                        examples=example_images,
-                        inputs=image_input,
-                        label="Example Plant Disease Images"
-                    )
-                else:
-                    gr.Markdown(
-                        """
-                        **No example images found.**
-
-                        To add example images:
-                        1. Create a folder: `ui/examples/`
-                        2. Add plant leaf images (.jpg, .png) to this folder
-                        3. Restart the app
-                        """
-                    )
 
             with gr.Tab("Batch Processing"):
                 gr.Markdown("### Upload multiple images for batch processing")
@@ -214,7 +233,7 @@ def create_interface():
                 batch_output = gr.Markdown(label="Batch Results")
 
                 batch_predict_btn.click(
-                    fn=app.predict_batch,
+                    # fn=app.predict_batch,
                     inputs=[batch_input, model_selector, confidence_slider],
                     outputs=batch_output
                 )

ui/classNames.txt

@@ -0,0 +1,39 @@
+Apple___Apple_scab
+Apple___Black_rot
+Apple___Cedar_apple_rust
+Apple___healthy
+Background_without_leaves
+Blueberry___healthy
+Cherry_(including_sour)_Powdery_mildew
+Cherry_(including_sour)_healthy
+Corn___Cercospora_leaf_spot Gray_leaf_spot
+Corn___Common_rust
+Corn___Northern_Leaf_Blight
+Corn___healthy
+Grape___Black_rot
+Grape__Esca(Black_Measles)
+Grape__Leaf_blight(Isariopsis_Leaf_Spot)
+Grape___healthy
+Orange__Haunglongbing(Citrus_greening)
+Peach___Bacterial_spot
+Peach___healthy
+Pepper,bell__Bacterial_spot
+Pepper,bell__healthy
+Potato___Early_blight
+Potato___Late_blight
+Potato___healthy
+Raspberry___healthy
+Soybean___healthy
+Squash___Powdery_mildew
+Strawberry___Leaf_scorch
+Strawberry___healthy
+Tomato___Bacterial_spot
+Tomato___Early_blight
+Tomato___Late_blight
+Tomato___Leaf_Mold
+Tomato___Septoria_leaf_spot
+Tomato__Spider_mites(Two-spotted_spider_mite)
+Tomato___Target_Spot
+Tomato___Tomato_Yellow_Leaf_Curl_Virus
+Tomato___Tomato_mosaic_virus
+Tomato___healthy

ui/config.py

@@ -5,11 +5,7 @@ MODEL_CONFIGS = {
     "CNN from Scratch": {
         "description": "Custom CNN model trained from scratch",
         "model_type": "cnn",
-        "clearml_task_id": "fe14662da63d45bf9208fdf9856d2fcc" 
-    },
-    "Transfer Learning (ResNet18)": {
-        "description": "Fine-tuned ResNet18 model",
-        "model_type": "resnet18",
-        "clearml_task_id": "SET_ME_TO_YOUR_RESNET_TASK_ID"
+        "clearml_task_id": "01345cf81fba4a2cac1176887bca9407" 
+        
     }
 }

ui/model_loader.py

@@ -2,9 +2,13 @@ import torch
 import sys
 from pathlib import Path
 import config
+import utils
+from clearml import Task
+from models.modelOne import modelOne
 
 sys.path.append(str(Path(__file__).parent.parent))
 
+MODEL_ARTIFACT_NAME = 'best_model'
 
 class ModelLoader:
     def __init__(self):
@@ -16,25 +20,48 @@ class ModelLoader:
         
         if not modelConfig:
             raise ValueError(f"ClearML configuration not found for model: {modelName}")
-
+        
         taskID = modelConfig['clearml_task_id']
-        modelType  = modelConfig['model_type']
 
         try:
-            print(f"attemtping to fetch '{modelName}' from clearML task: {taskID}")
+            print(f"Attempting to fetch '{modelName}' from ClearML task: {taskID}")
+
+            task = Task.get_task(task_id=taskID)
+            print("Available artifacts:", task.artifacts.keys())
+
+            # Fetch the artifact 'model_one.pt'
+            artifact = task.artifacts.get(MODEL_ARTIFACT_NAME)
+
+            if artifact is None:
+                raise RuntimeError(
+                    f"Artifact '{MODEL_ARTIFACT_NAME}' not found in ClearML task {taskID}"
+                )
+
+            modelPath = artifact.get_local_copy()
 
-            modelObject = Model(taskID=taskID)
-            modelPath = modelObject.get_local_copy()
+            if modelPath is None:
+                raise RuntimeError(
+                    f"Artifact '{MODEL_ARTIFACT_NAME}' could not be downloaded (returned None)"
+                )
 
-            model = self.loadRealModel(modelName, modelPath, modelType)
+            print(f"Weights downloaded to: {modelPath}")
+
+            # Load PyTorch model
+            model = modelOne(noOfClasses=39)
+            stateDict = torch.load(modelPath, map_location=self.device)
+            model.load_state_dict(stateDict)
+
+            model.to(self.device)
+            model.eval()
 
             return model
-        
+
         except Exception as e:
             print(f"Error loading from ClearML for {modelName}: {e}")
             raise RuntimeError(f"Failed to load model from ClearML: {e}")
+
     
-    def loadModel(self, modelName) :
+    def loadModel(self, modelName):
         if modelName in self.modelCache:
             return self.modelCache[modelName]
         

ui/utils.py

@@ -6,98 +6,72 @@ import torch
 import numpy as np
 from PIL import Image
 import torchvision.transforms as transforms
-import config
+import os
 
+IMAGE_SIZE = (256, 256)
+
+NORMALIZE_MEAN = [0.485, 0.456, 0.406]
+NORMALIZE_STD = [0.229, 0.224, 0.225]
+
+TOP_K_PREDICTIONS = 5
+CONFIDENCE_THRESHOLD = 0.01
 
-def preprocess_image(image, image_size=config.IMAGE_SIZE):
-    """
-    Preprocess image for model input
 
-    Args:
-        image: PIL Image or numpy array
-        image_size: Target size (height, width)
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
-    Returns:
-        Preprocessed tensor ready for model
+# Path to classNames.txt relative to this file
+CLASS_NAMES_FILE = os.path.join(BASE_DIR, "classNames.txt")
+
+with open(CLASS_NAMES_FILE, "r") as f:
+    CLASS_NAMES = [line.strip() for line in f.readlines() if line.strip()]
+
+def preprocess_image(image):
+    """
+    Preprocess image for model input
     """
-    # Convert to PIL Image if numpy array
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image.astype('uint8'))
 
-    # Convert RGBA to RGB if necessary
     if image.mode == 'RGBA':
         image = image.convert('RGB')
 
-    # Define preprocessing transforms
     transform = transforms.Compose([
-        transforms.Resize(image_size),
+        transforms.Resize(IMAGE_SIZE),
         transforms.ToTensor(),
-        transforms.Normalize(mean=config.NORMALIZE_MEAN, std=config.NORMALIZE_STD)
+        transforms.Normalize(NORMALIZE_MEAN, NORMALIZE_STD)
     ])
 
-    # Apply transforms
     tensor = transform(image)
+    return tensor.unsqueeze(0)
 
-    # Add batch dimension
-    tensor = tensor.unsqueeze(0)
 
-    return tensor
-
-
-def postprocess_predictions(logits, class_names=config.CLASS_NAMES, top_k=config.TOP_K_PREDICTIONS):
+def postprocess_predictions(logits, class_names=None, top_k=TOP_K_PREDICTIONS):
     """
-    Convert model logits to human-readable predictions
-
-    Args:
-        logits: Raw model output
-        class_names: List of class names
-        top_k: Number of top predictions to return
-
-    Returns:
-        Dictionary of predictions with confidences
+    Convert logits to formatted predictions
     """
-    # Convert logits to probabilities using softmax
-    probs = torch.nn.functional.softmax(logits, dim=1)
+    if class_names is None:
+        class_names = CLASS_NAMES
 
-    # Convert to numpy
+    probs = torch.nn.functional.softmax(logits, dim=1)
     probs = probs.cpu().detach().numpy()[0]
 
-    # Create predictions dictionary
     predictions = {name: float(prob) for name, prob in zip(class_names, probs)}
-
-    # Get top-k predictions
     top_predictions = sorted(predictions.items(), key=lambda x: x[1], reverse=True)[:top_k]
 
     return dict(top_predictions), predictions
 
 
-def format_prediction_for_display(predictions):
+def format_prediction_for_display(predictions, confidence_threshold=CONFIDENCE_THRESHOLD):
     """
-    Format predictions for Gradio display
-
-    Args:
-        predictions: Dictionary of class names and probabilities
-
-    Returns:
-        Dictionary formatted for Gradio Label component
+    Filter predictions for Gradio display
     """
-    # Filter out very low confidence predictions
-    filtered = {k: v for k, v in predictions.items() if v >= config.CONFIDENCE_THRESHOLD}
-
-    return filtered
+    return {k: v for k, v in predictions.items() if v >= confidence_threshold}
 
 
 def format_class_name(class_name):
     """
-    Format class name for better readability
-
-    Args:
-        class_name: Original class name (e.g., "Tomato___Late_blight")
-
-    Returns:
-        Formatted class name (e.g., "Tomato - Late blight")
+    Format class name into readable form
     """
-    # Replace underscores with spaces and split on ___
     parts = class_name.split("___")
 
     if len(parts) == 2:
@@ -105,74 +79,52 @@ def format_class_name(class_name):
         plant = plant.replace("_", " ")
         disease = disease.replace("_", " ")
         return f"{plant} - {disease}"
-    else:
-        return class_name.replace("_", " ")
+
+    return class_name.replace("_", " ")
 
 
 def get_disease_info(class_name):
     """
-    Get information about a disease (for future enhancement)
-
-    Args:
-        class_name: Disease class name
-
-    Returns:
-        Dictionary with disease information
+    Extract structured disease info from class name
     """
-    # This is a placeholder - you could expand this with actual disease information
     parts = class_name.split("___")
 
-    info = {
+    return {
         "plant": parts[0].replace("_", " ") if len(parts) > 0 else "Unknown",
         "disease": parts[1].replace("_", " ") if len(parts) > 1 else "Unknown",
         "is_healthy": "healthy" in class_name.lower(),
         "formatted_name": format_class_name(class_name)
     }
 
-    return info
-
 
 def batch_preprocess_images(images):
     """
-    Preprocess multiple images for batch prediction
-
-    Args:
-        images: List of PIL Images or numpy arrays
-
-    Returns:
-        Batched tensor ready for model
+    Preprocess a list of images into a batch tensor
     """
     tensors = [preprocess_image(img) for img in images]
-    batch = torch.cat(tensors, dim=0)
-    return batch
+    return torch.cat(tensors, dim=0)
 
 
 def create_confidence_label(predictions, top_k=5):
     """
-    Create a formatted string showing top predictions
-
-    Args:
-        predictions: Dictionary of predictions
-        top_k: Number of top predictions to show
-
-    Returns:
-        Formatted string
+    Render a formatted multiline prediction list
     """
     top_preds = sorted(predictions.items(), key=lambda x: x[1], reverse=True)[:top_k]
 
-    lines = []
-    for i, (class_name, prob) in enumerate(top_preds, 1):
-        formatted_name = format_class_name(class_name)
-        lines.append(f"{i}. {formatted_name}: {prob*100:.2f}%")
-
+    lines = [
+        f"{i}. {format_class_name(name)}: {prob*100:.2f}%"
+        for i, (name, prob) in enumerate(top_preds, 1)
+    ]
     return "\n".join(lines)
 
 
+def get_class_names():
+    """Return the loaded class names from the txt file."""
+    return CLASS_NAMES
+
 if __name__ == "__main__":
-    # Test utilities
     print("Testing utility functions...")
 
-    # Test class name formatting
     test_names = [
         "Tomato___Late_blight",
         "Apple___healthy",
@@ -183,7 +135,6 @@ if __name__ == "__main__":
     for name in test_names:
         print(f"  {name} -> {format_class_name(name)}")
 
-    # Test disease info
     print("\nDisease info:")
     for name in test_names:
         info = get_disease_info(name)
@@ -192,19 +143,8 @@ if __name__ == "__main__":
         print(f"    Disease: {info['disease']}")
         print(f"    Healthy: {info['is_healthy']}")
 
-    # Test image preprocessing
     print("\nImage preprocessing:")
     dummy_image = Image.new('RGB', (512, 512), color='red')
     tensor = preprocess_image(dummy_image)
     print(f"  Input size: {dummy_image.size}")
     print(f"  Output tensor shape: {tensor.shape}")
-
-    # Test mock predictions
-    print("\nMock predictions:")
-    from models.mock_model import create_mock_predictions
-    preds = create_mock_predictions(config.CLASS_NAMES)
-    top_preds, all_preds = postprocess_predictions(
-        torch.tensor([list(preds.values())]),
-        config.CLASS_NAMES
-    )
-    print(create_confidence_label(top_preds))