Remove backgrounds from cropped images

app.py

@@ -57,7 +57,7 @@ def get_scores(crops: List[PIL.Image.Image], query: str) -> torch.Tensor:
     txt_features = model.encode_text(token)
     img_features /= img_features.norm(dim=-1, keepdim=True)
     txt_features /= txt_features.norm(dim=-1, keepdim=True)
-    similarity = (100.0 * img_features @ txt_features.T).softmax(dim=0)
+    similarity = (100 * img_features @ txt_features.T).softmax(0)
     return similarity
 
 
@@ -82,9 +82,10 @@ def filter_masks(
         filtered_masks.append(mask)
 
         x, y, w, h = mask["bbox"]
-        crop = image[y: y + h, x: x + w]
+        masked = image * np.expand_dims(mask["segmentation"], -1)
+        crop = masked[y: y + h, x: x + w]
         crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
-        crop = PIL.Image.fromarray(np.uint8(crop * 255)).convert("RGB")
+        crop = PIL.Image.fromarray(crop * 255)
         crop.resize((CLIP_WIDTH, CLIP_HEIGHT))
         cropped_masks.append(crop)
 
@@ -141,7 +142,7 @@ def segment(
     )
     image = draw_masks(image, masks)
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    image = PIL.Image.fromarray(np.uint8(image)).convert("RGB")
+    image = PIL.Image.fromarray(image)
     return image
 
 
@@ -161,14 +162,14 @@ demo = gr.Interface(
         [
             0.9,
             0.8,
-            0.15,
+            0.30,
             os.path.join(os.path.dirname(__file__), "examples/dog.jpg"),
-            "A dog only",
+            "A dog",
         ],
         [
             0.9,
             0.8,
-            0.1,
+            0.05,
             os.path.join(os.path.dirname(__file__), "examples/city.jpg"),
             "A bridge on the water",
         ],
@@ -177,7 +178,7 @@ demo = gr.Interface(
             0.8,
             0.05,
             os.path.join(os.path.dirname(__file__), "examples/food.jpg"),
-            "",
+            "spoon",
         ],
         [
             0.9,