Upload folder using huggingface_hub

80e070e verified 5 months ago

4.06 kB

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)

	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)