Update app.py

app.py

@@ -19,8 +19,7 @@ def dice_loss(y_true, y_pred):
 
 # Load the model from the same directory as the script
 model_filename = "model.h5"  # Replace with your actual model filename
-model_path = os.path.join(os.path.dirname(__file__), model_filename) # Construct the full path
-
+model_path = os.path.join(os.path.dirname(__file__), model_filename)
 
 def load_model(model_path):
     try:
@@ -37,8 +36,8 @@ def perform_inference(image):
     if model is None:
         print("Model not loaded properly.")
         return None, None, None
-        
-     # Preprocess the image
+
+    # Preprocess the image
     original_shape = image.shape[:2]    
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     image_resized = cv2.resize(image, (256, 256))
@@ -50,18 +49,27 @@ def perform_inference(image):
     mask_resized = cv2.resize(mask, (original_shape[1], original_shape[0]))
     mask_binary = (mask_resized > 0.5).astype(np.uint8) * 255
 
+    # Find contours in the binary mask
+    contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
-    # Apply the mask to the original image (for better visualization)
+    # Apply the mask to the original image (heatmap for visualization)
     heatmap_img = cv2.applyColorMap(mask_binary, cv2.COLORMAP_JET)
     segmented_image = cv2.addWeighted(image, 0.7, heatmap_img, 0.3, 0)
+    segmented_image_with_box = segmented_image.copy() #make a copy for drawing box
 
-    #Convert back to BGR
-    segmented_image = cv2.cvtColor(segmented_image, cv2.COLOR_RGB2BGR)
 
-    # Convert results to PIL Images
+    # Get bounding boxes for all contours
+    for contour in contours:
+        x, y, w, h = cv2.boundingRect(contour)
+        cv2.rectangle(segmented_image_with_box, (x, y), (x + w, y + h), (0, 0, 255), 2)
+
+
+    #Convert back to BGR
+    segmented_image_with_box = cv2.cvtColor(segmented_image_with_box, cv2.COLOR_RGB2BGR)
+    #Return image with bounding box
     return (Image.fromarray(image),
-            Image.fromarray(mask_binary.astype(np.uint8)), #Mask is already multiplied by 255
-            Image.fromarray(segmented_image))
+            Image.fromarray(mask_binary.astype(np.uint8)),
+            Image.fromarray(segmented_image_with_box))
 
 
 # Gradio app