Add Krea2 atlas coverage report

krea2_tuning_report.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import Counter
+from pathlib import Path
 from typing import Any
 
 try:
@@ -66,6 +67,76 @@ def coverage_summary() -> dict[str, Any]:
     }
 
 
+def _catalog_atlas_root() -> Path:
+    catalog = krea2_pose_variant_catalog.load_catalog()
+    return Path(str(catalog.get("atlas_root") or ""))
+
+
+def _mapped_atlas_folders() -> dict[str, list[str]]:
+    mapped: dict[str, list[str]] = {}
+    for variant in krea2_pose_variant_catalog.variants():
+        key = str(variant.get("key") or "")
+        for folder in variant.get("atlas_folders") or []:
+            folder_name = str(folder)
+            if not folder_name:
+                continue
+            mapped.setdefault(folder_name, []).append(key)
+    return mapped
+
+
+def _is_background_or_control_folder(folder_name: str) -> bool:
+    lower = folder_name.lower()
+    return (
+        lower == "bg"
+        or lower == "woman"
+        or lower.endswith("_control")
+        or lower.endswith("_bg")
+        or lower.endswith("_control_bg")
+    )
+
+
+def atlas_folder_rows(atlas_root: str | Path | None = None) -> list[dict[str, Any]]:
+    root = Path(atlas_root) if atlas_root is not None else _catalog_atlas_root()
+    if not root.is_dir():
+        return []
+    mapped = _mapped_atlas_folders()
+    rows: list[dict[str, Any]] = []
+    for folder in sorted(root.iterdir(), key=lambda path: path.name.lower()):
+        if not folder.is_dir():
+            continue
+        folder_name = folder.name
+        if _is_background_or_control_folder(folder_name):
+            continue
+        image_count = sum(1 for _ in folder.glob("*.png"))
+        if image_count <= 0:
+            continue
+        control_folder = root / f"{folder_name}_control"
+        variant_keys = mapped.get(folder_name, [])
+        if not variant_keys and not control_folder.is_dir():
+            continue
+        rows.append(
+            {
+                "folder": folder_name,
+                "image_count": image_count,
+                "mapped": bool(variant_keys),
+                "variant_keys": list(variant_keys),
+                "control_folder": str(control_folder) if control_folder.is_dir() else "",
+            }
+        )
+    return rows
+
+
+def atlas_coverage_summary(atlas_root: str | Path | None = None) -> dict[str, Any]:
+    rows = atlas_folder_rows(atlas_root=atlas_root)
+    unmapped = [str(row.get("folder")) for row in rows if not row.get("mapped")]
+    return {
+        "pose_folder_count": len(rows),
+        "mapped_folder_count": len(rows) - len(unmapped),
+        "unmapped_folder_count": len(unmapped),
+        "unmapped_folders": unmapped,
+    }
+
+
 def next_test_plans() -> list[dict[str, Any]]:
     rows_by_key = {str(row.get("key")): row for row in coverage_rows()}
     plans: list[dict[str, Any]] = []
@@ -94,7 +165,7 @@ def next_test_plans() -> list[dict[str, Any]]:
     return plans
 
 
-def markdown_report() -> str:
+def markdown_report(atlas_root: str | Path | None = None) -> str:
     lines = [
         "# Krea2 Pose Variant Coverage",
         "",
@@ -132,4 +203,19 @@ def markdown_report() -> str:
                     *[f"  - {cue}" for cue in plan["avoid_cues"]],
                 ]
             )
+    atlas_summary = atlas_coverage_summary(atlas_root=atlas_root)
+    if atlas_summary["pose_folder_count"]:
+        unmapped = atlas_summary["unmapped_folders"]
+        lines.extend(
+            [
+                "",
+                "## Atlas Folder Coverage",
+                "",
+                f"- Pose folders: {atlas_summary['pose_folder_count']}",
+                f"- Mapped folders: {atlas_summary['mapped_folder_count']}",
+                f"- Unmapped folders: {atlas_summary['unmapped_folder_count']}",
+            ]
+        )
+        if unmapped:
+            lines.extend(["", "Unmapped atlas folders:", *[f"- {folder}" for folder in unmapped]])
     return "\n".join(lines)