Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from typing import Tuple | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from scipy.ndimage import gaussian_filter | |
| def _to_numpy(x: np.ndarray) -> np.ndarray: | |
| """Convert input to contiguous float32 numpy array.""" | |
| if hasattr(x, "detach") and hasattr(x, "cpu"): | |
| x = x.detach().cpu().numpy() | |
| x = np.asarray(x) | |
| if x.dtype != np.float32: | |
| x = x.astype(np.float32, copy=False) | |
| return x | |
| def _biome_palette() -> np.ndarray: | |
| """Return (31,3) float32 RGB palette for Koppen-Geiger classes with natural tones. | |
| Index 0 is Unknown. | |
| """ | |
| # Colors chosen for realism (earth tones), normalized to 0..1 | |
| lut = np.array( | |
| [ | |
| [0, 0, 0], # 0: Unknown | |
| [16, 86, 24], # 1: Af Tropical rainforest - deep green | |
| [38, 120, 40], # 2: Am Tropical monsoon - lush green | |
| [187, 212, 92], # 3: Aw Tropical savannah - yellow-green | |
| [227, 192, 122], # 4: BWh Arid desert hot - sand | |
| [217, 200, 163], # 5: BWk Arid desert cold - pale sand | |
| [210, 168, 90], # 6: BSh Steppe hot - ochre | |
| [203, 182, 136], # 7: BSk Steppe cold - tan | |
| [176, 156, 78], # 8: Csa Med dry hot summer - olive | |
| [162, 148, 84], # 9: Csb Med dry warm summer - olive-brown | |
| [148, 140, 104], # 10: Csc Med dry cold summer - muted olive | |
| [132, 178, 96], # 11: Cwa Temp dry winter hot summer - light green | |
| [112, 164, 96], # 12: Cwb Temp dry winter warm summer - green | |
| [96, 148, 96], # 13: Cwc Temp dry winter cold summer - darker green | |
| [124, 186, 84], # 14: Cfa Temp no dry hot summer - bright green | |
| [96, 168, 84], # 15: Cfb Temp no dry warm summer - temperate green | |
| [76, 140, 76], # 16: Cfc Temp no dry cold summer - dark green | |
| [120, 140, 160], # 17: Dsa Cold dry summer hot summer - cool grey-green | |
| [108, 130, 150], # 18: Dsb Cold dry summer warm summer - cool grey-green | |
| [96, 120, 140], # 19: Dsc Cold dry summer cold summer - cool slate | |
| [88, 112, 132], # 20: Dsd Cold dry summer very cold winter - slate | |
| [136, 152, 176], # 21: Dwa Cold dry winter hot summer - cool blue-grey | |
| [112, 136, 168], # 22: Dwb Cold dry winter warm summer - blue-grey | |
| [100, 120, 160], # 23: Dwc Cold dry winter cold summer - blue slate | |
| [84, 104, 140], # 24: Dwd Cold dry winter very cold winter - deep blue slate | |
| [120, 170, 120], # 25: Dfa Cold no dry hot summer - mixed forest | |
| [96, 150, 120], # 26: Dfb Cold no dry warm summer - boreal edge | |
| [72, 120, 110], # 27: Dfc Cold no dry cold summer - boreal | |
| [64, 96, 108], # 28: Dfd Cold no dry very cold winter - dark boreal | |
| [173, 180, 180], # 29: ET Polar tundra - grey-green tundra | |
| [230, 238, 244], # 30: EF Polar frost - ice/snow | |
| ], | |
| dtype=np.float32, | |
| ) / 255.0 | |
| return lut | |
| def get_relief_map( | |
| elevation: np.ndarray, | |
| climate: np.ndarray, | |
| biome: np.ndarray, | |
| flow: np.ndarray, | |
| *, | |
| azimuths: Tuple[float, float, float, float] = (315.0, 45.0, 135.0, 225.0), | |
| flow_threshold: float = 7, | |
| sigma_large: float = 6.0, | |
| sigma_small: float = 1.2, | |
| resolution: float = 90, | |
| rgb: np.ndarray | None = None, | |
| relief: float = 1.0, | |
| vmin: float | None = None, | |
| vmax: float | None = None, | |
| ) -> Tuple[plt.Figure, plt.Axes]: | |
| """Plot a GDAL-style shaded relief map using Matplotlib, with optional river overlay. | |
| Args: | |
| elevation: (H, W) float meters. | |
| climate: unused. | |
| biome: unused. | |
| flow: (H, W) flow accumulation; rivers drawn where flow > flow_threshold. | |
| flow_threshold: threshold for river mask. | |
| Returns: | |
| (fig, ax): Matplotlib figure and GeoAxes containing the rendered map. | |
| """ | |
| elev = _to_numpy(elevation) | |
| assert elev.ndim == 2, "elevation must be (H, W)" | |
| H, W = elev.shape | |
| flow_np = None | |
| if flow is not None: | |
| flow_np = _to_numpy(flow) | |
| assert flow_np.shape == elev.shape, "flow must be (H, W) matching elevation" | |
| # Hillshade (GDAL-style) parameters | |
| azimuth_deg = float(azimuths[0]) if isinstance(azimuths, (tuple, list)) and len(azimuths) > 0 else 315.0 | |
| altitude_deg = 45.0 # sun elevation angle | |
| # Replace NaNs before any processing | |
| elev_f32 = elev.astype(np.float32, copy=False) | |
| if np.isnan(elev_f32).any(): | |
| median_val = float(np.nanmedian(elev_f32)) if np.isfinite(np.nanmedian(elev_f32)) else 0.0 | |
| elev_f32 = np.nan_to_num(elev_f32, nan=median_val) | |
| def compute_hillshade(src: np.ndarray) -> np.ndarray: | |
| dy, dx = np.gradient(src) | |
| dy, dx = dy/(15 * resolution/90), dx/(15 * resolution/90) | |
| slope_rad = np.pi / 2.0 - np.arctan(np.hypot(dx, dy)) | |
| aspect_rad = np.arctan2(dy, -dx) | |
| az_rad = np.deg2rad(azimuth_deg) | |
| alt_rad = np.deg2rad(altitude_deg) | |
| hs = ( | |
| np.sin(alt_rad) * np.sin(slope_rad) | |
| + np.cos(alt_rad) * np.cos(slope_rad) * np.cos(az_rad - aspect_rad) | |
| ) | |
| return np.clip(hs, 0.0, 1.0).astype(np.float32) | |
| # Multi-scale hillshade: emphasize large landforms, suppress pixel-scale roughness | |
| elev_large = gaussian_filter(elev_f32, sigma=sigma_large) | |
| elev_small = gaussian_filter(elev_f32, sigma=sigma_small) | |
| hs_large = compute_hillshade(elev_large) | |
| hs_small = compute_hillshade(elev_small) | |
| hillshade = np.clip(0.75 * hs_large + 0.25 * hs_small, 0.0, 1.0) | |
| hillshade = np.power(hillshade, 0.85) # gentle gamma to lift broad features | |
| # Colorize elevation; this will be used where biome is unknown | |
| if rgb is None: | |
| land_elev = np.maximum(0, elev) | |
| if vmin is None or vmax is None: | |
| _vmin = float(np.nanmin(land_elev)) | |
| _vmax = float(np.nanmax(land_elev)) | |
| if not np.isfinite(_vmin) or not np.isfinite(_vmax) or _vmax == _vmin: | |
| _vmin, _vmax = 0.0, 1.0 | |
| else: | |
| _vmin, _vmax = max(0.0, float(vmin)), float(vmax) | |
| norm = (land_elev - _vmin) / (_vmax - _vmin + 1e-8) | |
| cmap = plt.get_cmap("terrain") | |
| # terrain cmap 0–0.25 is water-blue; when vmin=0 (absolute scale) map | |
| # land to the 0.25–1.0 range so sea level starts at lowland green. | |
| if _vmin == 0.0: | |
| norm_cmap = 0.25 + np.clip(norm ** 0.7, 0.0, 1.0) * 0.75 | |
| else: | |
| norm_cmap = np.clip(norm ** 0.7, 0.0, 1.0) | |
| rgb = cmap(norm_cmap)[..., :3].astype(np.float32) | |
| # Base RGB: prefer biome colors when available, otherwise elevation colormap | |
| base_rgb = rgb | |
| if biome is not None: | |
| b_idx = _to_numpy(biome).astype(np.int32, copy=False) | |
| if b_idx.shape == elev.shape: | |
| lut = _biome_palette() | |
| b_idx = np.clip(b_idx, 0, lut.shape[0] - 1) | |
| mask = b_idx > 0 | |
| if np.any(mask): | |
| biome_rgb = lut[b_idx] | |
| mask3 = mask[..., None] | |
| base_rgb = np.where(mask3, biome_rgb, base_rgb) | |
| # GDAL-like intensity blend (ambient term + directional light) | |
| intensity = 0.35 + 0.65 * hillshade # slightly higher ambient to reduce ragged contrast | |
| shaded_rgb = np.clip(base_rgb * (relief * intensity + (1 - relief))[..., None], 0.0, 1.0) | |
| shaded_rgb[np.isnan(elev)] = np.nan | |
| # Optional blue river overlay where flow exceeds threshold | |
| if flow_np is not None: | |
| river_mask = flow_np > float(flow_threshold) | |
| if np.any(river_mask): | |
| river_color = np.array([0.100, 0.450, 0.850], dtype=np.float32) | |
| river_alpha = 0.75 | |
| shaded_rgb[river_mask] = ( | |
| (1.0 - river_alpha) * shaded_rgb[river_mask] | |
| + river_alpha * river_color[None, :] | |
| ) | |
| # Ocean coloring: fade from light blue (coast) to dark blue (deep ocean) | |
| ocean_mask = elev_f32 < 0.0 | |
| if np.any(ocean_mask): | |
| depth = -elev_f32 # positive depth below sea level | |
| max_depth = 10_000.0 | |
| if max_depth > 0.0: | |
| t = np.zeros_like(elev_f32, dtype=np.float32) | |
| t[ocean_mask] = np.clip(depth[ocean_mask] / max_depth, 0.0, 1.0) | |
| # Bias toward deeper blue sooner near the coast | |
| t = t ** 0.7 | |
| t3 = t[..., None] | |
| # More saturated blues | |
| coast_color = np.array([0.68, 0.88, 1.00], dtype=np.float32) # lighter, bluer coast | |
| deep_color = np.array([0.00, 0.10, 0.45], dtype=np.float32) # deeper blue | |
| ocean_rgb = (1.0 - t3) * coast_color + t3 * deep_color | |
| shaded_rgb = np.where(ocean_mask[..., None], ocean_rgb, shaded_rgb) | |
| return shaded_rgb | |
| __all__ = ["get_relief_map"] | |