Update app.py

app.py

@@ -14,9 +14,71 @@ available_models = {
 }
 
 
+def segment_image(input_image):
+    # Resize the input image to 255x255
+    img = np.array(input_image)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    # Perform object detection and segmentation
+    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)
+
+    # Convert the resized mask to PIL format
+    mask_pil = Image.fromarray(resized_mask)
+    
+    for result in results:
+        boxes = result.boxes
+        bbox = boxes.xyxy.tolist()[0]
+    
+    sam_checkpoint = "sam_vit_h_4b8939.pth"
+    model_type = "vit_h"
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+    sam.to(device='cpu')
+    predictor = SamPredictor(sam)
+    predictor.set_image(img)
+    
+    
+    input_box = np.array(bbox)
+    masks_, _, _ = predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=input_box,
+        multimask_output=False)
+
+    fmask = masks_[0].astype(int)
+
+    resized_mask1 =cv2.resize(fmask[0], (target_width, target_height))
+    resized_mask1 = (resized_mask1 * 255).astype(np.uint8)
+
+    overlay_image1 = img.copy()
+    # Apply the resized mask to the overlay image
+    overlay_image1[resized_mask1 > 0] = [50, 50, 0]  # Overlay in green
+
+    # 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):
+'''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)
@@ -41,7 +103,7 @@ def segment_image(input_image, selected_model):
     # Convert the overlay image to PIL format
     overlay_pil = Image.fromarray(overlay_image)
 
-    return overlay_pil
+    return overlay_pil'''
 
 # Create the Gradio interface with a dropdown for model selection
 iface = gr.Interface(