Spaces:
Sleeping
Sleeping
| import os | |
| from typing import Dict, List, Tuple | |
| import numpy as np | |
| import rasterio | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from torch.utils.data import Dataset | |
| from config import DEFAULT_PATCH_SIZE, NUM_CHANNELS, NUM_CLASSES, IGNORE_INDEX, CLASS_NAMES | |
| DATASET_REPO = os.environ.get("DATASET_REPO", "") | |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") or None | |
| BAND_FILES = [f"H_{i}.tif" for i in range(1, 8)] | |
| TRAIN_MASK_FILE = "TRAINING.tif" | |
| VAL_MASK_FILE = "GROUND TRUTH.tif" | |
| # ── File helpers ───────────────────────────────────────────── | |
| def _get_path(filename: str) -> str: | |
| """Use local file if it exists, otherwise download from HuggingFace.""" | |
| if os.path.exists(filename): | |
| return filename | |
| if not DATASET_REPO: | |
| raise EnvironmentError( | |
| f"'{filename}' not found locally and DATASET_REPO is not set." | |
| ) | |
| return hf_hub_download( | |
| repo_id=DATASET_REPO, | |
| filename=filename, | |
| repo_type="dataset", | |
| token=HF_TOKEN, | |
| ) | |
| def _read_band(path: str) -> np.ndarray: | |
| with rasterio.open(path) as src: | |
| data = src.read(1).astype(np.float32) | |
| if src.nodata is not None: | |
| data[data == src.nodata] = np.nan | |
| return data | |
| def _read_mask_raw(path: str) -> Tuple[np.ndarray, object, str]: | |
| with rasterio.open(path) as src: | |
| data = src.read(1) | |
| nodata = src.nodata | |
| info = f"shape={src.shape} dtype={src.dtypes[0]} nodata={nodata}" | |
| return data, nodata, info | |
| def _normalize(image: np.ndarray) -> np.ndarray: | |
| out = np.zeros_like(image, dtype=np.float32) | |
| for b in range(image.shape[0]): | |
| band = image[b] | |
| finite = band[np.isfinite(band)] | |
| if len(finite) == 0: | |
| continue | |
| lo, hi = np.percentile(finite, 2), np.percentile(finite, 98) | |
| out[b] = np.clip((np.nan_to_num(band) - lo) / max(hi - lo, 1e-6), 0.0, 1.0) | |
| return out | |
| def _compute_ndvi(raw_image: np.ndarray) -> np.ndarray: | |
| """NDVI = (NIR - Red) / (NIR + Red). H_5=NIR (idx 4), H_4=Red (idx 3).""" | |
| nir = raw_image[4].astype(np.float32) | |
| red = raw_image[3].astype(np.float32) | |
| denom = nir + red | |
| return np.where(np.abs(denom) > 1e-6, (nir - red) / denom, 0.0).clip(-1.0, 1.0) | |
| def _compute_ndwi(raw_image: np.ndarray) -> np.ndarray: | |
| """NDWI = (Green - NIR) / (Green + NIR). H_3=Green (idx 2), H_5=NIR (idx 4).""" | |
| green = raw_image[2].astype(np.float32) | |
| nir = raw_image[4].astype(np.float32) | |
| denom = green + nir | |
| return np.where(np.abs(denom) > 1e-6, (green - nir) / denom, 0.0).clip(-1.0, 1.0) | |
| def _remap_mask(raw: np.ndarray, nodata_val) -> Tuple[np.ndarray, List[int]]: | |
| """ | |
| Map raw pixel values -> 0..NUM_CLASSES-1. | |
| Value 0 and nodata -> IGNORE_INDEX. | |
| """ | |
| if nodata_val is not None: | |
| nodata_px = raw == int(nodata_val) | |
| else: | |
| nodata_px = np.zeros(raw.shape, dtype=bool) | |
| ignore_px = nodata_px | (raw == 0) | |
| valid = ~ignore_px | |
| raw_unique = sorted(int(v) for v in np.unique(raw[valid])) | |
| mask = np.full(raw.shape, IGNORE_INDEX, dtype=np.int64) | |
| for cls_idx, raw_val in enumerate(raw_unique[:NUM_CLASSES]): | |
| mask[raw == raw_val] = cls_idx | |
| return mask, raw_unique | |
| def _extract_patches( | |
| image: np.ndarray, | |
| mask: np.ndarray, | |
| patch_size: int, | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| _, H_img, W_img = image.shape | |
| H_msk, W_msk = mask.shape | |
| H = min(H_img, H_msk) | |
| W = min(W_img, W_msk) | |
| stride = patch_size // 2 | |
| imgs, masks = [], [] | |
| def steps(size): | |
| s = list(range(0, size - patch_size + 1, stride)) | |
| if not s: | |
| s = [0] if size >= patch_size else [] | |
| elif s[-1] < size - patch_size: | |
| s.append(size - patch_size) | |
| return s | |
| for y in steps(H): | |
| for x in steps(W): | |
| pm = mask[y : y + patch_size, x : x + patch_size] | |
| pi = image[:, y : y + patch_size, x : x + patch_size] | |
| if pm.shape == (patch_size, patch_size) and (pm != IGNORE_INDEX).any(): | |
| imgs.append(pi) | |
| masks.append(pm) | |
| if not imgs: | |
| ph = min(patch_size, H) | |
| pw = min(patch_size, W) | |
| img_pad = np.zeros((image.shape[0], patch_size, patch_size), dtype=np.float32) | |
| mask_pad = np.full((patch_size, patch_size), IGNORE_INDEX, dtype=np.int64) | |
| img_pad[:, :ph, :pw] = image[:, :ph, :pw] | |
| mask_pad[:ph, :pw] = mask[:ph, :pw] | |
| imgs.append(img_pad) | |
| masks.append(mask_pad) | |
| return np.stack(imgs).astype(np.float32), np.stack(masks).astype(np.int64) | |
| # ── Spectral analysis ───────────────────────────────────────── | |
| def compute_spectral_signatures(full_image: np.ndarray, full_mask: np.ndarray) -> Dict: | |
| """Per-class mean and std across the 7 normalized bands, from labeled pixels.""" | |
| sigs = {} | |
| for cls_idx in range(NUM_CLASSES): | |
| px = full_mask == cls_idx | |
| if px.sum() == 0: | |
| sigs[cls_idx] = { | |
| "mean": np.zeros(NUM_CHANNELS, dtype=np.float32), | |
| "std": np.zeros(NUM_CHANNELS, dtype=np.float32), | |
| "n": 0, | |
| } | |
| else: | |
| vals = full_image[:, px] # (7, N) | |
| sigs[cls_idx] = { | |
| "mean": vals.mean(axis=1).astype(np.float32), | |
| "std": vals.std(axis=1).astype(np.float32), | |
| "n": int(px.sum()), | |
| } | |
| return sigs | |
| # ── Dataset class ───────────────────────────────────────────── | |
| class MultiSpectralDataset(Dataset): | |
| def __init__(self, images: np.ndarray, masks: np.ndarray): | |
| self.images = images.astype(np.float32) | |
| self.masks = masks.astype(np.int64) | |
| def __len__(self): | |
| return len(self.images) | |
| def __getitem__(self, idx: int): | |
| return torch.from_numpy(self.images[idx]), torch.from_numpy(self.masks[idx]) | |
| # ── Public API ──────────────────────────────────────────────── | |
| def load_data(patch_size: int = DEFAULT_PATCH_SIZE) -> Dict: | |
| # Read raw bands | |
| raw_bands = [_read_band(_get_path(f)) for f in BAND_FILES] | |
| raw_image = np.stack(raw_bands, axis=0) # (7, H, W) raw float32 | |
| # Compute spectral indices from raw values (ratio, so must use raw) | |
| ndvi = _compute_ndvi(raw_image) | |
| ndwi = _compute_ndwi(raw_image) | |
| # Normalize | |
| image = _normalize(raw_image) # (7, H, W) normalized [0,1] | |
| # Read masks | |
| raw_train, nd_train, info_train = _read_mask_raw(_get_path(TRAIN_MASK_FILE)) | |
| raw_val, nd_val, info_val = _read_mask_raw(_get_path(VAL_MASK_FILE)) | |
| # Remap | |
| train_mask, train_vals = _remap_mask(raw_train, nd_train) | |
| val_mask, val_vals = _remap_mask(raw_val, nd_val) | |
| if not train_vals: | |
| raise ValueError( | |
| f"TRAINING.tif has no labeled pixels. Info: {info_train} | " | |
| f"Unique raw values: {np.unique(raw_train).tolist()}" | |
| ) | |
| # Patches | |
| tr_imgs, tr_masks = _extract_patches(image, train_mask, patch_size) | |
| va_imgs, va_masks = _extract_patches(image, val_mask, patch_size) | |
| # Spectral signatures from training labels | |
| signatures = compute_spectral_signatures(image, train_mask) | |
| train_labeled = int((train_mask != IGNORE_INDEX).sum()) | |
| val_labeled = int((val_mask != IGNORE_INDEX).sum()) | |
| def _class_dist(mask, total): | |
| parts = [] | |
| for i, name in enumerate(CLASS_NAMES): | |
| n = int((mask == i).sum()) | |
| parts.append(f"{name}: {n:,} ({n / max(1, total) * 100:.1f}%)") | |
| return " | ".join(parts) | |
| status = "\n".join([ | |
| f"Train patches: **{len(tr_imgs)}** | Val patches: **{len(va_imgs)}** | Patch: **{patch_size}x{patch_size}**", | |
| f"Training labels: **{train_labeled:,}** px — {_class_dist(train_mask, train_labeled)}", | |
| f"Validation labels: **{val_labeled:,}** px — {_class_dist(val_mask, val_labeled)}", | |
| ]) | |
| return { | |
| "full_image": image, | |
| "full_train_mask": train_mask, | |
| "full_val_mask": val_mask, | |
| "ndvi": ndvi, | |
| "ndwi": ndwi, | |
| "signatures": signatures, | |
| "train_images": tr_imgs, | |
| "train_masks": tr_masks, | |
| "val_images": va_imgs, | |
| "val_masks": va_masks, | |
| "status": status, | |
| } | |