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)