Update app.py

app.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import base64
-import functools
 import imghdr
 import io
 import json
@@ -12,13 +11,12 @@ import re
 import shutil
 import stat
 import tempfile
-import time
 import zipfile
 from collections import Counter
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Tuple
 
 import gradio as gr
 import numpy as np
@@ -51,25 +49,15 @@ try:
     from cleanlab.pruning import get_noise_indices
 except ImportError:
     get_noise_indices = None
-try:
-    import torch
-except ImportError:
-    torch = None
 
 # ───────────────── Config & Constants ───────────────────────────────────────
 TMP_ROOT = Path(tempfile.gettempdir()) / "rf_datasets"
 TMP_ROOT.mkdir(parents=True, exist_ok=True)
+CPU_COUNT    = int(os.getenv("QC_CPU",      1))
+BATCH_SIZE   = int(os.getenv("QC_BATCH",    4))
+SAMPLE_LIMIT = int(os.getenv("QC_SAMPLE", 200))
 
-# Clean out any subfolders older than 1 hour
-_now = time.time()
-for sub in TMP_ROOT.iterdir():
-    if sub.is_dir() and (_now - sub.stat().st_mtime) > 3600:
-        shutil.rmtree(sub, ignore_errors=True)
-
-CPU_COUNT     = int(os.getenv("QC_CPU",      os.cpu_count() or 1))
-BATCH_SIZE    = int(os.getenv("QC_BATCH",    4))
-SAMPLE_LIMIT  = int(os.getenv("QC_SAMPLE", 200))
-DEFAULT_W     = {
+DEFAULT_W = {
     "Integrity":     0.25,
     "Class balance": 0.10,
     "Image quality": 0.15,
@@ -77,32 +65,26 @@ DEFAULT_W     = {
     "Model QA":      0.30,
     "Label issues":  0.10,
 }
+
 logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
 
 _model_cache: dict[str, YOLO] = {}
-autoinc = 0  # for ZIP unpack dirs
+autoinc = 0  # helper for tmp‑dir names
 
+# ────────────────────────────────────────────────────────────────────────────
+#  Data‑class & basic helpers
 # ────────────────────────────────────────────────────────────────────────────
 @dataclass
 class QCConfig:
     blur_thr:    float
     iou_thr:     float
     conf_thr:    float
-    weights:     Optional[str]
+    weights:     str | None
     cpu_count:   int = CPU_COUNT
     batch_size:  int = BATCH_SIZE
     sample_limit:int = SAMPLE_LIMIT
 
-@dataclass
-class DatasetInfo:
-    path: str
-    class_names: List[str]
-    splits: List[str]
-    name: str
-
-# ────────────────────────────────────────────────────────────────────────────
 def load_yaml(path: Path) -> Dict:
-    """Load a YAML file safely."""
     with path.open('r', encoding='utf-8') as f:
         return yaml.safe_load(f)
 
@@ -113,10 +95,8 @@ def load_class_names(yaml_path: Path) -> List[str]:
         return [names[k] for k in sorted(names, key=lambda x: int(x))]
     return list(names)
 
-@functools.lru_cache(maxsize=None)
-def parse_label_file(path: Path) -> List[Tuple[int,float,float,float,float]]:
-    """Parse a YOLO label file; empty or missing yields empty list."""
-    if not path.exists() or path.stat().st_size == 0:
+def parse_label_file(path: Path) -> list[tuple[int, float, float, float, float]]:
+    if not path or not path.exists() or path.stat().st_size == 0:
         return []
     try:
         arr = np.loadtxt(path, dtype=float)
@@ -136,8 +116,7 @@ def guess_image_dirs(root: Path) -> List[Path]:
     ]
     return [d for d in candidates if d.exists()]
 
-def gather_dataset(root: Path, yaml_path: Optional[Path]) -> Tuple[List[Path], List[Optional[Path]], Dict]:
-    """Return lists of image paths, label paths, and metadata."""
+def gather_dataset(root: Path, yaml_path: Path | None):
     if yaml_path is None:
         yamls = list(root.glob('*.yaml'))
         if not yamls:
@@ -155,16 +134,26 @@ def gather_dataset(root: Path, yaml_path: Optional[Path]) -> Tuple[List[Path], L
     ]
     return imgs, lbls, meta
 
-def get_model(weights: str) -> Optional[YOLO]:
-    """Load/cache a YOLO model, preferring GPU if available."""
+def get_model(weights: str) -> YOLO | None:
     if not weights or YOLO is None:
         return None
     if weights not in _model_cache:
-        device = "cuda" if torch and torch.cuda.is_available() else "cpu"
-        _model_cache[weights] = YOLO(weights, device=device)
+        _model_cache[weights] = YOLO(weights)
     return _model_cache[weights]
 
-# ────────────────────────────────────────────────────────────────────────────
+# ───────── Concurrency helpers & QC functions ───────────────────────────────
+def _quality_stat_args(args: Tuple[Path, float]) -> Tuple[Path, bool, bool, bool]:
+    path, thr = args
+    if cv2 is None:
+        return path, False, False, False
+    im = cv2.imread(str(path))
+    if im is None:
+        return path, False, False, False
+    gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
+    lap = cv2.Laplacian(gray, cv2.CV_64F).var()
+    mean = gray.mean()
+    return path, lap < thr, mean < 25, mean > 230
+
 def _is_corrupt(path: Path) -> bool:
     try:
         with Image.open(path) as im:
@@ -173,172 +162,171 @@ def _is_corrupt(path: Path) -> bool:
     except Exception:
         return True
 
-def _quality_stat_args(args: Tuple[str, float]) -> Tuple[str, bool, bool, bool]:
-    """Return (path, is_blurry, is_dark, is_bright)."""
-    path_str, thr = args
-    if cv2 is None:
-        return path_str, False, False, False
-    im = cv2.imread(path_str)
-    if im is None:
-        return path_str, False, False, False
-    gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
-    lap = cv2.Laplacian(gray, cv2.CV_64F).var()
-    mean = float(gray.mean())
-    return path_str, lap < thr, mean < 25, mean > 230
-
-def qc_integrity(imgs: List[Path], lbls: List[Optional[Path]], cfg: QCConfig) -> Dict:
-    """Check for missing labels & corrupt images."""
+def qc_integrity(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
     missing = [i for i, l in zip(imgs, lbls) if l is None]
     corrupt = []
     sample = imgs[:cfg.sample_limit]
-    with ProcessPoolExecutor(max_workers=cfg.cpu_count) as ex:
-        futures = {ex.submit(_is_corrupt, str(p)): p for p in sample}
-        for f in as_completed(futures):
+    with ThreadPoolExecutor(max_workers=cfg.cpu_count) as ex:
+        fut = {ex.submit(_is_corrupt, p): p for p in sample}
+        for f in as_completed(fut):
             if f.result():
-                corrupt.append(futures[f])
-    score = max(0.0, 100 - (len(missing) + len(corrupt)) / max(len(imgs),1) * 100)
+                corrupt.append(fut[f])
+    score = 100 - (len(missing) + len(corrupt)) / max(len(imgs), 1) * 100
     return {
         "name": "Integrity",
-        "score": score,
+        "score": max(score, 0),
         "details": {
             "missing_label_files": [str(p) for p in missing],
-            "corrupt_images":       [str(p) for p in corrupt],
+            "corrupt_images": [str(p) for p in corrupt],
         }
     }
 
-def qc_class_balance(lbls: List[Optional[Path]], cfg: QCConfig) -> Dict:
+def qc_class_balance(lbls: List[Path], cfg: QCConfig) -> Dict:
     counts, boxes = Counter(), []
     for l in lbls[:cfg.sample_limit]:
         bs = parse_label_file(l) if l else []
         boxes.append(len(bs))
         counts.update(int(b[0]) for b in bs)
     if not counts:
-        return {"name": "Class balance", "score": 0.0, "details": "No labels"}
+        return {"name": "Class balance", "score": 0, "details": "No labels"}
     bal = min(counts.values()) / max(counts.values()) * 100
     return {
         "name": "Class balance",
         "score": bal,
         "details": {
-            "class_counts":    dict(counts),
-            "boxes_per_image": {"min": min(boxes), "max": max(boxes), "mean": float(np.mean(boxes))}
+            "class_counts": dict(counts),
+            "boxes_per_image": {
+                "min": min(boxes),
+                "max": max(boxes),
+                "mean": float(np.mean(boxes))
+            }
         }
     }
 
 def qc_image_quality(imgs: List[Path], cfg: QCConfig) -> Dict:
     if cv2 is None:
-        return {"name": "Image quality", "score": 100.0, "details": "cv2 missing"}
+        return {"name": "Image quality", "score": 100, "details": "cv2 missing"}
     blurry, dark, bright = [], [], []
     sample = imgs[:cfg.sample_limit]
-    args = [(str(p), cfg.blur_thr) for p in sample]
-    with ProcessPoolExecutor(max_workers=cfg.cpu_count) as ex:
-        for path_str, isb, isd, isB in ex.map(_quality_stat_args, args):
-            if isb: blurry.append(path_str)
-            if isd: dark.append(path_str)
-            if isB: bright.append(path_str)
-    bad = len(set(blurry) | set(dark) | set(bright))
-    score = max(0.0, 100 - bad / max(len(sample),1) * 100)
+    with ThreadPoolExecutor(max_workers=cfg.cpu_count) as ex:
+        args = [(p, cfg.blur_thr) for p in sample]
+        for p, isb, isd, isB in ex.map(_quality_stat_args, args):
+            if isb: blurry.append(p)
+            if isd: dark.append(p)
+            if isB: bright.append(p)
+    bad = len({*blurry, *dark, *bright})
+    score = 100 - bad / max(len(sample), 1) * 100
     return {
         "name": "Image quality",
         "score": score,
-        "details": {"blurry": blurry, "dark": dark, "bright": bright}
+        "details": {
+            "blurry": [str(p) for p in blurry],
+            "dark": [str(p) for p in dark],
+            "bright": [str(p) for p in bright]
+        }
     }
 
 def qc_duplicates(imgs: List[Path], cfg: QCConfig) -> Dict:
-    if fastdup and len(imgs) > 50:
+    if fastdup is not None and len(imgs) > 50:
         try:
             fd = fastdup.create(
                 input_dir=str(Path(imgs[0]).parent.parent),
-                work_dir=str(TMP_ROOT/"fastdup")
+                work_dir=str(TMP_ROOT / "fastdup")
             )
             fd.run()
-            clusters = []
             try:
                 cc = fd.connected_components_grouped(sort_by="comp_size", ascending=False)
-                clusters = cc["files"].tolist()
+                clusters = cc["files"].tolist() if "files" in cc.columns else cc.groupby("component")["filename"].apply(list).tolist()
             except Exception:
                 clusters = fd.connected_components()
-            dup_count = sum(len(c)-1 for c in clusters)
-            score = max(0.0, 100 - dup_count/len(imgs)*100)
-            return {"name":"Duplicates","score":score,"details":{"groups":clusters[:50]}}
+            dup = sum(len(c) - 1 for c in clusters)
+            score = max(0.0, 100 - dup / len(imgs) * 100)
+            return {"name": "Duplicates", "score": score, "details": {"groups": clusters[:50]}}
         except Exception as e:
-            return {"name":"Duplicates","score":100.0,"details":{"error":str(e)}}
-    return {"name":"Duplicates","score":100.0,"details":{"note":"skipped"}}
-
-def _rel_iou(b1, b2) -> float:
-    x1,y1,w1,h1 = b1; x2,y2,w2,h2 = b2
-    xa1,ya1 = x1-w1/2, y1-h1/2
-    xa2,ya2 = x1+w1/2, y1+h1/2
-    xb1,yb1 = x2-w2/2, y2-h2/2
-    xb2,yb2 = x2+w2/2, y2+h2/2
-    ix = max(0, min(xa2,xb2)-max(xa1,xb1))
-    iy = max(0, min(ya2,yb2)-max(ya1,yb1))
-    inter = ix*iy
+            return {"name": "Duplicates", "score": 100.0, "details": {"fastdup_error": str(e)}}
+    return {"name": "Duplicates", "score": 100.0, "details": {"note": "skipped"}}
+
+def _rel_iou(b1, b2):
+    x1, y1, w1, h1 = b1
+    x2, y2, w2, h2 = b2
+    xa1, ya1 = x1 - w1/2, y1 - h1/2
+    xa2, ya2 = x1 + w1/2, y1 + h1/2
+    xb1, yb1 = x2 - w2/2, y2 - h2/2
+    xb2, yb2 = x2 + w2/2, y2 + h2/2
+    ix1 = max(xa1, xb1); iy1 = max(ya1, yb1)
+    ix2 = min(xa2, xb2); iy2 = min(ya2, yb2)
+    inter = max(ix2 - ix1, 0) * max(iy2 - iy1, 0)
     union = w1*h1 + w2*h2 - inter
     return inter/union if union else 0.0
 
-def qc_model_qa(imgs: List[Path], lbls: List[Optional[Path]], cfg: QCConfig) -> Dict:
+def qc_model_qa(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
     model = get_model(cfg.weights)
     if model is None:
-        return {"name":"Model QA","score":100.0,"details":"skipped"}
-    ious, mismatches = [], []
+        return {"name": "Model QA", "score": 100, "details": "skipped"}
+    ious, mism = [], []
     sample = imgs[:cfg.sample_limit]
     for i in range(0, len(sample), cfg.batch_size):
         batch = sample[i:i+cfg.batch_size]
-        with torch.no_grad() if torch else contextlib.nullcontext():
-            results = model.predict(batch, verbose=False, half=True, dynamic=True)
+        results = model.predict(batch, verbose=False, half=True, dynamic=True)
         for p, res in zip(batch, results):
-            gt = parse_label_file(Path(p).with_suffix('.txt').parent)  # adjust as needed
+            gt = parse_label_file(Path(p).parent.parent/'labels'/f"{Path(p).stem}.txt")
             for cls, x, y, w, h in gt:
-                best = max(
-                    ( _rel_iou((x,y,w,h), tuple(b)) 
-                      for b,c,conf in zip(res.boxes.xywh.cpu().numpy(),
-                                         res.boxes.cls.cpu().numpy(),
-                                         res.boxes.conf.cpu().numpy())
-                      if conf>=cfg.conf_thr and int(c)==cls
-                    ), default=0.0
-                )
+                best = 0.0
+                for b, c, conf in zip(
+                    res.boxes.xywh.cpu().numpy(),
+                    res.boxes.cls.cpu().numpy(),
+                    res.boxes.conf.cpu().numpy()
+                ):
+                    if conf < cfg.conf_thr or int(c) != cls:
+                        continue
+                    best = max(best, _rel_iou((x, y, w, h), tuple(b)))
                 ious.append(best)
                 if best < cfg.iou_thr:
-                    mismatches.append(str(p))
+                    mism.append(str(p))
     miou = float(np.mean(ious)) if ious else 1.0
-    return {"name":"Model QA","score":miou*100,"details":{"mean_iou":miou,"mismatches":mismatches[:50]}}
+    return {"name": "Model QA", "score": miou*100, "details": {"mean_iou": miou, "mismatches": mism[:50]}}
 
-def qc_label_issues(imgs: List[Path], lbls: List[Optional[Path]], cfg: QCConfig) -> Dict:
-    if not get_noise_indices:
-        return {"name":"Label issues","score":100.0,"details":"skipped"}
+def qc_label_issues(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
+    if get_noise_indices is None:
+        return {"name": "Label issues", "score": 100, "details": "skipped"}
     labels, idxs = [], []
     sample = imgs[:cfg.sample_limit]
-    for idx, p in enumerate(sample):
-        bs = parse_label_file(lbls[idx]) if lbls[idx] else []
-        for cls,*_ in bs:
-            labels.append(int(cls)); idxs.append(idx)
+    for i, p in enumerate(sample):
+        bs = parse_label_file(lbls[i]) if lbls[i] else []
+        for cls, *_ in bs:
+            labels.append(int(cls)); idxs.append(i)
     if not labels:
-        return {"name":"Label issues","score":100.0,"details":"no ground truth"}
-    arr = np.array(labels)
-    uniq = sorted(set(arr))
-    probs = np.eye(len(uniq))[np.searchsorted(uniq, arr)]
-    noise = get_noise_indices(labels=arr, probabilities=probs)
-    flagged = sorted({idxs[n] for n in noise})
-    files = [str(sample[i]) for i in flagged]
-    score = max(0.0, 100 - len(flagged)/len(labels)*100)
-    return {"name":"Label issues","score":score,"details":{"files":files[:50]}}
+        return {"name": "Label issues", "score": 100, "details": "no GT"}
+    labels_arr = np.array(labels)
+    uniq = sorted(set(labels_arr))
+    probs = np.eye(len(uniq))[np.searchsorted(uniq, labels_arr)]
+    noise = get_noise_indices(labels=labels_arr, probabilities=probs)
+    flags = sorted({idxs[n] for n in noise})
+    files = [str(sample[i]) for i in flags]
+    score = 100 - len(flags)/len(labels)*100
+    return {"name": "Label issues", "score": score, "details": {"files": files[:50]}}
 
 def aggregate(results: List[Dict]) -> float:
     return sum(DEFAULT_W[r["name"]]*r["score"] for r in results)
 
-def gather_class_counts(infos: List[DatasetInfo]) -> Counter[str]:
-    ctr = Counter()
-    for info in infos:
-        for split in info.splits:
-            labels_dir = Path(info.path)/split/'labels'
-            if not labels_dir.exists(): continue
-            for lp in labels_dir.rglob('*.txt'):
-                for cls, *_ in parse_label_file(lp):
-                    idx = int(cls)
-                    if 0 <= idx < len(info.class_names):
-                        ctr[info.class_names[idx]] += 1
-    return ctr
-
+# ───────── gathering actual per-class counts ────────────────────────────────
+def gather_class_counts(
+    dataset_info_list: List[Tuple[str, List[str], List[str], str]]
+) -> Counter[str]:
+    counts: Counter[str] = Counter()
+    for dloc, class_names, splits, _ in dataset_info_list:
+        for split in splits:
+            labels_dir = Path(dloc) / split / "labels"
+            if not labels_dir.exists():
+                continue
+            for lp in labels_dir.rglob("*.txt"):
+                for cls_id_float, *_ in parse_label_file(lp):
+                    idx = int(cls_id_float)
+                    if 0 <= idx < len(class_names):
+                        counts[class_names[idx]] += 1
+    return counts
+
+# ───────────────── Roboflow TXT‑loading logic ─────────────────────────────
 RF_RE = re.compile(r"https?://universe\.roboflow\.com/([^/]+)/([^/]+)/dataset/(\d+)")
 
 def download_rf_dataset(url: str, rf_api: Roboflow, dest: Path) -> Path:
@@ -346,17 +334,18 @@ def download_rf_dataset(url: str, rf_api: Roboflow, dest: Path) -> Path:
     if not m:
         raise ValueError(f"Bad RF URL: {url}")
     ws, proj, ver = m.groups()
-    ds_dir = dest/f"{ws}_{proj}_v{ver}"
+    ds_dir = dest / f"{ws}_{proj}_v{ver}"
     if ds_dir.exists():
         return ds_dir
     pr = rf_api.workspace(ws).project(proj)
     pr.version(int(ver)).download("yolov8", location=str(ds_dir))
     return ds_dir
 
+# ───────────────── run_quality & merge_datasets ────────────────────────────
 def run_quality(
     root: Path,
-    yaml_file: Optional[Path],
-    weights: Optional[Path],
+    yaml_file: Path | None,
+    weights: Path | None,
     cfg: QCConfig,
     run_dup: bool,
     run_modelqa: bool
@@ -366,167 +355,197 @@ def run_quality(
         qc_integrity(imgs, lbls, cfg),
         qc_class_balance(lbls, cfg),
         qc_image_quality(imgs, cfg),
-        qc_duplicates(imgs, cfg)      if run_dup     else {"name":"Duplicates","score":100.0,"details":"skipped"},
-        qc_model_qa(imgs, lbls, cfg)  if run_modelqa else {"name":"Model QA","score":100.0,"details":"skipped"},
-        qc_label_issues(imgs, lbls, cfg) if run_modelqa else {"name":"Label issues","score":100.0,"details":"skipped"},
+        qc_duplicates(imgs, cfg)      if run_dup     else {"name":"Duplicates","score":100,"details":"skipped"},
+        qc_model_qa(imgs, lbls, cfg)  if run_modelqa else {"name":"Model QA","score":100,"details":"skipped"},
+        qc_label_issues(imgs, lbls, cfg) if run_modelqa else {"name":"Label issues","score":100,"details":"skipped"},
     ]
     final = aggregate(results)
     md = [f"## **{meta.get('name', root.name)}** — Score {final:.1f}/100"]
     for r in results:
-        md += [
-            f"### {r['name']}  {r['score']:.1f}",
-            "<details><summary>details</summary>\n```json",
-            json.dumps(r["details"], indent=2),
-            "```\n</details>\n"
-        ]
-    # build class-balance df
-    cb = next(r for r in results if r["name"]=="Class balance")["details"]["class_counts"]
-    df = pd.DataFrame.from_dict(cb, orient='index', columns=['count'])
+        md.append(f"### {r['name']}  {r['score']:.1f}")
+        md.append("<details><summary>details</summary>\n```json")
+        md.append(json.dumps(r["details"], indent=2))
+        md.append("```\n</details>\n")
+    df = pd.DataFrame.from_dict(
+        next(r for r in results if r["name"] == "Class balance")["details"]["class_counts"],
+        orient="index", columns=["count"]
+    )
     df.index.name = "class"
     return "\n".join(md), df
 
 def merge_datasets(
-    infos: List[DatasetInfo],
+    dataset_info_list: List[Tuple[str, List[str], List[str], str]],
     class_map_df: pd.DataFrame,
     out_dir: Path = Path("merged_dataset"),
-    seed: int = 1234
+    seed: int = 1234,
 ) -> Path:
     random.seed(seed)
     if out_dir.exists():
-        shutil.rmtree(out_dir, onerror=lambda f,p,e: (os.chmod(p, stat.S_IWRITE), f(p)))
+        shutil.rmtree(out_dir, onerror=lambda f, p, _: (os.chmod(p, stat.S_IWRITE), f(p)))
     for sub in ("train/images","train/labels","valid/images","valid/labels"):
-        (out_dir/sub).mkdir(parents=True, exist_ok=True)
+        (out_dir / sub).mkdir(parents=True, exist_ok=True)
 
-    # build mapping
-    mapping = {
-        row.original_class: ("__REMOVED__" if row.remove else row.new_name)
-        for row in class_map_df.itertuples()
+    class_name_mapping = {
+        row["original_class"]: row["new_name"] if not row["remove"] else "__REMOVED__"
+        for _, row in class_map_df.iterrows()
     }
-    limits = {
-        row.new_name: row.max_images
-        for row in class_map_df.itertuples() if not row.remove
+    limits_per_merged = {
+        row["new_name"]: int(row["max_images"])
+        for _, row in class_map_df.iterrows()
+        if not row["remove"]
     }
-    active = sorted({v for v in mapping.values() if v!="__REMOVED__"})
-    id_map = {c:i for i,c in enumerate(active)}
-
-    # collect image→classes
-    img2cls, img2lbl = {}, {}
-    cls2imgs = {c:set() for c in active}
-    for info in infos:
-        for split in info.splits:
-            for lp in Path(info.path)/split/'labels'.rglob('*.txt'):
-                origs = [info.class_names[int(cls)] for cls,*_ in parse_label_file(lp)]
-                newset = {mapping[o] for o in origs if mapping.get(o) in active}
-                if not newset: continue
-                img_path = str(lp.parent.parent/'images'/f"{lp.stem}.jpg")
-                img2cls[img_path] = newset
-                img2lbl[img_path] = lp
-                for c in newset:
-                    cls2imgs[c].add(img_path)
-
-    # select under class limits
-    selected, counters = set(), {c:0 for c in active}
-    pool = list({img for imgs in cls2imgs.values() for img in imgs})
+    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)}
+
+    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}
+
+    for dloc, class_names_dataset, splits, _ in dataset_info_list:
+        for split in splits:
+            labels_root = Path(dloc) / split / "labels"
+            if not labels_root.exists():
+                continue
+            for lp in labels_root.rglob("*.txt"):
+                cls_set: set[str] = set()
+                for cls_id_float, *rest in parse_label_file(lp):
+                    idx = int(cls_id_float)
+                    if 0 <= idx < len(class_names_dataset):
+                        orig = class_names_dataset[idx]
+                        new = class_name_mapping.get(orig, orig)
+                        if new in active_classes:
+                            cls_set.add(new)
+                if not cls_set:
+                    continue
+                img_path = str(lp.parent.parent / "images" / f"{lp.stem}.jpg")
+                image_to_classes[img_path] = cls_set
+                image_to_label[img_path] = lp
+                for c in cls_set:
+                    class_to_images[c].add(img_path)
+
+    selected_images: set[str] = set()
+    counters = {c: 0 for c in active_classes}
+    pool = [img for imgs in class_to_images.values() for img in imgs]
     random.shuffle(pool)
+
     for img in pool:
-        cs = img2cls[img]
-        if any(counters[c]>=limits[c] for c in cs): continue
-        selected.add(img)
-        for c in cs: counters[c]+=1
-
-    # copy & rewrite labels
-    for img in selected:
-        split = "train" if random.random()<0.9 else "valid"
-        dst_im = out_dir/split/'images'/Path(img).name
-        dst_im.parent.mkdir(parents=True, exist_ok=True)
-        shutil.copy(img, dst_im)
-        lp = img2lbl[img]
-        dst_lbl = out_dir/split/'labels'/lp.name
+        cs = image_to_classes[img]
+        if any(counters[c] >= limits_per_merged.get(c, 0) for c in cs):
+            continue
+        selected_images.add(img)
+        for c in cs:
+            counters[c] += 1
+
+    for img in selected_images:
+        split = "train" if random.random() < 0.9 else "valid"
+        dst_img = out_dir / split / "images" / Path(img).name
+        dst_img.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy(img, dst_img)
+
+        lp_src = image_to_label[img]
+        dst_lbl = out_dir / split / "labels" / lp_src.name
         dst_lbl.parent.mkdir(parents=True, exist_ok=True)
-        new = []
-        for line in lp.read_text().splitlines():
+        lines = lp_src.read_text().splitlines()
+        new_lines: List[str] = []
+        for line in lines:
             parts = line.split()
             cid = int(parts[0])
-            orig = info.class_names[cid] if cid<len(info.class_names) else None
-            merged = mapping.get(orig)
-            if merged in active:
-                new.append(" ".join([str(id_map[merged])]+parts[1:]))
-        if new:
-            dst_lbl.write_text("\n".join(new))
+            orig = None
+            # find which dataset tuple this lp_src belongs to, to get class_names_dataset
+            for dloc, class_names_dataset, splits, _ in dataset_info_list:
+                if str(lp_src).startswith(dloc):
+                    orig = class_names_dataset[cid] if cid < len(class_names_dataset) else None
+                    break
+            merged = class_name_mapping.get(orig, orig) if orig else None
+            if merged and merged in active_classes:
+                new_id = id_map[merged]
+                new_lines.append(" ".join([str(new_id)] + parts[1:]))
+        if new_lines:
+            dst_lbl.write_text("\n".join(new_lines))
         else:
-            dst_im.unlink(missing_ok=True)
+            dst_img.unlink(missing_ok=True)
 
-    # write data.yaml
-    meta = {
+    data_yaml = {
         "path": str(out_dir.resolve()),
         "train": "train/images",
-        "val":   "valid/images",
-        "nc":    len(active),
-        "names": active
+        "val": "valid/images",
+        "nc": len(active_classes),
+        "names": active_classes,
     }
-    (out_dir/"data.yaml").write_text(yaml.safe_dump(meta))
+    (out_dir / "data.yaml").write_text(yaml.safe_dump(data_yaml))
     return out_dir
 
 # ════════════════════════════════════════════════════════════════════════════
-#                                 UI Layer
+#                                 UI LAYER
 # ════════════════════════════════════════════════════════════════════════════
 with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
-    gr.Markdown("# 🏹 **YOLO Dataset Toolkit**\n_Evaluate • Merge • Edit • Download_")
+    gr.Markdown("""
+    # 🏹 **YOLO Dataset Toolkit**  
+    _Evaluate • Merge • Edit • Download_
+    """)
 
-    # Evaluate Tab
+    # Evaluate Tab ...
     with gr.Tab("Evaluate"):
-        api_in      = gr.Textbox(label="Roboflow API key", type="password")
-        url_txt     = gr.File(label=".txt of RF dataset URLs", file_types=['.txt'])
-        zip_in      = gr.File(label="Dataset ZIP")
-        path_in     = gr.Textbox(label="Server path")
-        yaml_in     = gr.File(label="Custom YAML", file_types=['.yaml'])
-        weights_in  = gr.File(label="YOLO weights (.pt)")
-        blur_sl     = gr.Slider(0.0,500.0,value=100.0,label="Blur threshold")
-        iou_sl      = gr.Slider(0.0,1.0,value=0.5,label="IOU threshold")
-        conf_sl     = gr.Slider(0.0,1.0,value=0.25,label="Min detection confidence")
-        run_dup     = gr.Checkbox(label="Check duplicates", value=False)
-        run_modelqa = gr.Checkbox(label="Run Model QA & cleanlab", value=False)
-        run_eval    = gr.Button("Run Evaluation")
-        out_md      = gr.Markdown()
-        out_df      = gr.Dataframe()
+        api_in     = gr.Textbox(label="Roboflow API key", type="password")
+        url_txt    = gr.File(label=".txt of RF dataset URLs", file_types=['.txt'])
+        zip_in     = gr.File(label="Dataset ZIP")
+        path_in    = gr.Textbox(label="Server path")
+        yaml_in    = gr.File(label="Custom YAML", file_types=['.yaml'])
+        weights_in = gr.File(label="YOLO weights (.pt)")
+
+        blur_sl    = gr.Slider(0.0, 500.0, value=100.0, label="Blur threshold")
+        iou_sl     = gr.Slider(0.0, 1.0,   value=0.5,   label="IOU threshold")
+        conf_sl    = gr.Slider(0.0, 1.0,   value=0.25,  label="Min detection confidence")
+
+        run_dup    = gr.Checkbox(label="Check duplicates (fastdup)", value=False)
+        run_modelqa= gr.Checkbox(label="Run Model QA & cleanlab", value=False)
+
+        run_eval   = gr.Button("Run Evaluation")
+        out_md     = gr.Markdown()
+        out_df     = gr.Dataframe()
 
         def _evaluate_cb(
-            api_key, url_file, zip_file, server_path, yaml_file, weights,
-            blur_thr, iou_thr, conf_thr, dup, modelqa
+            api_key, url_txt, zip_file, server_path, yaml_file, weights,
+            blur_thr, iou_thr, conf_thr, run_dup, run_modelqa
         ):
             reports, dfs = [], []
-            cfg = QCConfig(blur_thr,iou_thr,conf_thr, weights.name if weights else None)
-            rf  = Roboflow(api_key) if api_key and Roboflow else None
+            cfg = QCConfig(blur_thr, iou_thr, conf_thr, weights.name if weights else None)
+            rf = Roboflow(api_key) if api_key and Roboflow else None
 
-            # Roboflow URLs
-            if url_file and rf:
-                for line in Path(url_file.name).read_text().splitlines():
+            if url_txt and rf:
+                for line in Path(url_txt.name).read_text().splitlines():
                     if not line.strip(): continue
                     try:
                         ds = download_rf_dataset(line, rf, TMP_ROOT)
-                        md, df = run_quality(Path(ds), None, Path(weights.name) if weights else None,
-                                              cfg, dup, modelqa)
+                        md, df = run_quality(
+                            ds, None,
+                            Path(weights.name) if weights else None,
+                            cfg, run_dup, run_modelqa
+                        )
                         reports.append(md); dfs.append(df)
                     except Exception as e:
                         reports.append(f"### {line}\n⚠️ {e}")
 
-            # ZIP upload
             if zip_file:
                 tmp = Path(tempfile.mkdtemp())
                 shutil.unpack_archive(zip_file.name, tmp)
-                md, df = run_quality(tmp,
-                                     Path(yaml_file.name) if yaml_file else None,
-                                     Path(weights.name) if weights else None,
-                                     cfg, dup, modelqa)
+                md, df = run_quality(
+                    tmp,
+                    Path(yaml_file.name) if yaml_file else None,
+                    Path(weights.name) if weights else None,
+                    cfg, run_dup, run_modelqa
+                )
                 reports.append(md); dfs.append(df)
                 shutil.rmtree(tmp, ignore_errors=True)
 
-            # Server path
             if server_path:
-                md, df = run_quality(Path(server_path),
-                                     Path(yaml_file.name) if yaml_file else None,
-                                     Path(weights.name) if weights else None,
-                                     cfg, dup, modelqa)
+                ds = Path(server_path)
+                md, df = run_quality(
+                    ds,
+                    Path(yaml_file.name) if yaml_file else None,
+                    Path(weights.name) if weights else None,
+                    cfg, run_dup, run_modelqa
+                )
                 reports.append(md); dfs.append(df)
 
             summary  = "\n---\n".join(reports) if reports else ""
@@ -535,24 +554,24 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
 
         run_eval.click(
             _evaluate_cb,
-            inputs=[api_in,url_txt,zip_in,path_in,yaml_in,weights_in,
-                    blur_sl,iou_sl,conf_sl,run_dup,run_modelqa],
-            outputs=[out_md,out_df]
+            inputs=[api_in, url_txt, zip_in, path_in, yaml_in, weights_in,
+                    blur_sl, iou_sl, conf_sl, run_dup, run_modelqa],
+            outputs=[out_md, out_df]
         )
 
-    # Merge/Edit Tab
+    # Merge / Edit Tab
     with gr.Tab("Merge / Edit"):
-        gr.Markdown("### 1️⃣ Load datasets")
-        rf_key  = gr.Textbox(label="Roboflow API key", type="password")
-        rf_urls = gr.File(label=".txt of RF URLs", file_types=['.txt'])
-        zips    = gr.Files(label="Dataset ZIPs")
-        load_btn= gr.Button("Load")
-        load_log= gr.Markdown()
-        ds_state= gr.State([])
+        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")
+        load_btn = gr.Button("Load datasets")
+        load_log = gr.Markdown()
+        ds_state = gr.State([])
 
         def _load_cb(rf_key, rf_urls_file, zip_files):
-            nonlocal autoinc
-            infos, logs = [], []
+            global autoinc
+            info_list, log_lines = [], []
             rf = Roboflow(rf_key) if rf_key and Roboflow else None
 
             if rf_urls_file and rf:
@@ -561,61 +580,62 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
                     if not url: continue
                     try:
                         ds = download_rf_dataset(url, rf, TMP_ROOT)
-                        names = load_class_names(Path(ds)/"data.yaml")
-                        splits= [s for s in ("train","valid","test") if (Path(ds)/s).exists()]
-                        infos.append(DatasetInfo(str(ds),names,splits,Path(ds).name))
-                        logs.append(f"✔️ Loaded RF dataset {Path(ds).name}")
+                        names = load_class_names(ds/"data.yaml")
+                        splits = [s for s in ("train","valid","test") if (ds/s).exists()]
+                        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:
-                        logs.append(f"⚠️ RF failed {url!r}: {e}")
+                        log_lines.append(f"⚠️ RF load failed for {url!r}: {e}")
 
             for f in zip_files or []:
                 autoinc += 1
-                tmp = TMP_ROOT/f"zip_{autoinc}"
+                tmp = TMP_ROOT / f"zip_{autoinc}"
                 tmp.mkdir(parents=True, exist_ok=True)
                 shutil.unpack_archive(f.name, tmp)
-                y = next(tmp.rglob("*.yaml"), None)
-                if y:
-                    names = load_class_names(y)
-                    splits= [s for s in ("train","valid","test") if (tmp/s).exists()]
-                    infos.append(DatasetInfo(str(tmp),names,splits,tmp.name))
-                    logs.append(f"✔️ Loaded ZIP {tmp.name}")
+                yaml_p = next(tmp.rglob("*.yaml"), None)
+                if yaml_p:
+                    names = load_class_names(yaml_p)
+                    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 infos, "\n".join(logs) or "No datasets loaded."
+            return info_list, "\n".join(log_lines) or "No datasets loaded."
 
-        load_btn.click(_load_cb, [rf_key, rf_urls, zips], [ds_state, load_log])
+        load_btn.click(_load_cb, [rf_key, rf_urls, zips_in], [ds_state, load_log])
 
-        gr.Markdown("### 2️⃣ Edit classes")
+        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   = gr.Button("Build table")
-        merge_btn = gr.Button("Merge")
-        zip_out   = gr.File(label="Download merged ZIP")
-        merge_log = gr.Markdown()
+        refresh_btn = gr.Button("Build class table from loaded datasets")
 
-        def _build_df(infos):
-            counts = gather_class_counts(infos)
-            names  = sorted(counts)
+        def _build_class_df(ds_info):
+            counts = gather_class_counts(ds_info)
+            all_names = sorted(counts.keys())
             return pd.DataFrame({
-                "original_class": names,
-                "new_name":       names,
-                "max_images":     [counts[n] for n in names],
-                "remove":         [False]*len(names),
+                "original_class": all_names,
+                "new_name":       all_names,
+                "max_images":     [counts[n] for n in all_names],
+                "remove":         [False]*len(all_names),
             })
 
-        refresh.click(_build_df, [ds_state], [class_df])
+        refresh_btn.click(_build_class_df, [ds_state], [class_df])
+
+        merge_btn = gr.Button("Merge datasets ✨")
+        zip_out   = gr.File(label="Download merged ZIP")
+        merge_log = gr.Markdown()
 
-        def _merge_cb(infos, df):
-            if not infos:
-                return None, "⚠️ Load datasets first."
-            out = merge_datasets(infos, df)
-            zipf = shutil.make_archive(str(out), "zip", out)
-            cnt  = len(list(Path(out).rglob("*.jpg")))
-            return zipf, f"✅ Merged to {out} ({cnt} images)"
+        def _merge_cb(ds_info, class_df):
+            if not ds_info:
+                return None, "⚠️ Load datasets first."
+            out_dir = merge_datasets(ds_info, class_df)
+            zip_path = shutil.make_archive(str(out_dir), "zip", out_dir)
+            count = len(list(Path(out_dir).rglob("*.jpg")))
+            return zip_path, f"✅ Merged dataset at **{out_dir}** with {count} images."
 
         merge_btn.click(_merge_cb, [ds_state, class_df], [zip_out, merge_log])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT",7860)))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))