Update app.py

app.py

@@ -147,123 +147,134 @@ optic_file = st.file_uploader("Upload Optical Image", type=["tiff"])
 mask_file = st.file_uploader("Upload Mask Image", type=["tiff"])
 
 num_samples = st.slider("Number of test samples to visualize", 1, 10, 1)
+if st.button("Run Inference"):
+    with st.spinner("Loading data and model..."):
+        if sar_file is not None and optic_file is not None and mask_file is not None:
+            st.success("All files uploaded successfully!")
+            st.write(f"Number of samples selected for visualization: {num_samples}")
 
-if sar_file is not None and optic_file is not None and mask_file is not None:
-    st.success("All files uploaded successfully!")
-    st.write(f"Number of samples selected for visualization: {num_samples}")
-    sar_path = save_uploaded_file(sar_file, suffix=".tif")
-    optic_path = save_uploaded_file(optic_file, suffix=".tif")
-    mask_path = save_uploaded_file(mask_file, suffix=".tif")
- 
-    sarImages = [sar_path]
-    opticImages = [optic_path]
-    masks = [mask_path]
-    model_path = "Residual_UNET_Bilinear.keras"
-    
-    
-    if st.button("Run Inference"):
-        with st.spinner("Loading data and model..."):
-    
+            # Save uploaded files
+            sar_path = save_uploaded_file(sar_file, suffix=".tif")
+            optic_path = save_uploaded_file(optic_file, suffix=".tif")
+            mask_path = save_uploaded_file(mask_file, suffix=".tif")
+
+            # Create image lists
+            sarImages = [sar_path]
+            opticImages = [optic_path]
+            masks = [mask_path]
+
+            # Model path
+            model_path = "Residual_UNET_Bilinear.keras"
+
+            # Read and normalize images
             sar_images = readImages(sarImages, typeData='s', width=WIDTH, height=HEIGHT)
             optic_images = readImages(opticImages, typeData='o', width=WIDTH, height=HEIGHT)
             masks = readImages(masks, typeData='m', width=WIDTH, height=HEIGHT)
-    
+
             sar_images = normalizeImages(sar_images, 's')
             optic_images = normalizeImages(optic_images, 'i')
-    
+
             # Load model
-            model = tf.keras.models.load_model(model_path,
-                                      custom_objects={"cce_dice_loss": cce_dice_loss, "dice_score": dice_score})
-    
+            model = tf.keras.models.load_model(
+                model_path,
+                custom_objects={"cce_dice_loss": cce_dice_loss, "dice_score": dice_score}
+            )
+
+            # Predict masks
             pred_masks = model.predict([optic_images, sar_images], verbose=0)
             is_multiclass = pred_masks.shape[-1] > 1
-    
+
             num_samples = min(num_samples, len(sar_images))
-    
-            # Plotting
+
+            # Plotting results
             fig, axes = plt.subplots(num_samples, 4, figsize=(21, 6 * num_samples))
-    
+
             for i in range(num_samples):
                 ax = axes[i] if num_samples > 1 else axes
-    
+
+                # Plot SAR image
                 ax[0].imshow(sar_images[i].squeeze(), cmap='gray')
                 ax[0].set_title(f"SAR Image {i+1}")
                 ax[0].axis('off')
-    
+
+                # Plot Optic image
                 ax[1].imshow(optic_images[i])
                 ax[1].set_title(f"Optic Image {i+1}")
                 ax[1].axis('off')
-    
+
+                # Plot Ground Truth Mask
                 if is_multiclass:
                     gt_color_mask = np.zeros((*masks[i].shape[:2], 3))
                     for j, color in enumerate(CLASS_COLORS):
-                        gt_color_mask += masks[i][:,:,j][:,:,np.newaxis] * np.array(color)
+                        gt_color_mask += masks[i][:, :, j][:, :, np.newaxis] * np.array(color)
                     ax[2].imshow(gt_color_mask)
                 else:
                     ax[2].imshow(masks[i], cmap='gray')
                 ax[2].set_title(f"Ground Truth Mask {i+1}")
                 ax[2].axis('off')
-    
+
+                # Plot Predicted Mask
                 if is_multiclass:
                     pred_color_mask = np.zeros((*pred_masks[i].shape[:2], 3))
                     for j, color in enumerate(CLASS_COLORS):
-                        pred_color_mask += pred_masks[i][:,:,j][:,:,np.newaxis] * np.array(color)
+                        pred_color_mask += pred_masks[i][:, :, j][:, :, np.newaxis] * np.array(color)
                     ax[3].imshow(pred_color_mask)
                 else:
                     ax[3].imshow(pred_masks[i], cmap='gray')
                 ax[3].set_title(f"Predicted Mask {i+1}")
                 ax[3].axis('off')
-    
+
             st.pyplot(fig)
-    
+
             # Define color for class 1: illegal mining
             red_color = [255, 0, 0]
-            
+
             # Convert optic_images to uint8 if needed
             if optic_images.dtype != np.uint8:
                 optic_images = (optic_images * 255).astype(np.uint8)
-            
-            # Create figure with subplots
+
+            # Plot overlays
             fig, axes = plt.subplots(num_samples, 4, figsize=(21, 6 * num_samples))
-            
+
             for i in range(num_samples):
                 ax = axes[i] if num_samples > 1 else axes
-            
+
                 # SAR image
                 ax[0].imshow(sar_images[i].squeeze(), cmap='gray')
                 ax[0].set_title(f"SAR Image {i+1}")
                 ax[0].axis('off')
-            
+
                 # Optic image
                 ax[1].imshow(optic_images[i])
                 ax[1].set_title(f"Optic Image {i+1}")
                 ax[1].axis('off')
-            
+
                 # Ground truth overlay
                 gt_overlay = optic_images[i].copy()
                 if is_multiclass:
                     gt_overlay[masks[i][:, :, 1] == 1] = red_color
                 else:
                     gt_overlay[masks[i].squeeze() == 1] = red_color
-            
+
                 ax[2].imshow(optic_images[i])
                 ax[2].imshow(gt_overlay, alpha=0.4)
                 ax[2].set_title(f"Ground Truth Overlay {i+1}")
                 ax[2].axis('off')
-            
+
                 # Predicted mask overlay
                 pred_overlay = optic_images[i].copy()
                 if is_multiclass:
                     pred_overlay[pred_masks[i][:, :, 1] > 0.5] = red_color
                 else:
                     pred_overlay[pred_masks[i].squeeze() > 0.5] = red_color
-            
+
                 ax[3].imshow(optic_images[i])
                 ax[3].imshow(pred_overlay, alpha=0.4)
                 ax[3].set_title(f"Predicted Overlay {i+1}")
                 ax[3].axis('off')
-            
+
             plt.tight_layout()
             st.pyplot(fig)
-else:
-    st.warning("Please upload all three .tiff files to proceed.")
+        else:
+            st.warning("Please upload all three .tiff files to proceed.")
+