Update app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import cv2
 import numpy as np
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 import base64
@@ -49,37 +48,44 @@ def find_similarity(base64_image, text_input):
 
 # Define a function for image segmentation
 def segment_image(input_image, text_input):
-    image_bytes = base64.b64decode(input_image)
-    image = Image.open(BytesIO(image_bytes))
-    
-    image = np.array(image)  # Remove the color mode conversion
-    mask_generator = SamAutomaticMaskGenerator(sam)
-    masks = mask_generator.generate(image)
-    
-    segmented_regions = []  # List to store segmented regions with similarity scores
+    try:
+        image_bytes = base64.b64decode(input_image)
+        image = Image.open(BytesIO(image_bytes))
+
+        # Convert the image to RGB color mode
+        image = image.convert("RGB")
+
+        mask_generator = SamAutomaticMaskGenerator(sam)
+        masks = mask_generator.generate(image)
 
-    for i, mask_dict in enumerate(masks):
-        mask_data = (mask_dict['segmentation'] * 255).astype(np.uint8)
-        segmented_region = cv2.bitwise_and(image, image, mask=mask_data)
+        segmented_regions = []  # List to store segmented regions with similarity scores
 
-        x, y, w, h = map(int, mask_dict['bbox'])
-        cropped_region = segmented_region[y:y+h, x:x+w]
+        for i, mask_dict in enumerate(masks):
+            mask_data = (mask_dict['segmentation'] * 255).astype(np.uint8)
+            segmented_region = Image.fromarray(mask_data)  # Convert mask to an image
 
-        # Convert to base64 image
-        _, buffer = cv2.imencode(".png", cropped_region)
-        segmented_image_base64 = base64.b64encode(buffer).decode()
+            x, y, w, h = map(int, mask_dict['bbox'])
+            cropped_region = image.crop((x, y, x + w, y + h))
 
-        # Calculate similarity for the segmented image
-        similarity = find_similarity(segmented_image_base64, text_input)
+            # Convert to base64 image
+            buffered = BytesIO()
+            cropped_region.save(buffered, format="PNG")
+            segmented_image_base64 = base64.b64encode(buffered.getvalue()).decode()
 
-        # Append the segmented image and its similarity score
-        segmented_regions.append({"image": segmented_image_base64, "similarity": similarity})
+            # Calculate similarity for the segmented image
+            similarity = find_similarity(segmented_image_base64, text_input)
 
-    # Sort the segmented images by similarity in descending order
-    segmented_regions.sort(key=lambda x: x["similarity"], reverse=True)
+            # Append the segmented image and its similarity score
+            segmented_regions.append({"image": segmented_image_base64, "similarity": similarity})
 
-    # Return the segmented images in descending order of similarity
-    return segmented_regions
+        # Sort the segmented images by similarity in descending order
+        segmented_regions.sort(key=lambda x: x["similarity"], reverse=True)
+
+        # Return the segmented images in descending order of similarity
+        return segmented_regions
+
+    except Exception as e:
+        return str(e)
 
 # Create Gradio components
 input_image = gr.Textbox(label="Base64 Image", lines=8)