coded

app.py

@@ -5,6 +5,7 @@ import clip
 import torch
 import numpy as np
 
+# preso spunto da https://github.com/maxi-w/CLIP-SAM/blob/main/main.ipynb
 mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="./models/sam_vit_h_4b8939.pth"))
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -46,39 +47,40 @@ def get_indices_of_values_above_threshold(values, threshold):
 
 
 def pred(search_string, img):
-    # original_image = img.copy()
+    original_image = img.copy()
 
-    # masks = mask_generator.generate(img)
-    # # Cut out all masks
-    # cropped_boxes = []
+    open_cv_image = np.array(img)[:, :, ::-1] 
+    masks = mask_generator.generate(open_cv_image)
+    # Cut out all masks
+    cropped_boxes = []
 
-    # for mask in masks:
-    #     cropped_boxes.append(segment_image(img, mask["segmentation"]).crop(convert_box_xywh_to_xyxy(mask["bbox"])))
+    for mask in masks:
+        cropped_boxes.append(segment_image(img, mask["segmentation"]).crop(convert_box_xywh_to_xyxy(mask["bbox"])))
 
-    # scores = retriev(cropped_boxes, "kiwi")
-    # indices = get_indices_of_values_above_threshold(scores, 0.05)
+    scores = retriev(cropped_boxes, search_string)
+    indices = get_indices_of_values_above_threshold(scores, 0.05)
 
-    # segmentation_masks = []
+    segmentation_masks = []
 
-    # for seg_idx in indices:
-    #     segmentation_mask_image = Image.fromarray(masks[seg_idx]["segmentation"].astype('uint8') * 255)
-    #     segmentation_masks.append(segmentation_mask_image)
+    for seg_idx in indices:
+        segmentation_mask_image = Image.fromarray(masks[seg_idx]["segmentation"].astype('uint8') * 255)
+        segmentation_masks.append(segmentation_mask_image)
 
-    # overlay_image = Image.new('RGBA', img.size, (0, 0, 0, 0))
-    # overlay_color = (255, 0, 0, 200)
+    overlay_image = Image.new('RGBA', img.size, (0, 0, 0, 0))
+    overlay_color = (255, 0, 0, 200)
 
-    # draw = ImageDraw.Draw(overlay_image)
-    # for segmentation_mask_image in segmentation_masks:
-    #     draw.bitmap((0, 0), segmentation_mask_image, fill=overlay_color)
+    draw = ImageDraw.Draw(overlay_image)
+    for segmentation_mask_image in segmentation_masks:
+        draw.bitmap((0, 0), segmentation_mask_image, fill=overlay_color)
 
-    # result_image = Image.alpha_composite(original_image.convert('RGBA'), overlay_image)
+    result_image = Image.alpha_composite(original_image.convert('RGBA'), overlay_image)
 
-    return img
+    return result_image, overlay_image
 
 iface = gr.Interface(
     fn=pred, 
     inputs=["text", gr.inputs.Image(type="pil")], 
-    outputs=gr.outputs.Image(type="pil"),
+    outputs=[gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil")],
     examples = [
         ["banana", "./imgs/test_1.jpg"],
         ["orange", "./imgs/test_1.jpg"],

packages.txt

@@ -0,0 +1 @@
+python3-opencv