Update app.py

app.py

@@ -1,4 +1,12 @@
+# app.py
+# Rolo: RT-DETRv2-only Training Dashboard (Supervisely ecosystem)
+# - No Ultralytics import or usage
+# - Auto-installs deps in HF Spaces
+# - Only supports models that ship with https://github.com/supervisely-ecosystem/RT-DETRv2
+
 import os
+import sys
+import subprocess
 import shutil
 import stat
 import yaml
@@ -11,33 +19,88 @@ import logging
 import requests
 import json
 from PIL import Image
+import torch
 import pandas as pd
-import matplotlib
-matplotlib.use("Agg")  # headless (HF Spaces)
 import matplotlib.pyplot as plt
 from threading import Thread
 from queue import Queue
-from huggingface_hub import HfApi, HfFolder
-import base64
-import subprocess
-import sys
+from glob import glob
 import time
-import glob
+import base64
 
 # --- Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# --- RT-DETRv2 backend defaults (Supervisely ecosystem) ---
-RTDETRV2_REPO_URL = "https://github.com/supervisely-ecosystem/RT-DETRv2"
-DEFAULT_REPO_DIR = os.path.join("third_party", "rtdetrv2")
+REPO_URL = "https://github.com/supervisely-ecosystem/RT-DETRv2"  # :contentReference[oaicite:1]{index=1}
+REPO_DIR = os.path.join(os.getcwd(), "third_party", "RT-DETRv2")
+PY_IMPL_DIR = os.path.join(REPO_DIR, "rtdetrv2_pytorch")  # contains the pytorch impl (models, training)
+WEIGHTS_DIR = os.path.join(PY_IMPL_DIR, "weights")
+
+# ------------------------------
+# Environment bootstrap (HF Spaces)
+# ------------------------------
+
+COMMON_REQUIREMENTS = [
+    "gradio>=4.36.1",
+    "roboflow>=1.1.28",
+    "pandas>=2.0.0",
+    "matplotlib>=3.7.0",
+    "pyyaml>=6.0.1",
+    "Pillow>=10.0.0",
+    "requests>=2.31.0",
+    "huggingface_hub>=0.22.0",
+]
+
+def pip_install(args):
+    logging.info(f"pip install {' '.join(args)}")
+    subprocess.check_call([sys.executable, "-m", "pip", "install"] + args)
+
+def ensure_repo_and_requirements():
+    os.makedirs(os.path.dirname(REPO_DIR), exist_ok=True)
+    if not os.path.exists(REPO_DIR):
+        logging.info(f"Cloning RT-DETRv2 repo to {REPO_DIR} ...")
+        subprocess.check_call(["git", "clone", "--depth", "1", REPO_URL, REPO_DIR])
+    else:
+        logging.info("RT-DETRv2 repo already present, pulling latest...")
+        try:
+            subprocess.check_call(["git", "-C", REPO_DIR, "pull", "--ff-only"])
+        except Exception:
+            logging.warning("Could not pull latest; continuing with current checkout.")
 
-RTDETRV2_MODELS = [
-    "rtdetrv2-l-640",  # labels only; match your config via the command template
-    "rtdetrv2-x-640"
+    # Install common libs
+    pip_install(COMMON_REQUIREMENTS)
+
+    # Install rtdetrv2_pytorch requirements if present
+    req_file = os.path.join(PY_IMPL_DIR, "requirements.txt")
+    if os.path.exists(req_file):
+        pip_install(["-r", req_file])
+    else:
+        logging.info("No rtdetrv2_pytorch/requirements.txt found; relying on common reqs.")
+
+# Do the bootstrap once at import time (HF Spaces-friendly).
+try:
+    ensure_repo_and_requirements()
+except Exception as e:
+    logging.exception("Bootstrap failed")
+    # Still allow UI to load so user can see the error
+    pass
+
+# ------------------------------
+# Model options (strictly from RT-DETRv2 repo)
+# ------------------------------
+# We expose only the canonical small/large/xlarge variants that ship with the repo.
+# If the repo adds/removes variants, you can read from weights dir dynamically.
+MODEL_CHOICES = [
+    ("rtdetrv2_s", "Small (default)"),
+    ("rtdetrv2_l", "Large"),
+    ("rtdetrv2_x", "X-Large")
 ]
-DEFAULT_MODEL = RTDETRV2_MODELS[0]
+DEFAULT_MODEL_KEY = "rtdetrv2_s"  # Small as default
+
+# ------------------------------
+# Utilities
+# ------------------------------
 
-# --- Utilities ---
 def handle_remove_readonly(func, path, exc_info):
     try:
         os.chmod(path, stat.S_IWRITE)
@@ -51,15 +114,11 @@ _ROBO_URL_RX = re.compile(
         (?:https?://)?(?:universe|app|www)?\.?roboflow\.com/
         (?P<ws>[A-Za-z0-9\-_]+)/
         (?P<proj>[A-Za-z0-9\-_]+)/?
-        (?:
-            (?:dataset/[^/]+/)?
-            (?:v?(?P<ver>\d+))?
-        )?
+        (?:(?:dataset/[^/]+/)?(?:v?(?P<ver>\d+))?)?
       |
         (?P<ws2>[A-Za-z0-9\-_]+)/(?P<proj2>[A-Za-z0-9\-_]+)(?:/(?:v)?(?P<ver2>\d+))?
     )$
-    """,
-    re.VERBOSE | re.IGNORECASE
+    """, re.VERBOSE | re.IGNORECASE
 )
 
 def parse_roboflow_url(s: str):
@@ -111,20 +170,25 @@ def _extract_class_names(data_yaml):
     names = data_yaml.get('names', None)
     if isinstance(names, dict):
         def _k(x):
-            try: return int(x)
-            except Exception: return str(x)
-        ordered = sorted(names.keys(), key=_k)
-        names_list = [names[k] for k in ordered]
+            try:
+                return int(x)
+            except Exception:
+                return str(x)
+        ordered_keys = sorted(names.keys(), key=_k)
+        names_list = [names[k] for k in ordered_keys]
     elif isinstance(names, list):
         names_list = names
     else:
         nc = data_yaml.get('nc', 0)
-        try: nc = int(nc)
-        except Exception: nc = 0
+        try:
+            nc = int(nc)
+        except Exception:
+            nc = 0
         names_list = [f"class_{i}" for i in range(nc)]
     return [str(x) for x in names_list]
 
 def download_dataset(api_key, workspace, project, version):
+    """Download a Roboflow dataset in YOLOv8 format (labels are compatible with our merger)."""
     try:
         rf = Roboflow(api_key=api_key)
         proj = rf.workspace(workspace).project(project)
@@ -143,16 +207,14 @@ def download_dataset(api_key, workspace, project, version):
         if len(class_names) != nc:
             logging.warning(f"[{project}-v{version}] names length ({len(class_names)}) != nc ({nc}); using normalized names.")
 
-        splits = [s for s in ['train', 'valid', 'test']
-                  if os.path.exists(os.path.join(dataset.location, s))]
-
+        splits = [s for s in ['train', 'valid', 'test'] if os.path.exists(os.path.join(dataset.location, s))]
         return dataset.location, class_names, splits, f"{project}-v{version}"
     except Exception as e:
         logging.error(f"Failed to download {workspace}/{project}/v{version}: {e}")
         return None, [], [], None
 
 def label_path_for(img_path: str) -> str:
-    split_dir = os.path.dirname(os.path.dirname(img_path))  # .../split
+    split_dir = os.path.dirname(os.path.dirname(img_path))
     base = os.path.splitext(os.path.basename(img_path))[0] + '.txt'
     return os.path.join(split_dir, 'labels', base)
 
@@ -161,6 +223,7 @@ def gather_class_counts(dataset_info, class_mapping):
         return {}
     final_names = set(v for v in class_mapping.values() if v is not None)
     counts = {name: 0 for name in final_names}
+
     for loc, names, splits, _ in dataset_info:
         id_to_name = {idx: class_mapping.get(n, None) for idx, n in enumerate(names)}
         for split in splits:
@@ -179,7 +242,7 @@ def gather_class_counts(dataset_info, class_mapping):
                         try:
                             cls_id = int(parts[0])
                             mapped = id_to_name.get(cls_id, None)
-                            if mapped in final_names:
+                            if mapped:
                                 found.add(mapped)
                         except Exception:
                             continue
@@ -197,7 +260,7 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
         os.makedirs(os.path.join(merged_dir, split, 'images'), exist_ok=True)
         os.makedirs(os.path.join(merged_dir, split, 'labels'), exist_ok=True)
 
-    active_classes = sorted(set([cls for cls, limit in class_limits.items() if limit > 0]))
+    active_classes = sorted({cls for cls, limit in class_limits.items() if limit > 0})
     final_class_map = {name: i for i, name in enumerate(active_classes)}
 
     all_images = []
@@ -216,6 +279,7 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
     current_counts = {cls: 0 for cls in active_classes}
     loc_to_names = {info[0]: info[1] for info in dataset_info}
 
+    # progress.tqdm is available on Gradio Progress objects
     for img_path, split, source_loc in progress.tqdm(all_images, desc="Analyzing images"):
         lbl_path = label_path_for(img_path)
         if not os.path.exists(lbl_path):
@@ -239,6 +303,7 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
 
         if not image_classes:
             continue
+
         if any(current_counts[c] >= class_limits[c] for c in image_classes):
             continue
 
@@ -288,128 +353,116 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
 
     return f"Dataset finalized with {len(selected_images)} images.", os.path.abspath(merged_dir)
 
-# --- Repo + deps helpers (auto-install for HF Spaces) ---
-
-def run_pip_install(args, desc="pip install"):
-    logging.info(f"{desc}: {args}")
-    cmd = [sys.executable, "-m", "pip", "install"] + args
-    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
-    logging.info(proc.stdout)
-    if proc.returncode != 0:
-        raise RuntimeError(f"{desc} failed with code {proc.returncode}")
+# ------------------------------
+# Training integration (RT-DETRv2 repo)
+# ------------------------------
 
-def ensure_repo(repo_dir: str, repo_url: str = RTDETRV2_REPO_URL):
-    if os.path.isdir(repo_dir) and os.path.isdir(os.path.join(repo_dir, ".git")):
-        return
-    os.makedirs(os.path.dirname(repo_dir), exist_ok=True)
-    logging.info(f"Cloning RT-DETRv2 repo into {repo_dir} ...")
-    subprocess.run(["git", "clone", "--depth", "1", repo_url, repo_dir], check=True)
-
-def ensure_python_deps(repo_dir: str):
+def detect_training_entrypoint():
     """
-    Auto-install dependencies (idempotent).
-    - Tries to install pinned basics that are often needed.
-    - If repo has requirements*.txt, install them.
-    - Creates a .deps_installed marker to skip on next run.
+    We try a couple of common patterns inside the Supervisely repo:
+      1) rtdetrv2_pytorch/train.py
+      2) tools/train.py
+    Returns (python_file, style) where style hints how to build args.
     """
-    marker = os.path.join(repo_dir, ".deps_installed")
-    if os.path.exists(marker):
-        logging.info("Dependencies already installed; skipping.")
-        return
-
-    # 1) Common essentials for vision training environments on HF Spaces
-    basics = [
-        "numpy<2",               # safer with many libs
-        "pillow",
-        "tqdm",
-        "pyyaml",
-        "matplotlib",
-        "pandas",
-        "scipy",
-        "opencv-python-headless",
-        "packaging",
-        "requests",
-        "pycocotools-windows; platform_system=='Windows'",
-        "pycocotools; platform_system!='Windows'",
-    ]
-    try:
-        run_pip_install(basics, desc="Installing common basics")
-    except Exception as e:
-        logging.warning(f"Basic installs had issues: {e}")
-
-    # 2) Repo requirements
-    req_files = []
-    for name in ["requirements.txt", "requirements-dev.txt", "requirements.in"]:
-        p = os.path.join(repo_dir, name)
-        if os.path.isfile(p):
-            req_files.append(p)
-
-    for rf in req_files:
-        try:
-            run_pip_install(["-r", rf], desc=f"Installing repo requirements from {rf}")
-        except Exception as e:
-            logging.warning(f"Installing {rf} failed: {e}")
-
-    # 3) Optional: torch if not present (CPU-only by default on Spaces)
-    try:
-        import torch  # noqa: F401
-    except Exception:
-        # Try a CPU-friendly torch; change version/cuda wheels if needed
-        try:
-            run_pip_install(["torch", "torchvision", "torchaudio", "--index-url", "https://download.pytorch.org/whl/cpu"], desc="Installing PyTorch (CPU)")
-        except Exception as e:
-            logging.warning(f"PyTorch installation failed/skipped: {e}")
-
-    # Mark done
-    with open(marker, "w") as f:
-        f.write("ok\n")
-
-def make_train_command(template: str, data_yaml: str, epochs: int, batch: int, imgsz: int,
-                       lr: float, optimizer: str, run_name: str, output_dir: str) -> str:
-    return template.format(
-        data_yaml=data_yaml,
-        epochs=int(epochs),
-        batch=int(batch),
-        imgsz=int(imgsz),
-        lr=float(lr),
-        optimizer=str(optimizer),
-        run_name=str(run_name),
-        output_dir=output_dir
-    )
-
-_METRIC_PATTERNS = [
-    (re.compile(r"mAP@0\.5[:/]?0\.95[^0-9]*([0-9]*\.?[0-9]+)"), "mAP50_95"),
-    (re.compile(r"mAP50[^0-9]*([0-9]*\.?[0-9]+)"), "mAP50"),
-    (re.compile(r"\bval[_/ ]?loss[^0-9\-]*([0-9]*\.?[0-9]+)"), "val_loss"),
-    (re.compile(r"\btrain[_/ ]?loss[^0-9\-]*([0-9]*\.?[0-9]+)"), "train_loss"),
-    (re.compile(r"\bepoch[^0-9]*([0-9]+)"), "epoch"),
-]
-
-def parse_metrics_from_line(line: str):
-    result = {}
-    for pat, key in _METRIC_PATTERNS:
-        m = pat.search(line)
-        if m:
-            val = m.group(1)
-            try:
-                result[key] = int(val) if key == "epoch" else float(val)
-            except Exception:
-                pass
-    return result
-
-def guess_final_weights(output_dir: str):
+    cand1 = os.path.join(PY_IMPL_DIR, "train.py")
+    cand2 = os.path.join(REPO_DIR, "tools", "train.py")
+    if os.path.exists(cand1):
+        return cand1, "pytorch_train"
+    if os.path.exists(cand2):
+        return cand2, "tools_train"
+    # Fallback: just try main.py if present
+    cand3 = os.path.join(REPO_DIR, "src", "main.py")
+    if os.path.exists(cand3):
+        return cand3, "app_main"
+    return None, None
+
+def build_command(entrypoint, style, dataset_path, model_key, run_name, epochs, batch, imgsz, lr, optimizer):
+    """
+    Build a best-guess command for the detected style.
+    Users never have to edit CLI; we do it for them.
+    We keep args conservative and standard (data, epochs, batch, img size).
+    """
+    data_yaml = os.path.join(dataset_path, "data.yaml")
+    out_dir = os.path.join("runs", "train", str(run_name))
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Some repos expect weight/model name; we pass model_key (e.g., rtdetrv2_s) and let their script resolve it.
+    # Learning rate / optimizer flags may differ; include only when style suggests they're supported.
+    if style == "pytorch_train":
+        # Hypothetical common args for a train.py inside rtdetrv2_pytorch
+        cmd = [
+            sys.executable, entrypoint,
+            "--data", data_yaml,
+            "--model", model_key,
+            "--epochs", str(int(epochs)),
+            "--batch", str(int(batch)),
+            "--imgsz", str(int(imgsz)),
+            "--project", os.path.abspath(out_dir)
+        ]
+        if lr is not None:
+            cmd += ["--lr", str(float(lr))]
+        if optimizer:
+            cmd += ["--optimizer", str(optimizer)]
+        return cmd, out_dir
+
+    if style == "tools_train":
+        # Alternate style (tools/train.py). We keep flags generic.
+        cmd = [
+            sys.executable, entrypoint,
+            "--data", data_yaml,
+            "--model", model_key,
+            "--epochs", str(int(epochs)),
+            "--batch-size", str(int(batch)),
+            "--imgsz", str(int(imgsz)),
+            "--project", os.path.abspath(out_dir),
+            "--name", "exp"
+        ]
+        if lr is not None:
+            cmd += ["--lr0", str(float(lr))]
+        if optimizer:
+            cmd += ["--optimizer", str(optimizer)]
+        return cmd, out_dir
+
+    if style == "app_main":
+        # If app_main exists, it may require an options file; we still try a generic mapping.
+        cmd = [
+            sys.executable, entrypoint,
+            "--data", data_yaml,
+            "--model", model_key,
+            "--epochs", str(int(epochs)),
+            "--batch", str(int(batch)),
+            "--imgsz", str(int(imgsz)),
+            "--output", os.path.abspath(out_dir)
+        ]
+        if lr is not None:
+            cmd += ["--lr", str(float(lr))]
+        if optimizer:
+            cmd += ["--optimizer", str(optimizer)]
+        return cmd, out_dir
+
+    raise gr.Error("Could not locate a training script inside RT-DETRv2 repo. Please check the repo layout.")
+
+def find_best_checkpoint(out_dir):
+    # Look for common patterns
     patterns = [
-        os.path.join(output_dir, "**", "best.*"),
-        os.path.join(output_dir, "**", "best_model.*"),
-        os.path.join(output_dir, "**", "checkpoint_best.*"),
+        os.path.join(out_dir, "**", "best*.pt"),
+        os.path.join(out_dir, "**", "best*.pth"),
+        os.path.join(out_dir, "**", "model_best*.pt"),
+        os.path.join(out_dir, "**", "model_best*.pth"),
     ]
     for p in patterns:
-        hits = glob.glob(p, recursive=True)
-        if hits:
-            return hits[0]
-    return None
+        files = sorted(glob(p, recursive=True))
+        if files:
+            return files[0]
+    # Fall back to latest .pt/.pth
+    any_ckpt = sorted(glob(os.path.join(out_dir, "**", "*.pt"), recursive=True) +
+                      glob(os.path.join(out_dir, "**", "*.pth"), recursive=True))
+    return any_ckpt[-1] if any_ckpt else None
+
+# ------------------------------
+# Gradio Handlers
+# ------------------------------
 
-# --- Gradio handlers ---
 def load_datasets_handler(api_key, url_file, progress=gr.Progress()):
     api_key = api_key or os.getenv("ROBOFLOW_API_KEY", "")
     if not api_key:
@@ -420,8 +473,7 @@ def load_datasets_handler(api_key, url_file, progress=gr.Progress()):
     with open(url_file.name, 'r', encoding='utf-8', errors='ignore') as f:
         urls = [line.strip() for line in f if line.strip()]
 
-    dataset_info = []
-    failures = []
+    dataset_info, failures = [], []
     for i, raw in enumerate(urls):
         progress((i + 1) / max(1, len(urls)), desc=f"Parsing {i+1}/{len(urls)}")
         ws, proj, ver = parse_roboflow_url(raw)
@@ -444,32 +496,33 @@ def load_datasets_handler(api_key, url_file, progress=gr.Progress()):
         msg = "No datasets were loaded successfully.\n" + "\n".join([f"- {u}: {why}" for u, why in failures[:10]])
         raise gr.Error(msg)
 
+    # Make sure names are strings before sorting to avoid mixed-type comparison
     all_names = sorted({str(n) for _, names, _, _ in dataset_info for n in names})
     class_map = {name: name for name in all_names}
+
     initial_counts = gather_class_counts(dataset_info, class_map)
-    df_data = [[name, name, initial_counts.get(name, 0), False] for name in all_names]
+    df = pd.DataFrame([[name, name, initial_counts.get(name, 0), False] for name in all_names],
+                      columns=["Original Name", "Rename To", "Max Images", "Remove"])
     status_text = "Datasets loaded successfully."
     if failures:
         status_text += f" ({len(dataset_info)} OK, {len(failures)} failed; see console logs)."
 
-    return status_text, dataset_info, gr.update(
-        value=pd.DataFrame(df_data, columns=["Original Name", "Rename To", "Max Images", "Remove"])
-    )
+    # Return the DataFrame value directly (works across Gradio versions)
+    return status_text, dataset_info, df
 
 def update_class_counts_handler(class_df, dataset_info):
     if class_df is None or not dataset_info:
         return None
+
     class_df = pd.DataFrame(class_df)
     mapping = {}
     for _, row in class_df.iterrows():
         orig = row["Original Name"]
-        if bool(row["Remove"]):
-            mapping[orig] = None
-        else:
-            mapping[orig] = row["Rename To"]
+        mapping[orig] = None if bool(row["Remove"]) else row["Rename To"]
 
     final_names = sorted(set(v for v in mapping.values() if v))
     counts = {k: 0 for k in final_names}
+
     for loc, names, splits, _ in dataset_info:
         id_to_final = {idx: mapping.get(n, None) for idx, n in enumerate(names)}
         for split in splits:
@@ -495,8 +548,7 @@ def update_class_counts_handler(class_df, dataset_info):
                 for m in found:
                     counts[m] += 1
 
-    summary_df = pd.DataFrame(list(counts.items()), columns=["Final Class Name", "Est. Total Images"])
-    return summary_df
+    return pd.DataFrame(list(counts.items()), columns=["Final Class Name", "Est. Total Images"])
 
 def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
     if not dataset_info:
@@ -505,8 +557,7 @@ def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
         raise gr.Error("Class data is missing.")
 
     class_df = pd.DataFrame(class_df)
-    class_mapping = {}
-    class_limits = {}
+    class_mapping, class_limits = {}, {}
     for _, row in class_df.iterrows():
         orig = row["Original Name"]
         if bool(row["Remove"]):
@@ -518,89 +569,100 @@ def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
     status, path = finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress)
     return status, path
 
-def training_handler_rtdetrv2(dataset_path, repo_dir, model_choice, run_name, epochs, batch, imgsz, lr, opt,
-                              cmd_template, progress=gr.Progress()):
+def training_handler(dataset_path, model_choice_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.")
 
-    # Clone + deps (idempotent)
-    try:
-        ensure_repo(repo_dir)
-        ensure_python_deps(repo_dir)
-    except subprocess.CalledProcessError as e:
-        raise gr.Error(f"Failed to clone repo: {e}")
-    except Exception as e:
-        raise gr.Error(f"Dependency setup failed: {e}")
+    # Verify repo entrypoint
+    entrypoint, style = detect_training_entrypoint()
+    if not entrypoint:
+        raise gr.Error("RT-DETRv2 training script not found in the repo. Please check repo contents.")
+
+    # Build and run command (users never touch CLI)
+    cmd, out_dir = build_command(
+        entrypoint=entrypoint,
+        style=style,
+        dataset_path=dataset_path,
+        model_key=model_choice_key,
+        run_name=run_name,
+        epochs=epochs,
+        batch=batch,
+        imgsz=imgsz,
+        lr=lr,
+        optimizer=opt
+    )
+    logging.info(f"Training command: {' '.join(cmd)}")
 
-    # Output dir
-    output_dir = os.path.join("runs", "train", str(run_name))
-    os.makedirs(output_dir, exist_ok=True)
+    # Live-run in a thread and stream logs
+    q = Queue()
 
-    data_yaml = os.path.join(dataset_path, "data.yaml")
-    if not os.path.isfile(data_yaml):
-        raise gr.Error(f"'data.yaml' was not found in: {dataset_path}")
-
-    # Build command from template
-    cmd = make_train_command(
-        template=cmd_template,
-        data_yaml=data_yaml,
-        epochs=int(epochs),
-        batch=int(batch),
-        imgsz=int(imgsz),
-        lr=float(lr),
-        optimizer=str(opt),
-        run_name=str(run_name),
-        output_dir=output_dir
-    )
+    def run_train():
+        try:
+            env = os.environ.copy()
+            env["PYTHONPATH"] = REPO_DIR + os.pathsep + env.get("PYTHONPATH", "")
+            proc = subprocess.Popen(cmd, cwd=REPO_DIR, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, text=True, env=env)
+            for line in proc.stdout:
+                q.put(line.rstrip())
+            proc.wait()
+            q.put(f"__EXITCODE__:{proc.returncode}")
+        except Exception as e:
+            q.put(f"__ERROR__:{e}")
+
+    Thread(target=run_train, daemon=True).start()
+
+    log_lines = []
+    last_epoch = 0
+    total_epochs = int(epochs)
+    while True:
+        line = q.get()
+        if line.startswith("__EXITCODE__"):
+            code = int(line.split(":", 1)[1])
+            if code != 0:
+                raise gr.Error(f"Training process exited with code {code}. Check logs above.")
+            break
+        if line.startswith("__ERROR__"):
+            raise gr.Error(f"Training failed: {line.split(':',1)[1]}")
+
+        log_lines.append(line)
+        # try to parse "Epoch X/Y" style hints for progress
+        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
 
-    logging.info(f"Running training command in {repo_dir}: {cmd}")
-    proc = subprocess.Popen(
-        cmd, cwd=repo_dir, shell=True,
-        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
-        bufsize=1, universal_newlines=True, env={**os.environ}
-    )
+        frac = min(max(last_epoch / max(1, total_epochs), 0.0), 1.0)
+        progress(frac, desc=f"Epoch {last_epoch}/{total_epochs}")
 
-    history = {k: [] for k in ['epoch', 'train_loss', 'val_loss', 'mAP50', 'mAP50_95']}
-    for line in iter(proc.stdout.readline, ''):
-        line = line.rstrip()
-        progress(0.0, desc=line[-120:])
-        metrics = parse_metrics_from_line(line)
-        if metrics:
-            for k, v in metrics.items():
-                history[k].append(v)
-
-            # plot loss
-            fig_loss = plt.figure()
-            ax_loss = fig_loss.add_subplot(111)
-            ax_loss.plot(history['epoch'], history['train_loss'], "o-", label='Train Loss')
-            ax_loss.plot(history['epoch'], history['val_loss'], "o-", label='Val Loss')
-            ax_loss.legend(); ax_loss.set_title("Loss")
-
-            # plot mAP
-            fig_map = plt.figure()
-            ax_map = fig_map.add_subplot(111)
-            ax_map.plot(history['epoch'], history['mAP50'], "o-", label='mAP@0.5')
-            ax_map.plot(history['epoch'], history['mAP50_95'], "o-", label='mAP@0.5:0.95')
-            ax_map.legend(); ax_map.set_title("mAP")
-
-            yield line[-200:], fig_loss, fig_map, None
-
-    proc.stdout.close()
-    ret = proc.wait()
-    if ret != 0:
-        raise gr.Error(f"Training process exited with code {ret}. Check console/logs for details.")
-
-    final_ckpt = guess_final_weights(output_dir)
-    if final_ckpt and os.path.isfile(final_ckpt):
-        yield "Training complete!", None, None, gr.File.update(value=final_ckpt, visible=True)
-    else:
-        yield ("Training finished. Could not auto-detect a 'best' checkpoint; "
-               "please check the output directory."), None, None, gr.update(visible=False)
+        # Light-weight plots (we won't have metrics dicts; just show empty placeholders so UI doesn't break)
+        fig_loss = plt.figure()
+        ax_loss = fig_loss.add_subplot(111)
+        ax_loss.set_title("Loss (see logs)")
+        fig_map = plt.figure()
+        ax_map = fig_map.add_subplot(111)
+        ax_map.set_title("mAP (see logs)")
+
+        yield "\n".join(log_lines[-30:]), fig_loss, fig_map, None
+
+    # Look for the best checkpoint
+    ckpt = find_best_checkpoint(out_dir)
+    if not ckpt or not os.path.exists(ckpt):
+        # try give user any artifact
+        alt = find_best_checkpoint("runs")
+        if not alt or not os.path.exists(alt):
+            raise gr.Error("Training finished, but checkpoint file was not found. See logs for details.")
+        ckpt = alt
+
+    yield "Training complete!", None, None, gr.File.update(value=ckpt, visible=True)
 
 def upload_handler(model_file, hf_token, hf_repo, gh_token, gh_repo, progress=gr.Progress()):
     if not model_file:
         raise gr.Error("No trained model file available to upload. Train a model first.")
 
+    from huggingface_hub import HfApi, HfFolder
+
     hf_status = "Skipped Hugging Face (credentials not provided)."
     if hf_token and hf_repo:
         progress(0, desc="Uploading to Hugging Face...")
@@ -624,6 +686,7 @@ def upload_handler(model_file, hf_token, hf_repo, gh_token, gh_repo, progress=gr
         try:
             if '/' not in gh_repo:
                 raise ValueError("GitHub repo must be in the form 'username/repo'.")
+
             username, repo_name = gh_repo.split('/')
             api_url = f"https://api.github.com/repos/{username}/{repo_name}/contents/{os.path.basename(model_file.name)}"
             headers = {"Authorization": f"token {gh_token}"}
@@ -635,9 +698,11 @@ def upload_handler(model_file, hf_token, hf_repo, gh_token, gh_repo, progress=gr
             sha = get_resp.json().get('sha') if get_resp.ok else None
 
             data = {"message": "Upload trained model from Rolo app", "content": content}
-            if sha: data["sha"] = sha
+            if sha:
+                data["sha"] = sha
 
             put_resp = requests.put(api_url, headers=headers, json=data, timeout=60)
+
             if put_resp.ok:
                 gh_status = f"Success! Model at: {put_resp.json()['content']['html_url']}"
             else:
@@ -649,24 +714,26 @@ def upload_handler(model_file, hf_token, hf_repo, gh_token, gh_repo, progress=gr
     progress(1)
     return hf_status, gh_status
 
-# --- Gradio UI ---
+# ------------------------------
+# Gradio UI
+# ------------------------------
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
-    gr.Markdown("# Rolo: RT-DETRv2 Training Dashboard (Auto-setup for Hugging Face)")
+    gr.Markdown("# Rolo: RT-DETRv2 Training (Supervisely ecosystem only)")
 
     dataset_info_state = gr.State([])
     final_dataset_path_state = gr.State(None)
 
     with gr.Tabs():
         with gr.TabItem("1. Prepare Datasets"):
-            gr.Markdown("Upload a `.txt` with Roboflow URLs or `workspace/project[/vN]` lines.")
+            gr.Markdown("### Load Roboflow Datasets\nProvide your Roboflow API key and upload a `.txt` file containing one Roboflow dataset URL or `workspace/project[/vN]` per line.")
             with gr.Row():
-                rf_api_key = gr.Textbox(label="Roboflow API Key (or set ROBOFLOW_API_KEY)", type="password", scale=2)
+                rf_api_key = gr.Textbox(label="Roboflow API Key (or set ROBOFLOW_API_KEY env)", type="password", scale=2)
                 rf_url_file = gr.File(label="Upload Roboflow URLs (.txt)", file_types=[".txt"], scale=1)
             load_btn = gr.Button("Load Datasets", variant="primary")
             dataset_status = gr.Textbox(label="Status", interactive=False)
 
         with gr.TabItem("2. Manage & Merge"):
-            gr.Markdown("Rename classes, set image limits, or remove them. Preview, then finalize.")
+            gr.Markdown("### Configure Classes and Finalize Dataset\nRename classes to merge them, set image limits, or remove them. Click **Update Counts** to preview, then **Finalize** to create the dataset.")
             with gr.Row():
                 class_df = gr.DataFrame(
                     headers=["Original Name", "Rename To", "Max Images", "Remove"],
@@ -684,43 +751,33 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
             finalize_status = gr.Textbox(label="Status", interactive=False)
 
         with gr.TabItem("3. Configure & Train"):
-            gr.Markdown("Set hyperparameters and the training command template.")
+            gr.Markdown("### Set Hyperparameters and Train the RT-DETRv2 Model")
             with gr.Row():
                 with gr.Column(scale=1):
-                    model_choice_dd = gr.Dropdown(
-                        label="Model Choice (label only; use your config in the template)",
-                        choices=RTDETRV2_MODELS, value=DEFAULT_MODEL
+                    model_file_dd = gr.Dropdown(
+                        label="Model (only RT-DETRv2 from Supervisely)",
+                        choices=[k for k, _ in MODEL_CHOICES],
+                        value=DEFAULT_MODEL_KEY
+                    )
+                    model_hints = gr.Markdown(
+                        "Choices: " +
+                        ", ".join([f"`{k}` ({label})" for k, label in MODEL_CHOICES])
                     )
                     run_name_tb = gr.Textbox(label="Run Name", value="rtdetrv2_run_1")
                     epochs_sl = gr.Slider(1, 500, 100, step=1, label="Epochs")
                     batch_sl = gr.Slider(1, 64, 16, step=1, label="Batch Size")
                     imgsz_num = gr.Number(label="Image Size", value=640)
                     lr_num = gr.Number(label="Learning Rate", value=0.001)
-                    opt_dd = gr.Dropdown(["Adam", "AdamW", "SGD"], value="AdamW", label="Optimizer")
-                    repo_dir_tb = gr.Textbox(label="RT-DETRv2 repo directory", value=DEFAULT_REPO_DIR)
-                    cmd_template_tb = gr.Textbox(
-                        label="Train command template",
-                        value=(
-                            "python tools/train.py "
-                            "--data {data_yaml} "
-                            "--epochs {epochs} "
-                            "--batch {batch} "
-                            "--imgsz {imgsz} "
-                            "--lr {lr} "
-                            "--optimizer {optimizer} "
-                            "--output {output_dir}"
-                        ),
-                        lines=4
-                    )
+                    opt_dd = gr.Dropdown(["Adam", "AdamW", "SGD"], value="Adam", label="Optimizer")
                     train_btn = gr.Button("Start Training", variant="primary")
                 with gr.Column(scale=2):
-                    train_status = gr.Textbox(label="Live Status / Logs", interactive=False)
-                    loss_plot = gr.Plot(label="Loss Curves")
-                    map_plot = gr.Plot(label="mAP Curves")
-                    final_model_file = gr.File(label="Download Trained Model (best.*)", interactive=False, visible=False)
+                    train_status = gr.Textbox(label="Live Logs (tail)", interactive=False, lines=12)
+                    loss_plot = gr.Plot(label="Loss")
+                    map_plot = gr.Plot(label="mAP")
+                    final_model_file = gr.File(label="Download Trained Model", interactive=False, visible=False)
 
         with gr.TabItem("4. Upload Model"):
-            gr.Markdown("Upload your best checkpoint to Hugging Face or GitHub.")
+            gr.Markdown("### Upload Your Trained Model")
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("#### Hugging Face")
@@ -735,6 +792,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
                 hf_status = gr.Textbox(label="Hugging Face Status", interactive=False)
                 gh_status = gr.Textbox(label="GitHub Status", interactive=False)
 
+    # Wire UI handlers
     load_btn.click(
         fn=load_datasets_handler,
         inputs=[rf_api_key, rf_url_file],
@@ -751,19 +809,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
         outputs=[finalize_status, final_dataset_path_state]
     )
     train_btn.click(
-        fn=training_handler_rtdetrv2,
-        inputs=[
-            final_dataset_path_state,  # dataset_path
-            repo_dir_tb,               # repo_dir (auto clone + pip install)
-            model_choice_dd,           # model_choice (label only)
-            run_name_tb,
-            epochs_sl,
-            batch_sl,
-            imgsz_num,
-            lr_num,
-            opt_dd,
-            cmd_template_tb
-        ],
+        fn=training_handler,
+        inputs=[final_dataset_path_state, model_file_dd, run_name_tb, epochs_sl, batch_sl, imgsz_num, lr_num, opt_dd],
         outputs=[train_status, loss_plot, map_plot, final_model_file]
     )
     upload_btn.click(
@@ -773,6 +820,6 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
     )
 
 if __name__ == "__main__":
-    # Hugging Face Spaces: set server name/port via env if needed.
-    # Example: app.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)), debug=True)
+    # Silence Ultralytics warnings if present in the env (we don't use Ultralytics at all)
+    os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")
     app.launch(debug=True)