Spaces:
Running on Zero
Running on Zero
File size: 2,633 Bytes
ad38b1a | 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 | import os
import logging
import numpy as np
from typing import Optional, Dict, Any, Tuple
from PIL import Image, UnidentifiedImageError
import rasterio
from rasterio.warp import reproject, Resampling
def load_aligned_dsm(rgb_path, dsm_path):
"""
Geospatially aligns a DSM to an RGB reference image.
Ensures identical CRS, bounding box, and pixel dimensions.
"""
with rasterio.open(rgb_path) as src_rgb:
rgb_transform = src_rgb.transform
rgb_crs = src_rgb.crs
h, w = src_rgb.height, src_rgb.width
with rasterio.open(dsm_path) as src_dsm:
# Create an empty array perfectly matching the RGB dimensions
aligned_dsm = np.zeros((h, w), dtype=np.float32)
# Reproject using actual geographic metadata, not just array stretching
reproject(
source=rasterio.band(src_dsm, 1),
destination=aligned_dsm,
src_transform=src_dsm.transform,
src_crs=src_dsm.crs,
dst_transform=rgb_transform,
dst_crs=rgb_crs,
resampling=Resampling.nearest # Nearest neighbor prevents interpolating fake elevations
)
return aligned_dsm
def load_and_validate_image(image_path: str) -> Optional[Tuple[Image.Image, Dict[str, Any]]]:
if not os.path.exists(image_path):
logging.error(f"[M1 PREPROCESS] File not found at path: {image_path}")
return None
# Default metadata dictionary
metadata = {"is_georeferenced": False, "crs": None, "transform": None}
try:
# Check if it is a georeferenced TIFF
if image_path.lower().endswith(('.tif', '.tiff')) and rasterio is not None:
with rasterio.open(image_path) as src:
# Read first 3 bands (RGB)
r, g, b = src.read(1), src.read(2), src.read(3)
# Stack and convert to PIL Image for the HuggingFace pipeline
rgb_array = np.dstack((r, g, b))
image = Image.fromarray(rgb_array)
metadata["is_georeferenced"] = True
metadata["crs"] = src.crs.to_string() if src.crs else None
metadata["transform"] = src.transform
logging.info("[M1 PREPROCESS] GeoTIFF loaded with spatial metadata.")
return image, metadata
# Fallback for standard PNG/JPG
image = Image.open(image_path).convert("RGB")
logging.info("[M1 PREPROCESS] Standard optical image loaded (No spatial metadata).")
return image, metadata
except Exception as e:
logging.error(f"[M1 PREPROCESS] Error: {str(e)}")
return None
|