Update app.py

app.py

@@ -16,9 +16,23 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 REPO_URL = "https://github.com/supervisely-ecosystem/RT-DETRv2"
 REPO_DIR = os.path.join(os.getcwd(), "third_party", "RT-DETRv2")
 PY_IMPL_DIR = os.path.join(REPO_DIR, "rtdetrv2_pytorch")  # Supervisely keeps PyTorch impl here
+
+# Core deps + your requested packages; pinned as lower-bounds to avoid downgrades
 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",
+    "gradio>=4.36.1",
+    "ultralytics>=8.2.0",
+    "roboflow>=1.1.28",
+    "requests>=2.31.0",
+    "huggingface_hub>=0.22.0",
+    "pandas>=2.0.0",
+    "matplotlib>=3.7.0",
+    "torch>=2.0.1",
+    "torchvision>=0.15.2",
+    "pyyaml>=6.0.1",
+    "Pillow>=10.0.0",
+    "supervisely>=6.0.0",     # <- fixes ModuleNotFoundError from repo trainer
+    "tensorboard>=2.13.0",    # convenience: sometimes used by forks
+    "pycocotools>=2.0.7",     # convenience: ensure wheels are present
 ]
 
 # === bootstrap (clone + pip) ===================================================
@@ -37,11 +51,21 @@ def ensure_repo_and_requirements():
         except Exception:
             logging.warning("git pull failed; continuing with current checkout")
 
+    # Make sure all our app/runtime deps (incl. supervisely & ultralytics) are present
     pip_install(COMMON_REQUIREMENTS)
+
+    # Then install repo-specific extras (pycocotools/tensorboard etc. if required)
     req_file = os.path.join(PY_IMPL_DIR, "requirements.txt")
     if os.path.exists(req_file):
         pip_install(["-r", req_file])
 
+    # Double-check supervisely importability; if not, try again explicitly.
+    try:
+        import supervisely  # noqa: F401
+    except Exception:
+        logging.warning("supervisely not importable after first pass; retrying install…")
+        pip_install(["supervisely>=6.0.0"])
+
 try:
     ensure_repo_and_requirements()
 except Exception:
@@ -141,15 +165,10 @@ def label_path_for(img_path: str) -> str:
 
 # === YOLOv8 -> COCO converter =================================================
 def yolo_to_coco(split_dir_images, split_dir_labels, class_names, out_json):
-    """
-    Convert YOLO txt labels to a COCO annotations json.
-    """
     images, annotations = [], []
     categories = [{"id": i, "name": n} for i, n in enumerate(class_names)]
     ann_id = 1
     img_id = 1
-
-    # Simple image size read (PIL); in Spaces this is fine.
     for fname in sorted(os.listdir(split_dir_images)):
         if not fname.lower().endswith((".jpg",".jpeg",".png")): continue
         img_path = os.path.join(split_dir_images, fname)
@@ -157,10 +176,8 @@ def yolo_to_coco(split_dir_images, split_dir_labels, class_names, out_json):
             with Image.open(img_path) as im:
                 w, h = im.size
         except Exception:
-            # skip unreadable images
             continue
         images.append({"id": img_id, "file_name": fname, "width": w, "height": h})
-
         label_file = os.path.join(split_dir_labels, os.path.splitext(fname)[0] + ".txt")
         if os.path.exists(label_file):
             with open(label_file, "r") as f:
@@ -169,7 +186,6 @@ def yolo_to_coco(split_dir_images, split_dir_labels, class_names, out_json):
                     if len(parts) < 5: continue
                     cls = int(float(parts[0]))
                     cx, cy, bw, bh = map(float, parts[1:5])
-                    # convert normalized (cx,cy,bw,bh) to x,y,w,h in pixels
                     x = (cx - bw/2.0) * w
                     y = (cy - bh/2.0) * h
                     ww = bw * w
@@ -185,16 +201,11 @@ def yolo_to_coco(split_dir_images, split_dir_labels, class_names, out_json):
                     })
                     ann_id += 1
         img_id += 1
-
     coco = {"images": images, "annotations": annotations, "categories": categories}
     os.makedirs(os.path.dirname(out_json), exist_ok=True)
     with open(out_json, "w") as f: json.dump(coco, f)
 
 def make_coco_annotations(merged_dir, class_names):
-    """
-    Build COCO jsons under merged_dir/annotations:
-      instances_train.json, instances_val.json, instances_test.json
-    """
     ann_dir = os.path.join(merged_dir, "annotations")
     os.makedirs(ann_dir, exist_ok=True)
     mapping = {"train": "instances_train.json", "valid": "instances_val.json", "test": "instances_test.json"}
@@ -317,33 +328,19 @@ def finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress=
             'names': active_classes
         }, f)
 
-    # also create COCO jsons for RT-DETRv2 training
     ann_dir = make_coco_annotations(merged_dir, active_classes)
     progress(0.98, desc="Finalizing...")
     return f"Dataset finalized with {len(selected_images)} images.", os.path.abspath(merged_dir)
 
 # === entrypoint + config detection/generation =================================
 def find_training_script(repo_root):
-    """
-    Recursively search for a tools/train.py (or train.py) suitable for RT-DETRv2.
-    """
     candidates = []
     for pat in ["**/tools/train.py", "**/train.py"]:
         candidates.extend(glob(os.path.join(repo_root, pat), recursive=True))
-    # Prefer ones inside rtdetrv2_pytorch
     candidates.sort(key=lambda p: (0 if "rtdetrv2_pytorch" in p else 1, len(p)))
     return candidates[0] if candidates else None
 
 def find_model_config_template(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"],
@@ -369,10 +366,6 @@ def find_model_config_template(model_key):
 
 def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
                       epochs, batch, imgsz, lr, optimizer):
-    """
-    Load the chosen repo config and patch only the keys that already exist.
-    This avoids schema mismatches between forks.
-    """
     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.")
 
@@ -385,12 +378,12 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         "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'
+        "val_img":    os.path.abspath(os.path.join(merged_dir, "valid", "images")),
         "test_img":   os.path.abspath(os.path.join(merged_dir, "test", "images")),
         "out_dir":    os.path.abspath(os.path.join("runs", "train", run_name)),
     }
 
-    # --- dataset block --------------------------------------------------------
+    # dataset block
     for root_key in ["dataset", "data"]:
         if root_key in cfg and isinstance(cfg[root_key], dict):
             ds = cfg[root_key]
@@ -401,17 +394,14 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
             ]:
                 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
+                            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
+                            ds[split][k] = paths[ip]; break
 
-    # --- num_classes ----------------------------------------------------------
+    # num_classes
     def set_num_classes(node, n):
         if not isinstance(node, dict): return False
         if "num_classes" in node:
@@ -426,7 +416,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
     else:
         cfg["model"] = {"num_classes": int(class_count)}
 
-    # --- epochs / image size --------------------------------------------------
+    # epochs / imgsz
     updated_epoch = False
     for key in ["max_epoch", "epochs", "num_epochs"]:
         if key in cfg:
@@ -442,7 +432,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         if key in cfg: cfg[key] = int(imgsz)
     if "input_size" not in cfg: cfg["input_size"] = int(imgsz)
 
-    # --- learning rate / optimizer / batch -----------------------------------
+    # lr / optimizer / batch
     if "solver" not in cfg or not isinstance(cfg["solver"], dict):
         cfg["solver"] = {}
     sol = cfg["solver"]
@@ -453,13 +443,12 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         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 -----------------------------------------------------------
+    # output dir
     if "output_dir" in cfg:
         cfg["output_dir"] = paths["out_dir"]
     elif "solver" in cfg:
@@ -467,7 +456,6 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
     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)
@@ -570,23 +558,19 @@ 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) 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) 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 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", [])]
     make_coco_annotations(dataset_path, class_names)
 
-    # 4) patch the base config safely (no custom schema assumptions)
     cfg_path = patch_base_config(
         base_cfg_path=base_cfg,
         merged_dir=dataset_path,
@@ -602,7 +586,6 @@ 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 command (no extra flags that might not exist)
     cmd = [sys.executable, train_script, "-c", os.path.abspath(cfg_path)]
     logging.info(f"Training command: {' '.join(cmd)}")
 
@@ -610,11 +593,9 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     def run_train():
         try:
             env = os.environ.copy()
-            # 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,
@@ -628,7 +609,7 @@ 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
+    first_lines = []
     while True:
         line = q.get()
         if line.startswith("__EXITCODE__"):