Update app.py

app.py

@@ -1,28 +1,3 @@
-# gradio_dataset_manager.py – Unified YOLO Dataset Toolkit (Evaluate + Merge/Edit)
-"""
-Gradio application that **combines** the functionality of the original evaluation‑only app and
-the Streamlit dataset‑merging dashboard.
-
-### Key features added
-* **Multi‑dataset loader**  
-  – Roboflow URLs (with automatic latest‑version fallback)  
-  – ZIP uploads (GitHub/Ultralytics releases, etc.)  
-  – Existing server paths
-* **Interactive class manager** (rename / remove / per‑class max images) using an **editable
-  Gradio `Dataframe`**. Multiple rows can share the same *New name* to merge classes.
-* **Dataset merger** with:
-  – Per‑class image limits  
-  – Consistent re‑indexing after renames/removals  
-  – Duplicate image avoidance  
-  – Final `data.yaml` + ready‑to‑train directory structure
-* One‑click **ZIP download** of the merged dataset
-* Original **Quality Evaluation** preserved (blur/IOU/conf sliders, fastdup, optional
-  Model‑QA/Cleanlab)
-
-> ⚠️ This is a standalone Python script – drop it into a Hugging Face **Space** or run with
-> `python gradio_dataset_manager.py`.  Requires the same pip deps you already list.
-"""
-
 from __future__ import annotations
 
 import base64
@@ -36,6 +11,7 @@ import re
 import shutil
 import stat
 import tempfile
+import zipfile
 from collections import Counter
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
@@ -178,18 +154,35 @@ def get_model(weights: str) -> YOLO | None:
 # -------------------- Roboflow helpers --------------------
 RF_RE = re.compile(r"https?://universe\.roboflow\.com/([^/]+)/([^/]+)/?([^/]*)")
 
-
-def parse_roboflow_url(url: str):
-    """Return (workspace, project, version|None) – tolerates many RF URL flavours."""
+def parse_roboflow_url(url: str) -> tuple[str, str, int | None]:
+    """
+    Return (workspace, project, version|None) – tolerates many RF URL flavours.
+    Any non‐positive or malformed version is treated as None.
+    """
     m = RF_RE.match(url.strip())
     if not m:
         return None, None, None
     ws, proj, tail = m.groups()
-    ver = None
+    ver: int | None = None
+
+    # explicit "dataset/<number>" in path
     if tail.startswith("dataset/"):
-        ver = tail.split("dataset/")[-1]
-    elif "?version=" in url:
-        ver = url.split("?version=")[-1]
+        try:
+            v = int(tail.split("dataset/", 1)[1])
+            if v > 0:
+                ver = v
+        except ValueError:
+            pass
+
+    # explicit "?version=<number>" in query
+    if ver is None and "?version=" in url:
+        try:
+            v = int(url.split("?version=", 1)[1])
+            if v > 0:
+                ver = v
+        except ValueError:
+            pass
+
     return ws, proj, ver
 
 
@@ -204,32 +197,44 @@ def get_latest_version(rf: Roboflow, ws: str, proj: str) -> str | None:
         return None
 
 
-def download_roboflow_dataset(url: str, rf_api_key: str, fmt: str = "yolov8") -> Tuple[Path, List[str], List[str]]:
+def download_roboflow_dataset(
+    url: str,
+    rf_api_key: str,
+    fmt: str = "yolov8",
+) -> Tuple[Path, List[str], List[str]]:
     """Return (dataset_location, class_names, splits). Caches by folder name."""
     if Roboflow is None:
         raise RuntimeError("`roboflow` pip package not installed")
+
     ws, proj, ver = parse_roboflow_url(url)
     if not (ws and proj):
-        raise ValueError("Bad Roboflow URL")
+        raise ValueError(f"Bad Roboflow URL: {url!r}")
 
     rf = Roboflow(api_key=rf_api_key)
-    if ver is None:
-        ver = get_latest_version(rf, ws, proj)
-        if ver is None:
+
+    # if no explicit version or invalid, fetch latest
+    if not ver or ver <= 0:
+        latest = get_latest_version(rf, ws, proj)
+        if latest is None:
             raise RuntimeError("Could not resolve latest Roboflow version")
+        try:
+            ver = int(latest)
+        except ValueError:
+            raise RuntimeError(f"Invalid latest version returned: {latest!r}")
 
     ds_dir = TMP_ROOT / f"{ws}_{proj}_v{ver}"
     if ds_dir.exists():
         yaml_path = ds_dir / "data.yaml"
         class_names = load_yaml(yaml_path).get("names", []) if yaml_path.exists() else []
-        splits = [s for s in ["train", "valid", "test"] if (ds_dir / s).exists()]
+        splits = [s for s in ("train","valid","test") if (ds_dir / s).exists()]
         return ds_dir, class_names, splits
 
     ds_dir.mkdir(parents=True, exist_ok=True)
-    rf.workspace(ws).project(proj).version(int(ver)).download(fmt, location=str(ds_dir))
+    rf.workspace(ws).project(proj).version(ver).download(fmt, location=str(ds_dir))
+
     yaml_path = ds_dir / "data.yaml"
     class_names = load_yaml(yaml_path).get("names", []) if yaml_path.exists() else []
-    splits = [s for s in ["train", "valid", "test"] if (ds_dir / s).exists()]
+    splits = [s for s in ("train","valid","test") if (ds_dir / s).exists()]
     return ds_dir, class_names, splits
 
 
@@ -252,8 +257,6 @@ def gather_class_counts(dataset_info_list, class_name_mapping):
     return dict(counts)
 
 
-# -- label‑file worker (same logic, no Streamlit)
-
 def _process_label_file(label_path: Path, class_names_dataset, class_name_mapping):
     im_name = label_path.stem + label_path.suffix.replace(".txt", ".jpg")
     img_classes = set()
@@ -265,10 +268,6 @@ def _process_label_file(label_path: Path, class_names_dataset, class_name_mappin
     return im_name, img_classes
 
 
-# ---------------------------------------------------------------------------
-#  merge_datasets():  **pure‑python** version (no streamlit session state)
-# ---------------------------------------------------------------------------
-
 def merge_datasets(
     dataset_info_list: List[Tuple[str, List[str], List[str], str]],
     class_map_df: pd.DataFrame,
@@ -284,21 +283,18 @@ def merge_datasets(
     (out_dir / "valid/images").mkdir(parents=True, exist_ok=True)
     (out_dir / "valid/labels").mkdir(parents=True, exist_ok=True)
 
-    # Build mapping dicts ----------------------------------------------------
     class_name_mapping = {
-        row["original_class"]: row["new_name"] if row["remove"] is False else "__REMOVED__"
+        row["original_class"]: row["new_name"] if not row["remove"] else "__REMOVED__"
         for _, row in class_map_df.iterrows()
     }
     limits_per_merged = {
         row["new_name"]: int(row["max_images"])
         for _, row in class_map_df.iterrows()
-        if row["remove"] is False
+        if not row["remove"]
     }
-    # active merged classes only
     active_classes = [c for c in sorted(set(class_name_mapping.values())) if c != "__REMOVED__"]
     id_map = {cls: idx for idx, cls in enumerate(active_classes)}
 
-    # Scan label files -------------------------------------------------------
     image_to_classes: dict[str, set[str]] = {}
     image_to_label: dict[str, Path] = {}
     class_to_images: dict[str, set[str]] = {c: set() for c in active_classes}
@@ -319,7 +315,6 @@ def merge_datasets(
                 for c in cls_set:
                     class_to_images[c].add(img_path)
 
-    # Select images respecting per‑class limits -----------------------------
     selected_images: set[str] = set()
     counters = {c: 0 for c in active_classes}
     shuffle_pool = [img for imgs in class_to_images.values() for img in imgs]
@@ -327,14 +322,12 @@ def merge_datasets(
 
     for img in shuffle_pool:
         cls_set = image_to_classes[img]
-        # skip if any class hit its limit
         if any(counters[c] >= limits_per_merged.get(c, 0) for c in cls_set):
             continue
         selected_images.add(img)
         for c in cls_set:
             counters[c] += 1
 
-    # Copy & re‑index --------------------------------------------------------
     for img in selected_images:
         split = "train" if random.random() < 0.9 else "valid"
         dst_img = out_dir / split / "images" / Path(img).name
@@ -352,7 +345,6 @@ def merge_datasets(
             if not parts:
                 continue
             cid = int(parts[0])
-            # find orig class name
             dloc_match = next((cl for dloc2, cl, _, _ in dataset_info_list if str(lp_src).startswith(dloc2)), None)
             if dloc_match is None:
                 continue
@@ -370,7 +362,6 @@ def merge_datasets(
         else:
             (out_dir / split / "images" / Path(img).name).unlink(missing_ok=True)
 
-    # Build data.yaml --------------------------------------------------------
     data_yaml = {
         "path": str(out_dir.resolve()),
         "train": "train/images",
@@ -384,12 +375,14 @@ def merge_datasets(
     return out_dir
 
 
-# Utility: zip a folder to bytes -------------------------------------------
-
 def zip_directory(folder: Path) -> bytes:
     buf = io.BytesIO()
-    with shutil.make_archive("dataset", "zip", folder) as _:
-        pass  # make_archive writes to disk – we avoid that (quick hack)
+    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+        for file in folder.rglob("*"):
+            zf.write(file, arcname=file.relative_to(folder))
+    buf.seek(0)
+    return buf.getvalue()
+
 
 # ════════════════════════════════════════════════════════════════════════════
 #                                 UI LAYER
@@ -420,7 +413,6 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
         out_md   = gr.Markdown()
         out_df   = gr.Dataframe(label="Class distribution")
 
-        # --- callback (identical logic from v3, omitted for brevity) ---
         def _evaluate_cb(api_key, url_txt, zip_file, server_path, yaml_file, weights,
                          blur_thr, iou_thr, conf_thr, run_dup, run_modelqa):
             return "Evaluation disabled in this trimmed snippet.", pd.DataFrame()
@@ -434,7 +426,7 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
 
     # ------------------------------ MERGE TAB -----------------------------
     with gr.Tab("Merge / Edit"):
-        gr.Markdown("""### 1️⃣ Load one or more datasets""")
+        gr.Markdown("### 1️⃣ Load one or more datasets")
         rf_key   = gr.Textbox(label="Roboflow API key", type="password")
         rf_urls  = gr.File(label=".txt of RF URLs", file_types=['.txt'])
         zips_in  = gr.Files(label="One or more dataset ZIPs")
@@ -446,14 +438,19 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
             global autoinc
             info_list = []
             log_lines = []
+
             # Roboflow URLs via txt
             if rf_urls_file is not None:
                 for url in Path(rf_urls_file.name).read_text().splitlines():
                     if not url.strip():
                         continue
-                    ds, names, splits = download_roboflow_dataset(url, rf_key)
-                    info_list.append((str(ds), names, splits, Path(ds).name))
-                    log_lines.append(f"✔️ RF dataset **{Path(ds).name}** loaded ({len(names)} classes)")
+                    try:
+                        ds, names, splits = download_roboflow_dataset(url, rf_key)
+                        info_list.append((str(ds), names, splits, Path(ds).name))
+                        log_lines.append(f"✔️ RF dataset **{Path(ds).name}** loaded ({len(names)} classes)")
+                    except Exception as e:
+                        log_lines.append(f"⚠️ RF load failed for {url!r}: {e}")
+
             # ZIPs
             for f in zip_files or []:
                 autoinc += 1
@@ -464,30 +461,33 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
                 if yaml_path is None:
                     continue
                 names = load_yaml(yaml_path).get("names", [])
-                splits = [s for s in ["train", "valid", "test"] if (tmp / s).exists()]
+                splits = [s for s in ("train","valid","test") if (tmp / s).exists()]
                 info_list.append((str(tmp), names, splits, tmp.name))
                 log_lines.append(f"✔️ ZIP **{tmp.name}** loaded")
+
             return info_list, "\n".join(log_lines) if log_lines else "No datasets loaded."
 
         load_btn.click(_load_cb, [rf_key, rf_urls, zips_in], [ds_state, load_log])
 
         # ------------- Class map editable table --------------------------
-        gr.Markdown("""### 2️⃣ Edit class mapping / limits / removal""")
-        class_df = gr.Dataframe(headers=["original_class", "new_name", "max_images", "remove"],
-                                datatype=["str", "str", "number", "bool"],
-                                interactive=True, elem_id="classdf")
+        gr.Markdown("### 2️⃣ Edit class mapping / limits / removal")
+        class_df = gr.Dataframe(
+            headers=["original_class", "new_name", "max_images", "remove"],
+            datatype=["str", "str", "number", "bool"],
+            interactive=True, elem_id="classdf"
+        )
         refresh_btn = gr.Button("Build class table from loaded datasets")
 
         def _build_class_df(ds_info):
             class_names_all = []
-            for _dloc, names, _spl, _n in ds_info:
+            for _dloc, names, _spl, _ in ds_info:
                 class_names_all.extend(names)
             class_names_all = sorted(set(class_names_all))
             df = pd.DataFrame({
                 "original_class": class_names_all,
                 "new_name":       class_names_all,
-                "max_images":     [99999]*len(class_names_all),
-                "remove":         [False]*len(class_names_all),
+                "max_images":     [99999] * len(class_names_all),
+                "remove":         [False]  * len(class_names_all),
             })
             return df
 
@@ -499,11 +499,14 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
         merge_log = gr.Markdown()
 
         def _merge_cb(ds_info, class_df):
-            if len(ds_info) == 0:
+            if not ds_info:
                 return None, "⚠️ Load datasets first."
-            out_dir = merge_datasets(ds_info, class_df)  # may be slow
+            out_dir = merge_datasets(ds_info, class_df)
             zip_path = shutil.make_archive(str(out_dir), "zip", out_dir)
-            return zip_path, f"✅ Merged dataset created at **{out_dir}** with {len(list(Path(out_dir).rglob('*.jpg')))} images."
+            return zip_path, (
+                f"✅ Merged dataset created at **{out_dir}** with "
+                f"{len(list(Path(out_dir).rglob('*.jpg')))} images."
+            )
 
         merge_btn.click(_merge_cb, [ds_state, class_df], [zip_out, merge_log])