Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py — Rolo: RT-DETRv2-only (Supervisely) trainer with auto COCO conversion & config
+# app.py — Rolo: RT-DETRv2-only (Supervisely) trainer with auto COCO conversion & safe config patching
 import os, sys, subprocess, shutil, stat, yaml, gradio as gr, re, random, logging, requests, json, base64, time
 from urllib.parse import urlparse
 from glob import glob
@@ -336,71 +336,141 @@ def find_training_script(repo_root):
 
 def find_model_config_template(model_key):
     """
-    Find a base config YAML in the repo that matches the chosen model key.
-    We look under any configs directory for a yaml containing 'rtdetrv2' and the model key.
+    Choose a native RT-DETRv2 config YAML from the Supervisely repo.
+
+    Heuristics:
+      - rtdetrv2_s -> r18 (Small)
+      - rtdetrv2_l -> r50 (Large)
+      - rtdetrv2_x -> r101 (X-Large)
+    Prefer files under rtdetrv2_pytorch/**/config(s) and with 'coco' in name.
     """
+    want_tokens = {
+        "rtdetrv2_s": ["rtdetrv2", "r18", "coco"],
+        "rtdetrv2_l": ["rtdetrv2", "r50", "coco"],
+        "rtdetrv2_x": ["rtdetrv2", "r101", "coco"],
+    }.get(model_key, ["rtdetrv2", "r18", "coco"])
+
     yamls = glob(os.path.join(REPO_DIR, "**", "*.yml"), recursive=True) + \
             glob(os.path.join(REPO_DIR, "**", "*.yaml"), recursive=True)
-    # prioritize files with both rtdetrv2 and the exact key in the name
+
     def score(p):
-        n = os.path.basename(p).lower()
+        pl = p.lower()
         s = 0
-        if "rtdetrv2" in n: s += 2
-        if model_key in n: s += 3
-        if "coco" in n: s += 1
+        if "/rtdetrv2_pytorch/" in pl: s += 4
+        if "/config" in pl: s += 3
+        for token in want_tokens:
+            if token in os.path.basename(pl): s += 3
+            if token in pl: s += 2
+        if "coco" in pl: s += 1
         return -s, len(p)
+
     yamls.sort(key=score)
     return yamls[0] if yamls else None
 
-def write_custom_config(base_cfg_path, merged_dir, class_count, model_key, run_name, epochs, batch, imgsz, lr, optimizer):
+def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
+                      epochs, batch, imgsz, lr, optimizer):
     """
-    Generate a small override config that points to our COCO jsons and sets key hyperparams.
-    This YAML gets merged by the repo's config system if it supports '_base_' includes;
-    otherwise, it still provides reasonable keys many RT-DETRv2 forks accept.
+    Load the chosen repo config and patch only the keys that already exist.
+    This avoids schema mismatches between forks.
     """
-    ann_dir = os.path.join(merged_dir, "annotations")
-    cfg_out_dir = os.path.join("generated_configs")
-    os.makedirs(cfg_out_dir, exist_ok=True)
-    out_path = os.path.join(cfg_out_dir, f"{run_name}.yaml")
+    if not base_cfg_path or not os.path.exists(base_cfg_path):
+        raise gr.Error("Could not locate a model config inside the RT-DETRv2 repo.")
+
+    with open(base_cfg_path, "r") as f:
+        cfg = yaml.safe_load(f)
 
-    # Try a broadly compatible structure (kept simple on purpose)
-    override = {
-        "_base_": os.path.relpath(base_cfg_path, start=cfg_out_dir) if base_cfg_path else None,
-        "model": {"name": model_key, "num_classes": int(class_count)},
-        "input_size": int(imgsz),
-        "max_epoch": int(epochs),
-        "solver": {
-            "base_lr": float(lr),
-            "optimizer": str(optimizer).lower(),  # "adam", "adamw", "sgd"
-            "batch_size": int(batch),
-        },
-        "dataset": {
-            "train": {
-                "name": "coco",
-                "ann_file": os.path.abspath(os.path.join(ann_dir, "instances_train.json")),
-                "img_prefix": os.path.abspath(os.path.join(merged_dir, "train", "images")),
-            },
-            "val": {
-                "name": "coco",
-                "ann_file": os.path.abspath(os.path.join(ann_dir, "instances_val.json")),
-                "img_prefix": os.path.abspath(os.path.join(merged_dir, "valid", "images")),
-            },
-            "test": {
-                "name": "coco",
-                "ann_file": os.path.abspath(os.path.join(ann_dir, "instances_test.json")),
-                "img_prefix": os.path.abspath(os.path.join(merged_dir, "test", "images")),
-            },
-        },
-        "output_dir": os.path.abspath(os.path.join("runs", "train", run_name)),
-        # some forks use these dataloader keys:
-        "train_dataloader": {"batch_size": int(batch)},
-        "val_dataloader": {"batch_size": int(batch)},
+    ann_dir = os.path.join(merged_dir, "annotations")
+    paths = {
+        "train_json": os.path.abspath(os.path.join(ann_dir, "instances_train.json")),
+        "val_json":   os.path.abspath(os.path.join(ann_dir, "instances_val.json")),
+        "test_json":  os.path.abspath(os.path.join(ann_dir, "instances_test.json")),
+        "train_img":  os.path.abspath(os.path.join(merged_dir, "train", "images")),
+        "val_img":    os.path.abspath(os.path.join(merged_dir, "valid", "images")),  # Roboflow uses 'valid'
+        "test_img":   os.path.abspath(os.path.join(merged_dir, "test", "images")),
+        "out_dir":    os.path.abspath(os.path.join("runs", "train", run_name)),
     }
-    # drop None values cleanly
-    if override["_base_"] is None:
-        del override["_base_"]
 
-    with open(out_path, "w") as f: yaml.safe_dump(override, f, sort_keys=False)
+    # --- dataset block --------------------------------------------------------
+    for root_key in ["dataset", "data"]:
+        if root_key in cfg and isinstance(cfg[root_key], dict):
+            ds = cfg[root_key]
+            for split, jf, ip in [
+                ("train", "train_json", "train_img"),
+                ("val",   "val_json",   "val_img"),
+                ("test",  "test_json",  "test_img"),
+            ]:
+                if split in ds and isinstance(ds[split], dict):
+                    ds[split]["name"] = ds[split].get("name", "coco")
+                    # Common key variants across forks:
+                    for k in ["ann_file", "ann_path", "annotation", "annotations"]:
+                        if k in ds[split] or k in ["ann_file", "ann_path"]:
+                            ds[split][k] = paths[jf]
+                            break
+                    for k in ["img_prefix", "img_dir", "image_root", "data_root"]:
+                        if k in ds[split] or k in ["img_prefix", "img_dir"]:
+                            ds[split][k] = paths[ip]
+                            break
+
+    # --- num_classes ----------------------------------------------------------
+    def set_num_classes(node, n):
+        if not isinstance(node, dict): return False
+        if "num_classes" in node:
+            node["num_classes"] = int(n); return True
+        for k, v in node.items():
+            if isinstance(v, dict) and set_num_classes(v, n): return True
+        return False
+
+    if "model" in cfg and isinstance(cfg["model"], dict):
+        if not set_num_classes(cfg["model"], class_count):
+            cfg["model"]["num_classes"] = int(class_count)
+    else:
+        cfg["model"] = {"num_classes": int(class_count)}
+
+    # --- epochs / image size --------------------------------------------------
+    updated_epoch = False
+    for key in ["max_epoch", "epochs", "num_epochs"]:
+        if key in cfg:
+            cfg[key] = int(epochs); updated_epoch = True; break
+    if "solver" in cfg and isinstance(cfg["solver"], dict):
+        for key in ["max_epoch", "epochs", "num_epochs"]:
+            if key in cfg["solver"]:
+                cfg["solver"][key] = int(epochs); updated_epoch = True; break
+    if not updated_epoch:
+        cfg["max_epoch"] = int(epochs)
+
+    for key in ["input_size", "img_size", "imgsz"]:
+        if key in cfg: cfg[key] = int(imgsz)
+    if "input_size" not in cfg: cfg["input_size"] = int(imgsz)
+
+    # --- learning rate / optimizer / batch -----------------------------------
+    if "solver" not in cfg or not isinstance(cfg["solver"], dict):
+        cfg["solver"] = {}
+    sol = cfg["solver"]
+    for key in ["base_lr", "lr", "learning_rate"]:
+        if key in sol:
+            sol[key] = float(lr); break
+    else:
+        sol["base_lr"] = float(lr)
+
+    sol["optimizer"] = str(optimizer).lower()
+
+    if "train_dataloader" in cfg and isinstance(cfg["train_dataloader"], dict):
+        cfg["train_dataloader"]["batch_size"] = int(batch)
+    else:
+        sol["batch_size"] = int(batch)
+
+    # --- output dir -----------------------------------------------------------
+    if "output_dir" in cfg:
+        cfg["output_dir"] = paths["out_dir"]
+    elif "solver" in cfg:
+        sol["output_dir"] = paths["out_dir"]
+    else:
+        cfg["output_dir"] = paths["out_dir"]
+
+    # --- write patched config -------------------------------------------------
+    cfg_out_dir = os.path.join("generated_configs"); os.makedirs(cfg_out_dir, exist_ok=True)
+    out_path = os.path.join(cfg_out_dir, f"{run_name}.yaml")
+    with open(out_path, "w") as f: yaml.safe_dump(cfg, f, sort_keys=False)
     return out_path
 
 def find_best_checkpoint(out_dir):
@@ -500,26 +570,27 @@ def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
 def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr, opt, progress=gr.Progress()):
     if not dataset_path: raise gr.Error("Finalize a dataset in Tab 2 before training.")
 
-    # 1) find training script (nested-safe)
+    # 1) training script (nested-safe)
     train_script = find_training_script(REPO_DIR)
     if not train_script:
         raise gr.Error("RT-DETRv2 training script not found inside the repo (looked for **/tools/train.py).")
 
-    # 2) pick a model config template from repo (best effort)
+    # 2) base config = a real model template from the repo
     base_cfg = find_model_config_template(model_key)
+    if not base_cfg:
+        raise gr.Error("Could not find a matching RT-DETRv2 config in the repo (S/L/X).")
 
-    # 3) read class names from our merged data.yaml to set num_classes + produce COCO JSONs
+    # 3) read classes + ensure COCO JSONs up to date
     data_yaml = os.path.join(dataset_path, "data.yaml")
     with open(data_yaml, "r") as f: dy = yaml.safe_load(f)
     class_names = [str(x) for x in dy.get("names", [])]
-    ann_dir = make_coco_annotations(dataset_path, class_names)
+    make_coco_annotations(dataset_path, class_names)
 
-    # 4) write a small override config that points to our data and injects hyper-params
-    cfg_path = write_custom_config(
+    # 4) patch the base config safely (no custom schema assumptions)
+    cfg_path = patch_base_config(
         base_cfg_path=base_cfg,
         merged_dir=dataset_path,
         class_count=len(class_names),
-        model_key=model_key,
         run_name=run_name,
         epochs=epochs,
         batch=batch,
@@ -531,18 +602,20 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     out_dir = os.path.abspath(os.path.join("runs", "train", run_name))
     os.makedirs(out_dir, exist_ok=True)
 
-    # 5) build & run the command (single-GPU by default, no manual CLI edits)
+    # 5) build & run command (no extra flags that might not exist)
     cmd = [sys.executable, train_script, "-c", os.path.abspath(cfg_path)]
-    # many forks accept optional flags; pass safe ones if present
-    if "--use-amp" in open(train_script).read():  # cheap check
-        cmd += ["--use-amp"]
     logging.info(f"Training command: {' '.join(cmd)}")
 
     q = Queue()
     def run_train():
         try:
             env = os.environ.copy()
-            env["PYTHONPATH"] = REPO_DIR + os.pathsep + env.get("PYTHONPATH", "")
+            # Ensure both repo root and pytorch impl are on PYTHONPATH
+            env["PYTHONPATH"] = os.pathsep.join(filter(None, [
+                PY_IMPL_DIR, REPO_DIR, env.get("PYTHONPATH", "")
+            ]))
+            # Disable wandb in Spaces by default
+            env.setdefault("WANDB_DISABLED", "true")
             proc = subprocess.Popen(cmd, cwd=os.path.dirname(train_script),
                                     stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                     bufsize=1, text=True, env=env)
@@ -555,25 +628,26 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     Thread(target=run_train, daemon=True).start()
 
     log_tail, last_epoch, total_epochs = [], 0, int(epochs)
+    first_lines = []  # capture early errors for nicer message
     while True:
         line = q.get()
         if line.startswith("__EXITCODE__"):
             code = int(line.split(":",1)[1])
-            if code != 0: raise gr.Error(f"Training exited with code {code}. See logs above.")
+            if code != 0:
+                head = "\n".join(first_lines[:60])
+                raise gr.Error(f"Training exited with code {code}.\nLast output:\n{head or 'No logs captured.'}")
             break
         if line.startswith("__ERROR__"):
             raise gr.Error(f"Training failed: {line.split(':',1)[1]}")
 
-        log_tail.append(line)
-        log_tail = log_tail[-30:]
+        if len(first_lines) < 120: first_lines.append(line)
+        log_tail.append(line); log_tail = log_tail[-40:]
 
         m = re.search(r"[Ee]poch\s+(\d+)\s*/\s*(\d+)", line)
         if m:
             try:
-                last_epoch = int(m.group(1))
-                total_epochs = max(total_epochs, int(m.group(2)))
-            except Exception:
-                pass
+                last_epoch = int(m.group(1)); total_epochs = max(total_epochs, int(m.group(2)))
+            except Exception: pass
         progress(min(max(last_epoch / max(1,total_epochs),0.0),1.0), desc=f"Epoch {last_epoch}/{total_epochs}")
 
         fig1 = plt.figure(); plt.title("Loss (see logs)")