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 # Lower our own process priority on Windows so the FastAPI web server # stays responsive in the browser while CPU training is running. if sys.platform == "win32": try: import ctypes BELOW_NORMAL_PRIORITY_CLASS = 0x00004000 handle = ctypes.windll.kernel32.GetCurrentProcess() ctypes.windll.kernel32.SetPriorityClass(handle, BELOW_NORMAL_PRIORITY_CLASS) print("[train.py] CPU priority set to BELOW_NORMAL to keep web server responsive.", flush=True) except Exception as _prio_err: print(f"[train.py] Could not set process priority: {_prio_err}", flush=True) 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 best available device: # 1. DirectML (AMD / Intel GPU on Windows via torch-directml) # 2. CUDA (NVIDIA GPU) # 3. CPU fallback try: import torch_directml device = torch_directml.device() print(f"Using DirectML (AMD GPU): {device}") except ImportError: if torch.cuda.is_available(): device = torch.device('cuda') print(f"Using CUDA (NVIDIA GPU): {torch.cuda.get_device_name(0)}") else: device = torch.device('cpu') print("No GPU found, using CPU.") # Automatically cap batch size on CPU to avoid freezing the web server if args.batch_size > 4: print(f"[CPU mode] Reducing batch size from {args.batch_size} to 4 to keep web server responsive.") args.batch_size = 4 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 = [] device_label = str(device) print(f"Starting training on {device_label} 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)