Update app.py

app.py

@@ -145,7 +145,7 @@ def parse_roboflow_url(s: str):
             version = None
             if len(p) >= 3:
                 v = p[2]
-                if v.lower().startswith('v') and v[1:].isdigit():  # <-- FIXED
+                if v.lower().startswith('v') and v[1:].isdigit():
                     version = int(v[1:])
                 elif v.isdigit():
                     version = int(v)
@@ -459,7 +459,7 @@ def _install_supervisely_logger_shim():
     """))
     return str(root)
 
-# ---- [UPDATED] robust sitecustomize shim with lazy import hook --------------------
+# ---- [!! CORRECTED !!] robust sitecustomize shim with lazy import hook --------------------
 def _install_workspace_shim_v3(dest_dir: str, module_default: str = "rtdetrv2_pytorch.src"):
     """
     sitecustomize shim that:
@@ -566,23 +566,6 @@ if not _try_patch_now():
         f.write(code)
     return sc_path
 
-# ---- Deprecated: on-disk workspace patch (no-op now) -------------------------
-def _patch_workspace_create(repo_root: str, module_default: str = "rtdetrv2_pytorch.src") -> str | None:
-    """
-    Deprecated: we no longer edit third-party files on disk.
-    The shim in sitecustomize.py handles cfg/_pymodule safely.
-    """
-    return None
-
-def _unpatch_workspace_create(repo_root: str):
-    ws_path = os.path.join(repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py")
-    bak_path = ws_path + ".bak"
-    if os.path.exists(bak_path):
-        try:
-            shutil.copy2(bak_path, ws_path)
-        except Exception:
-            pass
-
 def _ensure_checkpoint(model_key: str, out_dir: str) -> str | None:
     url = CKPT_URLS.get(model_key)
     if not url:
@@ -731,7 +714,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
                 "shuffle": bool(default_shuffle),
                 "num_workers": 2,
                 "drop_last": bool(dl_key == "train_dataloader"),
-                "collate_fn": {"type": "BatchImageCollateFunction"},  # correct spelling
+                "collate_fn": {"type": "BatchImageCollateFunction"},
                 "total_batch_size": int(batch),
             }
             cfg[dl_key] = block
@@ -751,7 +734,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         block.setdefault("shuffle", bool(default_shuffle))
         block.setdefault("drop_last", bool(dl_key == "train_dataloader"))
 
-        # ---- FORCE-FIX collate name even if it existed already
+        # ---- FORCE-FIX collate name typo even if it existed already
         cf = block.get("collate_fn", {})
         if isinstance(cf, dict):
             t = str(cf.get("type", ""))
@@ -805,7 +788,6 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         _maybe_set_model_field(cfg, "pretrain", p)
         _maybe_set_model_field(cfg, "pretrained", p)
 
-    # Defensive: if after keeping includes we still don't have a model block, add a stub
     if not cfg.get("model"):
         cfg["model"] = {"type": "RTDETR", "num_classes": int(class_count)}
 
@@ -961,21 +943,13 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
         try:
             train_cwd = os.path.dirname(train_script)
 
-            # --- NEW: create a temp dir for sitecustomize and put it FIRST on PYTHONPATH
             shim_dir = tempfile.mkdtemp(prefix="rtdetr_site_")
             _install_workspace_shim_v3(shim_dir, module_default="rtdetrv2_pytorch.src")
 
             env = os.environ.copy()
 
-            # Supervisely logger shim (can be later in path)
             sly_shim_root = _install_supervisely_logger_shim()
 
-            # Build PYTHONPATH — order matters!
-            #  1) shim_dir (so sitecustomize auto-imports)
-            #  2) train_cwd (belt & suspenders; makes local imports easy)
-            #  3) PY_IMPL_DIR + REPO_DIR (RT-DETRv2 code)
-            #  4) sly_shim_root (optional)
-            #  5) existing PYTHONPATH
             env["PYTHONPATH"] = os.pathsep.join(filter(None, [
                 shim_dir,
                 train_cwd,
@@ -987,8 +961,7 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
 
             env.setdefault("WANDB_DISABLED", "true")
             env.setdefault("RTDETR_PYMODULE", "rtdetrv2_pytorch.src")
-            env.setdefault("PYTHONUNBUFFERED", "1")  # nicer real-time logs
-            # Optional tiny guard: pick a single visible GPU if available
+            env.setdefault("PYTHONUNBUFFERED", "1")
             if torch.cuda.is_available():
                 env.setdefault("CUDA_VISIBLE_DEVICES", "0")
 
@@ -1147,7 +1120,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
             gr.Markdown("Pick RT-DETRv2 model, set hyper-params, press Start.")
             with gr.Row():
                 with gr.Column(scale=1):
-                    model_dd = gr.Dropdown(choices=[(label, k) for k, label in MODEL_CHOICES], value=DEFAULT_MODEL_KEY,
+                    # [UI IMPROVEMENT] Using (label, value) format for a better user experience
+                    model_dd = gr.Dropdown(choices=[(label, value) for value, label in MODEL_CHOICES],
+                                           value=DEFAULT_MODEL_KEY,
                                            label="Model (RT-DETRv2)")
                     run_name_tb = gr.Textbox(label="Run Name", value="rtdetrv2_run_1")
                     epochs_sl = gr.Slider(1, 500, 100, step=1, label="Epochs")