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()