Update app.py

app.py

@@ -217,12 +217,11 @@ model = YOLO(str(model_yaml))
 TRAINING_STATUS = "Starting Training... (Check Logs for details)"
 print("Starting Training...")
 
-# Define callback to update status
+# Define callbacks to update status
 def on_train_epoch_end(trainer):
     global TRAINING_STATUS
     current_epoch = trainer.epoch + 1
     total_epochs = trainer.epochs
-    # Try to get mAP if available, otherwise just epoch
     try:
         metrics = trainer.metrics
         map50 = metrics.get("metrics/mAP50(B)", 0)
@@ -230,7 +229,18 @@ def on_train_epoch_end(trainer):
     except:
         TRAINING_STATUS = f"Training in progress...<br>Epoch: {current_epoch}/{total_epochs}"
 
+def on_train_batch_end(trainer):
+    global TRAINING_STATUS
+    # Update status every 5 batches to avoid over-refreshing
+    if hasattr(trainer, 'pbar') and trainer.pbar and trainer.pbar.n % 5 == 0:
+        current_epoch = trainer.epoch + 1
+        total_epochs = trainer.epochs
+        current_batch = trainer.pbar.n
+        total_batch = trainer.pbar.total or "?"
+        TRAINING_STATUS = f"<b>Training Running...</b><br>Epoch: {current_epoch}/{total_epochs}<br>Batch: {current_batch}/{total_batch}<br><i>(CPU training is slow, please wait)</i>"
+
 model.add_callback("on_train_epoch_end", on_train_epoch_end)
+model.add_callback("on_train_batch_end", on_train_batch_end)
 
 # Train 200 epochs, CPU only
 results = model.train(