hoàn thành model swing-unet

2026-01-05 16:20:58 +07:00
parent 10219df149
commit 2b308ddb78
4 changed files with 509 additions and 95 deletions
@@ -0,0 +1,132 @@
+#!/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()