fix: pin supervision==0.21.0 and add map_location for HQSAM checkpoint

EfficientSAM/grounded_light_hqsam.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import numpy as np
 import supervision as sv
@@ -12,24 +13,29 @@
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # GroundingDINO config and checkpoint
-GROUNDING_DINO_CONFIG_PATH = "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py"
-GROUNDING_DINO_CHECKPOINT_PATH = "./groundingdino_swint_ogc.pth"
+
+BASE_PATH = "/home/saige/Workspace/Grounded-Segment-Anything/"
+GROUNDING_DINO_CONFIG_PATH = os.path.join(BASE_PATH,"GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py")
+GROUNDING_DINO_CHECKPOINT_PATH = os.path.join(BASE_PATH,"groundingdino_swint_ogc.pth")
 
 # Building GroundingDINO inference model
-grounding_dino_model = Model(model_config_path=GROUNDING_DINO_CONFIG_PATH, model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH)
+grounding_dino_model = Model(model_config_path=GROUNDING_DINO_CONFIG_PATH, model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH, device=DEVICE)
 
 # Building MobileSAM predictor
-HQSAM_CHECKPOINT_PATH = "./EfficientSAM/sam_hq_vit_tiny.pth"
-checkpoint = torch.load(HQSAM_CHECKPOINT_PATH)
+HQSAM_CHECKPOINT_PATH = os.path.join(BASE_PATH,"EfficientSAM/sam_hq_vit_tiny.pth")
+checkpoint = torch.load(HQSAM_CHECKPOINT_PATH, map_location=DEVICE)
 light_hqsam = setup_model()
 light_hqsam.load_state_dict(checkpoint, strict=True)
 light_hqsam.to(device=DEVICE)
 
 sam_predictor = SamPredictor(light_hqsam)
 
 
+grounding_dino_model.model.eval()
+sam_predictor.model.eval()
+
 # Predict classes and hyper-param for GroundingDINO
-SOURCE_IMAGE_PATH = "./EfficientSAM/LightHQSAM/example_light_hqsam.png"
+SOURCE_IMAGE_PATH = os.path.join(BASE_PATH,"./EfficientSAM/LightHQSAM/example_light_hqsam.png")
 CLASSES = ["bench"]
 BOX_THRESHOLD = 0.25
 TEXT_THRESHOLD = 0.25
@@ -56,7 +62,7 @@
 annotated_frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)
 
 # save the annotated grounding dino image
-cv2.imwrite("EfficientSAM/LightHQSAM/groundingdino_annotated_image.jpg", annotated_frame)
+cv2.imwrite("./dino_annotated_image.jpg", annotated_frame)
 
 
 # NMS post process
@@ -106,4 +112,4 @@ def segment(sam_predictor: SamPredictor, image: np.ndarray, xyxy: np.ndarray) ->
 annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
 
 # save the annotated grounded-sam image
-cv2.imwrite("EfficientSAM/LightHQSAM/grounded_light_hqsam_annotated_image.jpg", annotated_image)
+cv2.imwrite("./sam_annotated_image.jpg", annotated_image)

EfficientSAM/grounded_light_hqsam_server.py

@@ -0,0 +1,112 @@
+import os
+import cv2
+import numpy as np
+import supervision as sv
+
+import torch
+import torchvision
+
+from groundingdino.util.inference import Model
+from segment_anything import SamPredictor
+from LightHQSAM.setup_light_hqsam import setup_model
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# GroundingDINO config and checkpoint
+
+BASE_PATH = "/home/saige/Workspace/Grounded-Segment-Anything/"
+GROUNDING_DINO_CONFIG_PATH = os.path.join(BASE_PATH,"GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py")
+GROUNDING_DINO_CHECKPOINT_PATH = os.path.join(BASE_PATH,"groundingdino_swint_ogc.pth")
+
+# Building GroundingDINO inference model
+grounding_dino_model = Model(model_config_path=GROUNDING_DINO_CONFIG_PATH, model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH, device=DEVICE)
+
+# Building MobileSAM predictor
+HQSAM_CHECKPOINT_PATH = os.path.join(BASE_PATH,"EfficientSAM/sam_hq_vit_tiny.pth")
+checkpoint = torch.load(HQSAM_CHECKPOINT_PATH, map_location=DEVICE)
+light_hqsam = setup_model()
+light_hqsam.load_state_dict(checkpoint, strict=True)
+light_hqsam.to(device=DEVICE)
+
+sam_predictor = SamPredictor(light_hqsam)
+
+
+# Predict classes and hyper-param for GroundingDINO
+SOURCE_IMAGE_PATH = os.path.join(BASE_PATH,"./EfficientSAM/LightHQSAM/example_light_hqsam.png")
+CLASSES = ["bench"]
+BOX_THRESHOLD = 0.25
+TEXT_THRESHOLD = 0.25
+NMS_THRESHOLD = 0.8
+
+
+# load image
+image = cv2.imread(SOURCE_IMAGE_PATH)
+
+# detect objects
+detections = grounding_dino_model.predict_with_classes(
+    image=image,
+    classes=CLASSES,
+    box_threshold=BOX_THRESHOLD,
+    text_threshold=TEXT_THRESHOLD
+)
+
+# annotate image with detections
+box_annotator = sv.BoxAnnotator()
+labels = [
+    f"{CLASSES[class_id]} {confidence:0.2f}" 
+    for _, _, confidence, class_id, _, _
+    in detections]
+annotated_frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)
+
+# save the annotated grounding dino image
+cv2.imwrite("EfficientSAM/LightHQSAM/groundingdino_annotated_image.jpg", annotated_frame)
+
+
+# NMS post process
+print(f"Before NMS: {len(detections.xyxy)} boxes")
+nms_idx = torchvision.ops.nms(
+    torch.from_numpy(detections.xyxy), 
+    torch.from_numpy(detections.confidence), 
+    NMS_THRESHOLD
+).numpy().tolist()
+
+detections.xyxy = detections.xyxy[nms_idx]
+detections.confidence = detections.confidence[nms_idx]
+detections.class_id = detections.class_id[nms_idx]
+
+print(f"After NMS: {len(detections.xyxy)} boxes")
+
+# Prompting SAM with detected boxes
+def segment(sam_predictor: SamPredictor, image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
+    sam_predictor.set_image(image)
+    result_masks = []
+    for box in xyxy:
+        masks, scores, logits = sam_predictor.predict(
+            box=box,
+            multimask_output=False,
+            hq_token_only=True,
+        )
+        index = np.argmax(scores)
+        result_masks.append(masks[index])
+    return np.array(result_masks)
+
+
+# convert detections to masks
+detections.mask = segment(
+    sam_predictor=sam_predictor,
+    image=cv2.cvtColor(image, cv2.COLOR_BGR2RGB),
+    xyxy=detections.xyxy
+)
+
+# annotate image with detections
+box_annotator = sv.BoxAnnotator()
+mask_annotator = sv.MaskAnnotator()
+labels = [
+    f"{CLASSES[class_id]} {confidence:0.2f}" 
+    for _, _, confidence, class_id, _, _
+    in detections]
+annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
+annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
+
+# save the annotated grounded-sam image
+cv2.imwrite("EfficientSAM/LightHQSAM/grounded_light_hqsam_annotated_image.jpg", annotated_image)

GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py

@@ -16,7 +16,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint as checkpoint
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+#from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.layers import DropPath, to_2tuple, trunc_normal_
 
 from groundingdino.util.misc import NestedTensor
 

GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu

@@ -32,11 +32,25 @@ at::Tensor ms_deform_attn_cuda_forward(
     AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
     AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
 
+
+    // Original Code
+    /*
     AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
     AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor");
     AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor");
     AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor");
     AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor");
+    */
+
+    // API Change VAL.type().is_cuda() => VAL.is_cuda()
+
+    AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
+
+
 
     const int batch = value.size(0);
     const int spatial_size = value.size(1);
@@ -62,7 +76,8 @@ at::Tensor ms_deform_attn_cuda_forward(
     for (int n = 0; n < batch/im2col_step_; ++n)
     {
         auto columns = output_n.select(0, n);
-        AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] {
+        //AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] { // Old CUDA Code
+        AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_forward_cuda", ([&] {
             ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(),
                 value.data<scalar_t>() + n * im2col_step_ * per_value_size,
                 spatial_shapes.data<int64_t>(),
@@ -98,12 +113,26 @@ std::vector<at::Tensor> ms_deform_attn_cuda_backward(
     AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
     AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous");
 
+
+
+    // Original Code
+    /*
     AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
     AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor");
     AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor");
     AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor");
     AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor");
     AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor");
+    */
+
+    // API Change VAL.type().is_cuda() => VAL.is_cuda()
+    AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
+    AT_ASSERTM(grad_output.is_cuda(), "grad_output must be a CUDA tensor");
+
 
     const int batch = value.size(0);
     const int spatial_size = value.size(1);
@@ -132,7 +161,8 @@ std::vector<at::Tensor> ms_deform_attn_cuda_backward(
     for (int n = 0; n < batch/im2col_step_; ++n)
     {
         auto grad_output_g = grad_output_n.select(0, n);
-        AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] {
+        //AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] { // Old CUDA Code
+        AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_backward_cuda", ([&] {
             ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(),
                                     grad_output_g.data<scalar_t>(),
                                     value.data<scalar_t>() + n * im2col_step_ * per_value_size,
@@ -153,4 +183,4 @@ std::vector<at::Tensor> ms_deform_attn_cuda_backward(
     };
 }
 
-} // namespace groundingdino
+} // namespace groundingdino

GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py

@@ -8,7 +8,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from timm.models.layers import DropPath
+#from timm.models.layers import DropPath
+from timm.layers import DropPath
 
 
 class FeatureResizer(nn.Module):

requirements.txt

@@ -11,7 +11,7 @@ pycocotools
 PyYAML
 requests
 setuptools
-supervision
+supervision==0.21.0
 termcolor
 timm
 torch

server/EdgeSAM/common.py

@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Type
+
+
+class MLPBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        mlp_dim: int,
+        act: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin2(self.act(self.lin1(x)))
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
+class LayerNorm2d(nn.Module):
+    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+def val2list(x: list or tuple or any, repeat_time=1) -> list:
+    if isinstance(x, (list, tuple)):
+        return list(x)
+    return [x for _ in range(repeat_time)]
+
+
+def val2tuple(x: list or tuple or any, min_len: int = 1, idx_repeat: int = -1) -> tuple:
+    x = val2list(x)
+
+    # repeat elements if necessary
+    if len(x) > 0:
+        x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
+
+    return tuple(x)
+
+
+def list_sum(x: list) -> any:
+    return x[0] if len(x) == 1 else x[0] + list_sum(x[1:])
+
+
+def resize(
+        x: torch.Tensor,
+        size: any or None = None,
+        scale_factor=None,
+        mode: str = "bicubic",
+        align_corners: bool or None = False,
+) -> torch.Tensor:
+    if mode in ["bilinear", "bicubic"]:
+        return F.interpolate(
+            x,
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+        )
+    elif mode in ["nearest", "area"]:
+        return F.interpolate(x, size=size, scale_factor=scale_factor, mode=mode)
+    else:
+        raise NotImplementedError(f"resize(mode={mode}) not implemented.")
+
+
+class UpSampleLayer(nn.Module):
+    def __init__(
+            self,
+            mode="bicubic",
+            size=None,
+            factor=2,
+            align_corners=False,
+    ):
+        super(UpSampleLayer, self).__init__()
+        self.mode = mode
+        self.size = val2list(size, 2) if size is not None else None
+        self.factor = None if self.size is not None else factor
+        self.align_corners = align_corners
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return resize(x, self.size, self.factor, self.mode, self.align_corners)
+
+
+class OpSequential(nn.Module):
+    def __init__(self, op_list):
+        super(OpSequential, self).__init__()
+        valid_op_list = []
+        for op in op_list:
+            if op is not None:
+                valid_op_list.append(op)
+        self.op_list = nn.ModuleList(valid_op_list)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for op in self.op_list:
+            x = op(x)
+        return x