Upload sample_pairs.py with huggingface_hub

sample_pairs.py

@@ -2,10 +2,13 @@ from __future__ import annotations
 
 import argparse
 import json
+import logging
 import random
 from collections import defaultdict
 from pathlib import Path
 
+logger = logging.getLogger(__name__)
+
 
 def classify_transformation(pair: dict) -> str:
     age_changed = pair.get("source_age") != pair.get("target_age")
@@ -19,6 +22,13 @@ def classify_transformation(pair: dict) -> str:
     return "none"
 
 
+def _get_pathology_key(pair: dict) -> str:
+    labels = pair.get("pathology_labels", [])
+    if not labels:
+        return "unlabeled"
+    return ",".join(str(v) for v in sorted(labels))
+
+
 def stratified_sample(
     pairs: list[dict],
     n_per_stratum: int,
@@ -35,19 +45,52 @@ def stratified_sample(
     selected: list[int] = []
     for stratum in ["age_only", "sex_only", "intersectional"]:
         pool = by_stratum[stratum]
-        rng.shuffle(pool)
-        n = min(n_per_stratum, len(pool))
-        selected.extend(pool[:n])
-        if n < n_per_stratum:
-            print(
-                f"Warning: only {n} pairs available for {stratum} "
-                f"(requested {n_per_stratum})"
+        if not pool:
+            logger.warning("No pairs available for stratum %s", stratum)
+            continue
+
+        by_pathology: dict[str, list[int]] = defaultdict(list)
+        for idx in pool:
+            key = _get_pathology_key(pairs[idx])
+            by_pathology[key].append(idx)
+
+        pathology_keys = sorted(by_pathology.keys())
+        n_categories = len(pathology_keys)
+        per_category = n_per_stratum // n_categories
+        remainder = n_per_stratum % n_categories
+
+        stratum_selected: list[int] = []
+        overflow_pool: list[int] = []
+
+        for k_idx, key in enumerate(pathology_keys):
+            candidates = by_pathology[key]
+            rng.shuffle(candidates)
+            target = per_category + (1 if k_idx < remainder else 0)
+            taken = min(target, len(candidates))
+            stratum_selected.extend(candidates[:taken])
+            if taken < target:
+                overflow_pool.extend([])
+            if taken < len(candidates):
+                overflow_pool.extend(candidates[taken:])
+
+        if len(stratum_selected) < n_per_stratum and overflow_pool:
+            rng.shuffle(overflow_pool)
+            needed = n_per_stratum - len(stratum_selected)
+            stratum_selected.extend(overflow_pool[:needed])
+
+        if len(stratum_selected) < n_per_stratum:
+            logger.warning(
+                "Only %d pairs available for %s (requested %d)",
+                len(stratum_selected), stratum, n_per_stratum,
             )
 
+        selected.extend(stratum_selected)
+
     return selected
 
 
 def main() -> None:
+    logging.basicConfig(level=logging.INFO)
     parser = argparse.ArgumentParser(
         description="Stratified sampling of pairs for radiologist validation"
     )
@@ -67,18 +110,23 @@ def main() -> None:
         pair["_transform_type"] = classify_transformation(pair)
 
     counts: dict[str, int] = defaultdict(int)
+    pathology_counts: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
     for pair in sampled:
-        counts[pair["_transform_type"]] += 1
+        t = pair["_transform_type"]
+        counts[t] += 1
+        pathology_counts[t][_get_pathology_key(pair)] += 1
 
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
     with open(output_path, "w") as f:
         json.dump(sampled, f, indent=2)
 
-    print(f"Sampled {len(sampled)} pairs from {len(all_pairs)} total")
+    logger.info("Sampled %d pairs from %d total", len(sampled), len(all_pairs))
     for stratum, count in sorted(counts.items()):
-        print(f"  {stratum}: {count}")
-    print(f"Saved to: {output_path}")
+        logger.info("  %s: %d", stratum, count)
+        for pkey, pcount in sorted(pathology_counts[stratum].items()):
+            logger.info("    pathology [%s]: %d", pkey, pcount)
+    logger.info("Saved to: %s", output_path)
 
 
 if __name__ == "__main__":