Spaces:
Running
Running
File size: 6,805 Bytes
eebd638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """ConstructionChangeDataset — loads (tile_2022, tile_2024, mask_2022, mask_2024) pairs.
Directory layout expected (produced by PDOKDownloader)::
tile_dir/
<site_id>/
tile_2022.tif
tile_2024.tif
mask_2022.png ← binary mask from 2022 construction polygon
mask_2024.png ← optional; falls back to mask_2022 if absent
meta.json ← {"sample_type": "preexisting"|"new"|"negative"}
sample_type meanings:
preexisting — construction existed in 2022, may still be active in 2024
new — no construction in 2022 tile, but construction appeared by 2024
negative — non-construction tile (farmland, water, road, etc.)
"""
from __future__ import annotations
import json
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image
from data.transforms import get_train_transforms, get_inference_transforms
try:
import rasterio
_HAS_RASTERIO = True
except ImportError:
_HAS_RASTERIO = False
logger = logging.getLogger(__name__)
class ConstructionChangeDataset(Dataset):
"""Paired (2022, 2024) aerial tile dataset for construction change detection."""
def __init__(
self,
tile_dir: str,
image_size: int = 1024,
augment: bool = True,
min_fg_pixels: int = 500,
split: Optional[str] = None,
val_split: float = 0.15,
seed: int = 42,
) -> None:
"""
Args:
tile_dir: Root directory of downloaded tiles.
image_size: Spatial size fed to the model.
augment: If True use training augmentations, else inference-only.
min_fg_pixels: Skip samples whose mask_2022 has fewer foreground px.
split: "train" | "val" | None (use all).
val_split: Fraction of sites to hold out for validation.
seed: Random seed for the train/val split.
"""
self.tile_dir = Path(tile_dir)
self.image_size = image_size
self.augment = augment
self.transforms = (
get_train_transforms(image_size) if augment
else get_inference_transforms(image_size)
)
all_samples = self._scan(min_fg_pixels)
if split is not None:
rng = np.random.default_rng(seed)
indices = rng.permutation(len(all_samples))
n_val = max(1, int(len(all_samples) * val_split))
val_idx = set(indices[:n_val].tolist())
if split == "val":
all_samples = [all_samples[i] for i in val_idx]
else:
all_samples = [s for i, s in enumerate(all_samples) if i not in val_idx]
self.samples: List[Dict] = all_samples
logger.info("Dataset (%s split): %d samples", split or "all", len(self.samples))
# ------------------------------------------------------------------
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
s = self.samples[idx]
tile_2024 = self._load_image(s["tile_b"]) # query year
tile_2022 = self._load_image(s["tile_a"]) # reference year
mask_2022 = self._load_mask(s["mask_a"])
mask_2024 = self._load_mask(s["mask_b"])
result = self.transforms(
image=tile_2024,
image2=tile_2022,
mask=mask_2024,
mask2=mask_2022,
)
return {
"tile_2024": result["image"], # [3,H,W]
"tile_2022": result["image2"], # [3,H,W]
"mask_2024": result["mask"].unsqueeze(0).float() / 255.0, # [1,H,W] {0,1}
"mask_2022": result["mask2"].unsqueeze(0).float() / 255.0, # [1,H,W] {0,1}
"sample_type": s["sample_type"],
"site_id": s["site_id"],
}
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _scan(self, min_fg_pixels: int) -> List[Dict]:
samples = []
for site_dir in sorted(self.tile_dir.iterdir()):
if not site_dir.is_dir():
continue
meta_path = site_dir / "meta.json"
meta = {}
if meta_path.exists():
with open(meta_path) as f:
meta = json.load(f)
year_a = meta.get("year_a", 2022)
year_b = meta.get("year_b", 2024)
tile_a = site_dir / f"tile_{year_a}.tif"
tile_b = site_dir / f"tile_{year_b}.tif"
mask_a = site_dir / f"mask_{year_a}.png"
# mask_2024 is optional — fall back to mask_2022 for preexisting sites
mask_b = site_dir / f"mask_{year_b}.png"
if not mask_b.exists():
mask_b = mask_a
if not (tile_a.exists() and tile_b.exists() and mask_a.exists()):
continue
# Skip tiny masks
mask_arr = np.array(Image.open(mask_a))
if mask_arr.sum() < min_fg_pixels * 255:
logger.debug("Skipping %s — mask too small", site_dir.name)
continue
samples.append({
"site_id": site_dir.name,
"tile_a": str(tile_a),
"tile_b": str(tile_b),
"mask_a": str(mask_a),
"mask_b": str(mask_b),
"sample_type": meta.get("sample_type", "preexisting"),
})
return samples
@staticmethod
def _load_image(path: str) -> np.ndarray:
"""Load image as HxWx3 uint8 numpy array.
Handles both standard images (.jpg, .png) and GeoTIFFs (.tif).
"""
if path.endswith(".tif") or path.endswith(".tiff"):
if _HAS_RASTERIO:
with rasterio.open(path) as src:
# Read first 3 bands (RGB)
n = min(src.count, 3)
bands = src.read(list(range(1, n + 1))) # [C, H, W]
img = np.transpose(bands, (1, 2, 0)) # [H, W, C]
if img.dtype != np.uint8:
img = np.clip(img, 0, 255).astype(np.uint8)
return img
return np.array(Image.open(path).convert("RGB"))
@staticmethod
def _load_mask(path: str) -> np.ndarray:
"""Load binary mask as HxW uint8 (values 0 or 255)."""
mask = np.array(Image.open(path).convert("L"))
return (mask > 127).astype(np.uint8) * 255
|