133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Generate PNG previews for existing GeoTIFF prediction files
|
|
"""
|
|
|
|
import numpy as np
|
|
import rasterio
|
|
import matplotlib
|
|
matplotlib.use('Agg')
|
|
import matplotlib.pyplot as plt
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
def generate_png_preview(tif_file, output_png=None):
|
|
"""Generate PNG preview from GeoTIFF file"""
|
|
tif_path = Path(tif_file)
|
|
|
|
if not tif_path.exists():
|
|
print(f"❌ File not found: {tif_file}")
|
|
return False
|
|
|
|
# Determine output PNG path
|
|
if output_png is None:
|
|
output_png = tif_path.with_suffix('.png')
|
|
else:
|
|
output_png = Path(output_png)
|
|
|
|
try:
|
|
# Read GeoTIFF
|
|
with rasterio.open(tif_path) as src:
|
|
data = src.read(1)
|
|
|
|
print(f"📊 Data shape: {data.shape}, range: [{np.nanmin(data):.3f}, {np.nanmax(data):.3f}]")
|
|
|
|
# Determine if it's classification or NDVI based on filename
|
|
is_classification = 'classification' in tif_path.name.lower() or 'prediction' in tif_path.name.lower()
|
|
is_ndvi = 'ndvi' in tif_path.name.lower()
|
|
|
|
# Create figure
|
|
fig, ax = plt.subplots(figsize=(12, 10), dpi=150)
|
|
|
|
if is_ndvi:
|
|
# NDVI: use RdYlGn colormap, range -1 to 1
|
|
im = ax.imshow(data, cmap='RdYlGn', vmin=-1, vmax=1, interpolation='nearest')
|
|
ax.set_title(f'NDVI - {tif_path.stem}', fontsize=14, fontweight='bold')
|
|
cbar_label = 'NDVI'
|
|
elif is_classification:
|
|
# Classification: use tab20 colormap
|
|
im = ax.imshow(data, cmap='tab20', interpolation='nearest')
|
|
ax.set_title(f'Land Classification - {tif_path.stem}', fontsize=14, fontweight='bold')
|
|
cbar_label = 'Class'
|
|
else:
|
|
# Generic: use viridis
|
|
im = ax.imshow(data, cmap='viridis', interpolation='nearest')
|
|
ax.set_title(f'{tif_path.stem}', fontsize=14, fontweight='bold')
|
|
cbar_label = 'Value'
|
|
|
|
ax.set_xlabel('X (pixels)', fontsize=10)
|
|
ax.set_ylabel('Y (pixels)', fontsize=10)
|
|
|
|
# Add colorbar
|
|
cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
|
cbar.set_label(cbar_label, rotation=270, labelpad=15)
|
|
|
|
# For classification, try to set integer ticks
|
|
if is_classification:
|
|
try:
|
|
unique_vals = np.unique(data[~np.isnan(data)])
|
|
if len(unique_vals) < 20: # Only if not too many classes
|
|
cbar.set_ticks(unique_vals)
|
|
cbar.set_ticklabels([str(int(v)) for v in unique_vals])
|
|
except:
|
|
pass
|
|
|
|
# Add grid
|
|
ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
|
|
|
|
# Save PNG
|
|
plt.tight_layout()
|
|
plt.savefig(str(output_png), dpi=150, bbox_inches='tight')
|
|
plt.close(fig)
|
|
|
|
print(f"✅ Created PNG: {output_png}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error creating PNG: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
|
|
def generate_all_previews(predictions_dir="predictions"):
|
|
"""Generate PNG previews for all GeoTIFF files without PNGs"""
|
|
pred_path = Path(predictions_dir)
|
|
|
|
if not pred_path.exists():
|
|
print(f"❌ Directory not found: {predictions_dir}")
|
|
return
|
|
|
|
tif_files = list(pred_path.glob("*.tif"))
|
|
print(f"🔍 Found {len(tif_files)} GeoTIFF files")
|
|
|
|
generated = 0
|
|
skipped = 0
|
|
|
|
for tif_file in tif_files:
|
|
png_file = tif_file.with_suffix('.png')
|
|
|
|
if png_file.exists():
|
|
print(f"⏭️ Skipping {tif_file.name} (PNG already exists)")
|
|
skipped += 1
|
|
continue
|
|
|
|
print(f"\n🎨 Processing {tif_file.name}...")
|
|
if generate_png_preview(tif_file):
|
|
generated += 1
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"✅ Generated {generated} new PNG previews")
|
|
print(f"⏭️ Skipped {skipped} files (already have PNGs)")
|
|
print(f"{'='*60}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) > 1:
|
|
# Process specific file
|
|
tif_file = sys.argv[1]
|
|
generate_png_preview(tif_file)
|
|
else:
|
|
# Process all files in predictions directory
|
|
generate_all_previews()
|