Update app.py

app.py

@@ -1,9 +1,10 @@
 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 U-Net")
+st.title("🌊 Flood Segmentation with Heatmap")
 
 # 1. Load Model
 @st.cache_resource
@@ -24,36 +25,54 @@ if uploaded:
     # Normalisasi
     x = np.array(img_resized, dtype=np.float32) / 255.0
     
-    # --- PERBAIKAN KHUSUS UNTUK ERROR 3 VS 6 ---
-    # Model kamu butuh 6 channel (kemungkinan Before + After image).
-    # Karena cuma ada 1 gambar, kita gabung gambar ini dengan dirinya sendiri
-    # agar menjadi 6 channel (Mocking the second image).
+    # Stack gambar agar menjadi 6 channel (memenuhi syarat input model)
     x = np.concatenate([x, x], axis=-1) 
-    # Sekarang shape x adalah (256, 256, 6)
-    # -------------------------------------------
-
-    x = np.expand_dims(x, axis=0) # Menambah batch dimension menjadi (1, 256, 256, 6)
+    
+    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]
         
-        # 4. Post-processing
-        mask = (pred > 0.5).astype(np.uint8) * 255
-        mask_img = Image.fromarray(mask)
+        # --- 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)
 
-        # 5. Tampilkan Hasil
-        col1, col2 = st.columns(2)
+        # --------------------------------
+        
+        # 4. Tampilkan Hasil
+        col1, col2, col3 = st.columns(3)
+        
         with col1:
             st.write("Original Image")
-            st.image(img, use_container_width=True)
+            st.image(img_resized, use_container_width=True)
+            
         with col2:
-            st.write("Segmentation Mask")
-            st.image(mask_img, use_container_width=True)
+            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 saat prediksi: {e}")
-        st.write("Debug Info - Input Shape:", x.shape)
+        st.error(f"Terjadi error: {e}")