File size: 4,059 Bytes
80e070e | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import os
import random
import numpy as np
import rasterio
import matplotlib.pyplot as plt
from pathlib import Path
# Setup paths based on the script location
SCRIPT_DIR = Path(__file__).resolve().parent
BASE_DIR = SCRIPT_DIR.parent
TRAIN_FEAT_DIR = BASE_DIR / "patches_128" / "train" / "features"
TRAIN_MASK_DIR = BASE_DIR / "patches_128" / "train" / "masks"
OUT_DIR = SCRIPT_DIR / "plots"
def normalize_band(band, pmin=2, pmax=98):
"""Normalize a band for visualization using percentiles (robust to outliers)."""
# Exclude NaNs and extreme values when calculating percentiles
valid_pixels = band[~np.isnan(band)]
if len(valid_pixels) == 0:
return np.zeros_like(band)
b_min, b_max = np.percentile(valid_pixels, (pmin, pmax))
if b_max - b_min == 0:
return np.zeros_like(band)
band_norm = (band - b_min) / (b_max - b_min)
return np.clip(band_norm, 0, 1)
def main(num_samples=10):
# Ensure the output directory exists
OUT_DIR.mkdir(parents=True, exist_ok=True)
# Get a list of all feature files in the train split
feature_files = list(TRAIN_FEAT_DIR.glob("feature_*.tif"))
if not feature_files:
print(f"No feature files found in {TRAIN_FEAT_DIR}")
print("Please make sure you have run the patch building script.")
return
# Randomly select a few files to visualize
samples = random.sample(feature_files, min(num_samples, len(feature_files)))
# Optional: we can force it to pick at least some patches that contain a glacier
# by reading the metadata, but for a quick script we will just randomly sample.
print(f"Saving {len(samples)} visualization plots to {OUT_DIR} ...")
for feat_path in samples:
# Construct the matching mask path
mask_name = feat_path.name.replace("feature_", "mask_")
mask_path = TRAIN_MASK_DIR / mask_name
if not mask_path.exists():
print(f"Mask not found for {feat_path.name}, skipping.")
continue
# Read feature data
with rasterio.open(feat_path) as src:
# According to config.py, the bands are:
# 1: B2(Blue), 2: B3(Green), 3: B4(Red), 4: B8(NIR), 5: B11, 6: B12
# 7: NDVI, 8: NDWI, 9: NDSI, 10: Elev, 11: Slope, 12: Aspect
b = src.read(1)
g = src.read(2)
r = src.read(3)
ndvi = src.read(7)
elev = src.read(10)
# Read mask data
with rasterio.open(mask_path) as src:
mask = src.read(1)
# Normalize RGB for display
r_norm = normalize_band(r)
g_norm = normalize_band(g)
b_norm = normalize_band(b)
rgb_img = np.dstack((r_norm, g_norm, b_norm))
# Create the plot
fig, axes = plt.subplots(1, 4, figsize=(20, 5))
# 1. True Color (RGB)
axes[0].imshow(rgb_img)
axes[0].set_title("True Color (B4-B3-B2)")
axes[0].axis('off')
# 2. NDVI
im1 = axes[1].imshow(ndvi, cmap='RdYlGn', vmin=-1, vmax=1)
axes[1].set_title("NDVI")
axes[1].axis('off')
plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)
# 3. Elevation
im2 = axes[2].imshow(elev, cmap='terrain')
axes[2].set_title("Elevation (m)")
axes[2].axis('off')
plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04)
# 4. Ground Truth Mask
axes[3].imshow(mask, cmap='gray', interpolation='nearest')
axes[3].set_title("Ground Truth Mask")
axes[3].axis('off')
# Add a super title with the patch ID
patch_id = feat_path.stem.replace('feature_', '')
plt.suptitle(f"Patch: {patch_id}", fontsize=16)
plt.tight_layout()
# Save the figure
out_path = OUT_DIR / f"{patch_id}.png"
plt.savefig(out_path, dpi=150, bbox_inches='tight')
plt.close()
print(f"Saved {out_path.name}")
if __name__ == "__main__":
main(10)
|