File size: 6,528 Bytes
e854df7 | 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 187 188 189 190 191 192 193 194 195 196 197 198 199 | """
Preprocessing for SEN2SR inference.
Steps
-----
1. Normalize : uint16 → float32, divide by 10 000, clip to [0, 1].
2. Patch grid : compute (row, col) top-left coordinates for a sliding window
with configurable size and stride.
3. Extract : cut patches from the full array; optionally use multiprocessing
for the extraction step (useful when patch count is large and
copy overhead dominates GPU transfer).
Design note
-----------
Patch extraction is kept **separate** from model inference so that large images
(> 10 000 × 10 000 px) never have to be fully loaded into GPU memory.
The caller controls batch assembly.
"""
from __future__ import annotations
import multiprocessing as mp
import warnings
from dataclasses import dataclass, field
from functools import partial
from typing import List, Tuple
import numpy as np
from s2sr_pipe.utils.logging_utils import get_logger
logger = get_logger("preprocessing")
# ── Normalisation ─────────────────────────────────────────────────────────────
NORM_SCALE: float = 10_000.0
def normalize(array: np.ndarray) -> np.ndarray:
"""
Normalize a uint16 Sentinel-2 array to float32 in [0, 1].
Parameters
----------
array : np.ndarray — shape (C, H, W), dtype uint16 or float32.
Returns
-------
np.ndarray — shape (C, H, W), dtype float32, values in [0, 1].
"""
if array.dtype == np.float32 and float(array.max()) < 2.0:
warnings.warn(
"normalize() reçoit un tableau float32 avec des valeurs < 2.0 — "
"il est peut-être déjà normalisé. La division par 10000 va produire "
"des valeurs quasi nulles.",
UserWarning,
stacklevel=2,
)
out = array.astype(np.float32) / NORM_SCALE
np.clip(out, 0.0, 1.0, out=out)
return out
def denormalize(array: np.ndarray) -> np.ndarray:
"""
Reverse normalisation: float32 [0, 1] → uint16 [0, 10 000].
Parameters
----------
array : np.ndarray — shape (C, H, W), dtype float32.
Returns
-------
np.ndarray — shape (C, H, W), dtype uint16.
"""
out = array * NORM_SCALE
np.clip(out, 0, 10_000, out=out)
return out.astype(np.uint16)
# ── Patch grid ────────────────────────────────────────────────────────────────
@dataclass
class PatchConfig:
"""Configuration for sliding-window patch extraction."""
patch_size: int = 128
stride: int = 64
def __post_init__(self) -> None:
if self.stride > self.patch_size:
raise ValueError(
f"stride ({self.stride}) must be ≤ patch_size ({self.patch_size})"
)
if self.patch_size <= 0 or self.stride <= 0:
raise ValueError("patch_size and stride must be positive integers.")
def compute_patch_coords(
height: int,
width: int,
cfg: PatchConfig,
) -> List[Tuple[int, int]]:
"""
Compute (row, col) top-left coordinates for all patches.
The window always covers the full image: if the last patch would go
beyond the border it is snapped inward (no border reflection or padding).
Parameters
----------
height, width : Image spatial dimensions (after band alignment).
cfg : PatchConfig instance.
Returns
-------
List of (row, col) tuples.
"""
ps, st = cfg.patch_size, cfg.stride
if height < ps or width < ps:
raise ValueError(
f"L'image ({height}×{width} px) est plus petite que patch_size ({ps}). "
"Réduire patch_size ou utiliser une image plus grande."
)
rows = list(range(0, height - ps, st)) + [height - ps]
cols = list(range(0, width - ps, st)) + [width - ps]
# Deduplicate while preserving order
rows = sorted(set(rows))
cols = sorted(set(cols))
coords = [(r, c) for r in rows for c in cols]
logger.debug(
"Patch grid: %d rows x %d cols = %d patches "
"(patch_size=%d, stride=%d, image=%dx%d)",
len(rows),
len(cols),
len(coords),
ps,
st,
width,
height,
)
return coords
# ── Patch extraction ──────────────────────────────────────────────────────────
def _extract_one(
coord: Tuple[int, int],
array: np.ndarray,
patch_size: int,
) -> np.ndarray:
"""Extract a single (C, patch_size, patch_size) patch — pickleable worker."""
r, c = coord
return array[:, r : r + patch_size, c : c + patch_size].copy()
def extract_patches(
array: np.ndarray,
coords: List[Tuple[int, int]],
patch_size: int,
num_workers: int = 0,
) -> List[np.ndarray]:
"""
Extract all patches from *array* at positions given by *coords*.
Parameters
----------
array : (C, H, W) float32 normalised array.
coords : List of (row, col) top-left positions.
patch_size : Spatial size of each patch.
num_workers : Number of worker processes.
0 = single-threaded (safe for small images or debugging).
>0 = use multiprocessing.Pool (useful for very large images
with hundreds of thousands of patches).
Returns
-------
List[np.ndarray] — each element has shape (C, patch_size, patch_size).
Notes
-----
* Multiprocessing on large numpy arrays uses shared memory via fork (Linux).
On Windows/macOS, the array is serialised per chunk — use a smaller
num_workers or num_workers=0 if memory is limited on those platforms.
* For GPU-bound workloads the extraction is rarely the bottleneck;
num_workers=0 is the safe default.
"""
worker = partial(_extract_one, array=array, patch_size=patch_size)
if num_workers > 0:
logger.info("Extracting %d patches using %d workers ...", len(coords), num_workers)
with mp.Pool(num_workers) as pool:
patches = pool.map(worker, coords)
else:
logger.info("Extracting %d patches (single-threaded) ...", len(coords))
patches = [worker(c) for c in coords]
logger.info("Patch extraction complete.")
return patches
|