Blame for data_generator.py

"""
RockGlacierDataGenerator — loads pre-extracted 128×128 .npy patches
from datacollection/output/patches/{split}/features|masks.

Patches are already fixed-size (128×128), cleaned (no-data removed),
so this generator only needs to: normalize, augment, and batch.
"""

import os
import glob
import json
import numpy as np
import tensorflow as tf
import math


# Fallback per-band stats (measured on the OLD glacier-centric crops).
# The generator loads patches metadata/norm_stats.json when present.
GLOBAL_MEAN = np.array([
    0.324, 0.335, 0.344, 0.357, 0.251, 0.229,
    0.038, -0.061, 0.085, 5126.0, 17.4, 189.3
], dtype=np.float32).reshape(1, 1, 12)

GLOBAL_STD = np.array([
    0.183, 0.169, 0.160, 0.134, 0.058, 0.050,
    0.075, 0.096, 0.301, 282.6, 9.9, 110.8
], dtype=np.float32).reshape(1, 1, 12)


class RockGlacierDataGenerator(tf.keras.utils.Sequence):
    def __init__(self, features_dir, masks_dir, batch_size=8, target_size=(128, 128),
                 shuffle=True, augment=False, return_weight=False, radiometric=False,
                 band_indices=None, require_glacier=False, bg_keep_n=0,
                 min_glacier_pct=0.0, copy_paste=False, copy_paste_prob=0.5):
        self.features_dir = features_dir
        self.masks_dir = masks_dir
        self.batch_size = batch_size
        self.target_size = target_size
        self.shuffle = shuffle
        self.augment = augment
        # Optional band subset: list of integer band indices into the stored
        # feature stack. When set, features AND the loaded norm stats are sliced
        # to these bands, so the model trains on fewer channels.
        self.band_indices = np.asarray(band_indices) if band_indices is not None else None
        # When True, __getitem__ yields Y as a 2-channel tensor [label, validity],
        # validity=1 on real pixels and 0 on zero-padding, so the loss can ignore
        # padded pixels. Off by default to keep other consumers unchanged.
        self.return_weight = return_weight
        self.radiometric = radiometric
        self.require_glacier  = require_glacier
        self.bg_keep_n        = bg_keep_n
        self.min_glacier_pct  = min_glacier_pct
        # Copy-paste augmentation: paste a donor glacier's pixels onto the current
        # patch (only in augment mode). Synthesises new glacier arrangements —
        # the strongest augmentation for instance-scarce segmentation.
        self.copy_paste       = copy_paste
        self.copy_paste_prob  = copy_paste_prob

        # Support both .npy patches and .tif files
        self.feature_files = sorted(glob.glob(os.path.join(self.features_dir, '*.npy')))
        self.is_npy = len(self.feature_files) > 0

        if self.is_npy:
            # .npy mode: match by filename stem
            self.mask_files = []
            for fp in self.feature_files:
                stem = os.path.basename(fp)
                mp = os.path.join(self.masks_dir, stem)
                assert os.path.exists(mp), f"Missing mask for {stem}"
                self.mask_files.append(mp)
        else:
            # Legacy .tif mode
            self.feature_files = sorted(glob.glob(os.path.join(self.features_dir, 'feature_*.tif')))
            self.mask_files = sorted(glob.glob(os.path.join(self.masks_dir, 'mask_*.tif')))

        assert len(self.feature_files) == len(self.mask_files), \
            f"Mismatch: {len(self.feature_files)} features vs {len(self.mask_files)} masks"
        assert len(self.feature_files) > 0, f"No data found in {self.features_dir}"

        self.indices = np.arange(len(self.feature_files))

        if self.require_glacier:
            self._filter_to_glacier_patches(bg_keep_n=self.bg_keep_n,
                                            min_glacier_pct=self.min_glacier_pct)

        self.GLOBAL_MEAN, self.GLOBAL_STD = self._load_norm_stats()
        if self.band_indices is not None:
            self.GLOBAL_MEAN = self.GLOBAL_MEAN[..., self.band_indices]
            self.GLOBAL_STD = self.GLOBAL_STD[..., self.band_indices]

        print(f"[DataGenerator] Loaded {len(self.feature_files)} {'npy' if self.is_npy else 'tif'} "
              f"files from {self.features_dir}"
              + (f" | using {len(self.band_indices)} selected bands" if self.band_indices is not None else ""))

        # Donor pool for copy-paste = indices of patches that contain glacier
        self.glacier_pool = []
        if self.copy_paste and self.augment:
            for i, mp in enumerate(self.mask_files):
                if np.load(mp).max() > 0:
                    self.glacier_pool.append(i)
            print(f"[DataGenerator] copy-paste ON — donor pool = {len(self.glacier_pool)} "
                  f"glacier patches (prob={self.copy_paste_prob})")

        if self.shuffle:
            np.random.shuffle(self.indices)

    def _load_norm_donor(self, idx):
        """Load + band-select + NaN-fill + normalize a donor patch (same as main path)."""
        f = np.load(self.feature_files[idx]).astype(np.float32)
        m = np.load(self.mask_files[idx]).astype(np.float32)
        if self.band_indices is not None:
            f = f[..., self.band_indices]
        nod = (f == -9999) | ~np.isfinite(f)
        f = np.where(nod, 0.0, f)
        f = (f - self.GLOBAL_MEAN) / (self.GLOBAL_STD + 1e-7)
        f = np.where(nod, 0.0, f)
        return f, m

    def _apply_copy_paste(self, feature, mask):
        """Paste a random donor glacier's pixels onto (feature, mask), in-place-safe."""
        if not self.glacier_pool:
            return feature, mask
        idx = int(np.random.choice(self.glacier_pool))
        df, dm = self._load_norm_donor(idx)
        # random jitter on the donor so pastes are diverse
        k = np.random.randint(0, 4)
        if k:
            df = np.rot90(df, k); dm = np.rot90(dm, k)
        if np.random.rand() > 0.5:
            df = np.fliplr(df); dm = np.fliplr(dm)
        donor_glac = dm[..., 0] > 0.5                 # (H, W) boolean
        if donor_glac.any():
            feature = feature.copy()
            mask = mask.copy()
            feature[donor_glac] = df[donor_glac]       # overwrite features with donor glacier
            mask[..., 0][donor_glac] = 1.0             # mark those pixels as glacier
        return feature, mask

    def _filter_to_glacier_patches(self, bg_keep_n=0, min_glacier_pct=0.0):
        """Keep glacier patches + optionally N background patches for hard-negative mining.

        With 52% of patches containing no glacier, every other training step
        gave zero useful gradient for the minority class. require_glacier=True
        drops all background patches, but then the model never learns to suppress
        false positives on background terrain. bg_keep_n adds back a fixed number
        of background patches so the model sees both classes of scene.

        Result is cached in metadata/glacier_patches.txt so subsequent runs
        skip the mask scan entirely.
        """
        cache_path = os.path.normpath(
            os.path.join(self.masks_dir, os.pardir, os.pardir,
                         "metadata", "glacier_patches.txt"))

        if os.path.exists(cache_path):
            with open(cache_path) as f:
                glacier_stems = set(ln.strip() for ln in f if ln.strip())
            print(f"[DataGenerator] glacier patch list loaded from cache "
                  f"({len(glacier_stems)} patches)")
        else:
            print(f"[DataGenerator] scanning {len(self.mask_files)} masks "
                  f"for glacier pixels (one-time, result will be cached)...")
            glacier_stems = set()
            for mp in self.mask_files:
                if np.load(mp).max() > 0:
                    glacier_stems.add(os.path.basename(mp))
            os.makedirs(os.path.dirname(cache_path), exist_ok=True)
            with open(cache_path, "w") as f:
                f.write("\n".join(sorted(glacier_stems)))
            print(f"[DataGenerator] cached glacier patch list → {cache_path}")

        keep_f, keep_m = [], []
        bg_f,   bg_m   = [], []
        low_density    = 0
        for fp, mp in zip(self.feature_files, self.mask_files):
            if os.path.basename(mp) in glacier_stems:
                if min_glacier_pct > 0:
                    pct = 100.0 * float(np.load(mp).mean())
                    if pct < min_glacier_pct:
                        low_density += 1
                        continue          # drop noisy boundary-sliver patches
                keep_f.append(fp)
                keep_m.append(mp)
            else:
                bg_f.append(fp)
                bg_m.append(mp)

        if low_density > 0:
            print(f"[DataGenerator] dropped {low_density} glacier patches "
                  f"with <{min_glacier_pct:.1f}% glacier density (noise)")

        n_glacier = len(keep_f)

        # Hard-negative mining: add back N background patches so the model
        # learns to suppress false positives on background terrain.
        if bg_keep_n > 0 and bg_f:
            rng = np.random.default_rng(seed=42)
            n_bg = min(bg_keep_n, len(bg_f))
            idx  = rng.choice(len(bg_f), n_bg, replace=False)
            keep_f.extend([bg_f[i] for i in idx])
            keep_m.extend([bg_m[i] for i in idx])
            print(f"[DataGenerator] hard-negative mining: added {n_bg} background patches")

        n_total   = len(self.feature_files)
        n_dropped = n_total - len(keep_f)
        self.feature_files = keep_f
        self.mask_files    = keep_m
        self.indices       = np.arange(len(self.feature_files))
        print(f"[DataGenerator] {len(self.feature_files)} training patches "
              f"({n_glacier} glacier ≥{min_glacier_pct:.0f}% + "
              f"{len(keep_f)-n_glacier} background | dropped {n_dropped})")

    def _load_norm_stats(self):
        """Use metadata/norm_stats.json if it sits above features_dir;
        otherwise fall back to the module GLOBAL_MEAN/STD."""
        # features_dir = .../patches/train/features -> .../patches/metadata
        cand = os.path.join(self.features_dir, os.pardir, os.pardir, "metadata", "norm_stats.json")
        if os.path.exists(cand):
            with open(cand) as f:
                st = json.load(f)
            mean = np.array(st["mean"], np.float32).reshape(1, 1, -1)
            std = np.array(st["std"], np.float32).reshape(1, 1, -1)
            print(f"[DataGenerator] loaded norm stats from {os.path.normpath(cand)}")
            return mean, std
        print("[DataGenerator] norm_stats.json not found — using built-in GLOBAL_MEAN/STD")
        return GLOBAL_MEAN, GLOBAL_STD

    def __len__(self):
        return math.ceil(len(self.feature_files) / self.batch_size)

    def _load_pair(self, f_path, m_path):
        """Load a feature-mask pair from .npy or .tif."""
        if self.is_npy:
            feature = np.load(f_path).astype(np.float32)   # (128, 128, 12)
            mask = np.load(m_path).astype(np.float32)       # (128, 128, 1)
        else:
            import rasterio
            with rasterio.open(f_path) as f_src:
                feature = np.moveaxis(f_src.read(), 0, -1).astype(np.float32)  # (H, W, 12)
            with rasterio.open(m_path) as m_src:
                mask = np.expand_dims(m_src.read(1), axis=-1).astype(np.float32)  # (H, W, 1)
        return feature, mask

    def __getitem__(self, index):
        batch_indices = self.indices[index * self.batch_size:(index + 1) * self.batch_size]

        X_batch = []
        Y_batch = []

        for i in batch_indices:
            f_path = self.feature_files[i]
            m_path = self.mask_files[i]

            feature, mask = self._load_pair(f_path, m_path)

            # Select band subset (before normalization, which uses sliced stats)
            if self.band_indices is not None:
                feature = feature[..., self.band_indices]

            # Force mask to match feature shape (legacy .tif misalignment fix)
            H_f, W_f = feature.shape[:2]
            mask = mask[:H_f, :W_f, :]
            pad_mh, pad_mw = max(0, H_f - mask.shape[0]), max(0, W_f - mask.shape[1])
            if pad_mh > 0 or pad_mw > 0:
                mask = np.pad(mask, ((0, pad_mh), (0, pad_mw), (0, 0)), mode='constant')

            # Identify NoData pixels before normalization
            nodata_mask = (feature == -9999) | ~np.isfinite(feature)
            feature = np.where(nodata_mask, 0.0, feature)

            # Normalization
            feature = (feature - self.GLOBAL_MEAN) / (self.GLOBAL_STD + 1e-7)

            # Re-zero NoData pixels in normalized space so they equal the band mean
            # (normalized 0 = band mean), not a biased −mean/std artifact
            feature = np.where(nodata_mask, 0.0, feature)

            # Validity mask = 1 on real pixels (before any padding)
            valid = np.ones((feature.shape[0], feature.shape[1], 1), dtype=np.float32)

            # Padding to target_size if smaller (only needed for .tif mode)
            h, w = feature.shape[:2]
            th, tw = self.target_size

            pad_h = max(0, th - h)
            pad_w = max(0, tw - w)

            if pad_h > 0 or pad_w > 0:
                feature = np.pad(feature, ((0, pad_h), (0, pad_w), (0, 0)), mode='constant', constant_values=0)
                mask = np.pad(mask, ((0, pad_h), (0, pad_w), (0, 0)), mode='constant', constant_values=0)
                valid = np.pad(valid, ((0, pad_h), (0, pad_w), (0, 0)), mode='constant', constant_values=0)

            # Cropping if larger (only needed for .tif mode)
            h, w = feature.shape[:2]
            if h > th or w > tw:
                if self.augment:
                    h_start = np.random.randint(0, h - th + 1) if h > th else 0
                    w_start = np.random.randint(0, w - tw + 1) if w > tw else 0
                else:
                    h_start = (h - th) // 2 if h > th else 0
                    w_start = (w - tw) // 2 if w > tw else 0
                feature = feature[h_start:h_start + th, w_start:w_start + tw, :]
                mask = mask[h_start:h_start + th, w_start:w_start + tw, :]
                valid = valid[h_start:h_start + th, w_start:w_start + tw, :]

            # Augmentations (Geometric only)
            if self.augment:
                # Copy-paste: overlay a donor glacier before geometric transforms
                # so the combined patch is flipped/rotated as one.
                if self.copy_paste and np.random.rand() < self.copy_paste_prob:
                    feature, mask = self._apply_copy_paste(feature, mask)
                # Random Horizontal Flip
                if np.random.rand() > 0.5:
                    feature = np.fliplr(feature)
                    mask = np.fliplr(mask)
                    valid = np.fliplr(valid)
                # Random Vertical Flip
                if np.random.rand() > 0.5:
                    feature = np.flipud(feature)
                    mask = np.flipud(mask)
                    valid = np.flipud(valid)
                # Random Rotations (90, 180, 270)
                k = np.random.randint(0, 4)
                if k > 0:
                    feature = np.rot90(feature, k=k)
                    mask = np.rot90(mask, k=k)
                    valid = np.rot90(valid, k=k)

                # Radiometric augmentation (in normalized space): simulates
                # sensor/atmospheric variation to reduce overfitting.
                if self.radiometric:
                    feature = np.ascontiguousarray(feature, dtype=np.float32)
                    if np.random.rand() < 0.5:                       # per-band gain/offset
                        c = feature.shape[-1]
                        gain = np.random.uniform(0.9, 1.1, (1, 1, c)).astype(np.float32)
                        offset = np.random.uniform(-0.05, 0.05, (1, 1, c)).astype(np.float32)
                        feature = feature * gain + offset
                    if np.random.rand() < 0.3:                       # gaussian noise
                        feature = feature + np.random.normal(0, 0.05, feature.shape).astype(np.float32)

            X_batch.append(feature)
            if self.return_weight:
                Y_batch.append(np.concatenate([mask, valid], axis=-1))   # (H, W, 2)
            else:
                Y_batch.append(mask)

        return np.array(X_batch, dtype=np.float32), np.array(Y_batch, dtype=np.float32)

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.indices)