Merge branch 'develop'

.gitignore

@@ -1,10 +1,8 @@
-<<<<<<< HEAD
 .vscode/
 .venv/
 .vscode/
 .models/
 __pycache__/
-=======
 
 # Python environment
 venv/
@@ -18,4 +16,4 @@ __pycache__/
 
 # Generated files from data_preparation.py
 class_distribution.png
->>>>>>> 04cb88662062ef6b880c627546d067fa0cedfa8b
+

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3c19d6a5fea8043e6fda261763b7909aaed487b83991f29ca395b2ce7c8e591
+size 20532322

dataPrep/data_preparation.py

@@ -75,15 +75,15 @@ task.connect({
 })
 
 # ----- Load a subset from a given dataset & track with ClearML -----
-data_plants, subset_dataset, features = make_subset(
-    DATASET_LINK, DATASET_SUBSET_RATIO, task
+data_plants, prototyping_dataset, features, clearml_dataset = make_subset(
+    DATASET_LINK, DATASET_SUBSET_RATIO, clearml_logger
 )
 
 
 # ---- Exploratory data analysis (EDA) ----
 
 # Reformatting the label feature to understand bias
-labels_list = subset_dataset['label']
+labels_list = prototyping_dataset['label']
 df_labels = pd.Series(labels_list)
 label_count = df_labels.value_counts(sort=False)
 
@@ -112,7 +112,6 @@ 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}")
@@ -124,17 +123,16 @@ 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 Subset Dataset")
+plt.title("Class Distribution in Prototype Dataset")
 plt.xlabel("Class")
 plt.ylabel("Count")
 plt.tight_layout()
 
 clearml_logger.report_matplotlib_figure(
     title="EDA Class Distribution",
-    series="Subset Dataset",
+    series="Prototype Subset",
     figure=plt.gcf(),
     iteration=1
 )
@@ -152,7 +150,7 @@ if __name__ == "__main__":
     }
 
     prototype_loaders = make_dataset_loaders(
-        subset_dataset, SEED, BATCH_SIZE, TEST_SIZE, aug_config
+        prototyping_dataset, SEED, BATCH_SIZE, TEST_SIZE, aug_config
     )
 
     print("\n--- Handoff Test Successful ---")
@@ -176,9 +174,14 @@ 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

@@ -11,7 +11,7 @@ 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):
+def extract_latest_data_task(project_name: str = "Small Group Project", num_workers: int = 0):
 
   # --------- Get latest Data Preparation task from ClearML ---------
 

dataPrep/helpers/transforms_loaders.py

@@ -103,13 +103,15 @@ def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config, worke
         pin_memory=True,
         num_workers=workers
     )
+    class_names = dataset.features['label'].names
 
     print(f"\nWorkers used in DataLoaders: {workers}\n")
 
     dataset_loaders = {
         "train": train_loader,
         "val": val_loader,
-        "test": test_loader
+        "test": test_loader,
+        "classNames": class_names
     }
 
     return dataset_loaders

models/modelTwo.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BetterCNN(nn.Module):
+    def __init__(self, noOfClasses=39):
+        super(BetterCNN, self).__init__()
+
+        # 32 Channels 
+        # We use padding=1 to keep spatial size same before pooling
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(32)
+
+        # 64 Channels
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(64)
+
+        # 128 Channels      
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm2d(128)
+
+        # 256 Channels
+        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
+        self.bn4 = nn.BatchNorm2d(256)
+
+        # Pooling layer 
+        self.pool = nn.MaxPool2d(2, 2)
+
+        # Adaptive Pooling  
+        self.adaptive_pool = nn.AdaptiveAvgPool2d((4, 4))
+
+        # Classification Head
+        self.fc1 = nn.Linear(256 * 4 * 4, 1024)
+        self.dropout = nn.Dropout(0.5) # Dropout after Linear layer
+        
+        self.fc2 = nn.Linear(1024, 512)
+        self.fc3 = nn.Linear(512, noOfClasses)
+
+    def forward(self, x):
+        # Block 1
+        x = self.conv1(x)
+        x = self.bn1(x)     # BatchNorm 
+        x = F.relu(x)
+        x = self.pool(x)    
+
+        # Block 2
+        x = self.pool(F.relu(self.bn2(self.conv2(x))))
+
+        # Block 3
+        x = self.pool(F.relu(self.bn3(self.conv3(x))))
+
+        # Block 4
+        x = self.pool(F.relu(self.bn4(self.conv4(x))))
+        
+        # Adapt & Flatten
+        x = self.adaptive_pool(x) 
+        x = torch.flatten(x, 1)   # Flattens to (Batch, 4096)
+        
+        # Dense Layers
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)       # Regularization
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)           # No activation needed here (handled by CrossEntropyLoss)
+
+        return x

subset_indices.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:972615a5b506b5ee2490f61866c26a4a2f9e2498c0baedb195a2a0d10a62e76f
+size 111016

testingModel/helpers/evaluation.py

@@ -1,43 +1,88 @@
-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,
-    }
+import torch
+from torch.nn import CrossEntropyLoss
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+"""
+Evaluates a trained model on a dataloader that returns batches like:
+    batch["image"] -> Tensor [B, 3, 256, 256]
+    batch["label"] -> Tensor [B]
+"""
+def make_predictions(model, dataloader, device):
+
+    model.eval()
+    criterion = CrossEntropyLoss()
+
+    total_loss = 0
+    total_correct = 0
+    total_samples = 0
+
+    all_preds = []
+    all_labels = []
+
+    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)
+            preds = outputs.argmax(dim=1)
+
+            total_loss += loss.item() * images.size(0)
+            total_correct += (preds == labels).sum().item()
+            total_samples += labels.size(0)
+
+            # Accumulate all predictions and labels
+            all_preds.extend(preds.tolist())
+            all_labels.extend(labels.tolist())
+
+    accuracy = total_correct / total_samples
+    avg_loss = total_loss / total_samples
+
+    return {
+        "accuracy": accuracy,
+        "loss": avg_loss,
+        "predictions": np.array(all_preds),
+        "labels": np.array(all_labels),
+    }
+
+
+# Computes per-class accuracies
+def class_accuracies(labels, preds, num_classes):
+    correct = np.zeros(num_classes, dtype=int)
+    counts = np.zeros(num_classes, dtype=int)
+    accuracies = np.zeros(num_classes, dtype=float)
+
+    for true, pred in zip(labels, preds):
+        counts[true] += 1
+        if true == pred:
+            correct[true] += 1
+
+    # Calculate accuracies
+    for i in range(num_classes):
+        if counts[i] > 0:
+            accuracies[i] = round(correct[i] / counts[i], 4)
+        else:
+            accuracies[i] = 0.0 
+
+    return accuracies
+
+
+def plot_class_accuracies(accuracies, class_names):
+    fig, ax = plt.subplots(figsize=(12, 6))
+    
+    ax.set_title("Per-Class Accuracy")
+    ax.set_xlabel("Class")
+    ax.set_ylabel("Accuracy")
+    ax.set_ylim(0, 1.0)
+    ax.bar(class_names, accuracies)
+
+    plt.xticks(rotation=90)
+    plt.tight_layout()
+
+    return fig

testingModel/run_testing.py

@@ -1,76 +1,98 @@
-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()
+from clearml import Task
+from dataPrep.helpers.clearml_data import extract_latest_data_task
+
+import torch
+from models.modelOne import modelOne
+from models.modelTwo import BetterCNN
+from testingModel.helpers.evaluation import make_predictions, class_accuracies, plot_class_accuracies
+
+
+# -------------- 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 = "dca82d7c2f404c249f2e5325aaf77207"
+
+# 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 = BetterCNN(noOfClasses=testing_config["num_classes"])
+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"])
+
+# Compute per-class accuracy
+preds = subset_results["predictions"]
+labels = subset_results["labels"]
+class_acc = class_accuracies(
+    labels,
+    preds,
+    num_classes=testing_config["num_classes"]
+)
+
+# Plot with formatted class names
+class_names = subset_loaders['classNames']
+formatted_class_names = [" ".join(name.replace('_', ' ').split()) for name in class_names]
+acc_fig = plot_class_accuracies(class_acc, formatted_class_names)
+
+# Log accuracies plot to ClearML
+testing_logger.report_matplotlib_figure(
+    title="Subset Per-Class Accuracy",
+    series="Class Accuracy",
+    figure=acc_fig
+)
+
+
+# --------- 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/run_training.py

@@ -1,9 +1,9 @@
-import os
 from clearml import Task
 from dataPrep.helpers.clearml_data import extract_latest_data_task
 
 import torch
 from models.modelOne import modelOne
+from models.modelTwo import BetterCNN
 from trainingModel.helpers.Training import train_model
 
 
@@ -37,7 +37,7 @@ training_task.connect(training_config)
 
 
 # -------- Build the ML model --------
-model = modelOne(noOfClasses=training_config["num_classes"])
+model = BetterCNN(noOfClasses=training_config["num_classes"])
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 

ui/app.py

@@ -23,7 +23,7 @@ from config import *
 class PlantDiseaseApp:
     def __init__(self):
         self.model_loader = ModelLoader()
-        self.current_modelName = "CNN from Scratch"
+        self.current_modelName = list(config.MODEL_CONFIGS.keys())[0]
         self.model = self.model_loader.loadModel(self.current_modelName)
         self.flagged_predictions = []
         self.class_names = utils.get_class_names()
@@ -48,7 +48,7 @@ class PlantDiseaseApp:
         try:
             # Load model if needed
             if modelName != self.current_modelName:
-                self.model, self.class_names = self.model_loader.loadModel(modelName)
+                self.model = self.model_loader.loadModel(modelName)
                 self.current_modelName = modelName
 
             # Preprocess image
@@ -61,6 +61,10 @@ class PlantDiseaseApp:
             # Convert logits to probabilities
             probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy()[0]
 
+
+            predID = probs.argmax().item()
+            print("predicted index: " + str(predID))
+
             # Map to class names
             predictions = {name: float(prob) for name, prob in zip(self.class_names, probs)}
 
@@ -155,7 +159,7 @@ def create_interface():
         with gr.Row():
             model_selector = gr.Dropdown(
                 choices=list(config.MODEL_CONFIGS.keys()),
-                value="CNN from Scratch",
+                value="Shallow CNN",
                 label="Select Model",
                 info="Choose which model to use for predictions"
             )
@@ -218,25 +222,6 @@ def create_interface():
                     outputs=flag_output
                 )
 
-
-            with gr.Tab("Batch Processing"):
-                gr.Markdown("### Upload multiple images for batch processing")
-
-                batch_input = gr.File(
-                    label="Upload Multiple Images",
-                    file_count="multiple",
-                    type="filepath"
-                )
-
-                batch_predict_btn = gr.Button("Predict All", variant="primary")
-
-                batch_output = gr.Markdown(label="Batch Results")
-
-                batch_predict_btn.click(
-                    # fn=app.predict_batch,
-                    inputs=[batch_input, model_selector, confidence_slider],
-                    outputs=batch_output
-                )
             with gr.Tab("About"):
                 gr.Markdown(
                     """
@@ -258,7 +243,7 @@ def create_interface():
                     across 39 different plant disease categories.
 
                     ### Model Architecture
-                    - **CNN from Scratch**: Custom convolutional neural network
+                    - **Basic CNN**: Custom convolutional neural network
                     - **Transfer Learning**: Fine-tuned ResNet18 (if available)
 
                     ### Technology Stack

ui/config.py

@@ -2,10 +2,14 @@ CLEARML_PROJECT_NAME = "Plant Disease Classifier"
 CLEARML_TASK_NAME_DEFAULT = "CNN Training (Latest)"
 
 MODEL_CONFIGS = {
-    "CNN from Scratch": {
-        "description": "Custom CNN model trained from scratch",
-        "model_type": "cnn",
-        "clearml_task_id": "01345cf81fba4a2cac1176887bca9407" 
-        
+    "intermediate model": {
+        "description": "modelTwo trained on 20 epochs",
+        "class" : "betterCNN",
+        "clearml_task_id": "dca82d7c2f404c249f2e5325aaf77207"
+    },
+    "advanced model": {
+        "description" : "modleTwo trained on 30 epochs",
+        "class": "betterCNN",
+        "clearml_task_id": "c79a6939b46a4882a7fdaee117b1f32e"
     }
 }

ui/model_loader.py

@@ -2,12 +2,17 @@ import torch
 import sys
 from pathlib import Path
 import config
-import utils
 from clearml import Task
 from models.modelOne import modelOne
+from models.modelTwo import BetterCNN
 
 sys.path.append(str(Path(__file__).parent.parent))
 
+MODEL_CLASSES = {
+    "modelOne": modelOne,
+    "betterCNN": BetterCNN
+}
+
 MODEL_ARTIFACT_NAME = 'best_model'
 
 class ModelLoader:
@@ -22,6 +27,7 @@ class ModelLoader:
             raise ValueError(f"ClearML configuration not found for model: {modelName}")
         
         taskID = modelConfig['clearml_task_id']
+        className = modelConfig['class']
 
         try:
             print(f"Attempting to fetch '{modelName}' from ClearML task: {taskID}")
@@ -29,16 +35,13 @@ class ModelLoader:
             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()
-
             if modelPath is None:
                 raise RuntimeError(
                     f"Artifact '{MODEL_ARTIFACT_NAME}' could not be downloaded (returned None)"
@@ -46,8 +49,11 @@ class ModelLoader:
 
             print(f"Weights downloaded to: {modelPath}")
 
-            # Load PyTorch model
-            model = modelOne(noOfClasses=39)
+            # Load correct model class
+            ModelClass = MODEL_CLASSES[className]
+            model = ModelClass(noOfClasses=39)
+
+            # Load weights
             stateDict = torch.load(modelPath, map_location=self.device)
             model.load_state_dict(stateDict)
 

ui/utils.py

@@ -97,14 +97,6 @@ def get_disease_info(class_name):
     }
 
 
-def batch_preprocess_images(images):
-    """
-    Preprocess a list of images into a batch tensor
-    """
-    tensors = [preprocess_image(img) for img in images]
-    return torch.cat(tensors, dim=0)
-
-
 def create_confidence_label(predictions, top_k=5):
     """
     Render a formatted multiline prediction list
@@ -120,31 +112,4 @@ def create_confidence_label(predictions, top_k=5):
 
 def get_class_names():
     """Return the loaded class names from the txt file."""
-    return CLASS_NAMES
-
-if __name__ == "__main__":
-    print("Testing utility functions...")
-
-    test_names = [
-        "Tomato___Late_blight",
-        "Apple___healthy",
-        "Corn_(maize)___Common_rust_"
-    ]
-
-    print("\nClass name formatting:")
-    for name in test_names:
-        print(f"  {name} -> {format_class_name(name)}")
-
-    print("\nDisease info:")
-    for name in test_names:
-        info = get_disease_info(name)
-        print(f"  {name}:")
-        print(f"    Plant: {info['plant']}")
-        print(f"    Disease: {info['disease']}")
-        print(f"    Healthy: {info['is_healthy']}")
-
-    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}")
+    return CLASS_NAMES