Spaces:
Running on Zero
Running on Zero
| from pathlib import Path | |
| import numpy as np | |
| import rasterio | |
| from rasterio.transform import Affine | |
| from PIL import Image | |
| # ========================================================== | |
| # Read Sentinel-2 RGB bands | |
| # ========================================================== | |
| def read_sentinel_rgb( | |
| b02_path, | |
| b03_path, | |
| b04_path | |
| ): | |
| """ | |
| Read Sentinel-2 B02, B03 and B04 GeoTIFF bands. | |
| Band mapping: | |
| B04 -> Red | |
| B03 -> Green | |
| B02 -> Blue | |
| Returns | |
| ------- | |
| rgb : numpy.ndarray | |
| RGB array with shape (H, W, 3) | |
| profile : rasterio profile | |
| Original geospatial metadata. | |
| """ | |
| b02_path = Path(b02_path) | |
| b03_path = Path(b03_path) | |
| b04_path = Path(b04_path) | |
| with rasterio.open(b02_path) as src_b02: | |
| b02 = src_b02.read(1) | |
| profile = src_b02.profile.copy() | |
| transform = src_b02.transform | |
| crs = src_b02.crs | |
| width = src_b02.width | |
| height = src_b02.height | |
| with rasterio.open(b03_path) as src_b03: | |
| b03 = src_b03.read(1) | |
| with rasterio.open(b04_path) as src_b04: | |
| b04 = src_b04.read(1) | |
| # Make sure all bands have identical dimensions. | |
| if b02.shape != b03.shape or b02.shape != b04.shape: | |
| raise ValueError( | |
| "B02, B03 and B04 dimensions do not match." | |
| ) | |
| # ------------------------------------------------------ | |
| # Sentinel-2 RGB order: | |
| # Red = B04 | |
| # Green = B03 | |
| # Blue = B02 | |
| # ------------------------------------------------------ | |
| rgb = np.stack( | |
| [ | |
| b04, | |
| b03, | |
| b02 | |
| ], | |
| axis=-1 | |
| ) | |
| return rgb, profile | |
| # ========================================================== | |
| # Convert raw Sentinel data to visualization RGB | |
| # ========================================================== | |
| def normalize_rgb( | |
| rgb, | |
| lower_percentile=2, | |
| upper_percentile=98 | |
| ): | |
| """ | |
| Convert Sentinel-2 reflectance/DN values to uint8 RGB. | |
| Percentile normalization is performed independently | |
| for each channel. | |
| """ | |
| rgb = rgb.astype(np.float32) | |
| result = np.zeros_like( | |
| rgb, | |
| dtype=np.uint8 | |
| ) | |
| for channel in range(3): | |
| band = rgb[:, :, channel] | |
| low = np.percentile( | |
| band, | |
| lower_percentile | |
| ) | |
| high = np.percentile( | |
| band, | |
| upper_percentile | |
| ) | |
| if high <= low: | |
| result[:, :, channel] = np.clip( | |
| band, | |
| 0, | |
| 255 | |
| ).astype(np.uint8) | |
| else: | |
| normalized = ( | |
| (band - low) / | |
| (high - low) | |
| ) | |
| normalized = np.clip( | |
| normalized, | |
| 0, | |
| 1 | |
| ) | |
| result[:, :, channel] = ( | |
| normalized * 255 | |
| ).astype(np.uint8) | |
| return result | |
| # ========================================================== | |
| # GeoTIFF writer | |
| # ========================================================== | |
| def save_sr_geotiff( | |
| sr_image, | |
| source_profile, | |
| output_path, | |
| scale=4 | |
| ): | |
| """ | |
| Save a super-resolved RGB image as a GeoTIFF. | |
| The geographic extent remains unchanged while | |
| the pixel dimensions increase by `scale`. | |
| Parameters | |
| ---------- | |
| sr_image : PIL.Image.Image or numpy.ndarray | |
| source_profile : dict | |
| Rasterio profile from the original Sentinel image. | |
| output_path : str or Path | |
| scale : int | |
| Super-resolution scale factor. | |
| """ | |
| output_path = Path(output_path) | |
| output_path.parent.mkdir( | |
| parents=True, | |
| exist_ok=True | |
| ) | |
| # ------------------------------------------------------ | |
| # Convert image to numpy | |
| # ------------------------------------------------------ | |
| if isinstance( | |
| sr_image, | |
| Image.Image | |
| ): | |
| rgb = np.asarray( | |
| sr_image.convert("RGB") | |
| ) | |
| else: | |
| rgb = np.asarray( | |
| sr_image | |
| ) | |
| if rgb.ndim != 3 or rgb.shape[2] != 3: | |
| raise ValueError( | |
| "sr_image must have shape (H, W, 3)." | |
| ) | |
| rgb = np.clip( | |
| rgb, | |
| 0, | |
| 255 | |
| ).astype(np.uint8) | |
| height, width, _ = rgb.shape | |
| # ------------------------------------------------------ | |
| # Original transform | |
| # ------------------------------------------------------ | |
| original_transform = source_profile[ | |
| "transform" | |
| ] | |
| # ------------------------------------------------------ | |
| # Scale pixel size by 1/scale | |
| # | |
| # Geographic bounds stay unchanged. | |
| # ------------------------------------------------------ | |
| new_transform = ( | |
| original_transform * | |
| Affine.scale( | |
| 1.0 / scale, | |
| 1.0 / scale | |
| ) | |
| ) | |
| # ------------------------------------------------------ | |
| # Build output profile | |
| # ------------------------------------------------------ | |
| profile = source_profile.copy() | |
| profile.update( | |
| { | |
| "driver": "GTiff", | |
| "height": height, | |
| "width": width, | |
| "count": 3, | |
| "dtype": "uint8", | |
| "transform": new_transform, | |
| "compress": "deflate", | |
| "photometric": "RGB", | |
| "tiled": True, | |
| "blockxsize": 256, | |
| "blockysize": 256, | |
| } | |
| ) | |
| # Remove incompatible fields where necessary. | |
| profile.pop( | |
| "nodata", | |
| None | |
| ) | |
| # ------------------------------------------------------ | |
| # Write GeoTIFF | |
| # ------------------------------------------------------ | |
| # Remove block settings when not using tiled GeoTIFF | |
| profile.pop("blockxsize", None) | |
| profile.pop("blockysize", None) | |
| profile.pop("tiled", None) | |
| with rasterio.open( | |
| output_path, | |
| "w", | |
| **profile | |
| ) as dst: | |
| dst.write( | |
| rgb[:, :, 0], | |
| 1 | |
| ) | |
| dst.write( | |
| rgb[:, :, 1], | |
| 2 | |
| ) | |
| dst.write( | |
| rgb[:, :, 2], | |
| 3 | |
| ) | |
| return output_path | |
| # ========================================================== | |
| # Convenience function | |
| # ========================================================== | |
| def sentinel_rgb_from_bands( | |
| b02_path, | |
| b03_path, | |
| b04_path | |
| ): | |
| """ | |
| Read B02/B03/B04 and return a normalized PIL RGB image | |
| together with the source GeoTIFF profile. | |
| """ | |
| rgb, profile = read_sentinel_rgb( | |
| b02_path, | |
| b03_path, | |
| b04_path | |
| ) | |
| rgb_uint8 = normalize_rgb(rgb) | |
| image = Image.fromarray( | |
| rgb_uint8, | |
| mode="RGB" | |
| ) | |
| return image, profile | |