Spaces:
Paused
Paused
Download backend/preprocessing/raster_to_rgb.py from Sahil-nith/SatQuery-api: direct link, hf CLI and curl.
- Browser
- Download file 2.19 kB
-
https://huggingface.co/spaces/Sahil-nith/SatQuery-api/resolve/main/backend/preprocessing/raster_to_rgb.py
- Command line
-
hf download hf://spaces/Sahil-nith/SatQuery-api/backend/preprocessing/raster_to_rgb.py
-
curl -L -o raster_to_rgb.py https://huggingface.co/spaces/Sahil-nith/SatQuery-api/resolve/main/backend/preprocessing/raster_to_rgb.py
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") | |