diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..75ed18a9cb91092f5977593d7bfa226ac2507e9e 100644 --- a/.gitattributes +++ b/.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 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..8d2e7d29156b45714ef51d5580b20909d0b24354 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6dd57dd569973900f06365b8733898da7499926d --- /dev/null +++ b/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} +} +``` diff --git a/cfg_app.py b/cfg_app.py new file mode 100644 index 0000000000000000000000000000000000000000..d9ce134c0a2c3f9d864b5c39f9407b4068e4a685 --- /dev/null +++ b/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"] diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..b8ddc8a4c6e74df4a01269ecae2560d01fec599a --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/countgdpp.py b/countgdpp.py new file mode 100644 index 0000000000000000000000000000000000000000..2ebf0407ab0e492c1048376b0ca020c4d2293911 --- /dev/null +++ b/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()} diff --git a/datasets/__init__.py b/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b1864e13e91e60944a78de5b7903228cce165cdb --- /dev/null +++ b/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') diff --git a/datasets/coco.py b/datasets/coco.py new file mode 100644 index 0000000000000000000000000000000000000000..86debba3b18b49176686cce45bdd8bb8f5f2129e --- /dev/null +++ b/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)) diff --git a/datasets/coco_eval.py b/datasets/coco_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..74bbc0b32b5e4dbe40cddae7ff9add26e3f910f9 --- /dev/null +++ b/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 +################################################################# diff --git a/datasets/coco_panoptic.py b/datasets/coco_panoptic.py new file mode 100644 index 0000000000000000000000000000000000000000..b24f615c2faa14b422829e2edad996e2b5b84248 --- /dev/null +++ b/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 diff --git a/datasets/cocogrounding_eval.py b/datasets/cocogrounding_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..32f86a50a8f1daf36135edefd20574e022fe5ff6 --- /dev/null +++ b/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 +################################################################# diff --git a/datasets/data_util.py b/datasets/data_util.py new file mode 100644 index 0000000000000000000000000000000000000000..8019ccd7348f87640620c52b7a3f8a1c266859a7 --- /dev/null +++ b/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 + + + \ No newline at end of file diff --git a/datasets/dataset.py b/datasets/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..5655c00ccdbf0b09d51da8dfa3593a2bf9aee527 --- /dev/null +++ b/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() diff --git a/datasets/odvg.py b/datasets/odvg.py new file mode 100644 index 0000000000000000000000000000000000000000..c9d977c6eab2e076571a79278640ed7fcbc5ebaf --- /dev/null +++ b/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) diff --git a/datasets/panoptic_eval.py b/datasets/panoptic_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..9cb4f83409046a5c2a87643ee005e52a440aae74 --- /dev/null +++ b/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 diff --git a/datasets/random_crop.py b/datasets/random_crop.py new file mode 100644 index 0000000000000000000000000000000000000000..a23d936048e073423dd176f76607e9fc2b35bb48 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/datasets/sltransform.py b/datasets/sltransform.py new file mode 100644 index 0000000000000000000000000000000000000000..7d9d10bc20e159f59e46912057714c527a08d6cd --- /dev/null +++ b/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 \ No newline at end of file diff --git a/datasets/transforms.py b/datasets/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..ee597af70f26380db6de96e03389f7824e3a426c --- /dev/null +++ b/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 diff --git a/datasets/transforms_app.py b/datasets/transforms_app.py new file mode 100644 index 0000000000000000000000000000000000000000..e4ac7ac8afde6e7a3d3d9746dbd87d2bde2e6879 --- /dev/null +++ b/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 diff --git a/datasets/transforms_orig.py b/datasets/transforms_orig.py new file mode 100644 index 0000000000000000000000000000000000000000..42256862608332ad3caa288f5f7475090b166c11 --- /dev/null +++ b/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 diff --git a/example.py b/example.py new file mode 100644 index 0000000000000000000000000000000000000000..1ce9d12f308473bff9ccfa1e74e196333480f863 --- /dev/null +++ b/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") diff --git a/groundingdino/util/__init__.py b/groundingdino/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..168f9979a4623806934b0ff1102ac166704e7dec --- /dev/null +++ b/groundingdino/util/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved diff --git a/groundingdino/util/box_ops.py b/groundingdino/util/box_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..781068d294e576954edb4bd07b6e0f30e4e1bcd9 --- /dev/null +++ b/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() diff --git a/groundingdino/util/get_tokenlizer.py b/groundingdino/util/get_tokenlizer.py new file mode 100644 index 0000000000000000000000000000000000000000..77873bcf4e098c194ef6182b8cbfdc2aa4e3d0ed --- /dev/null +++ b/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)) diff --git a/groundingdino/util/inference.py b/groundingdino/util/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..d6e81d89db1c422bbf27e3c160ef0957dfa57223 --- /dev/null +++ b/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) diff --git a/groundingdino/util/logger.py b/groundingdino/util/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..18145f54c927abd59b95f3fa6e6da8002bc2ce97 --- /dev/null +++ b/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") diff --git a/groundingdino/util/misc.py b/groundingdino/util/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..d64b84ef24bea0c98e76824feb1903f6bfebe7a5 --- /dev/null +++ b/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 diff --git a/groundingdino/util/slconfig.py b/groundingdino/util/slconfig.py new file mode 100644 index 0000000000000000000000000000000000000000..672e72ed0b68a54c13ade66c9f146d2d542e97c6 --- /dev/null +++ b/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) diff --git a/groundingdino/util/slio.py b/groundingdino/util/slio.py new file mode 100644 index 0000000000000000000000000000000000000000..72c1f0f7b82cdc931d381feef64fe15815ba657e --- /dev/null +++ b/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') diff --git a/groundingdino/util/time_counter.py b/groundingdino/util/time_counter.py new file mode 100644 index 0000000000000000000000000000000000000000..0aedb2e4d61bfbe7571dca9d50053f0fedaa1359 --- /dev/null +++ b/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__) diff --git a/groundingdino/util/utils.py b/groundingdino/util/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8cf83ae03a7865bc48493be16e8b1b2d53a1b09f --- /dev/null +++ b/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) == '' + >>> assert repr(foo).startswith('>> 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) == '' + """ + + 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") diff --git a/groundingdino/util/visualizer.py b/groundingdino/util/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7a1b7b101e9b73f75f9136bc67f2063c7c1cf1c1 --- /dev/null +++ b/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"]) diff --git a/groundingdino/util/vl_utils.py b/groundingdino/util/vl_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c91bb02f584398f08a28e6b7719e2b99f6e28616 --- /dev/null +++ b/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 diff --git a/img/336.jpg b/img/336.jpg new file mode 100644 index 0000000000000000000000000000000000000000..17bbe4a39002d0305c43bc4135e907d9db1312e8 Binary files /dev/null and b/img/336.jpg differ diff --git a/img/inference-architecture.jpg b/img/inference-architecture.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8808f3d8689294941464a27298e2b26f2cb92990 --- /dev/null +++ b/img/inference-architecture.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d154e0e1ba2601ee5a426424f1e1f0efb83a84564d1564960e6e13aa63d3c7ea +size 107655 diff --git a/img/synthetic_blueberry.jpg b/img/synthetic_blueberry.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e0233bc0c7a31873a988cf73d06dfae9f19ced3c --- /dev/null +++ b/img/synthetic_blueberry.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f911a4f249c4de149a883b65ff77f549d5c848aebf40f1a246c28983983fa140 +size 1466911 diff --git a/img/synthetic_strawberry.jpg b/img/synthetic_strawberry.jpg new file mode 100644 index 0000000000000000000000000000000000000000..892e0305977385bf54da83bd315780fc21248307 --- /dev/null +++ b/img/synthetic_strawberry.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:861339ad69483128a7c50cc356f10d708e50e64d50beb11ae07438f11cc68220 +size 1377189 diff --git a/img/teaser.jpg b/img/teaser.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ed44234ef25a29aea9318a8037acb06b90ad4372 --- /dev/null +++ b/img/teaser.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bfe4ec0216403aa012c534a562dfd03a352b6cf51fff9f0fb531bda8a7753002 +size 424547 diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1e045381a5916f902e604c01a9ef7078ae5d2a65 --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:94eabe4e2595dac1e04cc1703641458b13c1c624ce10072ba57dd96b625a41f5 +size 937560192 diff --git a/models/GroundingDINO/__init__.py b/models/GroundingDINO/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2af819d61d589cfec2e0ca46612a7456f42b831a --- /dev/null +++ b/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 diff --git a/models/GroundingDINO/backbone/__init__.py b/models/GroundingDINO/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..76e4b272b479a26c63d120c818c140870cd8c287 --- /dev/null +++ b/models/GroundingDINO/backbone/__init__.py @@ -0,0 +1 @@ +from .backbone import build_backbone diff --git a/models/GroundingDINO/backbone/backbone.py b/models/GroundingDINO/backbone/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..c8340c723fad8e07e2fc62daaa3912487498814b --- /dev/null +++ b/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 diff --git a/models/GroundingDINO/backbone/position_encoding.py b/models/GroundingDINO/backbone/position_encoding.py new file mode 100644 index 0000000000000000000000000000000000000000..eac7e896bbe85a670824bfe8ef487d0535d5bd99 --- /dev/null +++ b/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 diff --git a/models/GroundingDINO/backbone/swin_transformer.py b/models/GroundingDINO/backbone/swin_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..affebbb7356ae2c01e23fa852d7105aab5b3a8f4 --- /dev/null +++ b/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) diff --git a/models/GroundingDINO/bertwarper.py b/models/GroundingDINO/bertwarper.py new file mode 100644 index 0000000000000000000000000000000000000000..f0cf9779b270e1aead32845006f8b881fcba37ad --- /dev/null +++ b/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 diff --git a/models/GroundingDINO/fuse_modules.py b/models/GroundingDINO/fuse_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..2753b3ddee43c7a9fe28d1824db5d786e7e1ad59 --- /dev/null +++ b/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) diff --git a/models/GroundingDINO/groundingdino.py b/models/GroundingDINO/groundingdino.py new file mode 100644 index 0000000000000000000000000000000000000000..52c466a15663bf4bebe191c08af2dc4fbbc6fcb6 --- /dev/null +++ b/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'] + # + + # 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 + + + diff --git a/models/GroundingDINO/groundingdino_app.py b/models/GroundingDINO/groundingdino_app.py new file mode 100644 index 0000000000000000000000000000000000000000..710c95f0e4d6cbebbd5dcc29c1c72ecec3369586 --- /dev/null +++ b/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'] + # + + # 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 + diff --git a/models/GroundingDINO/matcher.py b/models/GroundingDINO/matcher.py new file mode 100644 index 0000000000000000000000000000000000000000..daf43d8f1b92b2821773013dd0a36b8a969833f1 --- /dev/null +++ b/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)) diff --git a/models/GroundingDINO/mlp_loca.py b/models/GroundingDINO/mlp_loca.py new file mode 100644 index 0000000000000000000000000000000000000000..7d39aedb5072ed9ca7d55774304e7f40aea9a60d --- /dev/null +++ b/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)))) + ) diff --git a/models/GroundingDINO/ms_deform_attn.py b/models/GroundingDINO/ms_deform_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..e7c7610ba2f234ef54e79a35977db2fe81350c51 --- /dev/null +++ b/models/GroundingDINO/ms_deform_attn.py @@ -0,0 +1,415 @@ +# ------------------------------------------------------------------------ +# 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] +# ------------------------------------------------------------------------ +# Deformable DETR +# Copyright (c) 2020 SenseTime. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------------------------------ +# Modified from: +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/functions/ms_deform_attn_func.py +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/modules/ms_deform_attn.py +# https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/multi_scale_deform_attn.py +# ------------------------------------------------------------------------------------------------ + +import math +import warnings +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Function +from torch.autograd.function import once_differentiable +from torch.nn.init import constant_, xavier_uniform_ + +try: + # from groundingdino import _C + import MultiScaleDeformableAttention as _C +except Exception: + _C = None + warnings.warn("Custom CUDA ops not found; using the pure-PyTorch deformable attention (works on CPU and GPU).") + + +# helpers +def _is_power_of_2(n): + if (not isinstance(n, int)) or (n < 0): + raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n))) + return (n & (n - 1) == 0) and n != 0 + + +class MultiScaleDeformableAttnFunction(Function): + @staticmethod + def forward( + ctx, + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + im2col_step, + ): + ctx.im2col_step = im2col_step + output = _C.ms_deform_attn_forward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ctx.im2col_step, + ) + ctx.save_for_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ) + return output + + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + ( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ) = ctx.saved_tensors + grad_value, grad_sampling_loc, grad_attn_weight = _C.ms_deform_attn_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + grad_output, + ctx.im2col_step, + ) + + return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None + + +def multi_scale_deformable_attn_pytorch( + value: torch.Tensor, + value_spatial_shapes: torch.Tensor, + sampling_locations: torch.Tensor, + attention_weights: torch.Tensor, +) -> torch.Tensor: + + bs, _, num_heads, embed_dims = value.shape + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level, (H_, W_) in enumerate(value_spatial_shapes): + # bs, H_*W_, num_heads, embed_dims -> + # bs, H_*W_, num_heads*embed_dims -> + # bs, num_heads*embed_dims, H_*W_ -> + # bs*num_heads, embed_dims, H_, W_ + value_l_ = ( + value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_) + ) + # bs, num_queries, num_heads, num_points, 2 -> + # bs, num_heads, num_queries, num_points, 2 -> + # bs*num_heads, num_queries, num_points, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1) + # bs*num_heads, embed_dims, num_queries, num_points + sampling_value_l_ = F.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + # (bs, num_queries, num_heads, num_levels, num_points) -> + # (bs, num_heads, num_queries, num_levels, num_points) -> + # (bs, num_heads, 1, num_queries, num_levels*num_points) + attention_weights = attention_weights.transpose(1, 2).reshape( + bs * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(bs, num_heads * embed_dims, num_queries) + ) + return output.transpose(1, 2).contiguous() + + +class MultiScaleDeformableAttention(nn.Module): + """Multi-Scale Deformable Attention Module used in Deformable-DETR + + `Deformable DETR: Deformable Transformers for End-to-End Object Detection. + `_. + + Args: + embed_dim (int): The embedding dimension of Attention. Default: 256. + num_heads (int): The number of attention heads. Default: 8. + num_levels (int): The number of feature map used in Attention. Default: 4. + num_points (int): The number of sampling points for each query + in each head. Default: 4. + img2col_steps (int): The step used in image_to_column. Defualt: 64. + dropout (float): Dropout layer used in output. Default: 0.1. + batch_first (bool): if ``True``, then the input and output tensor will be + provided as `(bs, n, embed_dim)`. Default: False. `(n, bs, embed_dim)` + """ + + def __init__( + self, + embed_dim: int = 256, + num_heads: int = 8, + num_levels: int = 4, + num_points: int = 4, + img2col_step: int = 64, + batch_first: bool = False, + ): + super().__init__() + if embed_dim % num_heads != 0: + raise ValueError( + "embed_dim must be divisible by num_heads, but got {} and {}".format( + embed_dim, num_heads + ) + ) + head_dim = embed_dim // num_heads + + self.batch_first = batch_first + + if not _is_power_of_2(head_dim): + warnings.warn( + """ + You'd better set d_model in MSDeformAttn to make sure that + each dim of the attention head a power of 2, which is more efficient. + """ + ) + + self.im2col_step = img2col_step + self.embed_dim = embed_dim + self.num_heads = num_heads + self.num_levels = num_levels + self.num_points = num_points + self.sampling_offsets = nn.Linear(embed_dim, num_heads * num_levels * num_points * 2) + self.attention_weights = nn.Linear(embed_dim, num_heads * num_levels * num_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + + self.init_weights() + + def _reset_parameters(self): + return self.init_weights() + + def init_weights(self): + """ + Default initialization for Parameters of Module. + """ + constant_(self.sampling_offsets.weight.data, 0.0) + thetas = torch.arange(self.num_heads, dtype=torch.float32) * ( + 2.0 * math.pi / self.num_heads + ) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(self.num_heads, 1, 1, 2) + .repeat(1, self.num_levels, self.num_points, 1) + ) + for i in range(self.num_points): + grid_init[:, :, i, :] *= i + 1 + with torch.no_grad(): + self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1)) + constant_(self.attention_weights.weight.data, 0.0) + constant_(self.attention_weights.bias.data, 0.0) + xavier_uniform_(self.value_proj.weight.data) + constant_(self.value_proj.bias.data, 0.0) + xavier_uniform_(self.output_proj.weight.data) + constant_(self.output_proj.bias.data, 0.0) + + def freeze_sampling_offsets(self): + print("Freeze sampling offsets") + self.sampling_offsets.weight.requires_grad = False + self.sampling_offsets.bias.requires_grad = False + + def freeze_attention_weights(self): + print("Freeze attention weights") + self.attention_weights.weight.requires_grad = False + self.attention_weights.bias.requires_grad = False + + def forward( + self, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + value: Optional[torch.Tensor] = None, + query_pos: Optional[torch.Tensor] = None, + key_padding_mask: Optional[torch.Tensor] = None, + reference_points: Optional[torch.Tensor] = None, + spatial_shapes: Optional[torch.Tensor] = None, + level_start_index: Optional[torch.Tensor] = None, + **kwargs + ) -> torch.Tensor: + + """Forward Function of MultiScaleDeformableAttention + + Args: + query (torch.Tensor): Query embeddings with shape + `(num_query, bs, embed_dim)` + key (torch.Tensor): Key embeddings with shape + `(num_key, bs, embed_dim)` + value (torch.Tensor): Value embeddings with shape + `(num_key, bs, embed_dim)` + query_pos (torch.Tensor): The position embedding for `query`. Default: None. + key_padding_mask (torch.Tensor): ByteTensor for `query`, with shape `(bs, num_key)`, + indicating which elements within `key` to be ignored in attention. + reference_points (torch.Tensor): The normalized reference points + with shape `(bs, num_query, num_levels, 2)`, + all elements is range in [0, 1], top-left (0, 0), + bottom-right (1, 1), including padding are. + or `(N, Length_{query}, num_levels, 4)`, add additional + two dimensions `(h, w)` to form reference boxes. + spatial_shapes (torch.Tensor): Spatial shape of features in different levels. + With shape `(num_levels, 2)`, last dimension represents `(h, w)`. + level_start_index (torch.Tensor): The start index of each level. A tensor with + shape `(num_levels, )` which can be represented as + `[0, h_0 * w_0, h_0 * w_0 + h_1 * w_1, ...]`. + + Returns: + torch.Tensor: forward results with shape `(num_query, bs, embed_dim)` + """ + + if value is None: + value = query + + if query_pos is not None: + query = query + query_pos + + if not self.batch_first: + # change to (bs, num_query ,embed_dims) + query = query.permute(1, 0, 2) + value = value.permute(1, 0, 2) + + bs, num_query, _ = query.shape + bs, num_value, _ = value.shape + + assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value + + value = self.value_proj(value) + if key_padding_mask is not None: + value = value.masked_fill(key_padding_mask[..., None], float(0)) + value = value.view(bs, num_value, self.num_heads, -1) + sampling_offsets = self.sampling_offsets(query).view( + bs, num_query, self.num_heads, self.num_levels, self.num_points, 2 + ) + attention_weights = self.attention_weights(query).view( + bs, num_query, self.num_heads, self.num_levels * self.num_points + ) + attention_weights = attention_weights.softmax(-1) + attention_weights = attention_weights.view( + bs, + num_query, + self.num_heads, + self.num_levels, + self.num_points, + ) + + # bs, num_query, num_heads, num_levels, num_points, 2 + if reference_points.shape[-1] == 2: + offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif reference_points.shape[-1] == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets + / self.num_points + * reference_points[:, :, None, :, None, 2:] + * 0.5 + ) + else: + raise ValueError( + "Last dim of reference_points must be 2 or 4, but get {} instead.".format( + reference_points.shape[-1] + ) + ) + + if _C is not None and torch.cuda.is_available() and value.is_cuda: + halffloat = False + if value.dtype == torch.float16: + halffloat = True + value = value.float() + sampling_locations = sampling_locations.float() + attention_weights = attention_weights.float() + + output = MultiScaleDeformableAttnFunction.apply( + value, + spatial_shapes, + level_start_index, + sampling_locations, + attention_weights, + self.im2col_step, + ) + + if halffloat: + output = output.half() + else: + output = multi_scale_deformable_attn_pytorch( + value, spatial_shapes, sampling_locations, attention_weights + ) + + output = self.output_proj(output) + + if not self.batch_first: + output = output.permute(1, 0, 2) + + return output + + +def create_dummy_class(klass, dependency, message=""): + """ + When a dependency of a class is not available, create a dummy class which throws ImportError + when used. + + Args: + klass (str): name of the class. + dependency (str): name of the dependency. + message: extra message to print + Returns: + class: a class object + """ + err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, klass) + if message: + err = err + " " + message + + class _DummyMetaClass(type): + # throw error on class attribute access + def __getattr__(_, __): # noqa: B902 + raise ImportError(err) + + class _Dummy(object, metaclass=_DummyMetaClass): + # throw error on constructor + def __init__(self, *args, **kwargs): + raise ImportError(err) + + return _Dummy + + +def create_dummy_func(func, dependency, message=""): + """ + When a dependency of a function is not available, create a dummy function which throws + ImportError when used. + + Args: + func (str): name of the function. + dependency (str or list[str]): name(s) of the dependency. + message: extra message to print + Returns: + function: a function object + """ + err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, func) + if message: + err = err + " " + message + + if isinstance(dependency, (list, tuple)): + dependency = ",".join(dependency) + + def _dummy(*args, **kwargs): + raise ImportError(err) + + return _dummy diff --git a/models/GroundingDINO/positional_encoding_loca.py b/models/GroundingDINO/positional_encoding_loca.py new file mode 100644 index 0000000000000000000000000000000000000000..9ca20f8e000c1f134b6d5108fe5ab626292a5519 --- /dev/null +++ b/models/GroundingDINO/positional_encoding_loca.py @@ -0,0 +1,30 @@ +import torch +from torch import nn + + +class PositionalEncodingsFixed(nn.Module): + + def __init__(self, emb_dim, temperature=10000): + + super(PositionalEncodingsFixed, self).__init__() + + self.emb_dim = emb_dim + self.temperature = temperature + + def _1d_pos_enc(self, mask, dim): + temp = torch.arange(self.emb_dim // 2).float().to(mask.device) + temp = self.temperature ** (2 * (temp.div(2, rounding_mode='floor')) / self.emb_dim) + + enc = (~mask).cumsum(dim).float().unsqueeze(-1) / temp + enc = torch.stack([ + enc[..., 0::2].sin(), enc[..., 1::2].cos() + ], dim=-1).flatten(-2) + + return enc + + def forward(self, bs, h, w, device): + mask = torch.zeros(bs, h, w, dtype=torch.bool, requires_grad=False, device=device) + x = self._1d_pos_enc(mask, dim=2) + y = self._1d_pos_enc(mask, dim=1) + + return torch.cat([y, x], dim=3).permute(0, 3, 1, 2) diff --git a/models/GroundingDINO/transformer.py b/models/GroundingDINO/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..0fb6a7df3e33579e7126df1f6553f18e671f78b9 --- /dev/null +++ b/models/GroundingDINO/transformer.py @@ -0,0 +1,970 @@ +# ------------------------------------------------------------------------ +# 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 Transformer class. +# 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. +# ------------------------------------------------------------------------ + +from typing import Optional + +import torch +import torch.utils.checkpoint as checkpoint +from torch import Tensor, nn + +from groundingdino.util.misc import inverse_sigmoid + +from .fuse_modules import BiAttentionBlock +from .ms_deform_attn import MultiScaleDeformableAttention as MSDeformAttn +from .transformer_vanilla import TransformerEncoderLayer +from .utils import ( + MLP, + _get_activation_fn, + _get_clones, + gen_encoder_output_proposals, + gen_sineembed_for_position, + get_sine_pos_embed, +) + + +class Transformer(nn.Module): + def __init__( + self, + d_model=256, + nhead=8, + num_queries=300, + num_encoder_layers=6, + num_unicoder_layers=0, + num_decoder_layers=6, + dim_feedforward=2048, + dropout=0.0, + activation="relu", + normalize_before=False, + return_intermediate_dec=False, + query_dim=4, + num_patterns=0, + # for deformable encoder + num_feature_levels=1, + enc_n_points=4, + dec_n_points=4, + # init query + learnable_tgt_init=False, + # two stage + two_stage_type="no", # ['no', 'standard', 'early', 'combine', 'enceachlayer', 'enclayer1'] + embed_init_tgt=False, + # for text + use_text_enhancer=False, + use_fusion_layer=False, + use_checkpoint=False, + use_transformer_ckpt=False, + use_text_cross_attention=False, + text_dropout=0.1, + fusion_dropout=0.1, + fusion_droppath=0.0, + ): + super().__init__() + self.num_feature_levels = num_feature_levels + self.num_encoder_layers = num_encoder_layers + self.num_unicoder_layers = num_unicoder_layers + self.num_decoder_layers = num_decoder_layers + self.num_queries = num_queries + assert query_dim == 4 + + # choose encoder layer type + encoder_layer = DeformableTransformerEncoderLayer( + d_model, dim_feedforward, dropout, activation, num_feature_levels, nhead, enc_n_points + ) + + if use_text_enhancer: + text_enhance_layer = TransformerEncoderLayer( + d_model=d_model, + nhead=nhead // 2, + dim_feedforward=dim_feedforward // 2, + dropout=text_dropout, + ) + else: + text_enhance_layer = None + + if use_fusion_layer: + feature_fusion_layer = BiAttentionBlock( + v_dim=d_model, + l_dim=d_model, + embed_dim=dim_feedforward // 2, + num_heads=nhead // 2, + dropout=fusion_dropout, + drop_path=fusion_droppath, + ) + else: + feature_fusion_layer = None + + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + assert encoder_norm is None + self.encoder = TransformerEncoder( + encoder_layer, + num_encoder_layers, + d_model=d_model, + num_queries=num_queries, + text_enhance_layer=text_enhance_layer, + feature_fusion_layer=feature_fusion_layer, + use_checkpoint=use_checkpoint, + use_transformer_ckpt=use_transformer_ckpt, + ) + + # choose decoder layer type + decoder_layer = DeformableTransformerDecoderLayer( + d_model, + dim_feedforward, + dropout, + activation, + num_feature_levels, + nhead, + dec_n_points, + use_text_cross_attention=use_text_cross_attention, + ) + + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder( + decoder_layer, + num_decoder_layers, + decoder_norm, + return_intermediate=return_intermediate_dec, + d_model=d_model, + query_dim=query_dim, + num_feature_levels=num_feature_levels, + ) + + self.d_model = d_model + self.nhead = nhead + self.dec_layers = num_decoder_layers + self.num_queries = num_queries # useful for single stage model only + self.num_patterns = num_patterns + if not isinstance(num_patterns, int): + Warning("num_patterns should be int but {}".format(type(num_patterns))) + self.num_patterns = 0 + + if num_feature_levels > 1: + if self.num_encoder_layers > 0: + self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model)) + else: + self.level_embed = None + + self.learnable_tgt_init = learnable_tgt_init + assert learnable_tgt_init, "why not learnable_tgt_init" + self.embed_init_tgt = embed_init_tgt + if (two_stage_type != "no" and embed_init_tgt) or (two_stage_type == "no"): + self.tgt_embed = nn.Embedding(self.num_queries, d_model) + nn.init.normal_(self.tgt_embed.weight.data) + else: + self.tgt_embed = None + + # for 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 == "standard": + # anchor selection at the output of encoder + self.enc_output = nn.Linear(d_model, d_model) + self.enc_output_norm = nn.LayerNorm(d_model) + self.two_stage_wh_embedding = None + + if two_stage_type == "no": + self.init_ref_points(num_queries) # init self.refpoint_embed + + self.enc_out_class_embed = None + self.enc_out_bbox_embed = None + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + for m in self.modules(): + if isinstance(m, MSDeformAttn): + m._reset_parameters() + if self.num_feature_levels > 1 and self.level_embed is not None: + nn.init.normal_(self.level_embed) + + def get_valid_ratio(self, mask): + _, H, W = mask.shape + valid_H = torch.sum(~mask[:, :, 0], 1) + valid_W = torch.sum(~mask[:, 0, :], 1) + valid_ratio_h = valid_H.float() / H + valid_ratio_w = valid_W.float() / W + valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1) + return valid_ratio + + def init_ref_points(self, use_num_queries): + self.refpoint_embed = nn.Embedding(use_num_queries, 4) + + def forward(self, srcs, masks, refpoint_embed, pos_embeds, tgt, attn_mask=None, text_dict=None): + """ + Input: + - srcs: List of multi features [bs, ci, hi, wi] + - masks: List of multi masks [bs, hi, wi] + - refpoint_embed: [bs, num_dn, 4]. None in infer + - pos_embeds: List of multi pos embeds [bs, ci, hi, wi] + - tgt: [bs, num_dn, d_model]. None in infer + + """ + # prepare input for encoder + src_flatten = [] + mask_flatten = [] + lvl_pos_embed_flatten = [] + spatial_shapes = [] + for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)): + bs, c, h, w = src.shape + spatial_shape = (h, w) + spatial_shapes.append(spatial_shape) + + src = src.flatten(2).transpose(1, 2) # bs, hw, c + mask = mask.flatten(1) # bs, hw + pos_embed = pos_embed.flatten(2).transpose(1, 2) # bs, hw, c + if self.num_feature_levels > 1 and self.level_embed is not None: + lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1) + else: + lvl_pos_embed = pos_embed + lvl_pos_embed_flatten.append(lvl_pos_embed) + src_flatten.append(src) + mask_flatten.append(mask) + src_flatten = torch.cat(src_flatten, 1) # bs, \sum{hxw}, c + mask_flatten = torch.cat(mask_flatten, 1) # bs, \sum{hxw} + lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) # bs, \sum{hxw}, c + spatial_shapes = torch.as_tensor( + spatial_shapes, dtype=torch.long, device=src_flatten.device + ) + level_start_index = torch.cat( + (spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]) + ) + valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1) + + # two stage + enc_topk_proposals = enc_refpoint_embed = None + + ######################################################### + # Begin Encoder + ######################################################### + + memory, memory_text = self.encoder( + src_flatten, + pos=lvl_pos_embed_flatten, + level_start_index=level_start_index, + spatial_shapes=spatial_shapes, + valid_ratios=valid_ratios, + key_padding_mask=mask_flatten, + memory_text=text_dict["encoded_text"], + text_attention_mask=~text_dict["text_token_mask"], + # we ~ the mask . False means use the token; True means pad the token + position_ids=text_dict["position_ids"], + text_self_attention_masks=text_dict["text_self_attention_masks"], + ) + + ######################################################### + # End Encoder + # - memory: bs, \sum{hw}, c + # - mask_flatten: bs, \sum{hw} + # - lvl_pos_embed_flatten: bs, \sum{hw}, c + # - enc_intermediate_output: None or (nenc+1, bs, nq, c) or (nenc, bs, nq, c) + # - enc_intermediate_refpoints: None or (nenc+1, bs, nq, c) or (nenc, bs, nq, c) + ######################################################### + text_dict["encoded_text"] = memory_text + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # if memory.isnan().any() | memory.isinf().any(): + # import ipdb; ipdb.set_trace() + + + if self.two_stage_type == "standard": #把encoder的输出作为proposal + output_memory, output_proposals = gen_encoder_output_proposals( + memory, mask_flatten, spatial_shapes + ) + output_memory = self.enc_output_norm(self.enc_output(output_memory)) + + if text_dict is not None: + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory, text_dict) + else: + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory) + + topk_logits = enc_outputs_class_unselected.max(-1)[0] + enc_outputs_coord_unselected = ( + self.enc_out_bbox_embed(output_memory) + output_proposals + ) # (bs, \sum{hw}, 4) unsigmoid + topk = self.num_queries + + topk_proposals = torch.topk(topk_logits, topk, dim=1)[1] # bs, nq + + # gather boxes + refpoint_embed_undetach = torch.gather( + enc_outputs_coord_unselected, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4) + ) # unsigmoid + refpoint_embed_ = refpoint_embed_undetach.detach() + init_box_proposal = torch.gather( + output_proposals, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4) + ).sigmoid() # sigmoid + + # gather tgt + tgt_undetach = torch.gather( + output_memory, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, self.d_model) + ) + if self.embed_init_tgt: + tgt_ = ( + self.tgt_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, d_model + else: + tgt_ = tgt_undetach.detach() + + if refpoint_embed is not None: + refpoint_embed = torch.cat([refpoint_embed, refpoint_embed_], dim=1) + tgt = torch.cat([tgt, tgt_], dim=1) + else: + refpoint_embed, tgt = refpoint_embed_, tgt_ + + elif self.two_stage_type == "no": + tgt_ = ( + self.tgt_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, d_model + refpoint_embed_ = ( + self.refpoint_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, 4 + + if refpoint_embed is not None: + refpoint_embed = torch.cat([refpoint_embed, refpoint_embed_], dim=1) + tgt = torch.cat([tgt, tgt_], dim=1) + else: + refpoint_embed, tgt = refpoint_embed_, tgt_ + + if self.num_patterns > 0: + tgt_embed = tgt.repeat(1, self.num_patterns, 1) + refpoint_embed = refpoint_embed.repeat(1, self.num_patterns, 1) + tgt_pat = self.patterns.weight[None, :, :].repeat_interleave( + self.num_queries, 1 + ) # 1, n_q*n_pat, d_model + tgt = tgt_embed + tgt_pat + + init_box_proposal = refpoint_embed_.sigmoid() + + else: + raise NotImplementedError("unknown two_stage_type {}".format(self.two_stage_type)) + ######################################################### + # End preparing tgt + # - tgt: bs, NQ, d_model + # - refpoint_embed(unsigmoid): bs, NQ, d_model + ######################################################### + + ######################################################### + # Begin Decoder + ######################################################### + + #memory torch.Size([2, 16320, 256]) + + # import pdb;pdb.set_trace() + hs, references = self.decoder( + tgt=tgt.transpose(0, 1), + memory=memory.transpose(0, 1), + memory_key_padding_mask=mask_flatten, + pos=lvl_pos_embed_flatten.transpose(0, 1), + refpoints_unsigmoid=refpoint_embed.transpose(0, 1), + level_start_index=level_start_index, + spatial_shapes=spatial_shapes, + valid_ratios=valid_ratios, + tgt_mask=attn_mask, + memory_text=text_dict["encoded_text"], + text_attention_mask=~text_dict["text_token_mask"], + # we ~ the mask . False means use the token; True means pad the token + ) + ######################################################### + # End Decoder + # hs: n_dec, bs, nq, d_model + # references: n_dec+1, bs, nq, query_dim + ######################################################### + + ######################################################### + # Begin postprocess + ######################################################### + if self.two_stage_type == "standard": + hs_enc = tgt_undetach.unsqueeze(0) + ref_enc = refpoint_embed_undetach.sigmoid().unsqueeze(0) + else: + hs_enc = ref_enc = None + ######################################################### + # End postprocess + # hs_enc: (n_enc+1, bs, nq, d_model) or (1, bs, nq, d_model) or (n_enc, bs, nq, d_model) or None + # ref_enc: (n_enc+1, bs, nq, query_dim) or (1, bs, nq, query_dim) or (n_enc, bs, nq, d_model) or None + ######################################################### + + return hs, references, hs_enc, ref_enc, init_box_proposal + # hs: (n_dec, bs, nq, d_model) + # references: sigmoid coordinates. (n_dec+1, bs, bq, 4) + # hs_enc: (n_enc+1, bs, nq, d_model) or (1, bs, nq, d_model) or None + # ref_enc: sigmoid coordinates. \ + # (n_enc+1, bs, nq, query_dim) or (1, bs, nq, query_dim) or None + + +class TransformerEncoder(nn.Module): + def __init__( + self, + encoder_layer, + num_layers, + d_model=256, + num_queries=300, + enc_layer_share=False, + text_enhance_layer=None, + feature_fusion_layer=None, + use_checkpoint=False, + use_transformer_ckpt=False, + ): + """_summary_ + + Args: + encoder_layer (_type_): _description_ + num_layers (_type_): _description_ + norm (_type_, optional): _description_. Defaults to None. + d_model (int, optional): _description_. Defaults to 256. + num_queries (int, optional): _description_. Defaults to 300. + enc_layer_share (bool, optional): _description_. Defaults to False. + + """ + super().__init__() + # prepare layers + self.layers = [] + self.text_layers = [] + self.fusion_layers = [] + if num_layers > 0: + self.layers = _get_clones(encoder_layer, num_layers, layer_share=enc_layer_share) + + if text_enhance_layer is not None: + self.text_layers = _get_clones( + text_enhance_layer, num_layers, layer_share=enc_layer_share + ) + if feature_fusion_layer is not None: + self.fusion_layers = _get_clones( + feature_fusion_layer, num_layers, layer_share=enc_layer_share + ) + else: + self.layers = [] + del encoder_layer + + if text_enhance_layer is not None: + self.text_layers = [] + del text_enhance_layer + if feature_fusion_layer is not None: + self.fusion_layers = [] + del feature_fusion_layer + + self.query_scale = None + self.num_queries = num_queries + self.num_layers = num_layers + self.d_model = d_model + + self.use_checkpoint = use_checkpoint + self.use_transformer_ckpt = use_transformer_ckpt + + @staticmethod + def get_reference_points(spatial_shapes, valid_ratios, device): + reference_points_list = [] + for lvl, (H_, W_) in enumerate(spatial_shapes): + + ref_y, ref_x = torch.meshgrid( + torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device), + torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device), + ) + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + return reference_points + + def forward( + self, + # for images + src: Tensor, + pos: Tensor, + spatial_shapes: Tensor, + level_start_index: Tensor, + valid_ratios: Tensor, + key_padding_mask: Tensor, + # for texts + memory_text: Tensor = None, + text_attention_mask: Tensor = None, + pos_text: Tensor = None, + text_self_attention_masks: Tensor = None, + position_ids: Tensor = None, + ): + """ + Input: + - src: [bs, sum(hi*wi), 256] + - pos: pos embed for src. [bs, sum(hi*wi), 256] + - spatial_shapes: h,w of each level [num_level, 2] + - level_start_index: [num_level] start point of level in sum(hi*wi). + - valid_ratios: [bs, num_level, 2] + - key_padding_mask: [bs, sum(hi*wi)] + + - memory_text: bs, n_text, 256 + - text_attention_mask: bs, n_text + False for no padding; True for padding + - pos_text: bs, n_text, 256 + + - position_ids: bs, n_text + Intermedia: + - reference_points: [bs, sum(hi*wi), num_level, 2] + Outpus: + - output: [bs, sum(hi*wi), 256] + """ + + output = src + + # preparation and reshape + if self.num_layers > 0: + reference_points = self.get_reference_points( + spatial_shapes, valid_ratios, device=src.device + ) + + if self.text_layers: + # generate pos_text + bs, n_text, text_dim = memory_text.shape + if pos_text is None and position_ids is None: + pos_text = ( + torch.arange(n_text, device=memory_text.device) + .float() + .unsqueeze(0) + .unsqueeze(-1) + .repeat(bs, 1, 1) + ) + pos_text = get_sine_pos_embed(pos_text, num_pos_feats=256, exchange_xy=False) + if position_ids is not None: + pos_text = get_sine_pos_embed( + position_ids[..., None], num_pos_feats=256, exchange_xy=False + ) + + # main process + for layer_id, layer in enumerate(self.layers): + # if output.isnan().any() or memory_text.isnan().any(): + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + if self.fusion_layers: + if self.use_checkpoint: + output, memory_text = checkpoint.checkpoint( + self.fusion_layers[layer_id], + output, + memory_text, + key_padding_mask, + text_attention_mask, + ) + else: + output, memory_text = self.fusion_layers[layer_id]( + v=output, + l=memory_text, + attention_mask_v=key_padding_mask, + attention_mask_l=text_attention_mask, + ) + + if self.text_layers: + memory_text = self.text_layers[layer_id]( + src=memory_text.transpose(0, 1), + src_mask=~text_self_attention_masks, # note we use ~ for mask here + src_key_padding_mask=text_attention_mask, + pos=(pos_text.transpose(0, 1) if pos_text is not None else None), + ).transpose(0, 1) + + # main process + if self.use_transformer_ckpt: + output = checkpoint.checkpoint( + layer, + output, + pos, + reference_points, + spatial_shapes, + level_start_index, + key_padding_mask, + ) + else: + output = layer( + src=output, + pos=pos, + reference_points=reference_points, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + key_padding_mask=key_padding_mask, + ) + + return output, memory_text + + +class TransformerDecoder(nn.Module): + def __init__( + self, + decoder_layer, + num_layers, + norm=None, + return_intermediate=False, + d_model=256, + query_dim=4, + num_feature_levels=1, + ): + super().__init__() + if num_layers > 0: + self.layers = _get_clones(decoder_layer, num_layers) + else: + self.layers = [] + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + assert return_intermediate, "support return_intermediate only" + self.query_dim = query_dim + assert query_dim in [2, 4], "query_dim should be 2/4 but {}".format(query_dim) + self.num_feature_levels = num_feature_levels + + self.ref_point_head = MLP(query_dim // 2 * d_model, d_model, d_model, 2) + self.query_pos_sine_scale = None + + self.query_scale = None + self.bbox_embed = None + self.class_embed = None + + self.d_model = d_model + + self.ref_anchor_head = None + + def forward( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + refpoints_unsigmoid: Optional[Tensor] = None, # num_queries, bs, 2 + # for memory + level_start_index: Optional[Tensor] = None, # num_levels + spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + valid_ratios: Optional[Tensor] = None, + # for text + memory_text: Optional[Tensor] = None, + text_attention_mask: Optional[Tensor] = None, + ): + """ + Input: + - tgt: nq, bs, d_model + - memory: hw, bs, d_model + - pos: hw, bs, d_model + - refpoints_unsigmoid: nq, bs, 2/4 + - valid_ratios/spatial_shapes: bs, nlevel, 2 + """ + output = tgt + + intermediate = [] + reference_points = refpoints_unsigmoid.sigmoid() + ref_points = [reference_points] + + + + for layer_id, layer in enumerate(self.layers): + + if reference_points.shape[-1] == 4: + reference_points_input = ( + reference_points[:, :, None] + * torch.cat([valid_ratios, valid_ratios], -1)[None, :] + ) # nq, bs, nlevel, 4 + else: + assert reference_points.shape[-1] == 2 + reference_points_input = reference_points[:, :, None] * valid_ratios[None, :] + query_sine_embed = gen_sineembed_for_position( + reference_points_input[:, :, 0, :] + ) # nq, bs, 256*2 + + # conditional query + raw_query_pos = self.ref_point_head(query_sine_embed) # nq, bs, 256 + pos_scale = self.query_scale(output) if self.query_scale is not None else 1 + query_pos = pos_scale * raw_query_pos + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # if query_pos.isnan().any() | query_pos.isinf().any(): + # import ipdb; ipdb.set_trace() + + # main process + output = layer( + tgt=output, + tgt_query_pos=query_pos, + tgt_query_sine_embed=query_sine_embed, + tgt_key_padding_mask=tgt_key_padding_mask, + tgt_reference_points=reference_points_input, + memory_text=memory_text, + text_attention_mask=text_attention_mask, + memory=memory, + memory_key_padding_mask=memory_key_padding_mask, + memory_level_start_index=level_start_index, + memory_spatial_shapes=spatial_shapes, + memory_pos=pos, + self_attn_mask=tgt_mask, + cross_attn_mask=memory_mask, + ) + if output.isnan().any() | output.isinf().any(): + print(f"output layer_id {layer_id} is nan") + try: + num_nan = output.isnan().sum().item() + num_inf = output.isinf().sum().item() + print(f"num_nan {num_nan}, num_inf {num_inf}") + except Exception as e: + print(e) + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # import ipdb; ipdb.set_trace() + + # iter update + if self.bbox_embed is not None: + # box_holder = self.bbox_embed(output) + # box_holder[..., :self.query_dim] += inverse_sigmoid(reference_points) + # new_reference_points = box_holder[..., :self.query_dim].sigmoid() + + reference_before_sigmoid = inverse_sigmoid(reference_points) + delta_unsig = self.bbox_embed[layer_id](output) + outputs_unsig = delta_unsig + reference_before_sigmoid + new_reference_points = outputs_unsig.sigmoid() + + reference_points = new_reference_points.detach() + # if layer_id != self.num_layers - 1: + ref_points.append(new_reference_points) + + intermediate.append(self.norm(output)) + + # import pdb;pdb.set_trace() + + return [ + [itm_out.transpose(0, 1) for itm_out in intermediate], + [itm_refpoint.transpose(0, 1) for itm_refpoint in ref_points], + ] + + +class DeformableTransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + ): + super().__init__() + + # self attention + self.self_attn = MSDeformAttn( + embed_dim=d_model, + num_levels=n_levels, + num_heads=n_heads, + num_points=n_points, + batch_first=True, + ) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn) + self.dropout2 = nn.Dropout(dropout) + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout3 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(d_model) + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, src): + src2 = self.linear2(self.dropout2(self.activation(self.linear1(src)))) + src = src + self.dropout3(src2) + src = self.norm2(src) + return src + + def forward( + self, src, pos, reference_points, spatial_shapes, level_start_index, key_padding_mask=None + ): + # self attention + # import ipdb; ipdb.set_trace() + src2 = self.self_attn( + query=self.with_pos_embed(src, pos), + reference_points=reference_points, + value=src, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + key_padding_mask=key_padding_mask, + ) + src = src + self.dropout1(src2) + src = self.norm1(src) + + # ffn + src = self.forward_ffn(src) + + return src + + +class DeformableTransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + use_text_feat_guide=False, + use_text_cross_attention=False, + ): + super().__init__() + + # cross attention + self.cross_attn = MSDeformAttn( + embed_dim=d_model, + num_levels=n_levels, + num_heads=n_heads, + num_points=n_points, + batch_first=True, + ) + self.dropout1 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm1 = nn.LayerNorm(d_model) + + # cross attention text + if use_text_cross_attention: + self.ca_text = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.catext_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.catext_norm = nn.LayerNorm(d_model) + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.dropout2 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm2 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn, batch_dim=1) + self.dropout3 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout4 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm3 = nn.LayerNorm(d_model) + + self.key_aware_proj = None + self.use_text_feat_guide = use_text_feat_guide + assert not use_text_feat_guide + self.use_text_cross_attention = use_text_cross_attention + + def rm_self_attn_modules(self): + self.self_attn = None + self.dropout2 = None + self.norm2 = None + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, tgt): + with torch.cuda.amp.autocast(enabled=False): + tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout4(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward( + self, + # for tgt + tgt: Optional[Tensor], # nq, bs, d_model + tgt_query_pos: Optional[Tensor] = None, # pos for query. MLP(Sine(pos)) + tgt_query_sine_embed: Optional[Tensor] = None, # pos for query. Sine(pos) + tgt_key_padding_mask: Optional[Tensor] = None, + tgt_reference_points: Optional[Tensor] = None, # nq, bs, 4 + memory_text: Optional[Tensor] = None, # bs, num_token, d_model + text_attention_mask: Optional[Tensor] = None, # bs, num_token + # for memory + memory: Optional[Tensor] = None, # hw, bs, d_model + memory_key_padding_mask: Optional[Tensor] = None, + memory_level_start_index: Optional[Tensor] = None, # num_levels + memory_spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + memory_pos: Optional[Tensor] = None, # pos for memory + # sa + self_attn_mask: Optional[Tensor] = None, # mask used for self-attention + cross_attn_mask: Optional[Tensor] = None, # mask used for cross-attention + ): + """ + Input: + - tgt/tgt_query_pos: nq, bs, d_model + - + """ + assert cross_attn_mask is None + + # self attention + if self.self_attn is not None: + # import ipdb; ipdb.set_trace() + q = k = self.with_pos_embed(tgt, tgt_query_pos) + tgt2 = self.self_attn(q, k, tgt, attn_mask=self_attn_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + + if self.use_text_cross_attention: + tgt2 = self.ca_text( + self.with_pos_embed(tgt, tgt_query_pos), + memory_text.transpose(0, 1), + memory_text.transpose(0, 1), + key_padding_mask=text_attention_mask, + )[0] + tgt = tgt + self.catext_dropout(tgt2) + tgt = self.catext_norm(tgt) + + tgt2 = self.cross_attn( + query=self.with_pos_embed(tgt, tgt_query_pos).transpose(0, 1), + reference_points=tgt_reference_points.transpose(0, 1).contiguous(), + value=memory.transpose(0, 1), + spatial_shapes=memory_spatial_shapes, + level_start_index=memory_level_start_index, + key_padding_mask=memory_key_padding_mask, + ).transpose(0, 1) + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + + # ffn + tgt = self.forward_ffn(tgt) + + return tgt + + +def build_transformer(args): + return Transformer( + d_model=args.hidden_dim, + dropout=args.dropout, + nhead=args.nheads, + num_queries=args.num_queries, + dim_feedforward=args.dim_feedforward, + num_encoder_layers=args.enc_layers, + num_decoder_layers=args.dec_layers, + normalize_before=args.pre_norm, + return_intermediate_dec=True, + query_dim=args.query_dim, + activation=args.transformer_activation, + num_patterns=args.num_patterns, + num_feature_levels=args.num_feature_levels, + enc_n_points=args.enc_n_points, + dec_n_points=args.dec_n_points, + learnable_tgt_init=True, + # two stage + two_stage_type=args.two_stage_type, # ['no', 'standard', 'early'] + embed_init_tgt=args.embed_init_tgt, + use_text_enhancer=args.use_text_enhancer, + use_fusion_layer=args.use_fusion_layer, + use_checkpoint=args.use_checkpoint, + use_transformer_ckpt=args.use_transformer_ckpt, + use_text_cross_attention=args.use_text_cross_attention, + text_dropout=args.text_dropout, + fusion_dropout=args.fusion_dropout, + fusion_droppath=args.fusion_droppath, + ) diff --git a/models/GroundingDINO/transformer_loca.py b/models/GroundingDINO/transformer_loca.py new file mode 100644 index 0000000000000000000000000000000000000000..f4096b3665779238b661e8067ccd5342691d0077 --- /dev/null +++ b/models/GroundingDINO/transformer_loca.py @@ -0,0 +1,94 @@ +from .mlp_loca import MLP + +from torch import nn + + +class TransformerEncoder(nn.Module): + + def __init__( + self, + num_layers: int, + emb_dim: int, + num_heads: int, + dropout: float, + layer_norm_eps: float, + mlp_factor: int, + norm_first: bool, + activation: nn.Module, + norm: bool, + ): + + super(TransformerEncoder, self).__init__() + + self.layers = nn.ModuleList([ + TransformerEncoderLayer( + emb_dim, num_heads, dropout, layer_norm_eps, + mlp_factor, norm_first, activation + ) for _ in range(num_layers) + ]) + + self.norm = nn.LayerNorm(emb_dim, layer_norm_eps) if norm else nn.Identity() + + def forward(self, src, pos_emb, src_mask, src_key_padding_mask): + output = src + for layer in self.layers: + output = layer(output, pos_emb, src_mask, src_key_padding_mask) + return self.norm(output) + + +class TransformerEncoderLayer(nn.Module): + + def __init__( + self, + emb_dim: int, + num_heads: int, + dropout: float, + layer_norm_eps: float, + mlp_factor: int, + norm_first: bool, + activation: nn.Module, + ): + super(TransformerEncoderLayer, self).__init__() + + self.norm_first = norm_first + + self.norm1 = nn.LayerNorm(emb_dim, layer_norm_eps) + self.norm2 = nn.LayerNorm(emb_dim, layer_norm_eps) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.self_attn = nn.MultiheadAttention( + emb_dim, num_heads, dropout + ) + self.mlp = MLP(emb_dim, mlp_factor * emb_dim, dropout, activation) + + def with_emb(self, x, emb): + return x if emb is None else x + emb + + def forward(self, src, pos_emb, src_mask, src_key_padding_mask): + if self.norm_first: + src_norm = self.norm1(src) + q = k = src_norm + pos_emb + src = src + self.dropout1(self.self_attn( + query=q, + key=k, + value=src_norm, + attn_mask=src_mask, + key_padding_mask=src_key_padding_mask + )[0]) + + src_norm = self.norm2(src) + src = src + self.dropout2(self.mlp(src_norm)) + else: + q = k = src + pos_emb + src = self.norm1(src + self.dropout1(self.self_attn( + query=q, + key=k, + value=src, + attn_mask=src_mask, + key_padding_mask=src_key_padding_mask + )[0])) + + src = self.norm2(src + self.dropout2(self.mlp(src))) + + return src diff --git a/models/GroundingDINO/transformer_vanilla.py b/models/GroundingDINO/transformer_vanilla.py new file mode 100644 index 0000000000000000000000000000000000000000..10c0920c1a217af5bb3e1b13077568035ab3b7b5 --- /dev/null +++ b/models/GroundingDINO/transformer_vanilla.py @@ -0,0 +1,123 @@ +# ------------------------------------------------------------------------ +# 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] +# ------------------------------------------------------------------------ +# 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 +""" +DETR Transformer class. + +Copy-paste from torch.nn.Transformer with modifications: + * positional encodings are passed in MHattention + * extra LN at the end of encoder is removed + * decoder returns a stack of activations from all decoding layers +""" +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from .utils import ( + MLP, + _get_activation_fn, + _get_clones, + gen_encoder_output_proposals, + gen_sineembed_for_position, + sigmoid_focal_loss, +) + + +class TextTransformer(nn.Module): + def __init__(self, num_layers, d_model=256, nheads=8, dim_feedforward=2048, dropout=0.1): + super().__init__() + self.num_layers = num_layers + self.d_model = d_model + self.nheads = nheads + self.dim_feedforward = dim_feedforward + self.norm = None + + single_encoder_layer = TransformerEncoderLayer( + d_model=d_model, nhead=nheads, dim_feedforward=dim_feedforward, dropout=dropout + ) + self.layers = _get_clones(single_encoder_layer, num_layers) + + def forward(self, memory_text: torch.Tensor, text_attention_mask: torch.Tensor): + """ + + Args: + text_attention_mask: bs, num_token + memory_text: bs, num_token, d_model + + Raises: + RuntimeError: _description_ + + Returns: + output: bs, num_token, d_model + """ + + output = memory_text.transpose(0, 1) + + for layer in self.layers: + output = layer(output, src_key_padding_mask=text_attention_mask) + + if self.norm is not None: + output = self.norm(output) + + return output.transpose(0, 1) + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + self.nhead = nhead + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward( + self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + # repeat attn mask + if src_mask.dim() == 3 and src_mask.shape[0] == src.shape[1]: + # bs, num_q, num_k + src_mask = src_mask.repeat(self.nhead, 1, 1) + + q = k = self.with_pos_embed(src, pos) + + src2 = self.self_attn(q, k, value=src, attn_mask=src_mask)[0] + + # src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src diff --git a/models/GroundingDINO/utils.py b/models/GroundingDINO/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3c8e4b194290f178b45c68183c4113ebc768df39 --- /dev/null +++ b/models/GroundingDINO/utils.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 copy +import math + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +def _get_clones(module, N, layer_share=False): + # import ipdb; ipdb.set_trace() + if layer_share: + return nn.ModuleList([module for i in range(N)]) + else: + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def get_sine_pos_embed( + pos_tensor: torch.Tensor, + num_pos_feats: int = 128, + temperature: int = 10000, + exchange_xy: bool = True, +): + """generate sine position embedding from a position tensor + Args: + pos_tensor (torch.Tensor): shape: [..., n]. + num_pos_feats (int): projected shape for each float in the tensor. + temperature (int): temperature in the sine/cosine function. + exchange_xy (bool, optional): exchange pos x and pos y. \ + For example, input tensor is [x,y], the results will be [pos(y), pos(x)]. Defaults to True. + Returns: + pos_embed (torch.Tensor): shape: [..., n*num_pos_feats]. + """ + scale = 2 * math.pi + dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos_tensor.device) + dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats) + + def sine_func(x: torch.Tensor): + sin_x = x * scale / dim_t + sin_x = torch.stack((sin_x[..., 0::2].sin(), sin_x[..., 1::2].cos()), dim=3).flatten(2) + return sin_x + + pos_res = [sine_func(x) for x in pos_tensor.split([1] * pos_tensor.shape[-1], dim=-1)] + if exchange_xy: + pos_res[0], pos_res[1] = pos_res[1], pos_res[0] + pos_res = torch.cat(pos_res, dim=-1) + return pos_res + + +def gen_encoder_output_proposals( + memory: Tensor, memory_padding_mask: Tensor, spatial_shapes: Tensor, learnedwh=None +): + """ + Input: + - memory: bs, \sum{hw}, d_model + - memory_padding_mask: bs, \sum{hw} + - spatial_shapes: nlevel, 2 + - learnedwh: 2 + Output: + - output_memory: bs, \sum{hw}, d_model + - output_proposals: bs, \sum{hw}, 4 + """ + N_, S_, C_ = memory.shape + proposals = [] + _cur = 0 + for lvl, (H_, W_) in enumerate(spatial_shapes): + mask_flatten_ = memory_padding_mask[:, _cur : (_cur + H_ * W_)].view(N_, H_, W_, 1) + valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1) + valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1) + + # import ipdb; ipdb.set_trace() + + grid_y, grid_x = torch.meshgrid( + torch.linspace(0, H_ - 1, H_, dtype=torch.float32, device=memory.device), + torch.linspace(0, W_ - 1, W_, dtype=torch.float32, device=memory.device), + ) + grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) # H_, W_, 2 + + scale = torch.cat([valid_W.unsqueeze(-1), valid_H.unsqueeze(-1)], 1).view(N_, 1, 1, 2) + grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale + + if learnedwh is not None: + # import ipdb; ipdb.set_trace() + wh = torch.ones_like(grid) * learnedwh.sigmoid() * (2.0**lvl) + else: + wh = torch.ones_like(grid) * 0.05 * (2.0**lvl) + + # scale = torch.cat([W_[None].unsqueeze(-1), H_[None].unsqueeze(-1)], 1).view(1, 1, 1, 2).repeat(N_, 1, 1, 1) + # grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale + # wh = torch.ones_like(grid) / scale + proposal = torch.cat((grid, wh), -1).view(N_, -1, 4) + proposals.append(proposal) + _cur += H_ * W_ + # import ipdb; ipdb.set_trace() + output_proposals = torch.cat(proposals, 1) + output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all( + -1, keepdim=True + ) + output_proposals = torch.log(output_proposals / (1 - output_proposals)) # unsigmoid + output_proposals = output_proposals.masked_fill(memory_padding_mask.unsqueeze(-1), float("inf")) + output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf")) + + output_memory = memory + output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float(0)) + output_memory = output_memory.masked_fill(~output_proposals_valid, float(0)) + + # output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float('inf')) + # output_memory = output_memory.masked_fill(~output_proposals_valid, float('inf')) + + return output_memory, output_proposals + + +class RandomBoxPerturber: + def __init__( + self, x_noise_scale=0.2, y_noise_scale=0.2, w_noise_scale=0.2, h_noise_scale=0.2 + ) -> None: + self.noise_scale = torch.Tensor( + [x_noise_scale, y_noise_scale, w_noise_scale, h_noise_scale] + ) + + def __call__(self, refanchors: Tensor) -> Tensor: + nq, bs, query_dim = refanchors.shape + device = refanchors.device + + noise_raw = torch.rand_like(refanchors) + noise_scale = self.noise_scale.to(device)[:query_dim] + + new_refanchors = refanchors * (1 + (noise_raw - 0.5) * noise_scale) + return new_refanchors.clamp_(0, 1) + + +def sigmoid_focal_loss( + inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2, no_reduction=False +): + """ + Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. + Args: + inputs: A float tensor of arbitrary shape. + The predictions for each example. + targets: A float tensor with the same shape as inputs. Stores the binary + classification label for each element in inputs + (0 for the negative class and 1 for the positive class). + alpha: (optional) Weighting factor in range (0,1) to balance + positive vs negative examples. Default = -1 (no weighting). + gamma: Exponent of the modulating factor (1 - p_t) to + balance easy vs hard examples. + Returns: + Loss tensor + """ + prob = inputs.sigmoid() + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = prob * targets + (1 - prob) * (1 - targets) + loss = ce_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * loss + + if no_reduction: + return loss + + return loss.mean(1).sum() / num_boxes + + +class MLP(nn.Module): + """Very simple multi-layer perceptron (also called FFN)""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +def _get_activation_fn(activation, d_model=256, batch_dim=0): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + if activation == "prelu": + return nn.PReLU() + if activation == "selu": + return F.selu + + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def gen_sineembed_for_position(pos_tensor): + # n_query, bs, _ = pos_tensor.size() + # sineembed_tensor = torch.zeros(n_query, bs, 256) + scale = 2 * math.pi + dim_t = torch.arange(128, dtype=torch.float32, device=pos_tensor.device) + dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode='floor')) / 128) + x_embed = pos_tensor[:, :, 0] * scale + y_embed = pos_tensor[:, :, 1] * scale + 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=3).flatten(2) + pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) + if pos_tensor.size(-1) == 2: + pos = torch.cat((pos_y, pos_x), dim=2) + elif pos_tensor.size(-1) == 4: + w_embed = pos_tensor[:, :, 2] * scale + pos_w = w_embed[:, :, None] / dim_t + pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2) + + h_embed = pos_tensor[:, :, 3] * scale + pos_h = h_embed[:, :, None] / dim_t + pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2) + + pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) + else: + raise ValueError("Unknown pos_tensor shape(-1):{}".format(pos_tensor.size(-1))) + return pos + + +class ContrastiveEmbed(nn.Module): + def __init__(self, max_text_len=256): + """ + Args: + max_text_len: max length of text. + """ + super().__init__() + self.max_text_len = max_text_len + + def forward(self, x, text_dict): + """_summary_ + + Args: + x (_type_): _description_ + text_dict (_type_): _description_ + { + 'encoded_text': encoded_text, # bs, 195, d_model + 'text_token_mask': text_token_mask, # bs, 195 + # True for used tokens. False for padding tokens + } + Returns: + _type_: _description_ + """ + assert isinstance(text_dict, dict) + # print(x) #torch.Size([2, 16320, 256]) + # print(text_dict) + + # import pdb;pdb.set_trace() + y = text_dict["encoded_text"] #torch.Size([2, 195, 256]) + text_token_mask = text_dict["text_token_mask"] + + res = x @ y.transpose(-1, -2) + res.masked_fill_(~text_token_mask[:, None, :], float("-inf")) + # 接着，对res进行掩码操作，将未使用的文本token（即padding的token）对应的得分置为负无穷float("-inf")。这是为了在计算相似度时，排除padding部分的影响。 + + + # padding to max_text_len + new_res = torch.full((*res.shape[:-1], self.max_text_len), float("-inf"), device=res.device) + new_res[..., : res.shape[-1]] = res #torch.Size([2, 16320, 195]) + + return new_res diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b123bbf7ac35014e68a70c2056f1325389f0ecc --- /dev/null +++ b/models/__init__.py @@ -0,0 +1,10 @@ +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +from .GroundingDINO import build_groundingdino + +def build_model(args): + return build(args) diff --git a/models/registry.py b/models/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..f438c6e3918a84cc2004b5da9c1d79d18cfb3118 --- /dev/null +++ b/models/registry.py @@ -0,0 +1,58 @@ +# -*- coding: utf-8 -*- +# @Author: Yihao Chen +# @Date: 2021-08-16 16:03:17 +# @Last Modified by: Shilong Liu +# @Last Modified time: 2022-01-23 15:26 +# modified from mmcv + +import inspect +from functools import partial + + +class Registry(object): + + def __init__(self, name): + self._name = name + self._module_dict = dict() + + def __repr__(self): + format_str = self.__class__.__name__ + '(name={}, items={})'.format( + self._name, list(self._module_dict.keys())) + return format_str + + def __len__(self): + return len(self._module_dict) + + @property + def name(self): + return self._name + + @property + def module_dict(self): + return self._module_dict + + def get(self, key): + return self._module_dict.get(key, None) + + def registe_with_name(self, module_name=None, force=False): + return partial(self.register, module_name=module_name, force=force) + + def register(self, module_build_function, module_name=None, force=False): + """Register a module build function. + Args: + module (:obj:`nn.Module`): Module to be registered. + """ + if not inspect.isfunction(module_build_function): + raise TypeError('module_build_function must be a function, but got {}'.format( + type(module_build_function))) + if module_name is None: + module_name = module_build_function.__name__ + if not force and module_name in self._module_dict: + raise KeyError('{} is already registered in {}'.format( + module_name, self.name)) + self._module_dict[module_name] = module_build_function + + return module_build_function + +MODULE_BUILD_FUNCS = Registry('model build functions') + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..960df03a1fdcd69a59497739ff91fd125cc0c19c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,16 @@ +torch +torchvision +transformers<4.50 +huggingface_hub>=0.23 +safetensors +timm +addict +yapf==0.40.1 +termcolor +scipy +numpy +pillow +matplotlib +opencv-python-headless +pycocotools +pyyaml diff --git a/util/__init__.py b/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..168f9979a4623806934b0ff1102ac166704e7dec --- /dev/null +++ b/util/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved diff --git a/util/box_loss.py b/util/box_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..c0f94a5b5804a5e6a613785249d92f71594f4ab7 --- /dev/null +++ b/util/box_loss.py @@ -0,0 +1,112 @@ +# borrow from https://github.com/Zzh-tju/CIoU/blob/master/layers/modules/multibox_loss.py + +import torch, math + + + +def ciou(bboxes1, bboxes2): + bboxes1 = torch.sigmoid(bboxes1) + bboxes2 = torch.sigmoid(bboxes2) + rows = bboxes1.shape[0] + cols = bboxes2.shape[0] + cious = torch.zeros((rows, cols)) + if rows * cols == 0: + return cious + exchange = False + if bboxes1.shape[0] > bboxes2.shape[0]: + bboxes1, bboxes2 = bboxes2, bboxes1 + cious = torch.zeros((cols, rows)) + exchange = True + w1 = torch.exp(bboxes1[:, 2]) + h1 = torch.exp(bboxes1[:, 3]) + w2 = torch.exp(bboxes2[:, 2]) + h2 = torch.exp(bboxes2[:, 3]) + area1 = w1 * h1 + area2 = w2 * h2 + center_x1 = bboxes1[:, 0] + center_y1 = bboxes1[:, 1] + center_x2 = bboxes2[:, 0] + center_y2 = bboxes2[:, 1] + + inter_l = torch.max(center_x1 - w1 / 2,center_x2 - w2 / 2) + inter_r = torch.min(center_x1 + w1 / 2,center_x2 + w2 / 2) + inter_t = torch.max(center_y1 - h1 / 2,center_y2 - h2 / 2) + inter_b = torch.min(center_y1 + h1 / 2,center_y2 + h2 / 2) + inter_area = torch.clamp((inter_r - inter_l),min=0) * torch.clamp((inter_b - inter_t),min=0) + + c_l = torch.min(center_x1 - w1 / 2,center_x2 - w2 / 2) + c_r = torch.max(center_x1 + w1 / 2,center_x2 + w2 / 2) + c_t = torch.min(center_y1 - h1 / 2,center_y2 - h2 / 2) + c_b = torch.max(center_y1 + h1 / 2,center_y2 + h2 / 2) + + inter_diag = (center_x2 - center_x1)**2 + (center_y2 - center_y1)**2 + c_diag = torch.clamp((c_r - c_l),min=0)**2 + torch.clamp((c_b - c_t),min=0)**2 + + union = area1+area2-inter_area + u = (inter_diag) / c_diag + iou = inter_area / union + v = (4 / (math.pi ** 2)) * torch.pow((torch.atan(w2 / h2) - torch.atan(w1 / h1)), 2) + with torch.no_grad(): + S = (iou>0.5).float() + alpha= S*v/(1-iou+v) + cious = iou - u - alpha * v + cious = torch.clamp(cious,min=-1.0,max = 1.0) + if exchange: + cious = cious.T + return 1-cious + +def diou(bboxes1, bboxes2): + bboxes1 = torch.sigmoid(bboxes1) + bboxes2 = torch.sigmoid(bboxes2) + rows = bboxes1.shape[0] + cols = bboxes2.shape[0] + cious = torch.zeros((rows, cols)) + if rows * cols == 0: + return cious + exchange = False + if bboxes1.shape[0] > bboxes2.shape[0]: + bboxes1, bboxes2 = bboxes2, bboxes1 + cious = torch.zeros((cols, rows)) + exchange = True + w1 = torch.exp(bboxes1[:, 2]) + h1 = torch.exp(bboxes1[:, 3]) + w2 = torch.exp(bboxes2[:, 2]) + h2 = torch.exp(bboxes2[:, 3]) + area1 = w1 * h1 + area2 = w2 * h2 + center_x1 = bboxes1[:, 0] + center_y1 = bboxes1[:, 1] + center_x2 = bboxes2[:, 0] + center_y2 = bboxes2[:, 1] + + inter_l = torch.max(center_x1 - w1 / 2,center_x2 - w2 / 2) + inter_r = torch.min(center_x1 + w1 / 2,center_x2 + w2 / 2) + inter_t = torch.max(center_y1 - h1 / 2,center_y2 - h2 / 2) + inter_b = torch.min(center_y1 + h1 / 2,center_y2 + h2 / 2) + inter_area = torch.clamp((inter_r - inter_l),min=0) * torch.clamp((inter_b - inter_t),min=0) + + c_l = torch.min(center_x1 - w1 / 2,center_x2 - w2 / 2) + c_r = torch.max(center_x1 + w1 / 2,center_x2 + w2 / 2) + c_t = torch.min(center_y1 - h1 / 2,center_y2 - h2 / 2) + c_b = torch.max(center_y1 + h1 / 2,center_y2 + h2 / 2) + + inter_diag = (center_x2 - center_x1)**2 + (center_y2 - center_y1)**2 + c_diag = torch.clamp((c_r - c_l),min=0)**2 + torch.clamp((c_b - c_t),min=0)**2 + + union = area1+area2-inter_area + u = (inter_diag) / c_diag + iou = inter_area / union + dious = iou - u + dious = torch.clamp(dious,min=-1.0,max = 1.0) + if exchange: + dious = dious.T + return 1-dious + + +if __name__ == "__main__": + x = torch.rand(10, 4) + y = torch.rand(10,4) + import ipdb;ipdb.set_trace() + cxy = ciou(x, y) + dxy = diou(x, y) + print(cxy.shape, dxy.shape) diff --git a/util/box_ops.py b/util/box_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..99ba6d0f67d536aa77bbf787c5ef19aeb56a50df --- /dev/null +++ b/util/box_ops.py @@ -0,0 +1,138 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Utilities for bounding box manipulation and GIoU. +""" +import torch, os +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) + + + 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(), f"{boxes1}" + assert (boxes2[:, 2:] >= boxes2[:, :2]).all(), f"{boxes2}" + + 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() \ No newline at end of file diff --git a/util/coco_id2name.json b/util/coco_id2name.json new file mode 100644 index 0000000000000000000000000000000000000000..52ad2e96e7974593e13ee56cd3eff0d22c59f96d --- /dev/null +++ b/util/coco_id2name.json @@ -0,0 +1 @@ +{"1": "person", "2": "bicycle", "3": "car", "4": "motorcycle", "5": "airplane", "6": "bus", "7": "train", "8": "truck", "9": "boat", "10": "traffic light", "11": "fire hydrant", "13": "stop sign", "14": "parking meter", "15": "bench", "16": "bird", "17": "cat", "18": "dog", "19": "horse", "20": "sheep", "21": "cow", "22": "elephant", "23": "bear", "24": "zebra", "25": "giraffe", "27": "backpack", "28": "umbrella", "31": "handbag", "32": "tie", "33": "suitcase", "34": "frisbee", "35": "skis", "36": "snowboard", "37": "sports ball", "38": "kite", "39": "baseball bat", "40": "baseball glove", "41": "skateboard", "42": "surfboard", "43": "tennis racket", "44": "bottle", "46": "wine glass", "47": "cup", "48": "fork", "49": "knife", "50": "spoon", "51": "bowl", "52": "banana", "53": "apple", "54": "sandwich", "55": "orange", "56": "broccoli", "57": "carrot", "58": "hot dog", "59": "pizza", "60": "donut", "61": "cake", "62": "chair", "63": "couch", "64": "potted plant", "65": "bed", "67": "dining table", "70": "toilet", "72": "tv", "73": "laptop", "74": "mouse", "75": "remote", "76": "keyboard", "77": "cell phone", "78": "microwave", "79": "oven", "80": "toaster", "81": "sink", "82": "refrigerator", "84": "book", "85": "clock", "86": "vase", "87": "scissors", "88": "teddy bear", "89": "hair drier", "90": "toothbrush"} \ No newline at end of file diff --git a/util/get_param_dicts.py b/util/get_param_dicts.py new file mode 100644 index 0000000000000000000000000000000000000000..5f96351e74f98557932393904f019df03b2cec9f --- /dev/null +++ b/util/get_param_dicts.py @@ -0,0 +1,85 @@ +import json +import torch +import torch.nn as nn + + +def match_name_keywords(n: str, name_keywords: list): + out = False + for b in name_keywords: + if b in n: + out = True + break + return out + + +def get_param_dict(args, model_without_ddp: nn.Module): + try: + param_dict_type = args.param_dict_type + except: + param_dict_type = 'default' + assert param_dict_type in ['default', 'ddetr_in_mmdet', 'large_wd'] + + # by default + # import pdb;pdb.set_trace() + if param_dict_type == 'default': + param_dicts = [ + {"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]}, + { + "params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad], + "lr": args.lr_backbone, + } + ] + return param_dicts + + if param_dict_type == 'ddetr_in_mmdet': + param_dicts = [ + { + "params": + [p for n, p in model_without_ddp.named_parameters() + if not match_name_keywords(n, args.lr_backbone_names) and not match_name_keywords(n, args.lr_linear_proj_names) and p.requires_grad], + "lr": args.lr, + }, + { + "params": [p for n, p in model_without_ddp.named_parameters() + if match_name_keywords(n, args.lr_backbone_names) and p.requires_grad], + "lr": args.lr_backbone, + }, + { + "params": [p for n, p in model_without_ddp.named_parameters() + if match_name_keywords(n, args.lr_linear_proj_names) and p.requires_grad], + "lr": args.lr_linear_proj_mult, + } + ] + return param_dicts + + if param_dict_type == 'large_wd': + param_dicts = [ + { + "params": + [p for n, p in model_without_ddp.named_parameters() + if not match_name_keywords(n, ['backbone']) and not match_name_keywords(n, ['norm', 'bias']) and p.requires_grad], + }, + { + "params": [p for n, p in model_without_ddp.named_parameters() + if match_name_keywords(n, ['backbone']) and match_name_keywords(n, ['norm', 'bias']) and p.requires_grad], + "lr": args.lr_backbone, + "weight_decay": 0.0, + }, + { + "params": [p for n, p in model_without_ddp.named_parameters() + if match_name_keywords(n, ['backbone']) and not match_name_keywords(n, ['norm', 'bias']) and p.requires_grad], + "lr": args.lr_backbone, + "weight_decay": args.weight_decay, + }, + { + "params": + [p for n, p in model_without_ddp.named_parameters() + if not match_name_keywords(n, ['backbone']) and match_name_keywords(n, ['norm', 'bias']) and p.requires_grad], + "lr": args.lr, + "weight_decay": 0.0, + } + ] + + # print("param_dicts: {}".format(param_dicts)) + + return param_dicts \ No newline at end of file diff --git a/util/get_tokenlizer.py b/util/get_tokenlizer.py new file mode 100644 index 0000000000000000000000000000000000000000..77873bcf4e098c194ef6182b8cbfdc2aa4e3d0ed --- /dev/null +++ b/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)) diff --git a/util/logger.py b/util/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..99d2fd8ef20b2d1354b3b5d3e6d8b4987e5c8ce2 --- /dev/null +++ b/util/logger.py @@ -0,0 +1,74 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +import functools +import logging +import os +import sys +import colorlog + + +# 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) + formatter = colorlog.ColoredFormatter( + "%(log_color)s%(levelname)-8s%(reset)s %(log_color)s%(asctime)s | %(blue)s%(message)s", + datefmt=None, + reset=True, + log_colors={ + 'DEBUG': 'cyan', + 'INFO': 'green', + 'WARNING': 'yellow', + 'ERROR': 'red', + 'CRITICAL': 'red,bg_white' + } + ) + logger.propagate = False + + if abbrev_name is None: + abbrev_name = name + + # stdout logging: master only + if distributed_rank == 0: + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(logging.DEBUG) + 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(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") diff --git a/util/misc.py b/util/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..4a3d759dea39fe8ad01f04d2bb19842b995f53e3 --- /dev/null +++ b/util/misc.py @@ -0,0 +1,626 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Misc functions, including distributed helpers. + +Mostly copy-paste from torchvision references. +""" +import os +import random +import subprocess +import time +from collections import OrderedDict, defaultdict, deque +import datetime +import pickle +from typing import Optional, List + +import json, time +import numpy as np +import torch +import torch.distributed as dist +from torch import Tensor + +import colorsys + +# needed due to empty tensor bug in pytorch and torchvision 0.5 +import torchvision +__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): + return self.total / self.count + + @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) + + +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 + """ + 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 + + 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): + + 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' ... + + local_world_size = int(os.environ['WORLD_SIZE']) + 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, init_method=args.dist_url, + world_size=args.world_size, rank=args.rank) + print("Before torch.distributed.barrier()") + torch.distributed.barrier() + print("End torch.distributed.barrier()") + setup_for_distributed(args.rank == 0) + +def setup_distributed(args): + if 'WORLD_SIZE' in os.environ and os.environ['WORLD_SIZE'] != '': # 'RANK' in os.environ and + local_world_size = int(os.environ['WORLD_SIZE']) + 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_NTASKS']) + node_list = os.environ["SLURM_NODELIST"] + addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1") + if "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = "23233" + if "MASTER_ADDR" not in os.environ: + os.environ["MASTER_ADDR"] = addr + os.environ["WORLD_SIZE"] = str(args.world_size) + os.environ["LOCAL_RANK"] = str(args.local_rank) + os.environ["RANK"] = str(args.rank) + 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 + + args.distributed = True + torch.cuda.set_device(args.local_rank) + args.dist_backend = 'nccl' + torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, + world_size=args.world_size, rank=args.rank) + torch.distributed.barrier() + print(f' == distributed init (rank {args.rank}) done.') + + 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 + + +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., 360., 360. / num_colors): + hue = i/360. + lightness = (50 + np.random.rand() * 10)/100. + saturation = (90 + np.random.rand() * 10)/100. + 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 \ No newline at end of file diff --git a/util/plot_utils.py b/util/plot_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..af67acd9f7a3ad7ec61908adfa2c67acb734ca08 --- /dev/null +++ b/util/plot_utils.py @@ -0,0 +1,112 @@ +""" +Plotting utilities to visualize training logs. +""" +import torch +import pandas as pd +import numpy as np +import seaborn as sns +import matplotlib.pyplot as plt + +from pathlib import Path, PurePath + + +def plot_logs(logs, fields=('class_error', 'loss_bbox_unscaled', 'mAP'), ewm_col=0, log_name='log.txt'): + ''' + Function to plot specific fields from training log(s). Plots both training and test results. + + :: Inputs - logs = list containing Path objects, each pointing to individual dir with a log file + - fields = which results to plot from each log file - plots both training and test for each field. + - ewm_col = optional, which column to use as the exponential weighted smoothing of the plots + - log_name = optional, name of log file if different than default 'log.txt'. + + :: Outputs - matplotlib plots of results in fields, color coded for each log file. + - solid lines are training results, dashed lines are test results. + + ''' + func_name = "plot_utils.py::plot_logs" + + # verify logs is a list of Paths (list[Paths]) or single Pathlib object Path, + # convert single Path to list to avoid 'not iterable' error + + if not isinstance(logs, list): + if isinstance(logs, PurePath): + logs = [logs] + print(f"{func_name} info: logs param expects a list argument, converted to list[Path].") + else: + raise ValueError(f"{func_name} - invalid argument for logs parameter.\n \ + Expect list[Path] or single Path obj, received {type(logs)}") + + # Quality checks - verify valid dir(s), that every item in list is Path object, and that log_name exists in each dir + for i, dir in enumerate(logs): + if not isinstance(dir, PurePath): + raise ValueError(f"{func_name} - non-Path object in logs argument of {type(dir)}: \n{dir}") + if not dir.exists(): + raise ValueError(f"{func_name} - invalid directory in logs argument:\n{dir}") + # verify log_name exists + fn = Path(dir / log_name) + if not fn.exists(): + print(f"-> missing {log_name}. Have you gotten to Epoch 1 in training?") + print(f"--> full path of missing log file: {fn}") + return + + # load log file(s) and plot + dfs = [pd.read_json(Path(p) / log_name, lines=True) for p in logs] + + fig, axs = plt.subplots(ncols=len(fields), figsize=(16, 5)) + + for df, color in zip(dfs, sns.color_palette(n_colors=len(logs))): + for j, field in enumerate(fields): + if field == 'mAP': + coco_eval = pd.DataFrame( + np.stack(df.test_coco_eval_bbox.dropna().values)[:, 1] + ).ewm(com=ewm_col).mean() + axs[j].plot(coco_eval, c=color) + else: + df.interpolate().ewm(com=ewm_col).mean().plot( + y=[f'train_{field}', f'test_{field}'], + ax=axs[j], + color=[color] * 2, + style=['-', '--'] + ) + for ax, field in zip(axs, fields): + if field == 'mAP': + ax.legend([Path(p).name for p in logs]) + ax.set_title(field) + else: + ax.legend([f'train', f'test']) + ax.set_title(field) + + return fig, axs + +def plot_precision_recall(files, naming_scheme='iter'): + if naming_scheme == 'exp_id': + # name becomes exp_id + names = [f.parts[-3] for f in files] + elif naming_scheme == 'iter': + names = [f.stem for f in files] + else: + raise ValueError(f'not supported {naming_scheme}') + fig, axs = plt.subplots(ncols=2, figsize=(16, 5)) + for f, color, name in zip(files, sns.color_palette("Blues", n_colors=len(files)), names): + data = torch.load(f) + # precision is n_iou, n_points, n_cat, n_area, max_det + precision = data['precision'] + recall = data['params'].recThrs + scores = data['scores'] + # take precision for all classes, all areas and 100 detections + precision = precision[0, :, :, 0, -1].mean(1) + scores = scores[0, :, :, 0, -1].mean(1) + prec = precision.mean() + rec = data['recall'][0, :, 0, -1].mean() + print(f'{naming_scheme} {name}: mAP@50={prec * 100: 05.1f}, ' + + f'score={scores.mean():0.3f}, ' + + f'f1={2 * prec * rec / (prec + rec + 1e-8):0.3f}' + ) + axs[0].plot(recall, precision, c=color) + axs[1].plot(recall, scores, c=color) + + axs[0].set_title('Precision / Recall') + axs[0].legend(names) + axs[1].set_title('Scores / Recall') + axs[1].legend(names) + return fig, axs diff --git a/util/slconfig.py b/util/slconfig.py new file mode 100644 index 0000000000000000000000000000000000000000..8064727596d57ccbe8e9f94dbb0bba9be732812a --- /dev/null +++ b/util/slconfig.py @@ -0,0 +1,440 @@ +# ========================================================== +# Modified from mmcv +# ========================================================== +import os, sys +import os.path as osp +import ast +import tempfile +import shutil +from importlib import import_module + +from argparse import Action + +from addict import Dict +from yapf.yapflib.yapf_api import FormatCode + +import platform +MACOS, LINUX, WINDOWS = (platform.system() == x for x in ['Darwin', 'Linux', 'Windows']) # environment booleans + +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 WINDOWS: + 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] + """ + + 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 : + + 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): + + 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) + diff --git a/util/slio.py b/util/slio.py new file mode 100644 index 0000000000000000000000000000000000000000..8b8f4dad2441b8352ab7311dbf16019515441331 --- /dev/null +++ b/util/slio.py @@ -0,0 +1,173 @@ +# ========================================================== +# Modified from mmcv +# ========================================================== + +import json, pickle, yaml +try: + from yaml import CLoader as Loader, CDumper as Dumper +except ImportError: + from yaml import Loader, Dumper + +from pathlib import Path +from abc import ABCMeta, abstractmethod + +# =========================== +# 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') diff --git a/util/static_data_path.py b/util/static_data_path.py new file mode 100644 index 0000000000000000000000000000000000000000..88a075a848cb634113563b1bcaff6d7c5f6982e2 --- /dev/null +++ b/util/static_data_path.py @@ -0,0 +1,10 @@ +coco = dict( + train = dict( + img_folder = '/comp_robot/cv_public_dataset/COCO2017/val2017', + ann_file = '/comp_robot/cv_public_dataset/COCO2017/annotations/instances_val2017.json' + ), + val = dict( + img_folder = '/comp_robot/cv_public_dataset/COCO2017/val2017', + ann_file = '/comp_robot/cv_public_dataset/COCO2017/annotations/instances_val2017.json' + ) +) \ No newline at end of file diff --git a/util/time_counter.py b/util/time_counter.py new file mode 100644 index 0000000000000000000000000000000000000000..19dc2e640bb0f19e686a5078d8e4d7db7ddaad96 --- /dev/null +++ b/util/time_counter.py @@ -0,0 +1,60 @@ +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__) \ No newline at end of file diff --git a/util/utils.py b/util/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8852c785f85969fd635e64854499ff731e82da1e --- /dev/null +++ b/util/utils.py @@ -0,0 +1,471 @@ +from collections import OrderedDict +from copy import deepcopy +import json +import warnings + +import torch +import numpy as np + +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.**torch.linspace(0., max_freq, steps=N_freqs) + else: + freq_bands = torch.linspace(2.**0., 2.**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) + +import argparse +from util.slconfig import SLConfig +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) == '' + >>> assert repr(foo).startswith('>> 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) == '' + """ + + 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() + + 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. - 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__() + \ No newline at end of file diff --git a/util/vis_utils.py b/util/vis_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7fca5785e7bd6b54909edca083e9fa13e6be18a5 --- /dev/null +++ b/util/vis_utils.py @@ -0,0 +1,92 @@ +import cv2 +import numpy as np + +from util.utils import renorm +from util.misc import color_sys + +_color_getter = color_sys(100) + +# plot known and unknown box +def add_box_to_img(img, boxes, colorlist, brands=None): + """[summary] + + Args: + img ([type]): np.array, H,W,3 + boxes ([type]): list of list(4) + colorlist: list of colors. + brands: text. + + Return: + img: np.array. H,W,3. + """ + H, W = img.shape[:2] + for _i, (box, color) in enumerate(zip(boxes, colorlist)): + x, y, w, h = box[0] * W, box[1] * H, box[2] * W, box[3] * H + img = cv2.rectangle(img.copy(), (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), color, 2) + if brands is not None: + brand = brands[_i] + org = (int(x-w/2), int(y+h/2)) + font = cv2.FONT_HERSHEY_SIMPLEX + fontScale = 0.5 + thickness = 1 + img = cv2.putText(img.copy(), str(brand), org, font, + fontScale, color, thickness, cv2.LINE_AA) + return img + +def plot_dual_img(img, boxes, labels, idxs, probs=None): + """[summary] + + Args: + img ([type]): 3,H,W. tensor. + boxes (): tensor(Kx4) or list of tensor(1x4). + labels ([type]): list of ints. + idxs ([type]): list of ints. + probs (optional): listof floats. + + Returns: + img_classcolor: np.array. H,W,3. img with class-wise label. + img_seqcolor: np.array. H,W,3. img with seq-wise label. + """ + + boxes = [i.cpu().tolist() for i in boxes] + img = (renorm(img.cpu()).permute(1,2,0).numpy() * 255).astype(np.uint8) + # plot with class + class_colors = [_color_getter(i) for i in labels] + if probs is not None: + brands = ["{},{:.2f}".format(j,k) for j,k in zip(labels, probs)] + else: + brands = labels + img_classcolor = add_box_to_img(img, boxes, class_colors, brands=brands) + # plot with seq + seq_colors = [_color_getter((i * 11) % 100) for i in idxs] + img_seqcolor = add_box_to_img(img, boxes, seq_colors, brands=idxs) + return img_classcolor, img_seqcolor + + +def plot_raw_img(img, boxes, labels): + """[summary] + + Args: + img ([type]): 3,H,W. tensor. + boxes ([type]): Kx4. tensor + labels ([type]): K. tensor. + + return: + img: np.array. H,W,3. img with bbox annos. + + """ + img = (renorm(img.cpu()).permute(1,2,0).numpy() * 255).astype(np.uint8) + H, W = img.shape[:2] + for box, label in zip(boxes.tolist(), labels.tolist()): + x, y, w, h = box[0] * W, box[1] * H, box[2] * W, box[3] * H + + img = cv2.rectangle(img.copy(), (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), _color_getter(label), 2) + # add text + org = (int(x-w/2), int(y+h/2)) + font = cv2.FONT_HERSHEY_SIMPLEX + fontScale = 1 + thickness = 1 + img = cv2.putText(img.copy(), str(label), org, font, + fontScale, _color_getter(label), thickness, cv2.LINE_AA) + + return img \ No newline at end of file diff --git a/util/visualizer.py b/util/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..5eb2a28d6593a5767ec6b78ce4cb482af344879b --- /dev/null +++ b/util/visualizer.py @@ -0,0 +1,130 @@ +# -*- coding: utf-8 -*- +''' +@File : visualizer.py +@Time : 2022/04/05 11:39:33 +@Author : Shilong Liu +@Contact : liusl20@mail.tsinghua.edu.cn; slongliu86@gmail.com +Modified from COCO evaluator +''' + +import os, sys +from textwrap import wrap +import torch +import numpy as np +import cv2 +import datetime + +import matplotlib.pyplot as plt +from matplotlib.collections import PatchCollection +from matplotlib.patches import Polygon +from pycocotools import mask as maskUtils +from matplotlib import transforms + +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 + + +class COCOVisualizer(): + def __init__(self) -> None: + pass + + def visualize(self, img, tgt, caption=None, dpi=120, savedir=None, show_in_console=True): + """ + 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) + ax.imshow(img) + + self.addtgt(tgt) + if show_in_console: + plt.show() + + if savedir is not None: + if caption is None: + savename = '{}/{}-{}.png'.format(savedir, int(tgt['image_id']), str(datetime.datetime.now()).replace(' ', '-')) + else: + savename = '{}/{}-{}-{}.png'.format(savedir, caption, int(tgt['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): + """ + - tgt: dict. args: + - boxes: num_boxes, 4. xywh, [0,1]. + - box_label: num_boxes. + """ + assert 'boxes' in tgt + ax = plt.gca() + H, W = tgt['size'].tolist() + 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 '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) + +