Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import cv2
 import numpy as np
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 import base64
@@ -6,7 +7,6 @@ from PIL import Image
 from io import BytesIO
 import torch
 import clip
-import json
 
 # Load the segmentation model
 sam_checkpoint = "sam_vit_h_4b8939.pth"
@@ -18,10 +18,16 @@ model, preprocess = clip.load("ViT-B/32")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device).eval()
 
-def find_similarity(image, text_input):
+
+def find_similarity(base64_image, text_input):
     try:
+        # Decode the base64 image to bytes
+        image_bytes = base64.b64decode(base64_image)
+
+        # Convert the bytes to a PIL image
+        image = Image.open(BytesIO(image_bytes))
+
         # Preprocess the image
-        image = Image.fromarray(image)
         image = preprocess(image).unsqueeze(0).to(device)
 
         # Prepare input text
@@ -37,45 +43,43 @@ def find_similarity(image, text_input):
         text_features /= text_features.norm(dim=-1, keepdim=True)
         similarity = (text_features @ image_features.T).squeeze(0).cpu().numpy()
 
-        return similarity.tolist()  # Convert to a list
-
+        return similarity
     except Exception as e:
-        return json.dumps({"error": str(e)})
+        return str(e)
 
+# Define a function for image segmentation
 def segment_image(input_image, text_input):
-    try:
-        image_bytes = base64.b64decode(input_image)
-        image = Image.open(BytesIO(image_bytes))
-
-        # Convert the image to RGB color mode
-        image = image.convert("RGB")
+    image_bytes = base64.b64decode(input_image)
+    image = Image.open(BytesIO(image_bytes))
+    
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
+    mask_generator = SamAutomaticMaskGenerator(sam)
+    masks = mask_generator.generate(image)
 
-        mask_generator = SamAutomaticMaskGenerator(sam)
-        masks = mask_generator.generate(np.array(image))
+    segmented_regions = []  # List to store segmented regions with similarity scores
 
-        segmented_regions = []  # List to store segmented regions with similarity scores
+    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)
 
-        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
+        x, y, w, h = map(int, mask_dict['bbox'])
+        cropped_region = segmented_region[y:y+h, x:x+w]
 
-            x, y, w, h = map(int, mask_dict['bbox'])
-            cropped_region = image.crop((x, y, x + w, y + h))
+        # Convert to base64 image
+        _, buffer = cv2.imencode(".png", cv2.cvtColor(cropped_region, cv2.COLOR_BGR2RGB))
+        segmented_image_base64 = base64.b64encode(buffer).decode()
 
-            # Calculate similarity for the segmented image
-            similarity = find_similarity(np.array(cropped_region), text_input)
+        # Calculate similarity for the segmented image
+        similarity = find_similarity(segmented_image_base64, text_input)
 
-            # Append the segmented image and its similarity score
-            segmented_regions.append({"image": input_image, "similarity": similarity})
+        # Append the segmented image and its similarity score
+        segmented_regions.append({"image": segmented_image_base64, "similarity": similarity})
 
-        # Sort the segmented images by similarity in descending order
-        segmented_regions.sort(key=lambda x: x["similarity"], reverse=True)
+    # Sort the segmented images by similarity in descending order
+    segmented_regions.sort(key=lambda x: x["similarity"], reverse=True)
 
-        # Return the segmented images in a JSON format
-        return json.dumps(segmented_regions)
-
-    except Exception as e:
-        return json.dumps({"error": str(e)})
+    # Return the segmented images in descending order of similarity
+    return segmented_regions
 
 # Create Gradio components
 input_image = gr.Textbox(label="Base64 Image", lines=8)