ai-detection/train.py

import os
import sys
import json
import time
import argparse
import torch
import torchvision
from torch.utils.data import DataLoader
import albumentations as A
import numpy as np

from dataset import FruitVegetableDataset, collate_fn
from data_generator import build_dataset

STATUS_FILE = "training_status.json"

def update_status(status, epoch=0, total_epochs=0, train_loss=[], val_loss=[], progress=0.0, message=""):
    """
    Saves current training status to a JSON file so that the FastAPI server
    can read it and report it to the frontend UI in real-time.
    """
    status_data = {
        "status": status,
        "epoch": epoch,
        "total_epochs": total_epochs,
        "train_loss": train_loss,
        "val_loss": val_loss,
        "progress": progress,
        "message": message,
        "timestamp": time.time()
    }
    with open(STATUS_FILE, "w") as f:
        json.dump(status_data, f, indent=2)

def train(args):
    # Ensure dataset exists
    if not os.path.exists(args.dataset_dir) or len(os.listdir(os.path.join(args.dataset_dir, "images", "train"))) == 0:
        update_status("generating_data", message="Generating synthetic fruit & vegetable dataset...")
        print("Generating dataset...")
        build_dataset(base_dir=args.dataset_dir, train_count=200, val_count=50)
        
    update_status("training", epoch=0, total_epochs=args.epochs, progress=5.0, message="Initializing model and loading dataset...")
    
    # Define augmentations using Albumentations
    train_transform = A.Compose([
        A.Resize(320, 320),
        A.HorizontalFlip(p=0.5),
        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.15, rotate_limit=30, p=0.7),
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.7),
        A.GaussNoise(p=0.3),
    ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['category_ids']))
    
    val_transform = A.Compose([
        A.Resize(320, 320),
    ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['category_ids']))
    
    # Datasets and loaders
    train_dataset = FruitVegetableDataset(
        images_dir=os.path.join(args.dataset_dir, "images", "train"),
        labels_dir=os.path.join(args.dataset_dir, "labels", "train"),
        transform=train_transform
    )
    
    val_dataset = FruitVegetableDataset(
        images_dir=os.path.join(args.dataset_dir, "images", "val"),
        labels_dir=os.path.join(args.dataset_dir, "labels", "val"),
        transform=val_transform
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=0, # 0 works best on Windows to avoid multi-processing issues
        collate_fn=collate_fn,
        drop_last=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0,
        collate_fn=collate_fn
    )
    
    # Set number of classes dynamically (classes + 1 background)
    from data_generator import CLASSES
    num_classes = len(CLASSES) + 1
    
    print("Building SSDLite MobileNetV3 model...")
    # Load pretrained model with backbone weights
    model = torchvision.models.detection.ssdlite320_mobilenet_v3_large(
        num_classes=num_classes,
        weights_backbone=torchvision.models.MobileNet_V3_Large_Weights.DEFAULT
    )
    
    # Send to GPU if available, else CPU
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    model.to(device)
    
    # Optimizer and learning rate scheduler
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0005)
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
    
    train_loss_history = []
    val_loss_history = []
    
    print(f"Starting training on {device} for {args.epochs} epochs...")
    for epoch in range(1, args.epochs + 1):
        # 1. Training Phase
        model.train()
        epoch_train_losses = []
        
        for batch_idx, (images, targets) in enumerate(train_loader):
            images = [img.to(device) for img in images]
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
            
            # Forward pass: model returns a dict of losses in training mode
            loss_dict = model(images, targets)
            losses = sum(loss for loss in loss_dict.values())
            
            optimizer.zero_grad()
            losses.backward()
            optimizer.step()
            
            epoch_train_losses.append(losses.item())
            
        train_loss = np.mean(epoch_train_losses)
        train_loss_history.append(train_loss)
        
        # 2. Validation Phase (compute validation loss by running model in train mode with no grad)
        model.train() # Must remain in train mode to output losses
        epoch_val_losses = []
        with torch.no_grad():
            for images, targets in val_loader:
                images = [img.to(device) for img in images]
                targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
                
                loss_dict = model(images, targets)
                losses = sum(loss for loss in loss_dict.values())
                epoch_val_losses.append(losses.item())
                
        val_loss = np.mean(epoch_val_losses)
        val_loss_history.append(val_loss)
        
        # Step the scheduler
        lr_scheduler.step()
        
        # Update progress reporting
        progress = 10.0 + (epoch / args.epochs) * 80.0 # scale 10% to 90%
        msg = f"Epoch {epoch}/{args.epochs} - Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}"
        print(msg)
        update_status(
            status="training",
            epoch=epoch,
            total_epochs=args.epochs,
            train_loss=train_loss_history,
            val_loss=val_loss_history,
            progress=progress,
            message=msg
        )
        
    # Save the PyTorch weights
    os.makedirs("models", exist_ok=True)
    pytorch_model_path = os.path.join("models", "model.pt")
    torch.save(model.state_dict(), pytorch_model_path)
    
    msg = f"Training completed successfully! Saved PyTorch model to {pytorch_model_path}"
    print(msg)
    update_status(
        status="completed",
        epoch=args.epochs,
        total_epochs=args.epochs,
        train_loss=train_loss_history,
        val_loss=val_loss_history,
        progress=100.0,
        message=msg
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train SSDLite MobileNetV3 on Fruit & Vegetable Dataset")
    parser.add_argument("--dataset-dir", type=str, default="dataset", help="Path to dataset directory")
    parser.add_argument("--epochs", type=int, default=15, help="Number of training epochs")
    parser.add_argument("--batch-size", type=int, default=8, help="Batch size for training")
    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
    
    args = parser.parse_args()
    
    try:
        # Initial status write
        update_status("idle", message="Starting training script...")
        train(args)
    except Exception as e:
        err_msg = f"Error during training: {str(e)}"
        print(err_msg, file=sys.stderr)
        update_status("failed", message=err_msg)
        sys.exit(1)