Update app.py

app.py

@@ -92,6 +92,7 @@ def load_class_names(yaml_path: Path) -> List[str]:
     data = load_yaml(yaml_path)
     names = data.get("names", [])
     if isinstance(names, dict):
+        # Convert { "0": "Person", ... } to ["Person", ...] by key order
         return [names[k] for k in sorted(names, key=lambda x: int(x))]
     return list(names)
 
@@ -310,7 +311,23 @@ def qc_label_issues(imgs: List[Path], lbls: List[Path], cfg: QCConfig) -> Dict:
 def aggregate(results: List[Dict]) -> float:
     return sum(DEFAULT_W[r["name"]]*r["score"] for r in results)
 
-# ───────────────── Roboflow TXT‑loading logic for both tabs ─────────────────
+# ─────── New: gather actual per‑class image 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, *_ in parse_label_file(lp):
+                    if 0 <= cls_id < len(class_names):
+                        counts[class_names[cls_id]] += 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:
@@ -325,6 +342,7 @@ 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
 
+# ───────────────── run_quality & merge_datasets ────────────────────────────
 def run_quality(
     root: Path,
     yaml_file: Path | None,
@@ -350,7 +368,7 @@ def run_quality(
         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"],
+        next(r for r in results if r["name"] == "Class balance")["details"]["class_counts"],
         orient="index", columns=["count"]
     )
     df.index.name = "class"
@@ -461,6 +479,7 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
     _Evaluate • Merge • Edit • Download_
     """)
 
+    # ------------------------------ 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'])
@@ -535,6 +554,7 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
             outputs=[out_md, out_df]
         )
 
+    # ------------------------------ MERGE / EDIT TAB --------------------------
     with gr.Tab("Merge / Edit"):
         gr.Markdown("### 1️⃣ Load one or more datasets")
         rf_key   = gr.Textbox(label="Roboflow API key", type="password")
@@ -570,7 +590,7 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
                 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()]
+                    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")
 
@@ -587,11 +607,12 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
         refresh_btn = gr.Button("Build class table from loaded datasets")
 
         def _build_class_df(ds_info):
-            all_names = sorted({n for _, names, _, _ in ds_info for n in names})
+            counts = gather_class_counts(ds_info)
+            all_names = sorted(counts.keys())
             return pd.DataFrame({
                 "original_class": all_names,
                 "new_name":       all_names,
-                "max_images":     [99999]*len(all_names),
+                "max_images":     [counts[n] for n in all_names],
                 "remove":         [False]*len(all_names),
             })