Update app.py

app.py

@@ -8,7 +8,6 @@ st.title("🌊 Flood Segmentation with U-Net")
 # 1. Load Model
 @st.cache_resource
 def load_model():
-    # Load model mentah (Raw SavedModel)
     model = tf.saved_model.load("unet_savedmodel")
     return model
 
@@ -22,32 +21,39 @@ if uploaded:
     img = Image.open(uploaded).convert("RGB")
     img_resized = img.resize((256, 256))
     
+    # Normalisasi
     x = np.array(img_resized, dtype=np.float32) / 255.0
-    x = np.expand_dims(x, axis=0) # Menambah batch dimension
-
-    # 3. Prediksi (PERBAIKAN UTAMA DISINI)
-    # Ambil fungsi inferensi spesifik dari dalam model
-    infer = model.signatures["serving_default"]
-    
-    # Jalankan inferensi (input harus dijadikan Tensor constant)
-    output_dict = infer(tf.constant(x))
-    
-    # Hasilnya adalah dictionary (misal: {'conv2d_output': tensor}). 
-    # Kita ambil value pertama dari dictionary tersebut.
-    pred_tensor = list(output_dict.values())[0]
-    
-    # Konversi Tensor ke Numpy
-    pred = pred_tensor.numpy()[0, :, :, 0]
     
-    # 4. Post-processing
-    mask = (pred > 0.5).astype(np.uint8) * 255
-    mask_img = Image.fromarray(mask)
-
-    # 5. Tampilkan Hasil
-    col1, col2 = st.columns(2)
-    with col1:
-        st.write("Original Image")
-        st.image(img, use_container_width=True)
-    with col2:
-        st.write("Segmentation Mask")
-        st.image(mask_img, use_container_width=True)
+    # --- 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).
+    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)
+
+    # 3. Prediksi
+    try:
+        infer = model.signatures["serving_default"]
+        output_dict = infer(tf.constant(x))
+        pred_tensor = list(output_dict.values())[0]
+        pred = pred_tensor.numpy()[0, :, :, 0]
+        
+        # 4. Post-processing
+        mask = (pred > 0.5).astype(np.uint8) * 255
+        mask_img = Image.fromarray(mask)
+
+        # 5. Tampilkan Hasil
+        col1, col2 = st.columns(2)
+        with col1:
+            st.write("Original Image")
+            st.image(img, use_container_width=True)
+        with col2:
+            st.write("Segmentation Mask")
+            st.image(mask_img, use_container_width=True)
+            
+    except Exception as e:
+        st.error(f"Terjadi error saat prediksi: {e}")
+        st.write("Debug Info - Input Shape:", x.shape)