new models

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+models/sam_vit_h_4b8939.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,12 +1,84 @@
 import gradio as gr
+from segment_anything import build_sam, SamAutomaticMaskGenerator
+from PIL import Image, ImageDraw
+import clip
+import torch
+import numpy as np
+
+mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="./models/sam_vit_h_4b8939.pth"))
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model, preprocess = clip.load("ViT-B/32", device=device)
+
+def convert_box_xywh_to_xyxy(box):
+    x1 = box[0]
+    y1 = box[1]
+    x2 = box[0] + box[2]
+    y2 = box[1] + box[3]
+    return [x1, y1, x2, y2]
+
+def segment_image(image, segmentation_mask):
+    image_array = np.array(image)
+    segmented_image_array = np.zeros_like(image_array)
+    segmented_image_array[segmentation_mask] = image_array[segmentation_mask]
+    segmented_image = Image.fromarray(segmented_image_array)
+    black_image = Image.new("RGB", image.size, (0, 0, 0))
+    transparency_mask = np.zeros_like(segmentation_mask, dtype=np.uint8)
+    transparency_mask[segmentation_mask] = 255
+    transparency_mask_image = Image.fromarray(transparency_mask, mode='L')
+    black_image.paste(segmented_image, mask=transparency_mask_image)
+    return black_image
+
+@torch.no_grad()
+def retriev(elements: list[Image.Image], search_text: str) -> int:
+    preprocessed_images = [preprocess(image).to(device) for image in elements]
+    tokenized_text = clip.tokenize([search_text]).to(device)
+    stacked_images = torch.stack(preprocessed_images)
+    image_features = model.encode_image(stacked_images)
+    text_features = model.encode_text(tokenized_text)
+    image_features /= image_features.norm(dim=-1, keepdim=True)
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+    probs = 100. * image_features @ text_features.T
+    return probs[:, 0].softmax(dim=0)
+
+def get_indices_of_values_above_threshold(values, threshold):
+    return [i for i, v in enumerate(values) if v > threshold]
+
 
 def pred(search_string, img):
+    # original_image = img.copy()
+
+    # masks = mask_generator.generate(img)
+    # # 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"])))
+
+    # scores = retriev(cropped_boxes, "kiwi")
+    # indices = get_indices_of_values_above_threshold(scores, 0.05)
+
+    # 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)
+
+    # 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)
+
+    # result_image = Image.alpha_composite(original_image.convert('RGBA'), overlay_image)
+
     return img
 
 iface = gr.Interface(
     fn=pred, 
-    inputs=["text", "image"], 
-    outputs="image",
+    inputs=["text", gr.inputs.Image(type="pil")], 
+    outputs=gr.outputs.Image(type="pil"),
     examples = [
         ["banana", "./imgs/test_1.jpg"],
         ["orange", "./imgs/test_1.jpg"],

models/sam_vit_h_4b8939.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e
+size 2564550879

requirements.txt

@@ -1 +1,9 @@
-gradio
+gradio
+
+torch 
+opencv-python 
+Pillow
+
+git+https://github.com/openai/CLIP.git
+
+git+https://github.com/facebookresearch/segment-anything.git