import streamlit as st
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt  # <--- Library baru untuk membuat Heatmap
from PIL import Image

st.title("🌊 Flood Segmentation with Heatmap")

# 1. Load Model
@st.cache_resource
def load_model():
    model = tf.saved_model.load("unet_savedmodel")
    return model

with st.spinner('Loading Model...'):
    model = load_model()

uploaded = st.file_uploader("Upload flood image", type=["jpg", "jpeg", "png"])

if uploaded:
    # 2. Preprocessing Gambar
    img = Image.open(uploaded).convert("RGB")
    img_resized = img.resize((256, 256))
    
    # Normalisasi
    x = np.array(img_resized, dtype=np.float32) / 255.0
    
    # Stack gambar agar menjadi 6 channel (memenuhi syarat input model)
    x = np.concatenate([x, x], axis=-1) 
    
    x = np.expand_dims(x, axis=0) # Batch dimension

    # 3. Prediksi
    try:
        infer = model.signatures["serving_default"]
        output_dict = infer(tf.constant(x))
        pred_tensor = list(output_dict.values())[0]
        
        # Ambil nilai probabilitas mentah (0.0 sampai 1.0)
        pred = pred_tensor.numpy()[0, :, :, 0]
        
        # --- BAGIAN PEMBUATAN HEATMAP ---
        
        # 1. Buat colormap (Jet: Biru=Rendah, Merah=Tinggi)
        # plt.cm.jet mengembalikan RGBA (0-1), kita ambil RGB saja dan kali 255
        heatmap_colored = plt.cm.jet(pred)[:, :, :3]
        heatmap_img = (heatmap_colored * 255).astype(np.uint8)
        
        # 2. Buat Overlay (Gabungkan gambar asli dengan Heatmap)
        # Kita ambil gambar asli yang sudah di-resize
        original_img_arr = np.array(img_resized)
        
        # Rumus blending: 60% Gambar Asli + 40% Heatmap
        blended = (original_img_arr * 0.6 + heatmap_img * 0.4).astype(np.uint8)
        
        # Konversi ke format Image agar bisa ditampilkan Streamlit
        final_heatmap = Image.fromarray(blended)
        pure_heatmap = Image.fromarray(heatmap_img)

        # --------------------------------
        
        # 4. Tampilkan Hasil
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.write("Original Image")
            st.image(img_resized, use_container_width=True)
            
        with col2:
            st.write("Pure Heatmap")
            st.image(pure_heatmap, use_container_width=True, caption="Merah = Banjir")

        with col3:
            st.write("Overlay Result")
            st.image(final_heatmap, use_container_width=True, caption="Gabungan")
            
    except Exception as e:
        st.error(f"Terjadi error: {e}")