Add serving-oriented tiled inference and jittered training support

training/train_cold_tiled_mainline.py

@@ -28,6 +28,11 @@ def read_rows(path: Path) -> List[Dict[str, str]]:
         return list(csv.DictReader(fh))
 
 
+def write_json(path: Path, data: Dict[str, object]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+
+
 def integral_image(mask: np.ndarray) -> np.ndarray:
     return np.pad(mask.astype(np.int32), ((1, 0), (1, 0)), mode="constant").cumsum(0).cumsum(1)
 
@@ -44,16 +49,47 @@ def centered_tile_origin(y: int, x: int, tile_h: int, tile_w: int, h: int, w: in
     return int(top), int(left)
 
 
-def positive_tile_origins(mask: np.ndarray, tile_h: int, tile_w: int, max_tiles: int, rng: random.Random) -> List[Tuple[int, int]]:
+def containing_tile_origin(y: int, x: int, tile_h: int, tile_w: int, h: int, w: int, rng: random.Random) -> Tuple[int, int]:
+    max_top = max(h - tile_h, 0)
+    max_left = max(w - tile_w, 0)
+    top_min = max(0, y - tile_h + 1)
+    top_max = min(y, max_top)
+    left_min = max(0, x - tile_w + 1)
+    left_max = min(x, max_left)
+    if top_min > top_max or left_min > left_max:
+        return centered_tile_origin(y, x, tile_h, tile_w, h, w)
+    return int(rng.randint(top_min, top_max)), int(rng.randint(left_min, left_max))
+
+
+def positive_tile_origins(
+    mask: np.ndarray,
+    tile_h: int,
+    tile_w: int,
+    max_tiles: int,
+    rng: random.Random,
+    placement: str = "center",
+) -> List[Tuple[int, int]]:
     ys, xs = np.where(mask > 0.5)
     if ys.size == 0:
         return []
-    origins = list(
-        {
-            centered_tile_origin(int(y), int(x), tile_h, tile_w, mask.shape[0], mask.shape[1])
-            for y, x in zip(ys, xs)
-        }
-    )
+    origin_fn = containing_tile_origin if placement == "random_containing" else centered_tile_origin
+    origins = []
+    seen = set()
+    indices = list(range(int(ys.size)))
+    rng.shuffle(indices)
+    for i in indices:
+        y = int(ys[i])
+        x = int(xs[i])
+        if placement == "random_containing":
+            origin = origin_fn(y, x, tile_h, tile_w, mask.shape[0], mask.shape[1], rng)
+        else:
+            origin = origin_fn(y, x, tile_h, tile_w, mask.shape[0], mask.shape[1])
+        if origin in seen:
+            continue
+        seen.add(origin)
+        origins.append(origin)
+        if max_tiles > 0 and len(origins) >= max_tiles:
+            break
     rng.shuffle(origins)
     if max_tiles > 0:
         origins = origins[:max_tiles]
@@ -191,9 +227,103 @@ class UNetSmallFlex(nn.Module):
         return logits
 
 
+def normalization_config(config: Dict[str, object]) -> Dict[str, object]:
+    raw = config.get("input_normalization", {})
+    if raw is True:
+        return {"enabled": True}
+    if isinstance(raw, dict):
+        return raw
+    return {"enabled": False}
+
+
+def compute_input_normalization_stats(
+    rows: List[Dict[str, str]],
+    continuous_channel_indices: List[int],
+    eps: float,
+) -> Dict[str, object]:
+    count: Dict[int, int] = {int(idx): 0 for idx in continuous_channel_indices}
+    sum_x: Dict[int, float] = {int(idx): 0.0 for idx in continuous_channel_indices}
+    sum_x2: Dict[int, float] = {int(idx): 0.0 for idx in continuous_channel_indices}
+    channel_names: List[str] | None = None
+    static_names: List[str] | None = None
+    for row in rows:
+        sample = np.load(row["sample_path"], allow_pickle=True)
+        weather = ColdFeatureStore._sanitize(sample["weather"].astype(np.float32))
+        firewx = ColdFeatureStore._sanitize(sample["firewx"].astype(np.float32))
+        extra = []
+        if "firewx_valid" in sample:
+            extra.append(sample["firewx_valid"].astype(np.float32))
+        static_npz = np.load(row["static_path"], allow_pickle=True)
+        static = ColdFeatureStore._sanitize(static_npz["static"].astype(np.float32))
+        static_parts = []
+        if "static_valid" in static_npz:
+            static_parts.append(static_npz["static_valid"].astype(np.float32))
+        static_parts.append(static)
+        if channel_names is None:
+            weather_names = [str(v) for v in sample.get("weather_names", np.array([], dtype=object)).tolist()]
+            firewx_names = [str(v) for v in sample.get("firewx_names", np.array([], dtype=object)).tolist()]
+            extra_names = ["firewx_valid"] if "firewx_valid" in sample else []
+            static_names = [str(v) for v in static_npz.get("static_names", np.array([], dtype=object)).tolist()]
+            static_valid_names = ["static_valid"] if "static_valid" in static_npz else []
+            channel_names = weather_names + firewx_names + extra_names + static_valid_names + static_names
+        x = np.concatenate([weather, firewx, *extra, *static_parts], axis=0).astype(np.float32)
+        for idx in continuous_channel_indices:
+            arr = x[int(idx)].astype(np.float64, copy=False).ravel()
+            count[int(idx)] += int(arr.size)
+            sum_x[int(idx)] += float(arr.sum())
+            sum_x2[int(idx)] += float(np.square(arr).sum())
+    channels = []
+    for idx in continuous_channel_indices:
+        n = max(int(count[int(idx)]), 1)
+        mean = float(sum_x[int(idx)] / n)
+        variance = max(float(sum_x2[int(idx)] / n - mean * mean), 0.0)
+        std = float(max(math.sqrt(variance), eps))
+        channels.append(
+            {
+                "index": int(idx),
+                "name": channel_names[int(idx)] if channel_names and int(idx) < len(channel_names) else str(idx),
+                "mean": mean,
+                "std": std,
+                "count": int(count[int(idx)]),
+            }
+        )
+    return {
+        "enabled": True,
+        "method": "per-channel z-score",
+        "stats_source": "train split full maps",
+        "continuous_channel_indices": [int(idx) for idx in continuous_channel_indices],
+        "eps": float(eps),
+        "channels": channels,
+        "channel_names": channel_names or [],
+        "static_names": static_names or [],
+    }
+
+
+def load_or_compute_input_normalization_stats(
+    config: Dict[str, object],
+    train_rows: List[Dict[str, str]],
+    metric_dir: Path,
+) -> Dict[str, object] | None:
+    norm_cfg = normalization_config(config)
+    if not bool(norm_cfg.get("enabled", False)):
+        return None
+    stats_path_value = str(norm_cfg.get("stats_path", "")).strip()
+    if stats_path_value:
+        stats = load_json(Path(stats_path_value))
+    else:
+        stats = compute_input_normalization_stats(
+            rows=train_rows,
+            continuous_channel_indices=[int(v) for v in norm_cfg.get("continuous_channel_indices", list(range(10)))],
+            eps=float(norm_cfg.get("eps", 1e-6)),
+        )
+    write_json(metric_dir / "input_normalization_stats.json", stats)
+    return stats
+
+
 class ColdFeatureStore:
-    def __init__(self, rows: List[Dict[str, str]]):
+    def __init__(self, rows: List[Dict[str, str]], normalization_stats: Dict[str, object] | None = None):
         self.cache: Dict[str, Dict[str, np.ndarray]] = {}
+        self.normalization_stats = normalization_stats
         static_path = Path(rows[0]["static_path"])
         static_npz = np.load(static_path, allow_pickle=True)
         static = self._sanitize(static_npz["static"].astype(np.float32))
@@ -211,6 +341,7 @@ class ColdFeatureStore:
             if "firewx_valid" in sample:
                 extra.append(sample["firewx_valid"].astype(np.float32))
             x = np.concatenate([weather, firewx, *extra, static_x], axis=0).astype(np.float32)
+            x = self._apply_normalization(x)
             y = np.nan_to_num(sample["y_occ"].astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0)
             self.cache[str(row["sample_id"])] = {"x": x, "y": y}
 
@@ -220,6 +351,18 @@ class ColdFeatureStore:
         x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
         return x.astype(np.float32, copy=False)
 
+    def _apply_normalization(self, x: np.ndarray) -> np.ndarray:
+        if not self.normalization_stats:
+            return x
+        for item in self.normalization_stats.get("channels", []):
+            idx = int(item["index"])
+            if idx < 0 or idx >= x.shape[0]:
+                continue
+            mean = float(item["mean"])
+            std = max(float(item["std"]), float(self.normalization_stats.get("eps", 1e-6)))
+            x[idx] = (x[idx] - mean) / std
+        return x.astype(np.float32, copy=False)
+
     def get(self, sample_id: str) -> Dict[str, np.ndarray]:
         return self.cache[sample_id]
 
@@ -511,6 +654,7 @@ def build_train_tiles(
     max_positive_tiles_per_sample: int,
     min_negative_tiles_per_sample: int,
     neg_pos_ratio: float,
+    positive_tile_placement: str,
     rng: random.Random,
 ) -> List[Dict[str, object]]:
     tile_rows: List[Dict[str, object]] = []
@@ -523,6 +667,7 @@ def build_train_tiles(
             tile_w=tile_size,
             max_tiles=max_positive_tiles_per_sample,
             rng=rng,
+            placement=positive_tile_placement,
         )
         neg_count = max(min_negative_tiles_per_sample, int(math.ceil(len(pos_origins) * neg_pos_ratio)))
         neg_origins = negative_tile_origins(
@@ -678,7 +823,8 @@ def main() -> None:
     train_rows = read_rows(index_root / "splits" / "train.csv")
     val_rows = read_rows(index_root / "splits" / "val.csv")
     test_rows = read_rows(index_root / "splits" / "test.csv")
-    store = ColdFeatureStore(train_rows + val_rows + test_rows)
+    normalization_stats = load_or_compute_input_normalization_stats(config, train_rows, metric_dir)
+    store = ColdFeatureStore(train_rows + val_rows + test_rows, normalization_stats=normalization_stats)
 
     tile_rows = build_train_tiles(
         train_rows=train_rows,
@@ -687,6 +833,7 @@ def main() -> None:
         max_positive_tiles_per_sample=int(config.get("max_positive_tiles_per_sample", 64)),
         min_negative_tiles_per_sample=int(config.get("min_negative_tiles_per_sample", 4)),
         neg_pos_ratio=float(config.get("negative_to_positive_ratio", 2.0)),
+        positive_tile_placement=str(config.get("positive_tile_placement", "center")),
         rng=rng,
     )
     train_ds = TrainTileDataset(tile_rows=tile_rows, store=store, augment_flip=bool(config.get("augment_flip", True)), seed=int(config.get("seed", 7)))
@@ -747,6 +894,7 @@ def main() -> None:
         "positive_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "positive")),
         "negative_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "negative")),
         "input_channels": in_ch,
+        "positive_tile_placement": str(config.get("positive_tile_placement", "center")),
         "train_positive_rate": train_positive_rate,
         "metric_reference_positive_rate": metric_reference_positive_rate,
         "pos_weight_source": pos_weight_source,
@@ -755,6 +903,7 @@ def main() -> None:
         "loss_type": str(config.get("loss_type", "bce")),
         "train_target_mode": str(config.get("train_target_mode", "hard")),
         "train_target_radius": int(config.get("train_target_radius", 0)),
+        "input_normalization": normalization_stats if normalization_stats else {"enabled": False},
     }
     (metric_dir / "tile_summary.json").write_text(json.dumps(summary_seed, indent=2), encoding="utf-8")
 
@@ -874,6 +1023,7 @@ def main() -> None:
         "num_train_tiles": len(tile_rows),
         "positive_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "positive")),
         "negative_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "negative")),
+        "positive_tile_placement": str(config.get("positive_tile_placement", "center")),
         "train_positive_rate": train_positive_rate,
         "metric_reference_positive_rate": metric_reference_positive_rate,
         "pos_weight": pos_weight,
@@ -882,6 +1032,7 @@ def main() -> None:
         "aux_spatial_loss_weight": aux_spatial_loss_weight,
         "aux_positive_rate": aux_positive_rate,
         "init_checkpoint": str(init_checkpoint) if init_checkpoint else "",
+        "input_normalization": normalization_stats if normalization_stats else {"enabled": False},
         "best_val_pr_auc": float(val_best["pr_auc"]),
         "best_val_auroc": float(val_best["auroc"]),
         "best_test_pr_auc": float(test_best["pr_auc"]),