Update app.py

app.py

@@ -18,15 +18,10 @@ model, preprocess = clip.load("ViT-B/32")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device).eval()
 
-def find_similarity(base64_image, text_input):
+def find_similarity(image, text_input):
     try:
-        # Decode the base64 image to bytes
-        image_bytes = base64.b64decode(base64_image)
-
-        # Convert the bytes to a numpy array
-        image = np.array(Image.open(BytesIO(image_bytes)))
-
         # Preprocess the image
+        image = Image.fromarray(image)
         image = preprocess(image).unsqueeze(0).to(device)
 
         # Prepare input text
@@ -42,7 +37,8 @@ def find_similarity(base64_image, text_input):
         text_features /= text_features.norm(dim=-1, keepdim=True)
         similarity = (text_features @ image_features.T).squeeze(0).cpu().numpy()
 
-        return similarity
+        return similarity.tolist()  # Convert to a list
+
     except Exception as e:
         return json.dumps({"error": str(e)})
 
@@ -66,16 +62,11 @@ def segment_image(input_image, text_input):
             x, y, w, h = map(int, mask_dict['bbox'])
             cropped_region = image.crop((x, y, x + w, y + h))
 
-            # Convert to base64 image
-            buffered = BytesIO()
-            cropped_region.save(buffered, format="PNG")
-            segmented_image_base64 = base64.b64encode(buffered.getvalue()).decode()
-
             # Calculate similarity for the segmented image
-            similarity = find_similarity(segmented_image_base64, text_input)
+            similarity = find_similarity(np.array(cropped_region), text_input)
 
             # Append the segmented image and its similarity score
-            segmented_regions.append({"image": segmented_image_base64, "similarity": similarity})
+            segmented_regions.append({"image": input_image, "similarity": similarity})
 
         # Sort the segmented images by similarity in descending order
         segmented_regions.sort(key=lambda x: x["similarity"], reverse=True)