Update app.py

app.py

@@ -3,7 +3,7 @@ import cv2
 import numpy as np
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 import base64
-from gradio_client import Client
+from huggingface_hub import InferenceClient
 
 # Load the segmentation model
 sam_checkpoint = "sam_vit_h_4b8939.pth"
@@ -16,7 +16,8 @@ def segment_image(input_image):
     mask_generator = SamAutomaticMaskGenerator(sam)
     masks = mask_generator.generate(image)
 
-    segmented_regions = []  # List to store segmented regions
+    highest_cosine_value = -1
+    highest_cosine_base64 = ""
 
     for i, mask_dict in enumerate(masks):
         mask_data = (mask_dict['segmentation'] * 255).astype(np.uint8)
@@ -28,43 +29,22 @@ def segment_image(input_image):
         # Convert to base64 image
         _, buffer = cv2.imencode(".png", cv2.cvtColor(cropped_region, cv2.COLOR_BGR2RGB))
         segmented_image_base64 = base64.b64encode(buffer).decode()
-        segmented_regions.append(segmented_image_base64)  # Add to the list
 
-    return segmented_regions
+        # Call the API to get the cosine similarity
+        client = InferenceClient()
+        result = client.post(json={"inputs": segmented_image_base64}, model="https://ktllc-clip-model-inputbase64.hf.space/--replicas/mmz7z/")
 
-# Function to call the API and calculate cosine similarity
-def calculate_cosine_similarity(segmented_images):
-    highest_cosine = -1
-    highest_cosine_base64 = ""
-    
-    client = Client("https://ktllc-clip-model-inputbase64.hf.space/--replicas/mmz7z/")
-    
-    for base64_image in segmented_images:
-        # Call the API here using the base64 image
-        result = client.predict(base64_image, base64_image, api_name="/predict")
-        
-        cosine_value = result['similarity']
-        print(f"Base64 Image: {base64_image}, Cosine Similarity: {cosine_value}")
-            
-        if cosine_value > highest_cosine:
-            highest_cosine = cosine_value
-            highest_cosine_base64 = base64_image
+        cosine_similarity = result[0].get("score", 0.0)
+
+        if cosine_similarity > highest_cosine_value:
+            highest_cosine_value = cosine_similarity
+            highest_cosine_base64 = segmented_image_base64
 
-    print(f"Highest Cosine Similarity: {highest_cosine} (Base64 Image: {highest_cosine_base64})")
+    return highest_cosine_base64
 
 # Create Gradio components
 input_image = gr.inputs.Image()
-output_images = gr.outputs.JSON()
+output_image = gr.outputs.Image(type="pil")
 
 # Create a Gradio interface
-segmentation_interface = gr.Interface(fn=segment_image, inputs=input_image, outputs=output_images)
-
-# Launch the segmentation interface
-segmentation_interface.launch()
-
-# Get the segmented images from the segmentation interface
-segmented_images = segmentation_interface.run()
-segmentation_interface.close()
-
-# Call the API for each segmented image and calculate cosine similarity
-calculate_cosine_similarity(segmented_images)
+gr.Interface(fn=segment_image, inputs=input_image, outputs=output_image).launch()