Update app.py

app.py

@@ -8,7 +8,6 @@ from io import BytesIO
 import torch
 import clip
 
-
 # Load the segmentation model
 sam_checkpoint = "sam_vit_h_4b8939.pth"
 model_type = "vit_h"
@@ -48,6 +47,7 @@ def find_similarity(base64_image, text_input):
     except Exception as e:
         return str(e)
 
+
 def segment_image(input_image, text_input):
     image_bytes = base64.b64decode(input_image)
     image = Image.open(BytesIO(image_bytes))
@@ -76,6 +76,10 @@ def segment_image(input_image, text_input):
         x, y, w, h = map(int, mask_dict['bbox'])
         cropped_region = segmented_region[y:y+h, x:x+w]
 
+        if not cropped_region.size:
+            # If the cropped region is empty, return the input image as is
+            return input_image
+
         # Convert to base64 image
         _, buffer = cv2.imencode(".png", cv2.cvtColor(cropped_region, cv2.COLOR_BGR2RGB))
         segmented_image_base64 = base64.b64encode(buffer).decode()
@@ -95,10 +99,10 @@ def segment_image(input_image, text_input):
     # Return the segmented images in descending order of similarity
     return segmented_regions
 
+
 # Create Gradio components
 input_image = gr.Textbox(label="Base64 Image", lines=8)
 text_input = gr.Textbox(label="Text Input")  # Use Textbox with a label
-#output_images = gr.outputs.JSON()
 
 # Create a Gradio interface
 gr.Interface(fn=segment_image, inputs=[input_image, text_input], outputs="text").launch()