Update app.py

app.py

@@ -3,7 +3,6 @@ import shutil
 import stat
 import yaml
 import gradio as gr
-from ultralytics import YOLO  # Ultralytics RT-DETR runner
 from roboflow import Roboflow
 import re
 from urllib.parse import urlparse
@@ -12,34 +11,31 @@ import logging
 import requests
 import json
 from PIL import Image
-import torch
 import pandas as pd
 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
+import time
+import glob
 
 # --- Configuration ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Hardcode RT-DETR model configurations. All YOLO options are removed.
-RTDETR_MODELS = {
-    "detection": [
-        {
-            "filename": "rtdetr-l.pt",
-            "url": "https://github.com/ultralytics/assets/releases/download/v8.0.0/rtdetr-l.pt",
-            "description": "RT-DETR Large model (Default)"
-        },
-        {
-            "filename": "rtdetr-x.pt",
-            "url": "https://github.com/ultralytics/assets/releases/download/v8.0.0/rtdetr-x.pt",
-            "description": "RT-DETR Extra-Large model."
-        }
-    ]
-}
-DEFAULT_MODEL = "rtdetr-l.pt"
+# Defaults for RT-DETRv2 (Supervisely ecosystem) integration
+RTDETRV2_REPO_URL = "https://github.com/supervisely-ecosystem/RT-DETRv2"
+DEFAULT_REPO_DIR = os.path.join("third_party", "rtdetrv2")
 
+# You can still offer "model size" choices to hint the user which config to use,
+# but the actual command is controlled by the template.
+RTDETRV2_MODELS = [
+    "rtdetrv2-l-640",  # label only; adapt your command template to use real config/weights
+    "rtdetrv2-x-640"
+]
+DEFAULT_MODEL = RTDETRV2_MODELS[0]
 
 # ------------------------------
 # Utilities
@@ -53,19 +49,18 @@ def handle_remove_readonly(func, path, exc_info):
         pass
     func(path)
 
-
 _ROBO_URL_RX = re.compile(
     r"""
     ^(?:
-        (?:https?://)?(?:universe|app|www)?\.?roboflow\.com/    # Any roboflow host
-        (?P<ws>[A-Za-z0-9\-_]+)/                                # workspace
-        (?P<proj>[A-Za-z0-9\-_]+)/?                              # project
+        (?:https?://)?(?:universe|app|www)?\.?roboflow\.com/
+        (?P<ws>[A-Za-z0-9\-_]+)/
+        (?P<proj>[A-Za-z0-9\-_]+)/?
         (?:
-            (?:dataset/[^/]+/)?                                  # optional 'dataset/<fmt>/'
-            (?:v?(?P<ver>\d+))?                                  # optional version 'vN' or 'N'
+            (?:dataset/[^/]+/)?
+            (?:v?(?P<ver>\d+))?
         )?
       |
-        (?P<ws2>[A-Za-z0-9\-_]+)/(?P<proj2>[A-Za-z0-9\-_]+)(?:/(?:v)?(?P<ver2>\d+))?  # raw ws/proj[/vN]
+        (?P<ws2>[A-Za-z0-9\-_]+)/(?P<proj2>[A-Za-z0-9\-_]+)(?:/(?:v)?(?P<ver2>\d+))?
     )$
     """,
     re.VERBOSE | re.IGNORECASE
@@ -73,15 +68,14 @@ _ROBO_URL_RX = re.compile(
 
 def parse_roboflow_url(s: str):
     """
-    Accepts:
-      - https://universe.roboflow.com/<workspace>/<project>[/vN | /N]
-      - https://app.roboflow.com/<workspace>/<project>[/vN | /N]
-      - https://roboflow.com/<workspace>/<project>[/vN | /N]
-      - raw: <workspace>/<project>[/vN | /N]
+    Support:
+      - https://universe.roboflow.com/<workspace>/<project>[/vN]
+      - https://app.roboflow.com/<workspace>/<project>[/vN]
+      - https://roboflow.com/<workspace>/<project>[/vN]
+      - raw: <workspace>/<project>[/vN]
     Returns: (workspace, project, version_or_None)
     """
     s = s.strip()
-    # Fast path: try regex
     m = _ROBO_URL_RX.match(s)
     if m:
         ws = m.group('ws') or m.group('ws2')
@@ -89,14 +83,11 @@ def parse_roboflow_url(s: str):
         ver = m.group('ver') or m.group('ver2')
         return ws, proj, (int(ver) if ver else None)
 
-    # Fallback: parse like URL and split path
     parsed = urlparse(s)
     parts = [p for p in parsed.path.strip('/').split('/') if p]
     if len(parts) >= 2:
-        # Try to pull raw version from the 3rd part if it exists
         version = None
         if len(parts) >= 3:
-            # Accept 'vN' or 'N'
             vpart = parts[2]
             if vpart.lower().startswith('v') and vpart[1:].isdigit():
                 version = int(vpart[1:])
@@ -104,11 +95,9 @@ def parse_roboflow_url(s: str):
                 version = int(vpart)
         return parts[0], parts[1], version
 
-    # Fallback raw "ws/proj" without slashes in URL
     if '/' in s and 'roboflow' not in s:
         p = s.split('/')
         if len(p) >= 2:
-            # Accept trailing version if present
             version = None
             if len(p) >= 3:
                 v = p[2]
@@ -120,7 +109,6 @@ def parse_roboflow_url(s: str):
 
     return None, None, None
 
-
 def get_latest_version(api_key, workspace, project):
     """Gets the latest version number of a Roboflow project."""
     try:
@@ -132,15 +120,14 @@ def get_latest_version(api_key, workspace, project):
         logging.error(f"Could not get latest version for {workspace}/{project}: {e}")
         return None
 
-
-# --- NEW: normalize class names from data.yaml ---
+# --- Normalize class names from data.yaml ---
 def _extract_class_names(data_yaml):
     """
-    Return a list[str] of class names in index order.
-    Handles:
-      - list (possibly containing non-str types)
-      - dict with numeric keys (e.g., {0: 'cat', 1: 'dog'})
-      - fallback to ['class_0', ..., f'class_{nc-1}'] if names missing
+    Return list[str] of class names in index order.
+    Supports:
+      - list
+      - dict with numeric keys {0:'cat',1:'dog'}
+      - fallback to ['class_0', ...]
     """
     names = data_yaml.get('names', None)
 
@@ -150,8 +137,8 @@ def _extract_class_names(data_yaml):
                 return int(x)
             except Exception:
                 return str(x)
-        ordered_keys = sorted(names.keys(), key=_k)
-        names_list = [names[k] for k in ordered_keys]
+        ordered = sorted(names.keys(), key=_k)
+        names_list = [names[k] for k in ordered]
     elif isinstance(names, list):
         names_list = names
     else:
@@ -164,9 +151,8 @@ def _extract_class_names(data_yaml):
 
     return [str(x) for x in names_list]
 
-
 def download_dataset(api_key, workspace, project, version):
-    """Downloads a single dataset from Roboflow (yolov8 format works fine for RT-DETR)."""
+    """Download Roboflow dataset in 'yolov8' layout (works fine for RT-DETR variants)."""
     try:
         rf = Roboflow(api_key=api_key)
         proj = rf.workspace(workspace).project(project)
@@ -177,7 +163,6 @@ def download_dataset(api_key, workspace, project, version):
         with open(data_yaml_path, 'r') as f:
             data_yaml = yaml.safe_load(f)
 
-        # --- UPDATED: use normalized names and optional sanity log ---
         class_names = _extract_class_names(data_yaml)
         try:
             nc = int(data_yaml.get('nc', len(class_names)))
@@ -194,30 +179,25 @@ def download_dataset(api_key, workspace, project, version):
         logging.error(f"Failed to download {workspace}/{project}/v{version}: {e}")
         return None, [], [], None
 
-
 def label_path_for(img_path: str) -> str:
-    """Convert .../split/images/file.jpg -> .../split/labels/file.txt in a safe way."""
+    """Convert .../split/images/file.jpg -> .../split/labels/file.txt."""
     split_dir = os.path.dirname(os.path.dirname(img_path))  # .../split
     base = os.path.splitext(os.path.basename(img_path))[0] + '.txt'
     return os.path.join(split_dir, 'labels', base)
 
-
 def gather_class_counts(dataset_info, class_mapping):
     """
-    Count, per final class, how many images contain at least one instance of that class
-    (counted once per image). class_mapping maps original_name -> final_name.
+    Count per final class how many images contain that class at least once (counted once per image).
+    class_mapping: original_name -> final_name (or None if removed).
     """
     if not dataset_info:
         return {}
 
-    final_names = set(class_mapping.values())
+    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:
-        # Map from original idx -> mapped name (or None if removed later)
-        id_to_name = {}
-        for idx, n in enumerate(names):
-            id_to_name[idx] = class_mapping.get(n, None)
+        id_to_name = {idx: class_mapping.get(n, None) for idx, n in enumerate(names)}
 
         for split in splits:
             labels_dir = os.path.join(loc, split, 'labels')
@@ -244,9 +224,8 @@ def gather_class_counts(dataset_info, class_mapping):
 
     return counts
 
-
 def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=gr.Progress()):
-    """Core function to merge datasets based on user rules."""
+    """Merge datasets following mapping and per-class image limits."""
     merged_dir = 'rolo_merged_dataset'
     if os.path.exists(merged_dir):
         shutil.rmtree(merged_dir, onerror=handle_remove_readonly)
@@ -256,12 +235,10 @@ 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)
 
-    # Only classes with positive limits are active
-    active_classes = [cls for cls, limit in class_limits.items() if limit > 0]
-    active_classes = sorted(set(active_classes))
+    active_classes = sorted(set([cls for cls, limit in class_limits.items() if limit > 0]))
     final_class_map = {name: i for i, name in enumerate(active_classes)}
 
-    # Collect all candidate images
+    # Collect candidates
     all_images = []
     for loc, _, splits, _ in dataset_info:
         for split in splits:
@@ -276,8 +253,6 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
     progress(0.2, desc="Selecting images based on limits...")
     selected_images = []
     current_counts = {cls: 0 for cls in active_classes}
-
-    # Build a quick lookup: source_loc -> names list
     loc_to_names = {info[0]: info[1] for info in dataset_info}
 
     for img_path, split, source_loc in progress.tqdm(all_images, desc="Analyzing images"):
@@ -303,8 +278,6 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
 
         if not image_classes:
             continue
-
-        # Check limits
         if any(current_counts[c] >= class_limits[c] for c in image_classes):
             continue
 
@@ -319,7 +292,6 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
         out_lbl = os.path.join(merged_dir, split, 'labels', os.path.basename(lbl_path))
         shutil.copy(img_path, out_img)
 
-        # Determine source names by matching the parent dataset root
         source_loc = None
         for info in dataset_info:
             if img_path.startswith(info[0]):
@@ -355,6 +327,68 @@ 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)
 
+# ------------------------------
+# RT-DETRv2 backend helpers
+# ------------------------------
+
+def ensure_repo(repo_dir: str, repo_url: str = RTDETRV2_REPO_URL):
+    """Clone the repo into repo_dir if not present."""
+    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} ...")
+    cmd = ["git", "clone", "--depth", "1", repo_url, repo_dir]
+    subprocess.run(cmd, check=True)
+
+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 = [
+    # add more patterns if your repo prints differently
+    (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):
+    """
+    Try to locate a 'best' checkpoint in output_dir.
+    Supports .pt/.pth/.pdparams etc. Return first match or None.
+    """
+    patterns = [
+        os.path.join(output_dir, "**", "best.*"),
+        os.path.join(output_dir, "**", "best_model.*"),
+        os.path.join(output_dir, "**", "checkpoint_best.*"),
+    ]
+    for p in patterns:
+        hits = glob.glob(p, recursive=True)
+        if hits:
+            return hits[0]
+    return None
 
 # ------------------------------
 # Gradio UI Event Handlers
@@ -394,36 +428,29 @@ def load_datasets_handler(api_key, url_file, progress=gr.Progress()):
             failures.append((raw, f"DownloadError: {ws}/{proj}/v{ver}"))
 
     if not dataset_info:
-        # Show a compact failure report to the UI
         msg = "No datasets were loaded successfully.\n" + "\n".join([f"- {u}: {why}" for u, why in failures[:10]])
         raise gr.Error(msg)
 
-    # --- UPDATED: ensure all names are strings before sorting
+    # ensure names are strings before sorting
     all_names = sorted({str(n) for _, names, _, _ in dataset_info for n in names})
     class_map = {name: name for name in all_names}
 
-    # Initial preview uses "keep all" mapping
     initial_counts = gather_class_counts(dataset_info, class_map)
     df_data = [[name, name, initial_counts.get(name, 0), False] for name in all_names]
     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.DataFrame.update(
+    # FIX: gr.update(...) (not gr.DataFrame.update)
+    return status_text, dataset_info, gr.update(
         value=pd.DataFrame(df_data, columns=["Original Name", "Rename To", "Max Images", "Remove"])
     )
 
-
 def update_class_counts_handler(class_df, dataset_info):
-    """
-    Provides live feedback on class counts as the user edits the DataFrame.
-    We compute a mapping of original -> final (or None if removed), then count images
-    for each final name.
-    """
+    """Live preview of merged class counts given the current mapping/removals."""
     if class_df is None or not dataset_info:
         return None
 
-    # Build mapping original_name -> final_name or None if removed
     class_df = pd.DataFrame(class_df)
     mapping = {}
     for _, row in class_df.iterrows():
@@ -433,14 +460,11 @@ def update_class_counts_handler(class_df, dataset_info):
         else:
             mapping[orig] = row["Rename To"]
 
-    # Build final set
     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 = {}
-        for idx, n in enumerate(names):
-            id_to_final[idx] = mapping.get(n, None)
+        id_to_final = {idx: mapping.get(n, None) for idx, n in enumerate(names)}
 
         for split in splits:
             labels_dir = os.path.join(loc, split, 'labels')
@@ -468,15 +492,13 @@ def update_class_counts_handler(class_df, dataset_info):
     summary_df = pd.DataFrame(list(counts.items()), columns=["Final Class Name", "Est. Total Images"])
     return summary_df
 
-
 def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
-    """Handles the 'Finalize' button click."""
+    """Create the merged dataset directory with relabeled .txts and data.yaml."""
     if not dataset_info:
         raise gr.Error("Load datasets first in Tab 1.")
     if class_df is None:
         raise gr.Error("Class data is missing.")
 
-    # Mapping and limits
     class_df = pd.DataFrame(class_df)
     class_mapping = {}
     class_limits = {}
@@ -486,112 +508,109 @@ def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
             continue
         final_name = row["Rename To"]
         class_mapping[orig] = final_name
-        # Sum limits for final_name over any merged originals
         class_limits[final_name] = class_limits.get(final_name, 0) + int(row["Max Images"])
 
     status, path = finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress)
     return status, path
 
-
-def training_handler(dataset_path, model_filename, run_name, epochs, batch, imgsz, lr, opt, progress=gr.Progress()):
-    """Handles the training process with live feedback."""
+def training_handler_rtdetrv2(dataset_path, repo_dir, model_choice, run_name, epochs, batch, imgsz, lr, opt,
+                              cmd_template, progress=gr.Progress()):
+    """
+    Train using RT-DETRv2 repo via a configurable command template.
+    We stream logs, parse simple metrics when patterns match, and try to locate a best checkpoint on completion.
+    """
     if not dataset_path:
         raise gr.Error("Finalize a dataset in Tab 2 before training.")
 
-    # Ultralytics expects device string, e.g. '0' or 'cpu'
-    device_str = "0" if torch.cuda.is_available() else "cpu"
-
-    metrics_queue = Queue()
-
-    def on_epoch_end(trainer):
-        # Be defensive about metric keys
-        m = trainer.metrics or {}
-        metrics_queue.put({
-            'epoch': (trainer.epoch or 0) + 1,
-            'train_loss': m.get('train/loss') or m.get('loss'),
-            'val_loss': m.get('val/loss'),
-            'mAP50': m.get('metrics/mAP50(B)') or m.get('metrics/mAP50'),
-            'mAP50_95': m.get('metrics/mAP50-95(B)') or m.get('metrics/mAP50-95')
-        })
-
-    def train_thread_func():
-        try:
-            model_url = next(m['url'] for m in RTDETR_MODELS['detection'] if m['filename'] == model_filename)
-            weights_path = os.path.join('pretrained_models', model_filename)
-            if not os.path.exists(weights_path):
-                os.makedirs('pretrained_models', exist_ok=True)
-                r = requests.get(model_url, stream=True, timeout=60)
-                r.raise_for_status()
-                with open(weights_path, 'wb') as f:
-                    for chunk in r.iter_content(chunk_size=8192):
-                        f.write(chunk)
-
-            model = YOLO(weights_path)
-            model.add_callback("on_train_epoch_end", on_epoch_end)
-
-            model.train(
-                data=os.path.join(dataset_path, 'data.yaml'),
-                epochs=int(epochs),
-                batch=int(batch),
-                imgsz=int(imgsz),
-                lr0=float(lr),
-                optimizer=str(opt),
-                project='runs/train',
-                name=str(run_name),
-                exist_ok=True,
-                device=device_str
-            )
-            metrics_queue.put("done")
-        except Exception as e:
-            logging.exception("Training thread error")
-            metrics_queue.put(f"error: {e}")
+    # Make sure repo exists
+    try:
+        ensure_repo(repo_dir)
+    except subprocess.CalledProcessError as e:
+        raise gr.Error(f"Failed to clone RT-DETRv2 repo: {e}")
+
+    # Prepare output directory
+    output_dir = os.path.join("runs", "train", str(run_name))
+    os.makedirs(output_dir, exist_ok=True)
+
+    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 the command
+    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
+    )
 
-    Thread(target=train_thread_func, daemon=True).start()
+    # Launch training subprocess in repo_dir
+    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}  # inherit env (CUDA, etc.)
+    )
 
+    # Live metrics
     history = {k: [] for k in ['epoch', 'train_loss', 'val_loss', 'mAP50', 'mAP50_95']}
-    while True:
-        item = metrics_queue.get()
-        if isinstance(item, str):
-            if item == "done":
-                break
-            if item.startswith("error"):
-                raise gr.Error(f"Training failed: {item}")
-
-        # Append metrics
-        for key in ['epoch', 'train_loss', 'val_loss', 'mAP50', 'mAP50_95']:
-            val = item.get(key, None)
-            if val is not None:
-                history[key].append(val)
-
-        current_epoch = history['epoch'][-1] if history['epoch'] else 0
-        total_epochs = int(epochs)
-        frac = min(max(current_epoch / max(1, total_epochs), 0.0), 1.0)
-        progress(frac, desc=f"Epoch {current_epoch}/{total_epochs}")
-
-        # 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 f"Epoch {current_epoch}/{total_epochs} complete.", fig_loss, fig_map, None
-
-    final_path = os.path.join('runs', 'train', str(run_name), 'weights', 'best.pt')
-    if not os.path.exists(final_path):
-        raise gr.Error("Training finished, but 'best.pt' was not found.")
-
-    yield "Training complete!", None, None, gr.File.update(value=final_path, visible=True)
-
+    last_epoch = 0
+
+    # Stream logs and parse
+    for line in iter(proc.stdout.readline, ''):
+        line = line.rstrip()
+        # Update progress indeterminately (we don't know total epochs from logs generically)
+        if "epoch" in line.lower():
+            progress(0.0, desc=line[-120:])  # show last part of the line
+        else:
+            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 when we detect an epoch number or mAP/loss update
+            # 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")
+
+            # Emit an update to the UI (status text is the last log line)
+            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.")
+
+    # Try to locate a best checkpoint
+    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:
+        # Still complete, but we couldn't find a checkpoint automatically
+        yield "Training finished. Could not auto-detect 'best' checkpoint; please check the output directory.", None, None, gr.update(visible=False)
 
 def upload_handler(model_file, hf_token, hf_repo, gh_token, gh_repo, progress=gr.Progress()):
     """Handles model upload to Hugging Face and GitHub."""
@@ -649,12 +668,11 @@ def upload_handler(model_file, hf_token, hf_repo, gh_token, gh_repo, progress=gr
     progress(1)
     return hf_status, gh_status
 
-
 # ------------------------------
 # Gradio UI
 # ------------------------------
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
-    gr.Markdown("# Rolo: A Dedicated RT-DETR Training Dashboard")
+    gr.Markdown("# Rolo: RT-DETRv2 Training Dashboard (Supervisely Ecosystem Backend)")
 
     # State variables
     dataset_info_state = gr.State([])
@@ -689,38 +707,54 @@ 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 Train the RT-DETR Model")
+            gr.Markdown("### Set Hyperparameters and Train with RT-DETRv2")
             with gr.Row():
                 with gr.Column(scale=1):
-                    model_file_dd = gr.Dropdown(
-                        label="Select Pre-Trained RT-DETR Model",
-                        choices=[m["filename"] for m in RTDETR_MODELS["detection"]],
+                    model_choice_dd = gr.Dropdown(
+                        label="Model Choice (label only – adjust your command template to use the right config)",
+                        choices=RTDETRV2_MODELS,
                         value=DEFAULT_MODEL
                     )
-                    run_name_tb = gr.Textbox(label="Run Name", value="rtdetr_run_1")
+                    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, 32, 8, step=1, label="Batch Size")
+                    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="Adam", label="Optimizer")
+                    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
+                    )
                     train_btn = gr.Button("Start Training", variant="primary")
                 with gr.Column(scale=2):
-                    train_status = gr.Textbox(label="Live Status", interactive=False)
+                    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.pt)", interactive=False, visible=False)
+                    final_model_file = gr.File(label="Download Trained Model (best.*)", interactive=False, visible=False)
 
         with gr.TabItem("4. Upload Model"):
-            gr.Markdown("### Upload Your Trained Model\nAfter training, you can upload the `best.pt` file to Hugging Face and/or GitHub.")
+            gr.Markdown("### Upload Your Trained Model\nAfter training, you can upload the best checkpoint to Hugging Face and/or GitHub.")
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("#### Hugging Face")
                     hf_token = gr.Textbox(label="Hugging Face API Token", type="password")
-                    hf_repo = gr.Textbox(label="Hugging Face Repo ID", placeholder="e.g., username/my-rtdetr-model")
+                    hf_repo = gr.Textbox(label="Hugging Face Repo ID", placeholder="e.g., username/my-rtdetrv2-model")
                 with gr.Column():
                     gr.Markdown("#### GitHub")
                     gh_token = gr.Textbox(label="GitHub Personal Access Token", type="password")
-                    gh_repo = gr.Textbox(label="GitHub Repo", placeholder="e.g., username/my-rtdetr-repo")
+                    gh_repo = gr.Textbox(label="GitHub Repo", placeholder="e.g., username/my-rtdetrv2-repo")
             upload_btn = gr.Button("Upload Model", variant="primary")
             with gr.Row():
                 hf_status = gr.Textbox(label="Hugging Face Status", interactive=False)
@@ -743,8 +777,19 @@ 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,
-        inputs=[final_dataset_path_state, model_file_dd, run_name_tb, epochs_sl, batch_sl, imgsz_num, lr_num, opt_dd],
+        fn=training_handler_rtdetrv2,
+        inputs=[
+            final_dataset_path_state,  # dataset_path
+            repo_dir_tb,               # repo_dir
+            model_choice_dd,           # model_choice (label only)
+            run_name_tb,
+            epochs_sl,
+            batch_sl,
+            imgsz_num,
+            lr_num,
+            opt_dd,
+            cmd_template_tb
+        ],
         outputs=[train_status, loss_plot, map_plot, final_model_file]
     )
     upload_btn.click(
@@ -754,6 +799,5 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
     )
 
 if __name__ == "__main__":
-    # Tip: silence Ultralytics settings warning by setting env var:
-    #   export YOLO_CONFIG_DIR=/tmp/Ultralytics
+    # If Ultralytics warnings annoy you, set: export YOLO_CONFIG_DIR=/tmp/Ultralytics
     app.launch(debug=True)