Update app.py

app.py

@@ -17,34 +17,33 @@ sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model, preprocess = clip.load("ViT-B/32", device=device)
 
-
 def find_similarity(base64_image, text_input):
-    # 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 = preprocess(image).unsqueeze(0).to(device)
-
-    # Prepare input text
-    text_tokens = clip.tokenize([text_input]).to(device)
+    try:
+        # Decode the base64 image to bytes
+        image_bytes = base64.b64decode(base64_image)
 
-    # Encode image and text features
+        # Convert the bytes to a PIL image
+        image = Image.open(BytesIO(image_bytes))
 
+        # Preprocess the image
+        image = preprocess(image).unsqueeze(0).to(device)
 
-    with torch.no_grad():
-        image_features = model.encode_image(image)
-        text_features = model.encode_text(text_tokens)
+        # Prepare input text
+        text_tokens = clip.tokenize([text_input]).to(device)
 
-    # Normalize features and calculate similarity
-    image_features /= image_features.norm(dim=-1, keepdim=True)
-    text_features /= text_features.norm(dim=-1, keepdim=True)
-    similarity = (text_features @ image_features.T).squeeze(0).cpu().numpy()
+        # Encode image and text features
+        with torch.no_grad():
+            image_features = model.encode_image(image)
+            text_features = model.encode_text(text_tokens)
 
-    return similarity
+        # Normalize features and calculate similarity
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        similarity = (text_features @ image_features.T).squeeze(0).cpu().numpy()
 
+        return similarity
+    except Exception as e:
+        return str(e)
 
 # Define a function for image segmentation
 def segment_image(input_image, text_input):