Update app.py

app.py

@@ -13,15 +13,15 @@ available_models = {
     # Add more models as needed
 }
 
-
-def segment_image(input_image):
+def segment_image(input_image, selected_model):
     # Resize the input image to 255x255
     img = np.array(input_image)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
+    model = available_models[selected_model]
+    
     # Perform object detection and segmentation
     results = model(img)
-    mask =  results[0].masks.data.numpy()
+    mask = results[0].masks.data.numpy()
     target_height = img.shape[0]
     target_width = img.shape[1]
 
@@ -52,7 +52,6 @@ def segment_image(input_image):
     predictor = SamPredictor(sam)
     predictor.set_image(img)
     
-    
     input_box = np.array(bbox)
     masks_, _, _ = predictor.predict(
         point_coords=None,
@@ -62,7 +61,7 @@ def segment_image(input_image):
 
     fmask = masks_[0].astype(int)
 
-    resized_mask1 =cv2.resize(fmask, (target_width, target_height))
+    resized_mask1 = cv2.resize(fmask, (target_width, target_height))
     resized_mask1 = (resized_mask1 * 255).astype(np.uint8)
 
     overlay_image1 = img.copy()
@@ -72,54 +71,25 @@ def segment_image(input_image):
     # Convert the overlay image to PIL format
     overlay_pil1 = Image.fromarray(overlay_image1)
 
-    
     return overlay_pil, overlay_pil1  # Return both overlay image and mask
 
-
-
-# Create a function to perform image segmentation using the selected model
-'''def segment_image(input_image, selected_model):
-    # Resize the input image to 255x255
-    img = np.array(input_image)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-    # Perform object detection and segmentation using the selected model
-    model = available_models[selected_model]
-    results = model(img)
-    mask = results[0].masks.data.numpy()
-    target_height = img.shape[0]
-    target_width = img.shape[1]
-
-    # Resize the mask using OpenCV
-    resized_mask = cv2.resize(mask[0], (target_width, target_height))
-    resized_mask = (resized_mask * 255).astype(np.uint8)
-
-    # Create a copy of the original image
-    overlay_image = img.copy()
-
-    # Apply the resized mask to the overlay image
-    overlay_image[resized_mask > 0] = [50, 0, 0]  # Overlay in green
-
-    # Convert the overlay image to PIL format
-    overlay_pil = Image.fromarray(overlay_image)
-
-    return overlay_pil'''
-
 # Create the Gradio interface with a dropdown for model selection
 iface = gr.Interface(
     fn=segment_image,
     inputs=[
-        gr.inputs.Image(type="pil", label="Upload an image"),
-        gr.inputs.Dropdown(
+        gr.components.Image(type="pil", label="Upload an image"),
+        gr.components.Dropdown(
             choices=list(available_models.keys()),
             label="Select YOLO Model",
             default="X-ray"
         )
     ],
-    outputs=[gr.outputs.Image(type="numpy", label="Segmented Image"),
-             gr.outputs.Image(type="numpy", label="Segmentation Mask")],  # Add an output for the mask
+    outputs=[
+        gr.components.Image(type="pil", label="Segmented Image"),
+        gr.components.Image(type="pil", label="Segmentation Mask")
+    ],
     title="YOLOv8 with SAM 😃",
-    description='This software generates the segmentation mask for Aorta for Point of Care Ultrasound (POCUS) images'
+    description='This software generates the segmentation mask Medical images'
 )
 
 iface.launch()