""" 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)