Upload 5 files

app.py

@@ -4,7 +4,6 @@ import os
 import shutil
 import subprocess
 import sys
-import threading
 from collections import Counter
 from pathlib import Path
 
@@ -329,14 +328,18 @@ def build_coco():
     if not classes:
         raise RuntimeError("No classes configured.")
     items = [x for x in data["images"] if x.get("annotations")]
-    if len(items) < 2:
-        raise RuntimeError("Annotate at least 2 images before training.")
-
-    # Deterministic split; upload a reasonably shuffled dataset.
-    split = max(1, int(len(items) * 0.8))
-    if split == len(items):
-        split -= 1
-    train_items, val_items = items[:split], items[split:]
+    if not items:
+        raise RuntimeError("Annotate at least 1 image before training.")
+
+    # With only one annotated image, use it for both training and validation so
+    # the first training run is possible. With 2+ images, use an 80/20 split.
+    if len(items) == 1:
+        train_items, val_items = items, items
+    else:
+        split = max(1, int(len(items) * 0.8))
+        if split >= len(items):
+            split = len(items) - 1
+        train_items, val_items = items[:split], items[split:]
     category_id = {name: i + 1 for i, name in enumerate(classes)}
 
     def make_coco(selected):
@@ -397,7 +400,8 @@ def run_training(epochs, batch_size, learning_rate):
         ]
         result = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True)
         if result.returncode != 0:
-            raise RuntimeError((result.stderr or result.stdout)[-8000:])
+            details = result.stderr.strip() or result.stdout.strip() or f"training process exited with code {result.returncode}"
+            raise RuntimeError(details[-12000:])
         _model = None
         _training = {"running": False, "message": "training complete", "error": None}
     except Exception as e:
@@ -405,14 +409,21 @@ def run_training(epochs, batch_size, learning_rate):
 
 
 def start_training(epochs, batch_size, learning_rate):
+    """Start training from the Gradio event itself.
+
+    Calling the @spaces.GPU function directly is important on Hugging Face
+    ZeroGPU: starting it from a normal Python background thread can bypass the
+    GPU allocation context, making the button appear to do nothing.
+    """
     if _training["running"]:
-        return "Training is already running."
-    threading.Thread(
-        target=run_training,
-        args=(epochs, batch_size, learning_rate),
-        daemon=True,
-    ).start()
-    return "Training started in the background. Use Refresh Training Status."
+        return json.dumps(_training, indent=2)
+    # Basic validation before requesting GPU time.
+    annotated = sum(bool(x.get("annotations")) for x in load_dataset()["images"])
+    if annotated < 1:
+        _training = {"running": False, "message": "training failed", "error": "Annotate at least 1 image before training."}
+        return json.dumps(_training, indent=2)
+    run_training(epochs, batch_size, learning_rate)
+    return json.dumps(_training, indent=2)
 
 
 def training_status():

training/train.py

@@ -93,9 +93,9 @@ def main():
     val=COCODetectionDataset(Path(a.val_dir)/"images",Path(a.val_dir)/"annotations.json",proc)
 
     if len(train)==0 or len(val)==0:
-        raise ValueError("Both train and validation splits must contain annotated images.")
+        raise ValueError("Training and validation datasets must contain at least one image.")
     if len(train.category_id_to_label)!=len(classes) or len(val.category_id_to_label)!=len(classes):
-        raise ValueError("COCO categories do not match classes.txt.")
+        raise ValueError("COCO categories do not match classes.txt. Rebuild the dataset after saving the classes.")
 
     model=RTDetrForObjectDetection.from_pretrained(
         BASE_MODEL,num_labels=len(classes),id2label=id2label,label2id=label2id,