Update app.py

app.py

@@ -1,16 +1,3 @@
-"""
-app.py – Roboflow‑aware YOLOv8 Dataset Quality Evaluator (v3)
-
-Changelog (2025‑04‑17)
-──────────────────────
-• Fix RF URL regex to accept http/https
-• Top-level functions for parallel mapping (picklable)
-• Fastdup-only path in qc_duplicates (skips hashing fallback)
-• YOLO model caching
-• Config dataclass & Gradio sliders for blur, IOU, confidence
-• Cleanlab integration for label-issue detection
-"""
-
 from __future__ import annotations
 
 import imghdr
@@ -20,7 +7,7 @@ import re
 import shutil
 import tempfile
 from collections import Counter, defaultdict
-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, Tuple
@@ -30,7 +17,6 @@ import numpy as np
 import pandas as pd
 import yaml
 from PIL import Image
-from tqdm import tqdm
 
 # Optional heavy deps -------------------------------------------------------
 try:
@@ -61,8 +47,9 @@ except ImportError:
 # ───────────────── Config & Constants ───────────────────────────────────────
 TMP_ROOT = Path(tempfile.gettempdir()) / "rf_datasets"
 TMP_ROOT.mkdir(parents=True, exist_ok=True)
-CPU_COUNT = int(os.getenv("QC_CPU", max(1, (os.cpu_count() or 4) // 2)))
-BATCH_SIZE = int(os.getenv("QC_BATCH", 16))
+CPU_COUNT = int(os.getenv("QC_CPU", 1))       # force single-core by default
+BATCH_SIZE = int(os.getenv("QC_BATCH", 4))    # small batches
+SAMPLE_LIMIT = int(os.getenv("QC_SAMPLE", 200))
 
 DEFAULT_W = {
     "Integrity":     0.25,
@@ -83,6 +70,7 @@ class QCConfig:
     weights: str | None
     cpu_count: int = CPU_COUNT
     batch_size: int = BATCH_SIZE
+    sample_limit: int = SAMPLE_LIMIT
 
 # ─────────── Helpers & Caching ─────────────────────────────────────────────
 def load_yaml(path: Path) -> Dict:
@@ -128,7 +116,7 @@ def get_model(weights: str) -> YOLO | None:
         _model_cache[weights] = YOLO(weights)
     return _model_cache[weights]
 
-# ───────── Functions for parallel mapping ──────────────────────────────────
+# ───────── Functions for I/O-bound concurrency ─────────────────────────────
 def _quality_stat_args(args: Tuple[Path, float]) -> Tuple[Path, bool, bool, bool]:
     path, thr = args
     if cv2 is None:
@@ -151,25 +139,26 @@ def _is_corrupt(path: Path) -> bool:
 
 # ───────────────── Quality Checks ──────────────────────────────────────────
 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]
+    missing = [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}
+    with ThreadPoolExecutor(max_workers=cfg.cpu_count) as ex:
+        fut = {ex.submit(_is_corrupt, p): p for p in imgs[:cfg.sample_limit]}
         for f in as_completed(fut):
-            if f.result(): corrupt.append(fut[f])
-    score = 100 - (len(missing)+len(corrupt))/max(len(imgs),1)*100
-    return {"name":"Integrity","score":max(score,0),
+            if f.result():
+                corrupt.append(fut[f])
+    score = 100 - (len(missing)+len(corrupt)) / max(len(imgs), 1) * 100
+    return {"name":"Integrity","score":max(score, 0),
             "details":{"missing_label_files":[str(p) for p in missing],
                        "corrupt_images":[str(p) for p in corrupt]}}
 
 def qc_class_balance(lbls: List[Path], cfg: QCConfig) -> Dict:
     counts, boxes = Counter(), []
-    for l in lbls:
+    for l in lbls[:cfg.sample_limit]:
         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
+    bal = min(counts.values()) / max(counts.values()) * 100
     return {"name":"Class balance","score":bal,
             "details":{"class_counts":dict(counts),
                        "boxes_per_image":{
@@ -180,24 +169,21 @@ def qc_image_quality(imgs: List[Path], cfg: QCConfig) -> Dict:
     if cv2 is None:
         return {"name":"Image quality","score":100,"details":"cv2 missing"}
     blurry, dark, bright = [], [], []
-    with ProcessPoolExecutor(max_workers=cfg.cpu_count) as ex:
-        args = [(p, cfg.blur_thr) for p in imgs]
-        for p, isb, isd, isB in tqdm(
-            ex.map(_quality_stat_args, args), total=len(imgs),
-            desc="img-quality", leave=False
-        ):
+    sample = imgs[:cfg.sample_limit]
+    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(imgs),1)*100
+    bad = len({*blurry, *dark, *bright})
+    score = 100 - bad / max(len(sample), 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) -> Dict:
-    # fastdup-only path
     if fastdup is not None and len(imgs) > 50:
         try:
             fd = fastdup.create(
@@ -212,7 +198,6 @@ def qc_duplicates(imgs: List[Path], cfg: QCConfig) -> Dict:
         except Exception as e:
             return {"name":"Duplicates","score":100,
                     "details":{"fastdup_error":str(e)}}
-    # fallback skipped
     return {"name":"Duplicates","score":100,
             "details":{"note":"fastdup not available or small dataset"}}
 
@@ -221,20 +206,23 @@ def qc_model_qa(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
     if model is None:
         return {"name":"Model QA","score":100,"details":"skipped"}
     ious, mism = [], []
-    for i in range(0,len(imgs),cfg.batch_size):
-        batch = imgs[i:i+cfg.batch_size]
+    sample = imgs[:cfg.sample_limit]
+    for i in range(0, len(sample), cfg.batch_size):
+        batch = sample[i:i+cfg.batch_size]
         results = model.predict(batch, verbose=False, half=True, dynamic=True)
-        for p,res in zip(batch,results):
+        for p, res in zip(batch, results):
             gt = parse_label_file(p.parent.parent/'labels'/f"{p.stem}.txt")
             for cls,x,y,w,h in gt:
-                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)))
+                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: mism.append(str(p))
+                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,"mismatches":mism[:50]}}
@@ -244,9 +232,10 @@ def qc_label_issues(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
         return {"name":"Label issues","score":100,"details":"cleanlab missing"}
     labels, preds, idxs = [], [], []
     model = get_model(cfg.weights)
-    for i,(img,lbl) in enumerate(zip(imgs,lbls)):
+    sample = imgs[:cfg.sample_limit]
+    for i, (img, lbl) in enumerate(zip(sample, lbls[:cfg.sample_limit])):
         bs = parse_label_file(lbl) if lbl else []
-        for cls,*_ in bs:
+        for cls, *_ in bs:
             labels.append(int(cls)); idxs.append(i)
             res = model.predict([img], verbose=False)[0]
             pred_cls = int(res.boxes.cls.cpu().numpy()[0]) if len(res.boxes)>0 else -1
@@ -254,26 +243,25 @@ def qc_label_issues(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
     if not labels:
         return {"name":"Label issues","score":100,"details":"no GT"}
     labels_arr = np.array(labels)
-    # one-hot dummy
     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(imgs[i]) for i in flags]
+    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 _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
+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 aggregate(results: List[Dict]) -> float:
+def aggregate(results: List[Drawable]) -> float:
     return sum(DEFAULT_W[r['name']]*r['score'] for r in results)
 
 RF_RE = re.compile(r"https?://universe\.roboflow\.com/([^/]+)/([^/]+)/dataset/(\d+)")
@@ -289,23 +277,35 @@ def download_rf_dataset(url: str, rf_api: Roboflow, dest: Path) -> Path:
     pr.version(int(ver)).download("yolov8", location=str(ds_dir))
     return ds_dir
 
-def run_quality(root: Path, yaml_file: Path | None, weights: Path | None, cfg: QCConfig) -> Tuple[str,pd.DataFrame]:
-    imgs,lbls,meta = gather_dataset(root, yaml_file)
+def run_quality(root: Path, yaml_file: Path | None, weights: Path | None, cfg: QCConfig,
+                run_dup: bool, run_modelqa: bool) -> Tuple[str, pd.DataFrame]:
+    imgs, lbls, meta = gather_dataset(root, yaml_file)
     results = [
-        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),
+        qc_integrity(imgs, lbls, cfg),
+        qc_class_balance(lbls, cfg),
+        qc_image_quality(imgs, cfg)
     ]
+    # conditional duplicates
+    if run_dup:
+        results.append(qc_duplicates(imgs, cfg))
+    else:
+        results.append({"name":"Duplicates","score":100,"details":"skipped"})
+    # conditional model QA & label issues
+    if run_modelqa:
+        results.append(qc_model_qa(imgs, lbls, cfg))
+        results.append(qc_label_issues(imgs, lbls, cfg))
+    else:
+        results.append({"name":"Model QA","score":100,"details":"skipped"})
+        results.append({"name":"Label issues","score":100,"details":"skipped"})
     final = aggregate(results)
-    md = [f"## **{meta.get('name',root.name)}** — Score {final:.1f}/100"]
+
+    md = [f"## **{meta.get('name', root.name)}** — Score {final:.1f}/100"]
     for r in results:
-        md.append(f"### {r['name']}  {r['score']:.1f}")
+        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(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']
@@ -318,8 +318,8 @@ with gr.Blocks(title="YOLO Dataset Quality Evaluator v3") as demo:
 # YOLOv8 Dataset Quality Evaluator v3
 
 * Configurable blur, IOU & confidence thresholds
-* Cleanlab label-issue detection
-* Fastdup-only duplicates (no hashing fallback)
+* Optional duplicates (fastdup)
+* Optional Model QA & cleanlab label-issue detection
 * Model caching for speed
 """)
     with gr.Row():
@@ -332,47 +332,64 @@ with gr.Blocks(title="YOLO Dataset Quality Evaluator v3") as demo:
         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.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")
+        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")
+    with gr.Row():
+        run_dup     = gr.Checkbox(label="Check duplicates (fastdup)", value=False)
+        run_modelqa = gr.Checkbox(label="Run Model QA & cleanlab", value=False)
     run_btn = gr.Button("Evaluate")
     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):
+                 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)
+        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_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)
-                    md, df = run_quality(ds, None, Path(weights.name) if weights else None, cfg)
+                    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}")
+                    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)
+            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:
             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)
+            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)
         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],
+                  inputs=[api_in, url_txt, zip_in, Path, yaml_in, weights_in,
+                          blur_sl, iou_sl, conf_sl, run_dup, run_modelqa],
                   outputs=[out_md, out_df])
 
 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)))