app.py

import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import cv2
from PIL import Image
from transformers import SegformerForSemanticSegmentation
import os

# Konfigurasi Model
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
NUM_CLASSES = 10
MODEL_PATH = "best_model.pth"  # Pastikan file ini ada saat deploy di HF

CLASS_NAMES = [
    'background', 'building-flooded', 'building-non-flooded', 
    'grass', 'pool', 'road-flooded', 'road-non-flooded', 
    'tree', 'vehicle', 'water'
]

CLASS_COLORS = [
    (0, 0, 0),         # background
    (255, 0, 0),       # building-flooded (Red)
    (180, 120, 120),   # building-non-flooded (Brownish)
    (4, 250, 7),       # grass (Green)
    (255, 235, 0),     # pool (Yellow)
    (160, 150, 20),    # road-flooded (Dark Yellow)
    (140, 140, 140),   # road-non-flooded (Gray)
    (0, 82, 255),      # tree (Blue)
    (255, 0, 245),     # vehicle (Magenta)
    (61, 230, 250)     # water (Cyan)
]

# Inisialisasi Model Global
model = None
load_error = None

def build_model():
    m = SegformerForSemanticSegmentation.from_pretrained(
        "nvidia/mit-b4",
        num_labels=NUM_CLASSES,
        ignore_mismatched_sizes=True,
    )
    return m

def load_model():
    global model, load_error
    if model is None:
        try:
            print("Memuat arsitektur SegFormer-B4...")
            model = build_model()
            print(f"Memuat bobot dari {MODEL_PATH}...")
            if not os.path.exists(MODEL_PATH):
                raise FileNotFoundError(f"File {MODEL_PATH} tidak ditemukan. Path absolut: {os.path.abspath(MODEL_PATH)}")
            
            # Tambahkan pengecekan ukuran file
            file_size = os.path.getsize(MODEL_PATH)
            print(f"Ukuran file {MODEL_PATH}: {file_size} bytes")
            if file_size < 1000:
                raise ValueError(f"Ukuran file terlalu kecil ({file_size} bytes), mungkin ini hanya pointer LFS.")

            checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
            if isinstance(checkpoint, dict) and 'model' in checkpoint:
                state_dict = checkpoint['model']
            else:
                state_dict = checkpoint
            model.load_state_dict(state_dict)
            model.to(DEVICE)
            model.eval()
            print("Model berhasil dimuat!")
        except Exception as e:
            import traceback
            load_error = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
            print(f"Error memuat model: {load_error}")
            model = None

def tta_predict(img_tensor):
    """Test-Time Augmentation (Original, H-Flip, V-Flip, HV-Flip)"""
    with torch.no_grad():
        img = img_tensor.unsqueeze(0).to(DEVICE)
        img_hf = torch.flip(img, [3])
        img_vf = torch.flip(img, [2])
        img_hvf = torch.flip(img, [2, 3])
        
        logits1 = model(img).logits
        logits2 = torch.flip(model(img_hf).logits, [3])
        logits3 = torch.flip(model(img_vf).logits, [2])
        logits4 = torch.flip(model(img_hvf).logits, [2, 3])
        
        logits_avg = (logits1 + logits2 + logits3 + logits4) / 4.0
        
        # Bilinear upsample logits to target size (same as img shape)
        logits_avg = nn.functional.interpolate(logits_avg, size=img.shape[2:], mode='bilinear', align_corners=False)
        
        preds = torch.argmax(logits_avg, dim=1).squeeze(0).cpu().numpy()
        return preds

def predict(image):
    if image is None:
        return None, "Error: Silakan unggah gambar terlebih dahulu."
    
    # Load model if not loaded
    if model is None:
        load_model()
        if model is None:
            return image, f"Error: Gagal memuat model. Detail kesalahan:\n\n```\n{load_error}\n```"

    # Preprocess Image
    img_array = np.array(image)
    orig_h, orig_w = img_array.shape[:2]
    
    # Resize to model input
    img_resized = cv2.resize(img_array, (640, 480))
    # Normalize ImageNet stats
    img_normalized = img_resized.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img_normalized = (img_normalized - mean) / std
    
    img_tensor = torch.from_numpy(img_normalized).permute(2, 0, 1)
    
    # Predict
    mask_idx = tta_predict(img_tensor)
    
    # Resize mask back to original
    mask_resized = cv2.resize(mask_idx.astype(np.uint8), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
    
    # Create colored overlay
    overlay = np.zeros((orig_h, orig_w, 3), dtype=np.uint8)
    for i in range(1, NUM_CLASSES):  # Skip background
        overlay[mask_resized == i] = CLASS_COLORS[i]
        
    # Blend image and overlay
    alpha = 0.5
    blended = cv2.addWeighted(img_array, 1 - alpha, overlay, alpha, 0)
    
    # Hitung Statistik Piksel
    stats = "### Hasil Deteksi (Area)\n"
    total_px = orig_h * orig_w
    for i in range(1, NUM_CLASSES):
        px_count = np.sum(mask_resized == i)
        if px_count > 0:
            pct = (px_count / total_px) * 100
            stats += f"- **{CLASS_NAMES[i]}**: {pct:.2f}%\n"

    return Image.fromarray(blended), stats

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        # 🌊 Flood Decision Support System - Live Prediction
        **Tim Dolanan Data UNESA 2026** | Model: `SegFormer-B4`
        
        Unggah citra udara (drone) untuk melihat hasil segmentasi area banjir secara *real-time*.
        """
    )
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Upload Citra Udara")
            btn_predict = gr.Button("Analisis Gambar", variant="primary")
        
        with gr.Column():
            img_output = gr.Image(label="Hasil Prediksi (Overlay Mask)")
            txt_stats = gr.Markdown()
            
    btn_predict.click(fn=predict, inputs=img_input, outputs=[img_output, txt_stats])
    
    # Examples
    sample_images = [
        "samples/7083.jpg",
        "samples/7188.jpg",
        "samples/8334.jpg",
        "samples/8131.jpg",
        "samples/8796.jpg",
        "samples/8009.jpg",
        "samples/8817.jpg",
        "samples/6338.jpg",
        "samples/7109.jpg",
        "samples/7171.jpg"
    ]
    
    # Filter only existing sample files to avoid UI crash if any sample is missing
    existing_samples = [s for s in sample_images if os.path.exists(s)]
    
    if existing_samples:
        gr.Examples(
            examples=existing_samples,
            inputs=img_input,
            outputs=[img_output, txt_stats],
            fn=predict,
            cache_examples=False
        )

    gr.Markdown(
        """
        ---
        **Catatan:** Model ini memprediksi 10 kelas (termasuk *building-flooded*, *road-flooded*, *water*, dll). 
        Inferensi menggunakan *Test-Time Augmentation (TTA)* untuk akurasi maksimal.
        """
    )

if __name__ == "__main__":
    app.launch()