initial commit

train.py

@@ -0,0 +1,199 @@
+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)