Upload 3 files

data_prep.py

@@ -0,0 +1,146 @@
+import torch
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+from utils.config import load_config
+
+config = load_config()
+batch_size = config["batch_size"]
+num_workers = config["num_workers"]
+mean_nm = config["normalize_mean"]
+std_nm = config["normalize_std"]
+execute_remotely = config.get("execute_remotely", False)
+
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+
+# set dataset to clearml dataset if executing remotely or load from huggingface otherwise
+if config["execute_remotely"]:
+    from clearml import Dataset as ClearMLDataset
+    clearml_dataset = ClearMLDataset.get(dataset_id="0c3de7af2d98482dacf41633a0587845")
+    dataset_path = clearml_dataset.get_local_copy()
+    dataset = load_dataset(dataset_path)
+else:
+    dataset = load_dataset("DScomp380/plant_village", cache_dir="./data_cache")
+#split dataset into train(70%), and 30% remaining for val and test
+splits = dataset["train"].train_test_split(test_size=0.30, seed=42)
+train_split = splits["train"] #training set
+remaining = splits["test"]
+
+#split remaining 30% into val(15%) and test(15%)    
+val_test = remaining.train_test_split(test_size=0.5, seed=42)
+val_split = val_test["train"] #validation set
+test_split = val_test["test"] #test set
+
+preprocess_transform = transforms.Compose([
+    # resize images to 224x224, convert to tensor, and normalize
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=mean_nm, std=std_nm)
+])
+
+def preprocess_batch(batch):
+    batch["pixel_values"] = [preprocess_transform(img) for img in batch["image"]]
+    return batch
+
+if execute_remotely:
+    def train_transform_batch(batch):
+        batch["pixel_values"] = [preprocess_transform(img) for img in batch["image"]]
+        return batch
+    train_split = train_split.with_transform(train_transform_batch)
+    val_split = val_split.with_transform(train_transform_batch)
+    test_split = test_split.with_transform(train_transform_batch)
+else:
+    train_split = train_split.map(
+        preprocess_batch,
+        batched=True,
+        batch_size=100,
+        remove_columns=["image"],
+        cache_file_name="./data_cache/train_preprocessed.arrow"
+    )
+
+    val_split = val_split.map(
+        preprocess_batch,
+        batched=True,
+        batch_size=100,
+        remove_columns=["image"],
+        cache_file_name="./data_cache/val_preprocessed.arrow"
+    )
+
+    test_split = test_split.map(
+        preprocess_batch,
+        batched=True,
+        batch_size=100,
+        remove_columns=["image"],
+        cache_file_name="./data_cache/test_preprocessed.arrow"
+    )
+
+    train_split.set_format(type="torch", columns=["pixel_values", "label"])
+    val_split.set_format(type="torch", columns=["pixel_values", "label"])
+    test_split.set_format(type="torch", columns=["pixel_values", "label"])
+
+# augmentations
+train_augment = transforms.Compose([
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.3),
+    transforms.RandomRotation(degrees=15),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+    transforms.RandomApply([
+        transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0))
+    ], p=0.3),
+])
+
+
+def train_collate_fn(batch):
+    pixel_values = [item["pixel_values"] for item in batch]
+    labels = [item["label"] for item in batch]
+    
+    augmented = [train_augment(img) for img in pixel_values] # apply augmentation while training
+    
+    return {
+        "pixel_values": torch.stack(augmented),
+        "labels": torch.tensor(labels)
+    }
+
+def val_collate_fn(batch):
+    return {
+        "pixel_values": torch.stack([item["pixel_values"] for item in batch]),
+        "labels": torch.tensor([item["label"] for item in batch])
+    }
+
+# create DataLoaders for train, val, and test sets
+train_loader = DataLoader(
+    train_split, 
+    batch_size=batch_size, 
+    shuffle=True, 
+    num_workers=num_workers, 
+    pin_memory=True,
+    persistent_workers=True if num_workers > 0 else False,
+    collate_fn=train_collate_fn
+)
+
+val_loader = DataLoader(
+    val_split, 
+    batch_size=batch_size, 
+    shuffle=False, 
+    num_workers=num_workers, 
+    pin_memory=True,
+    persistent_workers=True if num_workers > 0 else False,
+    collate_fn=val_collate_fn
+)
+
+test_loader = DataLoader(
+    test_split, 
+    batch_size=batch_size, 
+    shuffle=False, 
+    num_workers=num_workers, 
+    pin_memory=True,
+    persistent_workers=True if num_workers > 0 else False,
+    collate_fn=val_collate_fn
+)
+
+if __name__ == "__main__":
+    print(f"Device: {device}")
+    print(f"Train samples: {len(train_split)}")
+    print(f"Val samples: {len(val_split)}")
+    print(f"Test samples: {len(test_split)}")
+    print(f"Batches per epoch: {len(train_loader)}")

test.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+from pathlib import Path
+from data_prep import test_loader, device
+from models.model import PlantCNN
+from utils.config import load_config
+from clearml import Task
+import numpy as np
+from utils.vis import visualize_preds, plot_cfm
+from tqdm.auto import tqdm
+
+
+def evaluate_on_test(model, loader, loss_fn, device, num_imgs):
+    model.eval()
+    all_labels = []
+    all_preds = []
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    imgs_to_display = []
+    lbls_to_display = []
+    prs_to_display = []
+    
+    with torch.no_grad():
+        for batch_idx, batch in enumerate(tqdm(loader, desc="Val", leave=False)):
+            images = batch["pixel_values"].to(device)
+            labels = batch["labels"].to(device)
+            
+            output = model(images)
+            loss = loss_fn(output, labels)
+            
+            running_loss += loss.item()*labels.size(0)
+            
+            _, preds = torch.max(output, dim=1)
+            correct += (preds==labels).sum().item()
+            total += labels.size(0)
+            
+            all_labels.extend(labels.cpu().numpy())
+            all_preds.extend(preds.cpu().numpy())
+            
+            if len(imgs_to_display) < num_imgs:
+                remaining = num_imgs - len(imgs_to_display)
+                for img, lbl, pr in zip(images[:remaining], preds[:remaining], preds[:remaining]):
+                    imgs_to_display.append(img.cpu())
+                    lbls_to_display.append(lbl.item())
+                    prs_to_display.append(pr.item())
+                    
+    test_loss = running_loss / total
+    test_acc = correct / total
+    return test_loss, test_acc, all_labels, all_preds, imgs_to_display, lbls_to_display, prs_to_display
+    
+
+def main():
+    config = load_config()
+    num_classes = config["num_classes"]
+    channels = config["channels"]
+    dropout = config["dropout"]
+    lr = config["lr"]
+    project_name = "GAP_plant_disease_classification"
+    model_name = "PlantCNN"
+    mean_nm = config["normalize_mean"]
+    std_nm = config["normalize_std"]
+    
+    task = Task.init(project_name=project_name, task_name=f"{model_name}_test")
+    task.connect(config)
+    task.add_tags([model_name, "test"])
+    logger = task.get_logger()
+    
+    dataset = test_loader.dataset
+    class_names = dataset.features["label"].names
+    
+    model = PlantCNN(num_classes=num_classes, channels=channels, dropout=dropout).to(device)
+    project_root = Path(__file__).resolve().parent
+    model_path = project_root / "saved_models" / "plant_cnn.pt"
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict)
+    
+    loss_fn = nn.CrossEntropyLoss()
+    
+    test_loss, test_acc, all_labels, all_preds, display_images, display_labels, display_preds = evaluate_on_test(model, test_loader,
+                                                                                                                 loss_fn, device,
+                                                                                                                 num_imgs=24)
+    
+    print("\nTest results:")
+    print(f"Test loss: {test_loss:.3f} | Test accuracy: {test_acc:.3f}")
+    logger.report_scalar("loss", "test", test_loss, 0)
+    logger.report_scalar("accuracy", "test", test_acc, 0)
+    
+    visualize_preds(display_images, display_labels, display_preds, logger, class_names, mean_nm, std_nm, num_images=24)
+    plot_cfm(all_labels, all_preds, logger, class_names, num_classes)
+    
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+from pathlib import Path
+from data_prep import test_loader, device
+from models.model import PlantCNN
+from utils.config import load_config
+from clearml import Task, InputModel
+import numpy as np
+from utils.vis import visualize_preds, plot_cfm
+from tqdm.auto import tqdm
+import ast
+
+
+def evaluate_on_test(model, loader, loss_fn, device, num_imgs):
+    model.eval()
+    all_labels = []
+    all_preds = []
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    imgs_to_display = []
+    lbls_to_display = []
+    prs_to_display = []
+    
+    with torch.no_grad():
+        for batch_idx, batch in enumerate(tqdm(loader, desc="Val", leave=False)):
+            images = batch["pixel_values"].to(device)
+            labels = batch["labels"].to(device)
+            
+            output = model(images)
+            loss = loss_fn(output, labels)
+            
+            running_loss += loss.item()*labels.size(0)
+            
+            _, preds = torch.max(output, dim=1)
+            correct += (preds==labels).sum().item()
+            total += labels.size(0)
+            
+            all_labels.extend(labels.cpu().numpy())
+            all_preds.extend(preds.cpu().numpy())
+            
+            if len(imgs_to_display) < num_imgs:
+                remaining = num_imgs - len(imgs_to_display)
+                for img, lbl, pr in zip(images[:remaining], preds[:remaining], preds[:remaining]):
+                    imgs_to_display.append(img.cpu())
+                    lbls_to_display.append(lbl.item())
+                    prs_to_display.append(pr.item())
+                    
+    test_loss = running_loss / total
+    test_acc = correct / total
+    return test_loss, test_acc, all_labels, all_preds, imgs_to_display, lbls_to_display, prs_to_display
+    
+
+
+
+def main():
+    project_name = "GAP_plant_disease_classification"
+    model_name = "PlantCNN"
+    
+    task = Task.init(project_name=project_name, task_name=f"{model_name}_test")
+    logger = task.get_logger()
+    
+    input_model = InputModel(model_id="b9308022b85e4eea952d78124d1ee597")
+    
+    training_task_id = input_model.task
+    training_task = Task.get_task(task_id=training_task_id)
+    
+    training_params = training_task.get_parameters()
+    print(f"Training parameters: {training_params}")
+    
+    num_classes = int(training_params.get("General/num_classes"))
+    channels = ast.literal_eval(training_params.get("General/channels"))
+    dropout = float(training_params.get("General/dropout"))
+    mean_nm = ast.literal_eval(training_params.get("General/normalize_mean"))
+    std_nm = ast.literal_eval(training_params.get("General/normalize_std"))
+    kernel_sizes = ast.literal_eval(training_params.get("General/kernel_sizes"))
+    
+    task.add_tags([model_name, "test"])
+    
+    dataset = test_loader.dataset
+    class_names = dataset.features["label"].names
+
+
+    model = PlantCNN(num_classes=num_classes, channels=channels, dropout=dropout, kernel_sizes=kernel_sizes).to(device)
+    model_path = input_model.get_local_copy()
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict)
+    
+    loss_fn = nn.CrossEntropyLoss()
+    
+    test_loss, test_acc, all_labels, all_preds, display_images, display_labels, display_preds = evaluate_on_test(model, test_loader,
+                                                                                                                 loss_fn, device,
+                                                                                                                 num_imgs=24)
+    
+    print("\nTest results:")
+    print(f"Test loss: {test_loss:.3f} | Test accuracy: {test_acc:.3f}")
+    logger.report_scalar("loss", "test", test_loss, 0)
+    logger.report_scalar("accuracy", "test", test_acc, 0)
+    
+    visualize_preds(display_images, display_labels, display_preds, logger, class_names, mean_nm, std_nm, num_images=24)
+    plot_cfm(all_labels, all_preds, logger, class_names, num_classes)
+    
+if __name__ == "__main__":
+    main()