Cap regional failure oversampling

scripts/build_failure_driven_yolo_sampling.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import argparse
 import hashlib
 import json
+import math
 from collections import Counter
 from pathlib import Path
 from typing import Any
@@ -38,6 +39,7 @@ def build_sampling(
     precision_positive_repeat: int = 1,
     context_positive_repeat: int = 5,
     context_negative_repeat: int = 6,
+    max_region_share: float = 0.65,
 ) -> tuple[list[str], dict[str, Any]]:
     if assessment.get("status") != "continue_training_loop":
         raise ValueError("Failure-driven sampling requires a failed assessment")
@@ -49,6 +51,8 @@ def build_sampling(
         context_negative_repeat,
     ) < 1:
         raise ValueError("Repeat factors must be positive")
+    if not 0 < max_region_share <= 1:
+        raise ValueError("max_region_share must be in (0, 1]")
 
     samples = {item["sample_slug"]: item for item in manifest["samples"]}
     gates = assessment["gates"]
@@ -91,8 +95,8 @@ def build_sampling(
         ):
             weak_precision_contexts.add(key)
 
-    image_paths: list[str] = []
-    repeat_counts: Counter[str] = Counter()
+    base_paths_by_region: dict[str, list[str]] = {}
+    extra_paths_by_region: dict[str, list[str]] = {}
     selected_samples: set[str] = set()
     protected_samples: set[str] = set()
     for tile in summary["tiles"]:
@@ -120,10 +124,34 @@ def build_sampling(
                 else negative_repeat
             )
         path = str(Path(tile["image_path"]).resolve())
-        image_paths.extend([path] * repeat)
-        repeat_counts[region] += repeat
+        base_paths_by_region.setdefault(region, []).append(path)
+        extra_paths_by_region.setdefault(region, []).extend([path] * (repeat - 1))
         selected_samples.add(tile["sample_slug"])
 
+    pre_cap_counts = Counter({
+        region: len(paths) + len(extra_paths_by_region[region])
+        for region, paths in base_paths_by_region.items()
+    })
+    capped_counts = Counter(pre_cap_counts)
+    if max_region_share < 1:
+        while capped_counts:
+            region, count = max(capped_counts.items(), key=lambda item: (item[1], item[0]))
+            total = sum(capped_counts.values())
+            if count / total <= max_region_share:
+                break
+            others = total - count
+            limit = math.floor(max_region_share / (1 - max_region_share) * others)
+            limit = max(limit, len(base_paths_by_region[region]))
+            if limit >= count:
+                break
+            capped_counts[region] = limit
+    image_paths: list[str] = []
+    for region in sorted(base_paths_by_region):
+        base_paths = base_paths_by_region[region]
+        extra_limit = capped_counts[region] - len(base_paths)
+        image_paths.extend(base_paths)
+        image_paths.extend(extra_paths_by_region[region][:extra_limit])
+    repeat_counts = Counter({region: capped_counts[region] for region in capped_counts})
     if not image_paths:
         raise ValueError("No train-only tiles selected")
     metadata = {
@@ -141,6 +169,7 @@ def build_sampling(
         "precision_positive_repeat": precision_positive_repeat,
         "context_positive_repeat": context_positive_repeat,
         "context_negative_repeat": context_negative_repeat,
+        "max_region_share": max_region_share,
         "source_train_tile_count": sum(
             1
             for tile in summary["tiles"]
@@ -148,6 +177,8 @@ def build_sampling(
         ),
         "sampled_train_entry_count": len(image_paths),
         "sampled_entries_by_region": dict(sorted(repeat_counts.items())),
+        "pre_cap_entries_by_region": dict(sorted(pre_cap_counts.items())),
+        "dropped_region_repeat_count": sum(pre_cap_counts.values()) - len(image_paths),
         "selected_train_sample_count": len(selected_samples),
         "protected_sample_count": len(protected_samples),
         "protected_samples_in_training": [],
@@ -166,6 +197,7 @@ def main() -> int:
     parser.add_argument("--precision-positive-repeat", type=int, default=1)
     parser.add_argument("--context-positive-repeat", type=int, default=5)
     parser.add_argument("--context-negative-repeat", type=int, default=6)
+    parser.add_argument("--max-region-share", type=float, default=0.65)
     args = parser.parse_args()
 
     summary = json.loads(args.summary.read_text(encoding="utf-8"))
@@ -180,6 +212,7 @@ def main() -> int:
         precision_positive_repeat=args.precision_positive_repeat,
         context_positive_repeat=args.context_positive_repeat,
         context_negative_repeat=args.context_negative_repeat,
+        max_region_share=args.max_region_share,
     )
     args.output_dir.mkdir(parents=True, exist_ok=True)
     train_list = args.output_dir / "train-failure-driven.txt"