First Commit

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *.jit filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/fruits.jpg filter=lfs diff=lfs merge=lfs -text
+images/street.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,70 @@
+import os
+import warnings
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+from lang_efficient_sam.LangEfficientSAM import LangEfficientSAM
+from lang_efficient_sam.utils.draw_image import draw_image
+
+warnings.filterwarnings("ignore")
+
+model = LangEfficientSAM()
+
+
+def predict(box_threshold, text_threshold, image_path, text_prompt):
+    print("Predicting... ", box_threshold, text_threshold, image_path, text_prompt)
+
+    image_pil = Image.open(image_path).convert("RGB")
+
+    masks, boxes, phrases, logits = model.predict(image_pil, text_prompt, box_threshold, text_threshold)
+
+    labels = [f"{phrase} {logit:.2f}" for phrase, logit in zip(phrases, logits)]
+
+    image_array = np.asarray(image_pil)
+    image = draw_image(image_array, masks, boxes, labels)
+    image = Image.fromarray(np.uint8(image)).convert("RGB")
+
+    return image
+
+
+title = "LangEfficientSAM"
+
+inputs = [
+    gr.Slider(0, 1, value=0.3, label="Box threshold"),
+    gr.Slider(0, 1, value=0.25, label="Text threshold"),
+    gr.Image(type="filepath", label='Image'),
+    gr.Textbox(lines=1, label="Text Prompt"),
+]
+
+outputs = [gr.Image(type="pil", label="Output Image")]
+
+examples = [
+    [
+        0.20,
+        0.20,
+        os.path.join(os.path.dirname(__file__), "images", "living.jpg"),
+        "fabric",
+    ],
+    [
+        0.36,
+        0.25,
+        os.path.join(os.path.dirname(__file__), "images", "fruits.jpg"),
+        "apple",
+    ],
+    [
+        0.20,
+        0.20,
+        os.path.join(os.path.dirname(__file__), "images", "street.jpg"),
+        "car",
+    ]
+]
+
+demo = gr.Interface(fn=predict,
+                    inputs=inputs,
+                    outputs=outputs,
+                    examples=examples,
+                    title=title)
+
+demo.launch(debug=False, share=False)

lang_efficient_sam/LangEfficientSAM.py

@@ -0,0 +1,107 @@
+import groundingdino.datasets.transforms as T
+import torch
+import numpy as np
+from groundingdino.models import build_model
+from groundingdino.util import box_ops
+from groundingdino.util.inference import predict
+from groundingdino.util.slconfig import SLConfig
+from groundingdino.util.utils import clean_state_dict
+from torchvision.transforms import ToTensor
+from huggingface_hub import hf_hub_download
+import time
+
+
+def load_model_hugging_face(repo_id, filename, ckpt_config_filename, device='cpu'):
+    cache_config_file = hf_hub_download(repo_id=repo_id, filename=ckpt_config_filename)
+
+    args = SLConfig.fromfile(cache_config_file)
+    model = build_model(args)
+    args.device = device
+
+    cache_file = hf_hub_download(repo_id=repo_id, filename=filename)
+    checkpoint = torch.load(cache_file, map_location='cpu')
+    log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False)
+    print(f"Model loaded from {cache_file} \n => {log}")
+    model.eval()
+    return model
+
+
+class LangEfficientSAM:
+    def __init__(self, device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
+        self.device = device
+        print("Device:", self.device)
+        if self.device == "cpu":
+            self.sam_efficient = torch.jit.load('./models/efficientsam_s_cpu.jit')
+        else:
+            self.sam_efficient = torch.jit.load('./models/efficientsam_s_gpu.jit')
+        ckpt_repo_id = "ShilongLiu/GroundingDINO"
+        ckpt_filename = "groundingdino_swinb_cogcoor.pth"
+        ckpt_config_filename = "GroundingDINO_SwinB.cfg.py"
+        self.groundingdino = load_model_hugging_face(ckpt_repo_id,
+                                                     ckpt_filename,
+                                                     ckpt_config_filename,
+                                                     self.device)
+
+    def predict_dino(self, image_pil, text_prompt, box_threshold, text_threshold):
+        start = time.time()
+        transform = T.Compose([
+            T.RandomResize([800], max_size=1333),
+            T.ToTensor(),
+            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+        image_transformed, _ = transform(image_pil, None)
+
+        boxes, logits, phrases = predict(model=self.groundingdino,
+                                         image=image_transformed,
+                                         caption=text_prompt,
+                                         box_threshold=box_threshold,
+                                         text_threshold=text_threshold,
+                                         device=self.device)
+        W, H = image_pil.size
+        boxes = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H])
+        # print("DINO time: ", time.time() - start)
+
+        return boxes, logits, phrases
+
+    def predict_sam(self, image, box):
+        start = time.time()
+        img_tensor = ToTensor()(image).to(device=self.device)
+        bbox = torch.reshape(box.clone().detach(), [1, 1, 2, 2]).to(device=self.device)
+        bbox_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2]).to(device=self.device)
+
+        predicted_logits, predicted_iou = self.sam_efficient(
+            img_tensor[None, ...],
+            bbox,
+            bbox_labels,
+        )
+        predicted_logits = predicted_logits.cpu()
+        all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
+        predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
+
+        max_predicted_iou = -1
+        selected_mask_using_predicted_iou = None
+        for m in range(all_masks.shape[0]):
+            curr_predicted_iou = predicted_iou[m]
+            if (
+                    curr_predicted_iou > max_predicted_iou
+                    or selected_mask_using_predicted_iou is None
+            ):
+                max_predicted_iou = curr_predicted_iou
+                selected_mask_using_predicted_iou = all_masks[m]
+
+        # print("SAM time: ", time.time() - start)
+        return selected_mask_using_predicted_iou
+
+    def predict(self, image_pil, text_prompt, box_threshold=0.3, text_threshold=0.25):
+        boxes, logits, phrases = self.predict_dino(image_pil, text_prompt, box_threshold, text_threshold)
+        # masks = torch.tensor([])
+        masks = []
+        if len(boxes) > 0:
+            for box in boxes:
+                mask = self.predict_sam(image_pil, box)
+                masks.append(mask)
+
+        masks = np.array(masks)
+        masks = torch.from_numpy(masks)
+
+        return masks, boxes, phrases, logits

lang_efficient_sam/utils/draw_image.py

@@ -0,0 +1,12 @@
+import torch
+from torchvision.utils import draw_bounding_boxes
+from torchvision.utils import draw_segmentation_masks
+
+
+def draw_image(image, masks, boxes, labels, alpha=0.4):
+    image = torch.from_numpy(image).permute(2, 0, 1)
+    if len(boxes) > 0:
+        image = draw_bounding_boxes(image, boxes, colors=['red'] * len(boxes), labels=labels, width=2)
+    if len(masks) > 0:
+        image = draw_segmentation_masks(image, masks=masks, colors=['cyan'] * len(masks), alpha=alpha)
+    return image.numpy().transpose(1, 2, 0)

models/efficientsam_s_cpu.jit

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

models/efficientsam_s_gpu.jit

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

requirements.txt

@@ -0,0 +1,8 @@
+gradio==4.18.0
+groundingdino-py==0.4.0
+matplotlib==3.8.2
+numpy==1.26.4
+opencv-python==4.9.0.80
+pillow==10.2.0
+torch==2.2.0
+torchvision==0.17.0