Upload CountGD++ weights and code

.gitattributes

@@ -33,3 +33,7 @@ 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
+img/inference-architecture.jpg filter=lfs diff=lfs merge=lfs -text
+img/synthetic_blueberry.jpg filter=lfs diff=lfs merge=lfs -text
+img/synthetic_strawberry.jpg filter=lfs diff=lfs merge=lfs -text
+img/teaser.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Niki Amini-Naieni
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,93 @@
+---
+license: mit
+language:
+- en
+library_name: pytorch
+pipeline_tag: zero-shot-object-detection
+tags:
+- CountGD++
+- computer-vision
+- counting
+- grounding-dino
+- model_hub_mixin
+- multi-modal
+- open-vocabulary
+- pytorch_model_hub_mixin
+---
+
+# CountGD++
+
+**CountGD++: Generalized Prompting for Open-World Counting** (CVPR 2026)
+
+A multi-modal open-world counting model. You say what to count with text,
+visual exemplars (boxes on example objects), or both. You can also say what
+**not** to count with negative text and negative exemplars. Exemplars can come
+from the image being counted or from a different image.
+
+[Paper](https://arxiv.org/abs/2512.23351) · [Code](https://github.com/niki-amini-naieni/CountGDPlusPlus)
+
+![teaser](img/teaser.jpg)
+
+This repository holds both the **weights** (`model.safetensors`, `config.json`)
+and the **inference code**, so nothing has to be downloaded from Google Drive.
+
+## Usage
+
+```bash
+pip install torch torchvision "transformers<4.50" timm addict "yapf==0.40.1" termcolor \
+    scipy matplotlib opencv-python-headless pycocotools pyyaml huggingface_hub safetensors
+```
+
+```python
+import sys
+from huggingface_hub import snapshot_download
+
+repo = snapshot_download("nikigoli/CountGDPlusPlus")   # code + weights
+sys.path.insert(0, repo)
+from countgdpp import CountGDPlusPlus, count
+
+model = CountGDPlusPlus.from_pretrained("nikigoli/CountGDPlusPlus").to("cuda").eval()
+
+# Text only
+print(count(model, f"{repo}/img/336.jpg", "strawberry")["count"])
+
+# Positive text + negative text
+print(count(model, f"{repo}/img/336.jpg", "strawberry", negative_texts=["blueberry"])["count"])
+
+# Text + visual exemplars drawn on another image, with a negative exemplar
+result = count(
+    model, f"{repo}/img/336.jpg", "strawberry",
+    exemplar_boxes=[[x1, y1, x2, y2]],                    # pixels, on exemplar_image
+    exemplar_image=f"{repo}/img/synthetic_strawberry.jpg",
+    negative_texts=["blueberry"],
+    negative_exemplar_boxes=[[[x1, y1, x2, y2]]],          # one list per negative text
+    negative_exemplar_images=[f"{repo}/img/synthetic_blueberry.jpg"],
+)
+print(result["count"], result["boxes"])   # boxes: normalized (cx, cy, w, h)
+```
+
+`example.py` is a runnable version that also saves a picture of the detections.
+
+Notes:
+- The text encoder (`bert-base-uncased`) is fetched from the Hugging Face Hub
+  by `transformers` the first time the model loads.
+- The compiled multi-scale deformable attention kernels are optional. Without
+  them the model uses an equivalent pure-PyTorch implementation on CPU or GPU.
+  To compile the kernels, see the GitHub repository.
+- Pseudo-exemplars and adaptive cropping used for some benchmark results are not
+  part of `count()`. See the test scripts in the GitHub repository.
+
+## Architecture
+
+![architecture](img/inference-architecture.jpg)
+
+## Citation
+
+```bibtex
+@InProceedings{AminiNaieni26b,
+  title={CountGD++: Generalized Prompting for Open-World Counting},
+  author={Amini-Naieni, N. and Zisserman, A.},
+  booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2026}
+}
+```

cfg_app.py

@@ -0,0 +1,118 @@
+data_aug_scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
+data_aug_max_size = 1333
+data_aug_scales2_resize = [400, 500, 600]
+data_aug_scales2_crop = [384, 600]
+data_aug_scale_overlap = None
+batch_size = 4
+modelname = 'groundingdino'
+backbone = "swin_B_384_22k"
+position_embedding = 'sine'
+pe_temperatureH = 20
+pe_temperatureW = 20
+return_interm_indices = [1, 2, 3]
+enc_layers = 6
+dec_layers = 6
+pre_norm = False
+dim_feedforward = 2048
+hidden_dim = 256
+dropout = 0.0
+nheads = 8
+num_queries = 900
+query_dim = 4
+num_patterns = 0
+num_feature_levels = 4
+enc_n_points = 4
+dec_n_points = 4
+two_stage_type = 'standard'
+two_stage_bbox_embed_share = False
+two_stage_class_embed_share = False
+transformer_activation = 'relu'
+dec_pred_bbox_embed_share = True
+dn_box_noise_scale = 1.0
+dn_label_noise_ratio = 0.5
+dn_label_coef = 1.0
+dn_bbox_coef = 1.0
+embed_init_tgt = True
+dn_labelbook_size = 91
+max_text_len = 256
+text_encoder_type = "bert-base-uncased"
+use_text_enhancer = True
+use_fusion_layer = True
+use_checkpoint = False
+use_transformer_ckpt = False
+use_text_cross_attention = True
+text_dropout = 0.0
+fusion_dropout = 0.0
+fusion_droppath = 0.1
+sub_sentence_present = True
+max_labels = 90                               # pos + neg
+lr = 0.0001                                   # base learning rate
+backbone_freeze_keywords = None               # only for gdino backbone
+freeze_keywords = ['backbone.0', 'bert']      # for whole model, e.g. ['backbone.0', 'bert'] for freeze visual encoder and text encoder
+lr_backbone = 1e-05                           # specific learning rate
+lr_backbone_names = ['backbone.0', 'bert']
+lr_linear_proj_mult = 1e-05
+lr_linear_proj_names = ['ref_point_head', 'sampling_offsets']
+weight_decay = 0.0001
+param_dict_type = 'ddetr_in_mmdet'
+ddetr_lr_param = False
+epochs = 30
+lr_drop = 10
+save_checkpoint_interval = 10
+clip_max_norm = 0.1
+onecyclelr = False
+multi_step_lr = False
+lr_drop_list = [10, 20]
+frozen_weights = None
+dilation = False
+pdetr3_bbox_embed_diff_each_layer = False
+pdetr3_refHW = -1
+random_refpoints_xy = False
+fix_refpoints_hw = -1
+dabdetr_yolo_like_anchor_update = False
+dabdetr_deformable_encoder = False
+dabdetr_deformable_decoder = False
+use_deformable_box_attn = False
+box_attn_type = 'roi_align'
+dec_layer_number = None
+decoder_layer_noise = False
+dln_xy_noise = 0.2
+dln_hw_noise = 0.2
+add_channel_attention = False
+add_pos_value = False
+two_stage_pat_embed = 0
+two_stage_add_query_num = 0
+two_stage_learn_wh = False
+two_stage_default_hw = 0.05
+two_stage_keep_all_tokens = False
+num_select = 900
+batch_norm_type = 'FrozenBatchNorm2d'
+masks = False
+aux_loss = True
+set_cost_class = 5.0
+set_cost_bbox = 1.0
+set_cost_giou = 0.0
+cls_loss_coef = 5.0
+bbox_loss_coef = 1.0
+giou_loss_coef = 0.0
+enc_loss_coef = 1.0
+interm_loss_coef = 1.0
+no_interm_box_loss = False
+mask_loss_coef = 1.0
+dice_loss_coef = 1.0
+focal_alpha = 0.25
+focal_gamma = 2.0
+decoder_sa_type = 'sa'
+matcher_type = 'HungarianMatcher'
+decoder_module_seq = ['sa', 'ca', 'ffn']
+nms_iou_threshold = -1
+dec_pred_class_embed_share = True
+match_unstable_error = True
+use_detached_boxes_dec_out = False
+dn_scalar = 100
+
+box_threshold = 0.23
+text_threshold = 0
+use_coco_eval = False
+label_list = ['alcohol bottle', 'baguette roll', 'ball', 'banana', 'bead', 'bee', 'birthday candle', 'biscuit', 'boat', 'bottle', 'bowl', 'box', 'bread roll', 'brick', 'buffalo', 'bun', 'calamari ring', 'can', 'candle', 'cap', 'car', 'cartridge', 'cassette', 'cement bag', 'cereal', 'chewing gum piece', 'chopstick', 'clam', 'coffee bean', 'coin', 'cotton ball', 'cow', 'crane', 'crayon', 'croissant', 'crow', 'cup', 'cupcake', 'cupcake holder', 'fish', 'gemstone', 'go game piece', 'goat', 'goldfish snack', 'goose', 'ice cream', 'ice cream cone', 'instant noodle', 'jade stone', 'jeans', 'kidney bean', 'kitchen towel', 'lighter', 'lipstick', 'm&m piece', 'macaron', 'match', 'meat skewer', 'mini blind', 'mosaic tile', 'naan bread', 'nail', 'nut', 'onion ring', 'orange', 'pearl', 'pen', 'pencil', 'penguin', 'pepper', 'person', 'pigeon', 'plate', 'polka dot tile', 'potato', 'rice bag', 'roof tile', 'screw', 'shoe', 'spoon', 'spring roll', 'stair', 'stapler pin', 'straw', 'supermarket shelf', 'swan', 'tomato', 'watermelon', 'window', 'zebra']
+val_label_list = ["apple", "candy piece", "carrom board piece", "cashew nut", "comic book", "crab cake", "deer", "egg", "elephant", "finger food", "green pea", "hot air balloon", "keyboard key", "lego", "marble", "marker", "nail polish", "potato chip", "red bean", "round dessert", "sauce bottle", "sea shell", "sheep", "ski", "stamp", "sticky note", "strawberry", "sunglasses", "tree log", "watch"]

config.json

@@ -0,0 +1,157 @@
+{
+  "add_channel_attention": false,
+  "add_pos_value": false,
+  "aux_loss": true,
+  "backbone": "swin_B_384_22k",
+  "backbone_freeze_keywords": null,
+  "batch_norm_type": "FrozenBatchNorm2d",
+  "batch_size": 4,
+  "bbox_loss_coef": 1.0,
+  "box_attn_type": "roi_align",
+  "box_threshold": 0.23,
+  "clip_max_norm": 0.1,
+  "cls_loss_coef": 5.0,
+  "dabdetr_deformable_decoder": false,
+  "dabdetr_deformable_encoder": false,
+  "dabdetr_yolo_like_anchor_update": false,
+  "data_aug_max_size": 1333,
+  "data_aug_scale_overlap": null,
+  "data_aug_scales": [
+    480,
+    512,
+    544,
+    576,
+    608,
+    640,
+    672,
+    704,
+    736,
+    768,
+    800
+  ],
+  "data_aug_scales2_crop": [
+    384,
+    600
+  ],
+  "data_aug_scales2_resize": [
+    400,
+    500,
+    600
+  ],
+  "ddetr_lr_param": false,
+  "dec_layer_number": null,
+  "dec_layers": 6,
+  "dec_n_points": 4,
+  "dec_pred_bbox_embed_share": true,
+  "dec_pred_class_embed_share": true,
+  "decoder_layer_noise": false,
+  "decoder_module_seq": [
+    "sa",
+    "ca",
+    "ffn"
+  ],
+  "decoder_sa_type": "sa",
+  "device": "cpu",
+  "dice_loss_coef": 1.0,
+  "dilation": false,
+  "dim_feedforward": 2048,
+  "dln_hw_noise": 0.2,
+  "dln_xy_noise": 0.2,
+  "dn_bbox_coef": 1.0,
+  "dn_box_noise_scale": 1.0,
+  "dn_label_coef": 1.0,
+  "dn_label_noise_ratio": 0.5,
+  "dn_labelbook_size": 91,
+  "dn_scalar": 100,
+  "dropout": 0.0,
+  "embed_init_tgt": true,
+  "enc_layers": 6,
+  "enc_loss_coef": 1.0,
+  "enc_n_points": 4,
+  "epochs": 30,
+  "fix_refpoints_hw": -1,
+  "focal_alpha": 0.25,
+  "focal_gamma": 2.0,
+  "freeze_keywords": [
+    "backbone.0",
+    "bert"
+  ],
+  "frozen_weights": null,
+  "fusion_dropout": 0.0,
+  "fusion_droppath": 0.1,
+  "giou_loss_coef": 0.0,
+  "hidden_dim": 256,
+  "interm_loss_coef": 1.0,
+  "lr": 0.0001,
+  "lr_backbone": 1e-05,
+  "lr_backbone_names": [
+    "backbone.0",
+    "bert"
+  ],
+  "lr_drop": 10,
+  "lr_drop_list": [
+    10,
+    20
+  ],
+  "lr_linear_proj_mult": 1e-05,
+  "lr_linear_proj_names": [
+    "ref_point_head",
+    "sampling_offsets"
+  ],
+  "mask_loss_coef": 1.0,
+  "masks": false,
+  "match_unstable_error": true,
+  "matcher_type": "HungarianMatcher",
+  "max_labels": 90,
+  "max_text_len": 256,
+  "modelname": "groundingdino",
+  "multi_step_lr": false,
+  "nheads": 8,
+  "nms_iou_threshold": -1,
+  "no_interm_box_loss": false,
+  "num_feature_levels": 4,
+  "num_patterns": 0,
+  "num_queries": 900,
+  "num_select": 900,
+  "onecyclelr": false,
+  "param_dict_type": "ddetr_in_mmdet",
+  "pdetr3_bbox_embed_diff_each_layer": false,
+  "pdetr3_refHW": -1,
+  "pe_temperatureH": 20,
+  "pe_temperatureW": 20,
+  "position_embedding": "sine",
+  "pre_norm": false,
+  "query_dim": 4,
+  "random_refpoints_xy": false,
+  "return_interm_indices": [
+    1,
+    2,
+    3
+  ],
+  "save_checkpoint_interval": 10,
+  "set_cost_bbox": 1.0,
+  "set_cost_class": 5.0,
+  "set_cost_giou": 0.0,
+  "sub_sentence_present": true,
+  "text_dropout": 0.0,
+  "text_encoder_type": "bert-base-uncased",
+  "text_threshold": 0,
+  "transformer_activation": "relu",
+  "two_stage_add_query_num": 0,
+  "two_stage_bbox_embed_share": false,
+  "two_stage_class_embed_share": false,
+  "two_stage_default_hw": 0.05,
+  "two_stage_keep_all_tokens": false,
+  "two_stage_learn_wh": false,
+  "two_stage_pat_embed": 0,
+  "two_stage_type": "standard",
+  "use_checkpoint": false,
+  "use_coco_eval": false,
+  "use_deformable_box_attn": false,
+  "use_detached_boxes_dec_out": false,
+  "use_fusion_layer": true,
+  "use_text_cross_attention": true,
+  "use_text_enhancer": true,
+  "use_transformer_ckpt": false,
+  "weight_decay": 0.0001
+}

countgdpp.py

@@ -0,0 +1,177 @@
+"""CountGD++ on the Hugging Face Hub.
+
+    from countgdpp import CountGDPlusPlus, count
+    model = CountGDPlusPlus.from_pretrained("nikigoli/CountGDPlusPlus").to("cuda").eval()
+    result = count(model, "img/336.jpg", "strawberry", negative_texts=["blueberry"])
+    print(result["count"])
+
+`CountGDPlusPlus` is the CountGD++ GroundingDINO model with Hugging Face's
+PyTorchModelHubMixin added, so `from_pretrained`, `save_pretrained` and
+`push_to_hub` work. Its state dict keys are identical to the original
+`countgd_plusplus.pth` checkpoint.
+"""
+import os
+import sys
+from types import SimpleNamespace
+from typing import List, Optional, Sequence, Union
+
+import torch
+from PIL import Image
+from huggingface_hub import PyTorchModelHubMixin
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+
+from util.slconfig import SLConfig  # noqa: E402
+from util.misc import nested_tensor_from_tensor_list  # noqa: E402
+import datasets.transforms_app as T  # noqa: E402
+from models.GroundingDINO.groundingdino_app import GroundingDINO  # noqa: E402
+from models.GroundingDINO.backbone import build_backbone  # noqa: E402
+from models.GroundingDINO.transformer import build_transformer  # noqa: E402
+
+# Default score threshold used in the paper's demo.
+CONF_THRESH = 0.23
+# BERT token id of "." — separates the positive prompt from negative prompts in the caption.
+_DOT_TOKEN_ID = 1012
+_DROP_KEYS = {"label_list", "val_label_list"}
+
+
+def default_config() -> dict:
+    """Model configuration read from cfg_app.py (the config used by the CountGD++ demo)."""
+    cfg = SLConfig.fromfile(os.path.join(_HERE, "cfg_app.py"))._cfg_dict.to_dict()
+    cfg = {k: v for k, v in cfg.items() if k not in _DROP_KEYS}
+    cfg["text_encoder_type"] = "bert-base-uncased"  # fetched from the Hub by transformers
+    cfg["device"] = "cpu"
+    return cfg
+
+
+class CountGDPlusPlus(
+    GroundingDINO,
+    PyTorchModelHubMixin,
+    library_name="pytorch",
+    repo_url="https://github.com/niki-amini-naieni/CountGDPlusPlus",
+    license="mit",
+    tags=["CountGD++", "computer-vision", "counting", "grounding-dino", "multi-modal", "open-vocabulary"],
+):
+    """CountGD++ model. Build with `CountGDPlusPlus.from_pretrained(repo_id)`."""
+
+    def __init__(self, config: Optional[dict] = None):
+        config = dict(config or default_config())
+        args = SimpleNamespace(**config)
+        backbone = build_backbone(args)
+        transformer = build_transformer(args)
+        # Same arguments as build_groundingdino() in groundingdino_app.py
+        super().__init__(
+            backbone,
+            transformer,
+            num_queries=args.num_queries,
+            aux_loss=args.aux_loss,
+            iter_update=True,
+            query_dim=4,
+            num_feature_levels=args.num_feature_levels,
+            nheads=args.nheads,
+            dec_pred_bbox_embed_share=args.dec_pred_bbox_embed_share,
+            two_stage_type=args.two_stage_type,
+            two_stage_bbox_embed_share=args.two_stage_bbox_embed_share,
+            two_stage_class_embed_share=args.two_stage_class_embed_share,
+            num_patterns=args.num_patterns,
+            dn_number=0,
+            dn_box_noise_scale=args.dn_box_noise_scale,
+            dn_label_noise_ratio=args.dn_label_noise_ratio,
+            dn_labelbook_size=args.dn_labelbook_size,
+            text_encoder_type=args.text_encoder_type,
+            sub_sentence_present=args.sub_sentence_present,
+            max_text_len=args.max_text_len,
+        )
+        self.countgd_config = config
+
+
+# ----------------------------------------------------------------------------
+# Inference (same pre/post-processing as the CountGD++ demo app)
+# ----------------------------------------------------------------------------
+def build_transform():
+    normalize = T.Compose([T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
+    return T.Compose([T.RandomResize([800], max_size=1333), normalize])
+
+
+def _load(image: Union[str, Image.Image]) -> Image.Image:
+    return (Image.open(image) if isinstance(image, str) else image).convert("RGB")
+
+
+def _prep_exemplars(transform, image, boxes):
+    boxes = torch.tensor(boxes, dtype=torch.float32) if boxes else torch.tensor([])
+    img, target = transform(image, {"exemplars": boxes})
+    return img, target["exemplars"]
+
+
+@torch.no_grad()
+def count(
+    model: CountGDPlusPlus,
+    image: Union[str, Image.Image],
+    text: str,
+    exemplar_boxes: Optional[Sequence[Sequence[float]]] = None,
+    exemplar_image: Optional[Union[str, Image.Image]] = None,
+    negative_texts: Optional[List[str]] = None,
+    negative_exemplar_boxes: Optional[List[Sequence[Sequence[float]]]] = None,
+    negative_exemplar_images: Optional[List[Union[str, Image.Image]]] = None,
+    conf_thresh: float = CONF_THRESH,
+):
+    """Count objects in `image`.
+
+    text: what to count, e.g. "strawberry".
+    exemplar_boxes: optional visual exemplars [[x1, y1, x2, y2], ...] in pixels,
+        drawn on `exemplar_image` (defaults to `image`).
+    negative_texts: optional things NOT to count, e.g. ["blueberry"].
+    negative_exemplar_boxes / negative_exemplar_images: optional, one entry per
+        negative text, same format as the positive ones.
+
+    Returns dict with "count", "boxes" (normalized cx, cy, w, h) and "scores".
+    """
+    device = next(model.parameters()).device
+    transform = build_transform()
+    image = _load(image)
+    pos_img = _load(exemplar_image) if exemplar_image is not None else image
+    negative_texts = list(negative_texts or [])
+    n_neg = len(negative_texts)
+    neg_boxes = list(negative_exemplar_boxes or [None] * n_neg)
+    neg_imgs = list(negative_exemplar_images or [None] * n_neg)
+    if len(neg_boxes) != n_neg or len(neg_imgs) != n_neg:
+        raise ValueError("Give one negative exemplar entry (or None) per negative text.")
+
+    input_image, _ = transform(image, None)
+    pos_exemplar_image, pos_exemplars = _prep_exemplars(transform, pos_img, exemplar_boxes)
+
+    neg_images, neg_exemplars = [], []
+    for boxes, img in zip(neg_boxes, neg_imgs):
+        img = _load(img) if img is not None else image
+        ex_img, ex = _prep_exemplars(transform, img, boxes)
+        neg_images.append(nested_tensor_from_tensor_list(ex_img.unsqueeze(0).to(device)))
+        neg_exemplars.append([ex.to(device)])
+
+    caption = text + " . " + "".join(t + " . " for t in negative_texts)
+    out = model(
+        nested_tensor_from_tensor_list(input_image.unsqueeze(0).to(device)),
+        nested_tensor_from_tensor_list(pos_exemplar_image.unsqueeze(0).to(device)),
+        [pos_exemplars.to(device)],
+        neg_images,
+        neg_exemplars,
+        captions=[caption],
+    )
+
+    input_ids = out["input_ids"][0]
+    logits = out["pred_logits"].sigmoid()[0]
+    boxes = out["pred_boxes"][0]
+    split = next(i for i, t in enumerate(input_ids) if int(t) == _DOT_TOKEN_ID)
+    pos_logits, neg_logits = logits[:, : split + 1], logits[:, split + 1 :]
+
+    # Stage 1: keep queries confident about the positive prompt.
+    keep = pos_logits.max(dim=-1).values > conf_thresh
+    boxes, scores = boxes[keep], logits[keep].max(dim=-1).values
+    pos_logits, neg_logits = pos_logits[keep], neg_logits[keep]
+    # Stage 2: drop queries that match a negative prompt better than the positive one.
+    if neg_logits.shape[-1] > 0:
+        keep = pos_logits.max(dim=-1).values > neg_logits.max(dim=-1).values
+        boxes, scores = boxes[keep], scores[keep]
+
+    return {"count": int(boxes.shape[0]), "boxes": boxes.cpu().numpy(), "scores": scores.cpu().numpy()}

datasets/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import torch.utils.data
+import torchvision
+from .coco import build as build_coco
+
+
+def get_coco_api_from_dataset(dataset):
+    for _ in range(10):
+        # if isinstance(dataset, torchvision.datasets.CocoDetection):
+        #     break
+        if isinstance(dataset, torch.utils.data.Subset):
+            dataset = dataset.dataset
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return dataset.coco
+
+
+def build_dataset(image_set, args, datasetinfo):
+    if datasetinfo["dataset_mode"] == 'coco':
+        return build_coco(image_set, args, datasetinfo)
+    if datasetinfo["dataset_mode"] == 'odvg':
+        from .odvg import build_odvg
+        return build_odvg(image_set, args, datasetinfo)
+    raise ValueError(f'dataset {args.dataset_file} not supported')

datasets/coco.py

@@ -0,0 +1,660 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+COCO dataset which returns image_id for evaluation.
+
+Mostly copy-paste from https://github.com/pytorch/vision/blob/13b35ff/references/detection/coco_utils.py
+"""
+if __name__=="__main__":
+    # for debug only
+    import os, sys
+    sys.path.append(os.path.dirname(sys.path[0]))
+from torchvision.datasets.vision import VisionDataset
+
+import json
+from pathlib import Path
+import random
+import os
+from typing import Any, Callable, List, Optional, Tuple
+
+from PIL import Image
+
+import torch
+import torch.utils.data
+import torchvision
+from pycocotools import mask as coco_mask
+
+from datasets.data_util import preparing_dataset
+import datasets.transforms_orig as T
+from util.box_ops import box_cxcywh_to_xyxy, box_iou
+
+__all__ = ['build']
+
+
+class label2compat():
+    def __init__(self) -> None:
+        self.category_map_str = {"1": 1, "2": 2, "3": 3, "4": 4, "5": 5, "6": 6, "7": 7, "8": 8, "9": 9, "10": 10, "11": 11, "13": 12, "14": 13, "15": 14, "16": 15, "17": 16, "18": 17, "19": 18, "20": 19, "21": 20, "22": 21, "23": 22, "24": 23, "25": 24, "27": 25, "28": 26, "31": 27, "32": 28, "33": 29, "34": 30, "35": 31, "36": 32, "37": 33, "38": 34, "39": 35, "40": 36, "41": 37, "42": 38, "43": 39, "44": 40, "46": 41, "47": 42, "48": 43, "49": 44, "50": 45, "51": 46, "52": 47, "53": 48, "54": 49, "55": 50, "56": 51, "57": 52, "58": 53, "59": 54, "60": 55, "61": 56, "62": 57, "63": 58, "64": 59, "65": 60, "67": 61, "70": 62, "72": 63, "73": 64, "74": 65, "75": 66, "76": 67, "77": 68, "78": 69, "79": 70, "80": 71, "81": 72, "82": 73, "84": 74, "85": 75, "86": 76, "87": 77, "88": 78, "89": 79, "90": 80}
+        self.category_map = {int(k):v for k,v in self.category_map_str.items()}
+
+    def __call__(self, target, img=None):
+        labels = target['labels']
+        res = torch.zeros(labels.shape, dtype=labels.dtype)
+        for idx, item in enumerate(labels):
+            res[idx] = self.category_map[item.item()] - 1
+        target['label_compat'] = res
+        if img is not None:
+            return target, img
+        else:
+            return target
+
+
+class label_compat2onehot():
+    def __init__(self, num_class=80, num_output_objs=1):
+        self.num_class = num_class
+        self.num_output_objs = num_output_objs
+        if num_output_objs != 1:
+            raise DeprecationWarning("num_output_objs!=1, which is only used for comparison")
+
+    def __call__(self, target, img=None):
+        labels = target['label_compat']
+        place_dict = {k:0 for k in range(self.num_class)}
+        if self.num_output_objs == 1:
+            res = torch.zeros(self.num_class)
+            for i in labels:
+                itm = i.item()
+                res[itm] = 1.0
+        else:
+            # compat with baseline
+            res = torch.zeros(self.num_class, self.num_output_objs)
+            for i in labels:
+                itm = i.item()
+                res[itm][place_dict[itm]] = 1.0
+                place_dict[itm] += 1
+        target['label_compat_onehot'] = res
+        if img is not None:
+            return target, img
+        else:
+            return target
+
+
+class box_label_catter():
+    def __init__(self):
+        pass
+
+    def __call__(self, target, img=None):
+        labels = target['label_compat']
+        boxes = target['boxes']
+        box_label = torch.cat((boxes, labels.unsqueeze(-1)), 1)
+        target['box_label'] = box_label
+        if img is not None:
+            return target, img
+        else:
+            return target
+
+
+class RandomSelectBoxlabels():
+    def __init__(self, num_classes, leave_one_out=False, blank_prob=0.8,
+                    prob_first_item = 0.0,
+                    prob_random_item = 0.0,
+                    prob_last_item = 0.8,
+                    prob_stop_sign = 0.2
+                ) -> None:
+        self.num_classes = num_classes
+        self.leave_one_out = leave_one_out
+        self.blank_prob = blank_prob
+
+        self.set_state(prob_first_item, prob_random_item, prob_last_item, prob_stop_sign)
+        
+
+    def get_state(self):
+        return [self.prob_first_item, self.prob_random_item, self.prob_last_item, self.prob_stop_sign]
+
+    def set_state(self, prob_first_item, prob_random_item, prob_last_item, prob_stop_sign):
+        sum_prob = prob_first_item + prob_random_item + prob_last_item + prob_stop_sign
+        assert sum_prob - 1 < 1e-6, \
+            f"Sum up all prob = {sum_prob}. prob_first_item:{prob_first_item}" \
+            + f"prob_random_item:{prob_random_item}, prob_last_item:{prob_last_item}" \
+            + f"prob_stop_sign:{prob_stop_sign}"
+
+        self.prob_first_item = prob_first_item
+        self.prob_random_item = prob_random_item
+        self.prob_last_item = prob_last_item
+        self.prob_stop_sign = prob_stop_sign
+        
+
+    def sample_for_pred_first_item(self, box_label: torch.FloatTensor):
+        box_label_known = torch.Tensor(0,5)
+        box_label_unknown = box_label
+        return box_label_known, box_label_unknown
+
+    def sample_for_pred_random_item(self, box_label: torch.FloatTensor):
+        n_select = int(random.random() * box_label.shape[0])
+        box_label = box_label[torch.randperm(box_label.shape[0])]
+        box_label_known = box_label[:n_select]
+        box_label_unknown = box_label[n_select:]
+        return box_label_known, box_label_unknown
+
+    def sample_for_pred_last_item(self, box_label: torch.FloatTensor):
+        box_label_perm = box_label[torch.randperm(box_label.shape[0])]
+        known_label_list = []
+        box_label_known = []
+        box_label_unknown = []
+        for item in box_label_perm:
+            label_i = item[4].item()
+            if label_i in known_label_list:
+                box_label_known.append(item)
+            else:
+                # first item
+                box_label_unknown.append(item)
+                known_label_list.append(label_i)
+        box_label_known = torch.stack(box_label_known) if len(box_label_known) > 0 else torch.Tensor(0,5)
+        box_label_unknown = torch.stack(box_label_unknown) if len(box_label_unknown) > 0 else torch.Tensor(0,5)
+        return box_label_known, box_label_unknown
+
+    def sample_for_pred_stop_sign(self, box_label: torch.FloatTensor):
+        box_label_unknown = torch.Tensor(0,5)
+        box_label_known = box_label
+        return box_label_known, box_label_unknown
+
+    def __call__(self, target, img=None):
+        box_label = target['box_label'] # K, 5
+
+        dice_number = random.random()
+
+        if dice_number < self.prob_first_item:
+            box_label_known, box_label_unknown = self.sample_for_pred_first_item(box_label)
+        elif dice_number < self.prob_first_item + self.prob_random_item:
+            box_label_known, box_label_unknown = self.sample_for_pred_random_item(box_label)
+        elif dice_number < self.prob_first_item + self.prob_random_item + self.prob_last_item:
+            box_label_known, box_label_unknown = self.sample_for_pred_last_item(box_label)
+        else:
+            box_label_known, box_label_unknown = self.sample_for_pred_stop_sign(box_label)
+
+        target['label_onehot_known'] = label2onehot(box_label_known[:,-1], self.num_classes)
+        target['label_onehot_unknown'] = label2onehot(box_label_unknown[:, -1], self.num_classes)
+        target['box_label_known'] = box_label_known
+        target['box_label_unknown'] = box_label_unknown
+
+        return target, img
+
+
+class RandomDrop():
+    def __init__(self, p=0.2) -> None:
+        self.p = p
+
+    def __call__(self, target, img=None):
+        known_box = target['box_label_known']
+        num_known_box = known_box.size(0)
+        idxs = torch.rand(num_known_box)
+        # indices = torch.randperm(num_known_box)[:int((1-self).p*num_known_box + 0.5 + random.random())]
+        target['box_label_known'] = known_box[idxs > self.p]
+        return target, img
+
+
+class BboxPertuber():
+    def __init__(self, max_ratio = 0.02, generate_samples = 1000) -> None:
+        self.max_ratio = max_ratio
+        self.generate_samples = generate_samples
+        self.samples = self.generate_pertube_samples()
+        self.idx = 0
+
+    def generate_pertube_samples(self):
+        import torch
+        samples = (torch.rand(self.generate_samples, 5) - 0.5) * 2 * self.max_ratio
+        return samples
+
+    def __call__(self, target, img):
+        known_box = target['box_label_known'] # Tensor(K,5), K known bbox
+        K = known_box.shape[0]
+        known_box_pertube = torch.zeros(K, 6) # 4:bbox, 1:prob, 1:label
+        if K == 0:
+            pass
+        else:
+            if self.idx + K > self.generate_samples:
+                self.idx = 0
+            delta = self.samples[self.idx: self.idx + K, :]
+            known_box_pertube[:, :4] = known_box[:, :4] + delta[:, :4]
+            iou = (torch.diag(box_iou(box_cxcywh_to_xyxy(known_box[:, :4]), box_cxcywh_to_xyxy(known_box_pertube[:, :4]))[0])) * (1 + delta[:, -1])
+            known_box_pertube[:, 4].copy_(iou)
+            known_box_pertube[:, -1].copy_(known_box[:, -1])
+
+        target['box_label_known_pertube'] = known_box_pertube
+        return target, img
+
+
+class RandomCutout():
+    def __init__(self, factor=0.5) -> None:
+        self.factor = factor
+
+    def __call__(self, target, img=None):
+        unknown_box = target['box_label_unknown']           # Ku, 5
+        known_box = target['box_label_known_pertube']       # Kk, 6
+        Ku = unknown_box.size(0)
+
+        known_box_add = torch.zeros(Ku, 6) # Ku, 6
+        known_box_add[:, :5] = unknown_box
+        known_box_add[:, 5].uniform_(0.5, 1) 
+        
+
+        known_box_add[:, :2] += known_box_add[:, 2:4] * (torch.rand(Ku, 2) - 0.5) / 2
+        known_box_add[:, 2:4] /= 2
+
+        target['box_label_known_pertube'] = torch.cat((known_box, known_box_add))
+        return target, img
+
+
+class RandomSelectBoxes():
+    def __init__(self, num_class=80) -> None:
+        Warning("This is such a slow function and will be deprecated soon!!!")
+        self.num_class = num_class
+
+    def __call__(self, target, img=None):
+        boxes = target['boxes']
+        labels = target['label_compat']
+
+        # transform to list of tensors
+        boxs_list = [[] for i in range(self.num_class)]
+        for idx, item in enumerate(boxes):
+            label = labels[idx].item()
+            boxs_list[label].append(item)
+        boxs_list_tensor = [torch.stack(i) if len(i) > 0 else torch.Tensor(0,4) for i in boxs_list]
+
+        # random selection
+        box_known = []
+        box_unknown = []
+        for idx, item in enumerate(boxs_list_tensor):
+            ncnt = item.shape[0]
+            nselect = int(random.random() * ncnt) # close in both sides, much faster than random.randint
+
+            item = item[torch.randperm(ncnt)]
+            # random.shuffle(item)
+            box_known.append(item[:nselect])
+            box_unknown.append(item[nselect:])
+
+        # box_known_tensor = [torch.stack(i) if len(i) > 0 else torch.Tensor(0,4) for i in box_known]
+        # box_unknown_tensor = [torch.stack(i) if len(i) > 0 else torch.Tensor(0,4) for i in box_unknown]
+        # print('box_unknown_tensor:', box_unknown_tensor)
+        target['known_box'] = box_known
+        target['unknown_box'] = box_unknown
+        return target, img
+
+
+def label2onehot(label, num_classes):
+    """
+    label: Tensor(K)
+    """
+    res = torch.zeros(num_classes)
+    for i in label:
+        itm = int(i.item())
+        res[itm] = 1.0
+    return res
+
+
+class MaskCrop():
+    def __init__(self) -> None:
+        pass
+
+    def __call__(self, target, img):
+        known_box = target['known_box']
+        h,w = img.shape[1:] # h,w
+        # imgsize = target['orig_size'] # h,w
+
+        scale = torch.Tensor([w, h, w, h])
+
+        # _cnt = 0
+        for boxes in known_box:
+            if boxes.shape[0] == 0:
+                continue
+            box_xyxy = box_cxcywh_to_xyxy(boxes) * scale
+            for box in box_xyxy:
+                x1, y1, x2, y2 = [int(i) for i in box.tolist()]
+                img[:, y1:y2, x1:x2] = 0
+                # _cnt += 1
+        # print("_cnt:", _cnt)
+        return target, img
+
+
+dataset_hook_register = {
+    'label2compat': label2compat,
+    'label_compat2onehot': label_compat2onehot,
+    'box_label_catter': box_label_catter,
+    'RandomSelectBoxlabels': RandomSelectBoxlabels,
+    'RandomSelectBoxes': RandomSelectBoxes,
+    'MaskCrop': MaskCrop,
+    'BboxPertuber': BboxPertuber,
+}
+
+
+class CocoDetection(torchvision.datasets.CocoDetection):
+    def __init__(self, img_folder, ann_file, transforms, return_masks, aux_target_hacks=None):
+        super(CocoDetection, self).__init__(img_folder, ann_file)
+        self._transforms = transforms
+        self.prepare = ConvertCocoPolysToMask(return_masks)
+        self.aux_target_hacks = aux_target_hacks
+
+    def change_hack_attr(self, hackclassname, attrkv_dict):
+        target_class = dataset_hook_register[hackclassname]
+        for item in self.aux_target_hacks:
+            if isinstance(item, target_class):
+                for k,v in attrkv_dict.items():
+                    setattr(item, k, v)
+
+    def get_hack(self, hackclassname):
+        target_class = dataset_hook_register[hackclassname]
+        for item in self.aux_target_hacks:
+            if isinstance(item, target_class):
+                return item
+            
+    def _load_image(self, id: int) -> Image.Image:
+        path = self.coco.loadImgs(id)[0]["file_name"]
+        abs_path = os.path.join(self.root, path)
+        return Image.open(abs_path).convert("RGB")
+
+    def __getitem__(self, idx):
+        """
+        Output:
+            - target: dict of multiple items
+                - boxes: Tensor[num_box, 4]. \
+                    Init type: x0,y0,x1,y1. unnormalized data.
+                    Final type: cx,cy,w,h. normalized data. 
+        """
+        try:
+            img, target = super(CocoDetection, self).__getitem__(idx)
+            
+        except:
+            print("Error idx: {}".format(idx))
+            idx += 1
+            img, target = super(CocoDetection, self).__getitem__(idx)
+        
+        image_id = self.ids[idx]
+        target = {'image_id': image_id, 'annotations': target}
+        exemp_count = 0
+        for instance in target['annotations']:
+            if instance['area'] != 4:
+                exemp_count += 1
+        assert exemp_count == 3
+        img, target = self.prepare(img, target)
+        target['exemplars'] = target['boxes'][-3:]
+        target['boxes'] = target['boxes'][:-3]
+        target['labels'] = target['labels'][:-3]
+        
+        
+        if self._transforms is not None:
+            img, target = self._transforms(img, target)
+
+        # convert to needed format
+        if self.aux_target_hacks is not None:
+            for hack_runner in self.aux_target_hacks:
+                target, img = hack_runner(target, img=img)
+
+        return img, target
+
+
+def convert_coco_poly_to_mask(segmentations, height, width):
+    masks = []
+    for polygons in segmentations:
+        rles = coco_mask.frPyObjects(polygons, height, width)
+        mask = coco_mask.decode(rles)
+        if len(mask.shape) < 3:
+            mask = mask[..., None]
+        mask = torch.as_tensor(mask, dtype=torch.uint8)
+        mask = mask.any(dim=2)
+        masks.append(mask)
+    if masks:
+        masks = torch.stack(masks, dim=0)
+    else:
+        masks = torch.zeros((0, height, width), dtype=torch.uint8)
+    return masks
+
+
+class ConvertCocoPolysToMask(object):
+    def __init__(self, return_masks=False):
+        self.return_masks = return_masks
+
+    def __call__(self, image, target):
+        w, h = image.size
+
+        image_id = target["image_id"]
+        image_id = torch.tensor([image_id])
+
+        anno = target["annotations"]
+
+        anno = [obj for obj in anno if 'iscrowd' not in obj or obj['iscrowd'] == 0]
+
+        boxes = [obj["bbox"] for obj in anno]
+        # guard against no boxes via resizing
+        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+        boxes[:, 2:] += boxes[:, :2]
+        boxes[:, 0::2].clamp_(min=0, max=w)
+        boxes[:, 1::2].clamp_(min=0, max=h)
+
+        classes = [obj["category_id"] for obj in anno]
+        classes = torch.tensor(classes, dtype=torch.int64)
+
+        if self.return_masks:
+            segmentations = [obj["segmentation"] for obj in anno]
+            masks = convert_coco_poly_to_mask(segmentations, h, w)
+
+        keypoints = None
+        if anno and "keypoints" in anno[0]:
+            keypoints = [obj["keypoints"] for obj in anno]
+            keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
+            num_keypoints = keypoints.shape[0]
+            if num_keypoints:
+                keypoints = keypoints.view(num_keypoints, -1, 3)
+
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        classes = classes[keep]
+        if self.return_masks:
+            masks = masks[keep]
+        if keypoints is not None:
+            keypoints = keypoints[keep]
+
+        target = {}
+        target["boxes"] = boxes
+        target["labels"] = classes
+        if self.return_masks:
+            target["masks"] = masks
+        target["image_id"] = image_id
+        if keypoints is not None:
+            target["keypoints"] = keypoints
+
+        # for conversion to coco api
+        area = torch.tensor([obj["area"] for obj in anno])
+        iscrowd = torch.tensor([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in anno])
+        target["area"] = area[keep]
+        target["iscrowd"] = iscrowd[keep]
+
+        target["orig_size"] = torch.as_tensor([int(h), int(w)])
+        target["size"] = torch.as_tensor([int(h), int(w)])
+
+        return image, target
+
+
+def make_coco_transforms(image_set, fix_size=False, strong_aug=False, args=None):
+
+    normalize = T.Compose([
+        T.ToTensor(),
+        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+
+    # config the params for data aug
+    scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
+    max_size = 1333
+    scales2_resize = [400, 500, 600]
+    scales2_crop = [384, 600]
+    
+    # update args from config files
+    scales = getattr(args, 'data_aug_scales', scales)
+    max_size = getattr(args, 'data_aug_max_size', max_size)
+    scales2_resize = getattr(args, 'data_aug_scales2_resize', scales2_resize)
+    scales2_crop = getattr(args, 'data_aug_scales2_crop', scales2_crop)
+
+    # resize them
+    data_aug_scale_overlap = getattr(args, 'data_aug_scale_overlap', None)
+    if data_aug_scale_overlap is not None and data_aug_scale_overlap > 0:
+        data_aug_scale_overlap = float(data_aug_scale_overlap)
+        scales = [int(i*data_aug_scale_overlap) for i in scales]
+        max_size = int(max_size*data_aug_scale_overlap)
+        scales2_resize = [int(i*data_aug_scale_overlap) for i in scales2_resize]
+        scales2_crop = [int(i*data_aug_scale_overlap) for i in scales2_crop]
+
+    datadict_for_print = {
+        'scales': scales,
+        'max_size': max_size,
+        'scales2_resize': scales2_resize,
+        'scales2_crop': scales2_crop
+    }
+    # print("data_aug_params:", json.dumps(datadict_for_print, indent=2))
+
+    if image_set == 'train':
+        if fix_size:
+            return T.Compose([
+                T.RandomHorizontalFlip(),
+                T.RandomResize([(max_size, max(scales))]),
+                # T.RandomResize([(512, 512)]),
+                normalize,
+            ])
+
+        if strong_aug:
+            import datasets.sltransform as SLT
+            
+            return T.Compose([
+                T.RandomHorizontalFlip(),
+                T.RandomSelect(
+                    T.RandomResize(scales, max_size=max_size),
+                    T.Compose([
+                        T.RandomResize(scales2_resize),
+                        T.RandomSizeCrop(*scales2_crop),
+                        T.RandomResize(scales, max_size=max_size),
+                    ])
+                ),
+                SLT.RandomSelectMulti([
+                    SLT.RandomCrop(),
+                    SLT.LightingNoise(),
+                    SLT.AdjustBrightness(2),
+                    SLT.AdjustContrast(2),
+                ]),
+                normalize,
+            ])
+        
+        return T.Compose([
+            T.RandomHorizontalFlip(),
+            T.RandomSelect(
+                T.RandomResize(scales, max_size=max_size),
+                T.Compose([
+                    T.RandomResize(scales2_resize),
+                    T.RandomSizeCrop(*scales2_crop),
+                    T.RandomResize(scales, max_size=max_size),
+                ])
+            ),
+            normalize,
+        ])
+
+    if image_set in ['val', 'eval_debug', 'train_reg', 'test']:
+
+        if os.environ.get("GFLOPS_DEBUG_SHILONG", False) == 'INFO':
+            print("Under debug mode for flops calculation only!!!!!!!!!!!!!!!!")
+            return T.Compose([
+                T.ResizeDebug((1280, 800)),
+                normalize,
+            ])   
+
+        return T.Compose([
+            T.RandomResize([max(scales)], max_size=max_size),
+            normalize,
+        ])
+
+
+
+    raise ValueError(f'unknown {image_set}')
+
+
+def get_aux_target_hacks_list(image_set, args):
+    if args.modelname in ['q2bs_mask', 'q2bs']:
+        aux_target_hacks_list = [
+            label2compat(), 
+            label_compat2onehot(), 
+            RandomSelectBoxes(num_class=args.num_classes)
+        ]
+        if args.masked_data and image_set == 'train':
+            # aux_target_hacks_list.append()
+            aux_target_hacks_list.append(MaskCrop())
+    elif args.modelname in ['q2bm_v2', 'q2bs_ce', 'q2op', 'q2ofocal', 'q2opclip', 'q2ocqonly']:
+        aux_target_hacks_list = [
+            label2compat(),
+            label_compat2onehot(),
+            box_label_catter(),
+            RandomSelectBoxlabels(num_classes=args.num_classes,
+                                    prob_first_item=args.prob_first_item,
+                                    prob_random_item=args.prob_random_item,
+                                    prob_last_item=args.prob_last_item,
+                                    prob_stop_sign=args.prob_stop_sign,
+                                    ),
+            BboxPertuber(max_ratio=0.02, generate_samples=1000),
+        ]
+    elif args.modelname in ['q2omask', 'q2osa']:
+        if args.coco_aug:
+            aux_target_hacks_list = [
+                label2compat(),
+                label_compat2onehot(),
+                box_label_catter(),
+                RandomSelectBoxlabels(num_classes=args.num_classes,
+                                        prob_first_item=args.prob_first_item,
+                                        prob_random_item=args.prob_random_item,
+                                        prob_last_item=args.prob_last_item,
+                                        prob_stop_sign=args.prob_stop_sign,
+                                        ),
+                RandomDrop(p=0.2),
+                BboxPertuber(max_ratio=0.02, generate_samples=1000),
+                RandomCutout(factor=0.5)
+            ]
+        else:
+            aux_target_hacks_list = [
+                label2compat(),
+                label_compat2onehot(),
+                box_label_catter(),
+                RandomSelectBoxlabels(num_classes=args.num_classes,
+                                        prob_first_item=args.prob_first_item,
+                                        prob_random_item=args.prob_random_item,
+                                        prob_last_item=args.prob_last_item,
+                                        prob_stop_sign=args.prob_stop_sign,
+                                        ),
+                BboxPertuber(max_ratio=0.02, generate_samples=1000),
+            ]
+    else:
+        aux_target_hacks_list = None
+
+    return aux_target_hacks_list
+
+
+def build(image_set, args, datasetinfo):
+    img_folder = datasetinfo["root"]
+    ann_file = datasetinfo["anno"]
+
+    # copy to local path
+    if os.environ.get('DATA_COPY_SHILONG') == 'INFO':
+        preparing_dataset(dict(img_folder=img_folder, ann_file=ann_file), image_set, args)
+
+    try:
+        strong_aug = args.strong_aug
+    except:
+        strong_aug = False
+    print(img_folder, ann_file)
+    dataset = CocoDetection(img_folder, ann_file, 
+            transforms=make_coco_transforms(image_set, fix_size=args.fix_size, strong_aug=strong_aug, args=args), 
+            return_masks=args.masks,
+            aux_target_hacks=None,
+        )
+    return dataset
+
+
+if __name__ == "__main__":
+    # Objects365 Val example
+    dataset_o365 = CocoDetection(
+            '/path/Objects365/train/',
+            "/path/Objects365/slannos/anno_preprocess_train_v2.json",
+            transforms=None,
+            return_masks=False,
+        )
+    print('len(dataset_o365):', len(dataset_o365))

datasets/coco_eval.py

@@ -0,0 +1,266 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+COCO evaluator that works in distributed mode.
+
+Mostly copy-paste from https://github.com/pytorch/vision/blob/edfd5a7/references/detection/coco_eval.py
+The difference is that there is less copy-pasting from pycocotools
+in the end of the file, as python3 can suppress prints with contextlib
+"""
+import os
+import contextlib
+import copy
+import numpy as np
+import torch
+
+from pycocotools.cocoeval import COCOeval
+from pycocotools.coco import COCO
+import pycocotools.mask as mask_util
+
+from util.misc import all_gather
+
+
+class CocoEvaluator(object):
+    def __init__(self, coco_gt, iou_types, useCats=True):
+        assert isinstance(iou_types, (list, tuple))
+        coco_gt = copy.deepcopy(coco_gt)
+        self.coco_gt = coco_gt
+
+        self.iou_types = iou_types
+        self.coco_eval = {}
+        for iou_type in iou_types:
+            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)
+            self.coco_eval[iou_type].useCats = useCats
+
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in iou_types}
+        self.useCats = useCats
+
+    def update(self, predictions):
+        img_ids = list(np.unique(list(predictions.keys())))
+        self.img_ids.extend(img_ids)
+
+        for iou_type in self.iou_types:
+            results = self.prepare(predictions, iou_type)
+
+            # suppress pycocotools prints
+            with open(os.devnull, 'w') as devnull:
+                with contextlib.redirect_stdout(devnull):
+                    coco_dt = COCO.loadRes(self.coco_gt, results) if results else COCO()
+            coco_eval = self.coco_eval[iou_type]
+
+            coco_eval.cocoDt = coco_dt
+            coco_eval.params.imgIds = list(img_ids)
+            coco_eval.params.useCats = self.useCats
+            img_ids, eval_imgs = evaluate(coco_eval)
+
+            self.eval_imgs[iou_type].append(eval_imgs)
+
+    def synchronize_between_processes(self):
+        for iou_type in self.iou_types:
+            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
+            create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])
+
+    def accumulate(self):
+        for coco_eval in self.coco_eval.values():
+            coco_eval.accumulate()
+
+    def summarize(self):
+        for iou_type, coco_eval in self.coco_eval.items():
+            print("IoU metric: {}".format(iou_type))
+            coco_eval.summarize()
+
+    def prepare(self, predictions, iou_type):
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        elif iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        elif iou_type == "keypoints":
+            return self.prepare_for_coco_keypoint(predictions)
+        else:
+            raise ValueError("Unknown iou type {}".format(iou_type))
+
+    def prepare_for_coco_detection(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            if not isinstance(prediction["scores"], list):
+                scores = prediction["scores"].tolist()
+            else:
+                scores = prediction["scores"]
+            if not isinstance(prediction["labels"], list):
+                labels = prediction["labels"].tolist()
+            else:
+                labels = prediction["labels"]
+
+        
+            try:
+                coco_results.extend(
+                    [
+                        {
+                            "image_id": original_id,
+                            "category_id": labels[k],
+                            "bbox": box,
+                            "score": scores[k],
+                        }
+                        for k, box in enumerate(boxes)
+                    ]
+                )
+            except:
+                import ipdb; ipdb.set_trace()
+        return coco_results
+
+    def prepare_for_coco_segmentation(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"]
+            labels = prediction["labels"]
+            masks = prediction["masks"]
+
+            masks = masks > 0.5
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            rles = [
+                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0]
+                for mask in masks
+            ]
+            for rle in rles:
+                rle["counts"] = rle["counts"].decode("utf-8")
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "segmentation": rle,
+                        "score": scores[k],
+                    }
+                    for k, rle in enumerate(rles)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_keypoint(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            keypoints = prediction["keypoints"]
+            keypoints = keypoints.flatten(start_dim=1).tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        'keypoints': keypoint,
+                        "score": scores[k],
+                    }
+                    for k, keypoint in enumerate(keypoints)
+                ]
+            )
+        return coco_results
+
+
+def convert_to_xywh(boxes):
+    xmin, ymin, xmax, ymax = boxes.unbind(1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
+
+
+def merge(img_ids, eval_imgs):
+    all_img_ids = all_gather(img_ids)
+    all_eval_imgs = all_gather(eval_imgs)
+
+    merged_img_ids = []
+    for p in all_img_ids:
+        merged_img_ids.extend(p)
+
+    merged_eval_imgs = []
+    for p in all_eval_imgs:
+        merged_eval_imgs.append(p)
+
+    merged_img_ids = np.array(merged_img_ids)
+    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
+
+    # keep only unique (and in sorted order) images
+    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
+    merged_eval_imgs = merged_eval_imgs[..., idx]
+
+    return merged_img_ids, merged_eval_imgs
+
+
+def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
+    img_ids, eval_imgs = merge(img_ids, eval_imgs)
+    img_ids = list(img_ids)
+    eval_imgs = list(eval_imgs.flatten())
+
+    coco_eval.evalImgs = eval_imgs
+    coco_eval.params.imgIds = img_ids
+    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
+
+
+#################################################################
+# From pycocotools, just removed the prints and fixed
+# a Python3 bug about unicode not defined
+#################################################################
+
+
+def evaluate(self):
+    '''
+    Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
+    :return: None
+    '''
+    p = self.params
+    # add backward compatibility if useSegm is specified in params
+    if p.useSegm is not None:
+        p.iouType = 'segm' if p.useSegm == 1 else 'bbox'
+        print('useSegm (deprecated) is not None. Running {} evaluation'.format(p.iouType))
+    p.imgIds = list(np.unique(p.imgIds))
+    if p.useCats:
+        p.catIds = list(np.unique(p.catIds))
+    p.maxDets = sorted(p.maxDets)
+    self.params = p
+
+    self._prepare()
+    # loop through images, area range, max detection number
+    catIds = p.catIds if p.useCats else [-1]
+
+    if p.iouType == 'segm' or p.iouType == 'bbox':
+        computeIoU = self.computeIoU
+    elif p.iouType == 'keypoints':
+        computeIoU = self.computeOks
+    self.ious = {
+        (imgId, catId): computeIoU(imgId, catId)
+        for imgId in p.imgIds
+        for catId in catIds}
+
+    evaluateImg = self.evaluateImg
+    maxDet = p.maxDets[-1]
+    evalImgs = [
+        evaluateImg(imgId, catId, areaRng, maxDet)
+        for catId in catIds
+        for areaRng in p.areaRng
+        for imgId in p.imgIds
+    ]
+    # this is NOT in the pycocotools code, but could be done outside
+    evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))
+    self._paramsEval = copy.deepcopy(self.params)
+
+    return p.imgIds, evalImgs
+
+#################################################################
+# end of straight copy from pycocotools, just removing the prints
+#################################################################

datasets/coco_panoptic.py

@@ -0,0 +1,99 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from panopticapi.utils import rgb2id
+from util.box_ops import masks_to_boxes
+
+from .coco import make_coco_transforms
+
+
+class CocoPanoptic:
+    def __init__(self, img_folder, ann_folder, ann_file, transforms=None, return_masks=True):
+        with open(ann_file, 'r') as f:
+            self.coco = json.load(f)
+
+        # sort 'images' field so that they are aligned with 'annotations'
+        # i.e., in alphabetical order
+        self.coco['images'] = sorted(self.coco['images'], key=lambda x: x['id'])
+        # sanity check
+        if "annotations" in self.coco:
+            for img, ann in zip(self.coco['images'], self.coco['annotations']):
+                assert img['file_name'][:-4] == ann['file_name'][:-4]
+
+        self.img_folder = img_folder
+        self.ann_folder = ann_folder
+        self.ann_file = ann_file
+        self.transforms = transforms
+        self.return_masks = return_masks
+
+    def __getitem__(self, idx):
+        ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
+        img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')
+        ann_path = Path(self.ann_folder) / ann_info['file_name']
+
+        img = Image.open(img_path).convert('RGB')
+        w, h = img.size
+        if "segments_info" in ann_info:
+            masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
+            masks = rgb2id(masks)
+
+            ids = np.array([ann['id'] for ann in ann_info['segments_info']])
+            masks = masks == ids[:, None, None]
+
+            masks = torch.as_tensor(masks, dtype=torch.uint8)
+            labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)
+
+        target = {}
+        target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
+        if self.return_masks:
+            target['masks'] = masks
+        target['labels'] = labels
+
+        target["boxes"] = masks_to_boxes(masks)
+
+        target['size'] = torch.as_tensor([int(h), int(w)])
+        target['orig_size'] = torch.as_tensor([int(h), int(w)])
+        if "segments_info" in ann_info:
+            for name in ['iscrowd', 'area']:
+                target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
+
+        if self.transforms is not None:
+            img, target = self.transforms(img, target)
+
+        return img, target
+
+    def __len__(self):
+        return len(self.coco['images'])
+
+    def get_height_and_width(self, idx):
+        img_info = self.coco['images'][idx]
+        height = img_info['height']
+        width = img_info['width']
+        return height, width
+
+
+def build(image_set, args):
+    img_folder_root = Path(args.coco_path)
+    ann_folder_root = Path(args.coco_panoptic_path)
+    assert img_folder_root.exists(), f'provided COCO path {img_folder_root} does not exist'
+    assert ann_folder_root.exists(), f'provided COCO path {ann_folder_root} does not exist'
+    mode = 'panoptic'
+    PATHS = {
+        "train": ("train2017", Path("annotations") / f'{mode}_train2017.json'),
+        "val": ("val2017", Path("annotations") / f'{mode}_val2017.json'),
+    }
+
+    img_folder, ann_file = PATHS[image_set]
+    img_folder_path = img_folder_root / img_folder
+    ann_folder = ann_folder_root / f'{mode}_{img_folder}'
+    ann_file = ann_folder_root / ann_file
+
+    dataset = CocoPanoptic(img_folder_path, ann_folder, ann_file,
+                           transforms=make_coco_transforms(image_set), return_masks=args.masks)
+
+    return dataset

datasets/cocogrounding_eval.py

@@ -0,0 +1,271 @@
+# ------------------------------------------------------------------------
+# Grounding DINO. Midified by Shilong Liu.
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Copyright (c) Aishwarya Kamath & Nicolas Carion. Licensed under the Apache License 2.0. All Rights Reserved
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+COCO evaluator that works in distributed mode.
+
+Mostly copy-paste from https://github.com/pytorch/vision/blob/edfd5a7/references/detection/coco_eval.py
+The difference is that there is less copy-pasting from pycocotools
+in the end of the file, as python3 can suppress prints with contextlib
+"""
+import contextlib
+import copy
+import os
+
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+from groundingdino.util.misc import all_gather
+
+
+class CocoGroundingEvaluator(object):
+    def __init__(self, coco_gt, iou_types, useCats=True):
+        assert isinstance(iou_types, (list, tuple))
+        coco_gt = copy.deepcopy(coco_gt)
+        self.coco_gt = coco_gt
+
+        self.iou_types = iou_types
+        self.coco_eval = {}
+        for iou_type in iou_types:
+            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)
+            self.coco_eval[iou_type].useCats = useCats
+
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in iou_types}
+        self.useCats = useCats
+
+    def update(self, predictions):
+        img_ids = list(np.unique(list(predictions.keys())))
+        self.img_ids.extend(img_ids)
+        # import pdb;pdb.set_trace()
+        for iou_type in self.iou_types:
+            results = self.prepare(predictions, iou_type)
+
+            # suppress pycocotools prints
+            with open(os.devnull, "w") as devnull:
+                with contextlib.redirect_stdout(devnull):
+                    coco_dt = COCO.loadRes(self.coco_gt, results) if results else COCO()
+
+            coco_eval = self.coco_eval[iou_type]
+
+            coco_eval.cocoDt = coco_dt
+            coco_eval.params.imgIds = list(img_ids)
+            coco_eval.params.useCats = self.useCats
+            img_ids, eval_imgs = evaluate(coco_eval)
+
+            self.eval_imgs[iou_type].append(eval_imgs)
+
+    def synchronize_between_processes(self):
+        for iou_type in self.iou_types:
+            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
+            create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])
+
+    def accumulate(self):
+        for coco_eval in self.coco_eval.values():
+            coco_eval.accumulate()
+
+    def summarize(self):
+        for iou_type, coco_eval in self.coco_eval.items():
+            print("IoU metric: {}".format(iou_type))
+            coco_eval.summarize()
+
+    def prepare(self, predictions, iou_type):
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        elif iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        elif iou_type == "keypoints":
+            return self.prepare_for_coco_keypoint(predictions)
+        else:
+            raise ValueError("Unknown iou type {}".format(iou_type))
+
+    def prepare_for_coco_detection(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "bbox": box,
+                        "score": scores[k],
+                    }
+                    for k, box in enumerate(boxes)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_segmentation(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"]
+            labels = prediction["labels"]
+            masks = prediction["masks"]
+
+            masks = masks > 0.5
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            rles = [
+                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0] 
+                for mask in masks
+            ]
+            for rle in rles:
+                rle["counts"] = rle["counts"].decode("utf-8")
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "segmentation": rle,
+                        "score": scores[k],
+                    }
+                    for k, rle in enumerate(rles)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_keypoint(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            keypoints = prediction["keypoints"]
+            keypoints = keypoints.flatten(start_dim=1).tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "keypoints": keypoint,
+                        "score": scores[k],
+                    }
+                    for k, keypoint in enumerate(keypoints)
+                ]
+            )
+        return coco_results
+
+
+def convert_to_xywh(boxes):
+    xmin, ymin, xmax, ymax = boxes.unbind(1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
+
+
+def merge(img_ids, eval_imgs):
+    all_img_ids = all_gather(img_ids)
+    all_eval_imgs = all_gather(eval_imgs)
+
+    merged_img_ids = []
+    for p in all_img_ids:
+        merged_img_ids.extend(p)
+
+    merged_eval_imgs = []
+    for p in all_eval_imgs:
+        merged_eval_imgs.append(p)
+
+    merged_img_ids = np.array(merged_img_ids)
+    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
+
+    # keep only unique (and in sorted order) images
+    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
+    merged_eval_imgs = merged_eval_imgs[..., idx]
+
+    return merged_img_ids, merged_eval_imgs
+
+
+def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
+    img_ids, eval_imgs = merge(img_ids, eval_imgs)
+    img_ids = list(img_ids)
+    eval_imgs = list(eval_imgs.flatten())
+
+    coco_eval.evalImgs = eval_imgs
+    coco_eval.params.imgIds = img_ids
+    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
+
+
+#################################################################
+# From pycocotools, just removed the prints and fixed
+# a Python3 bug about unicode not defined
+#################################################################
+
+
+def evaluate(self):
+    """
+    Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
+    :return: None
+    """
+    # tic = time.time()
+    # print('Running per image evaluation...')
+
+    # import pdb;pdb.set_trace()
+    p = self.params
+    # add backward compatibility if useSegm is specified in params
+    if p.useSegm is not None:
+        p.iouType = "segm" if p.useSegm == 1 else "bbox"
+        print("useSegm (deprecated) is not None. Running {} evaluation".format(p.iouType))
+    # print('Evaluate annotation type *{}*'.format(p.iouType))
+    p.imgIds = list(np.unique(p.imgIds))
+    if p.useCats:
+        p.catIds = list(np.unique(p.catIds))
+    p.maxDets = sorted(p.maxDets)
+    self.params = p
+
+    self._prepare()
+    # loop through images, area range, max detection number
+    catIds = p.catIds if p.useCats else [-1]
+
+    if p.iouType == "segm" or p.iouType == "bbox":
+        computeIoU = self.computeIoU
+    elif p.iouType == "keypoints":
+        computeIoU = self.computeOks
+    self.ious = {
+        (imgId, catId): computeIoU(imgId, catId) 
+        for imgId in p.imgIds 
+        for catId in catIds}
+
+    evaluateImg = self.evaluateImg
+    maxDet = p.maxDets[-1]
+    evalImgs = [
+        evaluateImg(imgId, catId, areaRng, maxDet) 
+        for catId in catIds 
+        for areaRng in p.areaRng 
+        for imgId in p.imgIds
+    ]
+    # this is NOT in the pycocotools code, but could be done outside
+    evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))
+    self._paramsEval = copy.deepcopy(self.params)
+    # toc = time.time()
+    # print('DONE (t={:0.2f}s).'.format(toc-tic))
+    return p.imgIds, evalImgs
+
+
+#################################################################
+# end of straight copy from pycocotools, just removing the prints
+#################################################################

datasets/data_util.py

@@ -0,0 +1,170 @@
+import os
+import os.path as osp
+import shutil
+import time
+import datetime
+
+import torch
+
+from util.slconfig import SLConfig
+
+class Error(OSError):
+    pass
+
+def slcopytree(src, dst, symlinks=False, ignore=None, copy_function=shutil.copyfile,
+             ignore_dangling_symlinks=False):
+    """
+    modified from shutil.copytree without copystat.
+    
+    Recursively copy a directory tree.
+
+    The destination directory must not already exist.
+    If exception(s) occur, an Error is raised with a list of reasons.
+
+    If the optional symlinks flag is true, symbolic links in the
+    source tree result in symbolic links in the destination tree; if
+    it is false, the contents of the files pointed to by symbolic
+    links are copied. If the file pointed by the symlink doesn't
+    exist, an exception will be added in the list of errors raised in
+    an Error exception at the end of the copy process.
+
+    You can set the optional ignore_dangling_symlinks flag to true if you
+    want to silence this exception. Notice that this has no effect on
+    platforms that don't support os.symlink.
+
+    The optional ignore argument is a callable. If given, it
+    is called with the `src` parameter, which is the directory
+    being visited by copytree(), and `names` which is the list of
+    `src` contents, as returned by os.listdir():
+
+        callable(src, names) -> ignored_names
+
+    Since copytree() is called recursively, the callable will be
+    called once for each directory that is copied. It returns a
+    list of names relative to the `src` directory that should
+    not be copied.
+
+    The optional copy_function argument is a callable that will be used
+    to copy each file. It will be called with the source path and the
+    destination path as arguments. By default, copy2() is used, but any
+    function that supports the same signature (like copy()) can be used.
+
+    """
+    errors = []
+    if os.path.isdir(src):
+        names = os.listdir(src)
+        if ignore is not None:
+            ignored_names = ignore(src, names)
+        else:
+            ignored_names = set()
+
+        os.makedirs(dst)
+        for name in names:
+            if name in ignored_names:
+                continue
+            srcname = os.path.join(src, name)
+            dstname = os.path.join(dst, name)
+            try:
+                if os.path.islink(srcname):
+                    linkto = os.readlink(srcname)
+                    if symlinks:
+                        # We can't just leave it to `copy_function` because legacy
+                        # code with a custom `copy_function` may rely on copytree
+                        # doing the right thing.
+                        os.symlink(linkto, dstname)
+                    else:
+                        # ignore dangling symlink if the flag is on
+                        if not os.path.exists(linkto) and ignore_dangling_symlinks:
+                            continue
+                        # otherwise let the copy occurs. copy2 will raise an error
+                        if os.path.isdir(srcname):
+                            slcopytree(srcname, dstname, symlinks, ignore,
+                                    copy_function)
+                        else:
+                            copy_function(srcname, dstname)
+                elif os.path.isdir(srcname):
+                    slcopytree(srcname, dstname, symlinks, ignore, copy_function)
+                else:
+                    # Will raise a SpecialFileError for unsupported file types
+                    copy_function(srcname, dstname)
+            # catch the Error from the recursive copytree so that we can
+            # continue with other files
+            except Error as err:
+                errors.extend(err.args[0])
+            except OSError as why:
+                errors.append((srcname, dstname, str(why)))
+    else:
+        copy_function(src, dst)
+
+    if errors:
+        raise Error(errors)
+    return dst
+
+def check_and_copy(src_path, tgt_path):
+    if os.path.exists(tgt_path):
+        return None
+
+    return slcopytree(src_path, tgt_path)
+
+
+def remove(srcpath):
+    if os.path.isdir(srcpath):
+        return shutil.rmtree(srcpath)
+    else:
+        return os.remove(srcpath)  
+
+
+def preparing_dataset(pathdict, image_set, args):
+    start_time = time.time()
+    dataset_file = args.dataset_file
+    data_static_info = SLConfig.fromfile('util/static_data_path.py')
+    static_dict = data_static_info[dataset_file][image_set]
+
+    copyfilelist = []
+    for k,tgt_v in pathdict.items():
+        if os.path.exists(tgt_v):
+            if args.local_rank == 0:
+                print("path <{}> exist. remove it!".format(tgt_v))
+                remove(tgt_v)
+            # continue
+        
+        if args.local_rank == 0:
+            src_v = static_dict[k]
+            assert isinstance(src_v, str)
+            if src_v.endswith('.zip'):
+                # copy
+                cp_tgt_dir = os.path.dirname(tgt_v)
+                filename = os.path.basename(src_v)
+                cp_tgt_path = os.path.join(cp_tgt_dir, filename)
+                print('Copy from <{}> to <{}>.'.format(src_v, cp_tgt_path))
+                os.makedirs(cp_tgt_dir, exist_ok=True)
+                check_and_copy(src_v, cp_tgt_path)          
+
+                # unzip
+                import zipfile
+                print("Starting unzip <{}>".format(cp_tgt_path))
+                with zipfile.ZipFile(cp_tgt_path, 'r') as zip_ref:
+                    zip_ref.extractall(os.path.dirname(cp_tgt_path))      
+
+                copyfilelist.append(cp_tgt_path)
+                copyfilelist.append(tgt_v)
+            else:
+                print('Copy from <{}> to <{}>.'.format(src_v, tgt_v))
+                os.makedirs(os.path.dirname(tgt_v), exist_ok=True)
+                check_and_copy(src_v, tgt_v)
+                copyfilelist.append(tgt_v)
+    
+    if len(copyfilelist) == 0:
+        copyfilelist = None
+    args.copyfilelist = copyfilelist
+        
+    if args.distributed:
+        torch.distributed.barrier()
+    total_time = time.time() - start_time
+    if copyfilelist:
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('Data copy time {}'.format(total_time_str))
+    return copyfilelist
+
+
+    

datasets/dataset.py

@@ -0,0 +1,44 @@
+from __future__ import print_function
+
+import torch
+import torchvision.datasets as datasets
+from torch.utils.data import Dataset
+from PIL import Image
+from .tsv_io import TSVFile
+import numpy as np
+import base64
+import io
+
+
+class TSVDataset(Dataset):
+    """ TSV dataset for ImageNet 1K training
+    """    
+    def __init__(self, tsv_file, transform=None, target_transform=None):
+        self.tsv = TSVFile(tsv_file)
+        self.transform = transform
+        self.target_transform = target_transform
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is class_index of the target class.
+        """
+        row = self.tsv.seek(index)
+        image_data = base64.b64decode(row[-1])
+        image = Image.open(io.BytesIO(image_data))
+        image = image.convert('RGB')
+        target = int(row[1])
+
+        if self.transform is not None:
+            img = self.transform(image)
+        else:
+            img = image
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+    def __len__(self):
+        return self.tsv.num_rows()

datasets/odvg.py

@@ -0,0 +1,263 @@
+from torchvision.datasets.vision import VisionDataset
+import os.path
+from typing import Callable, Optional
+import json
+from PIL import Image
+import torch
+import random
+import os, sys
+sys.path.append(os.path.dirname(sys.path[0]))
+
+import datasets.transforms as T
+
+class ODVGDataset(VisionDataset):
+    """
+    Args:
+        root (string): Root directory where images are downloaded to.
+        anno (string): Path to json annotation file.
+        label_map_anno (string):  Path to json label mapping file. Only for Object Detection
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.PILToTensor``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        transforms (callable, optional): A function/transform that takes input sample and its target as entry
+            and returns a transformed version.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        anno: str,
+        label_map_anno: str = None,
+        max_labels: int = 80,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        transforms: Optional[Callable] = None,
+    ) -> None:
+        super().__init__(root, transforms, transform, target_transform)
+        self.root = root
+        self.dataset_mode = "OD" if label_map_anno else "VG"
+        self.max_labels = max_labels
+        if self.dataset_mode == "OD":
+            self.load_label_map(label_map_anno)
+        self._load_metas(anno)
+        self.get_dataset_info()
+
+    def load_label_map(self, label_map_anno):
+        with open(label_map_anno, 'r') as file:
+            self.label_map = json.load(file)
+        self.label_index = set(self.label_map.keys())
+
+    def _load_metas(self, anno):
+        with  open(anno, 'r')as f:
+            self.metas = [json.loads(line) for line in f]
+
+    def get_dataset_info(self):
+        print(f"  == total images: {len(self)}")
+        if self.dataset_mode == "OD":
+            print(f"  == total labels: {len(self.label_map)}")
+
+    def __getitem__(self, index: int):
+        meta = self.metas[index]
+        rel_path = meta["filename"]
+        abs_path = os.path.join(self.root, rel_path)
+        if not os.path.exists(abs_path):
+            raise FileNotFoundError(f"{abs_path} not found.")
+        image = Image.open(abs_path).convert('RGB')
+        
+        w, h = image.size
+        if self.dataset_mode == "OD":
+            # Process the dots for the instances.
+            anno = meta["detection"]
+            instances = [obj for obj in anno["instances"]]
+            exemplar_labels = [str(obj["label"]) for obj in meta["exemplars"]]
+            boxes = [obj["bbox"] for obj in instances]
+            # generate vg_labels
+            # pos bbox labels
+            ori_classes = [str(obj["label"]) for obj in instances] + exemplar_labels
+            pos_labels = set(ori_classes)
+            # neg bbox labels 
+            neg_labels = self.label_index.difference(pos_labels)
+             
+            vg_labels = list(pos_labels)
+            num_to_add = min(len(neg_labels), self.max_labels-len(pos_labels))
+            if num_to_add > 0:
+                vg_labels.extend(random.sample(neg_labels, num_to_add))
+            
+            # shuffle
+            for i in range(len(vg_labels)-1, 0, -1):
+                j = random.randint(0, i)
+                vg_labels[i], vg_labels[j] = vg_labels[j], vg_labels[i]
+
+            caption_list = [self.label_map[lb] for lb in vg_labels]
+            caption_dict = {item:index for index, item in enumerate(caption_list)}
+
+            caption = ' . '.join(caption_list) + ' .'
+            classes = [caption_dict[self.label_map[str(obj["label"])]] for obj in instances]
+            boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+            classes = torch.tensor(classes, dtype=torch.int64)
+
+            # Process the exemplars for the instances.
+            exemplars = [obj for obj in meta["exemplars"]]
+            exemplar_boxes = [obj["bbox"] for obj in exemplars]
+            exemplar_boxes = torch.as_tensor(exemplar_boxes, dtype=torch.float32).reshape(-1, 4)
+            exemplar_classes = [caption_dict[self.label_map[str(obj["label"])]] for obj in exemplars]
+            exemplar_classes = torch.tensor(exemplar_classes, dtype=torch.int64)
+        elif self.dataset_mode == "VG":
+            anno = meta["grounding"]
+            instances = [obj for obj in anno["regions"]]
+            boxes = [obj["bbox"] for obj in instances]
+            caption_list = [obj["phrase"] for obj in instances]
+            c = list(zip(boxes, caption_list))
+            random.shuffle(c)
+            boxes[:], caption_list[:] = zip(*c)
+            uni_caption_list  = list(set(caption_list))
+            label_map = {}
+            for idx in range(len(uni_caption_list)):
+                label_map[uni_caption_list[idx]] = idx
+            classes = [label_map[cap] for cap in caption_list]
+            caption = ' . '.join(uni_caption_list) + ' .'
+            boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+            classes = torch.tensor(classes, dtype=torch.int64)
+            caption_list = uni_caption_list
+        target = {}
+        target["size"] = torch.as_tensor([int(h), int(w)])
+        target["cap_list"] = caption_list
+        target["caption"] = caption
+        target["boxes"] = boxes
+        target["labels"] = classes
+        target["exemplar_boxes"] = exemplar_boxes
+        target["exemplar_labels"] = exemplar_classes
+        
+        if self.transforms is not None:
+            image, target = self.transforms(image, target)
+
+        return image, target
+    
+
+    def __len__(self) -> int:
+        return len(self.metas)
+
+
+def make_coco_transforms(image_set, fix_size=False, strong_aug=False, args=None):
+
+    normalize = T.Compose([
+        T.ToTensor(),
+        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+
+    # config the params for data aug
+    scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
+    max_size = 1333
+    scales2_resize = [400, 500, 600]
+    scales2_crop = [384, 600]
+    
+    # update args from config files
+    scales = getattr(args, 'data_aug_scales', scales)
+    max_size = getattr(args, 'data_aug_max_size', max_size)
+    scales2_resize = getattr(args, 'data_aug_scales2_resize', scales2_resize)
+    scales2_crop = getattr(args, 'data_aug_scales2_crop', scales2_crop)
+
+    # resize them
+    data_aug_scale_overlap = getattr(args, 'data_aug_scale_overlap', None)
+    if data_aug_scale_overlap is not None and data_aug_scale_overlap > 0:
+        data_aug_scale_overlap = float(data_aug_scale_overlap)
+        scales = [int(i*data_aug_scale_overlap) for i in scales]
+        max_size = int(max_size*data_aug_scale_overlap)
+        scales2_resize = [int(i*data_aug_scale_overlap) for i in scales2_resize]
+        scales2_crop = [int(i*data_aug_scale_overlap) for i in scales2_crop]
+
+    # datadict_for_print = {
+    #     'scales': scales,
+    #     'max_size': max_size,
+    #     'scales2_resize': scales2_resize,
+    #     'scales2_crop': scales2_crop
+    # }
+    # print("data_aug_params:", json.dumps(datadict_for_print, indent=2))
+
+    if image_set == 'train':
+        if fix_size:
+            return T.Compose([
+                T.RandomHorizontalFlip(),
+                T.RandomResize([(max_size, max(scales))]),
+                normalize,
+            ])
+
+        if strong_aug:
+            import datasets.sltransform as SLT
+            
+            return T.Compose([
+                T.RandomHorizontalFlip(),
+                T.RandomSelect(
+                    T.RandomResize(scales, max_size=max_size),
+                    T.Compose([
+                        T.RandomResize(scales2_resize),
+                        T.RandomSizeCrop(*scales2_crop),
+                        T.RandomResize(scales, max_size=max_size),
+                    ])
+                ),
+                SLT.RandomSelectMulti([
+                    SLT.RandomCrop(),
+                    SLT.LightingNoise(),
+                    SLT.AdjustBrightness(2),
+                    SLT.AdjustContrast(2),
+                ]),
+                normalize,
+            ])
+        
+        return T.Compose([
+            T.RandomHorizontalFlip(),
+            T.RandomSelect(
+                T.RandomResize(scales, max_size=max_size),
+                T.Compose([
+                    T.RandomResize(scales2_resize),
+                    T.RandomSizeCrop(*scales2_crop),
+                    T.RandomResize(scales, max_size=max_size),
+                ])
+            ),
+            normalize,
+        ])
+
+    if image_set in ['val', 'eval_debug', 'train_reg', 'test']:
+
+        if os.environ.get("GFLOPS_DEBUG_SHILONG", False) == 'INFO':
+            print("Under debug mode for flops calculation only!!!!!!!!!!!!!!!!")
+            return T.Compose([
+                T.ResizeDebug((1280, 800)),
+                normalize,
+            ])   
+
+        return T.Compose([
+            T.RandomResize([max(scales)], max_size=max_size),
+            normalize,
+        ])
+
+    raise ValueError(f'unknown {image_set}')
+
+def build_odvg(image_set, args, datasetinfo):
+    img_folder = datasetinfo["root"]
+    ann_file = datasetinfo["anno"]
+    label_map = datasetinfo["label_map"] if "label_map" in datasetinfo else None
+    try:
+        strong_aug = args.strong_aug
+    except:
+        strong_aug = False
+    print(img_folder, ann_file, label_map)
+    dataset = ODVGDataset(img_folder, ann_file, label_map, max_labels=args.max_labels,
+            transforms=make_coco_transforms(image_set, fix_size=args.fix_size, strong_aug=strong_aug, args=args), 
+    )
+    return dataset
+
+
+if __name__=="__main__":
+    dataset_vg = ODVGDataset("path/GRIT-20M/data/","path/GRIT-20M/anno/grit_odvg_10k.jsonl",)
+    print(len(dataset_vg))
+    data = dataset_vg[random.randint(0, 100)] 
+    print(data)
+    dataset_od = ODVGDataset("pathl/V3Det/",
+        "path/V3Det/annotations/v3det_2023_v1_all_odvg.jsonl",
+        "path/V3Det/annotations/v3det_label_map.json",
+    )
+    print(len(dataset_od))
+    data = dataset_od[random.randint(0, 100)] 
+    print(data)

datasets/panoptic_eval.py

@@ -0,0 +1,44 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import json
+import os
+
+import util.misc as utils
+
+try:
+    from panopticapi.evaluation import pq_compute
+except ImportError:
+    pass
+
+
+class PanopticEvaluator(object):
+    def __init__(self, ann_file, ann_folder, output_dir="panoptic_eval"):
+        self.gt_json = ann_file
+        self.gt_folder = ann_folder
+        if utils.is_main_process():
+            if not os.path.exists(output_dir):
+                os.mkdir(output_dir)
+        self.output_dir = output_dir
+        self.predictions = []
+
+    def update(self, predictions):
+        for p in predictions:
+            with open(os.path.join(self.output_dir, p["file_name"]), "wb") as f:
+                f.write(p.pop("png_string"))
+
+        self.predictions += predictions
+
+    def synchronize_between_processes(self):
+        all_predictions = utils.all_gather(self.predictions)
+        merged_predictions = []
+        for p in all_predictions:
+            merged_predictions += p
+        self.predictions = merged_predictions
+
+    def summarize(self):
+        if utils.is_main_process():
+            json_data = {"annotations": self.predictions}
+            predictions_json = os.path.join(self.output_dir, "predictions.json")
+            with open(predictions_json, "w") as f:
+                f.write(json.dumps(json_data))
+            return pq_compute(self.gt_json, predictions_json, gt_folder=self.gt_folder, pred_folder=self.output_dir)
+        return None

datasets/random_crop.py

@@ -0,0 +1,135 @@
+import PIL #version 1.2.0
+import torch
+import os
+import torchvision.transforms.functional as F
+import numpy as np
+import random
+
+
+def intersect(boxes1, boxes2):
+    '''
+        Find intersection of every box combination between two sets of box
+        boxes1: bounding boxes 1, a tensor of dimensions (n1, 4)
+        boxes2: bounding boxes 2, a tensor of dimensions (n2, 4)
+        
+        Out: Intersection each of boxes1 with respect to each of boxes2, 
+             a tensor of dimensions (n1, n2)
+    '''
+    n1 = boxes1.size(0)
+    n2 = boxes2.size(0)
+    max_xy =  torch.min(boxes1[:, 2:].unsqueeze(1).expand(n1, n2, 2),
+                        boxes2[:, 2:].unsqueeze(0).expand(n1, n2, 2))
+    
+    min_xy = torch.max(boxes1[:, :2].unsqueeze(1).expand(n1, n2, 2),
+                       boxes2[:, :2].unsqueeze(0).expand(n1, n2, 2))
+    inter = torch.clamp(max_xy - min_xy , min=0)  # (n1, n2, 2)
+    return inter[:, :, 0] * inter[:, :, 1]  #(n1, n2)
+def find_IoU(boxes1, boxes2):
+    '''
+        Find IoU between every boxes set of boxes 
+        boxes1: a tensor of dimensions (n1, 4) (left, top, right , bottom)
+        boxes2: a tensor of dimensions (n2, 4)
+        
+        Out: IoU each of boxes1 with respect to each of boxes2, a tensor of 
+             dimensions (n1, n2)
+        
+        Formula: 
+        (box1 ∩ box2) / (box1 u box2) = (box1 ∩ box2) / (area(box1) + area(box2) - (box1 ∩ box2 ))
+    '''
+    inter = intersect(boxes1, boxes2)
+    area_boxes1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+    area_boxes2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+    
+    area_boxes1 = area_boxes1.unsqueeze(1).expand_as(inter) #(n1, n2)
+    area_boxes2 = area_boxes2.unsqueeze(0).expand_as(inter)  #(n1, n2)
+    union = (area_boxes1 + area_boxes2 - inter)
+    return inter / union
+
+
+def random_crop(image, boxes, labels, difficulties=None):
+    '''
+        image: A PIL image
+        boxes: Bounding boxes, a tensor of dimensions (#objects, 4)
+        labels: labels of object, a tensor of dimensions (#objects)
+        difficulties: difficulties of detect object, a tensor of dimensions (#objects)
+        
+        Out: cropped image , new boxes, new labels, new difficulties
+    '''
+    if type(image) == PIL.Image.Image:
+        image = F.to_tensor(image)
+    original_h = image.size(1)
+    original_w = image.size(2)
+    
+    while True:
+        mode = random.choice([0.1, 0.3, 0.5, 0.9, None])
+        
+        if mode is None:
+            return F.to_pil_image(image), boxes, labels, difficulties
+        
+        new_image = image
+        new_boxes = boxes
+        new_difficulties = difficulties
+        new_labels = labels
+        for _ in range(50):
+            # Crop dimensions: [0.3, 1] of original dimensions
+            new_h = random.uniform(0.3*original_h, original_h)
+            new_w = random.uniform(0.3*original_w, original_w)
+            
+            # Aspect ratio constraint b/t .5 & 2
+            if new_h/new_w < 0.5 or new_h/new_w > 2:
+                continue
+            
+            #Crop coordinate
+            left = random.uniform(0, original_w - new_w)
+            right = left + new_w
+            top = random.uniform(0, original_h - new_h)
+            bottom = top + new_h
+            crop = torch.FloatTensor([int(left), int(top), int(right), int(bottom)])
+            
+            # Calculate IoU  between the crop and the bounding boxes
+            overlap = find_IoU(crop.unsqueeze(0), boxes) #(1, #objects)
+            overlap = overlap.squeeze(0)
+
+            # If not a single bounding box has a IoU of greater than the minimum, try again
+            if overlap.shape[0] == 0:
+                continue
+            if overlap.max().item() < mode:
+                continue
+            
+            #Crop
+            new_image = image[:, int(top):int(bottom), int(left):int(right)] #(3, new_h, new_w)
+            
+            #Center of bounding boxes
+            center_bb = (boxes[:, :2] + boxes[:, 2:])/2.0
+            
+            #Find bounding box has been had center in crop
+            center_in_crop = (center_bb[:, 0] >left) * (center_bb[:, 0] < right
+                             ) *(center_bb[:, 1] > top) * (center_bb[:, 1] < bottom)    #( #objects)
+            
+            if not center_in_crop.any():
+                continue
+            
+            #take matching bounding box
+            new_boxes = boxes[center_in_crop, :]
+            
+            #take matching labels
+            new_labels = labels[center_in_crop]
+            
+            #take matching difficulities
+            if difficulties is not None:
+                new_difficulties = difficulties[center_in_crop]
+            else:
+                new_difficulties = None
+            
+            #Use the box left and top corner or the crop's
+            new_boxes[:, :2] = torch.max(new_boxes[:, :2], crop[:2])
+            
+            #adjust to crop
+            new_boxes[:, :2] -= crop[:2]
+            
+            new_boxes[:, 2:] = torch.min(new_boxes[:, 2:],crop[2:])
+            
+            #adjust to crop
+            new_boxes[:, 2:] -= crop[:2]
+            
+            return F.to_pil_image(new_image), new_boxes, new_labels, new_difficulties

datasets/sltransform.py

@@ -0,0 +1,247 @@
+# modified from https://github.com/anhtuan85/Data-Augmentation-for-Object-Detection/blob/master/augmentation.ipynb
+
+import PIL #version 1.2.0
+from PIL import Image #version 6.1.0
+import torch
+import os
+import torchvision.transforms.functional as F
+import numpy as np
+import random
+
+from .random_crop import random_crop
+from util.box_ops import box_cxcywh_to_xyxy, box_xyxy_to_cxcywh
+
+class AdjustContrast:
+    def __init__(self, contrast_factor):
+        self.contrast_factor = contrast_factor
+
+    def __call__(self, img, target):
+        """
+        img (PIL Image or Tensor): Image to be adjusted.
+        """
+        _contrast_factor = ((random.random() + 1.0) / 2.0) * self.contrast_factor
+        img = F.adjust_contrast(img, _contrast_factor)
+        return img, target
+
+class AdjustBrightness:
+    def __init__(self, brightness_factor):
+        self.brightness_factor = brightness_factor
+
+    def __call__(self, img, target):
+        """
+        img (PIL Image or Tensor): Image to be adjusted.
+        """
+        _brightness_factor = ((random.random() + 1.0) / 2.0) * self.brightness_factor
+        img = F.adjust_brightness(img, _brightness_factor)
+        return img, target
+
+def lighting_noise(image):
+    '''
+        color channel swap in image
+        image: A PIL image
+    '''
+    new_image = image
+    perms = ((0, 1, 2), (0, 2, 1), (1, 0, 2), 
+             (1, 2, 0), (2, 0, 1), (2, 1, 0))
+    swap = perms[random.randint(0, len(perms)- 1)]
+    new_image = F.to_tensor(new_image)
+    new_image = new_image[swap, :, :]
+    new_image = F.to_pil_image(new_image)
+    return new_image
+
+class LightingNoise:
+    def __init__(self) -> None:
+        pass
+
+    def __call__(self, img, target):
+        return lighting_noise(img), target
+
+
+def rotate(image, boxes, angle):
+    '''
+        Rotate image and bounding box
+        image: A Pil image (w, h)
+        boxes: A tensors of dimensions (#objects, 4)
+        
+        Out: rotated image (w, h), rotated boxes
+    '''
+    new_image = image.copy()
+    new_boxes = boxes.clone()
+    
+    #Rotate image, expand = True
+    w = image.width
+    h = image.height
+    cx = w/2
+    cy = h/2
+    new_image = new_image.rotate(angle, expand=True)
+    angle = np.radians(angle)
+    alpha = np.cos(angle)
+    beta = np.sin(angle)
+    #Get affine matrix
+    AffineMatrix = torch.tensor([[alpha, beta, (1-alpha)*cx - beta*cy],
+                                 [-beta, alpha, beta*cx + (1-alpha)*cy]])
+    
+    #Rotation boxes
+    box_width = (boxes[:,2] - boxes[:,0]).reshape(-1,1)
+    box_height = (boxes[:,3] - boxes[:,1]).reshape(-1,1)
+    
+    #Get corners for boxes
+    x1 = boxes[:,0].reshape(-1,1)
+    y1 = boxes[:,1].reshape(-1,1)
+    
+    x2 = x1 + box_width
+    y2 = y1 
+    
+    x3 = x1
+    y3 = y1 + box_height
+    
+    x4 = boxes[:,2].reshape(-1,1)
+    y4 = boxes[:,3].reshape(-1,1)
+    
+    corners = torch.stack((x1,y1,x2,y2,x3,y3,x4,y4), dim= 1)
+    # corners.reshape(-1, 8)    #Tensors of dimensions (#objects, 8)
+    corners = corners.reshape(-1,2) #Tensors of dimension (4* #objects, 2)
+    corners = torch.cat((corners, torch.ones(corners.shape[0], 1)), dim= 1) #(Tensors of dimension (4* #objects, 3))
+    
+    cos = np.abs(AffineMatrix[0, 0])
+    sin = np.abs(AffineMatrix[0, 1])
+    
+    nW = int((h * sin) + (w * cos))
+    nH = int((h * cos) + (w * sin))
+    AffineMatrix[0, 2] += (nW / 2) - cx
+    AffineMatrix[1, 2] += (nH / 2) - cy
+    
+
+    #Apply affine transform
+    rotate_corners = torch.mm(AffineMatrix, corners.t().to(torch.float64)).t()
+    rotate_corners = rotate_corners.reshape(-1,8)
+    
+    x_corners = rotate_corners[:,[0,2,4,6]]
+    y_corners = rotate_corners[:,[1,3,5,7]]
+    
+    #Get (x_min, y_min, x_max, y_max)
+    x_min, _ = torch.min(x_corners, dim= 1)
+    x_min = x_min.reshape(-1, 1)
+    y_min, _ = torch.min(y_corners, dim= 1)
+    y_min = y_min.reshape(-1, 1)
+    x_max, _ = torch.max(x_corners, dim= 1)
+    x_max = x_max.reshape(-1, 1)
+    y_max, _ = torch.max(y_corners, dim= 1)
+    y_max = y_max.reshape(-1, 1)
+    
+    new_boxes = torch.cat((x_min, y_min, x_max, y_max), dim= 1)
+    
+    scale_x = new_image.width / w
+    scale_y = new_image.height / h
+    
+    #Resize new image to (w, h)
+
+    new_image = new_image.resize((w, h))
+    
+    #Resize boxes
+    new_boxes /= torch.Tensor([scale_x, scale_y, scale_x, scale_y])
+    new_boxes[:, 0] = torch.clamp(new_boxes[:, 0], 0, w)
+    new_boxes[:, 1] = torch.clamp(new_boxes[:, 1], 0, h)
+    new_boxes[:, 2] = torch.clamp(new_boxes[:, 2], 0, w)
+    new_boxes[:, 3] = torch.clamp(new_boxes[:, 3], 0, h)
+    return new_image, new_boxes
+
+# def convert_xywh_to_xyxy(boxes: torch.Tensor):
+#     _boxes = boxes.clone()
+#     box_xy = _boxes[:, :2]
+#     box_wh = _boxes[:, 2:]
+#     box_x1y1 = box_xy - box_wh/2 
+#     box_x2y2 = box_xy + box_wh/2
+#     box_xyxy = torch.cat((box_x1y1, box_x2y2), dim=-1)
+#     return box_xyxy
+
+class Rotate:
+    def __init__(self, angle=10) -> None:
+        self.angle = angle
+
+    def __call__(self, img, target):
+        w,h = img.size
+        whwh = torch.Tensor([w, h, w, h])
+        boxes_xyxy = box_cxcywh_to_xyxy(target['boxes']) * whwh
+        img, boxes_new = rotate(img, boxes_xyxy, self.angle)
+        target['boxes'] = box_xyxy_to_cxcywh(boxes_new).to(boxes_xyxy.dtype) / (whwh + 1e-3)
+        return img, target
+
+
+class RandomCrop:
+    def __init__(self) -> None:
+        pass
+
+    def __call__(self, img, target):
+        w,h = img.size
+        try:
+            boxes_xyxy = target['boxes']
+            labels = target['labels']
+            img, new_boxes, new_labels, _ = random_crop(img, boxes_xyxy, labels)
+            target['boxes'] = new_boxes
+            target['labels'] = new_labels
+        except Exception as e:
+            pass
+        return img, target
+
+
+class RandomCropDebug:
+    def __init__(self) -> None:
+        pass
+
+    def __call__(self, img, target):
+        boxes_xyxy = target['boxes'].clone()
+        labels = target['labels'].clone()
+        img, new_boxes, new_labels, _ = random_crop(img, boxes_xyxy, labels)
+        target['boxes'] = new_boxes
+        target['labels'] = new_labels
+
+
+        return img, target
+        
+class RandomSelectMulti(object):
+    """
+    Randomly selects between transforms1 and transforms2,
+    """
+    def __init__(self, transformslist, p=-1):
+        self.transformslist = transformslist
+        self.p = p
+        assert p == -1
+
+    def __call__(self, img, target):
+        if self.p == -1:
+            return random.choice(self.transformslist)(img, target)
+
+
+class Albumentations:
+    def __init__(self):
+        import albumentations as A
+        self.transform = A.Compose([
+            A.Blur(p=0.01),
+            A.MedianBlur(p=0.01),
+            A.ToGray(p=0.01),
+            A.CLAHE(p=0.01),
+            A.RandomBrightnessContrast(p=0.005),
+            A.RandomGamma(p=0.005),
+            A.ImageCompression(quality_lower=75, p=0.005)],
+            bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
+
+    def __call__(self, img, target, p=1.0):
+        """
+        Input:
+            target['boxes']: xyxy, unnormalized data.
+        
+        """
+        boxes_raw = target['boxes']
+        labels_raw = target['labels']
+        img_np = np.array(img)
+        if self.transform and random.random() < p:
+            new_res = self.transform(image=img_np, bboxes=boxes_raw, class_labels=labels_raw)  # transformed
+            boxes_new = torch.Tensor(new_res['bboxes']).to(boxes_raw.dtype).reshape_as(boxes_raw)
+            img_np = new_res['image']
+            labels_new = torch.Tensor(new_res['class_labels']).to(labels_raw.dtype)
+        img_new = Image.fromarray(img_np)
+        target['boxes'] = boxes_new
+        target['labels'] = labels_new
+        
+        return img_new, target

datasets/transforms.py

@@ -0,0 +1,324 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Transforms and data augmentation for both image + bbox.
+"""
+import random
+
+import PIL
+import torch
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
+
+from util.box_ops import box_xyxy_to_cxcywh
+from util.misc import interpolate
+
+
+def crop(image, target, region):
+    cropped_image = F.crop(image, *region)
+
+    target = target.copy()
+    i, j, h, w = region
+
+    # should we do something wrt the original size?
+    target["size"] = torch.tensor([h, w])
+
+    fields = ["labels", "area"]
+    exemplar_fields = ["exemplar_labels", "exemplar_area", "exemplar_boxes"]
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        max_size = torch.as_tensor([w, h], dtype=torch.float32)
+        cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
+        cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
+        cropped_boxes = cropped_boxes.clamp(min=0)
+        area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
+        target["boxes"] = cropped_boxes.reshape(-1, 4)
+        target["area"] = area
+        fields.append("boxes")
+    if "exemplar_boxes" in target:
+        exemplar_boxes = target["exemplar_boxes"]
+        max_size = torch.as_tensor([w, h], dtype=torch.float32)
+        cropped_exemplar_boxes = exemplar_boxes - torch.as_tensor([j, i, j, i])
+        cropped_exemplar_boxes = torch.min(cropped_exemplar_boxes.reshape(-1, 2, 2), max_size)
+        cropped_exemplar_boxes = cropped_exemplar_boxes.clamp(min=0)
+        exemplar_area = (cropped_exemplar_boxes[:, 1, :] - cropped_exemplar_boxes[:, 0, :]).prod(dim=1)
+        target["exemplar_boxes"] = cropped_exemplar_boxes.reshape(-1, 4)
+        target["exemplar_area"] = exemplar_area
+
+    if "masks" in target:
+        # FIXME should we update the area here if there are no boxes?
+        target['masks'] = target['masks'][:, i:i + h, j:j + w]
+        fields.append("masks")
+
+    # remove elements for which the boxes or masks that have zero area
+    if "boxes" in target or "masks" in target:
+        # favor boxes selection when defining which elements to keep
+        # this is compatible with previous implementation
+        if "boxes" in target:
+            cropped_boxes = target['boxes'].reshape(-1, 2, 2)
+            keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
+        else:
+            keep = target['masks'].flatten(1).any(1)
+
+        for field in fields:
+            target[field] = target[field][keep]
+
+    # remove exemplar elements for which the boxes have zero area
+    if "exemplar_boxes" in target:
+        cropped_exemplar_boxes = target["exemplar_boxes"].reshape(-1, 2, 2)
+        keep = torch.all(cropped_exemplar_boxes[:, 1, :] > cropped_exemplar_boxes[:, 0, :], dim=1)
+        for field in exemplar_fields:
+            target[field] = target[field][keep]
+
+    return cropped_image, target
+
+
+def hflip(image, target):
+    flipped_image = F.hflip(image)
+
+    w, h = image.size
+
+    target = target.copy()
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0])
+        target["boxes"] = boxes
+
+    if "exemplar_boxes" in target:
+        exemplar_boxes = target["exemplar_boxes"]
+        exemplar_boxes = exemplar_boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0])
+        target["exemplar_boxes"] = exemplar_boxes
+
+    if "masks" in target:
+        target['masks'] = target['masks'].flip(-1)
+
+    return flipped_image, target
+
+
+def resize(image, target, size, max_size=None):
+    # size can be min_size (scalar) or (w, h) tuple
+
+    def get_size_with_aspect_ratio(image_size, size, max_size=None):
+        w, h = image_size
+        if max_size is not None:
+            min_original_size = float(min((w, h)))
+            max_original_size = float(max((w, h)))
+            if max_original_size / min_original_size * size > max_size:
+                size = int(round(max_size * min_original_size / max_original_size))
+
+        if (w <= h and w == size) or (h <= w and h == size):
+            return (h, w)
+
+        if w < h:
+            ow = size
+            oh = int(size * h / w)
+        else:
+            oh = size
+            ow = int(size * w / h)
+
+        return (oh, ow)
+
+    def get_size(image_size, size, max_size=None):
+        if isinstance(size, (list, tuple)):
+            return size[::-1]
+        else:
+            return get_size_with_aspect_ratio(image_size, size, max_size)
+
+    size = get_size(image.size, size, max_size)
+    rescaled_image = F.resize(image, size)
+
+    if target is None:
+        return rescaled_image, None
+
+    ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))
+    ratio_width, ratio_height = ratios
+
+    target = target.copy()
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
+        target["boxes"] = scaled_boxes
+
+    if "area" in target:
+        area = target["area"]
+        scaled_area = area * (ratio_width * ratio_height)
+        target["area"] = scaled_area
+
+    if "exemplar_boxes" in target:
+        exemplar_boxes = target["exemplar_boxes"]
+        scaled_exemplar_boxes = exemplar_boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
+        target["exemplar_boxes"] = scaled_exemplar_boxes
+    if "exemplar_area" in target:
+        exemplar_area = target["exemplar_area"]
+        scaled_exemplar_area = exemplar_area * (ratio_width * ratio_height)
+        target["exemplar_area"] = scaled_exemplar_area
+
+    h, w = size
+    target["size"] = torch.tensor([h, w])
+
+    if "masks" in target:
+        target['masks'] = interpolate(
+            target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
+
+    return rescaled_image, target
+
+
+def pad(image, target, padding):
+    # assumes that we only pad on the bottom right corners
+    padded_image = F.pad(image, (0, 0, padding[0], padding[1]))
+    if target is None:
+        return padded_image, None
+    target = target.copy()
+    # should we do something wrt the original size?
+    target["size"] = torch.tensor(padded_image.size[::-1])
+    if "masks" in target:
+        target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[0], 0, padding[1]))
+    return padded_image, target
+
+
+class ResizeDebug(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        return resize(img, target, self.size)
+
+
+class RandomCrop(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        region = T.RandomCrop.get_params(img, self.size)
+        return crop(img, target, region)
+
+
+class RandomSizeCrop(object):
+    def __init__(self, min_size: int, max_size: int):
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, img: PIL.Image.Image, target: dict):
+        w = random.randint(self.min_size, min(img.width, self.max_size))
+        h = random.randint(self.min_size, min(img.height, self.max_size))
+        region = T.RandomCrop.get_params(img, [h, w])
+        return crop(img, target, region)
+
+
+class CenterCrop(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        image_width, image_height = img.size
+        crop_height, crop_width = self.size
+        crop_top = int(round((image_height - crop_height) / 2.))
+        crop_left = int(round((image_width - crop_width) / 2.))
+        return crop(img, target, (crop_top, crop_left, crop_height, crop_width))
+
+
+class RandomHorizontalFlip(object):
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, target):
+        if random.random() < self.p:
+            return hflip(img, target)
+        return img, target
+
+
+class RandomResize(object):
+    def __init__(self, sizes, max_size=None):
+        assert isinstance(sizes, (list, tuple))
+        self.sizes = sizes
+        self.max_size = max_size
+
+    def __call__(self, img, target=None):
+        size = random.choice(self.sizes)
+        return resize(img, target, size, self.max_size)
+
+
+class RandomPad(object):
+    def __init__(self, max_pad):
+        self.max_pad = max_pad
+
+    def __call__(self, img, target):
+        pad_x = random.randint(0, self.max_pad)
+        pad_y = random.randint(0, self.max_pad)
+        return pad(img, target, (pad_x, pad_y))
+
+
+class RandomSelect(object):
+    """
+    Randomly selects between transforms1 and transforms2,
+    with probability p for transforms1 and (1 - p) for transforms2
+    """
+    def __init__(self, transforms1, transforms2, p=0.5):
+        self.transforms1 = transforms1
+        self.transforms2 = transforms2
+        self.p = p
+
+    def __call__(self, img, target):
+        if random.random() < self.p:
+            return self.transforms1(img, target)
+        return self.transforms2(img, target)
+
+
+class ToTensor(object):
+    def __call__(self, img, target):
+        return F.to_tensor(img), target
+
+
+class RandomErasing(object):
+
+    def __init__(self, *args, **kwargs):
+        self.eraser = T.RandomErasing(*args, **kwargs)
+
+    def __call__(self, img, target):
+        return self.eraser(img), target
+
+
+class Normalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image, target=None):
+        image = F.normalize(image, mean=self.mean, std=self.std)
+        if target is None:
+            return image, None
+        target = target.copy()
+        h, w = image.shape[-2:]
+
+        if "boxes" in target:
+            boxes = target["boxes"]
+            boxes = box_xyxy_to_cxcywh(boxes)
+            boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
+            target["boxes"] = boxes
+
+        if "exemplar_boxes" in target:
+            exemplar_boxes = target["exemplar_boxes"]
+            exemplar_boxes = box_xyxy_to_cxcywh(exemplar_boxes)
+            exemplar_boxes = exemplar_boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
+            target["exemplar_boxes"] = exemplar_boxes
+        return image, target
+
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += "    {0}".format(t)
+        format_string += "\n)"
+        return format_string

datasets/transforms_app.py

@@ -0,0 +1,338 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Transforms and data augmentation for both image + bbox.
+"""
+import random
+
+import PIL
+import torch
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
+
+from util.box_ops import box_xyxy_to_cxcywh
+from util.misc import interpolate
+
+
+def crop(image, target, region):
+    cropped_image = F.crop(image, *region)
+
+    target = target.copy()
+    i, j, h, w = region
+
+    # should we do something wrt the original size?
+    target["size"] = torch.tensor([h, w])
+
+    fields = ["labels", "area"]
+
+    # Crop exemplars.
+    exemplars = target["exemplars"]
+    max_size = torch.as_tensor([w, h], dtype=torch.float32)
+    # Shift exemplars to cropped region.
+    cropped_exemplars = exemplars - torch.as_tensor([j, i, j, i])
+    # Correct exemplar regions that go past new image boundary (too far right).
+    cropped_exemplars = torch.min(cropped_exemplars.reshape(-1, 2, 2), max_size)
+    # Correct exemplar regions that go past new image boundary (too far left).
+    cropped_exemplars = cropped_exemplars.clamp(min=0)
+    # Get new exemplar areas.
+    area_exemplars = (cropped_exemplars[:, 1, :] - cropped_exemplars[:, 0, :]).prod(
+        dim=1
+    )
+    # Update [target] with cropped exemplars.
+    target["exemplars"] = cropped_exemplars.reshape(-1, 4)
+    if "boxes" in target:
+        boxes = target["boxes"]
+        max_size = torch.as_tensor([w, h], dtype=torch.float32)
+        cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
+        cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
+        cropped_boxes = cropped_boxes.clamp(min=0)
+        area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
+        target["boxes"] = cropped_boxes.reshape(-1, 4)
+        target["area"] = area
+        fields.append("boxes")
+
+    if "masks" in target:
+        # FIXME should we update the area here if there are no boxes?
+        target["masks"] = target["masks"][:, i : i + h, j : j + w]
+        fields.append("masks")
+
+    # Remove exemplars that have zero area (due to cropping).
+    keep = area_exemplars > 0
+    target["exemplars"] = target["exemplars"][keep, :]
+
+    # remove elements for which the boxes or masks that have zero area
+    if "boxes" in target or "masks" in target:
+        # favor boxes selection when defining which elements to keep
+        # this is compatible with previous implementation
+        if "boxes" in target:
+            cropped_boxes = target["boxes"].reshape(-1, 2, 2)
+            keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
+        else:
+            keep = target["masks"].flatten(1).any(1)
+
+        for field in fields:
+            target[field] = target[field][keep]
+
+    return cropped_image, target
+
+
+def hflip(image, target):
+    flipped_image = F.hflip(image)
+
+    w, h = image.size
+
+    target = target.copy()
+
+    exemplars = target["exemplars"]
+    # Flip image across x-axis.
+    exemplars = exemplars[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1])
+    # Shift flipped image to (0, 0).
+    exemplars = exemplars + torch.as_tensor([w, 0, w, 0])
+    # Update [target] with horizontally flipped exemplars.
+    target["exemplars"] = exemplars
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor(
+            [-1, 1, -1, 1]
+        ) + torch.as_tensor([w, 0, w, 0])
+        target["boxes"] = boxes
+
+    if "masks" in target:
+        target["masks"] = target["masks"].flip(-1)
+
+    return flipped_image, target
+
+
+def resize(image, target, size, max_size=None):
+    # size can be min_size (scalar) or (w, h) tuple
+
+    def get_size_with_aspect_ratio(image_size, size, max_size=None):
+        w, h = image_size
+        if max_size is not None:
+            min_original_size = float(min((w, h)))
+            max_original_size = float(max((w, h)))
+            if max_original_size / min_original_size * size > max_size:
+                size = int(round(max_size * min_original_size / max_original_size))
+
+        if (w <= h and w == size) or (h <= w and h == size):
+            return (h, w)
+
+        if w < h:
+            ow = size
+            oh = int(size * h / w)
+        else:
+            oh = size
+            ow = int(size * w / h)
+
+        return (oh, ow)
+
+    def get_size(image_size, size, max_size=None):
+        if isinstance(size, (list, tuple)):
+            return size[::-1]
+        else:
+            return get_size_with_aspect_ratio(image_size, size, max_size)
+
+    try:
+        size = get_size(image.size, size, max_size)
+    except:
+        size = get_size((image.shape[-1], image.shape[-2]), size, max_size)
+    rescaled_image = F.resize(image, size)
+
+    if target is None:
+        return rescaled_image, None
+
+    ratios = tuple(
+        float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size)
+    )
+    ratio_width, ratio_height = ratios
+
+    target = target.copy()
+
+    # Rescale exemplars.
+    exemplars = target["exemplars"]
+    if exemplars.shape[-1] == 4:
+        scaled_exemplars = exemplars * torch.as_tensor(
+            [ratio_width, ratio_height, ratio_width, ratio_height]
+        )
+    else:
+        scaled_exemplars = exemplars
+
+    target["exemplars"] = scaled_exemplars
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        scaled_boxes = boxes * torch.as_tensor(
+            [ratio_width, ratio_height, ratio_width, ratio_height]
+        )
+        target["boxes"] = scaled_boxes
+
+    if "area" in target:
+        area = target["area"]
+        scaled_area = area * (ratio_width * ratio_height)
+        target["area"] = scaled_area
+
+    h, w = size
+    target["size"] = torch.tensor([h, w])
+
+    if "masks" in target:
+        target["masks"] = (
+            interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0]
+            > 0.5
+        )
+
+    return rescaled_image, target
+
+
+def pad(image, target, padding):
+    # assumes that we only pad on the bottom right corners
+    padded_image = F.pad(image, (0, 0, padding[0], padding[1]))
+    if target is None:
+        return padded_image, None
+    target = target.copy()
+    # should we do something wrt the original size?
+    target["size"] = torch.tensor(padded_image.size[::-1])
+    if "masks" in target:
+        target["masks"] = torch.nn.functional.pad(
+            target["masks"], (0, padding[0], 0, padding[1])
+        )
+    return padded_image, target
+
+
+class ResizeDebug(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        return resize(img, target, self.size)
+
+
+class RandomCrop(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        region = T.RandomCrop.get_params(img, self.size)
+        return crop(img, target, region)
+
+
+class RandomSizeCrop(object):
+    def __init__(self, min_size: int, max_size: int):
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, img: PIL.Image.Image, target: dict):
+        w = random.randint(self.min_size, min(img.width, self.max_size))
+        h = random.randint(self.min_size, min(img.height, self.max_size))
+        region = T.RandomCrop.get_params(img, [h, w])
+        return crop(img, target, region)
+
+
+class CenterCrop(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        image_width, image_height = img.size
+        crop_height, crop_width = self.size
+        crop_top = int(round((image_height - crop_height) / 2.0))
+        crop_left = int(round((image_width - crop_width) / 2.0))
+        return crop(img, target, (crop_top, crop_left, crop_height, crop_width))
+
+
+class RandomHorizontalFlip(object):
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, target):
+        if random.random() < self.p:
+            return hflip(img, target)
+        return img, target
+
+
+class RandomResize(object):
+    def __init__(self, sizes, max_size=None):
+        assert isinstance(sizes, (list, tuple))
+        self.sizes = sizes
+        self.max_size = max_size
+
+    def __call__(self, img, target=None):
+        size = random.choice(self.sizes)
+        return resize(img, target, size, self.max_size)
+
+
+class RandomPad(object):
+    def __init__(self, max_pad):
+        self.max_pad = max_pad
+
+    def __call__(self, img, target):
+        pad_x = random.randint(0, self.max_pad)
+        pad_y = random.randint(0, self.max_pad)
+        return pad(img, target, (pad_x, pad_y))
+
+
+class RandomSelect(object):
+    """
+    Randomly selects between transforms1 and transforms2,
+    with probability p for transforms1 and (1 - p) for transforms2
+    """
+
+    def __init__(self, transforms1, transforms2, p=0.5):
+        self.transforms1 = transforms1
+        self.transforms2 = transforms2
+        self.p = p
+
+    def __call__(self, img, target):
+        if random.random() < self.p:
+            return self.transforms1(img, target)
+        return self.transforms2(img, target)
+
+
+class ToTensor(object):
+    def __call__(self, img, target):
+        return F.to_tensor(img), target
+
+
+class RandomErasing(object):
+    def __init__(self, *args, **kwargs):
+        self.eraser = T.RandomErasing(*args, **kwargs)
+
+    def __call__(self, img, target):
+        return self.eraser(img), target
+
+
+class Normalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image, target=None):
+        image = F.normalize(image, mean=self.mean, std=self.std)
+        if target is None:
+            return image, None
+        target = target.copy()
+        h, w = image.shape[-2:]
+        # No normalization of exemplars needed, since they are used directly for cropping.
+        if "boxes" in target:
+            boxes = target["boxes"]
+            boxes = box_xyxy_to_cxcywh(boxes)
+            boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
+            target["boxes"] = boxes
+        return image, target
+
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += "    {0}".format(t)
+        format_string += "\n)"
+        return format_string

datasets/transforms_orig.py

@@ -0,0 +1,317 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Transforms and data augmentation for both image + bbox.
+"""
+import random
+
+import PIL
+import torch
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
+
+from util.box_ops import box_xyxy_to_cxcywh
+from util.misc import interpolate
+
+
+def crop(image, target, region):
+    cropped_image = F.crop(image, *region)
+
+    target = target.copy()
+    i, j, h, w = region
+
+    # should we do something wrt the original size?
+    target["size"] = torch.tensor([h, w])
+
+    fields = ["labels", "area"]
+
+    # Crop exemplars.
+    exemplars = target["exemplars"]
+    max_size = torch.as_tensor([w, h], dtype=torch.float32)
+    # Shift exemplars to cropped region.
+    cropped_exemplars = exemplars - torch.as_tensor([j, i, j, i])
+    # Correct exemplar regions that go past new image boundary (too far right).
+    cropped_exemplars = torch.min(cropped_exemplars.reshape(-1, 2, 2), max_size)
+    # Correct exemplar regions that go past new image boundary (too far left).
+    cropped_exemplars = cropped_exemplars.clamp(min=0)
+    # Get new exemplar areas.
+    area_exemplars = (cropped_exemplars[:, 1, :] - cropped_exemplars[:, 0, :]).prod(dim=1)
+    # Update [target] with cropped exemplars.
+    target["exemplars"] = cropped_exemplars.reshape(-1, 4)
+    if "boxes" in target:
+        boxes = target["boxes"]
+        max_size = torch.as_tensor([w, h], dtype=torch.float32)
+        cropped_boxes = boxes - torch.as_tensor([j, i, j, i])
+        cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)
+        cropped_boxes = cropped_boxes.clamp(min=0)
+        area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
+        target["boxes"] = cropped_boxes.reshape(-1, 4)
+        target["area"] = area
+        fields.append("boxes")
+
+    if "masks" in target:
+        # FIXME should we update the area here if there are no boxes?
+        target['masks'] = target['masks'][:, i:i + h, j:j + w]
+        fields.append("masks")
+
+    # Remove exemplars that have zero area (due to cropping).
+    keep = area_exemplars > 0
+    target["exemplars"] = target["exemplars"][keep, :]
+
+    # remove elements for which the boxes or masks that have zero area
+    if "boxes" in target or "masks" in target:
+        # favor boxes selection when defining which elements to keep
+        # this is compatible with previous implementation
+        if "boxes" in target:
+            cropped_boxes = target['boxes'].reshape(-1, 2, 2)
+            keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1)
+        else:
+            keep = target['masks'].flatten(1).any(1)
+
+        for field in fields:
+            target[field] = target[field][keep]
+
+    return cropped_image, target
+
+
+def hflip(image, target):
+    flipped_image = F.hflip(image)
+
+    w, h = image.size
+
+    target = target.copy()
+
+    exemplars = target["exemplars"]
+    # Flip image across x-axis.
+    exemplars = exemplars[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1])
+    # Shift flipped image to (0, 0).
+    exemplars = exemplars + torch.as_tensor([w, 0, w, 0])
+    # Update [target] with horizontally flipped exemplars.
+    target["exemplars"] = exemplars
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor([w, 0, w, 0])
+        target["boxes"] = boxes
+
+    if "masks" in target:
+        target['masks'] = target['masks'].flip(-1)
+
+    return flipped_image, target
+
+
+def resize(image, target, size, max_size=None):
+    # size can be min_size (scalar) or (w, h) tuple
+
+    def get_size_with_aspect_ratio(image_size, size, max_size=None):
+        w, h = image_size
+        if max_size is not None:
+            min_original_size = float(min((w, h)))
+            max_original_size = float(max((w, h)))
+            if max_original_size / min_original_size * size > max_size:
+                size = int(round(max_size * min_original_size / max_original_size))
+
+        if (w <= h and w == size) or (h <= w and h == size):
+            return (h, w)
+
+        if w < h:
+            ow = size
+            oh = int(size * h / w)
+        else:
+            oh = size
+            ow = int(size * w / h)
+
+        return (oh, ow)
+
+    def get_size(image_size, size, max_size=None):
+        if isinstance(size, (list, tuple)):
+            return size[::-1]
+        else:
+            return get_size_with_aspect_ratio(image_size, size, max_size)
+
+    size = get_size(image.size, size, max_size)
+    rescaled_image = F.resize(image, size)
+
+    if target is None:
+        return rescaled_image, None
+
+    ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))
+    ratio_width, ratio_height = ratios
+
+    target = target.copy()
+
+    # Rescale exemplars.
+    exemplars = target["exemplars"]
+    scaled_exemplars = exemplars * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
+    target["exemplars"] = scaled_exemplars
+
+    if "boxes" in target:
+        boxes = target["boxes"]
+        scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height])
+        target["boxes"] = scaled_boxes
+
+    if "area" in target:
+        area = target["area"]
+        scaled_area = area * (ratio_width * ratio_height)
+        target["area"] = scaled_area
+
+    h, w = size
+    target["size"] = torch.tensor([h, w])
+
+    if "masks" in target:
+        target['masks'] = interpolate(
+            target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
+
+    return rescaled_image, target
+
+
+def pad(image, target, padding):
+    # assumes that we only pad on the bottom right corners
+    padded_image = F.pad(image, (0, 0, padding[0], padding[1]))
+    if target is None:
+        return padded_image, None
+    target = target.copy()
+    # should we do something wrt the original size?
+    target["size"] = torch.tensor(padded_image.size[::-1])
+    if "masks" in target:
+        target['masks'] = torch.nn.functional.pad(target['masks'], (0, padding[0], 0, padding[1]))
+    return padded_image, target
+
+
+class ResizeDebug(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        return resize(img, target, self.size)
+
+
+class RandomCrop(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        region = T.RandomCrop.get_params(img, self.size)
+        return crop(img, target, region)
+
+
+class RandomSizeCrop(object):
+    def __init__(self, min_size: int, max_size: int):
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, img: PIL.Image.Image, target: dict):
+        w = random.randint(self.min_size, min(img.width, self.max_size))
+        h = random.randint(self.min_size, min(img.height, self.max_size))
+        region = T.RandomCrop.get_params(img, [h, w])
+        return crop(img, target, region)
+
+
+class CenterCrop(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, img, target):
+        image_width, image_height = img.size
+        crop_height, crop_width = self.size
+        crop_top = int(round((image_height - crop_height) / 2.))
+        crop_left = int(round((image_width - crop_width) / 2.))
+        return crop(img, target, (crop_top, crop_left, crop_height, crop_width))
+
+
+class RandomHorizontalFlip(object):
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, target):
+        if random.random() < self.p:
+            return hflip(img, target)
+        return img, target
+
+
+class RandomResize(object):
+    def __init__(self, sizes, max_size=None):
+        assert isinstance(sizes, (list, tuple))
+        self.sizes = sizes
+        self.max_size = max_size
+
+    def __call__(self, img, target=None):
+        size = random.choice(self.sizes)
+        return resize(img, target, size, self.max_size)
+
+
+class RandomPad(object):
+    def __init__(self, max_pad):
+        self.max_pad = max_pad
+
+    def __call__(self, img, target):
+        pad_x = random.randint(0, self.max_pad)
+        pad_y = random.randint(0, self.max_pad)
+        return pad(img, target, (pad_x, pad_y))
+
+
+class RandomSelect(object):
+    """
+    Randomly selects between transforms1 and transforms2,
+    with probability p for transforms1 and (1 - p) for transforms2
+    """
+    def __init__(self, transforms1, transforms2, p=0.5):
+        self.transforms1 = transforms1
+        self.transforms2 = transforms2
+        self.p = p
+
+    def __call__(self, img, target):
+        if random.random() < self.p:
+            return self.transforms1(img, target)
+        return self.transforms2(img, target)
+
+
+class ToTensor(object):
+    def __call__(self, img, target):
+        return F.to_tensor(img), target
+
+
+class RandomErasing(object):
+
+    def __init__(self, *args, **kwargs):
+        self.eraser = T.RandomErasing(*args, **kwargs)
+
+    def __call__(self, img, target):
+        return self.eraser(img), target
+
+
+class Normalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image, target=None):
+        image = F.normalize(image, mean=self.mean, std=self.std)
+        if target is None:
+            return image, None
+        target = target.copy()
+        h, w = image.shape[-2:]
+        # No normalization of exemplars needed, since they are used directly for cropping.
+        if "boxes" in target:
+            boxes = target["boxes"]
+            boxes = box_xyxy_to_cxcywh(boxes)
+            boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
+            target["boxes"] = boxes
+        return image, target
+
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += "    {0}".format(t)
+        format_string += "\n)"
+        return format_string

example.py

@@ -0,0 +1,33 @@
+"""Count strawberries (not blueberries) in img/336.jpg and save the detections.
+
+    python example.py [--repo nikigoli/CountGDPlusPlus] [--device cuda]
+"""
+import argparse
+import os
+
+import torch
+from PIL import Image, ImageDraw
+
+from countgdpp import CountGDPlusPlus, count
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--repo", default="nikigoli/CountGDPlusPlus")
+parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+args = parser.parse_args()
+
+model = CountGDPlusPlus.from_pretrained(args.repo).to(args.device).eval()
+
+image_path = os.path.join(HERE, "img", "336.jpg")
+result = count(model, image_path, "strawberry", negative_texts=["blueberry"])
+print("Predicted count:", result["count"])
+
+image = Image.open(image_path).convert("RGB")
+w, h = image.size
+draw = ImageDraw.Draw(image)
+for cx, cy, bw, bh in result["boxes"]:
+    draw.rectangle([(cx - bw / 2) * w, (cy - bh / 2) * h, (cx + bw / 2) * w, (cy + bh / 2) * h],
+                   outline="cyan", width=3)
+image.save("detections.png")
+print("Saved detections.png")

groundingdino/util/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

groundingdino/util/box_ops.py

@@ -0,0 +1,140 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Utilities for bounding box manipulation and GIoU.
+"""
+import torch
+from torchvision.ops.boxes import box_area
+
+
+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 box_xyxy_to_cxcywh(x):
+    x0, y0, x1, y1 = x.unbind(-1)
+    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
+    return torch.stack(b, dim=-1)
+
+
+# modified from torchvision to also return the union
+def box_iou(boxes1, boxes2):
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+
+    # import ipdb; ipdb.set_trace()
+    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
+    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,2]
+    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
+
+    union = area1[:, None] + area2 - inter
+
+    iou = inter / (union + 1e-6)
+    return iou, union
+
+
+def generalized_box_iou(boxes1, boxes2):
+    """
+    Generalized IoU from https://giou.stanford.edu/
+
+    The boxes should be in [x0, y0, x1, y1] format
+
+    Returns a [N, M] pairwise matrix, where N = len(boxes1)
+    and M = len(boxes2)
+    """
+    # degenerate boxes gives inf / nan results
+    # so do an early check
+    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
+    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
+    # except:
+    #     import ipdb; ipdb.set_trace()
+    iou, union = box_iou(boxes1, boxes2)
+
+    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
+    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,2]
+    area = wh[:, :, 0] * wh[:, :, 1]
+
+    return iou - (area - union) / (area + 1e-6)
+
+
+# modified from torchvision to also return the union
+def box_iou_pairwise(boxes1, boxes2):
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+
+    lt = torch.max(boxes1[:, :2], boxes2[:, :2])  # [N,2]
+    rb = torch.min(boxes1[:, 2:], boxes2[:, 2:])  # [N,2]
+
+    wh = (rb - lt).clamp(min=0)  # [N,2]
+    inter = wh[:, 0] * wh[:, 1]  # [N]
+
+    union = area1 + area2 - inter
+
+    iou = inter / union
+    return iou, union
+
+
+def generalized_box_iou_pairwise(boxes1, boxes2):
+    """
+    Generalized IoU from https://giou.stanford.edu/
+
+    Input:
+        - boxes1, boxes2: N,4
+    Output:
+        - giou: N, 4
+    """
+    # degenerate boxes gives inf / nan results
+    # so do an early check
+    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
+    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
+    assert boxes1.shape == boxes2.shape
+    iou, union = box_iou_pairwise(boxes1, boxes2)  # N, 4
+
+    lt = torch.min(boxes1[:, :2], boxes2[:, :2])
+    rb = torch.max(boxes1[:, 2:], boxes2[:, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # [N,2]
+    area = wh[:, 0] * wh[:, 1]
+
+    return iou - (area - union) / area
+
+
+def masks_to_boxes(masks):
+    """Compute the bounding boxes around the provided masks
+
+    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.
+
+    Returns a [N, 4] tensors, with the boxes in xyxy format
+    """
+    if masks.numel() == 0:
+        return torch.zeros((0, 4), device=masks.device)
+
+    h, w = masks.shape[-2:]
+
+    y = torch.arange(0, h, dtype=torch.float)
+    x = torch.arange(0, w, dtype=torch.float)
+    y, x = torch.meshgrid(y, x)
+
+    x_mask = masks * x.unsqueeze(0)
+    x_max = x_mask.flatten(1).max(-1)[0]
+    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
+
+    y_mask = masks * y.unsqueeze(0)
+    y_max = y_mask.flatten(1).max(-1)[0]
+    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
+
+    return torch.stack([x_min, y_min, x_max, y_max], 1)
+
+
+if __name__ == "__main__":
+    x = torch.rand(5, 4)
+    y = torch.rand(3, 4)
+    iou, union = box_iou(x, y)
+    import ipdb
+
+    ipdb.set_trace()

groundingdino/util/get_tokenlizer.py

@@ -0,0 +1,29 @@
+from transformers import AutoTokenizer, BertModel, BertTokenizer, RobertaModel, RobertaTokenizerFast
+import os
+
+def get_tokenlizer(text_encoder_type):
+    if not isinstance(text_encoder_type, str):
+        # print("text_encoder_type is not a str")
+        if hasattr(text_encoder_type, "text_encoder_type"):
+            text_encoder_type = text_encoder_type.text_encoder_type
+        elif text_encoder_type.get("text_encoder_type", False):
+            text_encoder_type = text_encoder_type.get("text_encoder_type")
+        elif os.path.isdir(text_encoder_type) and os.path.exists(text_encoder_type):
+            pass
+        else:
+            raise ValueError(
+                "Unknown type of text_encoder_type: {}".format(type(text_encoder_type))
+            )
+    print("final text_encoder_type: {}".format(text_encoder_type))
+    tokenizer = AutoTokenizer.from_pretrained(text_encoder_type)
+    print("load tokenizer done.")
+    return tokenizer
+
+
+def get_pretrained_language_model(text_encoder_type):
+    if text_encoder_type == "bert-base-uncased" or (os.path.isdir(text_encoder_type) and os.path.exists(text_encoder_type)):
+        return BertModel.from_pretrained(text_encoder_type)
+    if text_encoder_type == "roberta-base":
+        return RobertaModel.from_pretrained(text_encoder_type)
+
+    raise ValueError("Unknown text_encoder_type {}".format(text_encoder_type))

groundingdino/util/inference.py

@@ -0,0 +1,259 @@
+from typing import Tuple, List
+
+import cv2
+import numpy as np
+import supervision as sv
+import torch
+from PIL import Image
+from torchvision.ops import box_convert
+import bisect
+
+import groundingdino.datasets.transforms as T
+from groundingdino.models import build_model
+from groundingdino.util.misc import clean_state_dict
+from groundingdino.util.slconfig import SLConfig
+from groundingdino.util.utils import get_phrases_from_posmap
+
+# ----------------------------------------------------------------------------------------------------------------------
+# OLD API
+# ----------------------------------------------------------------------------------------------------------------------
+
+
+def preprocess_caption(caption: str) -> str:
+    result = caption.lower().strip()
+    if result.endswith("."):
+        return result
+    return result + "."
+
+
+def load_model(model_config_path: str, model_checkpoint_path: str, device: str = "cuda"):
+    args = SLConfig.fromfile(model_config_path)
+    args.device = device
+    model = build_model(args)
+    checkpoint = torch.load(model_checkpoint_path, map_location="cpu")
+    model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
+    model.eval()
+    return model
+
+
+def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
+    transform = T.Compose(
+        [
+            T.RandomResize([800], max_size=1333),
+            T.ToTensor(),
+            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    image_source = Image.open(image_path).convert("RGB")
+    image = np.asarray(image_source)
+    image_transformed, _ = transform(image_source, None)
+    return image, image_transformed
+
+
+def predict(
+        model,
+        image: torch.Tensor,
+        caption: str,
+        box_threshold: float,
+        text_threshold: float,
+        device: str = "cuda",
+        remove_combined: bool = False
+) -> Tuple[torch.Tensor, torch.Tensor, List[str]]:
+    caption = preprocess_caption(caption=caption)
+
+    model = model.to(device)
+    image = image.to(device)
+
+    with torch.no_grad():
+        outputs = model(image[None], captions=[caption])
+
+    prediction_logits = outputs["pred_logits"].cpu().sigmoid()[0]  # prediction_logits.shape = (nq, 256)
+    prediction_boxes = outputs["pred_boxes"].cpu()[0]  # prediction_boxes.shape = (nq, 4)
+
+    mask = prediction_logits.max(dim=1)[0] > box_threshold
+    logits = prediction_logits[mask]  # logits.shape = (n, 256)
+    boxes = prediction_boxes[mask]  # boxes.shape = (n, 4)
+
+    tokenizer = model.tokenizer
+    tokenized = tokenizer(caption)
+    
+    if remove_combined:
+        sep_idx = [i for i in range(len(tokenized['input_ids'])) if tokenized['input_ids'][i] in [101, 102, 1012]]
+        
+        phrases = []
+        for logit in logits:
+            max_idx = logit.argmax()
+            insert_idx = bisect.bisect_left(sep_idx, max_idx)
+            right_idx = sep_idx[insert_idx]
+            left_idx = sep_idx[insert_idx - 1]
+            phrases.append(get_phrases_from_posmap(logit > text_threshold, tokenized, tokenizer, left_idx, right_idx).replace('.', ''))
+    else:
+        phrases = [
+            get_phrases_from_posmap(logit > text_threshold, tokenized, tokenizer).replace('.', '')
+            for logit
+            in logits
+        ]
+
+    return boxes, logits.max(dim=1)[0], phrases
+
+
+def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str]) -> np.ndarray:
+    h, w, _ = image_source.shape
+    boxes = boxes * torch.Tensor([w, h, w, h])
+    xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
+    detections = sv.Detections(xyxy=xyxy)
+
+    labels = [
+        f"{phrase} {logit:.2f}"
+        for phrase, logit
+        in zip(phrases, logits)
+    ]
+
+    box_annotator = sv.BoxAnnotator()
+    annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR)
+    annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
+    return annotated_frame
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+# NEW API
+# ----------------------------------------------------------------------------------------------------------------------
+
+
+class Model:
+
+    def __init__(
+        self,
+        model_config_path: str,
+        model_checkpoint_path: str,
+        device: str = "cuda"
+    ):
+        self.model = load_model(
+            model_config_path=model_config_path,
+            model_checkpoint_path=model_checkpoint_path,
+            device=device
+        ).to(device)
+        self.device = device
+
+    def predict_with_caption(
+        self,
+        image: np.ndarray,
+        caption: str,
+        box_threshold: float = 0.35,
+        text_threshold: float = 0.25
+    ) -> Tuple[sv.Detections, List[str]]:
+        """
+        import cv2
+
+        image = cv2.imread(IMAGE_PATH)
+
+        model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH)
+        detections, labels = model.predict_with_caption(
+            image=image,
+            caption=caption,
+            box_threshold=BOX_THRESHOLD,
+            text_threshold=TEXT_THRESHOLD
+        )
+
+        import supervision as sv
+
+        box_annotator = sv.BoxAnnotator()
+        annotated_image = box_annotator.annotate(scene=image, detections=detections, labels=labels)
+        """
+        processed_image = Model.preprocess_image(image_bgr=image).to(self.device)
+        boxes, logits, phrases = predict(
+            model=self.model,
+            image=processed_image,
+            caption=caption,
+            box_threshold=box_threshold,
+            text_threshold=text_threshold, 
+            device=self.device)
+        source_h, source_w, _ = image.shape
+        detections = Model.post_process_result(
+            source_h=source_h,
+            source_w=source_w,
+            boxes=boxes,
+            logits=logits)
+        return detections, phrases
+
+    def predict_with_classes(
+        self,
+        image: np.ndarray,
+        classes: List[str],
+        box_threshold: float,
+        text_threshold: float
+    ) -> sv.Detections:
+        """
+        import cv2
+
+        image = cv2.imread(IMAGE_PATH)
+
+        model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH)
+        detections = model.predict_with_classes(
+            image=image,
+            classes=CLASSES,
+            box_threshold=BOX_THRESHOLD,
+            text_threshold=TEXT_THRESHOLD
+        )
+
+
+        import supervision as sv
+
+        box_annotator = sv.BoxAnnotator()
+        annotated_image = box_annotator.annotate(scene=image, detections=detections)
+        """
+        caption = ". ".join(classes)
+        processed_image = Model.preprocess_image(image_bgr=image).to(self.device)
+        boxes, logits, phrases = predict(
+            model=self.model,
+            image=processed_image,
+            caption=caption,
+            box_threshold=box_threshold,
+            text_threshold=text_threshold,
+            device=self.device)
+        source_h, source_w, _ = image.shape
+        detections = Model.post_process_result(
+            source_h=source_h,
+            source_w=source_w,
+            boxes=boxes,
+            logits=logits)
+        class_id = Model.phrases2classes(phrases=phrases, classes=classes)
+        detections.class_id = class_id
+        return detections
+
+    @staticmethod
+    def preprocess_image(image_bgr: np.ndarray) -> torch.Tensor:
+        transform = T.Compose(
+            [
+                T.RandomResize([800], max_size=1333),
+                T.ToTensor(),
+                T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+        image_pillow = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
+        image_transformed, _ = transform(image_pillow, None)
+        return image_transformed
+
+    @staticmethod
+    def post_process_result(
+            source_h: int,
+            source_w: int,
+            boxes: torch.Tensor,
+            logits: torch.Tensor
+    ) -> sv.Detections:
+        boxes = boxes * torch.Tensor([source_w, source_h, source_w, source_h])
+        xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
+        confidence = logits.numpy()
+        return sv.Detections(xyxy=xyxy, confidence=confidence)
+
+    @staticmethod
+    def phrases2classes(phrases: List[str], classes: List[str]) -> np.ndarray:
+        class_ids = []
+        for phrase in phrases:
+            for class_ in classes:
+                if class_ in phrase:
+                    class_ids.append(classes.index(class_))
+                    break
+            else:
+                class_ids.append(None)
+        return np.array(class_ids)

groundingdino/util/logger.py

@@ -0,0 +1,93 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import functools
+import logging
+import os
+import sys
+
+from termcolor import colored
+
+
+class _ColorfulFormatter(logging.Formatter):
+    def __init__(self, *args, **kwargs):
+        self._root_name = kwargs.pop("root_name") + "."
+        self._abbrev_name = kwargs.pop("abbrev_name", "")
+        if len(self._abbrev_name):
+            self._abbrev_name = self._abbrev_name + "."
+        super(_ColorfulFormatter, self).__init__(*args, **kwargs)
+
+    def formatMessage(self, record):
+        record.name = record.name.replace(self._root_name, self._abbrev_name)
+        log = super(_ColorfulFormatter, self).formatMessage(record)
+        if record.levelno == logging.WARNING:
+            prefix = colored("WARNING", "red", attrs=["blink"])
+        elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:
+            prefix = colored("ERROR", "red", attrs=["blink", "underline"])
+        else:
+            return log
+        return prefix + " " + log
+
+
+# so that calling setup_logger multiple times won't add many handlers
+@functools.lru_cache()
+def setup_logger(output=None, distributed_rank=0, *, color=True, name="imagenet", abbrev_name=None):
+    """
+    Initialize the detectron2 logger and set its verbosity level to "INFO".
+
+    Args:
+        output (str): a file name or a directory to save log. If None, will not save log file.
+            If ends with ".txt" or ".log", assumed to be a file name.
+            Otherwise, logs will be saved to `output/log.txt`.
+        name (str): the root module name of this logger
+
+    Returns:
+        logging.Logger: a logger
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+
+    if abbrev_name is None:
+        abbrev_name = name
+
+    plain_formatter = logging.Formatter(
+        "[%(asctime)s.%(msecs)03d]: %(message)s", datefmt="%m/%d %H:%M:%S"
+    )
+    # stdout logging: master only
+    if distributed_rank == 0:
+        ch = logging.StreamHandler(stream=sys.stdout)
+        ch.setLevel(logging.DEBUG)
+        if color:
+            formatter = _ColorfulFormatter(
+                colored("[%(asctime)s.%(msecs)03d]: ", "green") + "%(message)s",
+                datefmt="%m/%d %H:%M:%S",
+                root_name=name,
+                abbrev_name=str(abbrev_name),
+            )
+        else:
+            formatter = plain_formatter
+        ch.setFormatter(formatter)
+        logger.addHandler(ch)
+
+    # file logging: all workers
+    if output is not None:
+        if output.endswith(".txt") or output.endswith(".log"):
+            filename = output
+        else:
+            filename = os.path.join(output, "log.txt")
+        if distributed_rank > 0:
+            filename = filename + f".rank{distributed_rank}"
+        os.makedirs(os.path.dirname(filename), exist_ok=True)
+
+        fh = logging.StreamHandler(_cached_log_stream(filename))
+        fh.setLevel(logging.DEBUG)
+        fh.setFormatter(plain_formatter)
+        logger.addHandler(fh)
+
+    return logger
+
+
+# cache the opened file object, so that different calls to `setup_logger`
+# with the same file name can safely write to the same file.
+@functools.lru_cache(maxsize=None)
+def _cached_log_stream(filename):
+    return open(filename, "a")

groundingdino/util/misc.py

@@ -0,0 +1,717 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Misc functions, including distributed helpers.
+
+Mostly copy-paste from torchvision references.
+"""
+import colorsys
+import datetime
+import functools
+import io
+import json
+import os
+import pickle
+import subprocess
+import time
+from collections import OrderedDict, defaultdict, deque
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+# needed due to empty tensor bug in pytorch and torchvision 0.5
+import torchvision
+from torch import Tensor
+
+__torchvision_need_compat_flag = float(torchvision.__version__.split(".")[1]) < 7
+if __torchvision_need_compat_flag:
+    from torchvision.ops import _new_empty_tensor
+    from torchvision.ops.misc import _output_size
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        if d.shape[0] == 0:
+            return 0
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        if os.environ.get("SHILONG_AMP", None) == "1":
+            eps = 1e-4
+        else:
+            eps = 1e-6
+        return self.total / (self.count + eps)
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+@functools.lru_cache()
+def _get_global_gloo_group():
+    """
+    Return a process group based on gloo backend, containing all the ranks
+    The result is cached.
+    """
+
+    if dist.get_backend() == "nccl":
+        return dist.new_group(backend="gloo")
+
+    return dist.group.WORLD
+
+
+def all_gather_cpu(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    cpu_group = _get_global_gloo_group()
+
+    buffer = io.BytesIO()
+    torch.save(data, buffer)
+    data_view = buffer.getbuffer()
+    device = "cuda" if cpu_group is None else "cpu"
+    tensor = torch.ByteTensor(data_view).to(device)
+
+    # obtain Tensor size of each rank
+    local_size = torch.tensor([tensor.numel()], device=device, dtype=torch.long)
+    size_list = [torch.tensor([0], device=device, dtype=torch.long) for _ in range(world_size)]
+    if cpu_group is None:
+        dist.all_gather(size_list, local_size)
+    else:
+        print("gathering on cpu")
+        dist.all_gather(size_list, local_size, group=cpu_group)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+    assert isinstance(local_size.item(), int)
+    local_size = int(local_size.item())
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=device))
+    if local_size != max_size:
+        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device=device)
+        tensor = torch.cat((tensor, padding), dim=0)
+    if cpu_group is None:
+        dist.all_gather(tensor_list, tensor)
+    else:
+        dist.all_gather(tensor_list, tensor, group=cpu_group)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        tensor = torch.split(tensor, [size, max_size - size], dim=0)[0]
+        buffer = io.BytesIO(tensor.cpu().numpy())
+        obj = torch.load(buffer)
+        data_list.append(obj)
+
+    return data_list
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+
+    if os.getenv("CPU_REDUCE") == "1":
+        return all_gather_cpu(data)
+
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.tensor([tensor.numel()], device="cuda")
+    size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
+    if local_size != max_size:
+        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            # print(name, str(meter))
+            # import ipdb;ipdb.set_trace()
+            if meter.count > 0:
+                loss_str.append("{}: {}".format(name, str(meter)))
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None, logger=None):
+        if logger is None:
+            print_func = print
+        else:
+            print_func = logger.info
+
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
+        else:
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                ]
+            )
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            # import ipdb; ipdb.set_trace()
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print_func(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print_func(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print_func(
+            "{} Total time: {} ({:.4f} s / it)".format(
+                header, total_time_str, total_time / len(iterable)
+            )
+        )
+
+
+def get_sha():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
+
+    sha = "N/A"
+    diff = "clean"
+    branch = "N/A"
+    try:
+        sha = _run(["git", "rev-parse", "HEAD"])
+        subprocess.check_output(["git", "diff"], cwd=cwd)
+        diff = _run(["git", "diff-index", "HEAD"])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+
+def collate_fn(batch):
+    # import ipdb; ipdb.set_trace()
+    batch = list(zip(*batch))
+    batch[0] = nested_tensor_from_tensor_list(batch[0])
+    return tuple(batch)
+
+
+def _max_by_axis(the_list):
+    # type: (List[List[int]]) -> List[int]
+    maxes = the_list[0]
+    for sublist in the_list[1:]:
+        for index, item in enumerate(sublist):
+            maxes[index] = max(maxes[index], item)
+    return maxes
+
+
+class NestedTensor(object):
+    def __init__(self, tensors, mask: Optional[Tensor]):
+        self.tensors = tensors
+        self.mask = mask
+        if mask == "auto":
+            self.mask = torch.zeros_like(tensors).to(tensors.device)
+            if self.mask.dim() == 3:
+                self.mask = self.mask.sum(0).to(bool)
+            elif self.mask.dim() == 4:
+                self.mask = self.mask.sum(1).to(bool)
+            else:
+                raise ValueError(
+                    "tensors dim must be 3 or 4 but {}({})".format(
+                        self.tensors.dim(), self.tensors.shape
+                    )
+                )
+
+    def imgsize(self):
+        res = []
+        for i in range(self.tensors.shape[0]):
+            mask = self.mask[i]
+            maxH = (~mask).sum(0).max()
+            maxW = (~mask).sum(1).max()
+            res.append(torch.Tensor([maxH, maxW]))
+        return res
+
+    def to(self, device):
+        # type: (Device) -> NestedTensor # noqa
+        cast_tensor = self.tensors.to(device)
+        mask = self.mask
+        if mask is not None:
+            assert mask is not None
+            cast_mask = mask.to(device)
+        else:
+            cast_mask = None
+        return NestedTensor(cast_tensor, cast_mask)
+
+    def to_img_list_single(self, tensor, mask):
+        assert tensor.dim() == 3, "dim of tensor should be 3 but {}".format(tensor.dim())
+        maxH = (~mask).sum(0).max()
+        maxW = (~mask).sum(1).max()
+        img = tensor[:, :maxH, :maxW]
+        return img
+
+    def to_img_list(self):
+        """remove the padding and convert to img list
+
+        Returns:
+            [type]: [description]
+        """
+        if self.tensors.dim() == 3:
+            return self.to_img_list_single(self.tensors, self.mask)
+        else:
+            res = []
+            for i in range(self.tensors.shape[0]):
+                tensor_i = self.tensors[i]
+                mask_i = self.mask[i]
+                res.append(self.to_img_list_single(tensor_i, mask_i))
+            return res
+
+    @property
+    def device(self):
+        return self.tensors.device
+
+    def decompose(self):
+        return self.tensors, self.mask
+
+    def __repr__(self):
+        return str(self.tensors)
+
+    @property
+    def shape(self):
+        return {"tensors.shape": self.tensors.shape, "mask.shape": self.mask.shape}
+
+
+def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
+    # TODO make this more general
+    if tensor_list[0].ndim == 3:
+        if torchvision._is_tracing():
+            # nested_tensor_from_tensor_list() does not export well to ONNX
+            # call _onnx_nested_tensor_from_tensor_list() instead
+            return _onnx_nested_tensor_from_tensor_list(tensor_list)
+
+        # TODO make it support different-sized images
+        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
+        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
+        batch_shape = [len(tensor_list)] + max_size
+        b, c, h, w = batch_shape
+        dtype = tensor_list[0].dtype
+        device = tensor_list[0].device
+        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
+        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
+        for img, pad_img, m in zip(tensor_list, tensor, mask):
+            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
+            m[: img.shape[1], : img.shape[2]] = False
+    else:
+        raise ValueError("not supported")
+    return NestedTensor(tensor, mask)
+
+
+# _onnx_nested_tensor_from_tensor_list() is an implementation of
+# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
+@torch.jit.unused
+def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
+    max_size = []
+    for i in range(tensor_list[0].dim()):
+        max_size_i = torch.max(
+            torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)
+        ).to(torch.int64)
+        max_size.append(max_size_i)
+    max_size = tuple(max_size)
+
+    # work around for
+    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
+    # m[: img.shape[1], :img.shape[2]] = False
+    # which is not yet supported in onnx
+    padded_imgs = []
+    padded_masks = []
+    for img in tensor_list:
+        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
+        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
+        padded_imgs.append(padded_img)
+
+        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
+        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
+        padded_masks.append(padded_mask.to(torch.bool))
+
+    tensor = torch.stack(padded_imgs)
+    mask = torch.stack(padded_masks)
+
+    return NestedTensor(tensor, mask=mask)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if "WORLD_SIZE" in os.environ and os.environ["WORLD_SIZE"] != "":  # 'RANK' in os.environ and
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = args.local_rank = int(os.environ["LOCAL_RANK"])
+
+        # launch by torch.distributed.launch
+        # Single node
+        #   python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 1 --rank 0 ...
+        # Multi nodes
+        #   python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 2 --rank 0 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' ...
+        #   python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 2 --rank 1 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' ...
+        # args.rank = int(os.environ.get('OMPI_COMM_WORLD_RANK'))
+        # local_world_size = int(os.environ['GPU_PER_NODE_COUNT'])
+        # args.world_size = args.world_size * local_world_size
+        # args.gpu = args.local_rank = int(os.environ['LOCAL_RANK'])
+        # args.rank = args.rank * local_world_size + args.local_rank
+        print(
+            "world size: {}, rank: {}, local rank: {}".format(
+                args.world_size, args.rank, args.local_rank
+            )
+        )
+        print(json.dumps(dict(os.environ), indent=2))
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.local_rank = int(os.environ["SLURM_LOCALID"])
+        args.world_size = int(os.environ["SLURM_NPROCS"])
+
+        print(
+            "world size: {}, world rank: {}, local rank: {}, device_count: {}".format(
+                args.world_size, args.rank, args.local_rank, torch.cuda.device_count()
+            )
+        )
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        args.world_size = 1
+        args.rank = 0
+        args.local_rank = 0
+        return
+
+    print("world_size:{} rank:{} local_rank:{}".format(args.world_size, args.rank, args.local_rank))
+    args.distributed = True
+    torch.cuda.set_device(args.local_rank)
+    args.dist_backend = "nccl"
+    print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True)
+
+    torch.distributed.init_process_group(
+        backend=args.dist_backend,
+        world_size=args.world_size,
+        rank=args.rank,
+        init_method=args.dist_url,
+    )
+
+    print("Before torch.distributed.barrier()")
+    torch.distributed.barrier()
+    print("End torch.distributed.barrier()")
+    setup_for_distributed(args.rank == 0)
+
+
+@torch.no_grad()
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    if target.numel() == 0:
+        return [torch.zeros([], device=output.device)]
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+@torch.no_grad()
+def accuracy_onehot(pred, gt):
+    """_summary_
+
+    Args:
+        pred (_type_): n, c
+        gt (_type_): n, c
+    """
+    tp = ((pred - gt).abs().sum(-1) < 1e-4).float().sum()
+    acc = tp / gt.shape[0] * 100
+    return acc
+
+
+def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
+    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
+    """
+    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
+    This will eventually be supported natively by PyTorch, and this
+    class can go away.
+    """
+    if __torchvision_need_compat_flag < 0.7:
+        if input.numel() > 0:
+            return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)
+
+        output_shape = _output_size(2, input, size, scale_factor)
+        output_shape = list(input.shape[:-2]) + list(output_shape)
+        return _new_empty_tensor(input, output_shape)
+    else:
+        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)
+
+
+class color_sys:
+    def __init__(self, num_colors) -> None:
+        self.num_colors = num_colors
+        colors = []
+        for i in np.arange(0.0, 360.0, 360.0 / num_colors):
+            hue = i / 360.0
+            lightness = (50 + np.random.rand() * 10) / 100.0
+            saturation = (90 + np.random.rand() * 10) / 100.0
+            colors.append(
+                tuple([int(j * 255) for j in colorsys.hls_to_rgb(hue, lightness, saturation)])
+            )
+        self.colors = colors
+
+    def __call__(self, idx):
+        return self.colors[idx]
+
+
+def inverse_sigmoid(x, eps=1e-3):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1 / x2)
+
+
+def clean_state_dict(state_dict):
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        if k[:7] == "module.":
+            k = k[7:]  # remove `module.`
+        new_state_dict[k] = v
+    return new_state_dict

groundingdino/util/slconfig.py

@@ -0,0 +1,427 @@
+# ==========================================================
+# Modified from mmcv
+# ==========================================================
+import ast
+import os
+import os.path as osp
+import shutil
+import sys
+import tempfile
+from argparse import Action
+from importlib import import_module
+
+from addict import Dict
+from yapf.yapflib.yapf_api import FormatCode
+
+BASE_KEY = "_base_"
+DELETE_KEY = "_delete_"
+RESERVED_KEYS = ["filename", "text", "pretty_text", "get", "dump", "merge_from_dict"]
+
+
+def check_file_exist(filename, msg_tmpl='file "{}" does not exist'):
+    if not osp.isfile(filename):
+        raise FileNotFoundError(msg_tmpl.format(filename))
+
+
+class ConfigDict(Dict):
+    def __missing__(self, name):
+        raise KeyError(name)
+
+    def __getattr__(self, name):
+        try:
+            value = super(ConfigDict, self).__getattr__(name)
+        except KeyError:
+            ex = AttributeError(f"'{self.__class__.__name__}' object has no " f"attribute '{name}'")
+        except Exception as e:
+            ex = e
+        else:
+            return value
+        raise ex
+
+
+class SLConfig(object):
+    """
+    config files.
+    only support .py file as config now.
+
+    ref: mmcv.utils.config
+
+    Example:
+        >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1])))
+        >>> cfg.a
+        1
+        >>> cfg.b
+        {'b1': [0, 1]}
+        >>> cfg.b.b1
+        [0, 1]
+        >>> cfg = Config.fromfile('tests/data/config/a.py')
+        >>> cfg.filename
+        "/home/kchen/projects/mmcv/tests/data/config/a.py"
+        >>> cfg.item4
+        'test'
+        >>> cfg
+        "Config [path: /home/kchen/projects/mmcv/tests/data/config/a.py]: "
+        "{'item1': [1, 2], 'item2': {'a': 0}, 'item3': True, 'item4': 'test'}"
+    """
+
+    @staticmethod
+    def _validate_py_syntax(filename):
+        with open(filename) as f:
+            content = f.read()
+        try:
+            ast.parse(content)
+        except SyntaxError:
+            raise SyntaxError("There are syntax errors in config " f"file {filename}")
+
+    @staticmethod
+    def _file2dict(filename):
+        filename = osp.abspath(osp.expanduser(filename))
+        check_file_exist(filename)
+        if filename.lower().endswith(".py"):
+            with tempfile.TemporaryDirectory() as temp_config_dir:
+                temp_config_file = tempfile.NamedTemporaryFile(dir=temp_config_dir, suffix=".py")
+                temp_config_name = osp.basename(temp_config_file.name)
+                if os.name == 'nt':
+                    temp_config_file.close()
+                shutil.copyfile(filename, osp.join(temp_config_dir, temp_config_name))
+                temp_module_name = osp.splitext(temp_config_name)[0]
+                sys.path.insert(0, temp_config_dir)
+                SLConfig._validate_py_syntax(filename)
+                mod = import_module(temp_module_name)
+                sys.path.pop(0)
+                cfg_dict = {
+                    name: value for name, value in mod.__dict__.items() if not name.startswith("__")
+                }
+                # delete imported module
+                del sys.modules[temp_module_name]
+                # close temp file
+                temp_config_file.close()
+        elif filename.lower().endswith((".yml", ".yaml", ".json")):
+            from .slio import slload
+
+            cfg_dict = slload(filename)
+        else:
+            raise IOError("Only py/yml/yaml/json type are supported now!")
+
+        cfg_text = filename + "\n"
+        with open(filename, "r") as f:
+            cfg_text += f.read()
+
+        # parse the base file
+        if BASE_KEY in cfg_dict:
+            cfg_dir = osp.dirname(filename)
+            base_filename = cfg_dict.pop(BASE_KEY)
+            base_filename = base_filename if isinstance(base_filename, list) else [base_filename]
+
+            cfg_dict_list = list()
+            cfg_text_list = list()
+            for f in base_filename:
+                _cfg_dict, _cfg_text = SLConfig._file2dict(osp.join(cfg_dir, f))
+                cfg_dict_list.append(_cfg_dict)
+                cfg_text_list.append(_cfg_text)
+
+            base_cfg_dict = dict()
+            for c in cfg_dict_list:
+                if len(base_cfg_dict.keys() & c.keys()) > 0:
+                    raise KeyError("Duplicate key is not allowed among bases")
+                    # TODO Allow the duplicate key while warnning user
+                base_cfg_dict.update(c)
+
+            base_cfg_dict = SLConfig._merge_a_into_b(cfg_dict, base_cfg_dict)
+            cfg_dict = base_cfg_dict
+
+            # merge cfg_text
+            cfg_text_list.append(cfg_text)
+            cfg_text = "\n".join(cfg_text_list)
+
+        return cfg_dict, cfg_text
+
+    @staticmethod
+    def _merge_a_into_b(a, b):
+        """merge dict `a` into dict `b` (non-inplace).
+            values in `a` will overwrite `b`.
+            copy first to avoid inplace modification
+
+        Args:
+            a ([type]): [description]
+            b ([type]): [description]
+
+        Returns:
+            [dict]: [description]
+        """
+        # import ipdb; ipdb.set_trace()
+        if not isinstance(a, dict):
+            return a
+
+        b = b.copy()
+        for k, v in a.items():
+            if isinstance(v, dict) and k in b and not v.pop(DELETE_KEY, False):
+
+                if not isinstance(b[k], dict) and not isinstance(b[k], list):
+                    # if :
+                    # import ipdb; ipdb.set_trace()
+                    raise TypeError(
+                        f"{k}={v} in child config cannot inherit from base "
+                        f"because {k} is a dict in the child config but is of "
+                        f"type {type(b[k])} in base config. You may set "
+                        f"`{DELETE_KEY}=True` to ignore the base config"
+                    )
+                b[k] = SLConfig._merge_a_into_b(v, b[k])
+            elif isinstance(b, list):
+                try:
+                    _ = int(k)
+                except:
+                    raise TypeError(
+                        f"b is a list, " f"index {k} should be an int when input but {type(k)}"
+                    )
+                b[int(k)] = SLConfig._merge_a_into_b(v, b[int(k)])
+            else:
+                b[k] = v
+
+        return b
+
+    @staticmethod
+    def fromfile(filename):
+        cfg_dict, cfg_text = SLConfig._file2dict(filename)
+        return SLConfig(cfg_dict, cfg_text=cfg_text, filename=filename)
+
+    def __init__(self, cfg_dict=None, cfg_text=None, filename=None):
+        if cfg_dict is None:
+            cfg_dict = dict()
+        elif not isinstance(cfg_dict, dict):
+            raise TypeError("cfg_dict must be a dict, but " f"got {type(cfg_dict)}")
+        for key in cfg_dict:
+            if key in RESERVED_KEYS:
+                raise KeyError(f"{key} is reserved for config file")
+
+        super(SLConfig, self).__setattr__("_cfg_dict", ConfigDict(cfg_dict))
+        super(SLConfig, self).__setattr__("_filename", filename)
+        if cfg_text:
+            text = cfg_text
+        elif filename:
+            with open(filename, "r") as f:
+                text = f.read()
+        else:
+            text = ""
+        super(SLConfig, self).__setattr__("_text", text)
+
+    @property
+    def filename(self):
+        return self._filename
+
+    @property
+    def text(self):
+        return self._text
+
+    @property
+    def pretty_text(self):
+
+        indent = 4
+
+        def _indent(s_, num_spaces):
+            s = s_.split("\n")
+            if len(s) == 1:
+                return s_
+            first = s.pop(0)
+            s = [(num_spaces * " ") + line for line in s]
+            s = "\n".join(s)
+            s = first + "\n" + s
+            return s
+
+        def _format_basic_types(k, v, use_mapping=False):
+            if isinstance(v, str):
+                v_str = f"'{v}'"
+            else:
+                v_str = str(v)
+
+            if use_mapping:
+                k_str = f"'{k}'" if isinstance(k, str) else str(k)
+                attr_str = f"{k_str}: {v_str}"
+            else:
+                attr_str = f"{str(k)}={v_str}"
+            attr_str = _indent(attr_str, indent)
+
+            return attr_str
+
+        def _format_list(k, v, use_mapping=False):
+            # check if all items in the list are dict
+            if all(isinstance(_, dict) for _ in v):
+                v_str = "[\n"
+                v_str += "\n".join(
+                    f"dict({_indent(_format_dict(v_), indent)})," for v_ in v
+                ).rstrip(",")
+                if use_mapping:
+                    k_str = f"'{k}'" if isinstance(k, str) else str(k)
+                    attr_str = f"{k_str}: {v_str}"
+                else:
+                    attr_str = f"{str(k)}={v_str}"
+                attr_str = _indent(attr_str, indent) + "]"
+            else:
+                attr_str = _format_basic_types(k, v, use_mapping)
+            return attr_str
+
+        def _contain_invalid_identifier(dict_str):
+            contain_invalid_identifier = False
+            for key_name in dict_str:
+                contain_invalid_identifier |= not str(key_name).isidentifier()
+            return contain_invalid_identifier
+
+        def _format_dict(input_dict, outest_level=False):
+            r = ""
+            s = []
+
+            use_mapping = _contain_invalid_identifier(input_dict)
+            if use_mapping:
+                r += "{"
+            for idx, (k, v) in enumerate(input_dict.items()):
+                is_last = idx >= len(input_dict) - 1
+                end = "" if outest_level or is_last else ","
+                if isinstance(v, dict):
+                    v_str = "\n" + _format_dict(v)
+                    if use_mapping:
+                        k_str = f"'{k}'" if isinstance(k, str) else str(k)
+                        attr_str = f"{k_str}: dict({v_str}"
+                    else:
+                        attr_str = f"{str(k)}=dict({v_str}"
+                    attr_str = _indent(attr_str, indent) + ")" + end
+                elif isinstance(v, list):
+                    attr_str = _format_list(k, v, use_mapping) + end
+                else:
+                    attr_str = _format_basic_types(k, v, use_mapping) + end
+
+                s.append(attr_str)
+            r += "\n".join(s)
+            if use_mapping:
+                r += "}"
+            return r
+
+        cfg_dict = self._cfg_dict.to_dict()
+        text = _format_dict(cfg_dict, outest_level=True)
+        # copied from setup.cfg
+        yapf_style = dict(
+            based_on_style="pep8",
+            blank_line_before_nested_class_or_def=True,
+            split_before_expression_after_opening_paren=True,
+        )
+        text, _ = FormatCode(text, style_config=yapf_style, verify=True)
+
+        return text
+
+    def __repr__(self):
+        return f"Config (path: {self.filename}): {self._cfg_dict.__repr__()}"
+
+    def __len__(self):
+        return len(self._cfg_dict)
+
+    def __getattr__(self, name):
+        # # debug
+        # print('+'*15)
+        # print('name=%s' % name)
+        # print("addr:", id(self))
+        # # print('type(self):', type(self))
+        # print(self.__dict__)
+        # print('+'*15)
+        # if self.__dict__ == {}:
+        #     raise ValueError
+
+        return getattr(self._cfg_dict, name)
+
+    def __getitem__(self, name):
+        return self._cfg_dict.__getitem__(name)
+
+    def __setattr__(self, name, value):
+        if isinstance(value, dict):
+            value = ConfigDict(value)
+        self._cfg_dict.__setattr__(name, value)
+
+    def __setitem__(self, name, value):
+        if isinstance(value, dict):
+            value = ConfigDict(value)
+        self._cfg_dict.__setitem__(name, value)
+
+    def __iter__(self):
+        return iter(self._cfg_dict)
+
+    def dump(self, file=None):
+        # import ipdb; ipdb.set_trace()
+        if file is None:
+            return self.pretty_text
+        else:
+            with open(file, "w") as f:
+                f.write(self.pretty_text)
+
+    def merge_from_dict(self, options):
+        """Merge list into cfg_dict
+
+        Merge the dict parsed by MultipleKVAction into this cfg.
+
+        Examples:
+            >>> options = {'model.backbone.depth': 50,
+            ...            'model.backbone.with_cp':True}
+            >>> cfg = Config(dict(model=dict(backbone=dict(type='ResNet'))))
+            >>> cfg.merge_from_dict(options)
+            >>> cfg_dict = super(Config, self).__getattribute__('_cfg_dict')
+            >>> assert cfg_dict == dict(
+            ...     model=dict(backbone=dict(depth=50, with_cp=True)))
+
+        Args:
+            options (dict): dict of configs to merge from.
+        """
+        option_cfg_dict = {}
+        for full_key, v in options.items():
+            d = option_cfg_dict
+            key_list = full_key.split(".")
+            for subkey in key_list[:-1]:
+                d.setdefault(subkey, ConfigDict())
+                d = d[subkey]
+            subkey = key_list[-1]
+            d[subkey] = v
+
+        cfg_dict = super(SLConfig, self).__getattribute__("_cfg_dict")
+        super(SLConfig, self).__setattr__(
+            "_cfg_dict", SLConfig._merge_a_into_b(option_cfg_dict, cfg_dict)
+        )
+
+    # for multiprocess
+    def __setstate__(self, state):
+        self.__init__(state)
+
+    def copy(self):
+        return SLConfig(self._cfg_dict.copy())
+
+    def deepcopy(self):
+        return SLConfig(self._cfg_dict.deepcopy())
+
+
+class DictAction(Action):
+    """
+    argparse action to split an argument into KEY=VALUE form
+    on the first = and append to a dictionary. List options should
+    be passed as comma separated values, i.e KEY=V1,V2,V3
+    """
+
+    @staticmethod
+    def _parse_int_float_bool(val):
+        try:
+            return int(val)
+        except ValueError:
+            pass
+        try:
+            return float(val)
+        except ValueError:
+            pass
+        if val.lower() in ["true", "false"]:
+            return True if val.lower() == "true" else False
+        if val.lower() in ["none", "null"]:
+            return None
+        return val
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        options = {}
+        for kv in values:
+            key, val = kv.split("=", maxsplit=1)
+            val = [self._parse_int_float_bool(v) for v in val.split(",")]
+            if len(val) == 1:
+                val = val[0]
+            options[key] = val
+        setattr(namespace, self.dest, options)

groundingdino/util/slio.py

@@ -0,0 +1,177 @@
+# ==========================================================
+# Modified from mmcv
+# ==========================================================
+
+import json
+import pickle
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+import yaml
+
+try:
+    from yaml import CLoader as Loader, CDumper as Dumper
+except ImportError:
+    from yaml import Loader, Dumper
+
+
+# ===========================
+# Rigister handler
+# ===========================
+
+
+class BaseFileHandler(metaclass=ABCMeta):
+    @abstractmethod
+    def load_from_fileobj(self, file, **kwargs):
+        pass
+
+    @abstractmethod
+    def dump_to_fileobj(self, obj, file, **kwargs):
+        pass
+
+    @abstractmethod
+    def dump_to_str(self, obj, **kwargs):
+        pass
+
+    def load_from_path(self, filepath, mode="r", **kwargs):
+        with open(filepath, mode) as f:
+            return self.load_from_fileobj(f, **kwargs)
+
+    def dump_to_path(self, obj, filepath, mode="w", **kwargs):
+        with open(filepath, mode) as f:
+            self.dump_to_fileobj(obj, f, **kwargs)
+
+
+class JsonHandler(BaseFileHandler):
+    def load_from_fileobj(self, file):
+        return json.load(file)
+
+    def dump_to_fileobj(self, obj, file, **kwargs):
+        json.dump(obj, file, **kwargs)
+
+    def dump_to_str(self, obj, **kwargs):
+        return json.dumps(obj, **kwargs)
+
+
+class PickleHandler(BaseFileHandler):
+    def load_from_fileobj(self, file, **kwargs):
+        return pickle.load(file, **kwargs)
+
+    def load_from_path(self, filepath, **kwargs):
+        return super(PickleHandler, self).load_from_path(filepath, mode="rb", **kwargs)
+
+    def dump_to_str(self, obj, **kwargs):
+        kwargs.setdefault("protocol", 2)
+        return pickle.dumps(obj, **kwargs)
+
+    def dump_to_fileobj(self, obj, file, **kwargs):
+        kwargs.setdefault("protocol", 2)
+        pickle.dump(obj, file, **kwargs)
+
+    def dump_to_path(self, obj, filepath, **kwargs):
+        super(PickleHandler, self).dump_to_path(obj, filepath, mode="wb", **kwargs)
+
+
+class YamlHandler(BaseFileHandler):
+    def load_from_fileobj(self, file, **kwargs):
+        kwargs.setdefault("Loader", Loader)
+        return yaml.load(file, **kwargs)
+
+    def dump_to_fileobj(self, obj, file, **kwargs):
+        kwargs.setdefault("Dumper", Dumper)
+        yaml.dump(obj, file, **kwargs)
+
+    def dump_to_str(self, obj, **kwargs):
+        kwargs.setdefault("Dumper", Dumper)
+        return yaml.dump(obj, **kwargs)
+
+
+file_handlers = {
+    "json": JsonHandler(),
+    "yaml": YamlHandler(),
+    "yml": YamlHandler(),
+    "pickle": PickleHandler(),
+    "pkl": PickleHandler(),
+}
+
+# ===========================
+# load and dump
+# ===========================
+
+
+def is_str(x):
+    """Whether the input is an string instance.
+
+    Note: This method is deprecated since python 2 is no longer supported.
+    """
+    return isinstance(x, str)
+
+
+def slload(file, file_format=None, **kwargs):
+    """Load data from json/yaml/pickle files.
+
+    This method provides a unified api for loading data from serialized files.
+
+    Args:
+        file (str or :obj:`Path` or file-like object): Filename or a file-like
+            object.
+        file_format (str, optional): If not specified, the file format will be
+            inferred from the file extension, otherwise use the specified one.
+            Currently supported formats include "json", "yaml/yml" and
+            "pickle/pkl".
+
+    Returns:
+        The content from the file.
+    """
+    if isinstance(file, Path):
+        file = str(file)
+    if file_format is None and is_str(file):
+        file_format = file.split(".")[-1]
+    if file_format not in file_handlers:
+        raise TypeError(f"Unsupported format: {file_format}")
+
+    handler = file_handlers[file_format]
+    if is_str(file):
+        obj = handler.load_from_path(file, **kwargs)
+    elif hasattr(file, "read"):
+        obj = handler.load_from_fileobj(file, **kwargs)
+    else:
+        raise TypeError('"file" must be a filepath str or a file-object')
+    return obj
+
+
+def sldump(obj, file=None, file_format=None, **kwargs):
+    """Dump data to json/yaml/pickle strings or files.
+
+    This method provides a unified api for dumping data as strings or to files,
+    and also supports custom arguments for each file format.
+
+    Args:
+        obj (any): The python object to be dumped.
+        file (str or :obj:`Path` or file-like object, optional): If not
+            specified, then the object is dump to a str, otherwise to a file
+            specified by the filename or file-like object.
+        file_format (str, optional): Same as :func:`load`.
+
+    Returns:
+        bool: True for success, False otherwise.
+    """
+    if isinstance(file, Path):
+        file = str(file)
+    if file_format is None:
+        if is_str(file):
+            file_format = file.split(".")[-1]
+        elif file is None:
+            raise ValueError("file_format must be specified since file is None")
+    if file_format not in file_handlers:
+        raise TypeError(f"Unsupported format: {file_format}")
+
+    handler = file_handlers[file_format]
+    if file is None:
+        return handler.dump_to_str(obj, **kwargs)
+    elif is_str(file):
+        handler.dump_to_path(obj, file, **kwargs)
+    elif hasattr(file, "write"):
+        handler.dump_to_fileobj(obj, file, **kwargs)
+    else:
+        raise TypeError('"file" must be a filename str or a file-object')

groundingdino/util/time_counter.py

@@ -0,0 +1,62 @@
+import json
+import time
+
+
+class TimeCounter:
+    def __init__(self) -> None:
+        pass
+
+    def clear(self):
+        self.timedict = {}
+        self.basetime = time.perf_counter()
+
+    def timeit(self, name):
+        nowtime = time.perf_counter() - self.basetime
+        self.timedict[name] = nowtime
+        self.basetime = time.perf_counter()
+
+
+class TimeHolder:
+    def __init__(self) -> None:
+        self.timedict = {}
+
+    def update(self, _timedict: dict):
+        for k, v in _timedict.items():
+            if k not in self.timedict:
+                self.timedict[k] = AverageMeter(name=k, val_only=True)
+            self.timedict[k].update(val=v)
+
+    def final_res(self):
+        return {k: v.avg for k, v in self.timedict.items()}
+
+    def __str__(self):
+        return json.dumps(self.final_res(), indent=2)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, name, fmt=":f", val_only=False):
+        self.name = name
+        self.fmt = fmt
+        self.val_only = val_only
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        if self.val_only:
+            fmtstr = "{name} {val" + self.fmt + "}"
+        else:
+            fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})"
+        return fmtstr.format(**self.__dict__)

groundingdino/util/utils.py

@@ -0,0 +1,610 @@
+import argparse
+import json
+import warnings
+from collections import OrderedDict
+from copy import deepcopy
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer
+
+from groundingdino.util.slconfig import SLConfig
+
+
+def slprint(x, name="x"):
+    if isinstance(x, (torch.Tensor, np.ndarray)):
+        print(f"{name}.shape:", x.shape)
+    elif isinstance(x, (tuple, list)):
+        print("type x:", type(x))
+        for i in range(min(10, len(x))):
+            slprint(x[i], f"{name}[{i}]")
+    elif isinstance(x, dict):
+        for k, v in x.items():
+            slprint(v, f"{name}[{k}]")
+    else:
+        print(f"{name}.type:", type(x))
+
+
+def clean_state_dict(state_dict):
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        if k[:7] == "module.":
+            k = k[7:]  # remove `module.`
+        new_state_dict[k] = v
+    return new_state_dict
+
+
+def renorm(
+    img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+) -> torch.FloatTensor:
+    # img: tensor(3,H,W) or tensor(B,3,H,W)
+    # return: same as img
+    assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim()
+    if img.dim() == 3:
+        assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % (
+            img.size(0),
+            str(img.size()),
+        )
+        img_perm = img.permute(1, 2, 0)
+        mean = torch.Tensor(mean)
+        std = torch.Tensor(std)
+        img_res = img_perm * std + mean
+        return img_res.permute(2, 0, 1)
+    else:  # img.dim() == 4
+        assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % (
+            img.size(1),
+            str(img.size()),
+        )
+        img_perm = img.permute(0, 2, 3, 1)
+        mean = torch.Tensor(mean)
+        std = torch.Tensor(std)
+        img_res = img_perm * std + mean
+        return img_res.permute(0, 3, 1, 2)
+
+
+class CocoClassMapper:
+    def __init__(self) -> None:
+        self.category_map_str = {
+            "1": 1,
+            "2": 2,
+            "3": 3,
+            "4": 4,
+            "5": 5,
+            "6": 6,
+            "7": 7,
+            "8": 8,
+            "9": 9,
+            "10": 10,
+            "11": 11,
+            "13": 12,
+            "14": 13,
+            "15": 14,
+            "16": 15,
+            "17": 16,
+            "18": 17,
+            "19": 18,
+            "20": 19,
+            "21": 20,
+            "22": 21,
+            "23": 22,
+            "24": 23,
+            "25": 24,
+            "27": 25,
+            "28": 26,
+            "31": 27,
+            "32": 28,
+            "33": 29,
+            "34": 30,
+            "35": 31,
+            "36": 32,
+            "37": 33,
+            "38": 34,
+            "39": 35,
+            "40": 36,
+            "41": 37,
+            "42": 38,
+            "43": 39,
+            "44": 40,
+            "46": 41,
+            "47": 42,
+            "48": 43,
+            "49": 44,
+            "50": 45,
+            "51": 46,
+            "52": 47,
+            "53": 48,
+            "54": 49,
+            "55": 50,
+            "56": 51,
+            "57": 52,
+            "58": 53,
+            "59": 54,
+            "60": 55,
+            "61": 56,
+            "62": 57,
+            "63": 58,
+            "64": 59,
+            "65": 60,
+            "67": 61,
+            "70": 62,
+            "72": 63,
+            "73": 64,
+            "74": 65,
+            "75": 66,
+            "76": 67,
+            "77": 68,
+            "78": 69,
+            "79": 70,
+            "80": 71,
+            "81": 72,
+            "82": 73,
+            "84": 74,
+            "85": 75,
+            "86": 76,
+            "87": 77,
+            "88": 78,
+            "89": 79,
+            "90": 80,
+        }
+        self.origin2compact_mapper = {int(k): v - 1 for k, v in self.category_map_str.items()}
+        self.compact2origin_mapper = {int(v - 1): int(k) for k, v in self.category_map_str.items()}
+
+    def origin2compact(self, idx):
+        return self.origin2compact_mapper[int(idx)]
+
+    def compact2origin(self, idx):
+        return self.compact2origin_mapper[int(idx)]
+
+
+def to_device(item, device):
+    if isinstance(item, torch.Tensor):
+        return item.to(device)
+    elif isinstance(item, list):
+        return [to_device(i, device) for i in item]
+    elif isinstance(item, dict):
+        return {k: to_device(v, device) for k, v in item.items()}
+    else:
+        raise NotImplementedError(
+            "Call Shilong if you use other containers! type: {}".format(type(item))
+        )
+
+
+#
+def get_gaussian_mean(x, axis, other_axis, softmax=True):
+    """
+
+    Args:
+        x (float): Input images(BxCxHxW)
+        axis (int): The index for weighted mean
+        other_axis (int): The other index
+
+    Returns: weighted index for axis, BxC
+
+    """
+    mat2line = torch.sum(x, axis=other_axis)
+    # mat2line = mat2line / mat2line.mean() * 10
+    if softmax:
+        u = torch.softmax(mat2line, axis=2)
+    else:
+        u = mat2line / (mat2line.sum(2, keepdim=True) + 1e-6)
+    size = x.shape[axis]
+    ind = torch.linspace(0, 1, size).to(x.device)
+    batch = x.shape[0]
+    channel = x.shape[1]
+    index = ind.repeat([batch, channel, 1])
+    mean_position = torch.sum(index * u, dim=2)
+    return mean_position
+
+
+def get_expected_points_from_map(hm, softmax=True):
+    """get_gaussian_map_from_points
+        B,C,H,W -> B,N,2 float(0, 1) float(0, 1)
+        softargmax function
+
+    Args:
+        hm (float): Input images(BxCxHxW)
+
+    Returns:
+        weighted index for axis, BxCx2. float between 0 and 1.
+
+    """
+    # hm = 10*hm
+    B, C, H, W = hm.shape
+    y_mean = get_gaussian_mean(hm, 2, 3, softmax=softmax)  # B,C
+    x_mean = get_gaussian_mean(hm, 3, 2, softmax=softmax)  # B,C
+    # return torch.cat((x_mean.unsqueeze(-1), y_mean.unsqueeze(-1)), 2)
+    return torch.stack([x_mean, y_mean], dim=2)
+
+
+# Positional encoding (section 5.1)
+# borrow from nerf
+class Embedder:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+        self.create_embedding_fn()
+
+    def create_embedding_fn(self):
+        embed_fns = []
+        d = self.kwargs["input_dims"]
+        out_dim = 0
+        if self.kwargs["include_input"]:
+            embed_fns.append(lambda x: x)
+            out_dim += d
+
+        max_freq = self.kwargs["max_freq_log2"]
+        N_freqs = self.kwargs["num_freqs"]
+
+        if self.kwargs["log_sampling"]:
+            freq_bands = 2.0 ** torch.linspace(0.0, max_freq, steps=N_freqs)
+        else:
+            freq_bands = torch.linspace(2.0**0.0, 2.0**max_freq, steps=N_freqs)
+
+        for freq in freq_bands:
+            for p_fn in self.kwargs["periodic_fns"]:
+                embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
+                out_dim += d
+
+        self.embed_fns = embed_fns
+        self.out_dim = out_dim
+
+    def embed(self, inputs):
+        return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
+
+
+def get_embedder(multires, i=0):
+    import torch.nn as nn
+
+    if i == -1:
+        return nn.Identity(), 3
+
+    embed_kwargs = {
+        "include_input": True,
+        "input_dims": 3,
+        "max_freq_log2": multires - 1,
+        "num_freqs": multires,
+        "log_sampling": True,
+        "periodic_fns": [torch.sin, torch.cos],
+    }
+
+    embedder_obj = Embedder(**embed_kwargs)
+    embed = lambda x, eo=embedder_obj: eo.embed(x)
+    return embed, embedder_obj.out_dim
+
+
+class APOPMeter:
+    def __init__(self) -> None:
+        self.tp = 0
+        self.fp = 0
+        self.tn = 0
+        self.fn = 0
+
+    def update(self, pred, gt):
+        """
+        Input:
+            pred, gt: Tensor()
+        """
+        assert pred.shape == gt.shape
+        self.tp += torch.logical_and(pred == 1, gt == 1).sum().item()
+        self.fp += torch.logical_and(pred == 1, gt == 0).sum().item()
+        self.tn += torch.logical_and(pred == 0, gt == 0).sum().item()
+        self.tn += torch.logical_and(pred == 1, gt == 0).sum().item()
+
+    def update_cm(self, tp, fp, tn, fn):
+        self.tp += tp
+        self.fp += fp
+        self.tn += tn
+        self.tn += fn
+
+
+def inverse_sigmoid(x, eps=1e-5):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1 / x2)
+
+
+def get_raw_dict(args):
+    """
+    return the dicf contained in args.
+
+    e.g:
+        >>> with open(path, 'w') as f:
+                json.dump(get_raw_dict(args), f, indent=2)
+    """
+    if isinstance(args, argparse.Namespace):
+        return vars(args)
+    elif isinstance(args, dict):
+        return args
+    elif isinstance(args, SLConfig):
+        return args._cfg_dict
+    else:
+        raise NotImplementedError("Unknown type {}".format(type(args)))
+
+
+def stat_tensors(tensor):
+    assert tensor.dim() == 1
+    tensor_sm = tensor.softmax(0)
+    entropy = (tensor_sm * torch.log(tensor_sm + 1e-9)).sum()
+
+    return {
+        "max": tensor.max(),
+        "min": tensor.min(),
+        "mean": tensor.mean(),
+        "var": tensor.var(),
+        "std": tensor.var() ** 0.5,
+        "entropy": entropy,
+    }
+
+
+class NiceRepr:
+    """Inherit from this class and define ``__nice__`` to "nicely" print your
+    objects.
+
+    Defines ``__str__`` and ``__repr__`` in terms of ``__nice__`` function
+    Classes that inherit from :class:`NiceRepr` should redefine ``__nice__``.
+    If the inheriting class has a ``__len__``, method then the default
+    ``__nice__`` method will return its length.
+
+    Example:
+        >>> class Foo(NiceRepr):
+        ...    def __nice__(self):
+        ...        return 'info'
+        >>> foo = Foo()
+        >>> assert str(foo) == '<Foo(info)>'
+        >>> assert repr(foo).startswith('<Foo(info) at ')
+
+    Example:
+        >>> class Bar(NiceRepr):
+        ...    pass
+        >>> bar = Bar()
+        >>> import pytest
+        >>> with pytest.warns(None) as record:
+        >>>     assert 'object at' in str(bar)
+        >>>     assert 'object at' in repr(bar)
+
+    Example:
+        >>> class Baz(NiceRepr):
+        ...    def __len__(self):
+        ...        return 5
+        >>> baz = Baz()
+        >>> assert str(baz) == '<Baz(5)>'
+    """
+
+    def __nice__(self):
+        """str: a "nice" summary string describing this module"""
+        if hasattr(self, "__len__"):
+            # It is a common pattern for objects to use __len__ in __nice__
+            # As a convenience we define a default __nice__ for these objects
+            return str(len(self))
+        else:
+            # In all other cases force the subclass to overload __nice__
+            raise NotImplementedError(f"Define the __nice__ method for {self.__class__!r}")
+
+    def __repr__(self):
+        """str: the string of the module"""
+        try:
+            nice = self.__nice__()
+            classname = self.__class__.__name__
+            return f"<{classname}({nice}) at {hex(id(self))}>"
+        except NotImplementedError as ex:
+            warnings.warn(str(ex), category=RuntimeWarning)
+            return object.__repr__(self)
+
+    def __str__(self):
+        """str: the string of the module"""
+        try:
+            classname = self.__class__.__name__
+            nice = self.__nice__()
+            return f"<{classname}({nice})>"
+        except NotImplementedError as ex:
+            warnings.warn(str(ex), category=RuntimeWarning)
+            return object.__repr__(self)
+
+
+def ensure_rng(rng=None):
+    """Coerces input into a random number generator.
+
+    If the input is None, then a global random state is returned.
+
+    If the input is a numeric value, then that is used as a seed to construct a
+    random state. Otherwise the input is returned as-is.
+
+    Adapted from [1]_.
+
+    Args:
+        rng (int | numpy.random.RandomState | None):
+            if None, then defaults to the global rng. Otherwise this can be an
+            integer or a RandomState class
+    Returns:
+        (numpy.random.RandomState) : rng -
+            a numpy random number generator
+
+    References:
+        .. [1] https://gitlab.kitware.com/computer-vision/kwarray/blob/master/kwarray/util_random.py#L270  # noqa: E501
+    """
+
+    if rng is None:
+        rng = np.random.mtrand._rand
+    elif isinstance(rng, int):
+        rng = np.random.RandomState(rng)
+    else:
+        rng = rng
+    return rng
+
+
+def random_boxes(num=1, scale=1, rng=None):
+    """Simple version of ``kwimage.Boxes.random``
+
+    Returns:
+        Tensor: shape (n, 4) in x1, y1, x2, y2 format.
+
+    References:
+        https://gitlab.kitware.com/computer-vision/kwimage/blob/master/kwimage/structs/boxes.py#L1390
+
+    Example:
+        >>> num = 3
+        >>> scale = 512
+        >>> rng = 0
+        >>> boxes = random_boxes(num, scale, rng)
+        >>> print(boxes)
+        tensor([[280.9925, 278.9802, 308.6148, 366.1769],
+                [216.9113, 330.6978, 224.0446, 456.5878],
+                [405.3632, 196.3221, 493.3953, 270.7942]])
+    """
+    rng = ensure_rng(rng)
+
+    tlbr = rng.rand(num, 4).astype(np.float32)
+
+    tl_x = np.minimum(tlbr[:, 0], tlbr[:, 2])
+    tl_y = np.minimum(tlbr[:, 1], tlbr[:, 3])
+    br_x = np.maximum(tlbr[:, 0], tlbr[:, 2])
+    br_y = np.maximum(tlbr[:, 1], tlbr[:, 3])
+
+    tlbr[:, 0] = tl_x * scale
+    tlbr[:, 1] = tl_y * scale
+    tlbr[:, 2] = br_x * scale
+    tlbr[:, 3] = br_y * scale
+
+    boxes = torch.from_numpy(tlbr)
+    return boxes
+
+
+class ModelEma(torch.nn.Module):
+    def __init__(self, model, decay=0.9997, device=None):
+        super(ModelEma, self).__init__()
+        # make a copy of the model for accumulating moving average of weights
+        self.module = deepcopy(model)
+        self.module.eval()
+
+        # import ipdb; ipdb.set_trace()
+
+        self.decay = decay
+        self.device = device  # perform ema on different device from model if set
+        if self.device is not None:
+            self.module.to(device=device)
+
+    def _update(self, model, update_fn):
+        with torch.no_grad():
+            for ema_v, model_v in zip(
+                self.module.state_dict().values(), model.state_dict().values()
+            ):
+                if self.device is not None:
+                    model_v = model_v.to(device=self.device)
+                ema_v.copy_(update_fn(ema_v, model_v))
+
+    def update(self, model):
+        self._update(model, update_fn=lambda e, m: self.decay * e + (1.0 - self.decay) * m)
+
+    def set(self, model):
+        self._update(model, update_fn=lambda e, m: m)
+
+
+class BestMetricSingle:
+    def __init__(self, init_res=0.0, better="large") -> None:
+        self.init_res = init_res
+        self.best_res = init_res
+        self.best_ep = -1
+
+        self.better = better
+        assert better in ["large", "small"]
+
+    def isbetter(self, new_res, old_res):
+        if self.better == "large":
+            return new_res > old_res
+        if self.better == "small":
+            return new_res < old_res
+
+    def update(self, new_res, ep):
+        if self.isbetter(new_res, self.best_res):
+            self.best_res = new_res
+            self.best_ep = ep
+            return True
+        return False
+
+    def __str__(self) -> str:
+        return "best_res: {}\t best_ep: {}".format(self.best_res, self.best_ep)
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def summary(self) -> dict:
+        return {
+            "best_res": self.best_res,
+            "best_ep": self.best_ep,
+        }
+
+
+class BestMetricHolder:
+    def __init__(self, init_res=0.0, better="large", use_ema=False) -> None:
+        self.best_all = BestMetricSingle(init_res, better)
+        self.use_ema = use_ema
+        if use_ema:
+            self.best_ema = BestMetricSingle(init_res, better)
+            self.best_regular = BestMetricSingle(init_res, better)
+
+    def update(self, new_res, epoch, is_ema=False):
+        """
+        return if the results is the best.
+        """
+        if not self.use_ema:
+            return self.best_all.update(new_res, epoch)
+        else:
+            if is_ema:
+                self.best_ema.update(new_res, epoch)
+                return self.best_all.update(new_res, epoch)
+            else:
+                self.best_regular.update(new_res, epoch)
+                return self.best_all.update(new_res, epoch)
+
+    def summary(self):
+        if not self.use_ema:
+            return self.best_all.summary()
+
+        res = {}
+        res.update({f"all_{k}": v for k, v in self.best_all.summary().items()})
+        res.update({f"regular_{k}": v for k, v in self.best_regular.summary().items()})
+        res.update({f"ema_{k}": v for k, v in self.best_ema.summary().items()})
+        return res
+
+    def __repr__(self) -> str:
+        return json.dumps(self.summary(), indent=2)
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+
+def targets_to(targets: List[Dict[str, Any]], device):
+    """Moves the target dicts to the given device."""
+    excluded_keys = [
+        "questionId",
+        "tokens_positive",
+        "strings_positive",
+        "tokens",
+        "dataset_name",
+        "sentence_id",
+        "original_img_id",
+        "nb_eval",
+        "task_id",
+        "original_id",
+        "token_span",
+        "caption",
+        "dataset_type",
+    ]
+    return [
+        {k: v.to(device) if k not in excluded_keys else v for k, v in t.items()} for t in targets
+    ]
+
+
+def get_phrases_from_posmap(
+    posmap: torch.BoolTensor, tokenized: Dict, tokenizer: AutoTokenizer, left_idx: int = 0, right_idx: int = 255
+):
+    assert isinstance(posmap, torch.Tensor), "posmap must be torch.Tensor"
+    if posmap.dim() == 1:
+        posmap[0: left_idx + 1] = False
+        posmap[right_idx:] = False
+        non_zero_idx = posmap.nonzero(as_tuple=True)[0].tolist()
+        token_ids = [tokenized["input_ids"][i] for i in non_zero_idx]
+        return tokenizer.decode(token_ids)
+    else:
+        raise NotImplementedError("posmap must be 1-dim")

groundingdino/util/visualizer.py

@@ -0,0 +1,318 @@
+# -*- coding: utf-8 -*-
+"""
+@File    :   visualizer.py
+@Time    :   2022/04/05 11:39:33
+@Author  :   Shilong Liu 
+@Contact :   slongliu86@gmail.com
+"""
+
+import datetime
+import os
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib import transforms
+from matplotlib.collections import PatchCollection
+from matplotlib.patches import Polygon
+from pycocotools import mask as maskUtils
+
+
+def renorm(
+    img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+) -> torch.FloatTensor:
+    # img: tensor(3,H,W) or tensor(B,3,H,W)
+    # return: same as img
+    assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim()
+    if img.dim() == 3:
+        assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % (
+            img.size(0),
+            str(img.size()),
+        )
+        img_perm = img.permute(1, 2, 0)
+        mean = torch.Tensor(mean)
+        std = torch.Tensor(std)
+        img_res = img_perm * std + mean
+        return img_res.permute(2, 0, 1)
+    else:  # img.dim() == 4
+        assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % (
+            img.size(1),
+            str(img.size()),
+        )
+        img_perm = img.permute(0, 2, 3, 1)
+        mean = torch.Tensor(mean)
+        std = torch.Tensor(std)
+        img_res = img_perm * std + mean
+        return img_res.permute(0, 3, 1, 2)
+
+
+class ColorMap:
+    def __init__(self, basergb=[255, 255, 0]):
+        self.basergb = np.array(basergb)
+
+    def __call__(self, attnmap):
+        # attnmap: h, w. np.uint8.
+        # return: h, w, 4. np.uint8.
+        assert attnmap.dtype == np.uint8
+        h, w = attnmap.shape
+        res = self.basergb.copy()
+        res = res[None][None].repeat(h, 0).repeat(w, 1)  # h, w, 3
+        attn1 = attnmap.copy()[..., None]  # h, w, 1
+        res = np.concatenate((res, attn1), axis=-1).astype(np.uint8)
+        return res
+
+
+def rainbow_text(x, y, ls, lc, **kw):
+    """
+    Take a list of strings ``ls`` and colors ``lc`` and place them next to each
+    other, with text ls[i] being shown in color lc[i].
+
+    This example shows how to do both vertical and horizontal text, and will
+    pass all keyword arguments to plt.text, so you can set the font size,
+    family, etc.
+    """
+    t = plt.gca().transData
+    fig = plt.gcf()
+    plt.show()
+
+    # horizontal version
+    for s, c in zip(ls, lc):
+        text = plt.text(x, y, " " + s + " ", color=c, transform=t, **kw)
+        text.draw(fig.canvas.get_renderer())
+        ex = text.get_window_extent()
+        t = transforms.offset_copy(text._transform, x=ex.width, units="dots")
+
+    # #vertical version
+    # for s,c in zip(ls,lc):
+    #     text = plt.text(x,y," "+s+" ",color=c, transform=t,
+    #             rotation=90,va='bottom',ha='center',**kw)
+    #     text.draw(fig.canvas.get_renderer())
+    #     ex = text.get_window_extent()
+    #     t = transforms.offset_copy(text._transform, y=ex.height, units='dots')
+
+
+class COCOVisualizer:
+    def __init__(self, coco=None, tokenlizer=None) -> None:
+        self.coco = coco
+
+    def visualize(self, img, tgt, caption=None, dpi=180, savedir="vis"):
+        """
+        img: tensor(3, H, W)
+        tgt: make sure they are all on cpu.
+            must have items: 'image_id', 'boxes', 'size'
+        """
+        plt.figure(dpi=dpi)
+        plt.rcParams["font.size"] = "5"
+        ax = plt.gca()
+        img = renorm(img).permute(1, 2, 0)
+        # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO':
+        #     import ipdb; ipdb.set_trace()
+        ax.imshow(img)
+
+        self.addtgt(tgt)
+
+        if tgt is None:
+            image_id = 0
+        elif "image_id" not in tgt:
+            image_id = 0
+        else:
+            image_id = tgt["image_id"]
+
+        if caption is None:
+            savename = "{}/{}-{}.png".format(
+                savedir, int(image_id), str(datetime.datetime.now()).replace(" ", "-")
+            )
+        else:
+            savename = "{}/{}-{}-{}.png".format(
+                savedir, caption, int(image_id), str(datetime.datetime.now()).replace(" ", "-")
+            )
+        print("savename: {}".format(savename))
+        os.makedirs(os.path.dirname(savename), exist_ok=True)
+        plt.savefig(savename)
+        plt.close()
+
+    def addtgt(self, tgt):
+        """ """
+        if tgt is None or not "boxes" in tgt:
+            ax = plt.gca()
+
+            if "caption" in tgt:
+                ax.set_title(tgt["caption"], wrap=True)
+
+            ax.set_axis_off()
+            return
+
+        ax = plt.gca()
+        H, W = tgt["size"]
+        numbox = tgt["boxes"].shape[0]
+
+        color = []
+        polygons = []
+        boxes = []
+        for box in tgt["boxes"].cpu():
+            unnormbbox = box * torch.Tensor([W, H, W, H])
+            unnormbbox[:2] -= unnormbbox[2:] / 2
+            [bbox_x, bbox_y, bbox_w, bbox_h] = unnormbbox.tolist()
+            boxes.append([bbox_x, bbox_y, bbox_w, bbox_h])
+            poly = [
+                [bbox_x, bbox_y],
+                [bbox_x, bbox_y + bbox_h],
+                [bbox_x + bbox_w, bbox_y + bbox_h],
+                [bbox_x + bbox_w, bbox_y],
+            ]
+            np_poly = np.array(poly).reshape((4, 2))
+            polygons.append(Polygon(np_poly))
+            c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0]
+            color.append(c)
+
+        p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.1)
+        ax.add_collection(p)
+        p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2)
+        ax.add_collection(p)
+
+        if "strings_positive" in tgt and len(tgt["strings_positive"]) > 0:
+            assert (
+                len(tgt["strings_positive"]) == numbox
+            ), f"{len(tgt['strings_positive'])} = {numbox}, "
+            for idx, strlist in enumerate(tgt["strings_positive"]):
+                cate_id = int(tgt["labels"][idx])
+                _string = str(cate_id) + ":" + " ".join(strlist)
+                bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx]
+                # ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1})
+                ax.text(
+                    bbox_x,
+                    bbox_y,
+                    _string,
+                    color="black",
+                    bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1},
+                )
+
+        if "box_label" in tgt:
+            assert len(tgt["box_label"]) == numbox, f"{len(tgt['box_label'])} = {numbox}, "
+            for idx, bl in enumerate(tgt["box_label"]):
+                _string = str(bl)
+                bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx]
+                # ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1})
+                ax.text(
+                    bbox_x,
+                    bbox_y,
+                    _string,
+                    color="black",
+                    bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1},
+                )
+
+        if "caption" in tgt:
+            ax.set_title(tgt["caption"], wrap=True)
+            # plt.figure()
+            # rainbow_text(0.0,0.0,"all unicorns poop rainbows ! ! !".split(),
+            #         ['red', 'orange', 'brown', 'green', 'blue', 'purple', 'black'])
+
+        if "attn" in tgt:
+            # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO':
+            #     import ipdb; ipdb.set_trace()
+            if isinstance(tgt["attn"], tuple):
+                tgt["attn"] = [tgt["attn"]]
+            for item in tgt["attn"]:
+                attn_map, basergb = item
+                attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-3)
+                attn_map = (attn_map * 255).astype(np.uint8)
+                cm = ColorMap(basergb)
+                heatmap = cm(attn_map)
+                ax.imshow(heatmap)
+        ax.set_axis_off()
+
+    def showAnns(self, anns, draw_bbox=False):
+        """
+        Display the specified annotations.
+        :param anns (array of object): annotations to display
+        :return: None
+        """
+        if len(anns) == 0:
+            return 0
+        if "segmentation" in anns[0] or "keypoints" in anns[0]:
+            datasetType = "instances"
+        elif "caption" in anns[0]:
+            datasetType = "captions"
+        else:
+            raise Exception("datasetType not supported")
+        if datasetType == "instances":
+            ax = plt.gca()
+            ax.set_autoscale_on(False)
+            polygons = []
+            color = []
+            for ann in anns:
+                c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0]
+                if "segmentation" in ann:
+                    if type(ann["segmentation"]) == list:
+                        # polygon
+                        for seg in ann["segmentation"]:
+                            poly = np.array(seg).reshape((int(len(seg) / 2), 2))
+                            polygons.append(Polygon(poly))
+                            color.append(c)
+                    else:
+                        # mask
+                        t = self.imgs[ann["image_id"]]
+                        if type(ann["segmentation"]["counts"]) == list:
+                            rle = maskUtils.frPyObjects(
+                                [ann["segmentation"]], t["height"], t["width"]
+                            )
+                        else:
+                            rle = [ann["segmentation"]]
+                        m = maskUtils.decode(rle)
+                        img = np.ones((m.shape[0], m.shape[1], 3))
+                        if ann["iscrowd"] == 1:
+                            color_mask = np.array([2.0, 166.0, 101.0]) / 255
+                        if ann["iscrowd"] == 0:
+                            color_mask = np.random.random((1, 3)).tolist()[0]
+                        for i in range(3):
+                            img[:, :, i] = color_mask[i]
+                        ax.imshow(np.dstack((img, m * 0.5)))
+                if "keypoints" in ann and type(ann["keypoints"]) == list:
+                    # turn skeleton into zero-based index
+                    sks = np.array(self.loadCats(ann["category_id"])[0]["skeleton"]) - 1
+                    kp = np.array(ann["keypoints"])
+                    x = kp[0::3]
+                    y = kp[1::3]
+                    v = kp[2::3]
+                    for sk in sks:
+                        if np.all(v[sk] > 0):
+                            plt.plot(x[sk], y[sk], linewidth=3, color=c)
+                    plt.plot(
+                        x[v > 0],
+                        y[v > 0],
+                        "o",
+                        markersize=8,
+                        markerfacecolor=c,
+                        markeredgecolor="k",
+                        markeredgewidth=2,
+                    )
+                    plt.plot(
+                        x[v > 1],
+                        y[v > 1],
+                        "o",
+                        markersize=8,
+                        markerfacecolor=c,
+                        markeredgecolor=c,
+                        markeredgewidth=2,
+                    )
+
+                if draw_bbox:
+                    [bbox_x, bbox_y, bbox_w, bbox_h] = ann["bbox"]
+                    poly = [
+                        [bbox_x, bbox_y],
+                        [bbox_x, bbox_y + bbox_h],
+                        [bbox_x + bbox_w, bbox_y + bbox_h],
+                        [bbox_x + bbox_w, bbox_y],
+                    ]
+                    np_poly = np.array(poly).reshape((4, 2))
+                    polygons.append(Polygon(np_poly))
+                    color.append(c)
+
+            # p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4)
+            # ax.add_collection(p)
+            p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2)
+            ax.add_collection(p)
+        elif datasetType == "captions":
+            for ann in anns:
+                print(ann["caption"])

groundingdino/util/vl_utils.py

@@ -0,0 +1,100 @@
+import os
+import random
+from typing import List
+
+import torch
+
+
+def create_positive_map_from_span(tokenized, token_span, max_text_len=256):
+    """construct a map such that positive_map[i,j] = True iff box i is associated to token j
+    Input:
+        - tokenized:
+            - input_ids: Tensor[1, ntokens]
+            - attention_mask: Tensor[1, ntokens]
+        - token_span: list with length num_boxes.
+            - each item: [start_idx, end_idx]
+    """
+    positive_map = torch.zeros((len(token_span), max_text_len), dtype=torch.float)
+    for j, tok_list in enumerate(token_span):
+        for (beg, end) in tok_list:
+            beg_pos = tokenized.char_to_token(beg)
+            end_pos = tokenized.char_to_token(end - 1)
+            if beg_pos is None:
+                try:
+                    beg_pos = tokenized.char_to_token(beg + 1)
+                    if beg_pos is None:
+                        beg_pos = tokenized.char_to_token(beg + 2)
+                except:
+                    beg_pos = None
+            if end_pos is None:
+                try:
+                    end_pos = tokenized.char_to_token(end - 2)
+                    if end_pos is None:
+                        end_pos = tokenized.char_to_token(end - 3)
+                except:
+                    end_pos = None
+            if beg_pos is None or end_pos is None:
+                continue
+
+            assert beg_pos is not None and end_pos is not None
+            if os.environ.get("SHILONG_DEBUG_ONLY_ONE_POS", None) == "TRUE":
+                positive_map[j, beg_pos] = 1
+                break
+            else:
+                positive_map[j, beg_pos : end_pos + 1].fill_(1)
+
+    return positive_map / (positive_map.sum(-1)[:, None] + 1e-6)
+
+
+def build_captions_and_token_span(cat_list, force_lowercase):
+    """
+    Return:
+        captions: str
+        cat2tokenspan: dict
+            {
+                'dog': [[0, 2]],
+                ...
+            }
+    """
+
+    cat2tokenspan = {}
+    captions = ""
+    for catname in cat_list:
+        class_name = catname
+        if force_lowercase:
+            class_name = class_name.lower()
+        if "/" in class_name:
+            class_name_list: List = class_name.strip().split("/")
+            class_name_list.append(class_name)
+            class_name: str = random.choice(class_name_list)
+
+        tokens_positive_i = []
+        subnamelist = [i.strip() for i in class_name.strip().split(" ")]
+        for subname in subnamelist:
+            if len(subname) == 0:
+                continue
+            if len(captions) > 0:
+                captions = captions + " "
+            strat_idx = len(captions)
+            end_idx = strat_idx + len(subname)
+            tokens_positive_i.append([strat_idx, end_idx])
+            captions = captions + subname
+
+        if len(tokens_positive_i) > 0:
+            captions = captions + " ."
+            cat2tokenspan[class_name] = tokens_positive_i
+
+    return captions, cat2tokenspan
+
+
+def build_id2posspan_and_caption(category_dict: dict):
+    """Build id2pos_span and caption from category_dict
+
+    Args:
+        category_dict (dict): category_dict
+    """
+    cat_list = [item["name"].lower() for item in category_dict]
+    id2catname = {item["id"]: item["name"].lower() for item in category_dict}
+    caption, cat2posspan = build_captions_and_token_span(cat_list, force_lowercase=True)
+    id2posspan = {catid: cat2posspan[catname] for catid, catname in id2catname.items()}
+    return id2posspan, caption

model.safetensors

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

models/GroundingDINO/__init__.py

@@ -0,0 +1,15 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Conditional DETR
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Copied from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+
+from .groundingdino import build_groundingdino

models/GroundingDINO/backbone/__init__.py

@@ -0,0 +1 @@
+from .backbone import build_backbone

models/GroundingDINO/backbone/backbone.py

@@ -0,0 +1,221 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Conditional DETR
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Copied from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+
+"""
+Backbone modules.
+"""
+
+from typing import Dict, List
+
+import torch
+import torch.nn.functional as F
+import torchvision
+from torch import nn
+from torchvision.models._utils import IntermediateLayerGetter
+
+from groundingdino.util.misc import NestedTensor, clean_state_dict, is_main_process
+
+from .position_encoding import build_position_encoding
+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_indices: list,
+    ):
+        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)
+
+        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)}
+            )
+
+        # if len:
+        #     if use_stage1_feature:
+        #         return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
+        #     else:
+        #         return_layers = {"layer2": "0", "layer3": "1", "layer4": "2"}
+        # else:
+        #     return_layers = {'layer4': "0"}
+        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)
+        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)
+        # import ipdb; ipdb.set_trace()
+        return out
+
+
+class Backbone(BackboneBase):
+    """ResNet backbone with frozen BatchNorm."""
+
+    def __init__(
+        self,
+        name: str,
+        train_backbone: bool,
+        dilation: bool,
+        return_interm_indices: list,
+        batch_norm=FrozenBatchNorm2d,
+    ):
+        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=batch_norm,
+            )
+        else:
+            raise NotImplementedError("Why you can get here with name {}".format(name))
+        # num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
+        assert name not in ("resnet18", "resnet34"), "Only resnet50 and resnet101 are available."
+        assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]
+        num_channels_all = [256, 512, 1024, 2048]
+        num_channels = num_channels_all[4 - len(return_interm_indices) :]
+        super().__init__(backbone, train_backbone, num_channels, return_interm_indices)
+
+
+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):
+    """
+    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
+
+    """
+    position_embedding = build_position_encoding(args)
+    train_backbone = True
+    if not train_backbone:
+        raise ValueError("Please set lr_backbone > 0")
+    return_interm_indices = args.return_interm_indices
+    assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]]
+    args.backbone_freeze_keywords
+    use_checkpoint = getattr(args, "use_checkpoint", False)
+
+    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=False,
+            use_checkpoint=use_checkpoint,
+        )
+
+        bb_num_channels = backbone.num_features[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))
+    # import ipdb; ipdb.set_trace()
+    return model

models/GroundingDINO/backbone/position_encoding.py

@@ -0,0 +1,186 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# DINO
+# Copyright (c) 2022 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Conditional DETR
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Copied from DETR (https://github.com/facebookresearch/detr)
+# 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 groundingdino.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
+            # if os.environ.get("SHILONG_AMP", None) == '1':
+            #     eps = 1e-4
+            # else:
+            #     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 PositionEmbeddingSineHW(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, temperatureH=10000, temperatureW=10000, normalize=False, scale=None
+    ):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperatureH = temperatureH
+        self.temperatureW = temperatureW
+        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)
+
+        # import ipdb; ipdb.set_trace()
+
+        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_tx = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_tx = self.temperatureW ** (2 * (torch.div(dim_tx, 2, rounding_mode='floor')) / self.num_pos_feats)
+        pos_x = x_embed[:, :, :, None] / dim_tx
+
+        dim_ty = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_ty = self.temperatureH ** (2 * (torch.div(dim_ty, 2, rounding_mode='floor')) / self.num_pos_feats)
+        pos_y = y_embed[:, :, :, None] / dim_ty
+
+        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)
+
+        # import ipdb; ipdb.set_trace()
+
+        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
+    if args.position_embedding in ("v2", "sine"):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSineHW(
+            N_steps,
+            temperatureH=args.pe_temperatureH,
+            temperatureW=args.pe_temperatureW,
+            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/GroundingDINO/backbone/swin_transformer.py

@@ -0,0 +1,804 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# DINO
+# Copyright (c) 2022 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# --------------------------------------------------------
+# modified from https://github.com/SwinTransformer/Swin-Transformer-Object-Detection/blob/master/mmdet/models/backbones/swin_transformer.py
+# --------------------------------------------------------
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+from groundingdino.util.misc import NestedTensor
+
+
+class Mlp(nn.Module):
+    """Multilayer perceptron."""
+
+    def __init__(
+        self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    """Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        dim,
+        window_size,
+        num_heads,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+        )  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """Forward function.
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B_, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+
+        relative_position_bias = self.relative_position_bias_table[
+            self.relative_position_index.view(-1)
+        ].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
+        )  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(
+            2, 0, 1
+        ).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwinTransformerBlock(nn.Module):
+    """Swin Transformer Block.
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=7,
+        shift_size=0,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim,
+            window_size=to_2tuple(self.window_size),
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop
+        )
+
+        self.H = None
+        self.W = None
+
+    def forward(self, x, mask_matrix):
+        """Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+            mask_matrix: Attention mask for cyclic shift.
+        """
+        B, L, C = x.shape
+        H, W = self.H, self.W
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # pad feature maps to multiples of window size
+        pad_l = pad_t = 0
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = mask_matrix
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(
+            shifted_x, self.window_size
+        )  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(
+            -1, self.window_size * self.window_size, C
+        )  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class PatchMerging(nn.Module):
+    """Patch Merging Layer
+    Args:
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x, H, W):
+        """Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+
+        # padding
+        pad_input = (H % 2 == 1) or (W % 2 == 1)
+        if pad_input:
+            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+
+class BasicLayer(nn.Module):
+    """A basic Swin Transformer layer for one stage.
+    Args:
+        dim (int): Number of feature channels
+        depth (int): Depths of this stage.
+        num_heads (int): Number of attention head.
+        window_size (int): Local window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(
+        self,
+        dim,
+        depth,
+        num_heads,
+        window_size=7,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        downsample=None,
+        use_checkpoint=False,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList(
+            [
+                SwinTransformerBlock(
+                    dim=dim,
+                    num_heads=num_heads,
+                    window_size=window_size,
+                    shift_size=0 if (i % 2 == 0) else window_size // 2,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, H, W):
+        """Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+
+        # calculate attention mask for SW-MSA
+        Hp = int(np.ceil(H / self.window_size)) * self.window_size
+        Wp = int(np.ceil(W / self.window_size)) * self.window_size
+        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1
+        h_slices = (
+            slice(0, -self.window_size),
+            slice(-self.window_size, -self.shift_size),
+            slice(-self.shift_size, None),
+        )
+        w_slices = (
+            slice(0, -self.window_size),
+            slice(-self.window_size, -self.shift_size),
+            slice(-self.shift_size, None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(
+            img_mask, self.window_size
+        )  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
+            attn_mask == 0, float(0.0)
+        )
+
+        for blk in self.blocks:
+            blk.H, blk.W = H, W
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, attn_mask)
+            else:
+                x = blk(x, attn_mask)
+        if self.downsample is not None:
+            x_down = self.downsample(x, H, W)
+            Wh, Ww = (H + 1) // 2, (W + 1) // 2
+            return x, H, W, x_down, Wh, Ww
+        else:
+            return x, H, W, x, H, W
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding
+    Args:
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        self.patch_size = patch_size
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        """Forward function."""
+        # padding
+        _, _, H, W = x.size()
+        if W % self.patch_size[1] != 0:
+            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+        if H % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+
+        x = self.proj(x)  # B C Wh Ww
+        if self.norm is not None:
+            Wh, Ww = x.size(2), x.size(3)
+            x = x.flatten(2).transpose(1, 2)
+            x = self.norm(x)
+            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+        return x
+
+
+class SwinTransformer(nn.Module):
+    """Swin Transformer backbone.
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+    Args:
+        pretrain_img_size (int): Input image size for training the pretrained model,
+            used in absolute postion embedding. Default 224.
+        patch_size (int | tuple(int)): Patch size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        depths (tuple[int]): Depths of each Swin Transformer stage.
+        num_heads (tuple[int]): Number of attention head of each stage.
+        window_size (int): Window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
+        drop_rate (float): Dropout rate.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
+        out_indices (Sequence[int]): Output from which stages.
+        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
+            -1 means not freezing any parameters.
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+        dilation (bool): if True, the output size if 16x downsample, ow 32x downsample.
+    """
+
+    def __init__(
+        self,
+        pretrain_img_size=224,
+        patch_size=4,
+        in_chans=3,
+        embed_dim=96,
+        depths=[2, 2, 6, 2],
+        num_heads=[3, 6, 12, 24],
+        window_size=7,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.2,
+        norm_layer=nn.LayerNorm,
+        ape=False,
+        patch_norm=True,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=-1,
+        dilation=False,
+        use_checkpoint=False,
+    ):
+        super().__init__()
+
+        self.pretrain_img_size = pretrain_img_size
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+        self.dilation = dilation
+
+        # if use_checkpoint:
+        #     print("use_checkpoint!!!!!!!!!!!!!!!!!!!!!!!!")
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None,
+        )
+
+        # absolute position embedding
+        if self.ape:
+            pretrain_img_size = to_2tuple(pretrain_img_size)
+            patch_size = to_2tuple(patch_size)
+            patches_resolution = [
+                pretrain_img_size[0] // patch_size[0],
+                pretrain_img_size[1] // patch_size[1],
+            ]
+
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])
+            )
+            trunc_normal_(self.absolute_pos_embed, std=0.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+        ]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        # prepare downsample list
+        downsamplelist = [PatchMerging for i in range(self.num_layers)]
+        downsamplelist[-1] = None
+        num_features = [int(embed_dim * 2**i) for i in range(self.num_layers)]
+        if self.dilation:
+            downsamplelist[-2] = None
+            num_features[-1] = int(embed_dim * 2 ** (self.num_layers - 1)) // 2
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(
+                # dim=int(embed_dim * 2 ** i_layer),
+                dim=num_features[i_layer],
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
+                norm_layer=norm_layer,
+                # downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                downsample=downsamplelist[i_layer],
+                use_checkpoint=use_checkpoint,
+            )
+            self.layers.append(layer)
+
+        # num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+        self.num_features = num_features
+
+        # add a norm layer for each output
+        for i_layer in out_indices:
+            layer = norm_layer(num_features[i_layer])
+            layer_name = f"norm{i_layer}"
+            self.add_module(layer_name, layer)
+
+        self._freeze_stages()
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            self.patch_embed.eval()
+            for param in self.patch_embed.parameters():
+                param.requires_grad = False
+
+        if self.frozen_stages >= 1 and self.ape:
+            self.absolute_pos_embed.requires_grad = False
+
+        if self.frozen_stages >= 2:
+            self.pos_drop.eval()
+            for i in range(0, self.frozen_stages - 1):
+                m = self.layers[i]
+                m.eval()
+                for param in m.parameters():
+                    param.requires_grad = False
+
+    # def init_weights(self, pretrained=None):
+    #     """Initialize the weights in backbone.
+    #     Args:
+    #         pretrained (str, optional): Path to pre-trained weights.
+    #             Defaults to None.
+    #     """
+
+    #     def _init_weights(m):
+    #         if isinstance(m, nn.Linear):
+    #             trunc_normal_(m.weight, std=.02)
+    #             if isinstance(m, nn.Linear) and m.bias is not None:
+    #                 nn.init.constant_(m.bias, 0)
+    #         elif isinstance(m, nn.LayerNorm):
+    #             nn.init.constant_(m.bias, 0)
+    #             nn.init.constant_(m.weight, 1.0)
+
+    #     if isinstance(pretrained, str):
+    #         self.apply(_init_weights)
+    #         logger = get_root_logger()
+    #         load_checkpoint(self, pretrained, strict=False, logger=logger)
+    #     elif pretrained is None:
+    #         self.apply(_init_weights)
+    #     else:
+    #         raise TypeError('pretrained must be a str or None')
+
+    def forward_raw(self, x):
+        """Forward function."""
+        x = self.patch_embed(x)
+
+        Wh, Ww = x.size(2), x.size(3)
+        if self.ape:
+            # interpolate the position embedding to the corresponding size
+            absolute_pos_embed = F.interpolate(
+                self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
+            )
+            x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B Wh*Ww C
+        else:
+            x = x.flatten(2).transpose(1, 2)
+        x = self.pos_drop(x)
+
+        outs = []
+        for i in range(self.num_layers):
+            layer = self.layers[i]
+            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
+            # import ipdb; ipdb.set_trace()
+
+            if i in self.out_indices:
+                norm_layer = getattr(self, f"norm{i}")
+                x_out = norm_layer(x_out)
+
+                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+                outs.append(out)
+        # in:
+        #   torch.Size([2, 3, 1024, 1024])
+        # outs:
+        #   [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \
+        #       torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])]
+        return tuple(outs)
+
+    def forward(self, tensor_list: NestedTensor):
+
+        x = tensor_list.tensors
+
+        """Forward function."""
+        x = self.patch_embed(x)
+
+
+        Wh, Ww = x.size(2), x.size(3)
+        if self.ape:
+            # interpolate the position embedding to the corresponding size
+            absolute_pos_embed = F.interpolate(
+                self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
+            )
+            x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B Wh*Ww C
+        else:
+            x = x.flatten(2).transpose(1, 2)
+        x = self.pos_drop(x)
+
+        outs = []
+        for i in range(self.num_layers):
+            layer = self.layers[i]
+            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
+
+            if i in self.out_indices:
+                norm_layer = getattr(self, f"norm{i}")
+                x_out = norm_layer(x_out)
+
+                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+                outs.append(out)
+        # in:
+        #   torch.Size([2, 3, 1024, 1024])
+        # out:
+        #   [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \
+        #       torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])]
+
+        # collect for nesttensors
+        outs_dict = {}
+        for idx, out_i in enumerate(outs):
+            m = tensor_list.mask
+            assert m is not None
+            mask = F.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0]
+            outs_dict[idx] = NestedTensor(out_i, mask)
+
+        return outs_dict
+
+    def train(self, mode=True):
+        """Convert the model into training mode while keep layers freezed."""
+        super(SwinTransformer, self).train(mode)
+        self._freeze_stages()
+
+
+def build_swin_transformer(modelname, pretrain_img_size, **kw):
+    assert modelname in [
+        "swin_T_224_1k",
+        "swin_B_224_22k",
+        "swin_B_384_22k",
+        "swin_L_224_22k",
+        "swin_L_384_22k",
+    ]
+
+    model_para_dict = {
+        "swin_T_224_1k": dict(
+            embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7
+        ),
+        "swin_B_224_22k": dict(
+            embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=7
+        ),
+        "swin_B_384_22k": dict(
+            embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12
+        ),
+        "swin_L_224_22k": dict(
+            embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=7
+        ),
+        "swin_L_384_22k": dict(
+            embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12
+        ),
+    }
+    kw_cgf = model_para_dict[modelname]
+    kw_cgf.update(kw)
+    model = SwinTransformer(pretrain_img_size=pretrain_img_size, **kw_cgf)
+    return model
+
+
+if __name__ == "__main__":
+    model = build_swin_transformer("swin_L_384_22k", 384, dilation=True)
+    x = torch.rand(2, 3, 1024, 1024)
+    y = model.forward_raw(x)
+    import ipdb
+
+    ipdb.set_trace()
+    x = torch.rand(2, 3, 384, 384)
+    y = model.forward_raw(x)

models/GroundingDINO/bertwarper.py

@@ -0,0 +1,273 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from torch import Tensor, nn
+from torchvision.ops.boxes import nms
+from transformers import BertConfig, BertModel, BertPreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
+
+
+class BertModelWarper(nn.Module):
+    def __init__(self, bert_model):
+        super().__init__()
+        # self.bert = bert_modelc
+
+        self.config = bert_model.config
+        self.embeddings = bert_model.embeddings
+        self.encoder = bert_model.encoder
+        self.pooler = bert_model.pooler
+
+        self.get_extended_attention_mask = bert_model.get_extended_attention_mask
+        self.invert_attention_mask = bert_model.invert_attention_mask
+        self.get_head_mask = bert_model.get_head_mask
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        """
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.config.is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # past_key_values_length
+        past_key_values_length = (
+            past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        )
+
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                ((batch_size, seq_length + past_key_values_length)), device=device
+            )
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
+            attention_mask, input_shape, device
+        )
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+        # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO':
+        #     import ipdb; ipdb.set_trace()
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values_length,
+        )
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+class TextEncoderShell(nn.Module):
+    def __init__(self, text_encoder):
+        super().__init__()
+        self.text_encoder = text_encoder
+        self.config = self.text_encoder.config
+
+    def forward(self, **kw):
+        # feed into text encoder
+        return self.text_encoder(**kw)
+
+
+def generate_masks_with_special_tokens(tokenized, special_tokens_list, tokenizer):
+    """Generate attention mask between each pair of special tokens
+    Args:
+        input_ids (torch.Tensor): input ids. Shape: [bs, num_token]
+        special_tokens_mask (list): special tokens mask.
+    Returns:
+        torch.Tensor: attention mask between each special tokens.
+    """
+    input_ids = tokenized["input_ids"]
+    bs, num_token = input_ids.shape
+    # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens
+    special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool()
+    for special_token in special_tokens_list:
+        special_tokens_mask |= input_ids == special_token
+
+    # idxs: each row is a list of indices of special tokens
+    idxs = torch.nonzero(special_tokens_mask)
+
+    # generate attention mask and positional ids
+    attention_mask = (
+        torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1)
+    )
+    position_ids = torch.zeros((bs, num_token), device=input_ids.device)
+    previous_col = 0
+    for i in range(idxs.shape[0]):
+        row, col = idxs[i]
+        if (col == 0) or (col == num_token - 1):
+            attention_mask[row, col, col] = True
+            position_ids[row, col] = 0
+        else:
+            attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True
+            position_ids[row, previous_col + 1 : col + 1] = torch.arange(
+                0, col - previous_col, device=input_ids.device
+            )
+
+        previous_col = col
+
+    # # padding mask
+    # padding_mask = tokenized['attention_mask']
+    # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool()
+
+    return attention_mask, position_ids.to(torch.long)
+
+
+def generate_masks_with_special_tokens_and_transfer_map(tokenized, special_tokens_list, tokenizer):
+    """Generate attention mask between each pair of special tokens
+    Args:
+        input_ids (torch.Tensor): input ids. Shape: [bs, num_token]
+        special_tokens_mask (list): special tokens mask.
+    Returns:
+        torch.Tensor: attention mask between each special tokens.
+    """
+    input_ids = tokenized["input_ids"]
+    bs, num_token = input_ids.shape
+    # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens
+    special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool()
+    for special_token in special_tokens_list:
+        special_tokens_mask |= input_ids == special_token
+
+    # idxs: each row is a list of indices of special tokens
+    idxs = torch.nonzero(special_tokens_mask)
+
+    # generate attention mask and positional ids
+    attention_mask = (
+        torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1)
+    )
+    position_ids = torch.zeros((bs, num_token), device=input_ids.device)
+    cate_to_token_mask_list = [[] for _ in range(bs)]
+    previous_col = 0
+    for i in range(idxs.shape[0]):
+        row, col = idxs[i]
+        if (col == 0) or (col == num_token - 1):
+            attention_mask[row, col, col] = True
+            position_ids[row, col] = 0
+        else:
+            attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True
+            position_ids[row, previous_col + 1 : col + 1] = torch.arange(
+                0, col - previous_col, device=input_ids.device
+            )
+            c2t_maski = torch.zeros((num_token), device=input_ids.device).bool()
+            c2t_maski[previous_col + 1 : col] = True
+            cate_to_token_mask_list[row].append(c2t_maski)
+        previous_col = col
+
+    cate_to_token_mask_list = [
+        torch.stack(cate_to_token_mask_listi, dim=0)
+        for cate_to_token_mask_listi in cate_to_token_mask_list
+    ]
+
+    # # padding mask
+    # padding_mask = tokenized['attention_mask']
+    # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool()
+
+    return attention_mask, position_ids.to(torch.long), cate_to_token_mask_list

models/GroundingDINO/fuse_modules.py

@@ -0,0 +1,297 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import DropPath
+
+
+class FeatureResizer(nn.Module):
+    """
+    This class takes as input a set of embeddings of dimension C1 and outputs a set of
+    embedding of dimension C2, after a linear transformation, dropout and normalization (LN).
+    """
+
+    def __init__(self, input_feat_size, output_feat_size, dropout, do_ln=True):
+        super().__init__()
+        self.do_ln = do_ln
+        # Object feature encoding
+        self.fc = nn.Linear(input_feat_size, output_feat_size, bias=True)
+        self.layer_norm = nn.LayerNorm(output_feat_size, eps=1e-12)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, encoder_features):
+        x = self.fc(encoder_features)
+        if self.do_ln:
+            x = self.layer_norm(x)
+        output = self.dropout(x)
+        return output
+
+
+def l1norm(X, dim, eps=1e-8):
+    """L1-normalize columns of X"""
+    norm = torch.abs(X).sum(dim=dim, keepdim=True) + eps
+    X = torch.div(X, norm)
+    return X
+
+
+def l2norm(X, dim, eps=1e-8):
+    """L2-normalize columns of X"""
+    norm = torch.pow(X, 2).sum(dim=dim, keepdim=True).sqrt() + eps
+    X = torch.div(X, norm)
+    return X
+
+
+def func_attention(query, context, smooth=1, raw_feature_norm="softmax", eps=1e-8):
+    """
+    query: (n_context, queryL, d)
+    context: (n_context, sourceL, d)
+    """
+    batch_size_q, queryL = query.size(0), query.size(1)
+    batch_size, sourceL = context.size(0), context.size(1)
+
+    # Get attention
+    # --> (batch, d, queryL)
+    queryT = torch.transpose(query, 1, 2)
+
+    # (batch, sourceL, d)(batch, d, queryL)
+    # --> (batch, sourceL, queryL)
+    attn = torch.bmm(context, queryT)
+    if raw_feature_norm == "softmax":
+        # --> (batch*sourceL, queryL)
+        attn = attn.view(batch_size * sourceL, queryL)
+        attn = nn.Softmax()(attn)
+        # --> (batch, sourceL, queryL)
+        attn = attn.view(batch_size, sourceL, queryL)
+    elif raw_feature_norm == "l2norm":
+        attn = l2norm(attn, 2)
+    elif raw_feature_norm == "clipped_l2norm":
+        attn = nn.LeakyReLU(0.1)(attn)
+        attn = l2norm(attn, 2)
+    else:
+        raise ValueError("unknown first norm type:", raw_feature_norm)
+    # --> (batch, queryL, sourceL)
+    attn = torch.transpose(attn, 1, 2).contiguous()
+    # --> (batch*queryL, sourceL)
+    attn = attn.view(batch_size * queryL, sourceL)
+    attn = nn.Softmax()(attn * smooth)
+    # --> (batch, queryL, sourceL)
+    attn = attn.view(batch_size, queryL, sourceL)
+    # --> (batch, sourceL, queryL)
+    attnT = torch.transpose(attn, 1, 2).contiguous()
+
+    # --> (batch, d, sourceL)
+    contextT = torch.transpose(context, 1, 2)
+    # (batch x d x sourceL)(batch x sourceL x queryL)
+    # --> (batch, d, queryL)
+    weightedContext = torch.bmm(contextT, attnT)
+    # --> (batch, queryL, d)
+    weightedContext = torch.transpose(weightedContext, 1, 2)
+
+    return weightedContext, attnT
+
+
+class BiMultiHeadAttention(nn.Module):
+    def __init__(self, v_dim, l_dim, embed_dim, num_heads, dropout=0.1, cfg=None):
+        super(BiMultiHeadAttention, self).__init__()
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.v_dim = v_dim
+        self.l_dim = l_dim
+
+        assert (
+            self.head_dim * self.num_heads == self.embed_dim
+        ), f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
+        self.scale = self.head_dim ** (-0.5)
+        self.dropout = dropout
+
+        self.v_proj = nn.Linear(self.v_dim, self.embed_dim)
+        self.l_proj = nn.Linear(self.l_dim, self.embed_dim)
+        self.values_v_proj = nn.Linear(self.v_dim, self.embed_dim)
+        self.values_l_proj = nn.Linear(self.l_dim, self.embed_dim)
+
+        self.out_v_proj = nn.Linear(self.embed_dim, self.v_dim)
+        self.out_l_proj = nn.Linear(self.embed_dim, self.l_dim)
+
+        self.stable_softmax_2d = True
+        self.clamp_min_for_underflow = True
+        self.clamp_max_for_overflow = True
+
+        self._reset_parameters()
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.v_proj.weight)
+        self.v_proj.bias.data.fill_(0)
+        nn.init.xavier_uniform_(self.l_proj.weight)
+        self.l_proj.bias.data.fill_(0)
+        nn.init.xavier_uniform_(self.values_v_proj.weight)
+        self.values_v_proj.bias.data.fill_(0)
+        nn.init.xavier_uniform_(self.values_l_proj.weight)
+        self.values_l_proj.bias.data.fill_(0)
+        nn.init.xavier_uniform_(self.out_v_proj.weight)
+        self.out_v_proj.bias.data.fill_(0)
+        nn.init.xavier_uniform_(self.out_l_proj.weight)
+        self.out_l_proj.bias.data.fill_(0)
+
+    def forward(self, v, l, attention_mask_v=None, attention_mask_l=None):
+        """_summary_
+
+        Args:
+            v (_type_): bs, n_img, dim
+            l (_type_): bs, n_text, dim
+            attention_mask_v (_type_, optional): _description_. bs, n_img
+            attention_mask_l (_type_, optional): _description_. bs, n_text
+
+        Returns:
+            _type_: _description_
+        """
+        # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO':
+        #     import ipdb; ipdb.set_trace()
+        bsz, tgt_len, _ = v.size()
+
+        query_states = self.v_proj(v) * self.scale
+        key_states = self._shape(self.l_proj(l), -1, bsz)
+        value_v_states = self._shape(self.values_v_proj(v), -1, bsz)
+        value_l_states = self._shape(self.values_l_proj(l), -1, bsz)
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
+        key_states = key_states.view(*proj_shape)
+        value_v_states = value_v_states.view(*proj_shape)
+        value_l_states = value_l_states.view(*proj_shape)
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))  # bs*nhead, nimg, ntxt
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is {attn_weights.size()}"
+            )
+
+        if self.stable_softmax_2d:
+            attn_weights = attn_weights - attn_weights.max()
+
+        if self.clamp_min_for_underflow:
+            attn_weights = torch.clamp(
+                attn_weights, min=-50000
+            )  # Do not increase -50000, data type half has quite limited range
+        if self.clamp_max_for_overflow:
+            attn_weights = torch.clamp(
+                attn_weights, max=50000
+            )  # Do not increase 50000, data type half has quite limited range
+
+        attn_weights_T = attn_weights.transpose(1, 2)
+        attn_weights_l = attn_weights_T - torch.max(attn_weights_T, dim=-1, keepdim=True)[0]
+        if self.clamp_min_for_underflow:
+            attn_weights_l = torch.clamp(
+                attn_weights_l, min=-50000
+            )  # Do not increase -50000, data type half has quite limited range
+        if self.clamp_max_for_overflow:
+            attn_weights_l = torch.clamp(
+                attn_weights_l, max=50000
+            )  # Do not increase 50000, data type half has quite limited range
+
+        # mask vison for language
+        if attention_mask_v is not None:
+            attention_mask_v = (
+                attention_mask_v[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1)
+            )
+            attn_weights_l.masked_fill_(attention_mask_v, float("-inf"))
+
+        attn_weights_l = attn_weights_l.softmax(dim=-1)
+
+        # mask language for vision
+        if attention_mask_l is not None:
+            attention_mask_l = (
+                attention_mask_l[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1)
+            )
+            attn_weights.masked_fill_(attention_mask_l, float("-inf"))
+        attn_weights_v = attn_weights.softmax(dim=-1)
+
+        attn_probs_v = F.dropout(attn_weights_v, p=self.dropout, training=self.training)
+        attn_probs_l = F.dropout(attn_weights_l, p=self.dropout, training=self.training)
+
+        attn_output_v = torch.bmm(attn_probs_v, value_l_states)
+        attn_output_l = torch.bmm(attn_probs_l, value_v_states)
+
+        if attn_output_v.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output_v` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is {attn_output_v.size()}"
+            )
+
+        if attn_output_l.size() != (bsz * self.num_heads, src_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output_l` should be of size {(bsz, self.num_heads, src_len, self.head_dim)}, but is {attn_output_l.size()}"
+            )
+
+        attn_output_v = attn_output_v.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output_v = attn_output_v.transpose(1, 2)
+        attn_output_v = attn_output_v.reshape(bsz, tgt_len, self.embed_dim)
+
+        attn_output_l = attn_output_l.view(bsz, self.num_heads, src_len, self.head_dim)
+        attn_output_l = attn_output_l.transpose(1, 2)
+        attn_output_l = attn_output_l.reshape(bsz, src_len, self.embed_dim)
+
+        attn_output_v = self.out_v_proj(attn_output_v)
+        attn_output_l = self.out_l_proj(attn_output_l)
+
+        return attn_output_v, attn_output_l
+
+
+# Bi-Direction MHA (text->image, image->text)
+class BiAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        v_dim,
+        l_dim,
+        embed_dim,
+        num_heads,
+        dropout=0.1,
+        drop_path=0.0,
+        init_values=1e-4,
+        cfg=None,
+    ):
+        """
+        Inputs:
+            embed_dim - Dimensionality of input and attention feature vectors
+            hidden_dim - Dimensionality of hidden layer in feed-forward network
+                         (usually 2-4x larger than embed_dim)
+            num_heads - Number of heads to use in the Multi-Head Attention block
+            dropout - Amount of dropout to apply in the feed-forward network
+        """
+        super(BiAttentionBlock, self).__init__()
+
+        # pre layer norm
+        self.layer_norm_v = nn.LayerNorm(v_dim)
+        self.layer_norm_l = nn.LayerNorm(l_dim)
+        self.attn = BiMultiHeadAttention(
+            v_dim=v_dim, l_dim=l_dim, embed_dim=embed_dim, num_heads=num_heads, dropout=dropout
+        )
+
+        # add layer scale for training stability
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.gamma_v = nn.Parameter(init_values * torch.ones((v_dim)), requires_grad=True)
+        self.gamma_l = nn.Parameter(init_values * torch.ones((l_dim)), requires_grad=True)
+
+    def forward(self, v, l, attention_mask_v=None, attention_mask_l=None):
+        v = self.layer_norm_v(v)
+        l = self.layer_norm_l(l)
+        delta_v, delta_l = self.attn(
+            v, l, attention_mask_v=attention_mask_v, attention_mask_l=attention_mask_l
+        )
+        # v, l = v + delta_v, l + delta_l
+        v = v + self.drop_path(self.gamma_v * delta_v)
+        l = l + self.drop_path(self.gamma_l * delta_l)
+        return v, l
+
+    # def forward(self, v:List[torch.Tensor], l, attention_mask_v=None, attention_mask_l=None)

models/GroundingDINO/groundingdino.py

@@ -0,0 +1,1041 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Conditional DETR model and criterion classes.
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# ------------------------------------------------------------------------
+import copy
+from typing import List
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision.ops.boxes import nms
+from torchvision.ops import roi_align
+from transformers import AutoTokenizer, BertModel, BertTokenizer, RobertaModel, RobertaTokenizerFast
+
+from groundingdino.util import box_ops, get_tokenlizer
+from groundingdino.util.misc import (
+    NestedTensor,
+    accuracy,
+    get_world_size,
+    interpolate,
+    inverse_sigmoid,
+    is_dist_avail_and_initialized,
+    nested_tensor_from_tensor_list,
+)
+from groundingdino.util.utils import get_phrases_from_posmap
+from groundingdino.util.visualizer import COCOVisualizer
+from groundingdino.util.vl_utils import create_positive_map_from_span
+
+from ..registry import MODULE_BUILD_FUNCS
+from .backbone import build_backbone
+from .bertwarper import (
+    BertModelWarper,
+    generate_masks_with_special_tokens,
+    generate_masks_with_special_tokens_and_transfer_map,
+)
+from .transformer import build_transformer
+from .transformer_loca import TransformerEncoder
+from .positional_encoding_loca import PositionalEncodingsFixed
+from .utils import MLP, ContrastiveEmbed, sigmoid_focal_loss
+
+from .matcher import build_matcher
+import numpy as np
+
+
+
+class GroundingDINO(nn.Module):
+    """This is the Cross-Attention Detector module that performs object detection"""
+
+    def __init__(
+        self,
+        backbone,
+        transformer,
+        num_queries,
+        aux_loss=False,
+        iter_update=False,
+        query_dim=2,
+        num_feature_levels=1,
+        nheads=8,
+        # two stage
+        two_stage_type="no",  # ['no', 'standard']
+        dec_pred_bbox_embed_share=True,
+        two_stage_class_embed_share=True,
+        two_stage_bbox_embed_share=True,
+        num_patterns=0,
+        dn_number=100,
+        dn_box_noise_scale=0.4,
+        dn_label_noise_ratio=0.5,
+        dn_labelbook_size=100,
+        text_encoder_type="bert-base-uncased",
+        sub_sentence_present=True,
+        max_text_len=256,
+    ):
+        """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_queries: number of object queries, ie detection slot. This is the maximal number of objects
+                         Conditional 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
+        self.hidden_dim = hidden_dim = transformer.d_model
+        self.num_feature_levels = num_feature_levels
+        self.nheads = nheads
+        self.max_text_len = 256
+        self.sub_sentence_present = sub_sentence_present
+
+        # setting query dim
+        self.query_dim = query_dim
+        assert query_dim == 4
+
+        # visual exemplar cropping
+        self.feature_map_proj = nn.Conv2d(
+            (256 + 512 + 1024), hidden_dim, kernel_size=1
+        )
+        self.feature_map_encoder = TransformerEncoder(
+            3, hidden_dim, 8, 0.1, 1e-5,
+            8, True, nn.GELU, True
+        )
+        self.feature_map_pos_embed = PositionalEncodingsFixed(hidden_dim)
+
+        # for dn training
+        self.num_patterns = num_patterns
+        self.dn_number = dn_number
+        self.dn_box_noise_scale = dn_box_noise_scale
+        self.dn_label_noise_ratio = dn_label_noise_ratio
+        self.dn_labelbook_size = dn_labelbook_size
+
+        # bert
+        self.tokenizer = get_tokenlizer.get_tokenlizer(text_encoder_type)
+        self.bert = get_tokenlizer.get_pretrained_language_model(text_encoder_type)
+        self.bert.pooler.dense.weight.requires_grad_(False)
+        self.bert.pooler.dense.bias.requires_grad_(False)
+        self.bert = BertModelWarper(bert_model=self.bert)
+
+        self.feat_map = nn.Linear(self.bert.config.hidden_size, self.hidden_dim, bias=True)
+        nn.init.constant_(self.feat_map.bias.data, 0)
+        nn.init.xavier_uniform_(self.feat_map.weight.data)
+        # freeze
+
+        # special tokens
+        self.specical_tokens = self.tokenizer.convert_tokens_to_ids(["[CLS]", "[SEP]", ".", "?"])
+        self.dot_id = self.tokenizer.convert_tokens_to_ids(".")
+
+        # prepare input projection layers
+        if num_feature_levels > 1:
+            num_backbone_outs = len(backbone.num_channels)
+            input_proj_list = []
+            for _ in range(num_backbone_outs):
+                in_channels = backbone.num_channels[_]
+                input_proj_list.append(
+                    nn.Sequential(
+                        nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
+                        nn.GroupNorm(32, hidden_dim),
+                    )
+                )
+            for _ in range(num_feature_levels - num_backbone_outs):
+                input_proj_list.append(
+                    nn.Sequential(
+                        nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1),
+                        nn.GroupNorm(32, hidden_dim),
+                    )
+                )
+                in_channels = hidden_dim
+            self.input_proj = nn.ModuleList(input_proj_list)
+        else:
+            assert two_stage_type == "no", "two_stage_type should be no if num_feature_levels=1 !!!"
+            self.input_proj = nn.ModuleList(
+                [
+                    nn.Sequential(
+                        nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1),
+                        nn.GroupNorm(32, hidden_dim),
+                    )
+                ]
+            )
+
+        self.backbone = backbone
+        self.aux_loss = aux_loss
+        self.box_pred_damping = box_pred_damping = None
+
+        self.iter_update = iter_update
+        assert iter_update, "Why not iter_update?"
+
+        # prepare pred layers
+        self.dec_pred_bbox_embed_share = dec_pred_bbox_embed_share
+        # prepare class & box embed
+        _class_embed = ContrastiveEmbed()
+
+        _bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
+        nn.init.constant_(_bbox_embed.layers[-1].weight.data, 0)
+        nn.init.constant_(_bbox_embed.layers[-1].bias.data, 0)
+
+        if dec_pred_bbox_embed_share:
+            box_embed_layerlist = [_bbox_embed for i in range(transformer.num_decoder_layers)]
+        else:
+            box_embed_layerlist = [
+                copy.deepcopy(_bbox_embed) for i in range(transformer.num_decoder_layers)
+            ]
+        class_embed_layerlist = [_class_embed for i in range(transformer.num_decoder_layers)]
+        self.bbox_embed = nn.ModuleList(box_embed_layerlist)
+        self.class_embed = nn.ModuleList(class_embed_layerlist)
+        self.transformer.decoder.bbox_embed = self.bbox_embed
+        self.transformer.decoder.class_embed = self.class_embed
+
+        # two stage
+        self.two_stage_type = two_stage_type
+        assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format(
+            two_stage_type
+        )
+        if two_stage_type != "no":
+            if two_stage_bbox_embed_share:
+                assert dec_pred_bbox_embed_share
+                self.transformer.enc_out_bbox_embed = _bbox_embed
+            else:
+                self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed)
+
+            if two_stage_class_embed_share:
+                assert dec_pred_bbox_embed_share
+                self.transformer.enc_out_class_embed = _class_embed
+            else:
+                self.transformer.enc_out_class_embed = copy.deepcopy(_class_embed)
+
+            self.refpoint_embed = None
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        # init input_proj
+        for proj in self.input_proj:
+            nn.init.xavier_uniform_(proj[0].weight, gain=1)
+            nn.init.constant_(proj[0].bias, 0)
+
+    def init_ref_points(self, use_num_queries):
+        self.refpoint_embed = nn.Embedding(use_num_queries, self.query_dim)
+
+    def add_exemplar_tokens(self, tokenized, text_dict, exemplar_tokens, exemplar_labels):
+        input_ids = tokenized["input_ids"]
+        
+        device = input_ids.device
+        new_input_ids = []
+        encoded_text = text_dict["encoded_text"]
+        new_encoded_text = []
+        text_token_mask = text_dict["text_token_mask"]
+        new_text_token_mask = []
+        position_ids = text_dict["position_ids"]
+        text_self_attention_masks = text_dict["text_self_attention_masks"]
+
+        for sample_ind in range(len(exemplar_tokens)):
+            new_input_ids_sample = input_ids[sample_ind]
+            new_encoded_text_sample = encoded_text[sample_ind]
+            for exemplar_ind in range(len(exemplar_tokens[sample_ind])):
+                label = exemplar_labels[sample_ind][exemplar_ind]
+                exemplar = exemplar_tokens[sample_ind][exemplar_ind].unsqueeze(0) # add an L (# of boxes) dimension
+                label_count = -1
+                for token_ind in range(len(new_input_ids_sample)):
+                    input_id = new_input_ids_sample[token_ind]
+                    if (input_id not in self.specical_tokens) and (token_ind == 0 or (new_input_ids_sample[token_ind - 1] in self.specical_tokens)):
+                        label_count += 1
+                    elif token_ind > 0 \
+                         and new_input_ids_sample[token_ind] == self.dot_id \
+                         and new_input_ids_sample[token_ind - 1] in self.specical_tokens:
+                        label_count += 1
+                    if label_count == label:
+                        # Get the index where to insert the exemplar tokens.
+                        ind_to_insert_exemplar = token_ind
+                        while new_input_ids_sample[ind_to_insert_exemplar] not in self.specical_tokens:
+                            ind_to_insert_exemplar += 1
+                        break
+            
+                # * token indicates exemplar
+                new_input_ids_sample = torch.cat([new_input_ids_sample[:ind_to_insert_exemplar], torch.tensor([1008] * exemplar.shape[0]).to(device), new_input_ids_sample[ind_to_insert_exemplar:]])
+                new_encoded_text_sample = torch.cat([new_encoded_text_sample[:ind_to_insert_exemplar, :], exemplar, new_encoded_text_sample[ind_to_insert_exemplar:, :]])
+
+            new_text_token_mask_sample = torch.full((len(new_input_ids_sample),), True).to(device)
+            new_input_ids.append(new_input_ids_sample)
+            new_encoded_text.append(new_encoded_text_sample)
+            new_text_token_mask.append(new_text_token_mask_sample)
+
+        # Add padding for variable numbers of exemplars for each image.
+        pad_token_id = self.tokenizer.pad_token_id
+        bs = len(new_input_ids)
+        d_model = new_encoded_text[0].shape[-1]
+        device = input_ids.device
+        
+        lens = [ids.numel() for ids in new_input_ids]
+        Lmax = max(lens)
+        
+        padded_input_ids = torch.full((bs, Lmax), pad_token_id, dtype=torch.long, device=device)
+        padded_encoded_text = torch.zeros(bs, Lmax, d_model, dtype=new_encoded_text[0].dtype, device=device)
+        padded_text_token_mask = torch.zeros(bs, Lmax, dtype=torch.bool, device=device)  # False for PAD
+        
+        for i in range(bs):
+            L = min(lens[i], Lmax)
+            padded_input_ids[i, :L] = new_input_ids[i][:L]
+            padded_encoded_text[i, :L] = new_encoded_text[i][:L]
+            padded_text_token_mask[i, :L] = new_text_token_mask[i][:L]  # True for real tokens
+        
+        tokenized["input_ids"] = padded_input_ids
+        tokenized["attention_mask"] = padded_text_token_mask.long()  # HF expects 0/1
+        
+        text_self_attention_masks, position_ids, _ = generate_masks_with_special_tokens_and_transfer_map(
+            tokenized, self.specical_tokens, None
+        )
+
+        
+        text_dict = {
+            "encoded_text": padded_encoded_text,
+            "text_token_mask": padded_text_token_mask,   # True=valid, False=pad
+            "position_ids": position_ids,
+            "text_self_attention_masks": text_self_attention_masks,
+        }
+        return text_dict
+
+    def combine_features(self, features):
+        
+        (bs, c, h, w) = (features[0].decompose()[0].shape[-4], features[0].decompose()[0].shape[-3], features[0].decompose()[0].shape[-2], features[0].decompose()[0].shape[-1])
+        
+        x = torch.cat([
+            F.interpolate(feat.decompose()[0], size=(h, w), mode='bilinear', align_corners=True)
+            for feat in features
+        ], dim=1)
+        
+        x = self.feature_map_proj(x)
+        
+        #pos_emb = self.feature_map_pos_embed(bs, h, w, x.device)
+        
+        #pos_emb = pos_emb.flatten(2).permute(2, 0, 1)
+        
+        #x = x.flatten(2).permute(2, 0, 1)
+        
+        #x = self.feature_map_encoder(x, pos_emb, src_key_padding_mask=None, src_mask=None)
+        
+        #x = x.permute(1, 2, 0).reshape(-1, self.hidden_dim, h, w)
+        
+        return x
+
+
+    def forward(self, samples: NestedTensor, exemplar_boxes, exemplar_labels, targets: List = None, **kw):
+        """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]
+           - "pred_boxes": The normalized boxes coordinates for all queries, represented as
+                           (center_x, center_y, width, height). 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 targets is None:
+            captions = kw["captions"]
+        else:
+            captions = [t["caption"] for t in targets]
+        
+        # encoder texts
+
+        tokenized = self.tokenizer(captions, padding="longest", return_tensors="pt").to(
+            samples.device
+        )
+
+        one_hot_token = tokenized
+
+        (
+            text_self_attention_masks,
+            position_ids,
+            cate_to_token_mask_list,
+        ) = generate_masks_with_special_tokens_and_transfer_map(
+            tokenized, self.specical_tokens, self.tokenizer
+        )
+
+        if text_self_attention_masks.shape[1] > self.max_text_len:
+            text_self_attention_masks = text_self_attention_masks[
+                :, : self.max_text_len, : self.max_text_len
+            ]
+            position_ids = position_ids[:, : self.max_text_len]
+            tokenized["input_ids"] = tokenized["input_ids"][:, : self.max_text_len]
+            tokenized["attention_mask"] = tokenized["attention_mask"][:, : self.max_text_len]
+            tokenized["token_type_ids"] = tokenized["token_type_ids"][:, : self.max_text_len]
+
+        # extract text embeddings
+        if self.sub_sentence_present:
+            tokenized_for_encoder = {k: v for k, v in tokenized.items() if k != "attention_mask"}
+            tokenized_for_encoder["attention_mask"] = text_self_attention_masks
+            tokenized_for_encoder["position_ids"] = position_ids
+        else:
+            tokenized_for_encoder = tokenized
+
+        bert_output = self.bert(**tokenized_for_encoder)  # bs, 195, 768
+
+        encoded_text = self.feat_map(bert_output["last_hidden_state"])  # bs, 195, d_model
+        text_token_mask = tokenized.attention_mask.bool()  # bs, 195
+        # text_token_mask: True for nomask, False for mask
+        # text_self_attention_masks: True for nomask, False for mask
+
+        if encoded_text.shape[1] > self.max_text_len:
+            encoded_text = encoded_text[:, : self.max_text_len, :]
+            text_token_mask = text_token_mask[:, : self.max_text_len]
+            position_ids = position_ids[:, : self.max_text_len]
+            text_self_attention_masks = text_self_attention_masks[
+                :, : self.max_text_len, : self.max_text_len
+            ]
+        
+
+        text_dict = {
+            "encoded_text": encoded_text,  # bs, 195, d_model
+            "text_token_mask": text_token_mask,  # bs, 195
+            "position_ids": position_ids,  # bs, 195
+            "text_self_attention_masks": text_self_attention_masks,  # bs, 195,195
+        }
+
+        if isinstance(samples, (list, torch.Tensor)):
+            samples = nested_tensor_from_tensor_list(samples)
+        
+        features, poss = self.backbone(samples)
+        combined_features = self.combine_features(features)
+        
+        # Get visual exemplar tokens.
+        bs = len(exemplar_boxes)
+        exemplar_tokens = roi_align(combined_features, boxes=exemplar_boxes, output_size=(1, 1), spatial_scale=(1 / 8), aligned=True)
+        if exemplar_tokens.shape[0] > 0:
+            # Assemble boxes back into [bs, num_exemplars, 4] list 
+            exemplar_counts = [b.size(0) for b in exemplar_boxes]
+            exemplar_tokens = list(torch.split(exemplar_tokens, exemplar_counts, dim=0))
+            exemplar_tokens = [exemp.squeeze(-1).squeeze(-1) for exemp in exemplar_tokens]
+            text_dict = self.add_exemplar_tokens(tokenized, text_dict, exemplar_tokens, exemplar_labels)
+        
+        srcs = []
+        masks = []
+        for l, feat in enumerate(features):
+            src, mask = feat.decompose()
+            srcs.append(self.input_proj[l](src))
+            masks.append(mask)
+            assert mask is not None
+        if self.num_feature_levels > len(srcs):
+            _len_srcs = len(srcs)
+            for l in range(_len_srcs, self.num_feature_levels):
+                if l == _len_srcs:
+                    src = self.input_proj[l](features[-1].tensors)
+                else:
+                    src = self.input_proj[l](srcs[-1])
+                m = samples.mask
+                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
+                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
+                srcs.append(src)
+                masks.append(mask)
+                poss.append(pos_l)
+        
+        # Comment out the next two lines after the ablation study is complete
+        #new_attention_mask = torch.ones_like(text_dict["text_self_attention_masks"], dtype=torch.bool)
+        #text_dict["text_self_attention_masks"] = new_attention_mask
+        
+        input_query_bbox = input_query_label = attn_mask = dn_meta = None
+        hs, reference, hs_enc, ref_enc, init_box_proposal = self.transformer(
+            srcs, masks, input_query_bbox, poss, input_query_label, attn_mask, text_dict
+        )
+
+        
+        # deformable-detr-like anchor update
+        outputs_coord_list = []
+        for dec_lid, (layer_ref_sig, layer_bbox_embed, layer_hs) in enumerate(
+            zip(reference[:-1], self.bbox_embed, hs)
+        ):
+            layer_delta_unsig = layer_bbox_embed(layer_hs)
+            layer_outputs_unsig = layer_delta_unsig + inverse_sigmoid(layer_ref_sig)
+            layer_outputs_unsig = layer_outputs_unsig.sigmoid()
+            outputs_coord_list.append(layer_outputs_unsig)
+        outputs_coord_list = torch.stack(outputs_coord_list)
+
+
+        outputs_class = torch.stack(
+            [
+                layer_cls_embed(layer_hs, text_dict)
+                for layer_cls_embed, layer_hs in zip(self.class_embed, hs)
+            ]
+        )
+
+        out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord_list[-1]}
+
+        # Used to calculate losses
+        bs, len_td = text_dict['text_token_mask'].shape
+        out['text_mask']=torch.zeros(bs, self.max_text_len, dtype=torch.bool).to(
+            samples.device
+        )
+        for b in range(bs):
+            for j in range(len_td):
+                if text_dict['text_token_mask'][b][j] == True:
+                    out['text_mask'][b][j] = True
+
+        # for intermediate outputs
+        if self.aux_loss:
+            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord_list)
+        out['token']=one_hot_token
+        # # for encoder output
+        if hs_enc is not None:
+            # prepare intermediate outputs
+            interm_coord = ref_enc[-1]
+            interm_class = self.transformer.enc_out_class_embed(hs_enc[-1], text_dict)
+            out['interm_outputs'] = {'pred_logits': interm_class, 'pred_boxes': interm_coord}
+            out['interm_outputs_for_matching_pre'] = {'pred_logits': interm_class, 'pred_boxes': init_box_proposal}
+
+        # outputs['pred_logits'].shape
+        # torch.Size([4, 900, 256])
+
+        # outputs['pred_boxes'].shape
+        # torch.Size([4, 900, 4])
+
+        # outputs['text_mask'].shape
+        # torch.Size([256])
+
+        # outputs['text_mask']
+
+        # outputs['aux_outputs'][0].keys()
+        # dict_keys(['pred_logits', 'pred_boxes', 'one_hot', 'text_mask'])
+
+        # outputs['aux_outputs'][img_idx]
+
+        # outputs['token']
+        # <class 'transformers.tokenization_utils_base.BatchEncoding'>
+
+        # outputs['interm_outputs'].keys()
+        # dict_keys(['pred_logits', 'pred_boxes', 'one_hot', 'text_mask'])
+
+
+        # outputs['interm_outputs_for_matching_pre'].keys()
+        # dict_keys(['pred_logits', 'pred_boxes'])
+
+        # outputs['one_hot'].shape
+        # torch.Size([4, 900, 256])
+
+        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):
+    def __init__(self, matcher, weight_dict, focal_alpha,focal_gamma, losses):
+        """ Create the criterion.
+        Parameters:
+            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.
+            losses: list of all the losses to be applied. See get_loss for list of available losses.
+            focal_alpha: alpha in Focal Loss
+        """
+        super().__init__()
+        self.matcher = matcher
+        self.weight_dict = weight_dict
+        self.losses = losses
+        self.focal_alpha = focal_alpha
+        self.focal_gamma= focal_gamma
+
+    @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)
+        # Get the normalized exemplars for weak box supervision.
+        exemplars = [t['exemplar_boxes'] for t in targets]
+
+        # Match and isolate exemplar source and targets.
+        src_boxes_exemplar = []
+        target_boxes_exemplar = []
+        bs = len(targets)
+        for ind in range(bs):
+            # For each exemplar, find the corresponding target box.
+            # The 'matching' target box is the one with the closest center (note: dot annotations are not necessarily exactly at the center of the exemplar bounding box).
+            # We use the l2 distance.
+            target_boxes_per_img = target_boxes[idx[0] == ind, :]
+            src_boxes_per_img = src_boxes[idx[0] == ind, :]
+            exemplars_per_img = exemplars[ind]
+
+            for exemplar in exemplars_per_img:
+
+                if target_boxes_per_img.shape[0] > 0:
+                    # Handle case where dots are cropped out but exemplar remains.
+                    dist = torch.sqrt(torch.sum((exemplar[:2] - target_boxes_per_img[:, :2]) ** 2, dim=-1))
+                    _, min_index = torch.min(dist, dim=0)
+                    src_boxes_exemplar.append(src_boxes_per_img[min_index,:])
+                    target_boxes_exemplar.append(exemplar)
+
+        # Handle no exemplars case.
+        if len(src_boxes_exemplar) > 0:
+            src_boxes_exemplar = torch.stack(src_boxes_exemplar, dim=0)
+            target_boxes_exemplar = torch.stack(target_boxes_exemplar, dim=0)
+            loss_hw_bbox = F.l1_loss(src_boxes_exemplar[:, 2:], target_boxes_exemplar[:, 2:], reduction='none')
+            loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(
+            box_ops.box_cxcywh_to_xyxy(src_boxes_exemplar),
+            box_ops.box_cxcywh_to_xyxy(target_boxes_exemplar)))
+        else:
+            loss_hw_bbox = torch.tensor(0).to(src_boxes.device)
+            loss_giou = torch.tensor(0).to(src_boxes.device)
+            
+        loss_center_bbox = F.l1_loss(src_boxes[:, :2], target_boxes[:, :2], reduction='none')
+
+        losses = {}
+        losses['loss_bbox'] = (loss_center_bbox.sum() + loss_hw_bbox.sum()) / num_boxes
+        losses['loss_giou'] = loss_giou.sum() / num_boxes
+
+        # calculate the x,y and h,w loss
+        with torch.no_grad():
+            losses['loss_xy'] = loss_center_bbox.sum() / num_boxes
+            losses['loss_hw'] = loss_hw_bbox.sum() / num_boxes
+
+        return losses
+
+
+    def token_sigmoid_binary_focal_loss(self, outputs, targets, indices, num_boxes):
+        pred_logits=outputs['pred_logits']
+        new_targets=outputs['one_hot'].to(pred_logits.device)
+        text_mask=outputs['text_mask']
+
+        assert (new_targets.dim() == 3)
+        assert (pred_logits.dim() == 3)  # batch x from x to
+        
+        bs, n, _ = pred_logits.shape
+        alpha=self.focal_alpha
+        gamma=self.focal_gamma
+        if text_mask is not None:
+            # ODVG: each sample has different mask 
+            text_mask = text_mask.repeat(1, pred_logits.size(1)).view(outputs['text_mask'].shape[0],-1,outputs['text_mask'].shape[1])
+            pred_logits = torch.masked_select(pred_logits, text_mask)
+            new_targets = torch.masked_select(new_targets, text_mask)
+
+        new_targets=new_targets.float()
+        p = torch.sigmoid(pred_logits)
+        ce_loss = F.binary_cross_entropy_with_logits(pred_logits, new_targets, reduction="none")
+        p_t = p * new_targets + (1 - p) * (1 - new_targets)
+        loss = ce_loss * ((1 - p_t) ** gamma)
+
+        if alpha >= 0:
+            alpha_t = alpha * new_targets + (1 - alpha) * (1 - new_targets)
+            loss = alpha_t * loss
+
+        total_num_pos=0
+        for batch_indices in indices:
+            total_num_pos += len(batch_indices[0])
+        num_pos_avg_per_gpu = max(total_num_pos , 1.0)
+        loss=loss.sum()/num_pos_avg_per_gpu
+        
+        losses = {'loss_ce': loss}
+        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.token_sigmoid_binary_focal_loss,
+            'cardinality': self.loss_cardinality,
+            'boxes': self.loss_boxes,
+        }
+        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, cat_list, caption, return_indices=False):
+        """ 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
+            
+             return_indices: used for vis. if True, the layer0-5 indices will be returned as well.
+        """
+        device=next(iter(outputs.values())).device
+        one_hot = torch.zeros(outputs['pred_logits'].size(),dtype=torch.int64) # torch.Size([bs, 900, 256])
+        token = outputs['token'] 
+        
+        label_map_list = []
+        indices = []
+        for j in range(len(cat_list)): # bs
+            label_map=[]
+            for i in range(len(cat_list[j])):
+                label_id=torch.tensor([i])
+                per_label = create_positive_map_exemplar(token['input_ids'][j], label_id, [101, 102, 1012, 1029])
+                label_map.append(per_label)
+            label_map=torch.stack(label_map,dim=0).squeeze(1)
+            
+            label_map_list.append(label_map)
+        for j in range(len(cat_list)): # bs
+            for_match = {
+                "pred_logits" : outputs['pred_logits'][j].unsqueeze(0),
+                "pred_boxes" : outputs['pred_boxes'][j].unsqueeze(0)
+            }
+            
+            inds = self.matcher(for_match, [targets[j]], label_map_list[j])
+            indices.extend(inds)
+        # indices : 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)
+
+        # import pdb; pdb.set_trace()
+        tgt_ids = [v["labels"].cpu() for v in targets]
+        # len(tgt_ids) == bs
+        for i in range(len(indices)):
+            tgt_ids[i]=tgt_ids[i][indices[i][1]]
+            one_hot[i,indices[i][0]] = label_map_list[i][tgt_ids[i]].to(torch.long)
+        outputs['one_hot'] = one_hot
+        if return_indices:
+            indices0_copy = indices
+            indices_list = []
+
+        # Compute the average number of target boxes accross all nodes, for normalization purposes
+        num_boxes_list = [len(t["labels"]) for t in targets]
+        num_boxes = sum(num_boxes_list)
+        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=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 idx, aux_outputs in enumerate(outputs['aux_outputs']):
+                indices = []
+                for j in range(len(cat_list)): # bs
+                    aux_output_single = {
+                        'pred_logits' : aux_outputs['pred_logits'][j].unsqueeze(0),
+                        'pred_boxes': aux_outputs['pred_boxes'][j].unsqueeze(0)
+                    }
+                    inds = self.matcher(aux_output_single, [targets[j]], label_map_list[j])
+                    indices.extend(inds)
+                one_hot_aux = torch.zeros(outputs['pred_logits'].size(),dtype=torch.int64)
+                tgt_ids = [v["labels"].cpu() for v in targets]
+                for i in range(len(indices)):
+                    tgt_ids[i]=tgt_ids[i][indices[i][1]]
+                    one_hot_aux[i,indices[i][0]] = label_map_list[i][tgt_ids[i]].to(torch.long)
+                aux_outputs['one_hot'] = one_hot_aux
+                aux_outputs['text_mask'] = outputs['text_mask']
+                if return_indices:
+                    indices_list.append(indices)
+                for loss in self.losses:
+                    kwargs = {}
+                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)                
+                    l_dict = {k + f'_{idx}': v for k, v in l_dict.items()}
+                    losses.update(l_dict)
+
+        # interm_outputs loss
+        if 'interm_outputs' in outputs:
+            interm_outputs = outputs['interm_outputs']
+            indices = []
+            for j in range(len(cat_list)): # bs
+                interm_output_single = {
+                    'pred_logits' : interm_outputs['pred_logits'][j].unsqueeze(0),
+                    'pred_boxes': interm_outputs['pred_boxes'][j].unsqueeze(0)
+                }
+                inds = self.matcher(interm_output_single, [targets[j]], label_map_list[j])
+                indices.extend(inds)
+            one_hot_aux = torch.zeros(outputs['pred_logits'].size(),dtype=torch.int64)
+            tgt_ids = [v["labels"].cpu() for v in targets]
+            for i in range(len(indices)):
+                tgt_ids[i]=tgt_ids[i][indices[i][1]]
+                one_hot_aux[i,indices[i][0]] = label_map_list[i][tgt_ids[i]].to(torch.long)
+            interm_outputs['one_hot'] = one_hot_aux
+            interm_outputs['text_mask'] = outputs['text_mask']
+            if return_indices:
+                indices_list.append(indices)
+            for loss in self.losses:
+                kwargs = {}
+                l_dict = self.get_loss(loss, interm_outputs, targets, indices, num_boxes, **kwargs)
+                l_dict = {k + f'_interm': v for k, v in l_dict.items()}
+                losses.update(l_dict)
+
+        if return_indices:
+            indices_list.append(indices0_copy)
+            return losses, indices_list
+
+
+        return losses
+
+
+class PostProcess(nn.Module):
+    """ This module converts the model's output into the format expected by the coco api"""
+    def __init__(self, num_select=100,text_encoder_type='text_encoder_type', nms_iou_threshold=-1,use_coco_eval=False,args=None) -> None:
+        super().__init__()
+        self.num_select = num_select
+        self.tokenizer = get_tokenlizer.get_tokenlizer(text_encoder_type)
+        if args.use_coco_eval:
+            from pycocotools.coco import COCO
+            coco = COCO(args.coco_val_path)
+            category_dict = coco.loadCats(coco.getCatIds())
+            cat_list = [item['name'] for item in category_dict]
+        else:
+            cat_list=args.label_list
+        caption = " . ".join(cat_list) + ' .'
+        tokenized = self.tokenizer(caption, padding="longest", return_tensors="pt")
+        label_list = torch.arange(len(cat_list))
+        pos_map=create_positive_map(tokenized,label_list,cat_list,caption)
+        # build a mapping from label_id to pos_map
+        if args.use_coco_eval:
+            id_map = {0: 1, 1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: 7, 7: 8, 8: 9, 9: 10, 10: 11, 11: 13, 12: 14, 13: 15, 14: 16, 15: 17, 16: 18, 17: 19, 18: 20, 19: 21, 20: 22, 21: 23, 22: 24, 23: 25, 24: 27, 25: 28, 26: 31, 27: 32, 28: 33, 29: 34, 30: 35, 31: 36, 32: 37, 33: 38, 34: 39, 35: 40, 36: 41, 37: 42, 38: 43, 39: 44, 40: 46,
+                    41: 47, 42: 48, 43: 49, 44: 50, 45: 51, 46: 52, 47: 53, 48: 54, 49: 55, 50: 56, 51: 57, 52: 58, 53: 59, 54: 60, 55: 61, 56: 62, 57: 63, 58: 64, 59: 65, 60: 67, 61: 70, 62: 72, 63: 73, 64: 74, 65: 75, 66: 76, 67: 77, 68: 78, 69: 79, 70: 80, 71: 81, 72: 82, 73: 84, 74: 85, 75: 86, 76: 87, 77: 88, 78: 89, 79: 90}
+            new_pos_map = torch.zeros((91, 256))
+            for k, v in id_map.items():
+                new_pos_map[v] = pos_map[k]
+            pos_map=new_pos_map
+
+
+        self.nms_iou_threshold=nms_iou_threshold
+        self.positive_map = pos_map
+
+    @torch.no_grad()
+    def forward(self, outputs, target_sizes, not_to_xyxy=False, test=False):
+        """ 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
+        """
+        num_select = self.num_select
+        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']
+
+
+        prob_to_token = out_logits.sigmoid()
+        pos_maps = self.positive_map.to(prob_to_token.device)
+        for label_ind in range(len(pos_maps)):
+            if pos_maps[label_ind].sum() != 0:
+                pos_maps[label_ind]=pos_maps[label_ind]/pos_maps[label_ind].sum()
+
+        prob_to_label = prob_to_token @ pos_maps.T
+
+        assert len(out_logits) == len(target_sizes)
+        assert target_sizes.shape[1] == 2
+
+        prob = prob_to_label
+        topk_values, topk_indexes = torch.topk(prob.view(prob.shape[0], -1), num_select, dim=1)
+        scores = topk_values
+        topk_boxes = torch.div(topk_indexes, prob.shape[2], rounding_mode='trunc')
+        labels = topk_indexes % prob.shape[2]
+        if not_to_xyxy:
+            boxes = out_bbox
+        else:
+            boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+
+        # if test:
+        #     assert not not_to_xyxy
+        #     boxes[:,:,2:] = boxes[:,:,2:] - boxes[:,:,:2]
+        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1,1,4))
+        
+        # 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, :]
+
+        if self.nms_iou_threshold > 0:
+            item_indices = [nms(b, s, iou_threshold=self.nms_iou_threshold) for b,s in zip(boxes, scores)]
+
+            results = [{'scores': s[i], 'labels': l[i], 'boxes': b[i]} for s, l, b, i in zip(scores, labels, boxes, item_indices)]
+        else:
+            results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
+        results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
+        return results
+
+
+@MODULE_BUILD_FUNCS.registe_with_name(module_name="groundingdino")
+def build_groundingdino(args):
+    device = torch.device(args.device)
+    backbone = build_backbone(args)
+    transformer = build_transformer(args)
+
+    dn_labelbook_size = args.dn_labelbook_size
+    dec_pred_bbox_embed_share = args.dec_pred_bbox_embed_share
+    sub_sentence_present = args.sub_sentence_present
+
+    model = GroundingDINO(
+        backbone,
+        transformer,
+        num_queries=args.num_queries,
+        aux_loss=args.aux_loss,
+        iter_update=True,
+        query_dim=4,
+        num_feature_levels=args.num_feature_levels,
+        nheads=args.nheads,
+        dec_pred_bbox_embed_share=dec_pred_bbox_embed_share,
+        two_stage_type=args.two_stage_type,
+        two_stage_bbox_embed_share=args.two_stage_bbox_embed_share,
+        two_stage_class_embed_share=args.two_stage_class_embed_share,
+        num_patterns=args.num_patterns,
+        dn_number=0,
+        dn_box_noise_scale=args.dn_box_noise_scale,
+        dn_label_noise_ratio=args.dn_label_noise_ratio,
+        dn_labelbook_size=dn_labelbook_size,
+        text_encoder_type=args.text_encoder_type,
+        sub_sentence_present=sub_sentence_present,
+        max_text_len=args.max_text_len,
+    )
+
+
+
+    matcher = build_matcher(args)
+
+    # prepare weight dict
+    weight_dict = {'loss_ce': args.cls_loss_coef, 'loss_bbox': args.bbox_loss_coef}
+    weight_dict['loss_giou'] = args.giou_loss_coef
+    clean_weight_dict_wo_dn = copy.deepcopy(weight_dict)
+
+    
+
+    clean_weight_dict = copy.deepcopy(weight_dict)
+
+    # 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 clean_weight_dict.items()})
+        weight_dict.update(aux_weight_dict)
+
+    if args.two_stage_type != 'no':
+        interm_weight_dict = {}
+        try:
+            no_interm_box_loss = args.no_interm_box_loss
+        except:
+            no_interm_box_loss = False
+        _coeff_weight_dict = {
+            'loss_ce': 1.0,
+            'loss_bbox': 1.0 if not no_interm_box_loss else 0.0,
+            'loss_giou': 1.0 if not no_interm_box_loss else 0.0,
+        }
+        try:
+            interm_loss_coef = args.interm_loss_coef
+        except:
+            interm_loss_coef = 1.0
+        interm_weight_dict.update({k + f'_interm': v * interm_loss_coef * _coeff_weight_dict[k] for k, v in clean_weight_dict_wo_dn.items()})
+        weight_dict.update(interm_weight_dict)
+
+    # losses = ['labels', 'boxes', 'cardinality']
+    losses = ['labels', 'boxes']
+
+    criterion = SetCriterion(matcher=matcher, weight_dict=weight_dict,
+                             focal_alpha=args.focal_alpha, focal_gamma=args.focal_gamma,losses=losses
+                             )
+    criterion.to(device)
+    postprocessors = {'bbox': PostProcess(num_select=args.num_select  , text_encoder_type=args.text_encoder_type,nms_iou_threshold=args.nms_iou_threshold,args=args)}
+
+    return model, criterion, postprocessors
+
+def create_positive_map(tokenized, tokens_positive,cat_list,caption):
+    """construct a map such that positive_map[i,j] = True iff box i is associated to token j"""
+    positive_map = torch.zeros((len(tokens_positive), 256), dtype=torch.float)
+
+    for j,label in enumerate(tokens_positive):
+
+        start_ind = caption.find(cat_list[label])
+        end_ind = start_ind + len(cat_list[label]) - 1
+        beg_pos = tokenized.char_to_token(start_ind)
+        try:
+            end_pos = tokenized.char_to_token(end_ind)
+        except:
+            end_pos = None
+        if end_pos is None:
+            try:
+                end_pos = tokenized.char_to_token(end_ind - 1)
+                if end_pos is None:
+                    end_pos = tokenized.char_to_token(end_ind - 2)
+            except:
+                end_pos = None
+        # except Exception as e:
+        #     print("beg:", beg, "end:", end)
+        #     print("token_positive:", tokens_positive)
+        #     # print("beg_pos:", beg_pos, "end_pos:", end_pos)
+        #     raise e
+        # if beg_pos is None:
+        #     try:
+        #         beg_pos = tokenized.char_to_token(beg + 1)
+        #         if beg_pos is None:
+        #             beg_pos = tokenized.char_to_token(beg + 2)
+        #     except:
+        #         beg_pos = None
+        # if end_pos is None:
+        #     try:
+        #         end_pos = tokenized.char_to_token(end - 2)
+        #         if end_pos is None:
+        #             end_pos = tokenized.char_to_token(end - 3)
+        #     except:
+        #         end_pos = None
+        if beg_pos is None or end_pos is None:
+            continue
+        if beg_pos < 0 or end_pos < 0:
+            continue
+        if beg_pos > end_pos:
+            continue
+        # assert beg_pos is not None and end_pos is not None
+        positive_map[j,beg_pos: end_pos + 1].fill_(1)
+    return positive_map 
+
+
+def create_positive_map_exemplar(input_ids, label, special_tokens):
+    tokens_positive = torch.zeros(256, dtype=torch.float)
+    count = -1
+    for token_ind in range(len(input_ids)):
+        input_id = input_ids[token_ind]
+        if (input_id not in special_tokens) and (token_ind == 0 or (input_ids[token_ind - 1] in special_tokens)):
+            count += 1
+        if count == label:
+            ind_to_insert_ones = token_ind
+            
+            while input_ids[ind_to_insert_ones] not in special_tokens:
+                tokens_positive[ind_to_insert_ones] = 1
+                ind_to_insert_ones += 1
+            break
+    return tokens_positive
+
+
+

models/GroundingDINO/groundingdino_app.py

@@ -0,0 +1,1257 @@
+# ------------------------------------------------------------------------
+# Grounding DINO
+# url: https://github.com/IDEA-Research/GroundingDINO
+# Copyright (c) 2023 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Conditional DETR model and criterion classes.
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# ------------------------------------------------------------------------
+import copy
+from typing import List
+import torchvision.transforms.functional as vis_F
+from torchvision.transforms import InterpolationMode
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision.ops.boxes import nms
+from torchvision.ops import roi_align
+from transformers import (
+    AutoTokenizer,
+    BertModel,
+    BertTokenizer,
+    RobertaModel,
+    RobertaTokenizerFast,
+)
+
+from groundingdino.util import box_ops, get_tokenlizer
+from groundingdino.util.misc import (
+    NestedTensor,
+    accuracy,
+    get_world_size,
+    interpolate,
+    inverse_sigmoid,
+    is_dist_avail_and_initialized,
+    nested_tensor_from_tensor_list,
+)
+from groundingdino.util.utils import get_phrases_from_posmap
+from groundingdino.util.visualizer import COCOVisualizer
+from groundingdino.util.vl_utils import create_positive_map_from_span
+
+from ..registry import MODULE_BUILD_FUNCS
+from .backbone import build_backbone
+from .bertwarper import (
+    BertModelWarper,
+    generate_masks_with_special_tokens,
+    generate_masks_with_special_tokens_and_transfer_map,
+)
+from .transformer import build_transformer
+from .utils import MLP, ContrastiveEmbed, sigmoid_focal_loss
+
+from .matcher import build_matcher
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from groundingdino.util.visualizer import renorm
+
+
+def numpy_2_cv2(np_img):
+    if np.min(np_img) < 0:
+        raise Exception("image min is less than 0. Img min: " + str(np.min(np_img)))
+    if np.max(np_img) > 1:
+        raise Exception("image max is greater than 1. Img max: " + str(np.max(np_img)))
+    np_img = (np_img * 255).astype(np.uint8)
+    # Need to somehow ensure image is in RGB format. Note this line shows up in SAM demo: image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    cv2_image = np.asarray(np_img)
+    return cv2_image
+
+
+def vis_exemps(image, exemp, f_name):
+    plt.imshow(image)
+    plt.gca().add_patch(
+        Rectangle(
+            (exemp[0], exemp[1]),
+            exemp[2] - exemp[0],
+            exemp[3] - exemp[1],
+            edgecolor="red",
+            facecolor="none",
+            lw=1,
+        )
+    )
+    plt.savefig(f_name)
+    plt.close()
+
+
+class GroundingDINO(nn.Module):
+    """This is the Cross-Attention Detector module that performs object detection"""
+
+    def __init__(
+        self,
+        backbone,
+        transformer,
+        num_queries,
+        aux_loss=False,
+        iter_update=False,
+        query_dim=2,
+        num_feature_levels=1,
+        nheads=8,
+        # two stage
+        two_stage_type="no",  # ['no', 'standard']
+        dec_pred_bbox_embed_share=True,
+        two_stage_class_embed_share=True,
+        two_stage_bbox_embed_share=True,
+        num_patterns=0,
+        dn_number=100,
+        dn_box_noise_scale=0.4,
+        dn_label_noise_ratio=0.5,
+        dn_labelbook_size=100,
+        text_encoder_type="bert-base-uncased",
+        sub_sentence_present=True,
+        max_text_len=256,
+    ):
+        """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_queries: number of object queries, ie detection slot. This is the maximal number of objects
+                         Conditional 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
+        self.hidden_dim = hidden_dim = transformer.d_model
+        self.num_feature_levels = num_feature_levels
+        self.nheads = nheads
+        self.max_text_len = max_text_len
+        self.sub_sentence_present = sub_sentence_present
+
+        # setting query dim
+        self.query_dim = query_dim
+        assert query_dim == 4
+
+        # visual exemplar cropping
+        self.feature_map_proj = nn.Conv2d((256 + 512 + 1024), hidden_dim, kernel_size=1)
+
+        # for dn training
+        self.num_patterns = num_patterns
+        self.dn_number = dn_number
+        self.dn_box_noise_scale = dn_box_noise_scale
+        self.dn_label_noise_ratio = dn_label_noise_ratio
+        self.dn_labelbook_size = dn_labelbook_size
+
+        # bert
+        self.tokenizer = get_tokenlizer.get_tokenlizer(text_encoder_type)
+        self.bert = get_tokenlizer.get_pretrained_language_model(text_encoder_type)
+        self.bert.pooler.dense.weight.requires_grad_(False)
+        self.bert.pooler.dense.bias.requires_grad_(False)
+        self.bert = BertModelWarper(bert_model=self.bert)
+
+        self.feat_map = nn.Linear(
+            self.bert.config.hidden_size, self.hidden_dim, bias=True
+        )
+        nn.init.constant_(self.feat_map.bias.data, 0)
+        nn.init.xavier_uniform_(self.feat_map.weight.data)
+        # freeze
+
+        # special tokens
+        self.specical_tokens = self.tokenizer.convert_tokens_to_ids(
+            ["[CLS]", "[SEP]", ".", "?"]
+        )
+
+        # prepare input projection layers
+        if num_feature_levels > 1:
+            num_backbone_outs = len(backbone.num_channels)
+            input_proj_list = []
+            for _ in range(num_backbone_outs):
+                in_channels = backbone.num_channels[_]
+                input_proj_list.append(
+                    nn.Sequential(
+                        nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
+                        nn.GroupNorm(32, hidden_dim),
+                    )
+                )
+            for _ in range(num_feature_levels - num_backbone_outs):
+                input_proj_list.append(
+                    nn.Sequential(
+                        nn.Conv2d(
+                            in_channels, hidden_dim, kernel_size=3, stride=2, padding=1
+                        ),
+                        nn.GroupNorm(32, hidden_dim),
+                    )
+                )
+                in_channels = hidden_dim
+            self.input_proj = nn.ModuleList(input_proj_list)
+        else:
+            assert (
+                two_stage_type == "no"
+            ), "two_stage_type should be no if num_feature_levels=1 !!!"
+            self.input_proj = nn.ModuleList(
+                [
+                    nn.Sequential(
+                        nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1),
+                        nn.GroupNorm(32, hidden_dim),
+                    )
+                ]
+            )
+
+        self.backbone = backbone
+        self.aux_loss = aux_loss
+        self.box_pred_damping = box_pred_damping = None
+
+        self.iter_update = iter_update
+        assert iter_update, "Why not iter_update?"
+
+        # prepare pred layers
+        self.dec_pred_bbox_embed_share = dec_pred_bbox_embed_share
+        # prepare class & box embed
+        _class_embed = ContrastiveEmbed()
+
+        _bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
+        nn.init.constant_(_bbox_embed.layers[-1].weight.data, 0)
+        nn.init.constant_(_bbox_embed.layers[-1].bias.data, 0)
+
+        if dec_pred_bbox_embed_share:
+            box_embed_layerlist = [
+                _bbox_embed for i in range(transformer.num_decoder_layers)
+            ]
+        else:
+            box_embed_layerlist = [
+                copy.deepcopy(_bbox_embed)
+                for i in range(transformer.num_decoder_layers)
+            ]
+        class_embed_layerlist = [
+            _class_embed for i in range(transformer.num_decoder_layers)
+        ]
+        self.bbox_embed = nn.ModuleList(box_embed_layerlist)
+        self.class_embed = nn.ModuleList(class_embed_layerlist)
+        self.transformer.decoder.bbox_embed = self.bbox_embed
+        self.transformer.decoder.class_embed = self.class_embed
+
+        # two stage
+        self.two_stage_type = two_stage_type
+        assert two_stage_type in [
+            "no",
+            "standard",
+        ], "unknown param {} of two_stage_type".format(two_stage_type)
+        if two_stage_type != "no":
+            if two_stage_bbox_embed_share:
+                assert dec_pred_bbox_embed_share
+                self.transformer.enc_out_bbox_embed = _bbox_embed
+            else:
+                self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed)
+
+            if two_stage_class_embed_share:
+                assert dec_pred_bbox_embed_share
+                self.transformer.enc_out_class_embed = _class_embed
+            else:
+                self.transformer.enc_out_class_embed = copy.deepcopy(_class_embed)
+
+            self.refpoint_embed = None
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        # init input_proj
+        for proj in self.input_proj:
+            nn.init.xavier_uniform_(proj[0].weight, gain=1)
+            nn.init.constant_(proj[0].bias, 0)
+
+    def init_ref_points(self, use_num_queries):
+        self.refpoint_embed = nn.Embedding(use_num_queries, self.query_dim)
+
+    def add_exemplar_tokens(self, tokenized, text_dict, exemplar_tokens, labels):
+        input_ids = tokenized["input_ids"]
+
+        device = input_ids.device
+        new_input_ids = []
+        encoded_text = text_dict["encoded_text"]
+        new_encoded_text = []
+        text_token_mask = text_dict["text_token_mask"]
+        new_text_token_mask = []
+        position_ids = text_dict["position_ids"]
+        text_self_attention_masks = text_dict["text_self_attention_masks"]
+
+        for sample_ind in range(len(labels)):
+            label = labels[sample_ind][0]
+            exemplars = exemplar_tokens[sample_ind]
+            label_count = -1
+            assert len(input_ids[sample_ind]) == len(position_ids[sample_ind])
+            for token_ind in range(len(input_ids[sample_ind])):
+                input_id = input_ids[sample_ind][token_ind]
+                if (token_ind > 0 and input_ids[sample_ind][token_ind - 1] in self.specical_tokens):
+                    label_count += 1
+                if label_count == label:
+                    # Get the index where to insert the exemplar tokens.
+                    ind_to_insert_exemplar = token_ind
+                    while (
+                        input_ids[sample_ind][ind_to_insert_exemplar]
+                        not in self.specical_tokens
+                    ):
+                        ind_to_insert_exemplar += 1
+                    break
+
+            # Handle no text case.
+            if label_count == -1:
+                ind_to_insert_exemplar = 1
+            # * token indicates exemplar.
+            new_input_ids.append(
+                torch.cat(
+                    [
+                        input_ids[sample_ind][:ind_to_insert_exemplar],
+                        torch.tensor([1008] * exemplars.shape[0]).to(device),
+                        input_ids[sample_ind][ind_to_insert_exemplar:],
+                    ]
+                )
+            )
+            new_encoded_text.append(
+                torch.cat(
+                    [
+                        encoded_text[sample_ind][:ind_to_insert_exemplar, :],
+                        exemplars,
+                        encoded_text[sample_ind][ind_to_insert_exemplar:, :],
+                    ]
+                )
+            )
+            new_text_token_mask.append(
+                torch.full((len(new_input_ids[sample_ind]),), True).to(device)
+            )
+
+        tokenized["input_ids"] = torch.stack(new_input_ids)
+
+        (
+            text_self_attention_masks,
+            position_ids,
+            _,
+        ) = generate_masks_with_special_tokens_and_transfer_map(
+            tokenized, self.specical_tokens, None
+        )
+
+        return ({
+            "encoded_text": torch.stack(new_encoded_text),
+            "text_token_mask": torch.stack(new_text_token_mask),
+            "position_ids": position_ids,
+            "text_self_attention_masks": text_self_attention_masks,
+        }, tokenized)
+
+    def combine_features(self, features):
+        (bs, c, h, w) = (
+            features[0].decompose()[0].shape[-4],
+            features[0].decompose()[0].shape[-3],
+            features[0].decompose()[0].shape[-2],
+            features[0].decompose()[0].shape[-1],
+        )
+
+        x = torch.cat(
+            [
+                F.interpolate(
+                    feat.decompose()[0],
+                    size=(h, w),
+                    mode="bilinear",
+                    align_corners=True,
+                )
+                for feat in features
+            ],
+            dim=1,
+        )
+
+        x = self.feature_map_proj(x)
+
+        return x
+
+    def forward(
+        self,
+        input_images: NestedTensor,
+        input_image_pos_exemplars: NestedTensor,
+        pos_exemplars: List,
+        input_images_neg_exemplars: List[NestedTensor],
+        neg_exemplars: List,
+        targets: List = None,
+        cropped=False,
+        orig_img=None,
+        crop_width=0,
+        crop_height=0,
+        **kw,
+    ):
+        """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]
+           - "pred_boxes": The normalized boxes coordinates for all queries, represented as
+                           (center_x, center_y, width, height). 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 targets is None:
+            captions = kw["captions"]
+        else:
+            captions = [t["caption"] for t in targets]
+
+        # encoder texts
+
+        tokenized = self.tokenizer(captions, padding="longest", return_tensors="pt").to(
+            input_images.device
+        )
+
+        one_hot_token = tokenized
+
+        (
+            text_self_attention_masks,
+            position_ids,
+            cate_to_token_mask_list,
+        ) = generate_masks_with_special_tokens_and_transfer_map(
+            tokenized, self.specical_tokens, self.tokenizer
+        )
+
+        if text_self_attention_masks.shape[1] > self.max_text_len:
+            text_self_attention_masks = text_self_attention_masks[
+                :, : self.max_text_len, : self.max_text_len
+            ]
+            position_ids = position_ids[:, : self.max_text_len]
+            tokenized["input_ids"] = tokenized["input_ids"][:, : self.max_text_len]
+            tokenized["attention_mask"] = tokenized["attention_mask"][
+                :, : self.max_text_len
+            ]
+            tokenized["token_type_ids"] = tokenized["token_type_ids"][
+                :, : self.max_text_len
+            ]
+
+        # extract text embeddings
+        if self.sub_sentence_present:
+            tokenized_for_encoder = {
+                k: v for k, v in tokenized.items() if k != "attention_mask"
+            }
+            tokenized_for_encoder["attention_mask"] = text_self_attention_masks
+            tokenized_for_encoder["position_ids"] = position_ids
+        else:
+            tokenized_for_encoder = tokenized
+
+        bert_output = self.bert(**tokenized_for_encoder)  # bs, 195, 768
+
+        encoded_text = self.feat_map(
+            bert_output["last_hidden_state"]
+        )  # bs, 195, d_model
+        text_token_mask = tokenized.attention_mask.bool()  # bs, 195
+        # text_token_mask: True for nomask, False for mask
+        # text_self_attention_masks: True for nomask, False for mask
+
+        if encoded_text.shape[1] > self.max_text_len:
+            encoded_text = encoded_text[:, : self.max_text_len, :]
+            text_token_mask = text_token_mask[:, : self.max_text_len]
+            position_ids = position_ids[:, : self.max_text_len]
+            text_self_attention_masks = text_self_attention_masks[
+                :, : self.max_text_len, : self.max_text_len
+            ]
+
+        text_dict = {
+            "encoded_text": encoded_text,  # bs, 195, d_model
+            "text_token_mask": text_token_mask,  # bs, 195
+            "position_ids": position_ids,  # bs, 195
+            "text_self_attention_masks": text_self_attention_masks,  # bs, 195,195
+        }
+
+        if isinstance(input_images, (list, torch.Tensor)):
+            input_images = nested_tensor_from_tensor_list(input_images)
+
+        features, poss = self.backbone(input_images)
+
+        # Positive exemplars
+        features_pos_exemp, _ = self.backbone(input_image_pos_exemplars)
+        combined_pos_features = self.combine_features(features_pos_exemp)
+        # Get positive visual exemplar tokens.
+        bs = len(pos_exemplars)
+        num_pos_exemplars = pos_exemplars[0].shape[0]
+        if num_pos_exemplars > 0:
+            pos_exemplar_tokens = (
+                roi_align(
+                    combined_pos_features,
+                    boxes=pos_exemplars,
+                    output_size=(1, 1),
+                    spatial_scale=(1 / 8),
+                    aligned=True,
+                )
+                .squeeze(-1)
+                .squeeze(-1)
+                .reshape(bs, num_pos_exemplars, -1)
+            )
+        else:
+            pos_exemplar_tokens = None
+
+        if pos_exemplar_tokens is not None:
+            text_dict, tokenized = self.add_exemplar_tokens(
+                tokenized, text_dict, pos_exemplar_tokens, [torch.tensor([0]).to(input_images.device)]
+            )
+
+        # Negative exemplars
+        num_neg_exemplar_imgs = len(input_images_neg_exemplars)
+        for ind in range(num_neg_exemplar_imgs):
+            exemplar_img = input_images_neg_exemplars[ind]
+            exemplars = neg_exemplars[ind]
+            bs = 1
+            num_neg_exemplars = len(exemplars[0]) # Each set of exemplars has a batch dimension for compatibility with roialign
+            if exemplar_img is not None and num_neg_exemplars > 0:
+                features_neg_exemp, _ = self.backbone(exemplar_img)
+                combined_neg_features = self.combine_features(features_neg_exemp)
+                # Get negative visual exemplar tokens.
+                neg_exemplar_tokens = (
+                    roi_align(
+                        combined_neg_features,
+                        boxes=exemplars,
+                        output_size=(1, 1),
+                        spatial_scale=(1 / 8),
+                        aligned=True,
+                    )
+                    .squeeze(-1)
+                    .squeeze(-1)
+                    .reshape(bs, num_neg_exemplars, -1)
+                )
+                text_dict, tokenized = self.add_exemplar_tokens(
+                    tokenized, text_dict, neg_exemplar_tokens, [torch.tensor([ind + 1]).to(input_images.device)]
+                )
+
+        srcs = []
+        masks = []
+        for l, feat in enumerate(features):
+            src, mask = feat.decompose()
+            srcs.append(self.input_proj[l](src))
+            masks.append(mask)
+            assert mask is not None
+        if self.num_feature_levels > len(srcs):
+            _len_srcs = len(srcs)
+            for l in range(_len_srcs, self.num_feature_levels):
+                if l == _len_srcs:
+                    src = self.input_proj[l](features[-1].tensors)
+                else:
+                    src = self.input_proj[l](srcs[-1])
+                m = input_images.mask
+                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(
+                    torch.bool
+                )[0]
+                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
+                srcs.append(src)
+                masks.append(mask)
+                poss.append(pos_l)
+
+        input_query_bbox = input_query_label = attn_mask = dn_meta = None
+        hs, reference, hs_enc, ref_enc, init_box_proposal = self.transformer(
+            srcs, masks, input_query_bbox, poss, input_query_label, attn_mask, text_dict
+        )
+
+        # deformable-detr-like anchor update
+        outputs_coord_list = []
+        for dec_lid, (layer_ref_sig, layer_bbox_embed, layer_hs) in enumerate(
+            zip(reference[:-1], self.bbox_embed, hs)
+        ):
+            layer_delta_unsig = layer_bbox_embed(layer_hs)
+            layer_outputs_unsig = layer_delta_unsig + inverse_sigmoid(layer_ref_sig)
+            layer_outputs_unsig = layer_outputs_unsig.sigmoid()
+            outputs_coord_list.append(layer_outputs_unsig)
+        outputs_coord_list = torch.stack(outputs_coord_list)
+
+        outputs_class = torch.stack(
+            [
+                layer_cls_embed(layer_hs, text_dict)
+                for layer_cls_embed, layer_hs in zip(self.class_embed, hs)
+            ]
+        )
+
+        out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord_list[-1]}
+
+        # Used to calculate losses
+        bs, len_td = text_dict["text_token_mask"].shape
+        out["text_mask"] = torch.zeros(bs, self.max_text_len, dtype=torch.bool).to(
+            input_images.device
+        )
+        for b in range(bs):
+            for j in range(len_td):
+                if text_dict["text_token_mask"][b][j] == True:
+                    out["text_mask"][b][j] = True
+
+        # for intermediate outputs
+        if self.aux_loss:
+            out["aux_outputs"] = self._set_aux_loss(outputs_class, outputs_coord_list)
+        out["token"] = one_hot_token
+        # # for encoder output
+        if hs_enc is not None:
+            # prepare intermediate outputs
+            interm_coord = ref_enc[-1]
+            interm_class = self.transformer.enc_out_class_embed(hs_enc[-1], text_dict)
+            out["interm_outputs"] = {
+                "pred_logits": interm_class,
+                "pred_boxes": interm_coord,
+            }
+            out["interm_outputs_for_matching_pre"] = {
+                "pred_logits": interm_class,
+                "pred_boxes": init_box_proposal,
+            }
+
+        out["input_ids"] = tokenized["input_ids"]
+
+        # outputs['pred_logits'].shape
+        # torch.Size([4, 900, 256])
+
+        # outputs['pred_boxes'].shape
+        # torch.Size([4, 900, 4])
+
+        # outputs['text_mask'].shape
+        # torch.Size([256])
+
+        # outputs['text_mask']
+
+        # outputs['aux_outputs'][0].keys()
+        # dict_keys(['pred_logits', 'pred_boxes', 'one_hot', 'text_mask'])
+
+        # outputs['aux_outputs'][img_idx]
+
+        # outputs['token']
+        # <class 'transformers.tokenization_utils_base.BatchEncoding'>
+
+        # outputs['interm_outputs'].keys()
+        # dict_keys(['pred_logits', 'pred_boxes', 'one_hot', 'text_mask'])
+
+        # outputs['interm_outputs_for_matching_pre'].keys()
+        # dict_keys(['pred_logits', 'pred_boxes'])
+
+        # outputs['one_hot'].shape
+        # torch.Size([4, 900, 256])
+
+        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):
+    def __init__(self, matcher, weight_dict, focal_alpha, focal_gamma, losses):
+        """Create the criterion.
+        Parameters:
+            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.
+            losses: list of all the losses to be applied. See get_loss for list of available losses.
+            focal_alpha: alpha in Focal Loss
+        """
+        super().__init__()
+        self.matcher = matcher
+        self.weight_dict = weight_dict
+        self.losses = losses
+        self.focal_alpha = focal_alpha
+        self.focal_gamma = focal_gamma
+
+    @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[:, :2], target_boxes[:, :2], 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
+
+        # calculate the x,y and h,w loss
+        with torch.no_grad():
+            losses["loss_xy"] = loss_bbox[..., :2].sum() / num_boxes
+            losses["loss_hw"] = loss_bbox[..., 2:].sum() / num_boxes
+
+        return losses
+
+    def token_sigmoid_binary_focal_loss(self, outputs, targets, indices, num_boxes):
+        pred_logits = outputs["pred_logits"]
+        new_targets = outputs["one_hot"].to(pred_logits.device)
+        text_mask = outputs["text_mask"]
+
+        assert new_targets.dim() == 3
+        assert pred_logits.dim() == 3  # batch x from x to
+
+        bs, n, _ = pred_logits.shape
+        alpha = self.focal_alpha
+        gamma = self.focal_gamma
+        if text_mask is not None:
+            # ODVG: each sample has different mask
+            text_mask = text_mask.repeat(1, pred_logits.size(1)).view(
+                outputs["text_mask"].shape[0], -1, outputs["text_mask"].shape[1]
+            )
+            pred_logits = torch.masked_select(pred_logits, text_mask)
+            new_targets = torch.masked_select(new_targets, text_mask)
+
+        new_targets = new_targets.float()
+        p = torch.sigmoid(pred_logits)
+        ce_loss = F.binary_cross_entropy_with_logits(
+            pred_logits, new_targets, reduction="none"
+        )
+        p_t = p * new_targets + (1 - p) * (1 - new_targets)
+        loss = ce_loss * ((1 - p_t) ** gamma)
+
+        if alpha >= 0:
+            alpha_t = alpha * new_targets + (1 - alpha) * (1 - new_targets)
+            loss = alpha_t * loss
+
+        total_num_pos = 0
+        for batch_indices in indices:
+            total_num_pos += len(batch_indices[0])
+        num_pos_avg_per_gpu = max(total_num_pos, 1.0)
+        loss = loss.sum() / num_pos_avg_per_gpu
+
+        losses = {"loss_ce": loss}
+        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.token_sigmoid_binary_focal_loss,
+            "cardinality": self.loss_cardinality,
+            "boxes": self.loss_boxes,
+        }
+        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, cat_list, caption, return_indices=False):
+        """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
+             return_indices: used for vis. if True, the layer0-5 indices will be returned as well.
+        """
+        device = next(iter(outputs.values())).device
+        one_hot = torch.zeros(
+            outputs["pred_logits"].size(), dtype=torch.int64
+        )  # torch.Size([bs, 900, 256])
+        token = outputs["token"]
+
+        label_map_list = []
+        indices = []
+        for j in range(len(cat_list)):  # bs
+            label_map = []
+            for i in range(len(cat_list[j])):
+                label_id = torch.tensor([i])
+                per_label = create_positive_map_exemplar(
+                    token["input_ids"][j], label_id, [101, 102, 1012, 1029]
+                )
+                label_map.append(per_label)
+            label_map = torch.stack(label_map, dim=0).squeeze(1)
+
+            label_map_list.append(label_map)
+        for j in range(len(cat_list)):  # bs
+            for_match = {
+                "pred_logits": outputs["pred_logits"][j].unsqueeze(0),
+                "pred_boxes": outputs["pred_boxes"][j].unsqueeze(0),
+            }
+
+            inds = self.matcher(for_match, [targets[j]], label_map_list[j])
+            indices.extend(inds)
+        # indices : 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)
+
+        # import pdb; pdb.set_trace()
+        tgt_ids = [v["labels"].cpu() for v in targets]
+        # len(tgt_ids) == bs
+        for i in range(len(indices)):
+            tgt_ids[i] = tgt_ids[i][indices[i][1]]
+            one_hot[i, indices[i][0]] = label_map_list[i][tgt_ids[i]].to(torch.long)
+        outputs["one_hot"] = one_hot
+        if return_indices:
+            indices0_copy = indices
+            indices_list = []
+
+        # Compute the average number of target boxes accross all nodes, for normalization purposes
+        num_boxes_list = [len(t["labels"]) for t in targets]
+        num_boxes = sum(num_boxes_list)
+        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=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 idx, aux_outputs in enumerate(outputs["aux_outputs"]):
+                indices = []
+                for j in range(len(cat_list)):  # bs
+                    aux_output_single = {
+                        "pred_logits": aux_outputs["pred_logits"][j].unsqueeze(0),
+                        "pred_boxes": aux_outputs["pred_boxes"][j].unsqueeze(0),
+                    }
+                    inds = self.matcher(
+                        aux_output_single, [targets[j]], label_map_list[j]
+                    )
+                    indices.extend(inds)
+                one_hot_aux = torch.zeros(
+                    outputs["pred_logits"].size(), dtype=torch.int64
+                )
+                tgt_ids = [v["labels"].cpu() for v in targets]
+                for i in range(len(indices)):
+                    tgt_ids[i] = tgt_ids[i][indices[i][1]]
+                    one_hot_aux[i, indices[i][0]] = label_map_list[i][tgt_ids[i]].to(
+                        torch.long
+                    )
+                aux_outputs["one_hot"] = one_hot_aux
+                aux_outputs["text_mask"] = outputs["text_mask"]
+                if return_indices:
+                    indices_list.append(indices)
+                for loss in self.losses:
+                    kwargs = {}
+                    l_dict = self.get_loss(
+                        loss, aux_outputs, targets, indices, num_boxes, **kwargs
+                    )
+                    l_dict = {k + f"_{idx}": v for k, v in l_dict.items()}
+                    losses.update(l_dict)
+
+        # interm_outputs loss
+        if "interm_outputs" in outputs:
+            interm_outputs = outputs["interm_outputs"]
+            indices = []
+            for j in range(len(cat_list)):  # bs
+                interm_output_single = {
+                    "pred_logits": interm_outputs["pred_logits"][j].unsqueeze(0),
+                    "pred_boxes": interm_outputs["pred_boxes"][j].unsqueeze(0),
+                }
+                inds = self.matcher(
+                    interm_output_single, [targets[j]], label_map_list[j]
+                )
+                indices.extend(inds)
+            one_hot_aux = torch.zeros(outputs["pred_logits"].size(), dtype=torch.int64)
+            tgt_ids = [v["labels"].cpu() for v in targets]
+            for i in range(len(indices)):
+                tgt_ids[i] = tgt_ids[i][indices[i][1]]
+                one_hot_aux[i, indices[i][0]] = label_map_list[i][tgt_ids[i]].to(
+                    torch.long
+                )
+            interm_outputs["one_hot"] = one_hot_aux
+            interm_outputs["text_mask"] = outputs["text_mask"]
+            if return_indices:
+                indices_list.append(indices)
+            for loss in self.losses:
+                kwargs = {}
+                l_dict = self.get_loss(
+                    loss, interm_outputs, targets, indices, num_boxes, **kwargs
+                )
+                l_dict = {k + f"_interm": v for k, v in l_dict.items()}
+                losses.update(l_dict)
+
+        if return_indices:
+            indices_list.append(indices0_copy)
+            return losses, indices_list
+
+        return losses
+
+
+class PostProcess(nn.Module):
+    """This module converts the model's output into the format expected by the coco api"""
+
+    def __init__(
+        self,
+        num_select=100,
+        text_encoder_type="text_encoder_type",
+        nms_iou_threshold=-1,
+        use_coco_eval=False,
+        args=None,
+    ) -> None:
+        super().__init__()
+        self.num_select = num_select
+        self.tokenizer = get_tokenlizer.get_tokenlizer(text_encoder_type)
+        if args.use_coco_eval:
+            from pycocotools.coco import COCO
+
+            coco = COCO(args.coco_val_path)
+            category_dict = coco.loadCats(coco.getCatIds())
+            cat_list = [item["name"] for item in category_dict]
+        else:
+            cat_list = args.label_list
+        caption = " . ".join(cat_list) + " ."
+        tokenized = self.tokenizer(caption, padding="longest", return_tensors="pt")
+        label_list = torch.arange(len(cat_list))
+        pos_map = create_positive_map(tokenized, label_list, cat_list, caption)
+        # build a mapping from label_id to pos_map
+        if args.use_coco_eval:
+            id_map = {
+                0: 1,
+                1: 2,
+                2: 3,
+                3: 4,
+                4: 5,
+                5: 6,
+                6: 7,
+                7: 8,
+                8: 9,
+                9: 10,
+                10: 11,
+                11: 13,
+                12: 14,
+                13: 15,
+                14: 16,
+                15: 17,
+                16: 18,
+                17: 19,
+                18: 20,
+                19: 21,
+                20: 22,
+                21: 23,
+                22: 24,
+                23: 25,
+                24: 27,
+                25: 28,
+                26: 31,
+                27: 32,
+                28: 33,
+                29: 34,
+                30: 35,
+                31: 36,
+                32: 37,
+                33: 38,
+                34: 39,
+                35: 40,
+                36: 41,
+                37: 42,
+                38: 43,
+                39: 44,
+                40: 46,
+                41: 47,
+                42: 48,
+                43: 49,
+                44: 50,
+                45: 51,
+                46: 52,
+                47: 53,
+                48: 54,
+                49: 55,
+                50: 56,
+                51: 57,
+                52: 58,
+                53: 59,
+                54: 60,
+                55: 61,
+                56: 62,
+                57: 63,
+                58: 64,
+                59: 65,
+                60: 67,
+                61: 70,
+                62: 72,
+                63: 73,
+                64: 74,
+                65: 75,
+                66: 76,
+                67: 77,
+                68: 78,
+                69: 79,
+                70: 80,
+                71: 81,
+                72: 82,
+                73: 84,
+                74: 85,
+                75: 86,
+                76: 87,
+                77: 88,
+                78: 89,
+                79: 90,
+            }
+            new_pos_map = torch.zeros((91, 256))
+            for k, v in id_map.items():
+                new_pos_map[v] = pos_map[k]
+            pos_map = new_pos_map
+
+        self.nms_iou_threshold = nms_iou_threshold
+        self.positive_map = pos_map
+
+    @torch.no_grad()
+    def forward(self, outputs, target_sizes, not_to_xyxy=False, test=False):
+        """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
+        """
+        num_select = self.num_select
+        out_logits, out_bbox = outputs["pred_logits"], outputs["pred_boxes"]
+
+        prob_to_token = out_logits.sigmoid()
+        pos_maps = self.positive_map.to(prob_to_token.device)
+        for label_ind in range(len(pos_maps)):
+            if pos_maps[label_ind].sum() != 0:
+                pos_maps[label_ind] = pos_maps[label_ind] / pos_maps[label_ind].sum()
+
+        prob_to_label = prob_to_token @ pos_maps.T
+
+        assert len(out_logits) == len(target_sizes)
+        assert target_sizes.shape[1] == 2
+
+        prob = prob_to_label
+        topk_values, topk_indexes = torch.topk(
+            prob.view(prob.shape[0], -1), num_select, dim=1
+        )
+        scores = topk_values
+        topk_boxes = torch.div(topk_indexes, prob.shape[2], rounding_mode="trunc")
+        labels = topk_indexes % prob.shape[2]
+        if not_to_xyxy:
+            boxes = out_bbox
+        else:
+            boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+
+        # if test:
+        #     assert not not_to_xyxy
+        #     boxes[:,:,2:] = boxes[:,:,2:] - boxes[:,:,:2]
+        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
+
+        # 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, :]
+
+        if self.nms_iou_threshold > 0:
+            item_indices = [
+                nms(b, s, iou_threshold=self.nms_iou_threshold)
+                for b, s in zip(boxes, scores)
+            ]
+
+            results = [
+                {"scores": s[i], "labels": l[i], "boxes": b[i]}
+                for s, l, b, i in zip(scores, labels, boxes, item_indices)
+            ]
+        else:
+            results = [
+                {"scores": s, "labels": l, "boxes": b}
+                for s, l, b in zip(scores, labels, boxes)
+            ]
+        results = [
+            {"scores": s, "labels": l, "boxes": b}
+            for s, l, b in zip(scores, labels, boxes)
+        ]
+        return results
+
+
+@MODULE_BUILD_FUNCS.registe_with_name(module_name="groundingdino_app")
+def build_groundingdino(args):
+    device = torch.device(args.device)
+    backbone = build_backbone(args)
+    transformer = build_transformer(args)
+
+    dn_labelbook_size = args.dn_labelbook_size
+    dec_pred_bbox_embed_share = args.dec_pred_bbox_embed_share
+    sub_sentence_present = args.sub_sentence_present
+
+    model = GroundingDINO(
+        backbone,
+        transformer,
+        num_queries=args.num_queries,
+        aux_loss=args.aux_loss,
+        iter_update=True,
+        query_dim=4,
+        num_feature_levels=args.num_feature_levels,
+        nheads=args.nheads,
+        dec_pred_bbox_embed_share=dec_pred_bbox_embed_share,
+        two_stage_type=args.two_stage_type,
+        two_stage_bbox_embed_share=args.two_stage_bbox_embed_share,
+        two_stage_class_embed_share=args.two_stage_class_embed_share,
+        num_patterns=args.num_patterns,
+        dn_number=0,
+        dn_box_noise_scale=args.dn_box_noise_scale,
+        dn_label_noise_ratio=args.dn_label_noise_ratio,
+        dn_labelbook_size=dn_labelbook_size,
+        text_encoder_type=args.text_encoder_type,
+        sub_sentence_present=sub_sentence_present,
+        max_text_len=args.max_text_len,
+    )
+
+    matcher = build_matcher(args)
+
+    # prepare weight dict
+    weight_dict = {"loss_ce": args.cls_loss_coef, "loss_bbox": args.bbox_loss_coef}
+    weight_dict["loss_giou"] = args.giou_loss_coef
+    clean_weight_dict_wo_dn = copy.deepcopy(weight_dict)
+
+    clean_weight_dict = copy.deepcopy(weight_dict)
+
+    # 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 clean_weight_dict.items()}
+            )
+        weight_dict.update(aux_weight_dict)
+
+    if args.two_stage_type != "no":
+        interm_weight_dict = {}
+        try:
+            no_interm_box_loss = args.no_interm_box_loss
+        except:
+            no_interm_box_loss = False
+        _coeff_weight_dict = {
+            "loss_ce": 1.0,
+            "loss_bbox": 1.0 if not no_interm_box_loss else 0.0,
+            "loss_giou": 1.0 if not no_interm_box_loss else 0.0,
+        }
+        try:
+            interm_loss_coef = args.interm_loss_coef
+        except:
+            interm_loss_coef = 1.0
+        interm_weight_dict.update(
+            {
+                k + f"_interm": v * interm_loss_coef * _coeff_weight_dict[k]
+                for k, v in clean_weight_dict_wo_dn.items()
+            }
+        )
+        weight_dict.update(interm_weight_dict)
+
+    # losses = ['labels', 'boxes', 'cardinality']
+    losses = ["labels", "boxes"]
+
+    criterion = SetCriterion(
+        matcher=matcher,
+        weight_dict=weight_dict,
+        focal_alpha=args.focal_alpha,
+        focal_gamma=args.focal_gamma,
+        losses=losses,
+    )
+    criterion.to(device)
+    postprocessors = {
+        "bbox": PostProcess(
+            num_select=args.num_select,
+            text_encoder_type=args.text_encoder_type,
+            nms_iou_threshold=args.nms_iou_threshold,
+            args=args,
+        )
+    }
+
+    return model, criterion, postprocessors
+
+
+def create_positive_map(tokenized, tokens_positive, cat_list, caption):
+    """construct a map such that positive_map[i,j] = True iff box i is associated to token j"""
+    positive_map = torch.zeros((len(tokens_positive), 256), dtype=torch.float)
+
+    for j, label in enumerate(tokens_positive):
+        start_ind = caption.find(cat_list[label])
+        end_ind = start_ind + len(cat_list[label]) - 1
+        beg_pos = tokenized.char_to_token(start_ind)
+        try:
+            end_pos = tokenized.char_to_token(end_ind)
+        except:
+            end_pos = None
+        if end_pos is None:
+            try:
+                end_pos = tokenized.char_to_token(end_ind - 1)
+                if end_pos is None:
+                    end_pos = tokenized.char_to_token(end_ind - 2)
+            except:
+                end_pos = None
+        # except Exception as e:
+        #     print("beg:", beg, "end:", end)
+        #     print("token_positive:", tokens_positive)
+        #     # print("beg_pos:", beg_pos, "end_pos:", end_pos)
+        #     raise e
+        # if beg_pos is None:
+        #     try:
+        #         beg_pos = tokenized.char_to_token(beg + 1)
+        #         if beg_pos is None:
+        #             beg_pos = tokenized.char_to_token(beg + 2)
+        #     except:
+        #         beg_pos = None
+        # if end_pos is None:
+        #     try:
+        #         end_pos = tokenized.char_to_token(end - 2)
+        #         if end_pos is None:
+        #             end_pos = tokenized.char_to_token(end - 3)
+        #     except:
+        #         end_pos = None
+        if beg_pos is None or end_pos is None:
+            continue
+        if beg_pos < 0 or end_pos < 0:
+            continue
+        if beg_pos > end_pos:
+            continue
+        # assert beg_pos is not None and end_pos is not None
+        positive_map[j, beg_pos : end_pos + 1].fill_(1)
+    return positive_map
+
+
+def create_positive_map_exemplar(input_ids, label, special_tokens):
+    tokens_positive = torch.zeros(256, dtype=torch.float)
+    count = -1
+    for token_ind in range(len(input_ids)):
+        input_id = input_ids[token_ind]
+        if (input_id not in special_tokens) and (
+            token_ind == 0 or (input_ids[token_ind - 1] in special_tokens)
+        ):
+            count += 1
+        if count == label:
+            ind_to_insert_ones = token_ind
+
+            while input_ids[ind_to_insert_ones] not in special_tokens:
+                tokens_positive[ind_to_insert_ones] = 1
+                ind_to_insert_ones += 1
+            break
+    return tokens_positive
+

models/GroundingDINO/matcher.py

@@ -0,0 +1,217 @@
+# ------------------------------------------------------------------------
+# DINO
+# Copyright (c) 2022 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modules to compute the matching cost and solve the corresponding LSAP.
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# ------------------------------------------------------------------------
+
+
+import torch, os
+from torch import nn
+from scipy.optimize import linear_sum_assignment
+
+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, focal_alpha = 0.25):
+        """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"
+
+        self.focal_alpha = focal_alpha
+
+    @torch.no_grad()
+    def forward(self, outputs, targets, label_map):
+        """ 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).sigmoid()  # [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.
+        alpha = self.focal_alpha
+        gamma = 2.0
+
+        new_label_map=label_map[tgt_ids.cpu()]
+
+        neg_cost_class = (1 - alpha) * (out_prob ** gamma) * (-(1 - out_prob + 1e-8).log())
+        pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
+        new_label_map=new_label_map.to(pos_cost_class.device) 
+        cost_bbox = torch.cdist(out_bbox[:, :2], tgt_bbox[:, :2], p=1)
+
+        # cost_class=(pos_cost_class @ new_label_map.T - neg_cost_class@ new_label_map.T)
+        cost_class=[]
+        for idx_map in new_label_map:       
+            idx_map = idx_map / idx_map.sum()
+            cost_class.append(pos_cost_class @ idx_map - neg_cost_class@ idx_map)
+        if cost_class:
+            cost_class=torch.stack(cost_class,dim=0).T
+        else:
+            cost_class=torch.zeros_like(cost_bbox)
+        # Compute the L1 cost between boxes
+        
+
+        # Compute the giou cost betwen boxes
+        cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))
+        # import pdb;pdb.set_trace()
+        # 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()
+        C[torch.isnan(C)] = 0.0
+        C[torch.isinf(C)] = 0.0
+
+        sizes = [len(v["boxes"]) for v in targets]
+        try:
+            indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
+        except:
+            print("warning: use SimpleMinsumMatcher")
+            indices = []
+            device = C.device
+            for i, (c, _size) in enumerate(zip(C.split(sizes, -1), sizes)):
+                weight_mat = c[i]
+                idx_i = weight_mat.min(0)[1]
+                idx_j = torch.arange(_size).to(device)
+                indices.append((idx_i, idx_j))
+        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
+
+
+class SimpleMinsumMatcher(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, focal_alpha = 0.25):
+        """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"
+
+        self.focal_alpha = focal_alpha
+
+    @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).sigmoid()  # [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.
+        alpha = self.focal_alpha
+        gamma = 2.0
+        neg_cost_class = (1 - alpha) * (out_prob ** gamma) * (-(1 - out_prob + 1e-8).log())
+        pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
+        cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, 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)
+
+        sizes = [len(v["boxes"]) for v in targets]
+        indices = []
+        device = C.device
+        for i, (c, _size) in enumerate(zip(C.split(sizes, -1), sizes)):
+            weight_mat = c[i]
+            idx_i = weight_mat.min(0)[1]
+            idx_j = torch.arange(_size).to(device)
+            indices.append((idx_i, idx_j))
+
+        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
+
+
+def build_matcher(args):
+    assert args.matcher_type in ['HungarianMatcher', 'SimpleMinsumMatcher'], "Unknown args.matcher_type: {}".format(args.matcher_type)
+    if args.matcher_type == 'HungarianMatcher':
+        return HungarianMatcher(
+            cost_class=args.set_cost_class, cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou,
+            focal_alpha=args.focal_alpha
+        )
+    elif args.matcher_type == 'SimpleMinsumMatcher':
+        return SimpleMinsumMatcher(
+            cost_class=args.set_cost_class, cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou,
+            focal_alpha=args.focal_alpha
+        )    
+    else:
+        raise NotImplementedError("Unknown args.matcher_type: {}".format(args.matcher_type))

models/GroundingDINO/mlp_loca.py

@@ -0,0 +1,23 @@
+from torch import nn
+
+
+class MLP(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        dropout: float,
+        activation: nn.Module
+    ):
+        super(MLP, self).__init__()
+
+        self.linear1 = nn.Linear(input_dim, hidden_dim)
+        self.linear2 = nn.Linear(hidden_dim, input_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = activation()
+
+    def forward(self, x):
+        return (
+            self.linear2(self.dropout(self.activation(self.linear1(x))))
+        )