delete me

train.py

@@ -12,6 +12,18 @@ 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=""):
@@ -95,8 +107,25 @@ def train(args):
         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')
+    # 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
@@ -107,7 +136,8 @@ def train(args):
     train_loss_history = []
     val_loss_history = []
     
-    print(f"Starting training on {device} for {args.epochs} epochs...")
+    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()