Update app.py

app.py

@@ -1,5 +1,5 @@
 # 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
+import os, sys, subprocess, shutil, stat, yaml, gradio as gr, re, random, logging, requests, json, base64, time, pathlib, tempfile, textwrap
 from urllib.parse import urlparse
 from glob import glob
 from threading import Thread
@@ -74,9 +74,32 @@ except Exception:
     logging.exception("Bootstrap failed, UI will still load so you can see errors")
 
 # === model choices (restricted to Supervisely RT-DETRv2) ======================
-MODEL_CHOICES = [("rtdetrv2_s", "Small (default)"), ("rtdetrv2_l", "Large"), ("rtdetrv2_x", "X-Large")]
+# Exact mapping to configs and reference COCO checkpoints you provided
+MODEL_CHOICES = [
+    ("rtdetrv2_s",   "S (r18vd, 120e) — default"),
+    ("rtdetrv2_m",   "M (r34vd, 120e)"),
+    ("rtdetrv2_msp", "M* (r50vd_m, 7x)"),
+    ("rtdetrv2_l",   "L (r50vd, 6x)"),
+    ("rtdetrv2_x",   "X (r101vd, 6x)"),
+]
 DEFAULT_MODEL_KEY = "rtdetrv2_s"
 
+CONFIG_PATHS = {
+    "rtdetrv2_s":   "rtdetrv2_pytorch/configs/rtdetrv2/rtdetrv2_r18vd_120e_coco.yml",
+    "rtdetrv2_m":   "rtdetrv2_pytorch/configs/rtdetrv2/rtdetrv2_r34vd_120e_coco.yml",
+    "rtdetrv2_msp": "rtdetrv2_pytorch/configs/rtdetrv2/rtdetrv2_r50vd_m_7x_coco.yml",
+    "rtdetrv2_l":   "rtdetrv2_pytorch/configs/rtdetrv2/rtdetrv2_r50vd_6x_coco.yml",
+    "rtdetrv2_x":   "rtdetrv2_pytorch/configs/rtdetrv2/rtdetrv2_r101vd_6x_coco.yml",
+}
+
+CKPT_URLS = {
+    "rtdetrv2_s":   "https://github.com/lyuwenyu/storage/releases/download/v0.2/rtdetrv2_r18vd_120e_coco_rerun_48.1.pth",
+    "rtdetrv2_m":   "https://github.com/lyuwenyu/storage/releases/download/v0.1/rtdetrv2_r34vd_120e_coco_ema.pth",
+    "rtdetrv2_msp": "https://github.com/lyuwenyu/storage/releases/download/v0.1/rtdetrv2_r50vd_m_7x_coco_ema.pth",
+    "rtdetrv2_l":   "https://github.com/lyuwenyu/storage/releases/download/v0.1/rtdetrv2_r50vd_6x_coco_ema.pth",
+    "rtdetrv2_x":   "https://github.com/lyuwenyu/storage/releases/download/v0.1/rtdetrv2_r101vd_6x_coco_from_paddle.pth",
+}
+
 # === utilities ================================================================
 def handle_remove_readonly(func, path, exc_info):
     try:
@@ -372,15 +395,12 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
 
 # === entrypoint + config detection/generation =================================
 def find_training_script(repo_root):
-    # Hard-prefer the canonical path widely used in the repo/issues
     canonical = os.path.join(repo_root, "rtdetrv2_pytorch", "tools", "train.py")
     if os.path.exists(canonical):
         return canonical
-
     candidates = []
     for pat in ["**/tools/train.py", "**/train.py", "**/tools/train_net.py"]:
         candidates.extend(glob(os.path.join(repo_root, pat), recursive=True))
-    # Prefer anything inside rtdetrv2_pytorch, then shorter paths
     def _score(p):
         pl = p.replace("\\", "/").lower()
         return (0 if "rtdetrv2_pytorch" in pl else 1, len(p))
@@ -388,40 +408,13 @@ def find_training_script(repo_root):
     return candidates[0] if candidates else None
 
 def find_model_config_template(model_key):
-    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)
-
-    def score(p):
-        pl = p.lower()
-        s = 0
-        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
+    rel = CONFIG_PATHS.get(model_key)
+    if not rel:
+        return None
+    path = os.path.join(REPO_DIR, rel)
+    return path if os.path.exists(path) else None
 
 def _set_first_existing_key(d: dict, keys: list, value, fallback_key: str | None = None):
-    """
-    If any key from `keys` exists in dict `d`, set the first one found to `value`.
-    Otherwise, if `fallback_key` is given, create it with `value`.
-    Returns the key that was set, or None.
-    """
     for k in keys:
         if k in d:
             d[k] = value
@@ -431,8 +424,74 @@ def _set_first_existing_key(d: dict, keys: list, value, fallback_key: str | None
         return fallback_key
     return None
 
+def _set_first_existing_key_deep(cfg: dict, keys: list, value):
+    """
+    Try to set one of `keys` at top-level, under 'model', or under 'solver'.
+    """
+    for scope in [cfg, cfg.get("model", {}), cfg.get("solver", {})]:
+        if isinstance(scope, dict):
+            for k in keys:
+                if k in scope:
+                    scope[k] = value
+                    return True
+    # If nowhere found, set on model
+    if "model" not in cfg or not isinstance(cfg["model"], dict):
+        cfg["model"] = {}
+    cfg["model"][keys[0]] = value
+    return True
+
+def _install_supervisely_logger_shim():
+    """
+    Creates a minimal shim so `from supervisely.nn.training import train_logger` works.
+    """
+    base = pathlib.Path(tempfile.gettempdir()) / "sly_shim" / "supervisely" / "nn"
+    base.mkdir(parents=True, exist_ok=True)
+    for p in [base.parent.parent, base.parent, base]:
+        (p / "__init__.py").write_text("")
+    (base / "training.py").write_text(textwrap.dedent("""
+        # minimal shim for backward-compat with older Supervisely examples
+        class _TrainLogger:
+            def __init__(self): pass
+            def reset(self): pass
+            def log_metrics(self, metrics: dict, step: int | None = None): pass
+            def log_artifacts(self, *a, **k): pass
+            def log_image(self, *a, **k): pass
+        train_logger = _TrainLogger()
+    """))
+    return str(base.parent.parent.parent)  # .../sly_shim
+
+def _ensure_checkpoint(model_key: str, out_dir: str) -> str | None:
+    """
+    Download the reference COCO checkpoint for the selected model if not present.
+    Returns local path (or None if not available).
+    """
+    url = CKPT_URLS.get(model_key)
+    if not url:
+        return None
+    os.makedirs(out_dir, exist_ok=True)
+    fname = os.path.join(out_dir, os.path.basename(url))
+    if os.path.exists(fname) and os.path.getsize(fname) > 0:
+        return fname
+    logging.info(f"Downloading pretrained checkpoint for {model_key} from {url}")
+    try:
+        with requests.get(url, stream=True, timeout=60) as r:
+            r.raise_for_status()
+            with open(fname, "wb") as f:
+                for chunk in r.iter_content(chunk_size=1024 * 1024):
+                    if chunk:
+                        f.write(chunk)
+        return fname
+    except Exception as e:
+        logging.warning(f"Could not fetch checkpoint: {e}")
+        try:
+            if os.path.exists(fname):
+                os.remove(fname)
+        except Exception:
+            pass
+        return None
+
 def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
-                      epochs, batch, imgsz, lr, optimizer):
+                      epochs, batch, imgsz, lr, optimizer, pretrained_path: str | None):
     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.")
 
@@ -450,7 +509,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         "out_dir":    os.path.abspath(os.path.join("runs", "train", run_name)),
     }
 
-    # dataset block: set an existing alias if present, otherwise add a common key
+    # dataset block
     for root_key in ["dataset", "data"]:
         if root_key in cfg and isinstance(cfg[root_key], dict):
             ds = cfg[root_key]
@@ -525,7 +584,6 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
             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)
@@ -540,6 +598,14 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
     else:
         cfg["output_dir"] = paths["out_dir"]
 
+    # Set pretrained weights if available; try common keys at top/model/solver
+    if pretrained_path:
+        _set_first_existing_key_deep(
+            cfg,
+            keys=["pretrain", "pretrained", "weight", "weights", "pretrained_path"],
+            value=os.path.abspath(pretrained_path),
+        )
+
     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")
@@ -652,7 +718,7 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
 
     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).")
+        raise gr.Error("Could not find a matching RT-DETRv2 config in the repo (S/M/M*/L/X).")
 
     data_yaml = os.path.join(dataset_path, "data.yaml")
     with open(data_yaml, "r") as f:
@@ -660,6 +726,12 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     class_names = [str(x) for x in dy.get("names", [])]
     make_coco_annotations(dataset_path, class_names)
 
+    out_dir = os.path.abspath(os.path.join("runs", "train", run_name))
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Download matching COCO checkpoint for warm-start
+    pretrained_path = _ensure_checkpoint(model_key, out_dir)
+
     cfg_path = patch_base_config(
         base_cfg_path=base_cfg,
         merged_dir=dataset_path,
@@ -670,11 +742,9 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
         imgsz=imgsz,
         lr=lr,
         optimizer=opt,
+        pretrained_path=pretrained_path,
     )
 
-    out_dir = os.path.abspath(os.path.join("runs", "train", run_name))
-    os.makedirs(out_dir, exist_ok=True)
-
     cmd = [sys.executable, train_script, "-c", os.path.abspath(cfg_path)]
     logging.info(f"Training command: {' '.join(cmd)}")
 
@@ -685,6 +755,9 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
             env["PYTHONPATH"] = os.pathsep.join(filter(None, [
                 PY_IMPL_DIR, REPO_DIR, env.get("PYTHONPATH", "")
             ]))
+            # put our shim at the very front so the import always resolves
+            shim_root = _install_supervisely_logger_shim()
+            env["PYTHONPATH"] = os.pathsep.join([shim_root, env["PYTHONPATH"]])
             env.setdefault("WANDB_DISABLED", "true")
             proc = subprocess.Popen(cmd, cwd=os.path.dirname(train_script),
                                     stdout=subprocess.PIPE, stderr=subprocess.STDOUT,