SatQuery-api / backend /preprocessing /raster_to_rgb.py
Sahil-nith's picture
Upload 47 files
0127c21 verified
Raw History Blame Contribute Delete
2.19 kB
"""
Raster -> PIL RGB conversion.
Phase 5. Every specialist vision model we integrate (VQA now, captioning/
grounding/change/fusion later) needs a plain RGB image, but uploads can be
1-band SAR, multi-band optical GeoTIFF, or 3-band PNG/JPEG benchmark
images. This is the single shared place that turns "whatever Rasterio can
open" into "3-channel uint8 PIL.Image" so every models/*.py file stays
model-only and doesn't reimplement band handling.
Band selection is intentionally simple (first 3 bands, or band 1 repeated
for single-band SAR/panchromatic). This is a reasonable prototype default,
not a radiometric-correctness claim — do not present the visualization as
calibrated reflectance.
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
import rasterio
from PIL import Image
def _to_uint8(band: np.ndarray) -> np.ndarray:
"""Percentile-stretch a single band to 0-255 uint8 for display/model input."""
band = band.astype("float32")
finite = band[np.isfinite(band)]
if finite.size == 0:
return np.zeros_like(band, dtype="uint8")
lo, hi = np.percentile(finite, [2, 98])
if hi <= lo:
lo, hi = float(finite.min()), float(finite.max() or 1.0)
stretched = np.clip((band - lo) / max(hi - lo, 1e-6), 0, 1)
return (stretched * 255).astype("uint8")
# Public alias — Phase 7/8 need single-band percentile stretching directly
# (grayscale diffing, SAR backscatter) without going through load_as_rgb's
# 3-channel stacking. Kept as a thin re-export so there's one implementation.
to_uint8 = _to_uint8
def load_as_rgb(path: Path) -> Image.Image:
"""Open any raster Phase 3 accepted and return a 3-channel PIL RGB image."""
with rasterio.open(path) as src:
count = src.count
if count >= 3:
r, g, b = (src.read(i) for i in (1, 2, 3))
elif count == 1:
band = src.read(1)
r = g = b = band
else: # 2-band, uncommon — reuse band 1 for the missing channel
r = g = src.read(1)
b = src.read(2)
rgb = np.dstack([_to_uint8(r), _to_uint8(g), _to_uint8(b)])
return Image.fromarray(rgb, mode="RGB")