Spaces:
Running on Zero
Running on Zero
| 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 | |