Update app.py

app.py

@@ -1,15 +1,14 @@
-"""
-app.py – Roboflow‑aware YOLOv8 Dataset Quality Evaluator (v2)
-
-Changelog (2025‑04‑17)
-──────────────────────
-• **CPU‑bound loops parallelised** with `concurrent.futures.ProcessPoolExecutor`.
-• **Batch inference** in `qc_model_qa()` (GPU util ↑, latency ↓).
-• Optional **fastdup** path for duplicate detection (≈ 10× faster on large sets).
-• Faster NumPy‑based `parse_label_file()`.
-• Small refactors → clearer separation of stages & fewer globals.
-• Graceful degradation if heavy deps unavailable (cv2, imagehash, fastdup).
-• Tunable `CPU_COUNT` + env‑var guard for HF Spaces quota.
+ """
+ app.py – Roboflow‑aware YOLOv8 Dataset Quality Evaluator (v3)
+
+ Changelog (2025‑04‑17)
+ ──────────────────────
+ • **YOLO model caching** to avoid repeated loads
+ • **Exposed thresholds** (blur, IOU, confidence) as Gradio sliders
+ • **Config dataclass** for unified settings
+ • **Cleanlab integration** for label-issue detection
+ • **Parallel label parsing** and expanded caching
+ • **Adjusted QC weights** to include Label Issues stage
 """
 
 from __future__ import annotations
@@ -33,62 +32,65 @@ import yaml
 from PIL import Image
 from tqdm import tqdm
 
-# ───────────────────────────────────────── Heavy optional deps ──
+# Optional heavy deps
 try:
-    import cv2  # type: ignore
+    import cv2
 except ImportError:
     cv2 = None
-
 try:
-    import imagehash  # type: ignore
+    import imagehash
 except ImportError:
     imagehash = None
-
 try:
-    import fastdup  # type: ignore
+    import fastdup
 except ImportError:
     fastdup = None
-
 try:
-    from ultralytics import YOLO  # type: ignore
+    from ultralytics import YOLO
 except ImportError:
-    YOLO = None  # noqa: N806
-
+    YOLO = None
 try:
-    from roboflow import Roboflow  # type: ignore
+    from roboflow import Roboflow
 except ImportError:
-    Roboflow = None  # type: ignore
+    Roboflow = None
+try:
+    from cleanlab.pruning import get_noise_indices
+except ImportError:
+    get_noise_indices = None
 
-# ───────────────────────────────────────── Config & constants ──
+# ───────────────── Config & Constants ─────────────────
 TMP_ROOT = Path(tempfile.gettempdir()) / "rf_datasets"
 TMP_ROOT.mkdir(parents=True, exist_ok=True)
-
-# Limit CPU workers on HF Spaces (feel free to raise locally)
 CPU_COUNT = int(os.getenv("QC_CPU", max(1, (os.cpu_count() or 4) // 2)))
-BATCH = int(os.getenv("QC_BATCH", 16))
+BATCH_SIZE = int(os.getenv("QC_BATCH", 16))
 
 DEFAULT_W = {
-    "Integrity": 0.30,
-    "Class balance": 0.15,
-    "Image quality": 0.15,
-    "Duplicates": 0.10,
-    "Model QA": 0.30,
+    "Integrity":    0.25,
+    "Class balance":0.10,
+    "Image quality":0.15,
+    "Duplicates":   0.10,
+    "Model QA":     0.30,
+    "Label issues": 0.10,
 }
 
-@dataclass
-class DuplicateGroup:
-    hash_val: str
-    paths: List[Path]
-
-# ───────────────────────────────────────── Generic helpers ─────
+_model_cache: dict[str, YOLO] = {}
 
+@dataclass
+class QCConfig:
+    blur_thr: float
+    iou_thr: float
+    conf_thr: float
+    weights: str | None
+    cpu_count: int = CPU_COUNT
+    batch_size: int = BATCH_SIZE
+
+# ─────────────────── Helpers & Caching ───────────────────
 def load_yaml(path: Path) -> Dict:
-    with path.open(encoding="utf-8") as f:
+    with path.open('r', encoding='utf-8') as f:
         return yaml.safe_load(f)
 
-
 def parse_label_file(path: Path) -> list[tuple[int, float, float, float, float]]:
-    if not path.exists() or path.stat().st_size == 0:
+    if not path or not path.exists() or path.stat().st_size == 0:
         return []
     try:
         arr = np.loadtxt(path, dtype=float)
@@ -98,364 +100,285 @@ def parse_label_file(path: Path) -> list[tuple[int, float, float, float, float]]
     except Exception:
         return []
 
-
 def guess_image_dirs(root: Path) -> List[Path]:
-    subs = [
-        root / "images",
-        root / "train" / "images",
-        root / "valid" / "images",
-        root / "val" / "images",
-        root / "test" / "images",
-    ]
+    subs = [root / 'images', root / 'train' / 'images', root / 'valid' / 'images',
+            root / 'val' / 'images', root / 'test' / 'images']
     return [d for d in subs if d.exists()]
 
-
-def gather_dataset(root: Path, yaml_path: Path | None = None):
+def gather_dataset(root: Path, yaml_path: Path | None):
     if yaml_path is None:
-        yamls = list(root.glob("*.yaml"))
+        yamls = list(root.glob('*.yaml'))
         if not yamls:
             raise FileNotFoundError("Dataset YAML not found")
         yaml_path = yamls[0]
-
     meta = load_yaml(yaml_path)
     img_dirs = guess_image_dirs(root)
     if not img_dirs:
-        raise FileNotFoundError("images/ directory hierarchy missing")
-
-    imgs = [p for d in img_dirs for p in d.rglob("*.*") if imghdr.what(p) is not None]
-    labels_root = {d.parent / "labels" for d in img_dirs}
-    lbls = [next((lr / f"{p.stem}.txt" for lr in labels_root if (lr / f"{p.stem}.txt").exists()), None) for p in imgs]
+        raise FileNotFoundError("images/ directory missing")
+    imgs = [p for d in img_dirs for p in d.rglob('*.*') if imghdr.what(p)]
+    labels_root = {d.parent/'labels' for d in img_dirs}
+    lbls = [next((lr/f"{p.stem}.txt" for lr in labels_root if (lr/f"{p.stem}.txt").exists()), None)
+            for p in imgs]
     return imgs, lbls, meta
 
-# ───────────────────────────────────────── Quality checks ─────
-# Integrity -----------------------------------------------------
+# YOLO model caching
+
+def get_model(weights: str) -> YOLO | None:
+    if weights is None or YOLO is None:
+        return None
+    if weights not in _model_cache:
+        _model_cache[weights] = YOLO(weights)
+    return _model_cache[weights]
+
+# ───────────────────── Quality Checks ─────────────────────
 
 def _is_corrupt(path: Path) -> bool:
     try:
         with Image.open(path) as im:
             im.verify()
         return False
-    except Exception:
+    except:
         return True
 
-
-def qc_integrity(imgs: List[Path], lbls: List[Path]):
-    miss_lbl = [i for i, l in zip(imgs, lbls) if l is None]
-    corrupt: List[Path] = []
-    with ProcessPoolExecutor(max_workers=CPU_COUNT) as ex:
-        fut = {ex.submit(_is_corrupt, p): p for p in imgs}
-        for f in tqdm(as_completed(fut), total=len(fut), desc="integrity", leave=False):
-            if f.result():
-                corrupt.append(fut[f])
-
-    score = 100 - (len(miss_lbl) + len(corrupt)) / max(len(imgs), 1) * 100
-    return {
-        "name": "Integrity",
-        "score": max(score, 0),
-        "details": {
-            "missing_label_files": [str(p) for p in miss_lbl],
-            "corrupt_images": [str(p) for p in corrupt],
-        },
-    }
-
-# Class balance -------------------------------------------------
-
-def qc_class_balance(lbls: List[Path]):
-    cls_counts = Counter()
-    boxes_per_img = []
+def qc_integrity(imgs: List[Path], lbls: List[Path], cfg: QCConfig):
+    miss = [i for i,l in zip(imgs,lbls) if l is None]
+    corrupt = []
+    with ProcessPoolExecutor(max_workers=cfg.cpu_count) as ex:
+        fut = {ex.submit(_is_corrupt,p):p for p in imgs}
+        for f in as_completed(fut):
+            if f.result(): corrupt.append(fut[f])
+    score = 100 - (len(miss)+len(corrupt))/max(len(imgs),1)*100
+    return {"name":"Integrity","score":max(score,0),
+            "details":{"missing_label_files":[str(p) for p in miss],
+                        "corrupt_images":[str(p) for p in corrupt]}}
+
+def qc_class_balance(lbls: List[Path], cfg: QCConfig):
+    counts=Counter(); boxes=[]
     for l in lbls:
-        bs = parse_label_file(l) if l else []
-        boxes_per_img.append(len(bs))
-        cls_counts.update(b[0] for b in bs)
-
-    if not cls_counts:
-        return {"name": "Class balance", "score": 0, "details": "No labels"}
-    bal = (min(cls_counts.values()) / max(cls_counts.values())) * 100
-    return {
-        "name": "Class balance",
-        "score": bal,
-        "details": {
-            "class_counts": dict(cls_counts),
-            "boxes_per_image": {
-                "min": int(np.min(boxes_per_img)),
-                "max": int(np.max(boxes_per_img)),
-                "mean": float(np.mean(boxes_per_img)),
-            },
-        },
-    }
-
-# Image quality -------------------------------------------------
-
-def _quality_stat(path: Path, blur_thr: float):
-    im = cv2.imread(str(path)) if cv2 else None
-    if im is None:
-        return path, False, False, False
-    gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
-    lap = cv2.Laplacian(gray, cv2.CV_64F).var()
-    br = gray.mean()
-    return path, lap < blur_thr, br < 25, br > 230
-
-
-def qc_image_quality(imgs: List[Path], blur_thr: float = 100.0):
+        bs=parse_label_file(l) if l else []
+        boxes.append(len(bs)); counts.update(b[0] for b in bs)
+    if not counts:
+        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":int(np.min(boxes)),
+                                           "max":int(np.max(boxes)),
+                                           "mean":float(np.mean(boxes))}}}
+
+def _quality_stat(path:Path, blur_thr:float):
+    if cv2 is None: return path,False,False,False
+    im=cv2.imread(str(path));
+    gray=cv2.cvtColor(im,cv2.COLOR_BGR2GRAY)
+    lap=cv2.Laplacian(gray,cv2.CV_64F).var(); br=gray.mean()
+    return path, lap<blur_thr, br<25, br>230
+
+def qc_image_quality(imgs:List[Path], cfg:QCConfig):
     if cv2 is None:
-        return {"name": "Image quality", "score": 100, "details": "cv2 not installed"}
-
-    blurry: list[Path] = []
-    dark: list[Path] = []
-    bright: list[Path] = []
-
-    with ProcessPoolExecutor(max_workers=CPU_COUNT) as ex:
-        for p, is_blur, is_dark, is_bright in tqdm(
-            ex.map(lambda x: _quality_stat(x, blur_thr), imgs),
-            total=len(imgs),
-            desc="img‑quality",
-            leave=False,
-        ):
-            if is_blur:
-                blurry.append(p)
-            if is_dark:
-                dark.append(p)
-            if is_bright:
-                bright.append(p)
-
-    bad = len(set(blurry + dark + bright))
-    score = 100 - bad / max(len(imgs), 1) * 100
-    return {
-        "name": "Image quality",
-        "score": score,
-        "details": {
-            "blurry": [str(p) for p in blurry],
-            "dark": [str(p) for p in dark],
-            "bright": [str(p) for p in bright],
-        },
-    }
-
-# Duplicate images ---------------------------------------------
-
-def qc_duplicates(imgs: List[Path]):
-    # Fast path – use fastdup if installed & enough images
-    if fastdup is not None and len(imgs) > 50:
+        return {"name":"Image quality","score":100,"details":"cv2 missing"}
+    blurry, dark, bright = [],[],[];
+    with ProcessPoolExecutor(max_workers=cfg.cpu_count) as ex:
+        for p,isb,isd,isB in tqdm(ex.map(lambda x: _quality_stat(x,cfg.blur_thr), imgs),
+                                  total=len(imgs), desc='img-quality', leave=False):
+            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(imgs),1)*100
+    return {"name":"Image quality","score":score,
+            "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):
+    if fastdup and len(imgs)>50:
         try:
-            fd = fastdup.create(input_dir=str(Path(imgs[0]).parent.parent), work_dir=str(TMP_ROOT / "fastdup"))
-            fd.run()
-            clusters = fd.get_clusters()
-            dup = sum(len(c) - 1 for c in clusters)
-            score = 100 - dup / max(len(imgs), 1) * 100
-            return {
-                "name": "Duplicates",
-                "score": score,
-                "details": {"groups": clusters[:50]},
-            }
-        except Exception:
-            pass  # fallback to hash
-
+            fd=fastdup.create(input_dir=str(Path(imgs[0]).parent.parent),
+                              work_dir=str(TMP_ROOT/'fastdup'))
+            fd.run(); clusters=fd.get_clusters()
+            dup=sum(len(c)-1 for c in clusters)
+            return {"name":"Duplicates","score":100-dup/len(imgs)*100,
+                    "details":{"groups":clusters[:50]}}
+        except: pass
     if imagehash is None:
-        return {"name": "Duplicates", "score": 100, "details": "skipped (deps)"}
-
-    def _hash(p):
-        return str(imagehash.average_hash(Image.open(p)))
-
-    hashes: Dict[str, List[Path]] = defaultdict(list)
-    with ProcessPoolExecutor(max_workers=CPU_COUNT) as ex:
-        for h, p in tqdm(
-            zip(ex.map(_hash, imgs), imgs),
-            total=len(imgs),
-            desc="hashing",
-            leave=False,
-        ):
+        return {"name":"Duplicates","score":100,"details":"deps missing"}
+    hashes=defaultdict(list)
+    with ProcessPoolExecutor(max_workers=cfg.cpu_count) as ex:
+        for h,p in zip(ex.map(lambda x: str(imagehash.average_hash(Image.open(x))),imgs), imgs):
             hashes[h].append(p)
-
-    groups = [g for g in hashes.values() if len(g) > 1]
-    dup = sum(len(g) - 1 for g in groups)
-    score = 100 - dup / max(len(imgs), 1) * 100
-    return {
-        "name": "Duplicates",
-        "score": score,
-        "details": {"groups": [[str(p) for p in g] for g in groups[:50]]},
-    }
-
-# Model‑assisted QA --------------------------------------------
-
-def _rel_iou(b1, b2):
-    x1, y1, w1, h1 = b1
-    x2, y2, w2, h2 = b2
-    xa1, ya1, xa2, ya2 = x1 - w1 / 2, y1 - h1 / 2, x1 + w1 / 2, y1 + h1 / 2
-    xb1, yb1, xb2, yb2 = x2 - w2 / 2, y2 - h2 / 2, x2 + w2 / 2, y2 + h2 / 2
-    ix1, iy1, ix2, iy2 = max(xa1, xb1), max(ya1, yb1), min(xa2, xb2), 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], weights: str | None, lbls: List[Path], iou_thr: float = 0.5):
-    if weights is None or YOLO is None:
-        return {"name": "Model QA", "score": 100, "details": "skipped (no weights)"}
-
-    model = YOLO(weights)
+    groups=[g for g in hashes.values() if len(g)>1]
+    dup=sum(len(g)-1 for g in groups)
+    return {"name":"Duplicates","score":100-dup/len(imgs)*100,
+            "details":{"groups":[[str(p) for p in g] for g in groups[:50]]}}
+
+def _rel_iou(b1,b2):
+    x1,y1,w1,h1=b1; x2,y2,w2,h2=b2
+    xa1,ya1,xa2,ya2=x1-w1/2,y1-h1/2,x1+w1/2,y1+h1/2
+    xb1,yb1,xb2,yb2=x2-w2/2,y2-h2/2,x2+w2/2,y2+h2/2
+    ix1,iy1,ix2,iy2=max(xa1,xb1),max(ya1,yb1),min(xa2,xb2),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[Path], cfg:QCConfig):
+    model=get_model(cfg.weights)
+    if model is None:
+        return {"name":"Model QA","score":100,"details":"skipped"}
     ious, mism = [], []
-
-    for i in range(0, len(imgs), BATCH):
-        batch_paths = imgs[i : i + BATCH]
-        results = model.predict(batch_paths, verbose=False)
-        for p, res in zip(batch_paths, results):
-            gtb = parse_label_file(p.parent.parent / "labels" / f"{p.stem}.txt")
-            if not gtb:
-                continue
-            for cls, x, y, w, h in gtb:
-                best = 0.0
-                for b, c in zip(res.boxes.xywh.cpu().numpy(), res.boxes.cls.cpu().numpy()):
-                    if int(c) != cls:
-                        continue
-                    best = max(best, _rel_iou((x, y, w, h), tuple(b)))
+    for i in range(0,len(imgs),cfg.batch_size):
+        batch=imgs[i:i+cfg.batch_size]
+        results=model.predict(batch, verbose=False)
+        for p,res in zip(batch,results):
+            gt=parse_label_file(lbls[imgs.index(p)])
+            if not gt: continue
+            preds = res.boxes.xywh.cpu().numpy()
+            confs = res.boxes.conf.cpu().numpy()
+            classes = res.boxes.cls.cpu().numpy()
+            mask = confs >= cfg.conf_thr
+            preds, classes = preds[mask], classes[mask]
+            for cls,x,y,w,h in gt:
+                best=0.0
+                for b,c in zip(preds,classes):
+                    if int(c)!=cls: continue
+                    best=max(best,_rel_iou((x,y,w,h),tuple(b)))
                 ious.append(best)
-                if best < iou_thr:
+                if best < cfg.iou_thr:
                     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, "mismatched_images": mism[:50]},
-    }
-
-# Aggregate -----------------------------------------------------
-
+    miou=float(np.mean(ious)) if ious else 1.0
+    return {"name":"Model QA","score":miou*100,
+            "details":{"mean_iou":miou,"mismatches":mism[:50]}}
+
+def qc_label_issues(imgs:List[Path], lbls:List[Path], cfg:QCConfig):
+    if get_noise_indices is None or cfg.weights is None:
+        return {"name":"Label issues","score":100,"details":"skipped"}
+    model=get_model(cfg.weights)
+    if model is None:
+        return {"name":"Label issues","score":100,"details":"skipped"}
+    labels,preds,samps = [],[],[]
+    for i in range(0,len(imgs),cfg.batch_size):
+        batch=imgs[i:i+cfg.batch_size]
+        results=model.predict(batch, verbose=False)
+        for p,res in zip(batch,results):
+            gt=parse_label_file(lbls[imgs.index(p)])
+            for cls,x,y,w,h in gt:
+                labels.append(int(cls))
+                # find predicted class with highest IoU
+                best_i, best_c = 0.0, -1
+                for b,c in zip(res.boxes.xywh.cpu().numpy(), res.boxes.cls.cpu().numpy()):
+                    iou=_rel_iou((x,y,w,h),tuple(b))
+                    if iou>best_i:
+                        best_i, best_c = iou, int(c)
+                preds.append(best_c)
+                samps.append(p)
+    if not labels:
+        return {"name":"Label issues","score":100,"details":"no GT"}
+    noise_idx = get_noise_indices(np.array(labels), np.array(preds))
+    sus = list({str(samps[i]) for i in noise_idx})[:50]
+    score = 100 - len(noise_idx)/len(labels)*100
+    return {"name":"Label issues","score":score,
+            "details":{"suspect_images": sus}}
+
+# ─────────────────────── Aggregate & Run ──────────────────────
 def aggregate(scores):
-    return sum(DEFAULT_W.get(r["name"], 0) * r["score"] for r in scores)
+    return sum(DEFAULT_W.get(r['name'],0)*r['score'] for r in scores)
 
-# ───────────────────────────────────────── Roboflow helpers ────
 RF_RE = re.compile(r"https://universe\.roboflow\.com/([^/]+)/([^/]+)/dataset/(\d+)")
-
-def download_rf_dataset(url: str, rf_api: "Roboflow", dest: Path) -> Path:
-    m = RF_RE.match(url.strip())
-    if not m:
-        raise ValueError(f"Bad RF URL: {url}")
-
-    ws, proj, ver = m.groups()
-    ds_dir = dest / f"{ws}_{proj}_v{ver}"
-    if ds_dir.exists():
-        return ds_dir
-
-    project = rf_api.workspace(ws).project(proj)
-    project.version(int(ver)).download("yolov8", location=str(ds_dir))
-    return ds_dir
-
-# ───────────────────────────────────────── Main logic ──────────
-
-def run_quality(root: Path, yaml_override: Path | None, weights: Path | None):
-    imgs, lbls, meta = gather_dataset(root, yaml_override)
-    res = [
-        qc_integrity(imgs, lbls),
-        qc_class_balance(lbls),
-        qc_image_quality(imgs),
-        qc_duplicates(imgs),
-        qc_model_qa(imgs, str(weights) if weights else None, lbls),
+def download_rf_dataset(url:str, rf_api:Roboflow, dest:Path)->Path:
+    m=RF_RE.match(url.strip());
+    if not m: raise ValueError(f"Bad RF URL: {url}")
+    ws,proj,ver = m.groups()
+    ds = dest/f"{ws}_{proj}_v{ver}"
+    if ds.exists(): return ds
+    proj_obj = rf_api.workspace(ws).project(proj)
+    proj_obj.version(int(ver)).download('yolov8', location=str(ds))
+    return ds
+
+def run_quality(root:Path, yaml_override:Path|None, lbls:List[Path], imgs:List[Path], cfg:QCConfig):
+    res=[
+        qc_integrity(imgs,lbls,cfg),
+        qc_class_balance(lbls,cfg),
+        qc_image_quality(imgs,cfg),
+        qc_duplicates(imgs,cfg),
+        qc_model_qa(imgs,lbls,cfg),
+        qc_label_issues(imgs,lbls,cfg),
     ]
-    final = aggregate(res)
-
-    md = [f"## **{meta.get('name', root.name)}** — Score {final:.1f}/100"]
+    final=aggregate(res)
+    md=[f"## **{root.name}** — Score {final:.1f}/100"]
     for r in res:
-        md.append(f"### {r['name']}  {r['score']:.1f}")
-        md.append("<details><summary>details</summary>\n\n```json")
-        md.append(json.dumps(r["details"], indent=2))
+        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")
-    md_str = "\n".join(md)
-
-    cls_counts = res[1]["details"].get("class_counts", {})  # type: ignore[index]
-    df = pd.DataFrame.from_dict(cls_counts, orient="index", columns=["count"])
-    df.index.name = "class"
-    return md_str, df
-
-# ───────────────────────────────────────── Gradio UI ───────────
-
-def evaluate(
-    api_key: str,
-    url_txt: gr.File | None,
-    zip_file: gr.File | None,
-    server_path: str,
-    yaml_file: gr.File | None,
-    weights: gr.File | None,
-):
-    if not any([url_txt, zip_file, server_path]):
-        return "Upload a .txt of URLs or dataset ZIP/path", pd.DataFrame()
-
-    reports, dfs = [], []
-
-    # Roboflow batch ------------------------------------------
-    if url_txt:
-        if Roboflow is None:
-            return "`roboflow` not installed", pd.DataFrame()
-        if not api_key:
-            return "Enter Roboflow API key", pd.DataFrame()
-
-        rf = Roboflow(api_key=api_key.strip())
-        for line in Path(url_txt.name).read_text().splitlines():
-            if not line.strip():
-                continue
-            try:
-                ds_root = download_rf_dataset(line, rf, TMP_ROOT)
-                md, df = run_quality(ds_root, None, Path(weights.name) if weights else None)
-                reports.append(md)
-                dfs.append(df)
-            except Exception as e:
-                reports.append(f"### {line}\n\n⚠️ {e}")
-
-    # Manual ZIP ----------------------------------------------
-    if zip_file:
-        tmp_dir = Path(tempfile.mkdtemp())
-        shutil.unpack_archive(zip_file.name, tmp_dir)
-        md, df = run_quality(tmp_dir, Path(yaml_file.name) if yaml_file else None, Path(weights.name) if weights else None)
-        reports.append(md)
-        dfs.append(df)
-        shutil.rmtree(tmp_dir, ignore_errors=True)
-
-    # Manual 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)
-        reports.append(md)
-        dfs.append(df)
-
-    summary_md = "\n\n---\n\n".join(reports)
-    combined_df = pd.concat(dfs).groupby(level=0).sum() if dfs else pd.DataFrame()
-    return summary_md, combined_df
-
-# ───────────────────────────────────────── Launch  ────────────
-with gr.Blocks(title="YOLO Dataset Quality Evaluator") as demo:
-    gr.Markdown(
-        """
-# YOLOv8 Dataset Quality Evaluator
-
-### Roboflow batch  
-1. Paste your **Roboflow API key**  
-2. Upload a **.txt** file – one `https://universe.roboflow.com/.../dataset/x` per line
-
-### Manual  
-* Upload a dataset **ZIP** or type a dataset **path** on the server  
-* Optionally supply a custom **data.yaml** and/or a **YOLO .pt** weights file for model‑assisted QA
-"""
+    df = pd.DataFrame.from_dict(
+        next(r for r in res if r['name']=='Class balance')['details']['class_counts'],
+        orient='index', columns=['count']
     )
+    df.index.name='class'
+    return "\n".join(md), df
 
-    with gr.Row():
-        api_in = gr.Textbox(label="Roboflow API key", type="password", placeholder="rf_XXXXXXXXXXXXXXXX")
-        url_txt_in = gr.File(label=".txt of RF dataset URLs", file_types=[".txt"])
+# ─────────────────────── Gradio UI ──────────────────────
+with gr.Blocks(title="YOLO Dataset Quality Evaluator v3") as demo:
+    gr.Markdown("""
+# YOLOv8 Dataset Quality Evaluator v3
 
+* Tweaks: blur, IOU & confidence sliders; Cleanlab label issues; model caching
+""")
     with gr.Row():
-        zip_in = gr.File(label="Dataset ZIP")
-        path_in = gr.Textbox(label="Path on server", placeholder="/data/my_dataset")
-
+        api_in = gr.Textbox(label="Roboflow API key", type="password")
+        url_txt = gr.File(label=".txt of RF dataset URLs", file_types=['.txt'])
+    with gr.Row():
+        zip_in  = gr.File(label="Dataset ZIP")
+        path_in = gr.Textbox(label="Server path")
     with gr.Row():
-        yaml_in = gr.File(label="Custom YAML", file_types=[".yaml"])
+        yaml_in    = gr.File(label="Custom YAML", file_types=['.yaml'])
         weights_in = gr.File(label="YOLO weights (.pt)")
-
+    with gr.Row():
+        blur_sl  = gr.Slider(0,500,value=100,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_btn = gr.Button("Evaluate")
-    out_md = gr.Markdown()
-    out_df = gr.Dataframe()
-
-    run_btn.click(
-        evaluate,
-        inputs=[api_in, url_txt_in, zip_in, path_in, yaml_in, weights_in],
-        outputs=[out_md, out_df],
-    )
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))
+    out_md  = gr.Markdown()
+    out_df  = gr.Dataframe()
+
+    def evaluate(api_key, url_txt, zip_file, server_path, yaml_file, weights,
+                 blur_thr, iou_thr, conf_thr):
+        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
+        # Roboflow batch
+        if url_txt:
+            for line in Path(url_txt.name).read_text().splitlines():
+                if not line.strip(): continue
+                try:
+                    ds = download_rf_dataset(line, rf, TMP_ROOT)
+                    imgs,lbls,_ = gather_dataset(ds,None)
+                    md, df = run_quality(ds,None,lbls,imgs,cfg)
+                    reports.append(md); dfs.append(df)
+                except Exception as e:
+                    reports.append(f"### {line}\n⚠️ {e}")
+        # ZIP
+        if zip_file:
+            tmp=Path(tempfile.mkdtemp())
+            shutil.unpack_archive(zip_file.name,tmp)
+            imgs,lbls,_=gather_dataset(tmp,Path(yaml_file.name) if yaml_file else None)
+            md,df=run_quality(tmp,None,lbls,imgs,cfg)
+            reports.append(md); dfs.append(df)
+            shutil.rmtree(tmp)
+        # Server path
+        if server_path:
+            ds=Path(server_path)
+            imgs,lbls,_=gather_dataset(ds,Path(yaml_file.name) if yaml_file else None)
+            md,df=run_quality(ds,None,lbls,imgs,cfg)
+            reports.append(md); dfs.append(df)
+        summary='\n---\n'.join(reports)
+        combined = pd.concat(dfs).groupby(level=0).sum() if dfs else pd.DataFrame()
+        return summary, combined
+
+    run_btn.click(evaluate,
+                  inputs=[api_in, url_txt, zip_in, path_in, yaml_in, weights_in,
+                          blur_sl, iou_sl, conf_sl],
+                  outputs=[out_md, out_df])
+
+if __name__ == '__main__':
+    demo.launch(server_name='0.0.0.0', server_port=int(os.getenv('PORT',7860)))