Update app.py

app.py

@@ -92,7 +92,6 @@ 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)
 
@@ -187,7 +186,7 @@ def qc_class_balance(lbls: List[Path], cfg: QCConfig) -> Dict:
     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)
+        counts.update(int(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
@@ -237,8 +236,7 @@ def qc_duplicates(imgs: List[Path], cfg: QCConfig) -> Dict:
             fd.run()
             try:
                 cc = fd.connected_components_grouped(sort_by="comp_size", ascending=False)
-                clusters = (cc["files"].tolist() if "files" in cc.columns
-                            else cc.groupby("component")["filename"].apply(list).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 = sum(len(c) - 1 for c in clusters)
@@ -311,7 +309,7 @@ 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)
 
-# ─────── New: gather actual per‑class image counts ─────────────────────────
+# ───────── gathering actual per-class counts ────────────────────────────────
 def gather_class_counts(
     dataset_info_list: List[Tuple[str, List[str], List[str], str]]
 ) -> Counter[str]:
@@ -322,12 +320,13 @@ def gather_class_counts(
             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
+                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 ─────────────────────────────
+# ───────────────── 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:
@@ -408,10 +407,11 @@ def merge_datasets(
             if not labels_root.exists():
                 continue
             for lp in labels_root.rglob("*.txt"):
-                im_name, cls_set = lp.stem + ".jpg", set()
-                for cls_id, *rest in parse_label_file(lp):
-                    orig = class_names_dataset[int(cls_id)] if int(cls_id) < len(class_names_dataset) else None
-                    if orig:
+                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)
@@ -423,7 +423,7 @@ def merge_datasets(
                 for c in cls_set:
                     class_to_images[c].add(img_path)
 
-    selected_images = set()
+    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)
@@ -446,11 +446,16 @@ def merge_datasets(
         dst_lbl = out_dir / split / "labels" / lp_src.name
         dst_lbl.parent.mkdir(parents=True, exist_ok=True)
         lines = lp_src.read_text().splitlines()
-        new_lines = []
+        new_lines: List[str] = []
         for line in lines:
             parts = line.split()
             cid = int(parts[0])
-            orig = class_names_dataset[cid] if cid < len(class_names_dataset) else None
+            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]
@@ -479,7 +484,7 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
     _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'])
@@ -554,7 +559,7 @@ with gr.Blocks(css="#classdf td{min-width:120px}") as demo:
             outputs=[out_md, out_df]
         )
 
-    # ------------------------------ MERGE / EDIT TAB --------------------------
+    # 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")