train_model.py

"""
Plant Disease Detection - Model Training Script
================================================
Downloads the PlantVillage dataset and trains the CNN model with:
  - Data augmentation (flips, rotations, color jitter, etc.)
  - Train / Validation / Test splits (70% / 15% / 15%)
  - Learning rate scheduling & early stopping
  - TensorBoard logging (optional)
  - Best-model checkpointing

Usage:
    python train_model.py                       # train from scratch
    python train_model.py --resume              # resume from last checkpoint
    python train_model.py --epochs 30           # custom epoch count
    python train_model.py --data_dir ./my_data  # use your own dataset folder
"""

import os
import sys
import json
import time
import shutil
import random
import argparse
import zipfile
from pathlib import Path
from datetime import datetime

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split, Subset
from torchvision import datasets, transforms
from PIL import Image

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from models.cnn_model import CNN, idx_to_classes

# ──────────────────────────── Configuration ────────────────────────────

DEFAULT_CONFIG = {
    "num_classes": 39,
    "image_size": 224,
    "batch_size": 32,
    "epochs": 25,
    "learning_rate": 0.001,
    "weight_decay": 1e-4,
    "lr_step_size": 7,
    "lr_gamma": 0.1,
    "early_stopping_patience": 7,
    "train_ratio": 0.70,
    "val_ratio": 0.15,
    "test_ratio": 0.15,
    "num_workers": 2,
    "seed": 42,
}

# Classes expected (must match idx_to_classes in cnn_model.py)
CLASSES = [idx_to_classes[i] for i in range(len(idx_to_classes))]

# ──────────────────────────── Helpers ──────────────────────────────────

def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def download_dataset(dest_dir: str) -> str:
    """
    Attempt to download the PlantVillage dataset.
    Tries multiple methods:
      1. kagglehub (pip install kagglehub)
      2. kaggle CLI
      3. Manual instruction fallback
    Returns the path to the extracted dataset root that contains class folders.
    """
    dataset_root = os.path.join(dest_dir, "PlantVillage")
    if os.path.isdir(dataset_root) and len(os.listdir(dataset_root)) >= 30:
        print(f"[INFO] Dataset already present at {dataset_root} "
              f"({len(os.listdir(dataset_root))} class folders found).")
        return dataset_root

    os.makedirs(dest_dir, exist_ok=True)

    # ── Method 1: kagglehub ────────────────────────────────────────────
    try:
        import kagglehub  # type: ignore
        print("[INFO] Downloading dataset via kagglehub …")
        path = kagglehub.dataset_download("abdallahalidev/plantvillage-dataset")
        # The download may contain nested folders; find the one with class dirs
        for root, dirs, _ in os.walk(path):
            if len(dirs) >= 30:
                if root != dataset_root:
                    shutil.copytree(root, dataset_root, dirs_exist_ok=True)
                print(f"[INFO] Dataset ready at {dataset_root}")
                return dataset_root
    except Exception as e:
        print(f"[WARN] kagglehub download failed: {e}")

    # ── Method 2: kaggle CLI ───────────────────────────────────────────
    try:
        print("[INFO] Trying kaggle CLI download …")
        ret = os.system(
            f'kaggle datasets download -d abdallahalidev/plantvillage-dataset '
            f'-p "{dest_dir}" --unzip'
        )
        if ret == 0:
            for root, dirs, _ in os.walk(dest_dir):
                if len(dirs) >= 30:
                    if root != dataset_root:
                        shutil.copytree(root, dataset_root, dirs_exist_ok=True)
                    print(f"[INFO] Dataset ready at {dataset_root}")
                    return dataset_root
    except Exception:
        pass

    # ── Method 3: Manual fallback ──────────────────────────────────────
    print("\n" + "=" * 70)
    print("AUTOMATIC DOWNLOAD FAILED – please download manually:")
    print("=" * 70)
    print("1. Go to: https://www.kaggle.com/datasets/abdallahalidev/plantvillage-dataset")
    print("2. Click 'Download' (you may need a free Kaggle account)")
    print(f"3. Extract the ZIP so that class folders are inside:\n   {dataset_root}")
    print("   The structure should look like:")
    print("     PlantVillage/")
    print("       Apple___Apple_scab/")
    print("       Apple___Black_rot/")
    print("       ...")
    print("4. Re-run this script.\n")
    print("Alternatively, install kagglehub:  pip install kagglehub")
    print("Or set up Kaggle API credentials: https://github.com/Kaggle/kaggle-api#api-credentials")
    print("=" * 70 + "\n")
    sys.exit(1)


def find_dataset_root(data_dir: str) -> str:
    """Walk data_dir to find the folder that directly contains class subfolders."""
    # Check if data_dir itself has class folders
    subdirs = [d for d in os.listdir(data_dir)
               if os.path.isdir(os.path.join(data_dir, d))]
    if len(subdirs) >= 30:
        return data_dir

    # Search one level deeper
    for sub in subdirs:
        candidate = os.path.join(data_dir, sub)
        inner = [d for d in os.listdir(candidate)
                 if os.path.isdir(os.path.join(candidate, d))]
        if len(inner) >= 30:
            return candidate

    # Search two levels deeper
    for sub in subdirs:
        candidate = os.path.join(data_dir, sub)
        if os.path.isdir(candidate):
            for sub2 in os.listdir(candidate):
                candidate2 = os.path.join(candidate, sub2)
                if os.path.isdir(candidate2):
                    inner = [d for d in os.listdir(candidate2)
                             if os.path.isdir(os.path.join(candidate2, d))]
                    if len(inner) >= 30:
                        return candidate2

    return data_dir  # fallback


# ──────────────────────────── Transforms ───────────────────────────────

def get_transforms(image_size: int):
    """Return train and validation/test transforms with augmentation."""
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(image_size, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.RandomRotation(20),
        transforms.ColorJitter(brightness=0.2, contrast=0.2,
                               saturation=0.2, hue=0.1),
        transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225]),
    ])

    val_transform = transforms.Compose([
        transforms.Resize(image_size + 32),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225]),
    ])

    return train_transform, val_transform


# ──────────────────────────── Training ─────────────────────────────────

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (images, labels) in enumerate(loader):
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

        if (batch_idx + 1) % 50 == 0:
            print(f"    Batch {batch_idx+1}/{len(loader)} | "
                  f"Loss: {loss.item():.4f} | "
                  f"Acc: {100.*correct/total:.2f}%")

    epoch_loss = running_loss / total
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)

        running_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc


def main():
    parser = argparse.ArgumentParser(description="Train Plant Disease CNN")
    parser.add_argument("--data_dir", type=str, default=None,
                        help="Path to dataset root with class subfolders. "
                             "If omitted, the script downloads PlantVillage.")
    parser.add_argument("--epochs", type=int, default=DEFAULT_CONFIG["epochs"])
    parser.add_argument("--batch_size", type=int, default=DEFAULT_CONFIG["batch_size"])
    parser.add_argument("--lr", type=float, default=DEFAULT_CONFIG["learning_rate"])
    parser.add_argument("--resume", action="store_true",
                        help="Resume training from latest checkpoint")
    parser.add_argument("--output_dir", type=str, default="training_output",
                        help="Directory to save checkpoints and logs")
    args = parser.parse_args()

    cfg = DEFAULT_CONFIG.copy()
    cfg["epochs"] = args.epochs
    cfg["batch_size"] = args.batch_size
    cfg["learning_rate"] = args.lr

    set_seed(cfg["seed"])
    device = get_device()
    print(f"\n{'='*60}")
    print(f"  Plant Disease Detection – Model Training")
    print(f"{'='*60}")
    print(f"  Device       : {device}")
    print(f"  Epochs       : {cfg['epochs']}")
    print(f"  Batch size   : {cfg['batch_size']}")
    print(f"  Learning rate: {cfg['learning_rate']}")
    print(f"{'='*60}\n")

    # ── 1. Prepare dataset ─────────────────────────────────────────────
    if args.data_dir and os.path.isdir(args.data_dir):
        dataset_root = find_dataset_root(args.data_dir)
    else:
        data_cache = os.path.join(os.path.dirname(__file__), "dataset_cache")
        dataset_root = download_dataset(data_cache)
        dataset_root = find_dataset_root(dataset_root)

    print(f"[INFO] Using dataset at: {dataset_root}")
    subdirs = sorted([d for d in os.listdir(dataset_root)
                      if os.path.isdir(os.path.join(dataset_root, d))])
    print(f"[INFO] Found {len(subdirs)} class folders")
    num_classes = len(subdirs)
    if num_classes != cfg["num_classes"]:
        print(f"[WARN] Expected {cfg['num_classes']} classes but found {num_classes}. "
              f"Adjusting model to {num_classes} classes.")
        cfg["num_classes"] = num_classes

    # ── 2. Create transforms & datasets ────────────────────────────────
    train_transform, val_transform = get_transforms(cfg["image_size"])

    full_dataset = datasets.ImageFolder(dataset_root, transform=train_transform)

    # Verify class mapping matches
    print(f"[INFO] Dataset classes: {len(full_dataset.classes)}")
    print(f"[INFO] Total images  : {len(full_dataset)}")

    # Save class-to-index mapping
    os.makedirs(args.output_dir, exist_ok=True)
    class_map_path = os.path.join(args.output_dir, "class_mapping.json")
    with open(class_map_path, "w") as f:
        json.dump(full_dataset.class_to_idx, f, indent=2)
    print(f"[INFO] Class mapping saved to {class_map_path}")

    # ── 3. Split into train / val / test ───────────────────────────────
    total = len(full_dataset)
    train_size = int(total * cfg["train_ratio"])
    val_size = int(total * cfg["val_ratio"])
    test_size = total - train_size - val_size

    train_set, val_set, test_set = random_split(
        full_dataset, [train_size, val_size, test_size],
        generator=torch.Generator().manual_seed(cfg["seed"])
    )

    # Apply val transform to val/test sets (no augmentation)
    class TransformSubset(torch.utils.data.Dataset):
        def __init__(self, subset, transform):
            self.subset = subset
            self.transform = transform
        def __len__(self):
            return len(self.subset)
        def __getitem__(self, idx):
            img, label = self.subset[idx]
            # img is already transformed by train_transform, we need raw
            # Instead, we'll load from the path directly
            sample_idx = self.subset.indices[idx]
            path, label = self.subset.dataset.samples[sample_idx]
            img = Image.open(path).convert("RGB")
            img = self.transform(img)
            return img, label

    val_set_proper = TransformSubset(val_set, val_transform)
    test_set_proper = TransformSubset(test_set, val_transform)

    print(f"[INFO] Split: Train={train_size} | Val={val_size} | Test={test_size}")

    train_loader = DataLoader(train_set, batch_size=cfg["batch_size"],
                              shuffle=True, num_workers=cfg["num_workers"],
                              pin_memory=True)
    val_loader = DataLoader(val_set_proper, batch_size=cfg["batch_size"],
                            shuffle=False, num_workers=cfg["num_workers"],
                            pin_memory=True)
    test_loader = DataLoader(test_set_proper, batch_size=cfg["batch_size"],
                             shuffle=False, num_workers=cfg["num_workers"],
                             pin_memory=True)

    # ── 4. Build model ─────────────────────────────────────────────────
    model = CNN(cfg["num_classes"]).to(device)
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[INFO] Model parameters: {total_params:,} total, {trainable_params:,} trainable")

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=cfg["learning_rate"],
                           weight_decay=cfg["weight_decay"])
    scheduler = optim.lr_scheduler.StepLR(optimizer,
                                          step_size=cfg["lr_step_size"],
                                          gamma=cfg["lr_gamma"])

    # ── 5. Resume from checkpoint if requested ─────────────────────────
    start_epoch = 0
    best_val_acc = 0.0
    checkpoint_path = os.path.join(args.output_dir, "checkpoint_latest.pt")

    if args.resume and os.path.exists(checkpoint_path):
        print(f"[INFO] Resuming from {checkpoint_path}")
        ckpt = torch.load(checkpoint_path, map_location=device)
        model.load_state_dict(ckpt["model_state_dict"])
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
        start_epoch = ckpt["epoch"] + 1
        best_val_acc = ckpt.get("best_val_acc", 0.0)
        print(f"[INFO] Resumed at epoch {start_epoch}, best val acc: {best_val_acc:.2f}%")

    # ── 6. Training loop ──────────────────────────────────────────────
    history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
    patience_counter = 0
    training_start = time.time()

    print(f"\n{'─'*60}")
    print(f"  Starting training …")
    print(f"{'─'*60}\n")

    for epoch in range(start_epoch, cfg["epochs"]):
        epoch_start = time.time()
        current_lr = optimizer.param_groups[0]["lr"]
        print(f"Epoch [{epoch+1}/{cfg['epochs']}]  (LR: {current_lr:.6f})")

        # Train
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )

        # Validate
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)

        scheduler.step()

        elapsed = time.time() - epoch_start
        print(f"  Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
        print(f"  Val   Loss: {val_loss:.4f} | Val   Acc: {val_acc:.2f}%")
        print(f"  Time: {elapsed:.1f}s\n")

        history["train_loss"].append(train_loss)
        history["train_acc"].append(train_acc)
        history["val_loss"].append(val_loss)
        history["val_acc"].append(val_acc)

        # ── Save checkpoint ────────────────────────────────────────────
        checkpoint = {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "best_val_acc": max(best_val_acc, val_acc),
            "config": cfg,
        }
        torch.save(checkpoint, checkpoint_path)

        # Save best model
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_model_path = os.path.join(args.output_dir, "best_model.pt")
            torch.save(model.state_dict(), best_model_path)
            print(f"  ★ New best model saved! (Val Acc: {val_acc:.2f}%)\n")
            patience_counter = 0
        else:
            patience_counter += 1

        # ── Early stopping ─────────────────────────────────────────────
        if patience_counter >= cfg["early_stopping_patience"]:
            print(f"[INFO] Early stopping triggered after {patience_counter} "
                  f"epochs without improvement.")
            break

    total_time = time.time() - training_start
    print(f"\n{'='*60}")
    print(f"  Training Complete!")
    print(f"  Total time: {total_time/60:.1f} minutes")
    print(f"  Best Val Accuracy: {best_val_acc:.2f}%")
    print(f"{'='*60}\n")

    # ── 7. Test evaluation ─────────────────────────────────────────────
    print("Evaluating on test set …")
    best_model_path = os.path.join(args.output_dir, "best_model.pt")
    if os.path.exists(best_model_path):
        model.load_state_dict(torch.load(best_model_path, map_location=device))
    test_loss, test_acc = evaluate(model, test_loader, criterion, device)
    print(f"  Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%\n")

    # ── 8. Copy best model to models/ for use by the app ───────────────
    app_model_path = os.path.join(os.path.dirname(__file__),
                                  "models", "plant_disease_model_1.pt")
    backup_path = app_model_path + ".backup"

    if os.path.exists(app_model_path):
        shutil.copy2(app_model_path, backup_path)
        print(f"[INFO] Previous model backed up to {backup_path}")

    shutil.copy2(best_model_path, app_model_path)
    print(f"[INFO] Best model copied to {app_model_path}")

    # ── 9. Save training history ───────────────────────────────────────
    history_path = os.path.join(args.output_dir, "training_history.json")
    history["test_loss"] = test_loss
    history["test_acc"] = test_acc
    history["best_val_acc"] = best_val_acc
    history["total_time_minutes"] = total_time / 60
    history["config"] = cfg
    history["timestamp"] = datetime.now().isoformat()

    with open(history_path, "w") as f:
        json.dump(history, f, indent=2)
    print(f"[INFO] Training history saved to {history_path}")

    print(f"\n{'='*60}")
    print(f"  Summary")
    print(f"{'='*60}")
    print(f"  Best Val Accuracy : {best_val_acc:.2f}%")
    print(f"  Test Accuracy     : {test_acc:.2f}%")
    print(f"  Model saved to    : {app_model_path}")
    print(f"  Backup saved to   : {backup_path}")
    print(f"{'='*60}\n")


if __name__ == "__main__":
    main()