Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ 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/notebooks/images/groceries.jpg filter=lfs diff=lfs merge=lfs -text
+models/sam/notebooks/images/truck.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,31 @@
----
-title: Video Segmentation
-emoji: 👁
-colorFrom: blue
-colorTo: purple
-sdk: gradio
-sdk_version: 6.2.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Surgical-DeSAM
+emoji: 🔬
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+# Surgical-DeSAM
+
+Surgical instrument segmentation using DeSAM (Decoupled SAM) architecture.
+
+Upload a surgical image to get instance segmentation masks for instruments.
+
+## Model
+
+Based on [Surgical-DeSAM](https://github.com/YuyangSheng/Surgical-DeSAM) - integrates DETR for detection and SAM for segmentation.
+
+## Classes Detected
+- Bipolar Forceps
+- Prograsp Forceps  
+- Large Needle Driver
+- Monopolar Curved Scissors
+- Ultrasound Probe
+- Suction
+- Clip Applier
+- Stapler

app.py

@@ -1,7 +1,264 @@
-import gradio as gr
-
-def greet(name):
-    return "Hello " + name + "!!"
-
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+"""
+Surgical-DeSAM Gradio App for Hugging Face Spaces
+Uses ZeroGPU for inference
+"""
+import os
+import spaces
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from huggingface_hub import hf_hub_download
+
+# Model imports (will be copied to hf_space)
+from models.detr_seg import DETR, SAMModel
+from models.backbone import build_backbone
+from models.transformer import build_transformer
+from util.misc import NestedTensor
+
+# Configuration
+MODEL_REPO = os.environ.get("MODEL_REPO", "IFMedTech/surgical-desam-weights")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+INSTRUMENT_CLASSES = (
+    'bipolar_forceps', 'prograsp_forceps', 'large_needle_driver', 
+    'monopolar_curved_scissors', 'ultrasound_probe', 'suction', 
+    'clip_applier', 'stapler'
+)
+
+COLORS = [
+    [0, 114, 189], [217, 83, 25], [237, 177, 32],
+    [126, 47, 142], [119, 172, 48], [77, 190, 238],
+    [162, 20, 47], [76, 76, 76]
+]
+
+# Global model variables
+model = None
+seg_model = None
+device = None
+
+
+def download_weights():
+    """Download model weights from private HF repo"""
+    weights_dir = "weights"
+    os.makedirs(weights_dir, exist_ok=True)
+    
+    # Download DeSAM weights
+    desam_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename="surgical_desam_1024.pth",
+        token=HF_TOKEN,
+        local_dir=weights_dir
+    )
+    
+    # Download SAM weights
+    sam_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename="sam_vit_b_01ec64.pth",
+        token=HF_TOKEN,
+        local_dir=weights_dir
+    )
+    
+    # Download Swin backbone
+    swin_dir = "swin_backbone"
+    os.makedirs(swin_dir, exist_ok=True)
+    swin_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename="swin_base_patch4_window7_224_22kto1k.pth",
+        token=HF_TOKEN,
+        local_dir=swin_dir
+    )
+    
+    return desam_path, sam_path
+
+
+class Args:
+    """Mock args for model building"""
+    backbone = 'swin_B_224_22k'
+    dilation = False
+    position_embedding = 'sine'
+    hidden_dim = 256
+    dropout = 0.1
+    nheads = 8
+    dim_feedforward = 2048
+    enc_layers = 6
+    dec_layers = 6
+    pre_norm = False
+    num_queries = 100
+    aux_loss = False
+    lr_backbone = 1e-5
+    masks = False
+    dataset_file = 'endovis18'
+    device = 'cuda'
+    backbone_dir = './swin_backbone'
+
+
+def load_models():
+    """Load DETR and SAM models"""
+    global model, seg_model, device
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Download weights
+    desam_path, sam_path = download_weights()
+    
+    # Build model
+    args = Args()
+    args.device = str(device)
+    
+    backbone = build_backbone(args)
+    transformer = build_transformer(args)
+    
+    model = DETR(
+        backbone,
+        transformer,
+        num_classes=9,  # 8 classes + background
+        num_queries=args.num_queries,
+        aux_loss=args.aux_loss,
+    )
+    
+    # Load weights
+    checkpoint = torch.load(desam_path, map_location='cpu')
+    model.load_state_dict(checkpoint['model'], strict=False)
+    model.to(device)
+    model.eval()
+    
+    # Load SAM model
+    seg_model = SAMModel(device=device, ckpt_path=sam_path)
+    if 'seg_model' in checkpoint:
+        seg_model.load_state_dict(checkpoint['seg_model'])
+    seg_model.to(device)
+    seg_model.eval()
+    
+    print("Models loaded successfully!")
+
+
+def preprocess_image(image):
+    """Preprocess image for model input"""
+    # Resize to 1024x1024
+    img = cv2.resize(np.array(image), (1024, 1024))
+    img = img.astype(np.float32) / 255.0
+    
+    # Normalize
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    img = (img - mean) / std
+    
+    # Convert to tensor
+    img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float()
+    return img_tensor
+
+
+def box_cxcywh_to_xyxy(x):
+    """Convert boxes from center format to corner format"""
+    x_c, y_c, w, h = x.unbind(-1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
+         (x_c + 0.5 * w), (y_c + 0.5 * h)]
+    return torch.stack(b, dim=-1)
+
+
+def create_visualization(image, boxes, labels, masks, scores):
+    """Create visualization with boxes and masks"""
+    img = np.array(image).copy()
+    h, w = img.shape[:2]
+    
+    for i, (box, label, mask, score) in enumerate(zip(boxes, labels, masks, scores)):
+        if score < 0.3:
+            continue
+            
+        color = COLORS[label % len(COLORS)]
+        
+        # Draw mask
+        mask_resized = cv2.resize(mask, (w, h))
+        mask_bool = mask_resized > 0.5
+        overlay = img.copy()
+        overlay[mask_bool] = color
+        img = cv2.addWeighted(img, 0.6, overlay, 0.4, 0)
+        
+        # Draw box
+        x1, y1, x2, y2 = box.astype(int)
+        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+        
+        # Draw label
+        label_text = f"{INSTRUMENT_CLASSES[label]}: {score:.2f}"
+        cv2.putText(img, label_text, (x1, y1 - 10), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    
+    return Image.fromarray(img)
+
+
+@spaces.GPU
+def predict(image):
+    """Run inference on input image"""
+    global model, seg_model, device
+    
+    if model is None:
+        load_models()
+    
+    if image is None:
+        return None
+    
+    # Preprocess
+    img_tensor = preprocess_image(image).unsqueeze(0).to(device)
+    
+    # Create nested tensor
+    mask = torch.zeros((1, 1024, 1024), dtype=torch.bool, device=device)
+    samples = NestedTensor(img_tensor, mask)
+    
+    # Run detection
+    with torch.no_grad():
+        outputs, image_embeddings = model(samples)
+        
+        # Get predictions
+        probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
+        keep = probas.max(-1).values > 0.3
+        
+        if not keep.any():
+            return image  # No detections
+        
+        # Get boxes
+        boxes = outputs['pred_boxes'][0, keep]
+        scores = probas[keep].max(-1).values.cpu().numpy()
+        labels = probas[keep].argmax(-1).cpu().numpy()
+        
+        # Scale boxes to image size
+        h, w = image.size[1], image.size[0]
+        boxes_scaled = box_cxcywh_to_xyxy(boxes) * torch.tensor([w, h, w, h], device=device)
+        boxes_np = boxes_scaled.cpu().numpy()
+        
+        # Run segmentation
+        low_res_masks, pred_masks, _ = seg_model(
+            img_tensor, boxes, image_embeddings, 
+            sizes=(1024, 1024), add_noise=False
+        )
+        masks_np = pred_masks.cpu().numpy()
+    
+    # Create visualization
+    result = create_visualization(image, boxes_np, labels, masks_np, scores)
+    
+    return result
+
+
+# Create Gradio interface
+with gr.Blocks(title="Surgical-DeSAM") as demo:
+    gr.Markdown("# 🔬 Surgical-DeSAM")
+    gr.Markdown("Upload a surgical image to segment instruments.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            submit_btn = gr.Button("Segment", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Segmentation Result")
+    
+    submit_btn.click(fn=predict, inputs=input_image, outputs=output_image)
+    
+    gr.Examples(
+        examples=[],  # Add example images if available
+        inputs=input_image
+    )
+
+if __name__ == "__main__":
+    demo.launch()

models/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+from .detr_seg import build
+
+
+def build_model(args):
+    return build(args)

models/backbone.py

@@ -0,0 +1,213 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Backbone modules.
+"""
+from collections import OrderedDict
+
+import os
+import torch
+import torch.nn.functional as F
+import torchvision
+from torch import nn
+from torchvision.models._utils import IntermediateLayerGetter
+from typing import Dict, List
+
+from util.misc import NestedTensor, is_main_process, clean_state_dict
+
+from .position_encoding import build_position_encoding
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
+from .swin_transformer import build_swin_transformer
+
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed.
+
+    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
+    without which any other models than torchvision.models.resnet[18,34,50,101]
+    produce nans.
+    """
+
+    def __init__(self, n):
+        super(FrozenBatchNorm2d, self).__init__()
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        num_batches_tracked_key = prefix + 'num_batches_tracked'
+        if num_batches_tracked_key in state_dict:
+            del state_dict[num_batches_tracked_key]
+
+        super(FrozenBatchNorm2d, self)._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs)
+
+    def forward(self, x):
+        # move reshapes to the beginning
+        # to make it fuser-friendly
+        w = self.weight.reshape(1, -1, 1, 1)
+        b = self.bias.reshape(1, -1, 1, 1)
+        rv = self.running_var.reshape(1, -1, 1, 1)
+        rm = self.running_mean.reshape(1, -1, 1, 1)
+        eps = 1e-5
+        scale = w * (rv + eps).rsqrt()
+        bias = b - rm * scale
+        return x * scale + bias
+
+
+class BackboneBase(nn.Module):
+
+    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
+        super().__init__()
+        for name, parameter in backbone.named_parameters():
+            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
+                parameter.requires_grad_(False)
+        # if return_interm_layers:
+        #     return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
+        # else:
+        #     return_layers = {'layer4': "0"}
+        
+        # swin
+        return_interm_indices = [1, 2, 3]
+        return_layers = {}
+        for idx, layer_index in enumerate(return_interm_indices):
+            return_layers.update({"layer{}".format(5 - len(return_interm_indices) + idx): "{}".format(layer_index)})
+
+        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
+        self.num_channels = num_channels
+
+    def forward(self, tensor_list: NestedTensor):
+        xs = self.body(tensor_list.tensors)
+        # xs = OrderedDict()
+        # xs['0'] = self.sam_encoder(tensor_list.tensors)
+
+        out: Dict[str, NestedTensor] = {}
+        for name, x in xs.items():
+            m = tensor_list.mask
+            assert m is not None
+            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
+            out[name] = NestedTensor(x, mask)
+        return out
+
+
+class Backbone(BackboneBase):
+    """ResNet backbone with frozen BatchNorm."""
+    def __init__(self, name: str,
+                 train_backbone: bool,
+                 return_interm_layers: bool,
+                 dilation: bool):
+        if name in ['resnet18', 'resnet34', 'resnet50', 'resnet101']:
+            backbone = getattr(torchvision.models, name)(
+                replace_stride_with_dilation=[False, False, dilation],
+                pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
+            num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
+        
+        super().__init__(backbone, train_backbone, num_channels, return_interm_layers)
+
+
+class Joiner(nn.Sequential):
+    def __init__(self, backbone, position_embedding):
+        super().__init__(backbone, position_embedding)
+
+    def forward(self, tensor_list: NestedTensor):
+        xs = self[0](tensor_list)
+        out: List[NestedTensor] = []
+        pos = []
+        for name, x in xs.items():
+            out.append(x)
+            # position encoding
+            pos.append(self[1](x).to(x.tensors.dtype))
+
+        return out, pos
+
+
+# def build_backbone(args):
+#     position_embedding = build_position_encoding(args)
+
+#     # sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth")
+#     # position_embedding = sam.prompt_encoder.get_dense_pe() # (bs, 256, 64, 64)
+
+#     train_backbone = args.lr_backbone > 0
+#     return_interm_layers = args.masks
+#     backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)
+#     model = Joiner(backbone, position_embedding)
+#     model.num_channels = backbone.num_channels
+#     return model
+
+def build_backbone(args):
+    """
+    Useful args:
+        - backbone: backbone name
+        - lr_backbone: 
+        - dilation
+        - return_interm_indices: available: [0,1,2,3], [1,2,3], [3]
+        - backbone_freeze_keywords: 
+        - use_checkpoint: for swin only for now
+
+    """
+    backbone_freeze_keywords = None
+    use_checkpoint = True
+    position_embedding = build_position_encoding(args)
+    train_backbone = args.lr_backbone > 0
+    if not train_backbone:
+        raise ValueError("Please set lr_backbone > 0")
+    return_interm_indices = [1, 2, 3]
+    assert return_interm_indices in [[0,1,2,3], [1,2,3], [3]]
+    
+
+    if args.backbone in ['resnet50', 'resnet101']:
+        backbone = Backbone(args.backbone, train_backbone, args.dilation,   
+                                return_interm_indices,   
+                                batch_norm=FrozenBatchNorm2d)
+        bb_num_channels = backbone.num_channels
+    elif args.backbone in ['swin_T_224_1k', 'swin_B_224_22k', 'swin_B_384_22k', 'swin_L_224_22k', 'swin_L_384_22k']:
+        pretrain_img_size = int(args.backbone.split('_')[-2])
+        backbone = build_swin_transformer(args.backbone, \
+                    pretrain_img_size=pretrain_img_size, \
+                    out_indices=tuple(return_interm_indices), \
+                dilation=args.dilation, use_checkpoint=use_checkpoint)
+
+        # freeze some layers
+        if backbone_freeze_keywords is not None:
+            for name, parameter in backbone.named_parameters():
+                for keyword in backbone_freeze_keywords:
+                    if keyword in name:
+                        parameter.requires_grad_(False)
+                        break
+        if "backbone_dir" in args:
+            pretrained_dir = args.backbone_dir
+            PTDICT = {
+                'swin_T_224_1k': 'swin_tiny_patch4_window7_224.pth',
+                'swin_B_224_22k': 'swin_base_patch4_window7_224_22kto1k.pth',
+                'swin_B_384_22k': 'swin_base_patch4_window12_384.pth',
+                'swin_L_384_22k': 'swin_large_patch4_window12_384_22k.pth',
+            }
+            pretrainedpath = os.path.join(pretrained_dir, PTDICT[args.backbone])
+            checkpoint = torch.load(pretrainedpath, map_location='cpu')['model']
+            from collections import OrderedDict
+            def key_select_function(keyname):
+                if 'head' in keyname:
+                    return False
+                if args.dilation and 'layers.3' in keyname:
+                    return False
+                return True
+            _tmp_st = OrderedDict({k:v for k, v in clean_state_dict(checkpoint).items() if key_select_function(k)})
+            _tmp_st_output = backbone.load_state_dict(_tmp_st, strict=False)
+            print(str(_tmp_st_output))
+        bb_num_channels = backbone.num_features[4 - len(return_interm_indices):]
+    # elif args.backbone in ['convnext_xlarge_22k']:
+    #     backbone = build_convnext(modelname=args.backbone, pretrained=True, out_indices=tuple(return_interm_indices),backbone_dir=args.backbone_dir)
+    #     bb_num_channels = backbone.dims[4 - len(return_interm_indices):]
+    else:
+        raise NotImplementedError("Unknown backbone {}".format(args.backbone))
+    
+
+    assert len(bb_num_channels) == len(return_interm_indices), f"len(bb_num_channels) {len(bb_num_channels)} != len(return_interm_indices) {len(return_interm_indices)}"
+
+
+    model = Joiner(backbone, position_embedding)
+    model.num_channels = bb_num_channels
+    assert isinstance(bb_num_channels, List), "bb_num_channels is expected to be a List but {}".format(type(bb_num_channels))
+    return model

models/detr.py

@@ -0,0 +1,407 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+DETR model and criterion classes.
+"""
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from util import box_ops
+from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
+                       accuracy, get_world_size, interpolate,
+                       is_dist_avail_and_initialized)
+
+from .backbone import build_backbone
+from .matcher import build_matcher
+from .segmentation import (DETRsegm, PostProcessPanoptic, PostProcessSegm,
+                           dice_loss, sigmoid_focal_loss)
+from .transformer import build_transformer
+
+
+class DETR(nn.Module):
+    """ This is the DETR module that performs object detection """
+    def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):
+        """ Initializes the model.
+        Parameters:
+            backbone: torch module of the backbone to be used. See backbone.py
+            transformer: torch module of the transformer architecture. See transformer.py
+            num_classes: number of object classes
+            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
+                         DETR can detect in a single image. For COCO, we recommend 100 queries.
+            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
+        """
+        super().__init__()
+        self.num_queries = num_queries
+        self.transformer = transformer
+        hidden_dim = transformer.d_model # =args.hidden_dim 256
+        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
+        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
+        self.query_embed = nn.Embedding(num_queries, hidden_dim)
+        # self.input_proj = nn.ModuleList([
+        #         nn.Sequential(
+        #             nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1),
+        #             nn.GroupNorm(32, hidden_dim),
+        #         )])
+        self.input_proj = nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1)
+        # self.input_proj = nn.Conv2d(256, hidden_dim, kernel_size=1)
+        self.backbone = backbone
+        self.aux_loss = aux_loss
+
+    def forward(self, samples: NestedTensor):
+        """ The forward expects a NestedTensor, which consists of:
+               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
+               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
+
+            It returns a dict with the following elements:
+               - "pred_logits": the classification logits (including no-object) for all queries.
+                                Shape= [batch_size x num_queries x (num_classes + 1)]
+               - "pred_boxes": The normalized boxes coordinates for all queries, represented as
+                               (center_x, center_y, height, width). These values are normalized in [0, 1],
+                               relative to the size of each individual image (disregarding possible padding).
+                               See PostProcess for information on how to retrieve the unnormalized bounding box.
+               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
+                                dictionnaries containing the two above keys for each decoder layer.
+        """
+        if isinstance(samples, (list, torch.Tensor)):
+            samples = nested_tensor_from_tensor_list(samples)
+        features, pos = self.backbone(samples)
+        # print('features:', features[0].tensors.shape)
+        # print('pos:', pos[0].shape)
+        src, mask = features[-1].decompose()
+        # print('src shape:', src.shape, mask.shape)
+        assert mask is not None
+
+        hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]
+
+        outputs_class = self.class_embed(hs)
+        outputs_coord = self.bbox_embed(hs).sigmoid()
+        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
+        if self.aux_loss:
+            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
+        return out
+
+    @torch.jit.unused
+    def _set_aux_loss(self, outputs_class, outputs_coord):
+        # this is a workaround to make torchscript happy, as torchscript
+        # doesn't support dictionary with non-homogeneous values, such
+        # as a dict having both a Tensor and a list.
+        return [{'pred_logits': a, 'pred_boxes': b}
+                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
+
+
+class SetCriterion(nn.Module):
+    """ This class computes the loss for DETR.
+    The process happens in two steps:
+        1) we compute hungarian assignment between ground truth boxes and the outputs of the model
+        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
+    """
+    def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses):
+        """ Create the criterion.
+        Parameters:
+            num_classes: number of object categories, omitting the special no-object category
+            matcher: module able to compute a matching between targets and proposals
+            weight_dict: dict containing as key the names of the losses and as values their relative weight.
+            eos_coef: relative classification weight applied to the no-object category
+            losses: list of all the losses to be applied. See get_loss for list of available losses.
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        self.matcher = matcher
+        self.weight_dict = weight_dict
+        self.eos_coef = eos_coef
+        self.losses = losses
+        empty_weight = torch.ones(self.num_classes + 1)
+        empty_weight[-1] = self.eos_coef
+        self.register_buffer('empty_weight', empty_weight)
+
+    def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
+        """Classification loss (NLL)
+        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
+        """
+        assert 'pred_logits' in outputs
+        src_logits = outputs['pred_logits']
+
+        idx = self._get_src_permutation_idx(indices)
+        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
+        target_classes = torch.full(src_logits.shape[:2], self.num_classes,
+                                    dtype=torch.int64, device=src_logits.device)
+        target_classes[idx] = target_classes_o
+
+        target_classes_onehot = torch.zeros([src_logits.shape[0], src_logits.shape[1], src_logits.shape[2]+1],
+                                            dtype=src_logits.dtype, layout=src_logits.layout, device=src_logits.device)
+        target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)
+        target_classes_onehot = target_classes_onehot[:,:,:-1]
+        loss_ce = sigmoid_focal_loss(src_logits, target_classes_onehot, num_boxes, alpha=0.25, gamma=2) * src_logits.shape[1]
+
+        # loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
+        losses = {'loss_ce': loss_ce}
+
+        if log:
+            # TODO this should probably be a separate loss, not hacked in this one here
+            losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]
+        return losses
+
+    @torch.no_grad()
+    def loss_cardinality(self, outputs, targets, indices, num_boxes):
+        """ Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes
+        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients
+        """
+        pred_logits = outputs['pred_logits']
+        device = pred_logits.device
+        tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device)
+        # Count the number of predictions that are NOT "no-object" (which is the last class)
+        card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
+        card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())
+        losses = {'cardinality_error': card_err}
+        return losses
+
+    def loss_boxes(self, outputs, targets, indices, num_boxes):
+        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
+           targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
+           The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size.
+        """
+        assert 'pred_boxes' in outputs
+        idx = self._get_src_permutation_idx(indices)
+        src_boxes = outputs['pred_boxes'][idx]
+        target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)
+
+        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
+
+        losses = {}
+        losses['loss_bbox'] = loss_bbox.sum() / num_boxes
+
+        loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(
+            box_ops.box_cxcywh_to_xyxy(src_boxes),
+            box_ops.box_cxcywh_to_xyxy(target_boxes)))
+        losses['loss_giou'] = loss_giou.sum() / num_boxes
+        return losses
+
+    def loss_masks(self, outputs, targets, indices, num_boxes):
+        """Compute the losses related to the masks: the focal loss and the dice loss.
+           targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
+        """
+        assert "pred_masks" in outputs
+
+        src_idx = self._get_src_permutation_idx(indices)
+        tgt_idx = self._get_tgt_permutation_idx(indices)
+        src_masks = outputs["pred_masks"]
+        src_masks = src_masks[src_idx]
+        masks = [t["masks"] for t in targets]
+        # TODO use valid to mask invalid areas due to padding in loss
+        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
+        target_masks = target_masks.to(src_masks)
+        target_masks = target_masks[tgt_idx]
+
+        # upsample predictions to the target size
+        src_masks = interpolate(src_masks[:, None], size=target_masks.shape[-2:],
+                                mode="bilinear", align_corners=False)
+        src_masks = src_masks[:, 0].flatten(1)
+
+        target_masks = target_masks.flatten(1)
+        target_masks = target_masks.view(src_masks.shape)
+        losses = {
+            "loss_mask": sigmoid_focal_loss(src_masks, target_masks, num_boxes),
+            "loss_dice": dice_loss(src_masks, target_masks, num_boxes),
+        }
+        return losses
+
+    def _get_src_permutation_idx(self, indices):
+        # permute predictions following indices
+        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+        src_idx = torch.cat([src for (src, _) in indices])
+        return batch_idx, src_idx
+
+    def _get_tgt_permutation_idx(self, indices):
+        # permute targets following indices
+        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
+        return batch_idx, tgt_idx
+
+    def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs):
+        loss_map = {
+            'labels': self.loss_labels,
+            'cardinality': self.loss_cardinality,
+            'boxes': self.loss_boxes,
+            'masks': self.loss_masks
+        }
+        assert loss in loss_map, f'do you really want to compute {loss} loss?'
+        return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs)
+
+    def forward(self, outputs, targets):
+        """ This performs the loss computation.
+        Parameters:
+             outputs: dict of tensors, see the output specification of the model for the format
+             targets: list of dicts, such that len(targets) == batch_size.
+                      The expected keys in each dict depends on the losses applied, see each loss' doc
+        """
+        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}
+
+        # Retrieve the matching between the outputs of the last layer and the targets
+        indices = self.matcher(outputs_without_aux, targets)
+
+        # Compute the average number of target boxes accross all nodes, for normalization purposes
+        num_boxes = sum(len(t["labels"]) for t in targets)
+        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
+        if is_dist_avail_and_initialized():
+            torch.distributed.all_reduce(num_boxes)
+        num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()
+
+        # Compute all the requested losses
+        losses = {}
+        for loss in self.losses:
+            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))
+
+        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+        if 'aux_outputs' in outputs:
+            for i, aux_outputs in enumerate(outputs['aux_outputs']):
+                indices = self.matcher(aux_outputs, targets)
+                for loss in self.losses:
+                    if loss == 'masks':
+                        # Intermediate masks losses are too costly to compute, we ignore them.
+                        continue
+                    kwargs = {}
+                    if loss == 'labels':
+                        # Logging is enabled only for the last layer
+                        kwargs = {'log': False}
+                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)
+                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
+                    losses.update(l_dict)
+
+        return losses
+
+
+class PostProcess(nn.Module):
+    """ This module converts the model's output into the format expected by the coco api"""
+    @torch.no_grad()
+    def forward(self, outputs, target_sizes):
+        """ Perform the computation
+        Parameters:
+            outputs: raw outputs of the model
+            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
+                          For evaluation, this must be the original image size (before any data augmentation)
+                          For visualization, this should be the image size after data augment, but before padding
+        """
+        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']
+
+        assert len(out_logits) == len(target_sizes)
+        assert target_sizes.shape[1] == 2
+
+        prob = F.softmax(out_logits, -1)
+        scores, labels = prob[..., :-1].max(-1)
+
+        # convert to [x0, y0, x1, y1] format
+        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+        # and from relative [0, 1] to absolute [0, height] coordinates
+        img_h, img_w = target_sizes.unbind(1)
+        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
+        boxes = boxes * scale_fct[:, None, :]
+
+        # print('Originial output:')
+        # print('Labels:', labels, 'bbox:', boxes)
+
+        results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
+
+        return results
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+def build(args):
+    # the `num_classes` naming here is somewhat misleading.
+    # it indeed corresponds to `max_obj_id + 1`, where max_obj_id
+    # is the maximum id for a class in your dataset. For example,
+    # COCO has a max_obj_id of 90, so we pass `num_classes` to be 91.
+    # As another example, for a dataset that has a single class with id 1,
+    # you should pass `num_classes` to be 2 (max_obj_id + 1).
+    # For more details on this, check the following discussion
+    # https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223
+    num_classes = 8 if args.dataset_file == 'endovis17' else 91
+    if args.dataset_file == "coco_panoptic":
+        # for panoptic, we just add a num_classes that is large enough to hold
+        # max_obj_id + 1, but the exact value doesn't really matter
+        num_classes = 250
+    device = torch.device(args.device)
+
+    backbone = build_backbone(args)
+
+    transformer = build_transformer(args)
+    transformer.eval()
+    for param in transformer.parameters():
+        param.requires_grad = False
+
+    if args.model:
+        # pretrained_model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
+        # # save model weights
+        # torch.save(pretrained_model.state_dict(), 'detr_weights.pth')
+
+        # initialize model weights
+        model = DETR(
+            backbone,
+            transformer,
+            num_classes=num_classes,
+            num_queries=args.num_queries,
+            aux_loss=args.aux_loss,
+        )
+
+        weights = torch.load('detr_weights.pth')
+        # checkpoint = torch.load('checkpoint0040.pth', map_location='cpu')
+        # weights = checkpoint['model']
+        # delete specific layers in weights
+        exclude_keys = ['class_embed.weight', 'class_embed.bias', 'input_proj.weight']
+        for key in exclude_keys:
+            del weights[key]
+
+        model.load_state_dict(weights, strict=False)
+
+    else:
+        model = DETR(
+            backbone,
+            transformer,
+            num_classes=num_classes,
+            num_queries=args.num_queries,
+            aux_loss=args.aux_loss,
+        )
+   
+
+    if args.masks:
+        model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
+    matcher = build_matcher(args)
+    weight_dict = {'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef}
+    weight_dict['loss_giou'] = args.giou_loss_coef
+    if args.masks:
+        weight_dict["loss_mask"] = args.mask_loss_coef
+        weight_dict["loss_dice"] = args.dice_loss_coef
+    # TODO this is a hack
+    if args.aux_loss:
+        aux_weight_dict = {}
+        for i in range(args.dec_layers - 1):
+            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
+        weight_dict.update(aux_weight_dict)
+
+    losses = ['labels', 'boxes', 'cardinality']
+    if args.masks:
+        losses += ["masks"]
+    criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
+                             eos_coef=args.eos_coef, losses=losses)
+    criterion.to(device)
+    postprocessors = {'bbox': PostProcess()}
+    if args.masks:
+        postprocessors['segm'] = PostProcessSegm()
+        if args.dataset_file == "coco_panoptic":
+            is_thing_map = {i: i <= 90 for i in range(201)}
+            postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)
+
+    return model, criterion, postprocessors

models/detr_seg.py

@@ -0,0 +1,616 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+DETR model and criterion classes.
+"""
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from util import box_ops
+from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
+                       accuracy, get_world_size, interpolate,
+                       is_dist_avail_and_initialized)
+
+from .backbone import build_backbone
+from .matcher import build_matcher
+from .segmentation import (DETRsegm, PostProcessPanoptic, PostProcessSegm,
+                           dice_loss, sigmoid_focal_loss, dice_coefficient, focal_loss_masks)
+from .transformer import build_transformer
+from models.sam.segment_anything.modeling import ImageEncoderViT, PromptEncoder, MaskDecoder, TwoWayTransformer
+from segment_anything import sam_model_registry, SamPredictor
+
+from torch.nn import functional as FN
+from torch import nn
+
+import torch.distributions as dist
+import numpy as np
+import monai
+
+def add_noise_to_bbox(bbox, max_noise=20):
+    '''
+    args: bbox (N, 4)
+    '''
+    # Calculate standard deviation as 10% of the box sidelength
+    box_width = bbox[:, 2] - bbox[:, 0]
+    box_height = bbox[:, 3] - bbox[:, 1]
+    std_dev = 0.1 * torch.max(box_width, box_height)
+    
+    # Create normal distribution for generating noise
+    noise_dist = dist.Normal(0, std_dev) # (num_boxes, )
+    num_boxes = bbox.shape[0]
+
+    # Generate random noise for each coordinate
+    x1_noise = noise_dist.sample()
+    y1_noise = noise_dist.sample()
+    x2_noise = noise_dist.sample()
+    y2_noise = noise_dist.sample()
+    
+    # Clip noise to a maximum of 20 pixels
+    x1_noise = torch.clamp(x1_noise, -max_noise, max_noise)
+    y1_noise = torch.clamp(y1_noise, -max_noise, max_noise)
+    x2_noise = torch.clamp(x2_noise, -max_noise, max_noise)
+    y2_noise = torch.clamp(y2_noise, -max_noise, max_noise)
+    noise = torch.stack([x1_noise, y1_noise, x2_noise, y2_noise], dim=1)
+
+    # Add noise to the original coordinates
+    noisy_bbox = bbox + noise
+    
+    return noisy_bbox
+
+def box_cxcywh_to_xyxy(x):
+    x_c, y_c, w, h = x.unbind(1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
+         (x_c + 0.5 * w), (y_c + 0.5 * h)]
+    return torch.stack(b, dim=1)
+
+def rescale_bboxes(out_bbox, size):
+    img_w, img_h = size
+    b = box_cxcywh_to_xyxy(out_bbox.cpu())
+    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
+    return b
+
+def postprocess_masks(masks, input_size, original_size,) -> torch.Tensor:
+        masks = FN.interpolate(
+            masks,
+            input_size,
+            mode="bilinear",
+            align_corners=False,
+        )
+        masks = masks[..., :input_size[0], :input_size[1]]
+        masks = FN.interpolate(masks, original_size, mode="bilinear", align_corners=False)
+        return masks
+
+class SAMModel(nn.Module):
+    def __init__(self, device, model_type='vit_b', ckpt_path='sam_vit_b_01ec64.pth'):
+        super().__init__()
+        
+        sam = sam_model_registry[model_type](checkpoint=ckpt_path).to(device)
+        self.predictor = SamPredictor(sam)
+
+        self.sam_image_encoder = sam.image_encoder
+        self.sam_prompt_encoder = sam.prompt_encoder
+        self.sam_mask_decoder = sam.mask_decoder
+        self.device = device
+        self.upsample_layer = nn.ConvTranspose2d(
+                                                in_channels=256,      # Number of input channels (should match your input image)
+                                                out_channels=256,     # Number of output channels (same as input channels for no change)
+                                                kernel_size=4,      # Kernel size for the convolution
+                                                stride=2,           # Upsampling factor (doubles the spatial dimensions)
+                                                padding=1,          # Padding to maintain spatial dimensions
+                                                output_padding=0,   # Additional padding to adjust the output size
+                                            )
+        
+    def forward(self, batched_img, boxes, image_embeddings=None, sizes=(1024, 1024), add_noise=True):
+        if image_embeddings is None:
+            image_embeddings = self.sam_image_encoder(batched_img)[0].unsqueeze(0)
+        else:
+            image_embeddings = self.upsample_layer(image_embeddings) # (3, 32, 32)-> (3, 64, 64)
+            # print(image_embeddings.shape)
+            
+        gt_boxes = rescale_bboxes(boxes, sizes) # xyxy-format with shape (N, 4)
+        if add_noise:
+            noisy_boxes = add_noise_to_bbox(gt_boxes)
+        else:
+            noisy_boxes = gt_boxes
+        
+        transformed_boxes = self.predictor.transform.apply_boxes_torch(noisy_boxes, sizes).to(self.device)
+        if gt_boxes.shape[0] == 0:
+            transformed_boxes = None
+        sparse_embeddings, dense_embeddings = self.sam_prompt_encoder(points=None,
+                                                         boxes=transformed_boxes,
+                                                         masks=None)
+        # print(gt_boxes.shape, transformed_boxes)
+        # print(image_embeddings.shape, self.sam_prompt_encoder.get_dense_pe().shape)
+        # print(sparse_embeddings.shape, dense_embeddings.shape)
+        low_res_masks, iou_predictions = self.sam_mask_decoder(
+                                                                image_embeddings=image_embeddings,
+                                                                image_pe=self.sam_prompt_encoder.get_dense_pe(),
+                                                                sparse_prompt_embeddings=sparse_embeddings,
+                                                                dense_prompt_embeddings=dense_embeddings,
+                                                                multimask_output=False,
+                                                                # hq_token_only=False,
+                                                                # interm_embeddings=False
+                                                            )
+        pred_masks = postprocess_masks(low_res_masks, input_size=sizes, original_size=sizes)
+        return low_res_masks.reshape(-1, 256, 256), pred_masks.reshape(-1, sizes[0], sizes[1]), iou_predictions
+
+class CustomSAMModel(nn.Module):
+    def __init__(self, device, img_size=224, prompt_embed_dim=256, image_embedding_size=(14, 14), input_image_size=(224, 224)):
+        super().__init__()
+        self.device = device
+        self.sam_image_encoder = ImageEncoderViT(img_size=img_size)
+        self.sam_prompt_encoder = PromptEncoder(embed_dim=prompt_embed_dim, 
+                                                image_embedding_size=image_embedding_size, 
+                                                input_image_size=input_image_size,
+                                                mask_in_chans=16
+                                                )
+        self.sam_mask_decoder = MaskDecoder(transformer_dim=256,
+                                            transformer=TwoWayTransformer(depth=2,
+                                                                    embedding_dim=prompt_embed_dim,
+                                                                    mlp_dim=2048,
+                                                                    num_heads=8))
+        
+    def forward(self, batched_img, image_embeddings, boxes, sizes, predictor, add_noise=True):
+        image_embeddings = self.sam_image_encoder(batched_img)
+        
+        gt_boxes = rescale_bboxes(boxes, sizes) # xyxy-format
+        if add_noise:
+            noisy_boxes = add_noise_to_bbox(gt_boxes)
+        else:
+            noisy_boxes = gt_boxes
+        
+        transformed_boxes = predictor.transform.apply_boxes_torch(noisy_boxes, (224, 224)).to(self.device)
+        if gt_boxes.shape[0] == 0:
+            transformed_boxes = None
+        sparse_embeddings, dense_embeddings = self.sam_prompt_encoder(points=None,
+                                                         boxes=transformed_boxes,
+                                                         masks=None)
+        # print(gt_boxes.shape, transformed_boxes)
+        # print(image_embeddings.shape, self.sam_prompt_encoder.get_dense_pe().shape)
+        # print(sparse_embeddings.shape, dense_embeddings.shape)
+        low_res_masks, iou_predictions = self.sam_mask_decoder(
+                                                                image_embeddings=image_embeddings,
+                                                                image_pe=self.sam_prompt_encoder.get_dense_pe(),
+                                                                sparse_prompt_embeddings=sparse_embeddings,
+                                                                dense_prompt_embeddings=dense_embeddings,
+                                                                multimask_output=False,
+                                                            )
+        pred_masks = postprocess_masks(low_res_masks, input_size=(224, 224), original_size=(224, 224))
+        return low_res_masks.reshape(-1, 56, 56), pred_masks.reshape(-1, 224, 224), iou_predictions
+    
+class DETR(nn.Module):
+    """ This is the DETR module that performs object detection """
+    def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):
+        """ Initializes the model.
+        Parameters:
+            backbone: torch module of the backbone to be used. See backbone.py
+            transformer: torch module of the transformer architecture. See transformer.py
+            num_classes: number of object classes
+            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
+                         DETR can detect in a single image. For COCO, we recommend 100 queries.
+            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
+        """
+        super().__init__()
+        self.num_queries = num_queries
+        self.transformer = transformer
+        hidden_dim = transformer.d_model # =args.hidden_dim 256
+        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
+        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
+        self.query_embed = nn.Embedding(num_queries, hidden_dim)
+        # self.input_proj = nn.ModuleList([
+        #         nn.Sequential(
+        #             nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1),
+        #             nn.GroupNorm(32, hidden_dim),
+        #         )])
+        self.input_proj = nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1)
+        # self.input_proj = nn.Conv2d(256, hidden_dim, kernel_size=1)
+        self.backbone = backbone
+        self.aux_loss = aux_loss
+
+    def forward(self, samples: NestedTensor):
+        """ The forward expects a NestedTensor, which consists of:
+               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
+               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
+
+            It returns a dict with the following elements:
+               - "pred_logits": the classification logits (including no-object) for all queries.
+                                Shape= [batch_size x num_queries x (num_classes + 1)]
+               - "pred_boxes": The normalized boxes coordinates for all queries, represented as
+                               (center_x, center_y, height, width). These values are normalized in [0, 1],
+                               relative to the size of each individual image (disregarding possible padding).
+                               See PostProcess for information on how to retrieve the unnormalized bounding box.
+               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
+                                dictionnaries containing the two above keys for each decoder layer.
+        """
+        if isinstance(samples, (list, torch.Tensor)):
+            samples = nested_tensor_from_tensor_list(samples)
+        features, pos = self.backbone(samples)
+        # print('features:', features[0].tensors.shape)
+        # print('pos:', pos[0].shape)
+        src, mask = features[-1].decompose()
+        # print('src shape:', src.shape, mask.shape)
+        assert mask is not None
+
+        hs, image_embeddings = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])
+
+        outputs_class = self.class_embed(hs)
+        outputs_coord = self.bbox_embed(hs).sigmoid()
+        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
+        if self.aux_loss:
+            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
+        return out, image_embeddings
+
+    @torch.jit.unused
+    def _set_aux_loss(self, outputs_class, outputs_coord):
+        # this is a workaround to make torchscript happy, as torchscript
+        # doesn't support dictionary with non-homogeneous values, such
+        # as a dict having both a Tensor and a list.
+        return [{'pred_logits': a, 'pred_boxes': b}
+                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
+
+def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
+    """
+    Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
+    Args:
+        inputs: A float tensor of arbitrary shape.
+                The predictions for each example.
+        targets: A float tensor with the same shape as inputs. Stores the binary
+                 classification label for each element in inputs
+                (0 for the negative class and 1 for the positive class).
+        alpha: (optional) Weighting factor in range (0,1) to balance
+                positive vs negative examples. Default = -1 (no weighting).
+        gamma: Exponent of the modulating factor (1 - p_t) to
+               balance easy vs hard examples.
+    Returns:
+        Loss tensor
+    """
+    prob = inputs.sigmoid()
+    ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
+    p_t = prob * targets + (1 - prob) * (1 - targets)
+    loss = ce_loss * ((1 - p_t) ** gamma)
+
+    if alpha >= 0:
+        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
+        loss = alpha_t * loss
+
+    return loss.mean(1).sum() / num_boxes
+
+class SetCriterion(nn.Module):
+    """ This class computes the loss for DETR.
+    The process happens in two steps:
+        1) we compute hungarian assignment between ground truth boxes and the outputs of the model
+        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
+    """
+    def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses, use_matcher=True):
+        """ Create the criterion.
+        Parameters:
+            num_classes: number of object categories, omitting the special no-object category
+            matcher: module able to compute a matching between targets and proposals
+            weight_dict: dict containing as key the names of the losses and as values their relative weight.
+            eos_coef: relative classification weight applied to the no-object category
+            losses: list of all the losses to be applied. See get_loss for list of available losses.
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        self.matcher = matcher
+        self.weight_dict = weight_dict
+        self.eos_coef = eos_coef
+        self.losses = losses
+        empty_weight = torch.ones(self.num_classes + 1)
+        empty_weight[-1] = self.eos_coef
+        self.register_buffer('empty_weight', empty_weight)
+        self.use_matcher = use_matcher
+
+    def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
+        """Classification loss (NLL)
+        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
+        """
+        assert 'pred_logits' in outputs
+        src_logits = outputs['pred_logits']
+
+        idx = self._get_src_permutation_idx(indices)
+        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
+        target_classes = torch.full(src_logits.shape[:2], self.num_classes,
+                                    dtype=torch.int64, device=src_logits.device)
+        target_classes[idx] = target_classes_o
+
+        target_classes_onehot = torch.zeros([src_logits.shape[0], src_logits.shape[1], src_logits.shape[2]+1],
+                                            dtype=src_logits.dtype, layout=src_logits.layout, device=src_logits.device)
+        target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)
+        target_classes_onehot = target_classes_onehot[:,:,:-1]
+        loss_ce = sigmoid_focal_loss(src_logits, target_classes_onehot, num_boxes, alpha=0.25, gamma=2) * src_logits.shape[1]
+
+        # loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
+        losses = {'loss_ce': loss_ce}
+
+        if log:
+            # TODO this should probably be a separate loss, not hacked in this one here
+            losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]
+        return losses
+
+    @torch.no_grad()
+    def loss_cardinality(self, outputs, targets, indices, num_boxes):
+        """ Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes
+        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients
+        """
+        pred_logits = outputs['pred_logits']
+        device = pred_logits.device
+        tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device)
+        # Count the number of predictions that are NOT "no-object" (which is the last class)
+        card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
+        card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())
+        losses = {'cardinality_error': card_err}
+        return losses
+
+    def loss_boxes(self, outputs, targets, indices, num_boxes):
+        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
+           targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
+           The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size.
+        """
+        assert 'pred_boxes' in outputs
+        idx = self._get_src_permutation_idx(indices)
+        src_boxes = outputs['pred_boxes'][idx] # (N, 4)
+        target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)
+
+        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
+
+        losses = {}
+        losses['loss_bbox'] = loss_bbox.sum() / num_boxes
+
+        loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(
+            box_ops.box_cxcywh_to_xyxy(src_boxes),
+            box_ops.box_cxcywh_to_xyxy(target_boxes)))
+        losses['loss_giou'] = loss_giou.sum() / num_boxes
+        return losses
+
+    def loss_masks(self, outputs, targets, indices, num_boxes):
+        """Compute the losses related to the masks: the focal loss and the dice loss.
+           targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
+        """
+        assert "pred_masks" in outputs
+
+        # src_idx = self._get_src_permutation_idx(indices)
+        # tgt_idx = self._get_tgt_permutation_idx(indices)
+  
+        src_masks = outputs["pred_masks"].unsqueeze(0) # (bs, N, H, W)
+        # src_masks = src_masks[src_idx]
+        # masks = [t["masks"] for t in targets]
+        # # TODO use valid to mask invalid areas due to padding in loss
+        # target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
+        # target_masks = target_masks.to(src_masks)
+        target_masks = targets[0]['masks'].unsqueeze(0)
+        
+        # target_masks = target_masks[tgt_idx]
+
+        # # upsample predictions to the target size
+        # src_masks = interpolate(src_masks[:, None], size=target_masks.shape[-2:],
+        #                         mode="bilinear", align_corners=False)
+        # src_masks = src_masks[:, 0].flatten(1)
+
+        # target_masks = target_masks.flatten(1)
+        # target_masks = target_masks.view(src_masks.shape)
+
+        # ---------sam--------
+        # src_masks = outputs['pred_masks']
+        # target_masks = targets[0]['masks']
+        dice_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')
+        losses = {
+            "loss_mask": focal_loss_masks(src_masks.cpu(), target_masks.cpu(), num_boxes),
+            # "loss_dice": dice_loss(src_masks, target_masks.cpu(), num_boxes),
+            "loss_dice": dice_loss(src_masks.cpu(), target_masks.cpu()),
+        }
+        return losses
+
+    def _get_src_permutation_idx(self, indices):
+        # permute predictions following indices
+        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+        src_idx = torch.cat([src for (src, _) in indices])
+        return batch_idx, src_idx
+
+    def _get_tgt_permutation_idx(self, indices):
+        # permute targets following indices
+        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
+        return batch_idx, tgt_idx
+
+    def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs):
+        loss_map = {
+            'labels': self.loss_labels,
+            'cardinality': self.loss_cardinality,
+            'boxes': self.loss_boxes,
+            'masks': self.loss_masks
+        }
+        assert loss in loss_map, f'do you really want to compute {loss} loss?'
+        return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs)
+
+    def forward(self, outputs, targets):
+        """ This performs the loss computation.
+        Parameters:
+             outputs: dict of tensors, see the output specification of the model for the format
+             targets: list of dicts, such that len(targets) == batch_size.
+                      The expected keys in each dict depends on the losses applied, see each loss' doc
+        """
+        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}
+
+        # Retrieve the matching between the outputs of the last layer and the targets
+        if self.use_matcher:
+            indices = self.matcher(outputs_without_aux, targets)
+        else:
+            indices = None
+
+        # Compute the average number of target boxes accross all nodes, for normalization purposes
+        num_boxes = sum(len(t["labels"]) for t in targets)
+        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
+        if is_dist_avail_and_initialized():
+            torch.distributed.all_reduce(num_boxes)
+        num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()
+
+        # Compute all the requested losses
+        losses = {}
+        for loss in self.losses:
+            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))
+
+        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+        if 'aux_outputs' in outputs:
+            for i, aux_outputs in enumerate(outputs['aux_outputs']):
+                indices = self.matcher(aux_outputs, targets)
+                for loss in self.losses:
+                    if loss == 'masks':
+                        # Intermediate masks losses are too costly to compute, we ignore them.
+                        continue
+                    kwargs = {}
+                    if loss == 'labels':
+                        # Logging is enabled only for the last layer
+                        kwargs = {'log': False}
+                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)
+                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
+                    losses.update(l_dict)
+
+        return losses
+
+
+class PostProcess(nn.Module):
+    """ This module converts the model's output into the format expected by the coco api"""
+    @torch.no_grad()
+    def forward(self, outputs, target_sizes):
+        """ Perform the computation
+        Parameters:
+            outputs: raw outputs of the model
+            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
+                          For evaluation, this must be the original image size (before any data augmentation)
+                          For visualization, this should be the image size after data augment, but before padding
+        """
+        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']
+
+        assert len(out_logits) == len(target_sizes)
+        assert target_sizes.shape[1] == 2
+
+        prob = F.softmax(out_logits, -1)
+        scores, labels = prob[..., :-1].max(-1)
+
+        # convert to [x0, y0, x1, y1] format
+        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+        # and from relative [0, 1] to absolute [0, height] coordinates
+        img_h, img_w = target_sizes.unbind(1)
+        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
+        boxes = boxes * scale_fct[:, None, :]
+
+        # print('Originial output:')
+        # print('Labels:', labels, 'bbox:', boxes)
+
+        results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
+
+        return results
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+def build(args):
+    # the `num_classes` naming here is somewhat misleading.
+    # it indeed corresponds to `max_obj_id + 1`, where max_obj_id
+    # is the maximum id for a class in your dataset. For example,
+    # COCO has a max_obj_id of 90, so we pass `num_classes` to be 91.
+    # As another example, for a dataset that has a single class with id 1,
+    # you should pass `num_classes` to be 2 (max_obj_id + 1).
+    # For more details on this, check the following discussion
+    # https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223
+    num_classes = 8 if args.dataset_file == 'endovis17' else 91
+    if args.dataset_file == 'endovis18':
+        num_classes = 9
+    if args.dataset_file == "coco_panoptic":
+        # for panoptic, we just add a num_classes that is large enough to hold
+        # max_obj_id + 1, but the exact value doesn't really matter
+        num_classes = 250
+    device = torch.device(args.device)
+
+    backbone = build_backbone(args)
+
+    transformer = build_transformer(args)
+
+    if args.model:
+        # pretrained_model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
+        # # save model weights
+        # torch.save(pretrained_model.state_dict(), 'detr_weights.pth')
+
+        # initialize model weights
+        model = DETR(
+            backbone,
+            transformer,
+            num_classes=num_classes,
+            num_queries=args.num_queries,
+            aux_loss=args.aux_loss,
+        )
+
+        # weights = torch.load('detr_weights.pth')
+        # # checkpoint = torch.load('outputs/ckpt_best.pth')
+        # # weights = checkpoint['model']
+        # # delete specific layers in weights
+        # exclude_keys = ['class_embed.weight', 'class_embed.bias', 'input_proj.weight']
+        # # for key in exclude_keys:
+        # #     del weights[key]
+
+        # # model.load_state_dict(weights, strict=False)
+
+    else:
+        model = DETR(
+            backbone,
+            transformer,
+            num_classes=num_classes,
+            num_queries=args.num_queries,
+            aux_loss=args.aux_loss,
+        )
+   
+
+    if args.masks:
+        model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
+    matcher = build_matcher(args)
+    weight_dict = {'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef}
+    weight_dict['loss_giou'] = args.giou_loss_coef
+    if args.masks:
+        weight_dict["loss_mask"] = args.mask_loss_coef
+        weight_dict["loss_dice"] = args.dice_loss_coef
+    # TODO this is a hack
+    if args.aux_loss:
+        aux_weight_dict = {}
+        for i in range(args.dec_layers - 1):
+            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
+        weight_dict.update(aux_weight_dict)
+
+    losses = ['labels', 'boxes', 'cardinality']
+    if args.masks:
+        losses += ["masks"]
+    criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
+                             eos_coef=args.eos_coef, losses=losses)
+    
+    seg_losses = ['masks']
+    seg_losses += losses
+    seg_weight_dict = {'loss_mask': args.mask_loss_coef, 'loss_dice': args.dice_loss_coef, 'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef,
+                       'loss_giou': args.giou_loss_coef}
+    #TODO if use_matcher == True: use Hungarian matching for predicted boxes and GT boxes 
+    seg_criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=seg_weight_dict,
+                             eos_coef=args.eos_coef, losses=seg_losses, use_matcher=True)
+    seg_criterion.to(device)
+    
+    criterion.to(device)
+    postprocessors = {'bbox': PostProcess()}
+    if args.masks:
+        postprocessors['segm'] = PostProcessSegm()
+        if args.dataset_file == "coco_panoptic":
+            is_thing_map = {i: i <= 90 for i in range(201)}
+            postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)
+
+    return model, criterion, seg_criterion, postprocessors
+

models/matcher.py

@@ -0,0 +1,86 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Modules to compute the matching cost and solve the corresponding LSAP.
+"""
+import torch
+from scipy.optimize import linear_sum_assignment
+from torch import nn
+
+from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou
+
+
+class HungarianMatcher(nn.Module):
+    """This class computes an assignment between the targets and the predictions of the network
+
+    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
+    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
+    while the others are un-matched (and thus treated as non-objects).
+    """
+
+    def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1):
+        """Creates the matcher
+
+        Params:
+            cost_class: This is the relative weight of the classification error in the matching cost
+            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost
+            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost
+        """
+        super().__init__()
+        self.cost_class = cost_class
+        self.cost_bbox = cost_bbox
+        self.cost_giou = cost_giou
+        assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, "all costs cant be 0"
+
+    @torch.no_grad()
+    def forward(self, outputs, targets):
+        """ Performs the matching
+
+        Params:
+            outputs: This is a dict that contains at least these entries:
+                 "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
+                 "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates
+
+            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
+                 "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
+                           objects in the target) containing the class labels
+                 "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates
+
+        Returns:
+            A list of size batch_size, containing tuples of (index_i, index_j) where:
+                - index_i is the indices of the selected predictions (in order)
+                - index_j is the indices of the corresponding selected targets (in order)
+            For each batch element, it holds:
+                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
+        """
+        bs, num_queries = outputs["pred_logits"].shape[:2]
+
+        # We flatten to compute the cost matrices in a batch
+        out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
+        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]
+
+        # Also concat the target labels and boxes
+        tgt_ids = torch.cat([v["labels"] for v in targets])
+        tgt_bbox = torch.cat([v["boxes"] for v in targets])
+
+        # Compute the classification cost. Contrary to the loss, we don't use the NLL,
+        # but approximate it in 1 - proba[target class].
+        # The 1 is a constant that doesn't change the matching, it can be ommitted.
+        cost_class = -out_prob[:, tgt_ids]
+
+        # Compute the L1 cost between boxes
+        cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)
+
+        # Compute the giou cost betwen boxes
+        cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))
+
+        # Final cost matrix
+        C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou
+        C = C.view(bs, num_queries, -1).cpu()
+
+        sizes = [len(v["boxes"]) for v in targets]
+        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
+        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
+
+
+def build_matcher(args):
+    return HungarianMatcher(cost_class=args.set_cost_class, cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou)

models/position_encoding.py

@@ -0,0 +1,89 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+
+from util.misc import NestedTensor
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        mask = tensor_list.mask
+        assert mask is not None
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        return pos
+
+
+class PositionEmbeddingLearned(nn.Module):
+    """
+    Absolute pos embedding, learned.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.row_embed = nn.Embedding(50, num_pos_feats)
+        self.col_embed = nn.Embedding(50, num_pos_feats)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.row_embed.weight)
+        nn.init.uniform_(self.col_embed.weight)
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        h, w = x.shape[-2:]
+        i = torch.arange(w, device=x.device)
+        j = torch.arange(h, device=x.device)
+        x_emb = self.col_embed(i)
+        y_emb = self.row_embed(j)
+        pos = torch.cat([
+            x_emb.unsqueeze(0).repeat(h, 1, 1),
+            y_emb.unsqueeze(1).repeat(1, w, 1),
+        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
+        return pos
+
+
+def build_position_encoding(args):
+    N_steps = args.hidden_dim // 2 # 256 // 2 = 128
+    if args.position_embedding in ('v2', 'sine'):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
+    elif args.position_embedding in ('v3', 'learned'):
+        position_embedding = PositionEmbeddingLearned(N_steps)
+    else:
+        raise ValueError(f"not supported {args.position_embedding}")
+
+    return position_embedding

models/sam/CODE_OF_CONDUCT.md

@@ -0,0 +1,80 @@
+# Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to make participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+This Code of Conduct also applies outside the project spaces when there is a
+reasonable belief that an individual's behavior may have a negative impact on
+the project or its community.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at <opensource-conduct@fb.com>. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq

models/sam/CONTRIBUTING.md

@@ -0,0 +1,31 @@
+# Contributing to segment-anything
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+We actively welcome your pull requests.
+
+1. Fork the repo and create your branch from `main`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints, using the `linter.sh` script in the project's root directory. Linting requires `black==23.*`, `isort==5.12.0`, `flake8`, and `mypy`.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to segment-anything, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.

models/sam/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

models/sam/README.md

@@ -0,0 +1,171 @@
+# Segment Anything
+
+**[Meta AI Research, FAIR](https://ai.facebook.com/research/)**
+
+[Alexander Kirillov](https://alexander-kirillov.github.io/), [Eric Mintun](https://ericmintun.github.io/), [Nikhila Ravi](https://nikhilaravi.com/), [Hanzi Mao](https://hanzimao.me/), Chloe Rolland, Laura Gustafson, [Tete Xiao](https://tetexiao.com), [Spencer Whitehead](https://www.spencerwhitehead.com/), Alex Berg, Wan-Yen Lo, [Piotr Dollar](https://pdollar.github.io/), [Ross Girshick](https://www.rossgirshick.info/)
+
+[[`Paper`](https://ai.facebook.com/research/publications/segment-anything/)] [[`Project`](https://segment-anything.com/)] [[`Demo`](https://segment-anything.com/demo)] [[`Dataset`](https://segment-anything.com/dataset/index.html)] [[`Blog`](https://ai.facebook.com/blog/segment-anything-foundation-model-image-segmentation/)] [[`BibTeX`](#citing-segment-anything)]
+
+![SAM design](assets/model_diagram.png?raw=true)
+
+The **Segment Anything Model (SAM)** produces high quality object masks from input prompts such as points or boxes, and it can be used to generate masks for all objects in an image. It has been trained on a [dataset](https://segment-anything.com/dataset/index.html) of 11 million images and 1.1 billion masks, and has strong zero-shot performance on a variety of segmentation tasks.
+
+<p float="left">
+  <img src="assets/masks1.png?raw=true" width="37.25%" />
+  <img src="assets/masks2.jpg?raw=true" width="61.5%" /> 
+</p>
+
+## Installation
+
+The code requires `python>=3.8`, as well as `pytorch>=1.7` and `torchvision>=0.8`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.
+
+Install Segment Anything:
+
+```
+pip install git+https://github.com/facebookresearch/segment-anything.git
+```
+
+or clone the repository locally and install with
+
+```
+git clone git@github.com:facebookresearch/segment-anything.git
+cd segment-anything; pip install -e .
+```
+
+The following optional dependencies are necessary for mask post-processing, saving masks in COCO format, the example notebooks, and exporting the model in ONNX format. `jupyter` is also required to run the example notebooks.
+
+```
+pip install opencv-python pycocotools matplotlib onnxruntime onnx
+```
+
+## <a name="GettingStarted"></a>Getting Started
+
+First download a [model checkpoint](#model-checkpoints). Then the model can be used in just a few lines to get masks from a given prompt:
+
+```
+from segment_anything import SamPredictor, sam_model_registry
+sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
+predictor = SamPredictor(sam)
+predictor.set_image(<your_image>)
+masks, _, _ = predictor.predict(<input_prompts>)
+```
+
+or generate masks for an entire image:
+
+```
+from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
+sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
+mask_generator = SamAutomaticMaskGenerator(sam)
+masks = mask_generator.generate(<your_image>)
+```
+
+Additionally, masks can be generated for images from the command line:
+
+```
+python scripts/amg.py --checkpoint <path/to/checkpoint> --model-type <model_type> --input <image_or_folder> --output <path/to/output>
+```
+
+See the examples notebooks on [using SAM with prompts](/notebooks/predictor_example.ipynb) and [automatically generating masks](/notebooks/automatic_mask_generator_example.ipynb) for more details.
+
+<p float="left">
+  <img src="assets/notebook1.png?raw=true" width="49.1%" />
+  <img src="assets/notebook2.png?raw=true" width="48.9%" />
+</p>
+
+## ONNX Export
+
+SAM's lightweight mask decoder can be exported to ONNX format so that it can be run in any environment that supports ONNX runtime, such as in-browser as showcased in the [demo](https://segment-anything.com/demo). Export the model with
+
+```
+python scripts/export_onnx_model.py --checkpoint <path/to/checkpoint> --model-type <model_type> --output <path/to/output>
+```
+
+See the [example notebook](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb) for details on how to combine image preprocessing via SAM's backbone with mask prediction using the ONNX model. It is recommended to use the latest stable version of PyTorch for ONNX export.
+
+### Web demo
+
+The `demo/` folder has a simple one page React app which shows how to run mask prediction with the exported ONNX model in a web browser with multithreading. Please see [`demo/README.md`](https://github.com/facebookresearch/segment-anything/blob/main/demo/README.md) for more details.
+
+## <a name="Models"></a>Model Checkpoints
+
+Three model versions of the model are available with different backbone sizes. These models can be instantiated by running
+
+```
+from segment_anything import sam_model_registry
+sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
+```
+
+Click the links below to download the checkpoint for the corresponding model type.
+
+- **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)**
+- `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth)
+- `vit_b`: [ViT-B SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth)
+
+## Dataset
+
+See [here](https://ai.facebook.com/datasets/segment-anything/) for an overview of the datastet. The dataset can be downloaded [here](https://ai.facebook.com/datasets/segment-anything-downloads/). By downloading the datasets you agree that you have read and accepted the terms of the SA-1B Dataset Research License.
+
+We save masks per image as a json file. It can be loaded as a dictionary in python in the below format.
+
+```python
+{
+    "image"                 : image_info,
+    "annotations"           : [annotation],
+}
+
+image_info {
+    "image_id"              : int,              # Image id
+    "width"                 : int,              # Image width
+    "height"                : int,              # Image height
+    "file_name"             : str,              # Image filename
+}
+
+annotation {
+    "id"                    : int,              # Annotation id
+    "segmentation"          : dict,             # Mask saved in COCO RLE format.
+    "bbox"                  : [x, y, w, h],     # The box around the mask, in XYWH format
+    "area"                  : int,              # The area in pixels of the mask
+    "predicted_iou"         : float,            # The model's own prediction of the mask's quality
+    "stability_score"       : float,            # A measure of the mask's quality
+    "crop_box"              : [x, y, w, h],     # The crop of the image used to generate the mask, in XYWH format
+    "point_coords"          : [[x, y]],         # The point coordinates input to the model to generate the mask
+}
+```
+
+Image ids can be found in sa_images_ids.txt which can be downloaded using the above [link](https://ai.facebook.com/datasets/segment-anything-downloads/) as well.
+
+To decode a mask in COCO RLE format into binary:
+
+```
+from pycocotools import mask as mask_utils
+mask = mask_utils.decode(annotation["segmentation"])
+```
+
+See [here](https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/mask.py) for more instructions to manipulate masks stored in RLE format.
+
+## License
+
+The model is licensed under the [Apache 2.0 license](LICENSE).
+
+## Contributing
+
+See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md).
+
+## Contributors
+
+The Segment Anything project was made possible with the help of many contributors (alphabetical):
+
+Aaron Adcock, Vaibhav Aggarwal, Morteza Behrooz, Cheng-Yang Fu, Ashley Gabriel, Ahuva Goldstand, Allen Goodman, Sumanth Gurram, Jiabo Hu, Somya Jain, Devansh Kukreja, Robert Kuo, Joshua Lane, Yanghao Li, Lilian Luong, Jitendra Malik, Mallika Malhotra, William Ngan, Omkar Parkhi, Nikhil Raina, Dirk Rowe, Neil Sejoor, Vanessa Stark, Bala Varadarajan, Bram Wasti, Zachary Winstrom
+
+## Citing Segment Anything
+
+If you use SAM or SA-1B in your research, please use the following BibTeX entry.
+
+```
+@article{kirillov2023segany,
+  title={Segment Anything},
+  author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
+  journal={arXiv:2304.02643},
+  year={2023}
+}
+```

models/sam/linter.sh

@@ -0,0 +1,32 @@
+#!/bin/bash -e
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+{
+  black --version | grep -E "23\." > /dev/null
+} || {
+  echo "Linter requires 'black==23.*' !"
+  exit 1
+}
+
+ISORT_VERSION=$(isort --version-number)
+if [[ "$ISORT_VERSION" != 5.12* ]]; then
+  echo "Linter requires isort==5.12.0 !"
+  exit 1
+fi
+
+echo "Running isort ..."
+isort . --atomic
+
+echo "Running black ..."
+black -l 100 .
+
+echo "Running flake8 ..."
+if [ -x "$(command -v flake8)" ]; then
+  flake8 .
+else
+  python3 -m flake8 .
+fi
+
+echo "Running mypy..."
+
+mypy --exclude 'setup.py|notebooks' .

models/sam/notebooks/onnx_model_example.ipynb

@@ -0,0 +1,774 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "901c8ef3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Copyright (c) Meta Platforms, Inc. and affiliates."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1662bb7c",
+   "metadata": {},
+   "source": [
+    "# Produces masks from prompts using an ONNX model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7fcc21a0",
+   "metadata": {},
+   "source": [
+    "SAM's prompt encoder and mask decoder are very lightweight, which allows for efficient computation of a mask given user input. This notebook shows an example of how to export and use this lightweight component of the model in ONNX format, allowing it to run on a variety of platforms that support an ONNX runtime."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "86daff77",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
+       "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+       "</a>\n"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from IPython.display import display, HTML\n",
+    "display(HTML(\n",
+    "\"\"\"\n",
+    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
+    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+    "</a>\n",
+    "\"\"\"\n",
+    "))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "55ae4e00",
+   "metadata": {},
+   "source": [
+    "## Environment Set-up"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "109a5cc2",
+   "metadata": {},
+   "source": [
+    "If running locally using jupyter, first install `segment_anything` in your environment using the [installation instructions](https://github.com/facebookresearch/segment-anything#installation) in the repository. The latest stable versions of PyTorch and ONNX are recommended for this notebook. If running from Google Colab, set `using_colab=True` below and run the cell. In Colab, be sure to select 'GPU' under 'Edit'->'Notebook Settings'->'Hardware accelerator'."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "39b99fc4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "using_colab = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "296a69be",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if using_colab:\n",
+    "    import torch\n",
+    "    import torchvision\n",
+    "    print(\"PyTorch version:\", torch.__version__)\n",
+    "    print(\"Torchvision version:\", torchvision.__version__)\n",
+    "    print(\"CUDA is available:\", torch.cuda.is_available())\n",
+    "    import sys\n",
+    "    !{sys.executable} -m pip install opencv-python matplotlib onnx onnxruntime\n",
+    "    !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'\n",
+    "    \n",
+    "    !mkdir images\n",
+    "    !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg\n",
+    "        \n",
+    "    !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dc4a58be",
+   "metadata": {},
+   "source": [
+    "## Set-up"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42396e8d",
+   "metadata": {},
+   "source": [
+    "Note that this notebook requires both the `onnx` and `onnxruntime` optional dependencies, in addition to `opencv-python` and `matplotlib` for visualization."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c712610",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import numpy as np\n",
+    "import cv2\n",
+    "import matplotlib.pyplot as plt\n",
+    "from segment_anything import sam_model_registry, SamPredictor\n",
+    "from segment_anything.utils.onnx import SamOnnxModel\n",
+    "\n",
+    "import onnxruntime\n",
+    "from onnxruntime.quantization import QuantType\n",
+    "from onnxruntime.quantization.quantize import quantize_dynamic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f29441b9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def show_mask(mask, ax):\n",
+    "    color = np.array([30/255, 144/255, 255/255, 0.6])\n",
+    "    h, w = mask.shape[-2:]\n",
+    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
+    "    ax.imshow(mask_image)\n",
+    "    \n",
+    "def show_points(coords, labels, ax, marker_size=375):\n",
+    "    pos_points = coords[labels==1]\n",
+    "    neg_points = coords[labels==0]\n",
+    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
+    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
+    "    \n",
+    "def show_box(box, ax):\n",
+    "    x0, y0 = box[0], box[1]\n",
+    "    w, h = box[2] - box[0], box[3] - box[1]\n",
+    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))   "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bd0f6b2b",
+   "metadata": {},
+   "source": [
+    "## Export an ONNX model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1540f719",
+   "metadata": {},
+   "source": [
+    "Set the path below to a SAM model checkpoint, then load the model. This will be needed to both export the model and to calculate embeddings for the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76fc53f4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint = \"sam_vit_h_4b8939.pth\"\n",
+    "model_type = \"vit_h\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11bfc8aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sam = sam_model_registry[model_type](checkpoint=checkpoint)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "450c089c",
+   "metadata": {},
+   "source": [
+    "The script `segment-anything/scripts/export_onnx_model.py` can be used to export the necessary portion of SAM. Alternatively, run the following code to export an ONNX model. If you have already exported a model, set the path below and skip to the next section. Assure that the exported ONNX model aligns with the checkpoint and model type set above. This notebook expects the model was exported with the parameter `return_single_mask=True`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "38a8add8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_model_path = None  # Set to use an already exported model, then skip to the next section."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7da638ba",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import warnings\n",
+    "\n",
+    "onnx_model_path = \"sam_onnx_example.onnx\"\n",
+    "\n",
+    "onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
+    "\n",
+    "dynamic_axes = {\n",
+    "    \"point_coords\": {1: \"num_points\"},\n",
+    "    \"point_labels\": {1: \"num_points\"},\n",
+    "}\n",
+    "\n",
+    "embed_dim = sam.prompt_encoder.embed_dim\n",
+    "embed_size = sam.prompt_encoder.image_embedding_size\n",
+    "mask_input_size = [4 * x for x in embed_size]\n",
+    "dummy_inputs = {\n",
+    "    \"image_embeddings\": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),\n",
+    "    \"point_coords\": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),\n",
+    "    \"point_labels\": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),\n",
+    "    \"mask_input\": torch.randn(1, 1, *mask_input_size, dtype=torch.float),\n",
+    "    \"has_mask_input\": torch.tensor([1], dtype=torch.float),\n",
+    "    \"orig_im_size\": torch.tensor([1500, 2250], dtype=torch.float),\n",
+    "}\n",
+    "output_names = [\"masks\", \"iou_predictions\", \"low_res_masks\"]\n",
+    "\n",
+    "with warnings.catch_warnings():\n",
+    "    warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
+    "    warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
+    "    with open(onnx_model_path, \"wb\") as f:\n",
+    "        torch.onnx.export(\n",
+    "            onnx_model,\n",
+    "            tuple(dummy_inputs.values()),\n",
+    "            f,\n",
+    "            export_params=True,\n",
+    "            verbose=False,\n",
+    "            opset_version=17,\n",
+    "            do_constant_folding=True,\n",
+    "            input_names=list(dummy_inputs.keys()),\n",
+    "            output_names=output_names,\n",
+    "            dynamic_axes=dynamic_axes,\n",
+    "        )    "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c450cf1a",
+   "metadata": {},
+   "source": [
+    "If desired, the model can additionally be quantized and optimized. We find this improves web runtime significantly for negligible change in qualitative performance. Run the next cell to quantize the model, or skip to the next section otherwise."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "235d39fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_model_quantized_path = \"sam_onnx_quantized_example.onnx\"\n",
+    "quantize_dynamic(\n",
+    "    model_input=onnx_model_path,\n",
+    "    model_output=onnx_model_quantized_path,\n",
+    "    optimize_model=True,\n",
+    "    per_channel=False,\n",
+    "    reduce_range=False,\n",
+    "    weight_type=QuantType.QUInt8,\n",
+    ")\n",
+    "onnx_model_path = onnx_model_quantized_path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "927a928b",
+   "metadata": {},
+   "source": [
+    "## Example Image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6be6eb55",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = cv2.imread('images/truck.jpg')\n",
+    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7e9a27a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10,10))\n",
+    "plt.imshow(image)\n",
+    "plt.axis('on')\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "027b177b",
+   "metadata": {},
+   "source": [
+    "## Using an ONNX model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "778d4593",
+   "metadata": {},
+   "source": [
+    "Here as an example, we use `onnxruntime` in python on CPU to execute the ONNX model. However, any platform that supports an ONNX runtime could be used in principle. Launch the runtime session below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9689b1bf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ort_session = onnxruntime.InferenceSession(onnx_model_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7708ead6",
+   "metadata": {},
+   "source": [
+    "To use the ONNX model, the image must first be pre-processed using the SAM image encoder. This is a heavier weight process best performed on GPU. SamPredictor can be used as normal, then `.get_image_embedding()` will retreive the intermediate features."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26e067b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sam.to(device='cuda')\n",
+    "predictor = SamPredictor(sam)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ad3f0d6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor.set_image(image)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8a6f0f07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_embedding = predictor.get_image_embedding().cpu().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5e112f33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_embedding.shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6337b654",
+   "metadata": {},
+   "source": [
+    "The ONNX model has a different input signature than `SamPredictor.predict`. The following inputs must all be supplied. Note the special cases for both point and mask inputs. All inputs are `np.float32`.\n",
+    "* `image_embeddings`: The image embedding from `predictor.get_image_embedding()`. Has a batch index of length 1.\n",
+    "* `point_coords`: Coordinates of sparse input prompts, corresponding to both point inputs and box inputs. Boxes are encoded using two points, one for the top-left corner and one for the bottom-right corner. *Coordinates must already be transformed to long-side 1024.* Has a batch index of length 1.\n",
+    "* `point_labels`: Labels for the sparse input prompts. 0 is a negative input point, 1 is a positive input point, 2 is a top-left box corner, 3 is a bottom-right box corner, and -1 is a padding point. *If there is no box input, a single padding point with label -1 and coordinates (0.0, 0.0) should be concatenated.*\n",
+    "* `mask_input`: A mask input to the model with shape 1x1x256x256. This must be supplied even if there is no mask input. In this case, it can just be zeros.\n",
+    "* `has_mask_input`: An indicator for the mask input. 1 indicates a mask input, 0 indicates no mask input.\n",
+    "* `orig_im_size`: The size of the input image in (H,W) format, before any transformation. \n",
+    "\n",
+    "Additionally, the ONNX model does not threshold the output mask logits. To obtain a binary mask, threshold at `sam.mask_threshold` (equal to 0.0)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bf5a9f55",
+   "metadata": {},
+   "source": [
+    "### Example point input"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c0deef0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_point = np.array([[500, 375]])\n",
+    "input_label = np.array([1])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7256394c",
+   "metadata": {},
+   "source": [
+    "Add a batch index, concatenate a padding point, and transform."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4f69903e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
+    "onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
+    "\n",
+    "onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b188dc53",
+   "metadata": {},
+   "source": [
+    "Create an empty mask input and an indicator for no mask."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5cb52bcf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
+    "onnx_has_mask_input = np.zeros(1, dtype=np.float32)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a99c2cc5",
+   "metadata": {},
+   "source": [
+    "Package the inputs to run in the onnx model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1d7ea11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ort_inputs = {\n",
+    "    \"image_embeddings\": image_embedding,\n",
+    "    \"point_coords\": onnx_coord,\n",
+    "    \"point_labels\": onnx_label,\n",
+    "    \"mask_input\": onnx_mask_input,\n",
+    "    \"has_mask_input\": onnx_has_mask_input,\n",
+    "    \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b6409c9",
+   "metadata": {},
+   "source": [
+    "Predict a mask and threshold it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dc4cc082",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "masks, _, low_res_logits = ort_session.run(None, ort_inputs)\n",
+    "masks = masks > predictor.model.mask_threshold"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d778a8fb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "masks.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "badb1175",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10,10))\n",
+    "plt.imshow(image)\n",
+    "show_mask(masks, plt.gca())\n",
+    "show_points(input_point, input_label, plt.gca())\n",
+    "plt.axis('off')\n",
+    "plt.show() "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1f1d4d15",
+   "metadata": {},
+   "source": [
+    "### Example mask input"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b319da82",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_point = np.array([[500, 375], [1125, 625]])\n",
+    "input_label = np.array([1, 1])\n",
+    "\n",
+    "# Use the mask output from the previous run. It is already in the correct form for input to the ONNX model.\n",
+    "onnx_mask_input = low_res_logits"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b1823b37",
+   "metadata": {},
+   "source": [
+    "Transform the points as in the previous example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8885130f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
+    "onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
+    "\n",
+    "onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "28e47b69",
+   "metadata": {},
+   "source": [
+    "The `has_mask_input` indicator is now 1."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3ab4483a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_has_mask_input = np.ones(1, dtype=np.float32)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d3781955",
+   "metadata": {},
+   "source": [
+    "Package inputs, then predict and threshold the mask."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0c1ec096",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ort_inputs = {\n",
+    "    \"image_embeddings\": image_embedding,\n",
+    "    \"point_coords\": onnx_coord,\n",
+    "    \"point_labels\": onnx_label,\n",
+    "    \"mask_input\": onnx_mask_input,\n",
+    "    \"has_mask_input\": onnx_has_mask_input,\n",
+    "    \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
+    "}\n",
+    "\n",
+    "masks, _, _ = ort_session.run(None, ort_inputs)\n",
+    "masks = masks > predictor.model.mask_threshold"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e36554b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10,10))\n",
+    "plt.imshow(image)\n",
+    "show_mask(masks, plt.gca())\n",
+    "show_points(input_point, input_label, plt.gca())\n",
+    "plt.axis('off')\n",
+    "plt.show() "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2ef211d0",
+   "metadata": {},
+   "source": [
+    "### Example box and point input"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51e58d2e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_box = np.array([425, 600, 700, 875])\n",
+    "input_point = np.array([[575, 750]])\n",
+    "input_label = np.array([0])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6e119dcb",
+   "metadata": {},
+   "source": [
+    "Add a batch index, concatenate a box and point inputs, add the appropriate labels for the box corners, and transform. There is no padding point since the input includes a box input."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bfbe4911",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_box_coords = input_box.reshape(2, 2)\n",
+    "onnx_box_labels = np.array([2,3])\n",
+    "\n",
+    "onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]\n",
+    "onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)\n",
+    "\n",
+    "onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65edabd2",
+   "metadata": {},
+   "source": [
+    "Package inputs, then predict and threshold the mask."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2abfba56",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
+    "onnx_has_mask_input = np.zeros(1, dtype=np.float32)\n",
+    "\n",
+    "ort_inputs = {\n",
+    "    \"image_embeddings\": image_embedding,\n",
+    "    \"point_coords\": onnx_coord,\n",
+    "    \"point_labels\": onnx_label,\n",
+    "    \"mask_input\": onnx_mask_input,\n",
+    "    \"has_mask_input\": onnx_has_mask_input,\n",
+    "    \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
+    "}\n",
+    "\n",
+    "masks, _, _ = ort_session.run(None, ort_inputs)\n",
+    "masks = masks > predictor.model.mask_threshold"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8301bf33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10, 10))\n",
+    "plt.imshow(image)\n",
+    "show_mask(masks[0], plt.gca())\n",
+    "show_box(input_box, plt.gca())\n",
+    "show_points(input_point, input_label, plt.gca())\n",
+    "plt.axis('off')\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/sam/scripts/amg.py

@@ -0,0 +1,238 @@
+# 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 cv2  # type: ignore
+
+from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
+
+import argparse
+import json
+import os
+from typing import Any, Dict, List
+
+parser = argparse.ArgumentParser(
+    description=(
+        "Runs automatic mask generation on an input image or directory of images, "
+        "and outputs masks as either PNGs or COCO-style RLEs. Requires open-cv, "
+        "as well as pycocotools if saving in RLE format."
+    )
+)
+
+parser.add_argument(
+    "--input",
+    type=str,
+    required=True,
+    help="Path to either a single input image or folder of images.",
+)
+
+parser.add_argument(
+    "--output",
+    type=str,
+    required=True,
+    help=(
+        "Path to the directory where masks will be output. Output will be either a folder "
+        "of PNGs per image or a single json with COCO-style masks."
+    ),
+)
+
+parser.add_argument(
+    "--model-type",
+    type=str,
+    required=True,
+    help="The type of model to load, in ['default', 'vit_h', 'vit_l', 'vit_b']",
+)
+
+parser.add_argument(
+    "--checkpoint",
+    type=str,
+    required=True,
+    help="The path to the SAM checkpoint to use for mask generation.",
+)
+
+parser.add_argument("--device", type=str, default="cuda", help="The device to run generation on.")
+
+parser.add_argument(
+    "--convert-to-rle",
+    action="store_true",
+    help=(
+        "Save masks as COCO RLEs in a single json instead of as a folder of PNGs. "
+        "Requires pycocotools."
+    ),
+)
+
+amg_settings = parser.add_argument_group("AMG Settings")
+
+amg_settings.add_argument(
+    "--points-per-side",
+    type=int,
+    default=None,
+    help="Generate masks by sampling a grid over the image with this many points to a side.",
+)
+
+amg_settings.add_argument(
+    "--points-per-batch",
+    type=int,
+    default=None,
+    help="How many input points to process simultaneously in one batch.",
+)
+
+amg_settings.add_argument(
+    "--pred-iou-thresh",
+    type=float,
+    default=None,
+    help="Exclude masks with a predicted score from the model that is lower than this threshold.",
+)
+
+amg_settings.add_argument(
+    "--stability-score-thresh",
+    type=float,
+    default=None,
+    help="Exclude masks with a stability score lower than this threshold.",
+)
+
+amg_settings.add_argument(
+    "--stability-score-offset",
+    type=float,
+    default=None,
+    help="Larger values perturb the mask more when measuring stability score.",
+)
+
+amg_settings.add_argument(
+    "--box-nms-thresh",
+    type=float,
+    default=None,
+    help="The overlap threshold for excluding a duplicate mask.",
+)
+
+amg_settings.add_argument(
+    "--crop-n-layers",
+    type=int,
+    default=None,
+    help=(
+        "If >0, mask generation is run on smaller crops of the image to generate more masks. "
+        "The value sets how many different scales to crop at."
+    ),
+)
+
+amg_settings.add_argument(
+    "--crop-nms-thresh",
+    type=float,
+    default=None,
+    help="The overlap threshold for excluding duplicate masks across different crops.",
+)
+
+amg_settings.add_argument(
+    "--crop-overlap-ratio",
+    type=int,
+    default=None,
+    help="Larger numbers mean image crops will overlap more.",
+)
+
+amg_settings.add_argument(
+    "--crop-n-points-downscale-factor",
+    type=int,
+    default=None,
+    help="The number of points-per-side in each layer of crop is reduced by this factor.",
+)
+
+amg_settings.add_argument(
+    "--min-mask-region-area",
+    type=int,
+    default=None,
+    help=(
+        "Disconnected mask regions or holes with area smaller than this value "
+        "in pixels are removed by postprocessing."
+    ),
+)
+
+
+def write_masks_to_folder(masks: List[Dict[str, Any]], path: str) -> None:
+    header = "id,area,bbox_x0,bbox_y0,bbox_w,bbox_h,point_input_x,point_input_y,predicted_iou,stability_score,crop_box_x0,crop_box_y0,crop_box_w,crop_box_h"  # noqa
+    metadata = [header]
+    for i, mask_data in enumerate(masks):
+        mask = mask_data["segmentation"]
+        filename = f"{i}.png"
+        cv2.imwrite(os.path.join(path, filename), mask * 255)
+        mask_metadata = [
+            str(i),
+            str(mask_data["area"]),
+            *[str(x) for x in mask_data["bbox"]],
+            *[str(x) for x in mask_data["point_coords"][0]],
+            str(mask_data["predicted_iou"]),
+            str(mask_data["stability_score"]),
+            *[str(x) for x in mask_data["crop_box"]],
+        ]
+        row = ",".join(mask_metadata)
+        metadata.append(row)
+    metadata_path = os.path.join(path, "metadata.csv")
+    with open(metadata_path, "w") as f:
+        f.write("\n".join(metadata))
+
+    return
+
+
+def get_amg_kwargs(args):
+    amg_kwargs = {
+        "points_per_side": args.points_per_side,
+        "points_per_batch": args.points_per_batch,
+        "pred_iou_thresh": args.pred_iou_thresh,
+        "stability_score_thresh": args.stability_score_thresh,
+        "stability_score_offset": args.stability_score_offset,
+        "box_nms_thresh": args.box_nms_thresh,
+        "crop_n_layers": args.crop_n_layers,
+        "crop_nms_thresh": args.crop_nms_thresh,
+        "crop_overlap_ratio": args.crop_overlap_ratio,
+        "crop_n_points_downscale_factor": args.crop_n_points_downscale_factor,
+        "min_mask_region_area": args.min_mask_region_area,
+    }
+    amg_kwargs = {k: v for k, v in amg_kwargs.items() if v is not None}
+    return amg_kwargs
+
+
+def main(args: argparse.Namespace) -> None:
+    print("Loading model...")
+    sam = sam_model_registry[args.model_type](checkpoint=args.checkpoint)
+    _ = sam.to(device=args.device)
+    output_mode = "coco_rle" if args.convert_to_rle else "binary_mask"
+    amg_kwargs = get_amg_kwargs(args)
+    generator = SamAutomaticMaskGenerator(sam, output_mode=output_mode, **amg_kwargs)
+
+    if not os.path.isdir(args.input):
+        targets = [args.input]
+    else:
+        targets = [
+            f for f in os.listdir(args.input) if not os.path.isdir(os.path.join(args.input, f))
+        ]
+        targets = [os.path.join(args.input, f) for f in targets]
+
+    os.makedirs(args.output, exist_ok=True)
+
+    for t in targets:
+        print(f"Processing '{t}'...")
+        image = cv2.imread(t)
+        if image is None:
+            print(f"Could not load '{t}' as an image, skipping...")
+            continue
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        masks = generator.generate(image)
+
+        base = os.path.basename(t)
+        base = os.path.splitext(base)[0]
+        save_base = os.path.join(args.output, base)
+        if output_mode == "binary_mask":
+            os.makedirs(save_base, exist_ok=False)
+            write_masks_to_folder(masks, save_base)
+        else:
+            save_file = save_base + ".json"
+            with open(save_file, "w") as f:
+                json.dump(masks, f)
+    print("Done!")
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    main(args)

models/sam/scripts/export_onnx_model.py

@@ -0,0 +1,201 @@
+# 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
+
+from segment_anything import sam_model_registry
+from segment_anything.utils.onnx import SamOnnxModel
+
+import argparse
+import warnings
+
+try:
+    import onnxruntime  # type: ignore
+
+    onnxruntime_exists = True
+except ImportError:
+    onnxruntime_exists = False
+
+parser = argparse.ArgumentParser(
+    description="Export the SAM prompt encoder and mask decoder to an ONNX model."
+)
+
+parser.add_argument(
+    "--checkpoint", type=str, required=True, help="The path to the SAM model checkpoint."
+)
+
+parser.add_argument(
+    "--output", type=str, required=True, help="The filename to save the ONNX model to."
+)
+
+parser.add_argument(
+    "--model-type",
+    type=str,
+    required=True,
+    help="In ['default', 'vit_h', 'vit_l', 'vit_b']. Which type of SAM model to export.",
+)
+
+parser.add_argument(
+    "--return-single-mask",
+    action="store_true",
+    help=(
+        "If true, the exported ONNX model will only return the best mask, "
+        "instead of returning multiple masks. For high resolution images "
+        "this can improve runtime when upscaling masks is expensive."
+    ),
+)
+
+parser.add_argument(
+    "--opset",
+    type=int,
+    default=17,
+    help="The ONNX opset version to use. Must be >=11",
+)
+
+parser.add_argument(
+    "--quantize-out",
+    type=str,
+    default=None,
+    help=(
+        "If set, will quantize the model and save it with this name. "
+        "Quantization is performed with quantize_dynamic from onnxruntime.quantization.quantize."
+    ),
+)
+
+parser.add_argument(
+    "--gelu-approximate",
+    action="store_true",
+    help=(
+        "Replace GELU operations with approximations using tanh. Useful "
+        "for some runtimes that have slow or unimplemented erf ops, used in GELU."
+    ),
+)
+
+parser.add_argument(
+    "--use-stability-score",
+    action="store_true",
+    help=(
+        "Replaces the model's predicted mask quality score with the stability "
+        "score calculated on the low resolution masks using an offset of 1.0. "
+    ),
+)
+
+parser.add_argument(
+    "--return-extra-metrics",
+    action="store_true",
+    help=(
+        "The model will return five results: (masks, scores, stability_scores, "
+        "areas, low_res_logits) instead of the usual three. This can be "
+        "significantly slower for high resolution outputs."
+    ),
+)
+
+
+def run_export(
+    model_type: str,
+    checkpoint: str,
+    output: str,
+    opset: int,
+    return_single_mask: bool,
+    gelu_approximate: bool = False,
+    use_stability_score: bool = False,
+    return_extra_metrics=False,
+):
+    print("Loading model...")
+    sam = sam_model_registry[model_type](checkpoint=checkpoint)
+
+    onnx_model = SamOnnxModel(
+        model=sam,
+        return_single_mask=return_single_mask,
+        use_stability_score=use_stability_score,
+        return_extra_metrics=return_extra_metrics,
+    )
+
+    if gelu_approximate:
+        for n, m in onnx_model.named_modules():
+            if isinstance(m, torch.nn.GELU):
+                m.approximate = "tanh"
+
+    dynamic_axes = {
+        "point_coords": {1: "num_points"},
+        "point_labels": {1: "num_points"},
+    }
+
+    embed_dim = sam.prompt_encoder.embed_dim
+    embed_size = sam.prompt_encoder.image_embedding_size
+    mask_input_size = [4 * x for x in embed_size]
+    dummy_inputs = {
+        "image_embeddings": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),
+        "point_coords": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),
+        "point_labels": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),
+        "mask_input": torch.randn(1, 1, *mask_input_size, dtype=torch.float),
+        "has_mask_input": torch.tensor([1], dtype=torch.float),
+        "orig_im_size": torch.tensor([1500, 2250], dtype=torch.float),
+    }
+
+    _ = onnx_model(**dummy_inputs)
+
+    output_names = ["masks", "iou_predictions", "low_res_masks"]
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
+        warnings.filterwarnings("ignore", category=UserWarning)
+        with open(output, "wb") as f:
+            print(f"Exporting onnx model to {output}...")
+            torch.onnx.export(
+                onnx_model,
+                tuple(dummy_inputs.values()),
+                f,
+                export_params=True,
+                verbose=False,
+                opset_version=opset,
+                do_constant_folding=True,
+                input_names=list(dummy_inputs.keys()),
+                output_names=output_names,
+                dynamic_axes=dynamic_axes,
+            )
+
+    if onnxruntime_exists:
+        ort_inputs = {k: to_numpy(v) for k, v in dummy_inputs.items()}
+        # set cpu provider default
+        providers = ["CPUExecutionProvider"]
+        ort_session = onnxruntime.InferenceSession(output, providers=providers)
+        _ = ort_session.run(None, ort_inputs)
+        print("Model has successfully been run with ONNXRuntime.")
+
+
+def to_numpy(tensor):
+    return tensor.cpu().numpy()
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    run_export(
+        model_type=args.model_type,
+        checkpoint=args.checkpoint,
+        output=args.output,
+        opset=args.opset,
+        return_single_mask=args.return_single_mask,
+        gelu_approximate=args.gelu_approximate,
+        use_stability_score=args.use_stability_score,
+        return_extra_metrics=args.return_extra_metrics,
+    )
+
+    if args.quantize_out is not None:
+        assert onnxruntime_exists, "onnxruntime is required to quantize the model."
+        from onnxruntime.quantization import QuantType  # type: ignore
+        from onnxruntime.quantization.quantize import quantize_dynamic  # type: ignore
+
+        print(f"Quantizing model and writing to {args.quantize_out}...")
+        quantize_dynamic(
+            model_input=args.output,
+            model_output=args.quantize_out,
+            optimize_model=True,
+            per_channel=False,
+            reduce_range=False,
+            weight_type=QuantType.QUInt8,
+        )
+        print("Done!")

models/sam/segment_anything/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+from .build_sam import (
+    build_sam,
+    build_sam_vit_h,
+    build_sam_vit_l,
+    build_sam_vit_b,
+    sam_model_registry,
+)
+from .predictor import SamPredictor
+from .automatic_mask_generator import SamAutomaticMaskGenerator