[
{
"path": ".gitignore",
"content": "**/__pycache__/**\n**/wandb/**\n**/.vscode/**\n*.out\n*.err\n*.zip\n*.tar\n*.pth\nstdout_train.txt\nstdout_eval.txt\nstdout_visualize.txt\nhostfile\n**/olds/**\n**/old/**\n*.pt\n*.txt\nmiccai_generate_Dir\ngenerate.py\nvisualzie.py"
},
{
"path": "README.md",
"content": "\n## LGRNet: Local-Global Reciprocal Network for Video Polyp Segmentation [`Paper`](https://arxiv.org/abs/2407.05703) | [`BibTeX`](#citing) | [`Huggingface(UFUV Dataset)`](https://huggingface.co/datasets/huihuixu/uterine_fibroid_ultrasound_video_segmentation)\n\nHuihui Xu, Yijun Yang, Angelica Aviles-Rivero, Guang Yang, Jing Qin, and Lei Zhu\n\n[](https://paperswithcode.com/sota/video-polyp-segmentation-on-sun-seg-hard?p=lgrnet-local-global-reciprocal-network-for)\n\n[](https://paperswithcode.com/sota/video-polyp-segmentation-on-sun-seg-easy?p=lgrnet-local-global-reciprocal-network-for)\n\n[](https://paperswithcode.com/sota/video-polyp-segmentation-on-sun-seg-easy-1?p=lgrnet-local-global-reciprocal-network-for)\n\t\n[](https://paperswithcode.com/sota/video-polyp-segmentation-on-sun-seg-hard-1?p=lgrnet-local-global-reciprocal-network-for)\n\n\nThis is the official implmentation of LGRNet (MICCAI'24 Early Accept), which incorporates local **[Cyclic Neighborhoold Propagation](https://github.com/bio-mlhui/LGRNet/blob/main/models/encoder/neighborhood_qk.py#L57)** and global **[Hilbert Selective Scan](https://github.com/bio-mlhui/LGRNet/blob/main/models/encoder/ops/modules/frame_query_ss2d.py#L531)**. Together with the notion of **[Frame Bottleneck Queries](https://github.com/bio-mlhui/LGRNet/blob/main/models/encoder/localGlobal.py#L185)**, LGRNet can both efficiently and effectively aggregate the local-global temporal context, which achieves *state-of-the-art* on the public [Video Polyp Segmentation(VPS)](https://paperswithcode.com/task/video-polyp-segmentation) benchmark.\n\n
As an example for ultrasound video, a single frame is too noisy and insufficient for accurate lesion diagnosis. In practice, doctors need to check neighboring frames(local) and collect all visual clues (global) in the video to predict possible lesion region and filter out irrelevent surrounding issues.
\n\n
\n\n
\n\n
In CNP, each token takes the neighborhood tokens (defined by a kernel) in the cyclic frame as attention keys. CNP enables aggregating the local(cyclic) temporal information into one token. In Hilbert Selective Scan, a set of frame bottleneck queries are used to aggreate spatial information from each frame. Then, we use Hilbert Selective Scan to efficiently parse the global temporal context based on these bottleneck queries. The global temporal context is then propagated back to the feature maps by a Distribute layer. Based on Mask2Former, the decoder can output a set of different mask predictions with corresponding confidence score, which also facilitates comprehesive diagnosis.
\n\n\n## Items\n\n1. Installation: Please refer to [INSTALL.md](assets/INSTALL.md) for more details.\n2. Data preparation: Please refer to [DATA.md](assets/DATA.md) for more details.\n\n\n3. Training: \n\nChange PORT_NUM for DDP and make sure the $CURRENT_TASK is 'VIS':\n```\nexport CURRENT_TASK=VIS\nexport MASTER_ADDR=127.0.0.1\nexport MASTER_PORT=PORT_NUM\n```\n\nMake sure the $PT_PATH and $DATASET_PATH are correctly set during installation and preparing data.\n\nThe training on SUN-SEG is conducted using 2 4090-24GB GPUs:\n```\nCUDA_VISIBLE_DEVICES=0,1 TORCH_NUM_WORKERS=8 python main.py --config_file output/VIS/sunseg/pvt/pvt.py --trainer_mode train_attmpt\n```\n\n4. logs, checkpoints, predictions\n\n| Backbone| Dataset | Dice | mIou | log | ckpt | predictions |\n| :----: | :----: | :----: | :----: | :----: | :----: |:----: |\n| PVTv2-B2 | SUN-SEG-Train | -- | -- | [log](https://drive.google.com/file/d/17MTOYW73RLbvZS3BLFBZEphY_0JzN6er/view?usp=sharing) | [ckpt](https://drive.google.com/file/d/1D4YAIfFCCQIsDfKgSCr9tCw7vDAgqf76/view?usp=sharing) | --\n| PVTv2-B2 | SUN-SEG-Hard-Testing | 0.876 | 0.805 | [log](https://drive.google.com/file/d/1wdVMWMknSlURaBROWbMax4iS9V1Tbn9-/view?usp=sharing) |[ckpt](https://drive.google.com/file/d/1D4YAIfFCCQIsDfKgSCr9tCw7vDAgqf76/view?usp=sharing) | [mask predictions](https://drive.google.com/file/d/1V8CDMC87o7t4eyts4BVEwflDUrFpAOVX/view?usp=sharing)\n| PVTv2-B2 | SUN-SEG-Easy-Testing | 0.875 | 0.810 | [log](https://drive.google.com/file/d/1wdVMWMknSlURaBROWbMax4iS9V1Tbn9-/view?usp=sharing) |[ckpt](https://drive.google.com/file/d/1D4YAIfFCCQIsDfKgSCr9tCw7vDAgqf76/view?usp=sharing) | [mask predictions](https://drive.google.com/file/d/1V8CDMC87o7t4eyts4BVEwflDUrFpAOVX/view?usp=sharing)\n| PVTv2-B2 | SUN-SEG-Hard-Unseen-Testing | 0.865 | 0.792 | [log](https://drive.google.com/file/d/1obt_qvWCvslhRY-e4SrTJNS0r6Diad4e/view?usp=sharing) | [ckpt](https://drive.google.com/file/d/1D4YAIfFCCQIsDfKgSCr9tCw7vDAgqf76/view?usp=sharing) | [mask predictions](https://drive.google.com/file/d/1V8CDMC87o7t4eyts4BVEwflDUrFpAOVX/view?usp=sharing)\n| PVTv2-B2 | SUN-SEG-Easy-Unseen-Testing | 0.853 | 0.783 | [log](https://drive.google.com/file/d/1obt_qvWCvslhRY-e4SrTJNS0r6Diad4e/view?usp=sharing) | [ckpt](https://drive.google.com/file/d/1D4YAIfFCCQIsDfKgSCr9tCw7vDAgqf76/view?usp=sharing)| [mask predictions](https://drive.google.com/file/d/1V8CDMC87o7t4eyts4BVEwflDUrFpAOVX/view?usp=sharing)\n| Res2Net-50 | SUN-SEG-Hard-Testing | 0.841 | 0.765 | [log](https://drive.google.com/file/d/17pUxFMuHpPD_In5RVrJUsPFZGOgNFzb6/view?usp=sharing) |\n| Res2Net-50 | SUN-SEG-Easy-Testing | 0.843 | 0.774 | [log](https://drive.google.com/file/d/17pUxFMuHpPD_In5RVrJUsPFZGOgNFzb6/view?usp=sharing) |\n| PVTv2-B2 | CVC612V | 0.933 | 0.877 | [log](https://drive.google.com/file/d/1m36mJL0Fu3T9F73TqFGnFsWaCGh3JDeJ/view?usp=drive_link) |\n| PVTv2-B2 | CVC300TV | 0.916 | 0.852 | [log](https://drive.google.com/file/d/1m36mJL0Fu3T9F73TqFGnFsWaCGh3JDeJ/view?usp=drive_link) |\n| PVTv2-B2 | CVC612T | 0.875 | 0.814 | [log](https://drive.google.com/file/d/1m36mJL0Fu3T9F73TqFGnFsWaCGh3JDeJ/view?usp=drive_link) |\n\n\n5. Evaluate:\nEvaluating on SUN-SEG-Easy AND SUN-SEG-Hard using 1 4090-24GPU GPUS (**modify the ckpt_path to the absolute path**):\n```\nCUDA_VISIBLE_DEVICES=0 TORCH_NUM_WORKERS=8 python main.py --config_file output/VIS/sunseg/pvt/pvt.py --trainer_mode eval --eval_path ckpt_path\n```\n\n## citing\n```\n@article{xu2024lgrnet,\n title={LGRNet: Local-Global Reciprocal Network for Uterine Fibroid Segmentation in Ultrasound Videos},\n author={Xu, Huihui and Yang, Yijun and Aviles-Rivero, Angelica I and Yang, Guang and Qin, Jing and Zhu, Lei},\n journal={arXiv preprint arXiv:2407.05703},\n year={2024}\n}\n``` \n\n## Acknowledgments\n- Thanks [Gilbert](https://github.com/jakubcerveny/gilbert) for the implementation of Hilbert curve generation.\n- Thanks GPT4 for helping me constructing idea of Hilbert Filling Curve v.s. Zigzag curve"
},
{
"path": "assets/DATA.md",
"content": "\n\n\n\n\n\n# Data Preparation\n\n## UFUV (Private): \nplease email the second author for UFUV dataset if you want, I have no absolute power for UFUV\n\n## VPS (Public)\n\n### CVC/Kvasir/Mayo\nWe follow [PNS-Net](https://github.com/GewelsJI/PNS-Net) to download the CVC/Kvasir/Mayo dataset. The download link is same as [link](https://drive.google.com/file/d/1TyaRy4c4nHFDa3o2bOl4dP5Z7wes7HV2/view?usp=sharing)\n\nPut MICCAI-VPS-dataset.zip in $DATASET_PATH, then run following script to change the directory structure:\n```\ncd $DATASET_PATH\nunzip -qq MICCAI-VPS-dataset.zip\n\n# cd LGRNet directory\n# normalize the VPS data structure\npython handle_vps.py\n```\n\nNow the structure should be like:\n```\n${DATASET_PATH}\n -- MICCAI-VPS-dataset\n -- Kvasir-SEG\n -- *\n -- VPS-TestSet\n -- CVC-ColonDB-300\n -- *\n -- CVC-ClinicDB-612-Valid\n -- *\t\t\n -- CVC-ClinicDB-612-Test\n -- *\n -- VPS-TrainSet\n -- ASU-Mayo_Clinic\n -- Train\n -- *\n -- CVC-ClinicDB-612\t\n -- Train\n -- *\n -- CVC-ColonDB-300\n -- Train\n -- *\n```\nwhere * means the following structure:\n```\n-- Frame\n -- vid1\n -- img file\n-- GT\n -- vid1\n -- mask file\n```\n\n### SUN-SEG\n\nPlease follow https://github.com/GewelsJI/VPS/blob/main/docs/DATA_PREPARATION.md to email the author for SUN-SEG.\nPut part1, part2, annotation in $DATASET_PATH/SUN-SEG\n\n```\n# normalize the directory\nunzip -qq $DATASET_PATH/SUN-SEG/sundatabase_positive_part1.zip -d $DATASET_PATH/SUN-SEG/SUN-Positive\nunzip -qq $DATASET_PATH/SUN-SEG/sundatabase_positive_part2.zip -d $DATASET_PATH/SUN-SEG/SUN-Positive\ntar -xf $DATASET_PATH/SUN-SEG/SUN-SEG-Annotation.tar -C $DATASET_PATH/SUN-SEG/\n\nrm -rf $DATASET_PATH/SUN-SEG/SUN-SEG-Annotation/TestEasyDataset/Unseen/Frame\nfind $DATASET_PATH/SUN-SEG/SUN-SEG-Annotation -name \"._.DS_Store\" -type f -delete\npython reorganize_sunseg.py\n\n```\nNow the structure should be like:\n```\n${DATASET_PATH}\n -- SUN-SEG\n -- SUN-SEG-Annotation\n -- TrainDataset\n -- *\n -- TestEasyDataset\n -- combine\n -- *\n -- TestHardDataset\n -- combine\n -- *\n```\n\n\n\n"
},
{
"path": "assets/INSTALL.md",
"content": "# Install\n## Requirements\nWe test the codes in the following environments\n\n- CUDA 12.1\n- Python 3.10.13\n- Pytorch 2.1.1\n- Torchvison 0.16.1\n- detectron 0.6\n- mamba_ssm 1.2.0.post1\n- natten 0.15.1\n- timm 0.9.12\n\n## Install environment for LGRNet\n\n```\nconda create --name lgrnet python=3.10\nconda activate lgrnet\n\n# make sure CUDA-12.1 is installed and activated in env var.\n\n# install torch\npip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu121\n\n\n# install detectron2, building may take much time.\npython -m pip install 'git+https://github.com/facebookresearch/detectron2.git'\n\n# install mamba, see https://github.com/state-spaces/mamba\ncd .. \ngit clone https://github.com/state-spaces/mamba.git\ncd mamba\npip install . --no-build-isolation\ncd ../LGRNet\n\n# install natten, see https://github.com/SHI-Labs/NATTEN/blob/main/docs/install.md\npip install natten==0.15.1+torch210cu121 -f https://shi-labs.com/natten/wheels/cu121/torch2.1.0/natten-0.15.1%2Btorch210cu121-cp310-cp310-linux_x86_64.whl\n\n# misc\npip install albumentations==1.3.1\npip install Pygments\npip install imgaug\npip install timm==0.9.12\n\n# compile deform attention\ncd models/encoder/ops/\npython setup.py build install --user\n\n# download resnet/pvtv2 ckpt, our model uses the same backbone with WeakPoly(phttps://github.com/weijun88/WeakPolyp)\nwget -P $PT_PATH/pvt_v2/pvt_v2_b2.pth https://huggingface.co/huihuixu/lgrnet_ckpts/blob/main/pvt_v2_b2.pth \nwget -P $PT_PATH/res2net/res2net50_v1b_26w_4s-3cf99910.pth https://huggingface.co/huihuixu/lgrnet_ckpts/blob/main/res2net50_v1b_26w_4s-3cf99910.pth\n\n```\n"
},
{
"path": "assets/MODEL_ZOO.md",
"content": ""
},
{
"path": "data_schedule/__init__.py",
"content": "\nimport os\n\nif os.getenv('CURRENT_TASK') == 'VIS':\n from . import vis\nelse:\n raise ValueError()\n\ndef build_schedule(configs, model_input_mapper, model_input_collate_fn):\n import logging\n from functools import partial\n import detectron2.utils.comm as comm\n from torch.utils.data import DataLoader, ConcatDataset\n from .registry import MAPPER_REGISTRY, EVALUATOR_REGISTRY\n from detectron2.data import DatasetCatalog, DatasetFromList, MapDataset, MetadataCatalog\n from data_schedule.utils.sampler import Evaluate_ExactSampler_Distributed, Train_InfiniteSampler_Distributed\n datasets = {'train': [], 'evaluate': []}\n meta_idx_shift = 0 \n for mode in ['train', 'evaluate']:\n for dataset_name in configs['data'][mode].keys():\n dataset_assume_mode = MetadataCatalog.get(dataset_name).get('mode')\n if dataset_assume_mode != mode:\n logging.warning(f'default mode of {dataset_name} is {dataset_assume_mode} not {mode}')\n dataset_dicts = DatasetFromList(DatasetCatalog.get(dataset_name), copy=True, serialize=True)\n mapper = MAPPER_REGISTRY.get(configs['data'][mode][dataset_name]['mapper']['name'])(mode=mode,\n dataset_name=dataset_name, \n configs=configs,\n meta_idx_shift=meta_idx_shift if mode == 'train' else 0)\n meta_idx_shift += len(dataset_dicts)\n dataset = MapDataset(dataset_dicts, partial(composition, mappers=[mapper, \n partial(model_input_mapper, mode=mode)]))\n if mode == 'train':\n datasets[mode].append(dataset)\n else:\n datasets[mode].append((dataset_name, dataset))\n\n train_dataset = ConcatDataset(datasets['train'])\n logging.debug(f'Total number of training meta: {len(train_dataset)}')\n\n train_loader_splits = configs['optim']['splits']\n batch_sizes = configs['optim']['batch_sizes']\n splits = list(zip(train_loader_splits[:-1], train_loader_splits[1:]))\n assert len(splits) == (len(batch_sizes))\n inf_stream_fn = partial(infinite_indices,\n seed=configs['stream_idx_seed'], \n batch_sizes=configs['optim']['batch_sizes'],\n splits=configs['optim']['splits'],\n one_batch_two_epoch=configs['optim']['one_batch_two_epoch'],\n dataset_length=len(train_dataset),\n shuffle=True) \n train_samplers = []\n train_loaders = []\n for btch_size, (range_start, range_end) in zip(batch_sizes, splits):\n if range_end is not None:\n assert (range_end - range_start) % btch_size == 0, ''\n assert btch_size % comm.get_world_size() == 0, ''\n each_process_batch_size = int(btch_size / comm.get_world_size())\n loader_sampler = Train_InfiniteSampler_Distributed(inf_stream_fn=inf_stream_fn,\n start_idx=range_start,\n end_idx=range_end,)\n train_samplers.append(loader_sampler)\n train_loaders.append(DataLoader(train_dataset,\n batch_size=each_process_batch_size,\n sampler=loader_sampler,\n collate_fn=partial(model_input_collate_fn, mode='train'), \n num_workers=int(os.getenv('TORCH_NUM_WORKERS')),\n pin_memory=True,\n persistent_workers=True))\n\n evaluators = []\n for eval_dataset_name, eval_dataset in datasets['evaluate']:\n logging.debug(f'Number of evaluate meta in {eval_dataset_name}: {len(eval_dataset)}')\n loader = DataLoader(eval_dataset, \n batch_size=1, \n sampler=Evaluate_ExactSampler_Distributed(eval_dataset),\n collate_fn=partial(model_input_collate_fn, mode='evaluate'),\n num_workers=int(os.getenv('TORCH_NUM_WORKERS')),\n pin_memory=True,\n persistent_workers=True)\n \n evaluator = EVALUATOR_REGISTRY.get(configs['data']['evaluate'][eval_dataset_name]['evaluator']['name'])(configs=configs,\n dataset_name=eval_dataset_name,\n data_loader=loader)\n evaluators.append((eval_dataset_name, evaluator))\n\n return train_samplers, train_loaders, partial(evaluate_call, evaluators=evaluators)\n\ndef composition(data_dict, mappers):\n for mappper in mappers:\n data_dict = mappper(data_dict)\n if data_dict is None:\n return None\n return data_dict\n\ndef evaluate_call(evaluators, model, output_dir):\n import detectron2.utils.comm as comm\n ret = {}\n for eval_dataset_name, evaluator in evaluators:\n metric_dict = evaluator(model=model,output_dir=output_dir)\n if comm.is_main_process():\n for key, value in metric_dict.items():\n assert f'{key}_{eval_dataset_name}' not in ret\n ret[f'{key}_{eval_dataset_name}'] = value\n comm.synchronize()\n return ret\n\n\ndef _infinite_indices(seed, dataset_length, shuffle=True,):\n import torch\n g = torch.Generator()\n g.manual_seed(seed)\n while True:\n if shuffle:\n yield from torch.randperm(dataset_length, generator=g).tolist()\n else:\n yield from torch.arange(dataset_length).tolist()\n\ndef infinite_indices(seed, \n dataset_length, \n batch_sizes, \n splits, \n one_batch_two_epoch='just_use',\n shuffle=True): # 'abandon', 'just_use', 'pad'\n import torch\n import math\n g = torch.Generator()\n g.manual_seed(seed)\n\n split_ranges = list(zip(splits[:-1], splits[1:]))\n assert len(split_ranges) == (len(batch_sizes))\n stream = _infinite_indices(seed, dataset_length=dataset_length, shuffle=shuffle)\n\n stream_throw_cnt = 0\n cnt = 0\n for (range_start, range_end), btch_size in zip(split_ranges, batch_sizes):\n assert cnt == range_start\n if range_end == None:\n range_end = math.inf\n \n while cnt < range_end:\n epoch_milestone = ((stream_throw_cnt // dataset_length) + 1 ) * dataset_length\n if (stream_throw_cnt < epoch_milestone) and (stream_throw_cnt + btch_size > epoch_milestone) and (one_batch_two_epoch != 'just_use'):\n if one_batch_two_epoch == 'abandon':\n for _ in range(epoch_milestone - stream_throw_cnt):\n abandon = next(stream)\n stream_throw_cnt += 1\n\n elif one_batch_two_epoch == 'pad':\n diff = stream_throw_cnt + btch_size - epoch_milestone\n num_throw = btch_size - diff\n rand_idxs = torch.randperm(dataset_length, generator=g)[:diff].tolist()\n for _ in range(num_throw):\n cnt += 1\n stream_throw_cnt += 1\n yield next(stream)\n for idx in rand_idxs:\n cnt += 1\n yield idx\n else:\n raise ValueError()\n else:\n for _ in range(btch_size):\n cnt += 1\n stream_throw_cnt += 1\n yield next(stream) \n \n assert cnt == range_end\n"
},
{
"path": "data_schedule/registry.py",
"content": "from detectron2.utils.registry import Registry\n\n\nEVALUATOR_REGISTRY = Registry('EVALUATOR')\nMAPPER_REGISTRY = Registry('MAPPER')\n\nclass Mapper:\n def __init__(self, \n meta_idx_shift,\n dataset_meta,) -> None:\n self.meta_idx_shift = meta_idx_shift \n self.visualized_meta_idxs = dataset_meta.get('visualize_meta_idxs') \n\n def _call(self, data_dict):\n pass\n\n def __call__(self, data_dict):\n meta_idx = data_dict['meta_idx']\n ret = self._call(data_dict)\n if ret is None:\n return None\n ret['meta_idx'] = meta_idx + self.meta_idx_shift\n if meta_idx in self.visualized_meta_idxs:\n ret['visualize'] = True\n else:\n ret['visualize'] = False\n return ret\n \n\n\n\n\n\n"
},
{
"path": "data_schedule/utils/box_ops.py",
"content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n\"\"\"\nUtilities for bounding box manipulation and GIoU.\n\"\"\"\nimport torch\nfrom torchvision.ops.boxes import box_area\n\n\ndef box_cxcywh_to_xyxy(x):\n x_c, y_c, w, h = x.unbind(-1)\n b = [(x_c - 0.5 * w), (y_c - 0.5 * h),\n (x_c + 0.5 * w), (y_c + 0.5 * h)]\n return torch.stack(b, dim=-1)\n\n\ndef box_xyxy_to_cxcywh(x):\n x0, y0, x1, y1 = x.unbind(-1)\n assert ((x1 - x0) >= 0).all()\n assert ((y1 - y0) >= 0).all()\n b = [(x0 + x1) / 2, (y0 + y1) / 2,\n (x1 - x0), (y1 - y0)]\n return torch.stack(b, dim=-1)\n\n\n# modified from torchvision to also return the union\ndef box_iou(boxes1, boxes2):\n area1 = box_area(boxes1)\n area2 = box_area(boxes2)\n\n lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]\n rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]\n\n wh = (rb - lt).clamp(min=0) # [N,M,2]\n inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]\n\n union = area1[:, None] + area2 - inter\n\n iou = inter / union\n return iou, union\n\n\ndef generalized_box_iou(boxes1, boxes2):\n \"\"\"\n Generalized IoU from https://giou.stanford.edu/\n\n The boxes should be in [x0, y0, x1, y1] format\n\n Returns a [N, M] pairwise matrix, where N = len(boxes1)\n and M = len(boxes2)\n \"\"\"\n # degenerate boxes gives inf / nan results\n # so do an early check\n assert (boxes1[:, 2:] >= boxes1[:, :2]).all()\n assert (boxes2[:, 2:] >= boxes2[:, :2]).all()\n iou, union = box_iou(boxes1, boxes2)\n\n lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])\n rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])\n\n wh = (rb - lt).clamp(min=0) # [N,M,2]\n area = wh[:, :, 0] * wh[:, :, 1]\n\n return iou - (area - union) / area\n\n\ndef masks_to_boxes(masks):\n \"\"\"Compute the bounding boxes around the provided masks\n\n The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.\n\n Returns a [N, 4] tensors, with the boxes in xyxy format\n \"\"\"\n if masks.numel() == 0:\n return torch.zeros((0, 4), device=masks.device)\n\n h, w = masks.shape[-2:]\n\n y = torch.arange(0, h, dtype=torch.float)\n x = torch.arange(0, w, dtype=torch.float)\n y, x = torch.meshgrid(y, x)\n\n x_mask = (masks * x.unsqueeze(0))\n x_max = x_mask.flatten(1).max(-1)[0]\n x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]\n\n y_mask = (masks * y.unsqueeze(0))\n y_max = y_mask.flatten(1).max(-1)[0]\n y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]\n\n return torch.stack([x_min, y_min, x_max, y_max], 1)\n"
},
{
"path": "data_schedule/utils/sampler.py",
"content": "\nimport math\nimport torch.distributed as dist\nfrom typing import TypeVar, Optional, Iterator \nT_co = TypeVar('T_co', covariant=True)\nfrom torch.utils.data.distributed import DistributedSampler\nimport detectron2.utils.comm as comm\nfrom utils.misc import all_gather\nfrom torch.utils.data import Sampler\nimport torch\n\nimport logging\nimport itertools\n\nclass TrainRandomSampler_ByEpoch(Sampler[int]):\n def __init__(self, \n data_source,\n seed,\n ) -> None:\n self.data_source = data_source\n self.num_samples = len(self.data_source)\n self.seed = seed\n self.epoch = None\n\n def __iter__(self):\n seed = self.seed + self.epoch\n print(f'generating a new indices permutations for this epoch using seed {seed}')\n n = len(self.data_source)\n g = torch.Generator()\n g.manual_seed(seed)\n \n for _ in range(self.num_samples // n):\n yield from torch.randperm(n, generator=g).tolist()\n yield from torch.randperm(n, generator=g).tolist()[:self.num_samples % n]\n\n def __len__(self) -> int:\n return self.num_samples\n\n def set_epoch(self, epoch: int) -> None:\n r\"\"\"\n\n Args:\n epoch (int): Epoch number.\n \"\"\"\n self.epoch = epoch\n \n \n\nclass Train_InfiniteSampler_Distributed(Sampler[T_co]):\n def __init__(self, \n inf_stream_fn, \n start_idx: int = 0,\n end_idx = None,\n \n ): \n self.rank = comm.get_rank()\n self.num_replicas = comm.get_world_size()\n self.start_idx = start_idx \n self.end_idx = end_idx\n self.inf_stream_fn = inf_stream_fn\n\n def set_iter_first_sample_idx(self, idx):\n self.start_idx = idx\n\n def set_iter_last_sample_idx(self, idx):\n self.end_idx = idx\n\n def __iter__(self) -> Iterator[T_co]:\n logging.debug(f'在 infinite stream 上定位到{self.start_idx} 为开头')\n yield from itertools.islice(self.inf_stream_fn(), self.start_idx + self.rank, self.end_idx, self.num_replicas)\n\nclass Evaluate_ExactSampler_Distributed(Sampler[T_co]):\n def __init__(self, dataset) -> None:\n self.dataset = dataset\n self.rank = comm.get_rank()\n self.num_replicas = comm.get_world_size()\n indices = list(range(len(self.dataset))) \n self.indices = indices[self.rank:len(self.dataset):self.num_replicas]\n\n\n def __iter__(self):\n yield from self.indices\n\n def __len__(self):\n return len(self.indices)\n\n\nclass TrainRandomSampler_ByEpoch_Distributed(Sampler[T_co]):\n def __init__(self, \n dataset, num_replicas,\n rank,\n seed: int = 0) -> None:\n if rank >= num_replicas or rank < 0:\n raise ValueError(\"Invalid rank {}, rank should be in the interval\"\" [0, {}]\".format(rank, num_replicas - 1))\n self.dataset = dataset\n self.num_replicas = num_replicas\n self.rank = rank\n self.epoch = None\n\n self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) \n self.total_size = self.num_samples * self.num_replicas\n self.seed = seed\n\n def __iter__(self) -> Iterator[T_co]:\n seed = self.seed + self.epoch\n logging.debug(f'generating a new indices permutations for this epoch using seed {seed}')\n g = torch.Generator()\n g.manual_seed(self.seed + self.epoch)\n indices = torch.randperm(len(self.dataset), generator=g).tolist()\n\n # add extra samples to make it evenly divisible\n padding_size = self.total_size - len(indices)\n if padding_size <= len(indices):\n indices += indices[:padding_size]\n else:\n indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]\n assert len(indices) == self.total_size\n # subsample\n indices = indices[self.rank:self.total_size:self.num_replicas]\n assert len(indices) == self.num_samples\n self.epoch = None \n return iter(indices)\n\n def __len__(self) -> int:\n return self.num_samples\n\n def set_epoch(self, epoch: int) -> None:\n r\"\"\"\n Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas\n use a different random ordering for each epoch. Otherwise, the next iteration of this\n sampler will yield the same ordering.\n\n Args:\n epoch (int): Epoch number.\n \"\"\"\n self.epoch = epoch\n\n\nclass InferenceSampler(Sampler):\n \"\"\"\n Produce indices for inference across all workers.\n Inference needs to run on the __exact__ set of samples,\n therefore when the total number of samples is not divisible by the number of workers,\n this sampler produces different number of samples on different workers.\n \"\"\"\n\n def __init__(self, size: int):\n \"\"\"\n Args:\n size (int): the total number of data of the underlying dataset to sample from\n \"\"\"\n self._size = size\n assert size > 0\n self._rank = comm.get_rank()\n self._world_size = comm.get_world_size()\n self._local_indices = self._get_local_indices(size, self._world_size, self._rank)\n\n @staticmethod\n def _get_local_indices(total_size, world_size, rank):\n shard_size = total_size // world_size\n left = total_size % world_size\n shard_sizes = [shard_size + int(r < left) for r in range(world_size)]\n\n begin = sum(shard_sizes[:rank])\n end = min(sum(shard_sizes[: rank + 1]), total_size)\n return range(begin, end)\n\n def __iter__(self):\n yield from self._local_indices\n\n def __len__(self):\n return len(self._local_indices)"
},
{
"path": "data_schedule/utils/segmentation.py",
"content": "\nimport torch\n\ndef bounding_box_from_mask(mask):\n if not mask.any():\n return torch.zeros([4]).float()\n rows = torch.any(mask, dim=1) # h\n cols = torch.any(mask, dim=0) # w\n row_indexs = torch.where(rows)[0]\n rmin, rmax = row_indexs.min(), row_indexs.max()\n\n col_indexs = torch.where(cols)[0]\n cmin, cmax = col_indexs.min(), col_indexs.max()\n return torch.tensor([cmin, rmin, cmax, rmax]).float() # x1y1x2y2\n"
},
{
"path": "data_schedule/vis/__init__.py",
"content": "from . import polyp\nfrom . import mapper \nfrom . import evaluator_fast\nfrom . import vis_aug_eval \nfrom . import vis_aug_train\nfrom . import vis_frame_sampler\n\n\n\n\n"
},
{
"path": "data_schedule/vis/apis.py",
"content": "\nclass VIS_Dataset:\n \"\"\"\n \"\"\"\n\nclass VIS_Aug_CallbackAPI:\n \"\"\"\n \"\"\"\n\nclass VIS_Evaluator_OutAPI_EvalFn_API:\n \"\"\"\n \"\"\"\n\nclass VIS_TrainAPI_clipped_video:\n \"\"\"\n \"\"\"\n\nclass VIS_EvalAPI_clipped_video_request_ann:\n \"\"\"\n \"\"\"\n\nclass VIS_FrameSampler_InputOutput_API:\n \"\"\"\n \"\"\"\n\nclass GetFrames:\n \"\"\"\n \"\"\""
},
{
"path": "data_schedule/vis/evaluator_fast.py",
"content": "\nimport os\nfrom tqdm import tqdm\nfrom functools import partial\nimport torch\nimport detectron2.utils.comm as comm\nfrom utils.misc import to_device\nfrom detectron2.data import MetadataCatalog\nfrom data_schedule.registry import EVALUATOR_REGISTRY\n\nfrom .evaluator_utils import vis_metric_entrypoint\nfrom data_schedule.vis.apis import VIS_EvalAPI_clipped_video_request_ann, VIS_Aug_CallbackAPI, VIS_Evaluator_OutAPI_EvalFn_API\nfrom collections import defaultdict\n\n\n@EVALUATOR_REGISTRY.register()\nclass VIS_Evaluator_FrameFast:\n def __init__(self,\n dataset_name,\n data_loader,\n configs) -> None:\n self.dataset_name = dataset_name\n self.loader = data_loader\n frame_metrics = configs['data']['evaluate'][dataset_name]['evaluator']['frame_metrics']\n dataset_meta = MetadataCatalog.get(dataset_name)\n self.frame_metric_fns = []\n for metric_name, metric_config in frame_metrics:\n metric_fn = vis_metric_entrypoint(metric_name)\n metric_fn = partial(metric_fn, dataset_meta=dataset_meta, **metric_config)\n self.frame_metric_fns.append(metric_fn)\n \n self.eval_meta_keys = dataset_meta.get('eval_meta_keys') \n\n metrics_aggregator = configs['data']['evaluate'][dataset_name]['evaluator']['metrics_aggregator']\n self.eval_meta_keys = dataset_meta.get('eval_meta_keys') # { video_id: list[fnames] }\n self.metrics_aggregator = partial(vis_metric_entrypoint(metrics_aggregator[0]),\n dataset_meta=dataset_meta,\n eval_meta_keys=self.eval_meta_keys,\n **metrics_aggregator[1])\n\n def visualize_path(self, meta_idxs, visualize, evaluator_path):\n return [os.path.join(evaluator_path, f'meta_{meta_idx}') if vis else None for (meta_idx, vis) in zip(meta_idxs, visualize)]\n \n @torch.no_grad()\n def __call__(self, model, output_dir):\n evaluator_path = os.path.join(output_dir, f'eval_{self.dataset_name}')\n os.makedirs(evaluator_path, exist_ok=True)\n macs, params = None, None\n metrics_by_video_id_frame = defaultdict(dict) \n for batch_dict in tqdm(self.loader):\n VIS_EvalAPI_clipped_video_request_ann\n eval_metas = batch_dict.pop('metas')\n request_anns = eval_metas['request_ann'][0] # t, bool tensor\n frame_strs = eval_metas['frames'][0] # t', list[str]\n video_id = eval_metas['video_id'][0] # str\n assert request_anns.int().sum() == len(frame_strs)\n callback_fns = eval_metas['callback_fns'][0] # list[fn]\n visualize_path = self.visualize_path(meta_idxs=batch_dict['meta_idxs'], visualize=batch_dict['visualize'], \n evaluator_path=os.path.join(evaluator_path, 'visualize_model'))\n batch_dict['visualize_paths'] = visualize_path\n batch_dict = to_device(batch_dict, device=model.device)\n VIS_Aug_CallbackAPI\n # if macs is None:\n # from detectron2.utils.analysis import (\n # FlopCountAnalysis,\n # )\n # flops = FlopCountAnalysis(model, batch_dict, inference_func=lambda model, *inputs: model.sample(*inputs))\n # total_flops = flops.total()\n # # counts = flops.by_operator()\n # logging.debug(f'macs: {total_flops/ (10**9) / len(request_anns)}')\n model_outputs = model.sample(batch_dict) \n predictions = {\n 'video': model_outputs['video'][0], # t 3 h w \n 'pred_masks': [haosen for idx, haosen in enumerate(model_outputs['pred_masks'][0]) if request_anns[idx]], # list[nt h w], t'\n 'pred_class': [haosen for idx, haosen in enumerate(model_outputs['pred_class'][0]) if request_anns[idx]], # list[nt c], t', \n } \n if 'pred_boxes' in model_outputs: \n predictions.update({'pred_boxes': [haosen for idx, haosen in enumerate(model_outputs['pred_boxes'][0]) if request_anns[idx]]}) # # list[nt 4], t,\n for cardib in callback_fns:\n predictions = cardib(predictions) \n pred_masks = predictions['pred_masks']\n pred_class = predictions['pred_class']\n assert len(frame_strs) == len(pred_masks)\n\n for idx, (fname, fmk, fclass) in enumerate(zip(frame_strs, pred_masks, pred_class)):\n VIS_Evaluator_OutAPI_EvalFn_API\n frame_pred = {'masks': fmk, 'classes': fclass.tolist(), 'video_id': video_id, 'frame_name': fname}\n if 'pred_boxes' in predictions:\n frame_pred.update({'boxes': predictions['pred_boxes'][idx]})\n\n meta_key_metrics = {} \n for metric_fn in self.frame_metric_fns:\n metric_values = metric_fn(frame_pred=frame_pred, output_dir=evaluator_path)\n for key, value in metric_values.items():\n assert key not in meta_key_metrics\n meta_key_metrics[key] = value\n\n assert fname not in metrics_by_video_id_frame[video_id]\n metrics_by_video_id_frame[video_id][fname] = meta_key_metrics\n\n metrics_by_video_id_frame = comm.gather(dict(metrics_by_video_id_frame), dst=0)\n eval_metrics = {}\n if comm.is_main_process():\n metrics_by_video = {} \n for video_id in tqdm(self.eval_meta_keys.keys(), desc='gathering different processes'):\n video_id_metrics = [haosen[video_id] for haosen in metrics_by_video_id_frame if video_id in haosen]\n\n video_id_frame_names = [list(haosen.keys()) for haosen in video_id_metrics]\n merged_video_id_frame_names = [item for sublist in video_id_frame_names for item in sublist]\n assert len(set(merged_video_id_frame_names)) == len(merged_video_id_frame_names),''\n assert set(merged_video_id_frame_names).issubset(set(self.eval_meta_keys[video_id]))\n assert set(self.eval_meta_keys[video_id]).issubset(set(merged_video_id_frame_names))\n\n # perframe metrics frame: predictions\n vid_frame_metrics = video_id_metrics[0]\n for haosen in video_id_metrics[1:]:\n vid_frame_metrics.update(haosen)\n metrics_by_video[video_id] = vid_frame_metrics\n\n eval_metrics = self.metrics_aggregator(metrics_by_video)\n comm.synchronize() \n return eval_metrics\n\n"
},
{
"path": "data_schedule/vis/evaluator_utils.py",
"content": "_vis_metric_entrypoints = {}\n\ndef register_vis_metric(fn):\n vis_metric_name = fn.__name__\n if vis_metric_name in _vis_metric_entrypoints:\n raise ValueError(f'vis_metric name {vis_metric_name} has been registered')\n _vis_metric_entrypoints[vis_metric_name] = fn\n\n return fn\n\ndef vis_metric_entrypoint(vis_metric_name):\n try:\n return _vis_metric_entrypoints[vis_metric_name]\n except KeyError as e:\n print(f'vis_metric Name {vis_metric_name} not found')\n\nimport numpy as np\n_EPS = np.spacing(1)\n_TYPE = np.float64\n\ndef _prepare_data(pred: np.ndarray, gt: np.ndarray) -> tuple:\n gt = gt > 128\n pred = pred / 255\n if pred.max() != pred.min():\n pred = (pred - pred.min()) / (pred.max() - pred.min())\n return pred, gt\n\n\n\n\nclass Smeasure(object):\n def __init__(self, length, alpha: float = 0.5):\n self.sms = []\n self.alpha = alpha\n\n def step(self, pred: np.ndarray, gt: np.ndarray, idx):\n pred, gt = _prepare_data(pred=pred, gt=gt)\n\n sm = self.cal_sm(pred, gt)\n self.sms.append(sm)\n\n def cal_sm(self, pred: np.ndarray, gt: np.ndarray) -> float:\n y = np.mean(gt)\n if y == 0:\n sm = 1 - np.mean(pred)\n elif y == 1:\n sm = np.mean(pred)\n else:\n sm = self.alpha * self.object(pred, gt) + (1 - self.alpha) * self.region(pred, gt)\n sm = max(0, sm)\n return sm\n\n def object(self, pred: np.ndarray, gt: np.ndarray) -> float:\n fg = pred * gt\n bg = (1 - pred) * (1 - gt)\n u = np.mean(gt)\n object_score = u * self.s_object(fg, gt) + (1 - u) * self.s_object(bg, 1 - gt)\n return object_score\n\n def s_object(self, pred: np.ndarray, gt: np.ndarray) -> float:\n x = np.mean(pred[gt == 1])\n sigma_x = np.std(pred[gt == 1], ddof=1)\n score = 2 * x / (np.power(x, 2) + 1 + sigma_x + _EPS)\n return score\n\n def region(self, pred: np.ndarray, gt: np.ndarray) -> float:\n x, y = self.centroid(gt)\n part_info = self.divide_with_xy(pred, gt, x, y)\n w1, w2, w3, w4 = part_info['weight']\n pred1, pred2, pred3, pred4 = part_info['pred']\n gt1, gt2, gt3, gt4 = part_info['gt']\n score1 = self.ssim(pred1, gt1)\n score2 = self.ssim(pred2, gt2)\n score3 = self.ssim(pred3, gt3)\n score4 = self.ssim(pred4, gt4)\n\n return w1 * score1 + w2 * score2 + w3 * score3 + w4 * score4\n\n def centroid(self, matrix: np.ndarray) -> tuple:\n \"\"\"\n To ensure consistency with the matlab code, one is added to the centroid coordinate,\n so there is no need to use the redundant addition operation when dividing the region later,\n because the sequence generated by ``1:X`` in matlab will contain ``X``.\n :param matrix: a bool data array\n :return: the centroid coordinate\n \"\"\"\n h, w = matrix.shape\n area_object = np.count_nonzero(matrix)\n if area_object == 0:\n x = np.round(w / 2)\n y = np.round(h / 2)\n else:\n # More details can be found at: https://www.yuque.com/lart/blog/gpbigm\n y, x = np.argwhere(matrix).mean(axis=0).round()\n return int(x) + 1, int(y) + 1\n\n def divide_with_xy(self, pred: np.ndarray, gt: np.ndarray, x, y) -> dict:\n h, w = gt.shape\n area = h * w\n\n gt_LT = gt[0:y, 0:x]\n gt_RT = gt[0:y, x:w]\n gt_LB = gt[y:h, 0:x]\n gt_RB = gt[y:h, x:w]\n\n pred_LT = pred[0:y, 0:x]\n pred_RT = pred[0:y, x:w]\n pred_LB = pred[y:h, 0:x]\n pred_RB = pred[y:h, x:w]\n\n w1 = x * y / area\n w2 = y * (w - x) / area\n w3 = (h - y) * x / area\n w4 = 1 - w1 - w2 - w3\n\n return dict(gt=(gt_LT, gt_RT, gt_LB, gt_RB),\n pred=(pred_LT, pred_RT, pred_LB, pred_RB),\n weight=(w1, w2, w3, w4))\n\n def ssim(self, pred: np.ndarray, gt: np.ndarray) -> float:\n h, w = pred.shape\n N = h * w\n\n x = np.mean(pred)\n y = np.mean(gt)\n\n sigma_x = np.sum((pred - x) ** 2) / (N - 1)\n sigma_y = np.sum((gt - y) ** 2) / (N - 1)\n sigma_xy = np.sum((pred - x) * (gt - y)) / (N - 1)\n\n alpha = 4 * x * y * sigma_xy\n beta = (x ** 2 + y ** 2) * (sigma_x + sigma_y)\n\n if alpha != 0:\n score = alpha / (beta + _EPS)\n elif alpha == 0 and beta == 0:\n score = 1\n else:\n score = 0\n return score\n\n def get_results(self):\n sm = np.mean(np.array(self.sms, dtype=_TYPE))\n return dict(Smeasure=sm)\n\n\nimport torch \nimport os\nfrom PIL import Image\n\n@register_vis_metric\ndef mask_dice_iou(frame_pred, dataset_meta, **kwargs):\n video_id = frame_pred['video_id']\n frame_name = frame_pred['frame_name']\n masks = frame_pred['masks'] # nq h w\n get_frames_gt_mask_fn = dataset_meta.get('get_frames_gt_mask_fn')\n scores = torch.tensor(frame_pred['classes']) # nq c\n foreground_scores = scores[:, :-1].sum(-1) # nq\n max_idx = foreground_scores.argmax()\n pred_mask = masks[max_idx].int() # h w\n\n gt_mask, _ = get_frames_gt_mask_fn(video_id=video_id, frames=[frame_name]) # 1 h w\n gt_mask = gt_mask[0].int() # h w\n\n inter, union = (pred_mask*gt_mask).sum(), (pred_mask+gt_mask).sum()\n dice = (2*inter+1)/(union+1)\n iou = (inter+1)/(union-inter+1)\n\n return {'dice': dice, 'iou': iou}\n\n\n@register_vis_metric\ndef mask_dice_iou_sen_mae_smeasure(frame_pred, dataset_meta, **kwargs):\n video_id = frame_pred['video_id']\n frame_name = frame_pred['frame_name']\n masks = frame_pred['masks'] # nq h w\n get_frames_gt_mask_fn = dataset_meta.get('get_frames_gt_mask_fn')\n scores = torch.tensor(frame_pred['classes']) # nq c\n foreground_scores = scores[:, :-1].sum(-1) # nq\n max_idx = foreground_scores.argmax()\n pred_mask = masks[max_idx].int() # h w\n\n gt_mask, _ = get_frames_gt_mask_fn(video_id=video_id, frames=[frame_name]) # 1 h w\n gt_mask = gt_mask[0].int() # h w\n\n # tp, tp*2 + fp + fn\n inter, union = (pred_mask*gt_mask).sum(), (pred_mask+gt_mask).sum()\n dice = (2*inter+1)/(union+1) # 2*tp / tp + tp + fp + fn\n iou = (inter+1)/(union-inter+1) # tp / tp + fp + fn\n\n\n tp = (pred_mask * gt_mask).sum().float()\n fp = (pred_mask.sum() - tp).float()\n fn = (gt_mask.sum() - tp).float()\n tn = (pred_mask.shape[0] * pred_mask.shape[1] - (tp + fp + fn)).float()\n their_dice = tp * 2 / (tp + fp + fn + tp)\n their_iou = tp / (tp + fp + fn)\n # their_spe = tn / (tn + fp)\n their_sen = tp / (tp + fn)\n their_mae = (pred_mask.float() - gt_mask.float()).abs().mean()\n \n Np = gt_mask.sum()\n Nn = gt_mask.shape[0] * gt_mask.shape[1] - Np\n \n null = Smeasure(length=1, alpha=0.5)\n null.step(pred=(pred_mask.float() * 255 ).numpy(), gt=(gt_mask.float() * 255).numpy(), idx=None)\n their_smeasure = torch.tensor(null.get_results()['Smeasure']).float()\n return {'dice': dice, 'iou': iou,\n 'their_dice': their_dice,\n 'their_iou': their_iou,\n 'their_sen': their_sen,\n 'their_mae_abs': their_mae,\n 'their_smeasure': their_smeasure,\n \n 'tp': tp, # true positive\n 'fp': fp, # false positive\n 'fn': fn, # false negative\n 'tn': tn, # true negative\n 'Np': Np, # positive accumulation\n 'Nn': Nn} # negative accumulation\n\n@register_vis_metric\ndef web(frame_pred, output_dir, **kwargs):\n\n os.makedirs(os.path.join(output_dir, 'web'), exist_ok=True) \n video_id = frame_pred['video_id']\n frame_name = frame_pred['frame_name']\n masks = frame_pred['masks'] # nq h w\n\n scores = torch.tensor(frame_pred['classes']) # nq c\n foreground_scores = scores[:, :-1].sum(-1) # nq\n max_idx = foreground_scores.argmax()\n pred_mask = masks[max_idx].int() # h w\n\n mask = Image.fromarray(255 * pred_mask.int().numpy()).convert('L')\n save_path = os.path.join(output_dir, 'web', video_id)\n\n os.makedirs(save_path, exist_ok=True)\n png_path = os.path.join(save_path, f'{frame_name}.png')\n if os.path.exists(png_path):\n os.remove(png_path)\n mask.save(png_path)\n return {}\n\n\n"
},
{
"path": "data_schedule/vis/fibroid/__init__.py",
"content": "# 注册fibrois数据集\nfrom . import fibroid_dataset\n\n# 注册fibroid评估标准\nfrom . import evals"
},
{
"path": "data_schedule/vis/fibroid/evals.py",
"content": "from data_schedule.vis.evaluator_utils import register_vis_metric\nimport os\nfrom glob import glob\nfrom tqdm import tqdm\nimport shutil\nfrom functools import partial\nfrom PIL import Image\nimport numpy as np\nimport torch\nimport detectron2.utils.comm as comm\nimport logging\nimport pycocotools.mask as mask_util\nfrom pycocotools.mask import decode as decode_rle\nimport data_schedule.vis.fibroid.metrics as metrics\n\n@register_vis_metric\ndef fibroid_other_medi(model_preds, \n dataset_meta,\n **kwargs):\n assert comm.is_main_process()\n\n iou_by_test_sample = []\n dice_by_test_sample = []\n preds_by_test_sample = []\n gt_by_test_sample = []\n get_frames_gt_mask_fn = dataset_meta.get('get_frames_gt_mask_fn')\n\n for pred in model_preds:\n video_id = pred['video_id'] # str\n frame_name = pred['frame_name'] # list[str], t'\n masks = pred['masks']# list[rle], nq\n scores = pred['scores'] # nq\n\n max_idx = torch.tensor(scores).argmax()\n pred_mask = masks[max_idx] # rle\n pred_mask = decode_rle(pred_mask)\n pred_mask = torch.as_tensor(pred_mask, dtype=torch.uint8).contiguous() # h w\n\n gt_mask, _ = get_frames_gt_mask_fn(video_id=video_id, frames=[frame_name]) # 1 h w\n gt_mask = gt_mask[0].int() # 0/1\n\n preds_by_test_sample.append(pred_mask)\n gt_by_test_sample.append(gt_mask)\n\n tp, fp, fn, tn = metrics.get_stats(pred_mask[None, None, ...], gt_mask[None, None, ...], \n mode='binary')\n iou_score = metrics.iou_score(tp, fp, fn, tn, reduction='micro')\n dice = metrics.dice(tp, fp, fn, tn, reduction='micro')\n\n iou_by_test_sample.append(iou_score)\n dice_by_test_sample.append(dice)\n\n mean_iou = torch.tensor(iou_by_test_sample).mean()\n mean_dice = torch.tensor(dice_by_test_sample).mean()\n\n preds_by_test_sample = torch.stack(preds_by_test_sample, dim=0).unsqueeze(1) # N 1 h w\n gt_by_test_sample = torch.stack(gt_by_test_sample, dim=0).unsqueeze(1) # N 1 h w\n\n tp, fp, fn, tn = metrics.get_stats(preds_by_test_sample, gt_by_test_sample, \n mode='binary')\n overall_iou = metrics.iou_score(tp, fp, fn, tn, reduction='micro') \n recall = metrics.recall(tp, fp, fn, tn, reduction='micro-imagewise') \n precision = metrics.precision(tp, fp, fn, tn, reduction='micro-imagewise')\n\n all_medi = {\n 'mean_iou': mean_iou,\n 'dice': mean_dice,\n 'overall_iou': overall_iou, # J/overallIoU\n 'recall': recall,\n 'precision': precision,\n 'F': 2 * precision * recall / (precision + recall)\n } \n return all_medi \n \n\nfrom collections import defaultdict\n\n# by_vid, by_frame\niou_dict = defaultdict(dict)\n\n@register_vis_metric\ndef fibroid_mask_dice_iou(frame_pred, dataset_meta, **kwargs):\n video_id = frame_pred['video_id']\n frame_name = frame_pred['frame_name']\n masks = frame_pred['masks'] # nq h w\n get_frames_gt_mask_fn = dataset_meta.get('get_frames_gt_mask_fn')\n scores = torch.tensor(frame_pred['classes']) # nq c, 保证c是2\n foreground_scores = scores[:, :-1].sum(-1) # nq\n max_idx = foreground_scores.argmax()\n pred_mask = masks[max_idx].int() # h w\n\n gt_mask, _ = get_frames_gt_mask_fn(video_id=video_id, frames=[frame_name]) # 1 h w\n gt_mask = gt_mask[0].int() # h w\n\n inter, union = (pred_mask*gt_mask).sum(), (pred_mask+gt_mask).sum()\n dice = (2*inter+1)/(union+1)\n iou = (inter+1)/(union-inter+1)\n iou_dict[video_id][frame_name] = iou\n if iou > 0.6:\n print(f'video_id: {video_id}, frame: {frame_name}: dice {dice}, iou {iou}')\n return {'dice': dice, 'iou': iou}\n\n@register_vis_metric\ndef fibroid_metric_aggregator(metrics_by_vid_frame, dataset_meta, eval_meta_keys, **kwargs):\n # output: eval_metrics\n # video: frame_name: metric/ vid_metrics\n\n eval_metrics = {}\n # video, frame_name\n # perframe metrics\n metric_names = metrics_by_vid_frame[list(eval_meta_keys.keys())[0]][eval_meta_keys[list(eval_meta_keys.keys())[0]][0]]\n for taylor_swift in metric_names:\n eval_metrics[taylor_swift] = torch.tensor([metrics_by_vid_frame[video][frame][taylor_swift] for video in eval_meta_keys.keys() for frame in eval_meta_keys[video]]).mean()\n \n # metrics by each video\n mean_iou_by_each_video = {}\n mean_dice_by_each_video = {}\n for video in eval_meta_keys:\n mean_iou_by_each_video[video] = torch.tensor([metrics_by_vid_frame[video][fname]['iou'] for fname in eval_meta_keys[video]]).mean()\n mean_dice_by_each_video[video] = torch.tensor([metrics_by_vid_frame[video][fname]['dice'] for fname in eval_meta_keys[video]]).mean()\n \n mean_iou_by_each_video = dict(sorted(mean_iou_by_each_video.items(), key=lambda x: x[1]))\n mean_dice_by_each_video = dict(sorted(mean_dice_by_each_video.items(), key=lambda x: x[1]))\n logging.debug(f'mean_iou_by_each_video: {mean_iou_by_each_video}')\n logging.debug(f'mean_dice_by_each_video: {mean_dice_by_each_video}')\n return eval_metrics\n\n"
},
{
"path": "data_schedule/vis/fibroid/fibroid_dataset.py",
"content": "from typing import Optional, Union\nimport json\nimport os\nfrom functools import partial\nimport numpy as np\nimport torch\nimport logging\nfrom tqdm import tqdm\nimport copy\nfrom detectron2.data import DatasetCatalog, MetadataCatalog\nfrom collections import defaultdict\nfrom data_schedule.vis.apis import VIS_Dataset\nfrom .fibroid_utils import get_frames, get_frames_mask, SET_NAME_TO_DIR,\\\n SET_NAME, SET_NAME_TO_NUM_VIDEOS, SET_NAME_TO_MODE, SET_NAME_TO_PREFIX, SET_NAME_TO_GT_TYPE\n\ndef fibroid_train(step_size, # none / int; 0, 6, 13, 19 ...\n split_dataset_name,\n video_ids,\n video_to_frames):\n logging.debug(f'{split_dataset_name} Generating metas...') \n metas = []\n for vid_id in tqdm(video_ids):\n all_frames = sorted(video_to_frames[vid_id])\n if step_size is None: \n metas.append({\n 'video_id': vid_id,\n 'all_frames' : all_frames,\n 'meta_idx': len(metas),\n 'all_objs': {1: {'class_label': 0,}} # 语义分割\n }) \n else:\n for frame_idx in range(0, len(all_frames), step_size):\n metas.append({\n 'video_id': vid_id,\n 'frame_idx': frame_idx,\n 'all_frames': all_frames,\n 'all_objs': {1: {'class_label': 0,}},\n 'meta_idx': len(metas)\n }) \n\n logging.debug(f'{split_dataset_name} Total metas: [{len(metas)}]')\n return metas\n\ndef fibroid_evaluate(eval_video_ids,\n split_dataset_name,\n step_size,\n video_to_frames):\n if (step_size is not None) and (step_size > 1):\n logging.warning('为什么 evaluate的时候step size大于1呢')\n raise ValueError()\n metas = []\n for video_id in eval_video_ids:\n VIS_Dataset\n all_frames = sorted(video_to_frames[video_id])\n if step_size == None:\n metas.append({\n 'video_id': video_id,\n 'all_frames': all_frames,\n 'meta_idx': len(metas)\n }) \n \n else: \n for frame_idx in range(0, len(all_frames), step_size):\n metas.append({\n 'video_id': video_id,\n 'frame_idx': frame_idx,\n 'all_frames': all_frames,\n 'meta_idx': len(metas)\n }) \n\n logging.debug(f'{split_dataset_name} Total metas: [{len(metas)}]') \n return metas\n\n\n_root = os.getenv('DATASET_PATH')\nroot = os.path.join(_root, 'uterus_myoma/Dataset')\nvisualize_meta_idxs = defaultdict(list)\nvisualize_meta_idxs['fibroid_train_step[6]'] = [] \nvisualize_meta_idxs['fibroid_train'] = [] \nvisualize_meta_idxs['fibroid_train_ste[1]'] = [] \nvisualize_meta_idxs['fibroid_validate'] = []\nvisualize_meta_idxs['fibroid_validate_step[1]'] = []\nvisualize_meta_idxs['weakPolyP_fibroid_validate_step[1]'] = [] \n\nfibroid_meta = {\n 'thing_classes': ['rumor', 'not rumor'],\n 'thing_colors': [(255., 140., 0.), (0., 255., 0.)],\n}\n\nfor name in SET_NAME:\n set_dir = SET_NAME_TO_DIR[name] \n set_dir = os.path.join(root, set_dir)\n num_videos = SET_NAME_TO_NUM_VIDEOS[name]\n\n video_ids = os.listdir(os.path.join(set_dir, 'Frame'))\n assert len(video_ids) == num_videos\n\n video_to_frames = {\n vid: sorted([png[:-4] for png in os.listdir(os.path.join(set_dir, 'Frame', vid)) if png.endswith('.png')])\\\n for vid in video_ids\n }\n mode = SET_NAME_TO_MODE[name]\n prefix = SET_NAME_TO_PREFIX[name]\n\n if mode == 'train':\n train_meta = copy.deepcopy(fibroid_meta)\n gt_type = SET_NAME_TO_GT_TYPE[name]\n train_meta.update({\n 'mode': 'train',\n 'get_frames_fn': partial(get_frames, frames_path=os.path.join(set_dir, 'Frame')),\n 'get_frames_mask_fn': partial(get_frames_mask, mask_path=os.path.join(set_dir, gt_type),),\n 'get_frames_gt_mask_fn': partial(get_frames_mask, mask_path=os.path.join(root, os.path.join(set_dir, 'GT')),),\n })\n # train\n for step_size in [1, 6, None]:\n step_identifer = '' if step_size is None else f'_step[{step_size}]'\n split_name = f'{prefix}{step_identifer}'\n train_meta.update({'name': split_name})\n DatasetCatalog.register(split_name, partial(fibroid_train,\n video_ids=video_ids, \n split_dataset_name=split_name,\n step_size=step_size,\n video_to_frames=video_to_frames,)) \n MetadataCatalog.get(split_name).set(**train_meta, \n step_size=step_size,\n visualize_meta_idxs=visualize_meta_idxs[split_name]) \n\n\n elif mode == 'evaluate':\n prefix = SET_NAME_TO_PREFIX[name]\n validate_meta = copy.deepcopy(fibroid_meta)\n\n validate_meta.update({\n 'mode': 'evaluate',\n 'get_frames_fn': partial(get_frames, frames_path=os.path.join(root, os.path.join(set_dir, 'Frame'))),\n 'eval_set_name': SET_NAME_TO_DIR[name],\n 'get_frames_gt_mask_fn': partial(get_frames_mask, mask_path=os.path.join(root, os.path.join(set_dir, 'GT')),),\n 'eval_meta_keys': video_to_frames\n })\n # validate\n for step_size in [1, None,]:\n step_identifer = '' if step_size is None else f'_step[{step_size}]'\n split_name = f'{prefix}{step_identifer}'\n validate_meta.update({'name': split_name})\n DatasetCatalog.register(split_name, partial(fibroid_evaluate,\n eval_video_ids=video_ids, \n split_dataset_name=split_name, \n step_size=step_size,\n video_to_frames=video_to_frames)) \n MetadataCatalog.get(split_name).set(**validate_meta, step_size=step_size,\n visualize_meta_idxs=visualize_meta_idxs[split_name])\n \n\n\n \n\n"
},
{
"path": "data_schedule/vis/fibroid/fibroid_utils.py",
"content": "import wandb\nimport plotly.express as px\nimport logging\nimport os\nimport numpy as np\nimport torch\nimport json\nfrom joblib import Parallel, delayed\nimport multiprocessing\nimport torch.distributed as dist\nimport detectron2.utils.comm as comm\n\nimport pycocotools.mask as mask_util\nfrom pycocotools.mask import encode, area\n\nfrom data_schedule.utils.segmentation import bounding_box_from_mask\nfrom data_schedule.utils.video_clips import generate_windows_of_video\nfrom glob import glob\nfrom PIL import Image\n\ndef get_frames(frames_path, video_id, frames):\n return [Image.open(os.path.join(frames_path, video_id, f'{f}.png'),).convert('RGB') for f in frames]\n\n# t' h w, int, obj_ids ; has_ann t\ndef get_frames_mask(mask_path, video_id, frames):\n masks = [Image.open(os.path.join(mask_path, video_id, f'{f}.png')).convert('L') for f in frames]\n masks = [np.array(mk) for mk in masks]\n masks = torch.stack([torch.from_numpy(mk) for mk in masks], dim=0) # t h w\n masks = (masks > 0).int()\n return masks, torch.ones(len(frames)).bool()\n\n\nSET_NAME = [\n 'fibroid_train', \n 'fibroid_validate',\n 'weakpolyp_train',\n\n 'fibroid_validate_temp7',\n 'fibroid_train_temp7',\n # 'weakpolyp_fibroid_train_temp7',\n\n 'fibroid_validate_temp8',\n 'fibroid_train_temp8',\n 'weakpolyp_fibroid_train_temp8'\n ]\n\nSET_NAME_TO_DIR = {\n 'fibroid_train': 'temp/train',\n 'fibroid_validate': 'temp/test',\n 'weakpolyp_train': 'temp/uterus_myoma_WeakPolyP_temp/train',\n\n\n 'fibroid_validate_temp7': 'temp7/test',\n 'fibroid_train_temp7': 'temp7/train',\n 'weakpolyp_fibroid_train_temp7': 'temp7/uterus_myoma_WeakPolyP_temp7/train',\n\n 'fibroid_validate_temp8': 'temp8/test',\n 'fibroid_train_temp8': 'temp8/train',\n 'weakpolyp_fibroid_train_temp8': 'temp8/uterus_myoma_WeakPolyP_temp8/train',\n\n}\n\nSET_NAME_TO_NUM_VIDEOS = {\n 'fibroid_train': 80,\n 'fibroid_validate': 20,\n 'weakpolyp_train': 80,\n\n\n 'fibroid_train_temp7': 85,\n 'fibroid_validate_temp7': 15,\n 'weakpolyp_fibroid_train_temp7': 85 , \n\n\n 'fibroid_train_temp8': 83,\n 'fibroid_validate_temp8': 17,\n 'weakpolyp_fibroid_train_temp8': 83 \n}\n\nSET_NAME_TO_MODE = {\n 'fibroid_train': 'train',\n 'fibroid_validate': 'evaluate',\n 'weakpolyp_train': 'train',\n\n 'fibroid_train_temp7': 'train',\n 'fibroid_validate_temp7': 'evaluate',\n 'weakpolyp_fibroid_train_temp7': 'train', \n\n 'fibroid_train_temp8': 'train',\n 'fibroid_validate_temp8': 'evaluate',\n 'weakpolyp_fibroid_train_temp8': 'train' \n}\n\nSET_NAME_TO_PREFIX = {\n 'fibroid_train': 'fibroid_train',\n 'fibroid_validate': 'fibroid_validate',\n 'weakpolyp_train': 'weakpolyp_fibroid_train',\n\n 'fibroid_train_temp7': 'fibroid_train_temp7',\n 'fibroid_validate_temp7': 'fibroid_validate_temp7',\n 'weakpolyp_fibroid_train_temp7': 'weakpolyp_fibroid_train_temp7' ,\n\n 'fibroid_train_temp8': 'fibroid_train_temp8',\n 'fibroid_validate_temp8': 'fibroid_validate_temp8',\n 'weakpolyp_fibroid_train_temp8': 'weakpolyp_fibroid_train_temp8' \n\n}\n\nSET_NAME_TO_GT_TYPE = {\n 'fibroid_train': 'GT',\n 'fibroid_validate': 'GT',\n 'weakpolyp_train': 'Box',\n\n 'fibroid_train_temp7': 'GT',\n 'fibroid_validate_temp7': 'GT',\n 'weakpolyp_fibroid_train_temp7': 'Box', \n\n 'fibroid_train_temp8': 'GT',\n 'fibroid_validate_temp8': 'GT',\n 'weakpolyp_fibroid_train_temp8': 'Box' \n}\n"
},
{
"path": "data_schedule/vis/fibroid/metrics.py",
"content": "import warnings\nfrom typing import Optional, List, Tuple, Union\nimport torch\n\n\"\"\"Various metrics based on Type I and Type II errors.\n\nReferences:\n https://en.wikipedia.org/wiki/Confusion_matrix\n\n\nExample:\n\n .. code-block:: python\n\n import segmentation_models_pytorch as smp\n\n # lets assume we have multilabel prediction for 3 classes\n output = torch.rand([10, 3, 256, 256])\n target = torch.rand([10, 3, 256, 256]).round().long()\n\n # first compute statistics for true positives, false positives, false negative and\n # true negative \"pixels\"\n tp, fp, fn, tn = smp.metrics.get_stats(output, target, mode='multilabel', threshold=0.5)\n\n # then compute metrics with required reduction (see metric docs)\n iou_score = smp.metrics.iou_score(tp, fp, fn, tn, reduction=\"micro\")\n f1_score = smp.metrics.f1_score(tp, fp, fn, tn, reduction=\"micro\")\n f2_score = smp.metrics.fbeta_score(tp, fp, fn, tn, beta=2, reduction=\"micro\")\n accuracy = smp.metrics.accuracy(tp, fp, fn, tn, reduction=\"macro\")\n recall = smp.metrics.recall(tp, fp, fn, tn, reduction=\"micro-imagewise\")\n\n\"\"\"\n\n\n__all__ = [\n \"get_stats\",\n \"fbeta_score\",\n \"f1_score\",\n \"iou_score\",\n \"accuracy\",\n \"precision\",\n \"recall\",\n \"sensitivity\",\n \"specificity\",\n \"balanced_accuracy\",\n \"positive_predictive_value\",\n \"negative_predictive_value\",\n \"false_negative_rate\",\n \"false_positive_rate\",\n \"false_discovery_rate\",\n \"false_omission_rate\",\n \"positive_likelihood_ratio\",\n \"negative_likelihood_ratio\",\n]\n\n\n###################################################################################################\n# Statistics computation (true positives, false positives, false negatives, false positives)\n###################################################################################################\n\n\ndef get_stats(\n output: Union[torch.LongTensor, torch.FloatTensor],\n target: torch.LongTensor,\n mode: str,\n ignore_index: Optional[int] = None,\n threshold: Optional[Union[float, List[float]]] = None,\n num_classes: Optional[int] = None,\n) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]:\n \"\"\"Compute true positive, false positive, false negative, true negative 'pixels'\n for each image and each class.\n\n Args:\n output (Union[torch.LongTensor, torch.FloatTensor]): Model output with following\n shapes and types depending on the specified ``mode``:\n\n 'binary'\n shape (N, 1, ...) and ``torch.LongTensor`` or ``torch.FloatTensor``\n\n 'multilabel'\n shape (N, C, ...) and ``torch.LongTensor`` or ``torch.FloatTensor``\n\n 'multiclass'\n shape (N, ...) and ``torch.LongTensor``\n\n target (torch.LongTensor): Targets with following shapes depending on the specified ``mode``:\n\n 'binary'\n shape (N, 1, ...)\n\n 'multilabel'\n shape (N, C, ...)\n\n 'multiclass'\n shape (N, ...)\n\n mode (str): One of ``'binary'`` | ``'multilabel'`` | ``'multiclass'``\n ignore_index (Optional[int]): Label to ignore on for metric computation.\n **Not** supproted for ``'binary'`` and ``'multilabel'`` modes. Defaults to None.\n threshold (Optional[float, List[float]]): Binarization threshold for\n ``output`` in case of ``'binary'`` or ``'multilabel'`` modes. Defaults to None.\n num_classes (Optional[int]): Number of classes, necessary attribute\n only for ``'multiclass'`` mode. Class values should be in range 0..(num_classes - 1).\n If ``ignore_index`` is specified it should be outside the classes range, e.g. ``-1`` or\n ``255``.\n\n Raises:\n ValueError: in case of misconfiguration.\n\n Returns:\n Tuple[torch.LongTensor]: true_positive, false_positive, false_negative,\n true_negative tensors (N, C) shape each.\n\n \"\"\"\n\n if torch.is_floating_point(target):\n raise ValueError(f\"Target should be one of the integer types, got {target.dtype}.\")\n\n if torch.is_floating_point(output) and threshold is None:\n raise ValueError(\n f\"Output should be one of the integer types if ``threshold`` is not None, got {output.dtype}.\"\n )\n\n if torch.is_floating_point(output) and mode == \"multiclass\":\n raise ValueError(f\"For ``multiclass`` mode ``output`` should be one of the integer types, got {output.dtype}.\")\n\n if mode not in {\"binary\", \"multiclass\", \"multilabel\"}:\n raise ValueError(f\"``mode`` should be in ['binary', 'multiclass', 'multilabel'], got mode={mode}.\")\n\n if mode == \"multiclass\" and threshold is not None:\n raise ValueError(\"``threshold`` parameter does not supported for this 'multiclass' mode\")\n\n if output.shape != target.shape:\n raise ValueError(\n \"Dimensions should match, but ``output`` shape is not equal to ``target`` \"\n + f\"shape, {output.shape} != {target.shape}\"\n )\n\n if mode != \"multiclass\" and ignore_index is not None:\n raise ValueError(f\"``ignore_index`` parameter is not supproted for '{mode}' mode\")\n\n if mode == \"multiclass\" and num_classes is None:\n raise ValueError(\"``num_classes`` attribute should be not ``None`` for 'multiclass' mode.\")\n\n if ignore_index is not None and 0 <= ignore_index <= num_classes - 1:\n raise ValueError(\n f\"``ignore_index`` should be outside the class values range, but got class values in range \"\n f\"0..{num_classes - 1} and ``ignore_index={ignore_index}``. Hint: if you have ``ignore_index = 0``\"\n f\"consirder subtracting ``1`` from your target and model output to make ``ignore_index = -1``\"\n f\"and relevant class values started from ``0``.\"\n )\n\n if mode == \"multiclass\":\n tp, fp, fn, tn = _get_stats_multiclass(output, target, num_classes, ignore_index)\n else:\n if threshold is not None:\n output = torch.where(output >= threshold, 1, 0)\n target = torch.where(target >= threshold, 1, 0)\n tp, fp, fn, tn = _get_stats_multilabel(output, target)\n\n return tp, fp, fn, tn\n\n\n@torch.no_grad()\ndef _get_stats_multiclass(\n output: torch.LongTensor,\n target: torch.LongTensor,\n num_classes: int,\n ignore_index: Optional[int],\n) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]:\n\n batch_size, *dims = output.shape\n num_elements = torch.prod(torch.tensor(dims)).long()\n\n if ignore_index is not None:\n ignore = target == ignore_index\n output = torch.where(ignore, -1, output)\n target = torch.where(ignore, -1, target)\n ignore_per_sample = ignore.view(batch_size, -1).sum(1)\n\n tp_count = torch.zeros(batch_size, num_classes, dtype=torch.long)\n fp_count = torch.zeros(batch_size, num_classes, dtype=torch.long)\n fn_count = torch.zeros(batch_size, num_classes, dtype=torch.long)\n tn_count = torch.zeros(batch_size, num_classes, dtype=torch.long)\n\n for i in range(batch_size):\n target_i = target[i]\n output_i = output[i]\n mask = output_i == target_i\n matched = torch.where(mask, target_i, -1)\n tp = torch.histc(matched.float(), bins=num_classes, min=0, max=num_classes - 1)\n fp = torch.histc(output_i.float(), bins=num_classes, min=0, max=num_classes - 1) - tp\n fn = torch.histc(target_i.float(), bins=num_classes, min=0, max=num_classes - 1) - tp\n tn = num_elements - tp - fp - fn\n if ignore_index is not None:\n tn = tn - ignore_per_sample[i]\n tp_count[i] = tp.long()\n fp_count[i] = fp.long()\n fn_count[i] = fn.long()\n tn_count[i] = tn.long()\n\n return tp_count, fp_count, fn_count, tn_count\n\n\n@torch.no_grad()\ndef _get_stats_multilabel(\n output: torch.LongTensor,\n target: torch.LongTensor,\n) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]:\n\n batch_size, num_classes, *dims = target.shape\n output = output.view(batch_size, num_classes, -1)\n target = target.view(batch_size, num_classes, -1)\n\n tp = (output * target).sum(2)\n fp = output.sum(2) - tp\n fn = target.sum(2) - tp\n tn = torch.prod(torch.tensor(dims)) - (tp + fp + fn)\n\n return tp, fp, fn, tn\n\n\n###################################################################################################\n# Metrics computation\n###################################################################################################\n\n\ndef _handle_zero_division(x, zero_division):\n nans = torch.isnan(x)\n if torch.any(nans) and zero_division == \"warn\":\n warnings.warn(\"Zero division in metric calculation!\")\n value = zero_division if zero_division != \"warn\" else 0\n value = torch.tensor(value, dtype=x.dtype).to(x.device)\n x = torch.where(nans, value, x)\n return x\n\n\ndef _compute_metric(\n metric_fn,\n tp,\n fp,\n fn,\n tn,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division=\"warn\",\n **metric_kwargs,\n) -> float:\n\n if class_weights is None and reduction is not None and \"weighted\" in reduction:\n raise ValueError(f\"Class weights should be provided for `{reduction}` reduction\")\n\n class_weights = class_weights if class_weights is not None else 1.0\n class_weights = torch.tensor(class_weights).to(tp.device)\n class_weights = class_weights / class_weights.sum()\n\n if reduction == \"micro\":\n tp = tp.sum()\n fp = fp.sum()\n fn = fn.sum()\n tn = tn.sum()\n score = metric_fn(tp, fp, fn, tn, **metric_kwargs)\n\n elif reduction == \"macro\":\n tp = tp.sum(0)\n fp = fp.sum(0)\n fn = fn.sum(0)\n tn = tn.sum(0)\n score = metric_fn(tp, fp, fn, tn, **metric_kwargs)\n score = _handle_zero_division(score, zero_division)\n score = (score * class_weights).mean()\n\n elif reduction == \"weighted\":\n tp = tp.sum(0)\n fp = fp.sum(0)\n fn = fn.sum(0)\n tn = tn.sum(0)\n score = metric_fn(tp, fp, fn, tn, **metric_kwargs)\n score = _handle_zero_division(score, zero_division)\n score = (score * class_weights).sum()\n\n elif reduction == \"micro-imagewise\":\n tp = tp.sum(1)\n fp = fp.sum(1)\n fn = fn.sum(1)\n tn = tn.sum(1)\n score = metric_fn(tp, fp, fn, tn, **metric_kwargs)\n score = _handle_zero_division(score, zero_division)\n score = score.mean()\n\n elif reduction == \"macro-imagewise\" or reduction == \"weighted-imagewise\":\n score = metric_fn(tp, fp, fn, tn, **metric_kwargs)\n score = _handle_zero_division(score, zero_division)\n score = (score.mean(0) * class_weights).mean()\n\n elif reduction == \"none\" or reduction is None:\n score = metric_fn(tp, fp, fn, tn, **metric_kwargs)\n score = _handle_zero_division(score, zero_division)\n\n else:\n raise ValueError(\n \"`reduction` should be in [micro, macro, weighted, micro-imagewise,\"\n + \"macro-imagesize, weighted-imagewise, none, None]\"\n )\n\n return score\n\n\n# Logic for metric computation, all metrics are with the same interface\n\n\ndef _fbeta_score(tp, fp, fn, tn, beta=1):\n beta_tp = (1 + beta**2) * tp\n beta_fn = (beta**2) * fn\n score = beta_tp / (beta_tp + beta_fn + fp)\n return score\n\n\ndef _iou_score(tp, fp, fn, tn):\n return tp / (tp + fp + fn)\n\n\ndef _accuracy(tp, fp, fn, tn):\n return (tp + tn) / (tp + fp + fn + tn)\n\n\ndef _sensitivity(tp, fp, fn, tn):\n return tp / (tp + fn)\n\n\ndef _specificity(tp, fp, fn, tn):\n return tn / (tn + fp)\n\n\ndef _balanced_accuracy(tp, fp, fn, tn):\n return (_sensitivity(tp, fp, fn, tn) + _specificity(tp, fp, fn, tn)) / 2\n\ndef _dice(tp, fp, fn, tn):\n return tp * 2 / (tp + fp + fn + tp)\n\ndef _positive_predictive_value(tp, fp, fn, tn):\n return tp / (tp + fp)\n\n\ndef _negative_predictive_value(tp, fp, fn, tn):\n return tn / (tn + fn)\n\n\ndef _false_negative_rate(tp, fp, fn, tn):\n return fn / (fn + tp)\n\n\ndef _false_positive_rate(tp, fp, fn, tn):\n return fp / (fp + tn)\n\n\ndef _false_discovery_rate(tp, fp, fn, tn):\n return 1 - _positive_predictive_value(tp, fp, fn, tn)\n\n\ndef _false_omission_rate(tp, fp, fn, tn):\n return 1 - _negative_predictive_value(tp, fp, fn, tn)\n\n\ndef _positive_likelihood_ratio(tp, fp, fn, tn):\n return _sensitivity(tp, fp, fn, tn) / _false_positive_rate(tp, fp, fn, tn)\n\n\ndef _negative_likelihood_ratio(tp, fp, fn, tn):\n return _false_negative_rate(tp, fp, fn, tn) / _specificity(tp, fp, fn, tn)\n\n\ndef fbeta_score(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n beta: float = 1.0,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"F beta score\"\"\"\n return _compute_metric(\n _fbeta_score,\n tp,\n fp,\n fn,\n tn,\n beta=beta,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef f1_score(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"F1 score\"\"\"\n return _compute_metric(\n _fbeta_score,\n tp,\n fp,\n fn,\n tn,\n beta=1.0,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\ndef dice(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"IoU score or Jaccard index\"\"\" # noqa\n return _compute_metric(\n _dice,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\ndef iou_score(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"IoU score or Jaccard index\"\"\" # noqa\n return _compute_metric(\n _iou_score,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef accuracy(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Accuracy\"\"\"\n return _compute_metric(\n _accuracy,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef sensitivity(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Sensitivity, recall, hit rate, or true positive rate (TPR)\"\"\"\n return _compute_metric(\n _sensitivity,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef specificity(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Specificity, selectivity or true negative rate (TNR)\"\"\"\n return _compute_metric(\n _specificity,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef balanced_accuracy(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Balanced accuracy\"\"\"\n return _compute_metric(\n _balanced_accuracy,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef positive_predictive_value(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Precision or positive predictive value (PPV)\"\"\"\n return _compute_metric(\n _positive_predictive_value,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef negative_predictive_value(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Negative predictive value (NPV)\"\"\"\n return _compute_metric(\n _negative_predictive_value,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef false_negative_rate(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Miss rate or false negative rate (FNR)\"\"\"\n return _compute_metric(\n _false_negative_rate,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef false_positive_rate(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Fall-out or false positive rate (FPR)\"\"\"\n return _compute_metric(\n _false_positive_rate,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef false_discovery_rate(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"False discovery rate (FDR)\"\"\" # noqa\n return _compute_metric(\n _false_discovery_rate,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef false_omission_rate(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"False omission rate (FOR)\"\"\" # noqa\n return _compute_metric(\n _false_omission_rate,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef positive_likelihood_ratio(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Positive likelihood ratio (LR+)\"\"\"\n return _compute_metric(\n _positive_likelihood_ratio,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\ndef negative_likelihood_ratio(\n tp: torch.LongTensor,\n fp: torch.LongTensor,\n fn: torch.LongTensor,\n tn: torch.LongTensor,\n reduction: Optional[str] = None,\n class_weights: Optional[List[float]] = None,\n zero_division: Union[str, float] = 1.0,\n) -> torch.Tensor:\n \"\"\"Negative likelihood ratio (LR-)\"\"\"\n return _compute_metric(\n _negative_likelihood_ratio,\n tp,\n fp,\n fn,\n tn,\n reduction=reduction,\n class_weights=class_weights,\n zero_division=zero_division,\n )\n\n\n_doc = \"\"\"\n\n Args:\n tp (torch.LongTensor): tensor of shape (N, C), true positive cases\n fp (torch.LongTensor): tensor of shape (N, C), false positive cases\n fn (torch.LongTensor): tensor of shape (N, C), false negative cases\n tn (torch.LongTensor): tensor of shape (N, C), true negative cases\n reduction (Optional[str]): Define how to aggregate metric between classes and images:\n\n - 'micro'\n Sum true positive, false positive, false negative and true negative pixels over\n all images and all classes and then compute score.\n\n - 'macro'\n Sum true positive, false positive, false negative and true negative pixels over\n all images for each label, then compute score for each label separately and average labels scores.\n This does not take label imbalance into account.\n\n - 'weighted'\n Sum true positive, false positive, false negative and true negative pixels over\n all images for each label, then compute score for each label separately and average\n weighted labels scores.\n\n - 'micro-imagewise'\n Sum true positive, false positive, false negative and true negative pixels for **each image**,\n then compute score for **each image** and average scores over dataset. All images contribute equally\n to final score, however takes into accout class imbalance for each image.\n\n - 'macro-imagewise'\n Compute score for each image and for each class on that image separately, then compute average score\n on each image over labels and average image scores over dataset. Does not take into account label\n imbalance on each image.\n\n - 'weighted-imagewise'\n Compute score for each image and for each class on that image separately, then compute weighted average\n score on each image over labels and average image scores over dataset.\n\n - 'none' or ``None``\n Same as ``'macro-imagewise'``, but without any reduction.\n\n For ``'binary'`` case ``'micro' = 'macro' = 'weighted'`` and\n ``'micro-imagewise' = 'macro-imagewise' = 'weighted-imagewise'``.\n\n Prefixes ``'micro'``, ``'macro'`` and ``'weighted'`` define how the scores for classes will be aggregated,\n while postfix ``'imagewise'`` defines how scores between the images will be aggregated.\n\n class_weights (Optional[List[float]]): list of class weights for metric\n aggregation, in case of `weighted*` reduction is chosen. Defaults to None.\n zero_division (Union[str, float]): Sets the value to return when there is a zero division,\n i.e. when all predictions and labels are negative. If set to “warn”, this acts as 0,\n but warnings are also raised. Defaults to 1.\n\n Returns:\n torch.Tensor: if ``'reduction'`` is not ``None`` or ``'none'`` returns scalar metric,\n else returns tensor of shape (N, C)\n\n References:\n https://en.wikipedia.org/wiki/Confusion_matrix\n\"\"\"\n\nfbeta_score.__doc__ += _doc\nf1_score.__doc__ += _doc\niou_score.__doc__ += _doc\naccuracy.__doc__ += _doc\nsensitivity.__doc__ += _doc\nspecificity.__doc__ += _doc\nbalanced_accuracy.__doc__ += _doc\npositive_predictive_value.__doc__ += _doc\nnegative_predictive_value.__doc__ += _doc\nfalse_negative_rate.__doc__ += _doc\nfalse_positive_rate.__doc__ += _doc\nfalse_discovery_rate.__doc__ += _doc\nfalse_omission_rate.__doc__ += _doc\npositive_likelihood_ratio.__doc__ += _doc\nnegative_likelihood_ratio.__doc__ += _doc\n\nprecision = positive_predictive_value\nrecall = sensitivity"
},
{
"path": "data_schedule/vis/mapper.py",
"content": "\nimport json\nimport os\nfrom typing import List\nimport copy\nfrom functools import partial\nimport random\nimport numpy as np\nimport torch\nimport logging\nfrom einops import rearrange\nfrom detectron2.data import MetadataCatalog\n\nfrom data_schedule.registry import MAPPER_REGISTRY\nfrom .mapper_utils import VIS_TrainMapper, VIS_EvalMapper\nfrom .vis_frame_sampler import VIS_FRAMES_SAMPLER_REGISTRY\nfrom data_schedule.vis.apis import VIS_Dataset, VIS_Aug_CallbackAPI,\\\n VIS_TrainAPI_clipped_video, VIS_EvalAPI_clipped_video_request_ann\n\n@MAPPER_REGISTRY.register()\nclass VIS_Video_EvalMapper(VIS_EvalMapper):\n def __init__(self,\n configs,\n dataset_name,\n mode,\n meta_idx_shift,\n ): \n assert mode == 'evaluate'\n dataset_meta = MetadataCatalog.get(dataset_name)\n assert dataset_meta.get('step_size') == None\n mapper_config = configs['data'][mode][dataset_name]['mapper']\n super().__init__(meta_idx_shift=meta_idx_shift,\n dataset_meta=dataset_meta,\n mapper_config=mapper_config)\n\n def _call(self, data_dict):\n VIS_Dataset\n video_id, all_frames = data_dict['video_id'], data_dict['all_frames']\n video_frames = self.get_frames_fn(video_id=video_id, frames=all_frames)\n aug_ret = {\n 'video': video_frames,\n 'callback_fns': []\n }\n VIS_Aug_CallbackAPI\n aug_ret = self.augmentation(aug_ret)\n video = aug_ret.pop('video')\n callback_fns = aug_ret.pop('callback_fns')[::-1]\n VIS_EvalAPI_clipped_video_request_ann\n return {\n 'video_dict': {'video': video},\n 'meta': {\n 'video_id': video_id,\n 'frames': all_frames,\n 'request_ann': torch.ones(len(all_frames)).bool(),\n 'callback_fns': callback_fns \n }\n }\n\n@MAPPER_REGISTRY.register()\nclass VIS_Video_or_Step_To_Clip_TrainMapper(VIS_TrainMapper):\n def __init__(self,\n dataset_name,\n configs,\n mode, \n meta_idx_shift,\n ): \n assert mode == 'train'\n dataset_meta = MetadataCatalog.get(dataset_name)\n assert dataset_meta.get('name') == dataset_name\n mapper_config = configs['data'][mode][dataset_name]['mapper']\n super().__init__(meta_idx_shift=meta_idx_shift,\n dataset_meta=dataset_meta,\n mapper_config=mapper_config)\n \n self.frames_sampler = VIS_FRAMES_SAMPLER_REGISTRY.get(\\\n mapper_config['frames_sampler']['name'])(sampler_configs=mapper_config['frames_sampler'],\n dataset_meta=dataset_meta)\n\n def _call(self, data_dict):\n VIS_Dataset\n video_id, all_frames, all_objs = data_dict['video_id'], data_dict['all_frames'], data_dict['all_objs']\n frame_idx = data_dict['frame_idx'] if 'frame_idx' in data_dict else None\n\n all_obj_ids = list(all_objs.keys()) # [1, 2, 5, 4]\n assert len(list(set(all_obj_ids))) == len(all_obj_ids)\n class_labels = torch.tensor([all_objs[key]['class_label'] for key in all_obj_ids]) # [8, 10, 20 34]\n\n re_sample = True\n sampled_counts = 0\n while re_sample:\n sampled_frames = self.frames_sampler(all_frames=all_frames, frame_idx=frame_idx, video_id=video_id)\n # t' h w, int, obj_ids ; has_ann t\n frames_mask, has_ann = self.get_frames_mask_fn(video_id=video_id, frames=sampled_frames)\n appear_objs = frames_mask.unique() # [0, 1, 2]\n assert set(appear_objs.tolist()).issubset(set([0] + all_obj_ids))\n re_sample = (len(list(set(appear_objs.tolist()) & set(all_obj_ids))) == 0)\n # 只要出现某些个物体就行\n sampled_counts += 1\n if sampled_counts > 2:\n logging.error('sampled two much times')\n raise RuntimeError()\n \n frames_mask = torch.stack([frames_mask == obj_id for obj_id in all_obj_ids], dim=0) # N t' h w, bool\n video_frames = self.get_frames_fn(video_id=video_id, frames=sampled_frames) \n width, height = video_frames[0].size\n aug_ret = {\n 'video': video_frames,\n 'masks': frames_mask, # N t' h w\n 'has_ann': has_ann, # t\n 'classes': class_labels, # N\n }\n VIS_Aug_CallbackAPI\n aug_ret = self.augmentation(aug_ret)\n video = aug_ret.pop('video')\n frame_targets = self.map_to_frame_targets(aug_ret)\n if self.clip_global_targets_map_to_local_targets:\n aug_ret = self.map_global_targets_to_local_targets(aug_ret)\n\n VIS_TrainAPI_clipped_video\n ret = {}\n ret['video_dict'] = {'video': video}\n ret['targets'] = aug_ret\n ret['frame_targets'] = frame_targets\n return ret\n\n\n\n\n\n\n\n\n"
},
{
"path": "data_schedule/vis/mapper_utils.py",
"content": "from .vis_aug_utils import VIS_EVAL_AUG_REGISTRY, VIS_TRAIN_AUG_REGISTRY\nimport torch\nfrom copy import deepcopy as dcopy\nfrom data_schedule.registry import Mapper\nimport copy\nfrom data_schedule.vis.apis import VIS_TrainAPI_clipped_video\n\nclass VIS_Mapper(Mapper):\n def __init__(self, \n meta_idx_shift,\n dataset_meta,) -> None:\n super().__init__(meta_idx_shift=meta_idx_shift, dataset_meta=dataset_meta)\n self.get_frames_fn = dataset_meta.get('get_frames_fn')\n\nclass VIS_TrainMapper(VIS_Mapper):\n\n def __init__(self, \n meta_idx_shift, \n dataset_meta,\n mapper_config) -> None:\n super().__init__(meta_idx_shift, dataset_meta)\n self.get_frames_mask_fn = dataset_meta.get('get_frames_mask_fn') \n self.clip_global_targets_map_to_local_targets = mapper_config['clip_global_targets_map_to_local_targets'] \n self.augmentation = VIS_TRAIN_AUG_REGISTRY.get(mapper_config['augmentation']['name'])(mapper_config['augmentation'])\n\n def map_to_frame_targets(self, clip_targets):\n VIS_TrainAPI_clipped_video\n clip_rets = copy.deepcopy(clip_targets)\n masks = clip_rets['masks'].transpose(0, 1).contiguous() # t' N h w\n class_labels = clip_rets['classes'] # [10, 32, 10, 4]\n has_box = 'boxes' in clip_rets\n if has_box:\n boxes = clip_rets['boxes'].transpose(0, 1).contiguous() # t' N 4\n assert len(masks) == len(boxes)\n ret = []\n for idx, frame_mk in enumerate(masks):\n frame_targets = {\n 'masks': frame_mk.unsqueeze(1), # N 1 h w\n 'classes': class_labels, # N\n }\n if has_box:\n frame_targets.update({'boxes': boxes[idx].unsqueeze(1)}) # N 1 4\n if self.clip_global_targets_map_to_local_targets:\n frame_targets = self.map_global_targets_to_local_targets(frame_targets)\n frame_targets['masks'] = frame_targets['masks'].squeeze(1)\n if has_box:\n frame_targets['boxes'] = frame_targets['boxes'].squeeze(1)\n ret.append(frame_targets)\n return ret\n\n def map_global_targets_to_local_targets(self, ret):\n VIS_TrainAPI_clipped_video\n masks = ret['masks'] # N t' h w\n global_obj_appear = masks.flatten(1).any(-1) # N [True, False, True, False, False, False, True]\n ret['masks'] = ret['masks'][global_obj_appear]\n ret['classes'] = ret['classes'][global_obj_appear]\n if 'boxes' in ret:\n ret['boxes'] = ret['boxes'][global_obj_appear] # n t' 4\n return ret\n \n\nclass VIS_EvalMapper(VIS_Mapper):\n def __init__(self, \n meta_idx_shift, \n dataset_meta,\n mapper_config) -> None:\n super().__init__(meta_idx_shift, dataset_meta)\n assert mapper_config['augmentation']['name'] in ['WeakPolyP_EvalAug', 'Visha_EvalAug']\n self.augmentation = VIS_EVAL_AUG_REGISTRY.get(mapper_config['augmentation']['name'])(mapper_config['augmentation'])\n \n\n\n"
},
{
"path": "data_schedule/vis/polyp/__init__.py",
"content": "\nfrom . import polyp_dataset\nfrom . import evals\n"
},
{
"path": "data_schedule/vis/polyp/evals.py",
"content": "from data_schedule.vis.evaluator_utils import register_vis_metric\nimport os\nimport torch\nimport detectron2.utils.comm as comm\nimport logging\nimport subprocess\n\n\n@register_vis_metric\ndef polyp_metric_aggregator(metrics_by_vid_frame, dataset_meta, eval_meta_keys, **kwargs):\n # output: eval_metrics\n # video: frame_name: metric/ vid_metrics\n\n eval_metrics = {}\n # video, frame_name\n # perframe metrics\n metric_names = metrics_by_vid_frame[list(eval_meta_keys.keys())[0]][eval_meta_keys[list(eval_meta_keys.keys())[0]][0]]\n for taylor_swift in metric_names:\n eval_metrics[taylor_swift] = torch.tensor([metrics_by_vid_frame[video][frame][taylor_swift] for video in eval_meta_keys.keys() for frame in eval_meta_keys[video]]).mean()\n \n # metrics by each video\n mean_iou_by_each_video = {}\n mean_dice_by_each_video = {}\n for billie_eilish in eval_meta_keys:\n mean_iou_by_each_video[billie_eilish] = torch.tensor([metrics_by_vid_frame[billie_eilish][fname]['iou'] for fname in eval_meta_keys[billie_eilish]]).mean()\n mean_dice_by_each_video[billie_eilish] = torch.tensor([metrics_by_vid_frame[billie_eilish][fname]['dice'] for fname in eval_meta_keys[billie_eilish]]).mean()\n \n mean_iou_by_each_video = dict(sorted(mean_iou_by_each_video.items(), key=lambda x: x[1]))\n mean_dice_by_each_video = dict(sorted(mean_dice_by_each_video.items(), key=lambda x: x[1])) \n logging.debug(f'mean_iou_by_each_video: {mean_iou_by_each_video}')\n logging.debug(f'mean_dice_by_each_video: {mean_dice_by_each_video}')\n \n return eval_metrics\n\n\n"
},
{
"path": "data_schedule/vis/polyp/polyp_dataset.py",
"content": "from typing import Optional, Union\nimport json\nimport os\nfrom functools import partial\nimport numpy as np\nimport torch\nimport logging\nfrom tqdm import tqdm\nimport copy\nfrom detectron2.data import DatasetCatalog, MetadataCatalog\nfrom collections import defaultdict\nfrom .polyp_utils import get_frames, get_frames_mask, SET_NAME_TO_DIR, SET_NAME, SET_NAME_TO_NUM_VIDEOS, SET_NAME_TO_MODE, SET_NAME_TO_PREFIX\n\ndef polyp_train(step_size,\n split_dataset_name,\n video_ids,\n video_to_frames,\n root_path):\n\n logging.debug(f'{split_dataset_name} Generating metas...') \n metas = []\n for vid_id in tqdm(video_ids):\n all_frames = sorted(video_to_frames[vid_id])\n poly_class = 0\n if step_size is None: \n metas.append({\n 'video_id': vid_id,\n 'all_frames' : all_frames,\n 'all_objs': { 1: {'class_label': poly_class} },\n 'meta_idx': len(metas)\n }) \n else:\n for frame_idx in range(0, len(all_frames), step_size):\n metas.append({\n 'video_id': vid_id,\n 'frame_idx': frame_idx,\n 'all_frames': all_frames,\n 'all_objs': { 1: {'class_label': poly_class} },\n 'meta_idx': len(metas)\n }) \n\n logging.debug(f'{split_dataset_name} Total metas: [{len(metas)}]')\n return metas\n\n\ndef polyp_evaluate(eval_video_ids,\n split_dataset_name,\n step_size,\n video_to_frames):\n if (step_size is not None) and (step_size > 1):\n logging.warning('why?')\n raise ValueError()\n metas = []\n for video_id in eval_video_ids:\n all_frames = sorted(video_to_frames[video_id])\n if step_size == None:\n metas.append({\n 'video_id': video_id,\n 'all_frames': all_frames,\n 'meta_idx': len(metas)\n }) \n \n else: \n for frame_idx in range(0, len(all_frames), step_size):\n metas.append({\n 'video_id': video_id,\n 'frame_idx': frame_idx,\n 'all_frames': all_frames,\n 'meta_idx': len(metas)\n }) \n\n logging.debug(f'{split_dataset_name} Total metas: [{len(metas)}]') \n return metas\n\n_root = os.getenv('DATASET_PATH')\n\nroot = os.path.join(_root, 'SUN/SUN-SEG2')\n\nvisualize_meta_idxs = defaultdict(list)\nvisualize_meta_idxs['polyp_train_step[6]'] = [] \nvisualize_meta_idxs['polyp_train'] = [] \nvisualize_meta_idxs['polyp_hard_unseen'] = []\nvisualize_meta_idxs['polyp_hard_seen'] = []\nvisualize_meta_idxs['polyp_easy_unseen'] = []\nvisualize_meta_idxs['polyp_easy_seen'] = []\n\npolyp_meta = {\n 'thing_classes': ['polyp', 'not polyp'],\n 'thing_colors': [(255., 140., 0.), (0., 255., 0.)],\n 'root': root\n}\n\nfor name in SET_NAME:\n set_dir = SET_NAME_TO_DIR[name]\n set_dir = os.path.join(root, set_dir)\n num_videos = SET_NAME_TO_NUM_VIDEOS[name]\n\n video_ids = os.listdir(os.path.join(set_dir, 'Frame'))\n assert len(video_ids) == num_videos\n\n video_to_frames = {\n vid: sorted([png[:-4] for png in os.listdir(os.path.join(set_dir, 'Frame', vid)) if png.endswith('.jpg')])\\\n for vid in video_ids\n }\n mode = SET_NAME_TO_MODE[name]\n if mode == 'train':\n prefix = SET_NAME_TO_PREFIX[name]\n train_meta = copy.deepcopy(polyp_meta)\n train_meta.update({\n 'mode': 'train',\n 'get_frames_fn': partial(get_frames, frames_path=os.path.join(set_dir, 'Frame')),\n 'get_frames_mask_fn': partial(get_frames_mask, mask_path=os.path.join(set_dir, 'GT'),),\n })\n\n # train\n for step_size in [1, 3, 6, 9, 12, None]:\n step_identifer = '' if step_size is None else f'_step[{step_size}]'\n split_name = f'{prefix}{step_identifer}'\n train_meta.update({'name': split_name})\n DatasetCatalog.register(split_name, partial(polyp_train,\n video_ids=video_ids, \n split_dataset_name=split_name,\n step_size=step_size,\n video_to_frames=video_to_frames,\n root_path=set_dir)) \n MetadataCatalog.get(split_name).set(**train_meta, \n step_size=step_size,\n visualize_meta_idxs=visualize_meta_idxs[split_name]) \n elif mode == 'evaluate':\n prefix = SET_NAME_TO_PREFIX[name]\n validate_meta = copy.deepcopy(polyp_meta)\n\n validate_meta.update({\n 'mode': 'evaluate',\n 'get_frames_fn': partial(get_frames, frames_path=os.path.join(root, os.path.join(set_dir, 'Frame'))),\n 'eval_set_name': SET_NAME_TO_DIR[name],\n 'get_frames_gt_mask_fn': partial(get_frames_mask, mask_path=os.path.join(root, os.path.join(set_dir, 'GT')),),\n 'eval_meta_keys': video_to_frames\n })\n for step_size in [1, None,]:\n step_identifer = '' if step_size is None else f'_step[{step_size}]'\n split_name = f'{prefix}{step_identifer}'\n validate_meta.update({'name': split_name})\n DatasetCatalog.register(split_name, partial(polyp_evaluate,\n eval_video_ids=video_ids, \n split_dataset_name=split_name, \n step_size=step_size,\n video_to_frames=video_to_frames)) \n MetadataCatalog.get(split_name).set(**validate_meta, step_size=step_size,\n visualize_meta_idxs=visualize_meta_idxs[split_name])\n"
},
{
"path": "data_schedule/vis/polyp/polyp_utils.py",
"content": "\nimport os\nimport numpy as np\nimport torch\nfrom PIL import Image\n\nSET_NAME = ['polyp_train', \n 'polyp_hard_seen_validate', \n 'polyp_hard_unseen_validate', \n 'polyp_easy_seen_validate', \n 'polyp_easy_unseen_validate',\n\n 'polyp_hard_validate',\n 'polyp_easy_validate',\n 'Kvasir-train',\n 'Mayo-train',\n '300-train',\n '612-train',\n '300-tv',\n '612-test',\n '612-val'\n ]\n\nSET_NAME_TO_DIR = {\n 'polyp_train': 'TrainDataset',\n 'polyp_hard_seen_validate': 'TestHardDataset/Seen',\n 'polyp_hard_unseen_validate': 'TestHardDataset/Unseen',\n 'polyp_easy_seen_validate': 'TestEasyDataset/Seen',\n 'polyp_easy_unseen_validate': 'TestEasyDataset/Unseen',\n \n 'polyp_hard_validate': 'TestHardDataset/Combine',\n 'polyp_easy_validate': 'TestEasyDataset/Combine',\n \n 'Kvasir-train': 'MICCAI-VPS-dataset/Kvasir-SEG',\n 'Mayo-train': 'MICCAI-VPS-dataset/VPS-TrainSet/ASU-Mayo_Clinic/Train',\n '300-train': 'MICCAI-VPS-dataset/VPS-TrainSet/CVC-ColonDB-300/Train',\n '612-train': 'MICCAI-VPS-dataset/VPS-TrainSet/CVC-ClinicDB-612/Train',\n '300-tv': 'MICCAI-VPS-dataset/VPS-TestSet/CVC-ColonDB-300',\n '612-test': 'MICCAI-VPS-dataset/VPS-TestSet/CVC-ClinicDB-612-Test',\n '612-val': 'MICCAI-VPS-dataset/VPS-TestSet/CVC-ClinicDB-612-Valid'\n}\n\nSET_NAME_TO_NUM_VIDEOS = {\n 'polyp_train': 112,\n 'polyp_hard_seen_validate': 17,\n 'polyp_hard_unseen_validate': 37,\n 'polyp_easy_seen_validate': 33,\n 'polyp_easy_unseen_validate': 86,\n 'polyp_hard_validate': 54,\n 'polyp_easy_validate': 119,\n 'Kvasir-train': 1,\n 'Mayo-train': 10,\n '300-train': 6,\n '612-train': 18,\n '300-tv': 6,\n '612-test': 5,\n '612-val': 5 \n}\n\nSET_NAME_TO_MODE = {\n 'polyp_train': 'train',\n 'polyp_hard_seen_validate': 'evaluate',\n 'polyp_hard_unseen_validate': 'evaluate',\n 'polyp_easy_seen_validate': 'evaluate',\n 'polyp_easy_unseen_validate': 'evaluate',\n 'polyp_hard_validate': 'evaluate',\n 'polyp_easy_validate': 'evaluate',\n 'Kvasir-train': 'train',\n 'Mayo-train': 'train',\n '300-train': 'train',\n '612-train': 'train',\n '300-tv': 'evaluate',\n '612-test': 'evaluate',\n '612-val': 'evaluate' \n}\n\nSET_NAME_TO_PREFIX = {\n 'polyp_train': 'polyp_train',\n 'polyp_hard_seen_validate': 'polyp_hard_seen_validate',\n 'polyp_hard_unseen_validate': 'polyp_hard_unseen_validate',\n 'polyp_easy_seen_validate': 'polyp_easy_seen_validate',\n 'polyp_easy_unseen_validate': 'polyp_easy_unseen_validate',\n 'polyp_hard_validate': 'polyp_hard_validate',\n 'polyp_easy_validate': 'polyp_easy_validate',\n 'Kvasir-train': 'Kvasir-train',\n 'Mayo-train': 'Mayo-train',\n '300-train': '300-train',\n '612-train': '612-train',\n '300-tv': '300-tv',\n '612-test': '612-test',\n '612-val': '612-val' \n}\n\nCLASS_TO_ID = {\n 'high_grade_adenoma':0, \n 'hyperplastic_polyp':1, \n 'invasive_cancer':2,\n 'low_grade_adenoma':3, \n 'sessile_serrated_lesion':4,\n 'traditional_serrated_adenoma':5\n}\n\ndef get_frames(frames_path, video_id, frames):\n return [Image.open(os.path.join(frames_path, video_id, f'{f}.jpg')).convert('RGB') for f in frames]\ndef get_frames_mask(mask_path, video_id, frames):\n # masks = [cv2.imread(os.path.join(mask_path, video_id, f'{f}.jpg')) for f in frames]\n if os.path.exists(os.path.join(mask_path, video_id, f'{frames[0]}.png')):\n masks = [Image.open(os.path.join(mask_path, video_id, f'{f}.png')).convert('L') for f in frames]\n elif os.path.exists(os.path.join(mask_path, video_id, f'{frames[0]}.jpg')):\n masks = [Image.open(os.path.join(mask_path, video_id, f'{f}.jpg')).convert('L') for f in frames]\n else:\n raise ValueError()\n masks = [np.array(mk) for mk in masks]\n masks = torch.stack([torch.from_numpy(mk) for mk in masks], dim=0) # t h w\n # assert set(masks.unique().tolist()) == set([0, 255]), f'{masks.unique().tolist()}'\n masks = (masks > 0).int()\n return masks, torch.ones(len(frames)).bool()\n\n\n"
},
{
"path": "data_schedule/vis/vis_aug_eval.py",
"content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport random\nimport torch\nimport torchvision.transforms.functional as F\nfrom data_schedule.vis.apis import VIS_Aug_CallbackAPI\nfrom .vis_aug_utils import get_tgt_size\n\nfrom .vis_aug_utils import VIS_EVAL_AUG_REGISTRY\n\nclass RandomResize:\n def __init__(self, sizes, max_size=None):\n assert isinstance(sizes, (list, tuple))\n self.sizes = sizes\n self.max_size = max_size\n\n def __call__(self, ret):\n video = ret['video']\n orig_size = video[0].size # w h\n tgt_size = get_tgt_size(video[0].size, random.choice(self.sizes), self.max_size) # h w\n\n resized_video = [F.resize(frame, tgt_size) for frame in video]\n ratio_width, ratio_height = tuple(float(s) / float(s_orig) for s, s_orig in zip(tgt_size[::-1], orig_size))\n ret['video'] = resized_video\n\n if 'callback_fns' in ret:\n VIS_Aug_CallbackAPI\n ret['callback_fns'].append(RandomResize(sizes=[orig_size], max_size=None))\n\n if \"pred_masks\" in ret:\n assert (len(self.sizes) == 1) and (self.max_size == None)\n VIS_Aug_CallbackAPI\n pred_masks = ret['pred_masks'] # list[nt h w], t\n pred_masks = [torch.nn.functional.interpolate(mk.unsqueeze(0).float(), tgt_size, mode='nearest')[0].bool()\n for mk in pred_masks]\n ret['pred_masks'] = pred_masks # list[nt h w], t\n\n if \"pred_boxes\" in ret:\n VIS_Aug_CallbackAPI\n pred_boxes = ret[\"pred_boxes\"] # list[nt 4], t\n scaled_boxes = [bx * (torch.tensor([ratio_width, ratio_height, ratio_width, ratio_height])[None, :])\n for bx in pred_boxes]\n ret[\"pred_boxes\"] = scaled_boxes\n\n return ret\n \nclass VideoToPIL:\n def __call__(self, ret):\n video = ret['video'] # t 3 h w ->\n assert video.dtype == torch.float and (video.max() <= 1) and (video.min() >=0) \n pil_video = [F.to_pil_image(frame, mode='RGB') for frame in video] # 3 h w, float, 0-1\n ret['video'] = pil_video\n assert 'callback_fns' not in ret\n return ret\n\n\nclass VideoToTensor:\n def __call__(self, ret):\n video = ret['video']\n tensor_video = torch.stack([F.to_tensor(frame) for frame in video], dim=0) # t 3 h w, float, 0-1\n ret['video'] = tensor_video\n\n if 'callback_fns' in ret:\n VIS_Aug_CallbackAPI\n ret['callback_fns'].append(VideoToPIL())\n\n return ret\n \n\n@VIS_EVAL_AUG_REGISTRY.register()\nclass WeakPolyP_EvalAug:\n def __init__(self, configs) -> None:\n self.resize = RandomResize(\n sizes=[[352, 352]],\n )\n self.tensor_video = VideoToTensor()\n\n def __call__(self, ret):\n VIS_Aug_CallbackAPI\n ret = self.resize(ret)\n ret = self.tensor_video(ret) \n return ret\n"
},
{
"path": "data_schedule/vis/vis_aug_train.py",
"content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport random\nfrom PIL import Image\nimport torch\nimport torchvision.transforms.functional as F\nfrom einops import rearrange\nfrom copy import deepcopy as dcopy\nfrom data_schedule.vis.apis import VIS_Aug_CallbackAPI\nimport albumentations as A\nimport numpy as np\nfrom data_schedule.utils.segmentation import bounding_box_from_mask\nfrom .vis_aug_utils import VIS_TRAIN_AUG_REGISTRY, pil_torch_to_numpy, numpy_to_pil_torch\n\nimport copy\n \nimport imgaug.augmenters as iaa\nfrom imgaug.augmentables.segmaps import SegmentationMapsOnImage\nfrom imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage\nimport imgaug\nfrom datetime import datetime\n\nclass RandomRotate90:\n def __init__(self) -> None:\n self.album_aug = A.ReplayCompose(\n [A.RandomRotate90(0.5)]\n )\n \n def __call__(self, ret):\n video = ret['video'] \n masks = ret['masks'] \n has_ann = ret['has_ann']\n # list[PIL], n t' h w -> \n # list[h w 3, 255rgb], t\n # list[list[h w, 01uint8]] t\n video, masks = pil_torch_to_numpy(video=video, masks=masks, has_ann=has_ann)\n replay = self.album_aug(image=video[0], mask=[masks[0][0]])['replay']\n auged_video = []\n auged_mask = []\n for vid, mk in zip(video, masks):\n ret = self.album_aug.replay(replay, image=vid, mask=mk)\n auged_video.append(ret['image'])\n auged_mask.append(ret['mask'])\n \n auged_video, auged_mask = numpy_to_pil_torch(video=auged_video, auged_mask=auged_mask, has_ann=has_ann)\n\n ret['video'] = auged_video\n ret['mask'] = auged_mask\n\n return ret\n\nclass ComputeBox:\n def __call__(self, ret):\n W, H = ret['video'][0].size\n N, T = ret['masks'].shape[:2] # n t' h w\n boxes = torch.stack([bounding_box_from_mask(mask) for mask in copy.deepcopy(ret['masks']).flatten(0, 1)], dim=0) # Nt' 4\n boxes = rearrange(boxes, '(N T) c -> N T c', N=N, T=T)\n boxes[:, :, 0::2].clamp_(min=0, max=W)\n boxes[:, :, 1::2].clamp_(min=0, max=H)\n\n ret['boxes'] = boxes\n\n return ret\n\nclass VideoToTensor:\n def __call__(self, ret):\n video = ret['video']\n tensor_video = torch.stack([F.to_tensor(frame) for frame in video], dim=0) # t 3 h w, float, 0-1\n ret['video'] = tensor_video\n return ret\n\nclass Compose:\n def __init__(self, transforms):\n self.transforms = transforms\n\n def __call__(self, ret):\n for t in self.transforms:\n ret = t(ret)\n return ret\n\n def __repr__(self):\n format_string = self.__class__.__name__ + \"(\"\n for t in self.transforms:\n format_string += \"\\n\"\n format_string += \" {0}\".format(t)\n format_string += \"\\n)\"\n return format_string\n\n\n@VIS_TRAIN_AUG_REGISTRY.register()\nclass WeakPolyP_TrainAug:\n def __init__(self, configs) -> None:\n self.transform = A.ReplayCompose([\n A.Resize(352, 352),\n A.HorizontalFlip(p=0.5),\n A.VerticalFlip(p=0.5),\n A.RandomRotate90(p=0.5),\n ])\n\n self.tensor_video = VideoToTensor()\n self.add_box = ComputeBox()\n\n def __call__(self, ret):\n VIS_Aug_CallbackAPI\n video = ret['video'] \n masks = ret['masks'] # n t' h w\n has_ann = ret['has_ann'] # t\n # list[PIL] -> list[h w 3, 0-1float], t\n # n t' h w -> list[list[h w, 01uint8], 没有annotation的帧box是空] t\n video, masks = pil_torch_to_numpy(video=video, masks=masks, has_ann=has_ann)\n\n replay = self.transform(image=video[0], masks=[masks[0][0]])['replay']\n auged_video = []\n auged_mask = []\n for vid, mk in zip(video, masks):\n auged_each_frame = self.transform.replay(replay, image=vid, masks=mk)\n auged_video.append(auged_each_frame['image'])\n auged_mask.append(auged_each_frame['masks']) # list[h w, 01uint8]\n \n auged_video, auged_mask = numpy_to_pil_torch(video=auged_video, masks=auged_mask, has_ann=has_ann)\n\n ret['video'] = auged_video\n ret['masks'] = auged_mask\n \n ret = self.add_box(ret)\n ret = self.tensor_video(ret)\n\n return ret\n\n\n@VIS_TRAIN_AUG_REGISTRY.register()\nclass WeakPolyP_TrainAug_RotateImageToClip:\n def __init__(self, configs) -> None:\n self.ImageToSeqAugmenter = ImageToSeqAugmenter(perspective=True, affine=True, motion_blur=True,\n rotation_range=(-20, 20), perspective_magnitude=0.08,\n hue_saturation_range=(-5, 5), brightness_range=(-40, 40),\n motion_blur_prob=0.25, motion_blur_kernel_sizes=(9, 11),\n translate_range=(-0.1, 0.1))\n self.num_frames = configs['num_frames']\n \n self.transform = A.ReplayCompose([\n A.Resize(352, 352),\n A.HorizontalFlip(p=0.5),\n A.VerticalFlip(p=0.5),\n A.RandomRotate90(p=0.5),\n ])\n self.tensor_video = VideoToTensor()\n self.add_box = ComputeBox()\n\n def apply_random_sequence_shuffle(self, images, instance_masks):\n perm = list(range(self.num_frames))\n random.shuffle(perm)\n images = [images[i] for i in perm]\n instance_masks = [instance_masks[i] for i in perm]\n return images, instance_masks\n \n def __call__(self, ret):\n VIS_Aug_CallbackAPI\n video = ret['video'] # list[pil], t\n masks = ret['masks'] # n t' h w\n has_ann = ret['has_ann'] # t\n\n # list[PIL] -> list[h w 3, uint8], t\n # n t' h w -> list[list[h w], n, uint8], t\n seq_images, seq_instance_masks = pil_torch_to_numpy(video=video, masks=masks, has_ann=has_ann, float_image=False)\n assert len(seq_images) == 1 and len(seq_instance_masks) == 1\n static_img, static_mask = seq_images[0], seq_instance_masks[0]\n for t in range(self.num_frames - 1):\n im_trafo, instance_masks_trafo = self.ImageToSeqAugmenter(static_img, static_mask) # h w 3, uint8; list[h w], n, uint8\n seq_images.append(np.uint8(im_trafo))\n seq_instance_masks.append(instance_masks_trafo)\n # list[h w 3], t ; # list[list[h w, 01uint8]] t \n seq_images, seq_instance_masks = self.apply_random_sequence_shuffle(seq_images, seq_instance_masks) \n has_ann = torch.ones(self.num_frames).bool() # T\n seq_images = [np.float32(haosen) / 255.0 for haosen in seq_images] # list[h w 3, 0-1float], t\n replay = self.transform(image=seq_images[0], masks=[seq_instance_masks[0][0]])['replay']\n auged_video = []\n auged_mask = []\n for vid, mk in zip(seq_images, seq_instance_masks):\n auged_each_frame = self.transform.replay(replay, image=vid, masks=mk)\n auged_video.append(auged_each_frame['image'])\n auged_mask.append(auged_each_frame['masks']) # list[h w, 01uint8]\n \n auged_video, auged_mask = numpy_to_pil_torch(video=auged_video, masks=auged_mask, has_ann=has_ann) # n t h w\n # [haosen.save(f'./test{idx}.png') for idx, haosen in enumerate(auged_video)]\n # import matplotlib.pyplot as plt\n # [plt.imsave( f'./mask{idx}.png', auged_mask[0][idx].float().numpy()) for idx in range(len(auged_mask[0]))]\n ret['video'] = auged_video\n ret['masks'] = auged_mask\n ret['has_ann'] = has_ann\n \n ret = self.add_box(ret)\n ret = self.tensor_video(ret)\n\n return ret\n \n\nclass ImageToSeqAugmenter(object):\n def __init__(self, perspective=True, affine=True, motion_blur=True,\n brightness_range=(-50, 50), hue_saturation_range=(-15, 15), perspective_magnitude=0.12,\n scale_range=1.0, translate_range={\"x\": (-0.15, 0.15), \"y\": (-0.15, 0.15)}, rotation_range=(-20, 20),\n motion_blur_kernel_sizes=(7, 9), motion_blur_prob=0.5, seed=2024):\n\n self.basic_augmenter = iaa.SomeOf((1, None), [\n iaa.Add(brightness_range),\n iaa.AddToHueAndSaturation(hue_saturation_range)\n ]\n )\n\n transforms = []\n if perspective:\n transforms.append(iaa.PerspectiveTransform(perspective_magnitude))\n if affine:\n transforms.append(iaa.Affine(scale=scale_range,\n translate_percent=translate_range,\n rotate=rotation_range,\n order=1, # cv2.INTER_LINEAR\n backend='auto'))\n transforms = iaa.Sequential(transforms)\n transforms = [transforms]\n\n if motion_blur:\n blur = iaa.Sometimes(motion_blur_prob, iaa.OneOf(\n [\n iaa.MotionBlur(ksize)\n for ksize in motion_blur_kernel_sizes\n ]\n ))\n transforms.append(blur)\n\n self.frame_shift_augmenter = iaa.Sequential(transforms)\n self.seed = seed\n @staticmethod\n def condense_masks(instance_masks):\n condensed_mask = np.zeros_like(instance_masks[0], dtype=np.int8)\n for instance_id, mask in enumerate(instance_masks, 1):\n condensed_mask = np.where(mask, instance_id, condensed_mask)\n\n return condensed_mask\n\n @staticmethod\n def expand_masks(condensed_mask, num_instances):\n return [(condensed_mask == instance_id).astype(np.uint8) for instance_id in range(1, num_instances + 1)]\n\n def __call__(self, image, masks=None, boxes=None): # n h w\n det_augmenter = self.frame_shift_augmenter.to_deterministic()\n if masks is not None:\n masks_np, is_binary_mask = [], []\n boxs_np = []\n\n for mask in masks:\n \n if isinstance(mask, np.ndarray):\n masks_np.append(mask.astype(np.bool_))\n is_binary_mask.append(False)\n else:\n raise ValueError(\"Invalid mask type: {}\".format(type(mask)))\n\n num_instances = len(masks_np)\n masks_np = SegmentationMapsOnImage(self.condense_masks(masks_np), shape=image.shape[:2])\n # boxs_np = BoundingBoxesOnImage(boxs_np, shape=image.shape[:2])\n seed = int(datetime.now().strftime('%M%S%f')[-8:])\n imgaug.seed(seed)\n aug_image, aug_masks = det_augmenter(image=self.basic_augmenter(image=image) , segmentation_maps=masks_np)\n imgaug.seed(seed)\n invalid_pts_mask = det_augmenter(image=np.ones(image.shape[:2] + (1,), np.uint8)).squeeze(2)\n aug_masks = self.expand_masks(aug_masks.get_arr(), num_instances)\n # aug_boxes = aug_boxes.remove_out_of_image().clip_out_of_image()\n aug_masks = [mask for mask, is_bm in zip(aug_masks, is_binary_mask)]\n return aug_image, aug_masks #, aug_boxes.to_xyxy_array()\n\n else:\n masks = [SegmentationMapsOnImage(np.ones(image.shape[:2], np.bool), shape=image.shape[:2])]\n aug_image, invalid_pts_mask = det_augmenter(image=image, segmentation_maps=masks)\n return aug_image, invalid_pts_mask.get_arr() == 0\n\n"
},
{
"path": "data_schedule/vis/vis_aug_utils.py",
"content": "from detectron2.utils.registry import Registry\nimport torch\nimport numpy as np\nimport torchvision.transforms.functional as F\nfrom PIL import Image\nfrom einops import rearrange, reduce, repeat\nVIS_EVAL_AUG_REGISTRY = Registry('VIS_EVAL_AUG')\nVIS_TRAIN_AUG_REGISTRY = Registry('VIS_TRAIN_AUG')\n\n\ndef get_size_with_aspect_ratio(image_size, size, max_size=None):\n w, h = image_size\n if max_size is not None:\n min_original_size = float(min((w, h)))\n max_original_size = float(max((w, h)))\n if max_original_size / min_original_size * size > max_size:\n size = int(round(max_size * min_original_size / max_original_size))\n\n if (w <= h and w == size) or (h <= w and h == size):\n return (h, w)\n\n if w < h:\n ow = size\n oh = int(size * h / w)\n else:\n oh = size\n ow = int(size * w / h)\n\n return (oh, ow)\n\ndef get_tgt_size(image_size, size, max_size=None):\n if isinstance(size, (list, tuple)):\n return size[::-1]\n else:\n return get_size_with_aspect_ratio(image_size, size, max_size)\n\ndef pil_torch_to_numpy(video, masks, has_ann, float_image=True):\n # n t' h w\n # list[pil_image, rgb], t\n # t\n N, T = masks.shape[:2]\n has_ann_idx = torch.where(has_ann)[0] # time_idx\n # list[Image], t -> list[h w 3, 255uint8], t\n masks = masks.permute(1, 0, 2, 3).contiguous().unbind(0) # list[n h w] t'\n numpy_masks = [[]] * len(has_ann) # list[list[h w, 01_uint8], n], t\n assert len(has_ann_idx) == len(masks)\n for fmask, taylor in zip(masks, has_ann_idx): # n h w\n fnumpy_masks = []\n for mk in fmask.unbind(0): # h w\n fnumpy_masks.append(mk.numpy().astype(np.uint8))\n numpy_masks[taylor] = fnumpy_masks\n \n if float_image:\n # list[h w 3, 0-1float], t\n video = [F.to_tensor(frame).permute(1,2,0).numpy() for frame in video]\n else:\n # uint8\n video = [np.array(frame) for frame in video]\n \n return video, numpy_masks\n\ndef numpy_to_pil_torch(video, masks, has_ann):\n # numpy, numpy -> torch, torch\n # list[h w 3, 0-1float], t\n H, W = video[0].shape[:2]\n T = has_ann.int().sum()\n video = [Image.fromarray(np.uint8(aug_vid * 255), mode=\"RGB\") for aug_vid in video]\n # t'n h w\n torch_masks = torch.stack([torch.from_numpy(obj_mk).bool() for frame_mk in masks for obj_mk in frame_mk], dim=0)\n torch_masks = rearrange(torch_masks, '(T N) h w -> N T h w', T=T) # n t' h w\n\n return video, torch_masks \n\n\n\n\n\n\n\n"
},
{
"path": "data_schedule/vis/vis_frame_sampler.py",
"content": "\nfrom detectron2.utils.registry import Registry\nimport random\nimport numpy as np\nimport torch\nimport logging\nfrom detectron2.utils import comm\n\nVIS_FRAMES_SAMPLER_REGISTRY = Registry('VIS_FRAMES_SAMPLER')\n\nimport random\n\n@VIS_FRAMES_SAMPLER_REGISTRY.register()\nclass Naive_ReferenceFrame_FrameSampler:\n def __init__(self, sampler_configs, dataset_meta, **kwargs): \n self.reference_frame_step_size = dataset_meta.get('step_size')\n\n self.clip_sizes = list(sampler_configs['clip_sizes']) # list[int]\n self.clip_distribute = sampler_configs['clip_distribute'] # dense, sparse, local_global\n self.clip_position = sampler_configs['clip_position'] # former, center, latter\n\n if max(self.clip_sizes) > self.reference_frame_step_size:\n if comm.is_main_process():\n logging.warning('')\n\n def __call__(self, \n frame_idx=None,\n all_frames=None, # list[str]\n **kwargs):\n random_clip_size = random.choice(self.clip_sizes)\n video_len = len(all_frames)\n sample_indx = [frame_idx]\n if (self.clip_position == 'center') and (self.clip_distribute == 'local_global'):\n if random_clip_size != 1:\n sample_id_before = random.randint(1, 3)\n sample_id_after = random.randint(1, 3)\n local_indx = [max(0, frame_idx - sample_id_before), min(video_len - 1, frame_idx + sample_id_after)]\n sample_indx.extend(local_indx)\n if random_clip_size > 3:\n all_inds = list(range(video_len))\n global_inds = all_inds[:min(sample_indx)] + all_inds[max(sample_indx):]\n global_n = random_clip_size - len(sample_indx)\n if len(global_inds) > global_n:\n select_id = random.sample(range(len(global_inds)), global_n)\n for s_id in select_id:\n sample_indx.append(global_inds[s_id])\n elif video_len >= global_n: \n select_id = random.sample(range(video_len), global_n)\n for s_id in select_id:\n sample_indx.append(all_inds[s_id])\n else:\n select_id = random.sample(range(video_len), global_n - video_len) + list(range(video_len)) \n for s_id in select_id: \n sample_indx.append(all_inds[s_id])\n \n elif (self.clip_position == 'center') and (self.clip_distribute == 'dense'):\n half_size = (random_clip_size - 1) // 2\n sample_indx += list(range(frame_idx - half_size, frame_idx))\n sample_indx += list(range(frame_idx+1, half_size + frame_idx + 1))\n\n if len(sample_indx) < random_clip_size: \n sample_indx = [min(sample_indx)] + sample_indx\n assert len(sample_indx) == random_clip_size\n sample_indx = torch.tensor(sample_indx)\n sample_indx = sample_indx.clamp_(min=0, max=video_len-1)\n sample_indx = sample_indx.tolist()\n else:\n raise ValueError()\n sample_indx.sort()\n sampled_frames = [all_frames[idx] for idx in sample_indx]\n return sampled_frames\n\n"
},
{
"path": "handle_vps.py",
"content": "\nimport cv2\nimport numpy as np\nimport os\nimport shutil\nfrom PIL import Image\nimport torch\nfrom tqdm import tqdm\ndataset_root = os.getenv('DATASET_PATH')\n# the original IVPS is the union of Kvasir and per-frame Mayo/CVC\nall_images = os.listdir(f'{dataset_root}/MICCAI-VPS-dataset/IVPS-TrainSet/Frame')\nka_images = [b for b in all_images if b.startswith('K')]\nassert len(ka_images) == 1000\n\nall_gts = os.listdir(f'{dataset_root}/MICCAI-VPS-dataset/IVPS-TrainSet/GT')\nka_gts = [b for b in all_gts if b.startswith('K')]\nassert len(ka_gts) == 1000\n\nos.makedirs(os.path.join(f'{dataset_root}/MICCAI-VPS-dataset/Kvasir-SEG/Frame/1'),exist_ok=True)\nos.makedirs(os.path.join(f'{dataset_root}/MICCAI-VPS-dataset/Kvasir-SEG/GT/1'),exist_ok=True)\n\nfor image_id in tqdm(ka_images):\n shutil.copy(os.path.join(f'{dataset_root}/MICCAI-VPS-dataset/IVPS-TrainSet/Frame', f'{image_id}'),\n os.path.join(f'{dataset_root}/MICCAI-VPS-dataset/Kvasir-SEG/Frame/1', f'{image_id}'),)\n \nfor image_id in tqdm(ka_gts):\n shutil.copy(os.path.join(f'{dataset_root}/MICCAI-VPS-dataset/IVPS-TrainSet/GT', f'{image_id}'),\n os.path.join(f'{dataset_root}/MICCAI-VPS-dataset/Kvasir-SEG/GT/1', f'{image_id}'),) \n\n# normalize train directory\nfor base_path in [f'{dataset_root}/MICCAI-VPS-dataset/VPS-TrainSet/CVC-ColonDB-300/Train',\n f'{dataset_root}/MICCAI-VPS-dataset/VPS-TrainSet/ASU-Mayo_Clinic/Train',\n f'{dataset_root}/MICCAI-VPS-dataset/VPS-TrainSet/CVC-ClinicDB-612/Train']:\n video_ids = os.listdir(base_path)\n frame_path = os.path.join(base_path, 'Frame')\n gt_path = os.path.join(base_path, 'GT')\n os.makedirs(frame_path, exist_ok=True)\n os.makedirs(gt_path, exist_ok=True)\n # Iterate through each video ID directory\n for vid in video_ids:\n shutil.copytree(os.path.join(base_path, vid, 'Frame'), os.path.join(frame_path, vid))\n shutil.copytree(os.path.join(base_path, vid, 'GT'), os.path.join(gt_path, vid))\n # TODO: dangerous: remove if you want\n\n# remove non-mask frames of each training set\nSET_NAME = [\n 'Kvasir-train',\n 'Mayo-train',\n '300-train',\n '612-train',\n]\n\nSET_NAME_TO_DIR = {\n 'Kvasir-train': 'MICCAI-VPS-dataset/Kvasir-SEG',\n 'Mayo-train': 'MICCAI-VPS-dataset/VPS-TrainSet/ASU-Mayo_Clinic/Train',\n '300-train': 'MICCAI-VPS-dataset/VPS-TrainSet/CVC-ColonDB-300/Train',\n '612-train': 'MICCAI-VPS-dataset/VPS-TrainSet/CVC-ClinicDB-612/Train',\n}\n\nSET_NAME_TO_NUM_VIDEOS = {\n 'Kvasir-train': 1,\n 'Mayo-train': 10,\n '300-train': 6,\n '612-train': 18,\n '300-tv': 6,\n '612-test': 5,\n '612-val': 5 \n}\n\n\nSET_NAME_TO_PREFIX = {\n 'Kvasir-train': 'Kvasir-train',\n 'Mayo-train': 'Mayo-train',\n '300-train': '300-train',\n '612-train': '612-train',\n}\n\n\nroot = os.getenv('DATASET_PATH')\ndef get_frames_mask(mask_path, video_id, frames):\n # masks = [cv2.imread(os.path.join(mask_path, video_id, f'{f}.jpg')) for f in frames]\n if os.path.exists(os.path.join(mask_path, video_id, f'{frames[0]}.png')):\n masks = [Image.open(os.path.join(mask_path, video_id, f'{f}.png')).convert('L') for f in frames]\n elif os.path.exists(os.path.join(mask_path, video_id, f'{frames[0]}.jpg')):\n masks = [Image.open(os.path.join(mask_path, video_id, f'{f}.jpg')).convert('L') for f in frames]\n else:\n raise ValueError()\n masks = [np.array(mk) for mk in masks]\n masks = torch.stack([torch.from_numpy(mk) for mk in masks], dim=0) # t h w\n # assert set(masks.unique().tolist()) == set([0, 255]), f'{masks.unique().tolist()}'\n masks = (masks > 0).int()\n return masks, torch.ones(len(frames)).bool()\n \nnum_delted_frames = 0\nfor train_set_name in SET_NAME:\n set_dir = SET_NAME_TO_DIR[train_set_name]\n frames_dir = os.path.join(root, set_dir, 'Frame')\n mask_dir = os.path.join(root, set_dir, 'GT')\n\n video_ids = os.listdir(frames_dir)\n for vid in tqdm(video_ids):\n frames = [haosen[:-4] for haosen in os.listdir(os.path.join(frames_dir, vid))]\n frame_has_fore = [get_frames_mask(mask_dir, vid, [haosen])[0].any() for haosen in tqdm(frames)] # list[t]\n assert len(frame_has_fore) == len(frames)\n num_delted_frames += (~ torch.tensor(frame_has_fore)).int().sum()\n for haosen, frame_name in tqdm(zip(frame_has_fore, frames)):\n if not haosen:\n os.remove(os.path.join(frames_dir, vid, f'{frame_name}.jpg'))\n\n if os.path.exists(os.path.join(mask_dir, vid, f'{frame_name}.jpg')):\n os.remove(os.path.join(mask_dir, vid, f'{frame_name}.jpg'))\n elif os.path.exists(os.path.join(mask_dir, vid, f'{frame_name}.png')):\n os.remove(os.path.join(mask_dir, vid, f'{frame_name}.png')) \n else:\n raise ValueError()\n\nprint(f'should be {num_delted_frames}/1546.') # should be 1546\n\n\n\n\n"
},
{
"path": "main.py",
"content": "import os\nimport argparse\nimport logging\nimport importlib\nfrom trainers import task_to_trainer\nimport detectron2.utils.comm as comm\nfrom termcolor import colored\nimport logging\nimport yaml\nimport torch\nfrom utils.misc import setup_for_distributed\n\ndef _highlight(code, filename):\n try:\n import pygments\n except ImportError:\n return code\n\n from pygments.lexers import Python3Lexer, YamlLexer\n from pygments.formatters import Terminal256Formatter\n\n lexer = Python3Lexer() if filename.endswith(\".py\") else YamlLexer()\n code = pygments.highlight(code, lexer, Terminal256Formatter(style=\"monokai\"))\n return code\n\nclass _ColorfulFormatter(logging.Formatter):\n def __init__(self, *args, **kwargs):\n self._root_name = kwargs.pop(\"root_name\") + \".\"\n self._abbrev_name = kwargs.pop(\"abbrev_name\", \"\")\n if len(self._abbrev_name):\n self._abbrev_name = self._abbrev_name + \".\"\n super(_ColorfulFormatter, self).__init__(*args, **kwargs)\n def formatMessage(self, record):\n record.name = record.name.replace(self._root_name, self._abbrev_name)\n message = record.message\n # message, asctime, name, filename = record.message, record.asctime, record.name, record.filename\n log = super(_ColorfulFormatter, self).formatMessage(record)\n if (record.levelno == logging.WARNING) or (record.levelno == logging.ERROR) or (record.levelno == logging.CRITICAL):\n colored_message = colored(message, \"red\", attrs=[\"blink\", \"underline\"])\n elif record.levelno == logging.DEBUG:\n colored_message = colored(message, \"yellow\", attrs=[\"blink\", \"underline\"])\n else: # INFO/NOTSET\n colored_message = colored(message, \"white\") \n return log + colored_message\n\ndef set_logging_file(output_dir, file_name, mode='a'):\n handler1 = logging.StreamHandler()\n handler2 = logging.FileHandler(os.path.join(output_dir, file_name), mode=mode)\n formatter = _ColorfulFormatter(\n colored(\"[%(asctime)s %(name)s %(filename)s]: \", \"green\"),\n datefmt=\"%m/%d %H:%M:%S\",\n root_name=os.path.join(output_dir, file_name),\n abbrev_name=str('grey'),\n )\n handler1.setFormatter(formatter)\n handler2.setFormatter(formatter)\n logger = logging.getLogger()\n logger.addHandler(handler1)\n logger.addHandler(handler2)\n logger.setLevel(logging.DEBUG)\n\n\ndef init_process_group_and_set_device(world_size, process_id, device_id):\n \"\"\"\n This function needs to be called on each spawned process to initiate learning using DistributedDataParallel.\n The function initiates the process' process group and assigns it a single GPU to use during training.\n \"\"\"\n torch.cuda.set_device(device_id)\n device = torch.device(f'cuda:{device_id}')\n if world_size > 1:\n torch.distributed.init_process_group(\n torch.distributed.Backend.NCCL,\n world_size=world_size,\n rank=process_id\n )\n comm.create_local_process_group(world_size)\n torch.distributed.barrier(device_ids=[device_id])\n setup_for_distributed(process_id == 0)\n return device\n\ndef run(rank, configs, world_size):\n os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'\n os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n os.environ['PYDEVD_WARN_SLOW_RESOLVE_TIMEOUT'] = \"4\"\n os.environ['PYDEVD_DISABLE_FILE_VALIDATION'] = \"1\"\n os.environ[\"DGLBACKEND\"] = \"pytorch\"\n logging.getLogger('penman').setLevel(logging.WARNING) \n logging.getLogger('PIL').setLevel(logging.WARNING) \n logging.getLogger('PIL.PngImagePlugin').setLevel(logging.WARNING) \n logging.getLogger('matplotlib').setLevel(logging.WARNING)\n logging.getLogger('urllib3').setLevel(logging.WARNING)\n logging.getLogger('h5py').setLevel(logging.WARNING)\n init_process_group_and_set_device(world_size, process_id=rank, device_id=rank)\n if comm.is_main_process():\n mode = configs['trainer_mode']\n out_dir = configs['out_dir']\n if mode == 'eval':\n num_of_eval_times = len([eval_txt for eval_txt in os.listdir(out_dir) if eval_txt.endswith('eval.txt')])\n set_logging_file(out_dir, f\"eval.txt\", mode='a')\n path = os.path.join(out_dir, f\"config_eval.yaml\")\n else:\n num_of_train_times = len([train_txt for train_txt in os.listdir(out_dir) if train_txt.endswith('train.txt')])\n if 'resume' in mode:\n set_logging_file(out_dir, f\"train.txt\", mode='a')\n else:\n set_logging_file(out_dir, f\"train.txt\", mode='w')\n path = os.path.join(out_dir, f\"config_train.yaml\")\n \n logging.debug(\"Running with full config:\\n{}\".format(_highlight(yaml.dump(configs, default_flow_style=False), \".yaml\")))\n with open(path, \"w\") as f:\n f.write(yaml.dump(configs, default_flow_style=False))\n logging.debug(\"Full config saved to {}\".format(path)) \n comm.synchronize()\n trainer = task_to_trainer[configs['task']](configs=configs)\n comm.synchronize()\n if configs['trainer_mode'] == 'eval':\n eval_ckpts = configs['eval_ckpts']\n for lunch in eval_ckpts:\n trainer.load_ckpt(lunch, load_model=True, load_schedule=True, load_random=False, load_optimize=False)\n trainer.evaluate()\n\n else:\n if configs['trainer_mode'] == 'train_resume':\n ckpt_dirs = os.listdir(configs['out_dir'])\n ckpt_dirs = sorted([a for a in ckpt_dirs if a.startswith('epc')], key=lambda x:int(x.split('sap[')[-1][:-1]))\n trainer_ckpt = '/'.join([configs['out_dir'], ckpt_dirs[-1], 'ckpt.pth.tar'])\n trainer.load_ckpt(trainer_ckpt, load_model=True, load_schedule=True, load_random=True, load_optimize=True)\n trainer.train()\n\nif __name__==\"__main__\":\n\n parser = argparse.ArgumentParser()\n parser.add_argument('--config_file', type=str, required=True)\n parser.add_argument('--trainer_mode', type=str, default='train_attmpt') \n parser.add_argument('--eval_path', type=str, default='') \n args = parser.parse_args()\n task, group, config, config2 = args.config_file.split('/')[-4:]\n assert config == config2[:-3]\n config_file = '.'.join(['output', task, group, config, config])\n configs = importlib.import_module(config_file).trainer_configs\n configs['task'], configs['group'], configs['config'] = task, group, config\n configs['out_dir'] = os.path.join('./', 'output', task, group, config)\n configs['trainer_mode'] = args.trainer_mode\n\n if configs['trainer_mode'] == 'eval':\n eval_ckpts = []\n eval_path = args.eval_path\n assert eval_path != '', f'eval path is none'\n \n if os.path.isfile(eval_path):\n eval_ckpts.append(eval_path)\n\n elif os.path.isdir(eval_path):\n ckpt_dirs = os.listdir(eval_path) \n ckpt_dirs = [taylor for taylor in ckpt_dirs if os.path.isdir(os.path.join(eval_path, taylor))]\n # epc[1]_iter[5000]_sap[60009]\n ckpt_dirs = sorted([billie for billie in ckpt_dirs if billie.startswith('epc')], key=lambda x:int(x.split('sap[')[-1][:-1]))\n eval_ckpts = [os.path.join(eval_path, cd, f'ckpt.pth.tar') for cd in ckpt_dirs]\n eval_ckpts = [eval_c for eval_c in eval_ckpts if os.path.exists(eval_c)]\n else:\n raise ValueError()\n configs['eval_ckpts'] = eval_ckpts\n else:\n pass\n \n gpu_ids = list(os.environ['CUDA_VISIBLE_DEVICES'].split(','))\n assert len(set(gpu_ids)) == len(gpu_ids)\n gpu_ids = list(range(len(gpu_ids)))\n \n if len(gpu_ids) > 1:\n torch.multiprocessing.spawn(run, nprocs=len(gpu_ids), args=(configs, len(gpu_ids)))\n elif len(gpu_ids) == 1:\n run(rank=0, configs=configs, world_size=len(gpu_ids))"
},
{
"path": "models/VIS/BackboneEncoderDecoder_WithScaleConsistency.py",
"content": "import matplotlib.pyplot as plt\nfrom typing import Any, Optional, List, Dict, Set, Callable\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom data_schedule import build_schedule\nfrom torch import Tensor\nfrom einops import repeat, rearrange, reduce\nfrom functools import partial\nfrom einops.layers.torch import Rearrange\nfrom torch import einsum\nimport numpy as np\nimport logging\nfrom data_schedule.vis.apis import VIS_TrainAPI_clipped_video, VIS_Aug_CallbackAPI\nfrom data_schedule.vis.apis import VIS_EvalAPI_clipped_video_request_ann\nimport torchvision.transforms.functional as Trans_F\nimport copy\nfrom models.registry import register_model\nfrom models.optimization.optimizer import get_optimizer\nfrom models.optimization.scheduler import build_scheduler \nfrom models.backbone.utils import VideoMultiscale_Shape\nfrom detectron2.modeling import BACKBONE_REGISTRY, META_ARCH_REGISTRY\n\nclass BackboneEncoderDecoder_WithScaleConsistency(nn.Module):\n def __init__(\n self,\n configs,\n pixel_mean = [0.485, 0.456, 0.406],\n pixel_std = [0.229, 0.224, 0.225],):\n super().__init__()\n self.register_buffer(\"pixel_mean\", torch.Tensor(pixel_mean).view(-1, 1, 1), False) # 3 1 1\n self.register_buffer(\"pixel_std\", torch.Tensor(pixel_std).view(-1, 1, 1), False)\n self.loss_weight = configs['model']['loss_weight']\n video_backbone_configs = configs['model']['video_backbone'] \n video_backbone_cls = BACKBONE_REGISTRY.get(video_backbone_configs['name'])\n self.video_backbone = video_backbone_cls(video_backbone_configs)\n self.max_stride = self.video_backbone.max_stride\n\n self.fusion_encoder = META_ARCH_REGISTRY.get(configs['model']['fusion']['name'])(configs['model']['fusion'],\n multiscale_shapes=self.video_backbone.multiscale_shapes)\n \n same_dim_multiscale_shapes = VideoMultiscale_Shape.set_multiscale_same_dim(shape_by_dim=self.video_backbone.multiscale_shapes,\n same_dim=configs['model']['fusion']['d_model']) \n \n self.decoder = META_ARCH_REGISTRY.get(configs['model']['decoder']['name'])(configs['model']['decoder'],\n multiscale_shapes=same_dim_multiscale_shapes)\n if configs['model']['fusion']['name'] == 'Video_Deform2D_DividedTemporal_MultiscaleEncoder_v2':\n self.fusion_encoder.hack_ref(query_norm=self.decoder.temporal_query_norm, mask_mlp=self.decoder.query_mask)\n \n self.test_clip_size = configs['model']['test_clip_size']\n @property\n def device(self):\n return self.pixel_mean.device\n \n def model_preds(self, videos, video_aux_dict,):\n if (not self.training) and (self.test_clip_size is not None):\n nf = videos.shape[2]\n clip_outputs = [] # list[dict]\n for start_idx in range(0, nf, self.test_clip_size):\n multiscales = self.video_backbone(x=videos[:, :, start_idx:(start_idx + self.test_clip_size)]) # b c t h w\n multiscales = self.fusion_encoder(multiscales, video_aux_dict=video_aux_dict) \n clip_outputs.append(self.decoder(multiscales, video_aux_dict=video_aux_dict)[-1]) # b t nq h w\n return [{\n 'pred_masks': torch.cat([haosen['pred_masks'] for haosen in clip_outputs], dim=1), # b t n h w\n 'pred_class': torch.cat([haosen['pred_class'] for haosen in clip_outputs], dim=1),\n }]\n # b 3 t h w -> b 3 t h w\n multiscales = self.video_backbone(x=videos) # b c t h w\n multiscales = self.fusion_encoder(multiscales, video_aux_dict=video_aux_dict) \n return self.decoder(multiscales, video_aux_dict=video_aux_dict)\n\n def forward(self, batch_dict):\n assert self.training\n VIS_TrainAPI_clipped_video\n videos = batch_dict['video_dict']['videos'] \n targets = batch_dict['targets']\n batch_size, nf = videos.shape[:2]\n videos = (videos - self.pixel_mean) / self.pixel_std\n size1 = np.random.choice([256, 288, 320, 352, 384, 416, 448])\n vid_1 = F.interpolate(videos.flatten(0, 1), size=size1, mode='bilinear')\n vid_1 = rearrange(vid_1, '(b T) c h w -> b c T h w',b=batch_size, T=nf)\n pred1 = self.model_preds(vid_1, video_aux_dict=batch_dict['video_dict']) # {pred_masks: b 1 t h w}\n pred1_loss = self.decoder.compute_loss(pred1, targets=targets, frame_targets=batch_dict['frame_targets'],\n video_aux_dict=batch_dict['video_dict'])\n loss_value_dict = {key: pred1_loss[key] for key in list(self.loss_weight.keys())}\n return loss_value_dict, self.loss_weight\n\n @torch.no_grad()\n def sample(self, batch_dict):\n assert not self.training\n VIS_EvalAPI_clipped_video_request_ann\n videos = batch_dict['video_dict']['videos'] # b t 3 h w, 0-1\n orig_t, _, orig_h, orig_w = batch_dict['video_dict']['orig_sizes'][0]\n videos = (videos - self.pixel_mean) / self.pixel_std\n assert videos.shape[0] == 1\n batch_size, T, _, H, W = videos.shape\n videos = videos.permute(0, 2, 1,3,4) # b c t h w\n decoder_output = self.model_preds(videos, video_aux_dict=batch_dict['video_dict']) # {pred_masks: b 1 t h w}\n if isinstance(decoder_output, list):\n decoder_output = decoder_output[-1]\n pred_masks = decoder_output['pred_masks'][0] # T n h w\n pred_masks = F.interpolate(pred_masks, size=(H, W), mode='bilinear') > 0 # T n h w\n pred_masks = pred_masks[:orig_t, :, :orig_h, :orig_w] # T n h w\n #\n pred_classes = decoder_output['pred_class'][0][:orig_t, :,:] # T n c, probability\n pred_classes = pred_classes.cpu().unbind(0) # list[n c], T\n pred_masks = pred_masks.cpu().unbind(0) # list[n h w], T\n\n VIS_Aug_CallbackAPI\n orig_video = videos[0][:, :orig_t, :orig_h, :orig_w].permute(1,0,2,3) # T 3 h w\n orig_video = Trans_F.normalize(orig_video, [0, 0, 0], 1 / self.pixel_std)\n orig_video = Trans_F.normalize(orig_video, -self.pixel_mean, [1, 1, 1]).cpu()\n\n return {\n 'video': [orig_video], # [t 3 h w], 1\n 'pred_masks': [pred_masks], # [list[n h w], t, bool], 1\n 'pred_class': [pred_classes], # [list[n c], t, probability], 1\n }\n\n @staticmethod\n def get_optim_params_group(model, configs):\n weight_decay_norm = configs['optim']['weight_decay_norm']\n weight_decay_embed = configs['optim']['weight_decay_embed']\n\n defaults = {}\n defaults['lr'] = configs['optim']['base_lr']\n defaults['weight_decay'] = configs['optim']['weight_decay']\n\n norm_module_types = (\n torch.nn.BatchNorm1d,\n torch.nn.BatchNorm2d,\n torch.nn.BatchNorm3d,\n torch.nn.SyncBatchNorm,\n # NaiveSyncBatchNorm inherits from BatchNorm2d\n torch.nn.GroupNorm,\n torch.nn.InstanceNorm1d,\n torch.nn.InstanceNorm2d,\n torch.nn.InstanceNorm3d,\n torch.nn.LayerNorm,\n torch.nn.LocalResponseNorm,\n ) \n params: List[Dict[str, Any]] = []\n memo: Set[torch.nn.parameter.Parameter] = set()\n log_lr_group_idx = {'backbone':None, 'base':None}\n\n for module_name, module in model.named_modules():\n for module_param_name, value in module.named_parameters(recurse=False):\n if not value.requires_grad:\n continue\n # Avoid duplicating parameters\n if value in memo:\n continue\n memo.add(value)\n hyperparams = copy.copy(defaults)\n if \"video_backbone\" in module_name:\n hyperparams[\"lr\"] = hyperparams[\"lr\"] * configs['optim']['backbone_lr_multiplier'] \n if log_lr_group_idx['backbone'] is None:\n log_lr_group_idx['backbone'] = len(params)\n\n else:\n if log_lr_group_idx['base'] is None:\n log_lr_group_idx['base'] = len(params)\n \n # pos_embed, norm, embedding的weight decay特殊对待\n if (\n \"relative_position_bias_table\" in module_param_name\n or \"absolute_pos_embed\" in module_param_name\n ):\n logging.debug(f'setting weight decay of {module_name}.{module_param_name} to zero')\n hyperparams[\"weight_decay\"] = 0.0\n if isinstance(module, norm_module_types):\n hyperparams[\"weight_decay\"] = weight_decay_norm\n if isinstance(module, torch.nn.Embedding):\n hyperparams[\"weight_decay\"] = weight_decay_embed\n params.append({\"params\": [value], **hyperparams})\n \n return params, log_lr_group_idx\n\n@register_model\ndef backbone_encoder_decoder_withScaleConsistency(configs, device):\n from .aux_mapper import AUXMapper_v1\n model = BackboneEncoderDecoder_WithScaleConsistency(configs)\n model.to(device)\n params_group, log_lr_group_idx = BackboneEncoderDecoder_WithScaleConsistency.get_optim_params_group(model=model, configs=configs)\n to_train_num_parameters = len([n for n, p in model.named_parameters() if p.requires_grad])\n assert len(params_group) == to_train_num_parameters, f''\n optimizer = get_optimizer(params_group, configs)\n\n scheduler = build_scheduler(configs=configs, optimizer=optimizer)\n model_input_mapper = AUXMapper_v1(configs['model']['input_aux'])\n\n train_samplers, train_loaders, eval_function = build_schedule(configs, \n model_input_mapper.mapper, \n partial(model_input_mapper.collate, max_stride=model.max_stride))\n\n\n return model, optimizer, scheduler, train_samplers, train_loaders, log_lr_group_idx, eval_function\n\n\n"
},
{
"path": "models/VIS/__init__.py",
"content": "from . import BackboneEncoderDecoder_WithScaleConsistency\nfrom .. import modality_input_mappers\nfrom .. import backbone\nfrom .. import decoder\nfrom .. import encoder"
},
{
"path": "models/VIS/aux_mapper.py",
"content": "\nimport torch\n\nfrom torch.nn import functional as F\nfrom models.registry import register_model\nfrom data_schedule.utils.box_ops import box_xyxy_to_cxcywh\nfrom models.registry import MODELITY_INPUT_MAPPER_REGISTRY\nfrom data_schedule.vis.apis import VIS_TrainAPI_clipped_video\nfrom data_schedule.vis.apis import VIS_EvalAPI_clipped_video_request_ann\nfrom utils.misc import nested_tensor_from_videos_list_with_stride\n\nclass AUXMapper_v1:\n def __init__(self, aux_configs):\n video_auxes = aux_configs['video_auxes']\n\n video_auxes_names = [config['name'] for config in video_auxes]\n assert len(list(set(video_auxes_names))) == len(video_auxes_names), '每个aux的名字必须不一样'\n self.video_auxes_names = video_auxes_names\n self.video_auxes = [MODELITY_INPUT_MAPPER_REGISTRY.get(config['name'])(config) for config in video_auxes]\n\n self.targets_auxes = None\n\n def mapper(self, data_dict, mode,):\n if mode == 'train':\n VIS_TrainAPI_clipped_video\n video = data_dict['video_dict']['video']\n for aux, aux_name in zip(self.video_auxes, self.video_auxes_names):\n data_dict['video_dict'][aux_name] = aux.mapper(video)\n \n elif mode == 'evaluate':\n VIS_EvalAPI_clipped_video_request_ann\n video = data_dict['video_dict']['video']\n for aux, aux_name in zip(self.video_auxes, self.video_auxes_names):\n data_dict['video_dict'][aux_name] = aux.mapper(video) \n else:\n raise ValueError()\n \n return data_dict\n\n def collate(self, batch_dict, mode, max_stride):\n if mode == 'train':\n VIS_TrainAPI_clipped_video\n video_dict = self.collate_video_dict(batch_dict, max_stride=max_stride)\n targets = [sample['targets'] for sample in batch_dict]\n frame_has_ann = [clip_tgt['has_ann'] for clip_tgt in targets] # list[t], b\n frame_targets = [sample['frame_targets'] for sample in batch_dict]\n\n _, pad_T, _, pad_H, pad_W = video_dict['videos'].shape\n targets = self.collate_targets(targets=targets, pad_H=pad_H, pad_W=pad_W, pad_T=pad_T)\n frame_targets = self.collate_frame_targets(frame_targets=frame_targets, \n frame_has_ann=frame_has_ann,\n pad_H=pad_H, pad_W=pad_W, pad_T=pad_T)\n \n ret = {\n 'video_dict': video_dict,\n 'targets': targets,\n 'frame_targets': frame_targets,\n 'meta_idxs': [sample['meta_idx'] for sample in batch_dict],\n 'visualize': [sample['visualize'] for sample in batch_dict], \n } \n \n elif mode == 'evaluate':\n VIS_EvalAPI_clipped_video_request_ann\n assert len(batch_dict) == 1\n video_dict = self.collate_video_dict(batch_dict, max_stride=max_stride) # 不pad\n metas = [sample['meta'] for sample in batch_dict]\n\n collated_metas = {}\n for key in metas[0].keys():\n collated_metas[key] = [mt[key] for mt in metas]\n ret = {\n 'video_dict': video_dict,\n 'metas': collated_metas,\n 'meta_idxs': [sample['meta_idx'] for sample in batch_dict],\n 'visualize': [sample['visualize'] for sample in batch_dict], \n } \n debug_data = False\n if debug_data:\n self.visualize_input_target_for_debug_data(ret) # ./test.png\n return ret\n\n def collate_video_dict(self, batch_dict, max_stride):\n videos = [sample['video_dict']['video'] for sample in batch_dict] # list[ti 3 hi wi] -> b T 3 H W\n orig_sizes = [list(vid.shape) for vid in videos] # t 3 h w\n if type(max_stride) == int: # temporal max stride 为1, spatial max stride\n pad_stride = [1, max_stride]\n if (type(max_stride) == list) and (len(max_stride) == 2):\n pad_stride = max_stride\n videos = nested_tensor_from_videos_list_with_stride(videos, max_stride=pad_stride).tensors # b t c h w\n video_dicts = {'videos': videos, 'orig_sizes': orig_sizes}\n\n for aux_name, aux in zip(self.video_auxes_names, self.video_auxes):\n auxes = [sample['video_dict'][aux_name] for sample in batch_dict] # list[dict] / list[tensor]\n collated_auxes = aux.collate(auxes, batch_videos=videos) # list[dict] / tensor\n if isinstance(auxes[0], dict):\n keys = collated_auxes.keys()\n for key in keys:\n assert key not in video_dicts\n video_dicts[key] = collated_auxes[key]\n else:\n video_dicts[aux_name] = collated_auxes\n\n return video_dicts\n\n def collate_frame_targets(self, frame_targets, frame_has_ann, pad_H, pad_W, pad_T): # \n VIS_TrainAPI_clipped_video\n ret = {}\n has_ann = torch.stack([F.pad(ha.float(), pad=(0, pad_T - len(ha)), value=0.).bool() for ha in frame_has_ann], dim=0).flatten() # bT\n ret['has_ann'] = has_ann\n masks = [ftarget['masks'] for sample in frame_targets for ftarget in sample] # list[ni h w], bt'\n masks = [F.pad(m.float(), pad=(0, pad_W-m.shape[-1], 0, pad_H-m.shape[-2])).bool() for m in masks] # list[ni H W], bt'\n ret['masks'] = masks # list[ni h w], bt'\n\n classes = [ftarget['classes'] for sample in frame_targets for ftarget in sample] # list[ni], bt'\n ret['classes'] = classes\n \n if 'boxes' in frame_targets[0][0]:\n boxes = [ftarget['boxes'] for sample in frame_targets for ftarget in sample] # list[ni 4], x1y1x2y2, bt'\n boxes = [box_xyxy_to_cxcywh(bx) for bx in boxes]\n boxes = [bx / torch.tensor([pad_W, pad_H, pad_W, pad_H], dtype=bx.dtype) for bx in boxes] # 0-1\n ret['boxes'] = boxes # list[ni 4], bt' \n return ret\n\n def collate_targets(self, targets, pad_H, pad_W, pad_T):\n VIS_TrainAPI_clipped_video\n has_ann = [sample['has_ann'] for sample in targets] # list[t], bool\n has_ann = torch.stack([F.pad(ha.float(), pad=(0, pad_T - len(ha)), value=0.).bool() for ha in has_ann], dim=0) # b T\n masks = [sample['masks'] for sample in targets] \n masks = [F.pad(m.float(), pad=(0, pad_W-m.shape[-1], 0, pad_H-m.shape[-2]), value=0.).bool() \\\n for m in masks] # list[ni T' H W]\n classes = [sample['classes'] for sample in targets]\n ret = {\n 'masks': masks, # list[ni T' h w]\n 'has_ann': has_ann, # b T\n 'classes': classes, # list[ni], b\n } \n if 'boxes' in targets[0]:\n boxes = [sample['boxes'] for sample in targets] # list[ni T' 4], x1y1x2y2\n boxes = [box_xyxy_to_cxcywh(bx) for bx in boxes]\n boxes = [bx / torch.tensor([pad_W, pad_H, pad_W, pad_H], dtype=torch.float) for bx in boxes] # 0-1\n ret.update({'boxes': boxes,})\n return ret\n\n def visualize_input_target_for_debug_data(self, ret):\n videos = ret['video_dict']['videos'] # b T 3 H W\n pass\n"
},
{
"path": "models/__init__.py",
"content": "import os\nfrom .registry import model_entrypoint\n\nif os.getenv('CURRENT_TASK') == 'VIS':\n from . import VIS\nelse:\n raise ValueError()\n\n\n\n\n\n\n\n\n\n"
},
{
"path": "models/backbone/__init__.py",
"content": "\nfrom . import res2net, pvtv2"
},
{
"path": "models/backbone/pvtv2.py",
"content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom timm.models.layers import DropPath\nfrom timm.models.registry import register_model\n\nclass DWConv(nn.Module):\n def __init__(self, dim):\n super(DWConv, self).__init__()\n self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, groups=dim)\n\n def forward(self, x, H, W):\n B,N,C = x.shape\n x = x.transpose(1, 2).view(B, C, H, W)\n x = self.dwconv(x)\n x = x.flatten(2).transpose(1, 2)\n return x\n\nclass Mlp(nn.Module):\n def __init__(self, in_features, hidden_features):\n super().__init__()\n self.fc1 = nn.Linear(in_features, hidden_features)\n self.dwconv = DWConv(hidden_features)\n self.fc2 = nn.Linear(hidden_features, in_features)\n\n def forward(self, x, H, W):\n x = self.fc1(x)\n x = F.gelu(self.dwconv(x, H, W))\n x = self.fc2(x)\n return x\n\nclass Attention(nn.Module):\n def __init__(self, dim, num_heads, sr_ratio):\n super().__init__()\n assert dim % num_heads == 0, f\"dim {dim} should be divided by num_heads {num_heads}.\"\n\n self.num_heads = num_heads\n self.scale = (dim//num_heads)**(-0.5)\n self.q = nn.Linear(dim, dim)\n self.kv = nn.Linear(dim, dim*2)\n self.proj = nn.Linear(dim, dim)\n self.sr_ratio = sr_ratio\n if sr_ratio > 1:\n self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)\n self.norm = nn.LayerNorm(dim)\n\n def forward(self, x, H, W):\n B, N, C = x.shape\n q = self.q(x).reshape(B, N, self.num_heads, C//self.num_heads).permute(0, 2, 1, 3)\n\n if self.sr_ratio > 1:\n x_ = x.permute(0, 2, 1).reshape(B, C, H, W)\n x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)\n x_ = self.norm(x_)\n kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n else:\n kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n \n k, v = kv[0], kv[1]\n attn = (q @ k.transpose(-2, -1)) * self.scale\n attn = attn.softmax(dim=-1)\n x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n x = self.proj(x)\n return x\n\nclass Block(nn.Module):\n def __init__(self, dim, num_heads, mlp_ratio, drop_path, sr_ratio):\n super().__init__()\n self.norm1 = nn.LayerNorm(dim, eps=1e-6)\n self.attn = Attention(dim, num_heads=num_heads, sr_ratio=sr_ratio)\n self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n self.norm2 = nn.LayerNorm(dim, eps=1e-6)\n self.mlp = Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio))\n\n def forward(self, x, H, W):\n x = x + self.drop_path(self.attn(self.norm1(x), H, W))\n x = x + self.drop_path(self.mlp(self.norm2(x), H, W))\n return x\n\nclass OverlapPatchEmbed(nn.Module):\n def __init__(self, patch_size, stride, in_chans, embed_dim):\n super().__init__()\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride, padding=(patch_size//2, patch_size//2))\n self.norm = nn.LayerNorm(embed_dim)\n\n def forward(self, x):\n x = self.proj(x)\n B,C,H,W = x.shape\n x = x.flatten(2).transpose(1, 2)\n x = self.norm(x)\n return x, H, W\n\nclass PVT(nn.Module):\n def __init__(self, embed_dims, mlp_ratios, depths, snapshot, sr_ratios=[8, 4, 2, 1]):\n super().__init__()\n self.depths = depths\n self.snapshot = snapshot\n # patch_embed\n self.patch_embed1 = OverlapPatchEmbed(patch_size=7, stride=4, in_chans=3, embed_dim=embed_dims[0])\n self.patch_embed2 = OverlapPatchEmbed(patch_size=3, stride=2, in_chans=embed_dims[0], embed_dim=embed_dims[1])\n self.patch_embed3 = OverlapPatchEmbed(patch_size=3, stride=2, in_chans=embed_dims[1], embed_dim=embed_dims[2])\n self.patch_embed4 = OverlapPatchEmbed(patch_size=3, stride=2, in_chans=embed_dims[2], embed_dim=embed_dims[3])\n # transformer encoder\n dpr = [x.item() for x in torch.linspace(0, 0.1, sum(depths))] # stochastic depth decay rule\n cur = 0\n self.block1 = nn.ModuleList([Block(dim=embed_dims[0], num_heads=1, mlp_ratio=mlp_ratios[0], drop_path=dpr[cur + i], sr_ratio=sr_ratios[0]) for i in range(depths[0])])\n self.norm1 = nn.LayerNorm(embed_dims[0], eps=1e-6)\n\n cur += depths[0]\n self.block2 = nn.ModuleList([Block(dim=embed_dims[1], num_heads=2, mlp_ratio=mlp_ratios[1], drop_path=dpr[cur + i], sr_ratio=sr_ratios[1]) for i in range(depths[1])])\n self.norm2 = nn.LayerNorm(embed_dims[1], eps=1e-6)\n\n cur += depths[1]\n self.block3 = nn.ModuleList([Block(dim=embed_dims[2], num_heads=5, mlp_ratio=mlp_ratios[2], drop_path=dpr[cur + i], sr_ratio=sr_ratios[2]) for i in range(depths[2])])\n self.norm3 = nn.LayerNorm(embed_dims[2], eps=1e-6)\n\n cur += depths[2]\n self.block4 = nn.ModuleList([Block(dim=embed_dims[3], num_heads=8, mlp_ratio=mlp_ratios[3], drop_path=dpr[cur + i], sr_ratio=sr_ratios[3]) for i in range(depths[3])])\n self.norm4 = nn.LayerNorm(embed_dims[3], eps=1e-6)\n\n state_dict:dict = torch.load(self.snapshot, map_location='cpu')\n state_dict.pop(\"head.weight\")\n state_dict.pop(\"head.bias\")\n self.load_state_dict(state_dict, strict=True)\n del state_dict\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better\n\n def forward(self, x):\n B = x.shape[0]\n # stage 1\n out1, H, W = self.patch_embed1(x)\n for i, blk in enumerate(self.block1):\n out1 = blk(out1, H, W)\n out1 = self.norm1(out1).reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n\n # stage 2\n out2, H, W = self.patch_embed2(out1)\n for i, blk in enumerate(self.block2):\n out2 = blk(out2, H, W)\n out2 = self.norm2(out2).reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n\n # stage 3\n out3, H, W = self.patch_embed3(out2)\n for i, blk in enumerate(self.block3):\n out3 = blk(out3, H, W)\n out3 = self.norm3(out3).reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n\n # stage 4\n out4, H, W = self.patch_embed4(out3)\n for i, blk in enumerate(self.block4):\n out4 = blk(out4, H, W)\n out4 = self.norm4(out4).reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n return out1, out2, out3, out4\n\n\nfrom detectron2.modeling import BACKBONE_REGISTRY\nfrom einops import rearrange, reduce, repeat\nfrom .utils import VideoMultiscale_Shape, ImageMultiscale_Shape\nimport os\nimport time\n\n\n@BACKBONE_REGISTRY.register()\nclass PVT_V2(nn.Module):\n def __init__(self, configs) -> None:\n super().__init__()\n pt_path = os.getenv('PT_PATH')\n pvt_v2 = PVT(embed_dims=[64, 128, 320, 512], mlp_ratios=[8, 8, 4, 4], \n depths=[3, 4, 6, 3], snapshot=os.path.join(pt_path, 'pvt_v2/pvt_v2_b2.pth'))\n self.pvt_v2 = pvt_v2\n\n freeze = configs['freeze']\n\n if freeze:\n for p in self.parameters():\n p.requires_grad_(False)\n\n self.multiscale_shapes = {}\n for name, spatial_stride, dim in zip(['res2', 'res3', 'res4', 'res5'], \n [4, 8, 16, 32],\n [64, 128, 320, 512]):\n self.multiscale_shapes[name] = ImageMultiscale_Shape(spatial_stride=spatial_stride, dim=dim)\n self.max_stride = 32\n\n def forward(self, x):\n\n if not self.training:\n batch_feats = []\n for haosen in x:\n feats = self.pvt_v2(haosen.unsqueeze(0))\n batch_feats.append(feats)\n batch_feats = list(zip(*batch_feats)) # 4\n batch_feats = [torch.cat(haosen, dim=0) for haosen in batch_feats] # list[bt c h w]\n ret = {}\n names = ['res2', 'res3', 'res4', 'res5']\n for name, feat in zip(names, batch_feats):\n ret[name] = feat\n return ret \n else:\n layer_outputs = self.pvt_v2(x)\n ret = {}\n names = ['res2', 'res3', 'res4', 'res5']\n for name, feat in zip(names, layer_outputs):\n ret[name] = feat\n return ret\n\n\n def num_parameters(self):\n return sum(p.numel() for p in self.parameters() if p.requires_grad)\n\n@BACKBONE_REGISTRY.register()\nclass Video2D_PVT_V2(nn.Module):\n def __init__(self, configs) -> None:\n super().__init__()\n self.image_homo = PVT_V2(configs=configs)\n\n self.multiscale_shapes = {}\n for name, temporal_stride, spatial_stride, dim in zip(['res2', 'res3', 'res4', 'res5'], \n [1, 1, 1, 1], \n [4, 8, 16, 32],\n [64, 128, 320, 512]):\n self.multiscale_shapes[name] = VideoMultiscale_Shape(temporal_stride=temporal_stride, \n spatial_stride=spatial_stride, dim=dim)\n self.max_stride = [1, 32]\n\n \n def forward(self, x):\n batch_size, _, T = x.shape[:3]\n x = rearrange(x, 'b c t h w -> (b t) c h w').contiguous()\n layer_outputs = self.image_homo(x)\n\n layer_outputs = {key: rearrange(value.contiguous(), '(b t) c h w -> b c t h w',b=batch_size, t=T).contiguous() \\\n for key, value in layer_outputs.items()}\n return layer_outputs\n \n def num_parameters(self):\n return sum(p.numel() for p in self.parameters() if p.requires_grad)"
},
{
"path": "models/backbone/res2net.py",
"content": "import math\nimport torch\nimport torch.nn as nn\nimport os\nimport torch.nn.functional as F\nfrom detectron2.modeling import BACKBONE_REGISTRY\nfrom einops import rearrange, reduce, repeat\nfrom .utils import VideoMultiscale_Shape\n\nclass Bottle2neck(nn.Module):\n expansion = 4\n def __init__(self, inplanes, planes, stride=1, downsample=None, baseWidth=26, scale=4, stype='normal'):\n super(Bottle2neck, self).__init__()\n width = int(math.floor(planes*(baseWidth/64.0)))\n self.conv1 = nn.Conv2d(inplanes, width*scale, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(width*scale)\n self.nums = 1 if scale == 1 else scale - 1\n if stype == 'stage':\n self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)\n convs, bns = [], []\n for i in range(self.nums):\n convs.append(nn.Conv2d(width, width, kernel_size=3, stride=stride, padding=1, bias=False))\n bns.append(nn.BatchNorm2d(width))\n self.convs = nn.ModuleList(convs)\n self.bns = nn.ModuleList(bns)\n self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)\n self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n self.downsample = downsample\n self.stype = stype\n self.scale = scale\n self.width = width\n\n def forward(self, x):\n out = F.relu(self.bn1(self.conv1(x)), inplace=True)\n spx = torch.split(out, self.width, 1)\n for i in range(self.nums):\n sp = spx[i] if i == 0 or self.stype == 'stage' else sp + spx[i]\n sp = self.convs[i](sp)\n sp = F.relu(self.bns[i](sp), inplace=True)\n out = sp if i == 0 else torch.cat((out, sp), 1)\n if self.scale != 1 and self.stype == 'normal':\n out = torch.cat((out, spx[self.nums]), 1)\n elif self.scale != 1 and self.stype == 'stage':\n out = torch.cat((out, self.pool(spx[self.nums])), 1)\n\n out = self.bn3(self.conv3(out))\n if self.downsample is not None:\n x = self.downsample(x)\n return F.relu(out+x, inplace=True)\n\nclass Res2Net(nn.Module):\n def __init__(self, layers, snapshot, baseWidth=26, scale=4):\n super(Res2Net, self).__init__()\n self.inplanes = 64\n self.snapshot = snapshot\n self.baseWidth = baseWidth\n self.scale = scale\n self.conv1 = nn.Sequential(\n nn.Conv2d(3, 32, 3, 2, 1, bias=False),\n nn.BatchNorm2d(32),\n nn.ReLU(inplace=True),\n nn.Conv2d(32, 32, 3, 1, 1, bias=False),\n nn.BatchNorm2d(32),\n nn.ReLU(inplace=True),\n nn.Conv2d(32, 64, 3, 1, 1, bias=False)\n )\n self.bn1 = nn.BatchNorm2d(64)\n self.layer1 = self._make_layer(Bottle2neck, 64, layers[0])\n self.layer2 = self._make_layer(Bottle2neck, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(Bottle2neck, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(Bottle2neck, 512, layers[3], stride=2)\n\n state_dict:dict = torch.load(self.snapshot, map_location='cpu')\n self.load_state_dict(state_dict, strict=False)\n del state_dict\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True, count_include_pad=False),\n nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=1, bias=False),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = [block(self.inplanes, planes, stride, downsample=downsample, stype='stage', baseWidth=self.baseWidth, scale=self.scale)]\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes, baseWidth=self.baseWidth, scale=self.scale))\n return nn.Sequential(*layers)\n\n def forward(self, x):\n out1 = F.relu(self.bn1(self.conv1(x)), inplace=True)\n out1 = F.max_pool2d(out1, kernel_size=3, stride=2, padding=1)\n out2 = self.layer1(out1)\n out3 = self.layer2(out2)\n out4 = self.layer3(out3)\n out5 = self.layer4(out4)\n return out2, out3, out4, out5\n\n def initialize(self):\n self.load_state_dict(torch.load(self.snapshot), strict=False)\n\n@BACKBONE_REGISTRY.register()\nclass Res2Net_50_EachFrame(nn.Module):\n def __init__(self, configs) -> None:\n super().__init__()\n pt_path = os.getenv('PT_PATH')\n res2net = Res2Net([3, 4, 6, 3], os.path.join(pt_path, 'res2net/res2net50_v1b_26w_4s-3cf99910.pth'))\n self.res2net = res2net\n\n freeze = configs['freeze']\n\n if freeze:\n for p in self.parameters():\n p.requires_grad_(False)\n\n self.multiscale_shapes = {}\n for name, temporal_stride, spatial_stride, dim in zip(['res2', 'res3', 'res4', 'res5'], \n [1, 1, 1, 1], \n [4, 8, 16, 32],\n [256, 512, 1024, 2048]):\n self.multiscale_shapes[name] = VideoMultiscale_Shape(temporal_stride=temporal_stride, \n spatial_stride=spatial_stride, dim=dim)\n\n self.max_stride = [1, 32]\n\n \n def forward(self, x):\n batch_size, _, T = x.shape[:3]\n x = rearrange(x, 'b c t h w -> (b t) c h w')\n layer_outputs = self.res2net(x)\n\n ret = {}\n names = ['res2', 'res3', 'res4', 'res5']\n for name, feat in zip(names, layer_outputs):\n ret[name] = rearrange(feat.contiguous(), '(b t) c h w -> b c t h w',b=batch_size, t=T)\n\n return ret\n \n\n def num_parameters(self):\n return sum(p.numel() for p in self.parameters() if p.requires_grad) \n "
},
{
"path": "models/backbone/utils.py",
"content": "\nclass VideoMultiscale_Shape:\n def __init__(self, temporal_stride, spatial_stride, dim) -> None:\n self.temporal_stride = temporal_stride\n self.spatial_stride = spatial_stride\n self.dim = dim\n \n @staticmethod\n def set_multiscale_same_dim(shape_by_dim, same_dim):\n return {\n key: VideoMultiscale_Shape(temporal_stride=value.temporal_stride,\n spatial_stride=value.spatial_stride,\n dim=same_dim) for key,value in shape_by_dim.items()\n }\n\nclass ImageMultiscale_Shape:\n def __init__(self, spatial_stride, dim) -> None:\n self.spatial_stride = spatial_stride\n self.dim = dim\n \n\n\n"
},
{
"path": "models/decoder/__init__.py",
"content": "\nfrom . import mask2former_video"
},
{
"path": "models/decoder/mask2former_video.py",
"content": "\n# multi-scale features, b c h w -> module -> obj queries, predictions, b nq c\nimport torch.nn as nn\nfrom models.layers.decoder_layers import CrossAttentionLayer, SelfAttentionLayer, FFNLayer\nfrom models.layers.anyc_trans import MLP\nimport torch.nn.functional as F\nimport torch\nimport copy\nfrom models.layers.utils import zero_module, _get_clones\nfrom models.layers.position_encoding import build_position_encoding\nfrom einops import rearrange, reduce, repeat\nfrom scipy.optimize import linear_sum_assignment\nfrom models.layers.matching import batch_dice_loss, batch_sigmoid_ce_loss, batch_sigmoid_focal_loss, dice_loss, ce_mask_loss\nfrom detectron2.modeling import META_ARCH_REGISTRY\nimport detectron2.utils.comm as comm\nimport data_schedule.utils.box_ops as box_ops\nfrom models.layers.utils import zero_module\nfrom utils.misc import is_dist_avail_and_initialized\nfrom collections import defaultdict\nfrom detectron2.projects.point_rend.point_features import point_sample\nfrom torch.cuda.amp import autocast\nfrom detectron2.projects.point_rend.point_features import (\n get_uncertain_point_coords_with_randomness,\n point_sample,\n)\ndef calculate_uncertainty(logits):\n \"\"\"\n We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the\n foreground class in `classes`.\n Args:\n logits (Tensor): A tensor of shape (R, 1, ...) for class-specific or\n class-agnostic, where R is the total number of predicted masks in all images and C is\n the number of foreground classes. The values are logits.\n Returns:\n scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with\n the most uncertain locations having the highest uncertainty score.\n \"\"\"\n assert logits.shape[1] == 1\n gt_class_logits = logits.clone()\n return -(torch.abs(gt_class_logits))\n\n\nclass Video_SetMatchingLoss(nn.Module):\n def __init__(self, \n loss_config,\n num_classes,) -> None:\n super().__init__()\n self.num_classes = num_classes # n=1 / n=0 / n>1\n self.matching_metrics = loss_config['matching_metrics'] # mask: mask/dice; point_sample_mask: ..\n self.losses = loss_config['losses'] \n\n self.aux_layer_weights = loss_config['aux_layer_weights'] # int/list\n empty_weight = torch.ones(self.num_classes + 1)\n empty_weight[-1] = loss_config['background_cls_eos']\n self.register_buffer('empty_weight', empty_weight)\n # self.register_buffer('small_obj_weight', torch.tensor(loss_config['small_obj_weight']).float())\n self._warmup_iters = 2000\n self.register_buffer(\"_iter\", torch.zeros([1]))\n\n @property\n def device(self,):\n return self.empty_weight.device\n \n def compute_loss(self, \n model_outs, \n targets,\n video_aux_dict,\n **kwargs):\n # list[n t' h w], batch\n if 'masks' in targets:\n num_objs = sum([haosen.flatten(1).any(-1).int().sum().item() for haosen in targets['masks']])\n # list[n t' 4], batch\n elif 'boxes' in targets:\n # n t' 2 -> n t -> n\n num_objs = sum([(haosen[:, :, 2:] > 0).all(-1).any(-1).int().sum().item() for haosen in targets['boxes']])\n else:\n raise ValueError('targets里没有boxes/masks, 需要确定数量')\n \n num_objs = torch.as_tensor([num_objs], dtype=torch.float, device=self.device)\n if is_dist_avail_and_initialized():\n torch.distributed.all_reduce(num_objs)\n num_objs = torch.clamp(num_objs / comm.get_world_size(), min=1).item()\n if isinstance(self.aux_layer_weights, list):\n assert len(self.aux_layer_weights) == (len(model_outs) - 1)\n else:\n self.aux_layer_weights = [self.aux_layer_weights] * (len(model_outs) - 1)\n\n layer_weights = self.aux_layer_weights + [1.]\n\n loss_values = {\n 'mask_dice':0., 'mask_ce':0.,\n 'box_l1': 0., 'box_giou': 0.,\n 'class_ce':0.,\n\n 'mask_dice_smobj':0., 'mask_ce_smobj':0.,\n 'boxMask_dice':0., 'boxMask_ce':0.,\n }\n \n if ('mask_ce_dice' in self.matching_metrics) or ('mask_ce_dice' in self.losses):\n # mask interpolate\n tgt_mask_shape = targets['masks'][0].shape[-2:] # list[n t H W], b\n for layer_idx in range(len(model_outs)):\n # b t nq h w\n batch_size, nf = model_outs[layer_idx]['pred_masks'].shape[:2]\n model_outs[layer_idx]['pred_masks'] = rearrange(F.interpolate(model_outs[layer_idx]['pred_masks'].flatten(0, 1),\n size=tgt_mask_shape, mode='bilinear', align_corners=False),\n '(b t) n h w -> b t n h w',b=batch_size, t=nf)\n for taylor, layer_out in zip(layer_weights, model_outs):\n if taylor != 0:\n matching_indices = self.matching(layer_out, targets)\n for loss in self.losses:\n loss_extra_param = self.losses[loss]\n if loss == 'mask_dice_ce' :\n loss_dict = self.loss_mask_dice_ce(layer_out, targets, matching_indices, num_objs,\n loss_extra_param=loss_extra_param)\n elif loss == 'class_ce':\n loss_dict = self.loss_class_ce(layer_out, targets, matching_indices, num_objs,\n loss_extra_param=loss_extra_param)\n elif loss == 'point_mask_dice_ce':\n loss_dict = self.loss_point_mask_dice_ce(layer_out, targets, matching_indices, num_objs,\n loss_extra_param=loss_extra_param)\n else:\n raise ValueError()\n \n for key, value in loss_dict.items():\n loss_values[key] = loss_values[key] + value\n \n return loss_values \n\n @torch.no_grad()\n def matching(self, layer_out, targets):\n batch_size = len(targets['masks']) if 'masks' in targets else len(targets['boxes'])\n indices = [] \n has_ann = targets['has_ann']\n for i in range(batch_size):\n C = 0.\n\n if 'class_prob' in self.matching_metrics:\n out_cls = layer_out['pred_class'][i].softmax(-1) # nq c\n tgt_cls = targets['classes'][i] # n\n cost_class = - out_cls[:, tgt_cls] # nq n\n C += self.matching_metrics['class_prob']['prob'] * cost_class\n\n if 'mask_dice_ce' in self.matching_metrics:\n out_mask = layer_out['pred_masks'][i][has_ann[i]].permute(1, 0, 2, 3).contiguous() # nq t' h w\n tgt_mask = targets['masks'][i].to(out_mask) # ni t' H W\n cost_mask = batch_sigmoid_ce_loss(out_mask.flatten(1), tgt_mask.flatten(1)) \n cost_dice = batch_dice_loss(out_mask.flatten(1), tgt_mask.flatten(1))\n C += self.matching_metrics['mask_dice_ce']['ce'] * cost_mask + \\\n self.matching_metrics['mask_dice_ce']['dice'] * cost_dice\n\n \n if 'point_mask_dice_ce' in self.matching_metrics:\n out_mask = layer_out['pred_masks'][i][has_ann[i]].permute(1, 0, 2, 3).contiguous() # nq t' h w\n tgt_mask = targets['masks'][i].to(out_mask)# ni t' H W\n nf = out_mask.shape[1]\n\n out_mask = out_mask.flatten(0, 1)[:, None]\n tgt_mask = tgt_mask.flatten(0, 1)[:, None]\n # all masks share the same set of points for efficient matching!\n point_coords = torch.rand(1, self.matching_metrics['point_mask_dice_ce']['num_points'],\n 2, device=self.device)\n # get gt labels\n tgt_mask = point_sample(\n tgt_mask,\n point_coords.repeat(tgt_mask.shape[0], 1, 1),\n align_corners=False,\n ).squeeze(1) # nqt s\n tgt_mask = rearrange(tgt_mask, '(nq t) s -> nq t s',t=nf)\n\n out_mask = point_sample(\n out_mask,\n point_coords.repeat(out_mask.shape[0], 1, 1),\n align_corners=False,\n ).squeeze(1) # nit s\n out_mask = rearrange(out_mask, '(nq t) s -> nq t s',t=nf)\n with autocast(enabled=False):\n out_mask = out_mask.float().flatten(1) # nq num_points\n tgt_mask = tgt_mask.float().flatten(1)\n cost_mask = batch_sigmoid_ce_loss(out_mask, tgt_mask)\n cost_dice = batch_dice_loss(out_mask, tgt_mask)\n C += self.matching_metrics['point_mask_dice_ce']['ce'] * cost_mask + \\\n self.matching_metrics['point_mask_dice_ce']['dice'] * cost_dice\n \n C = C.cpu()\n indices.append(linear_sum_assignment(C))\n\n return [\n (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))\n for i, j in indices\n ]\n\n def loss_mask_dice_ce(self, outputs, targets, indices, num_objs, loss_extra_param):\n has_ann = targets['has_ann'] # b t\n src_masks = outputs['pred_masks'].permute(0, 2, 1, 3, 4).contiguous() # b nq t h w\n tgt_masks = targets['masks'] # list[n t' h w]\n # list[nq t' h w] -> n_sigma t' h w\n src_masks = torch.cat([t[J][:, haosen] for t, (J, _), haosen in zip(src_masks, indices, has_ann)],dim=0) \n tgt_masks = torch.cat([t[J] for t, (_, J) in zip(tgt_masks, indices)], dim=0)\n tgt_masks = tgt_masks.to(src_masks)\n losses = {\n \"mask_ce\": ce_mask_loss(src_masks.flatten(0, 1).flatten(1), tgt_masks.flatten(0, 1).flatten(1), num_boxes=num_objs),\n \"mask_dice\": dice_loss(src_masks.flatten(0, 1).flatten(1), tgt_masks.flatten(0, 1).flatten(1), num_boxes=num_objs),\n }\n return losses \n\n def loss_point_mask_dice_ce(self, outputs, targets, indices, num_objs, loss_extra_param):\n has_ann = targets['has_ann'] # b t\n src_masks = outputs['pred_masks'].permute(0, 2, 1, 3, 4).contiguous() # b nq t h w\n tgt_masks = targets['masks'] # list[n t' h w]\n # list[nq t' h w] -> n_sigma t' h w\n src_masks = torch.cat([t[J][:, haosen] for t, (J, _), haosen in zip(src_masks, indices, has_ann)],dim=0) \n tgt_masks = torch.cat([t[J] for t, (_, J) in zip(tgt_masks, indices)], dim=0)\n tgt_masks = tgt_masks.to(src_masks)\n nf = src_masks.shape[1]\n\n # No need to upsample predictions as we are using normalized coordinates :)\n # NT x 1 x H x W\n src_masks = src_masks.flatten(0, 1).unsqueeze(1).contiguous() # nt' 1 h w\n target_masks = tgt_masks.flatten(0, 1).unsqueeze(1).contiguous() \n\n with torch.no_grad():\n # sample point_coords\n point_coords = get_uncertain_point_coords_with_randomness(\n src_masks,\n lambda logits: calculate_uncertainty(logits),\n loss_extra_param['num_points'],\n loss_extra_param['oversample_ratio'],\n loss_extra_param['importance_sample_ratio'],\n )\n # get gt labels\n point_labels = point_sample(\n target_masks,\n point_coords,\n align_corners=False,\n ).squeeze(1) # nt' s\n\n point_logits = point_sample(\n src_masks,\n point_coords,\n align_corners=False,\n ).squeeze(1) # nt' s\n\n # point_logits = rearrange(point_logits, '(n t) s -> n (t s)',t=nf)\n # point_labels = rearrange(point_labels, '(n t) s -> n (t s)',t=nf)\n\n losses = {\n \"mask_dice\": ce_mask_loss(point_logits, point_labels, num_objs),\n \"mask_ce\": dice_loss(point_logits, point_labels, num_objs),\n }\n\n del src_masks\n del target_masks\n return losses \n\n def loss_class_ce(self, outputs, targets, indices, num_objs, loss_extra_param):\n \"\"\"Classification loss (NLL)\n targets dicts must contain the key \"labels\" containing a tensor of dim [nb_target_boxes]\n \"\"\"\n src_logits = outputs[\"pred_class\"].float() # b nq c\n\n idx = self._get_src_permutation_idx(indices)\n # list[n], b -> bn\n target_classes_o = torch.cat([t[J] for t, (_, J) in zip(targets['classes'], indices)]) \n target_classes = torch.full(\n src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=self.device\n )\n target_classes[idx] = target_classes_o\n\n loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)\n losses = {\"class_ce\": loss_ce}\n return losses\n\n def loss_box_l1_giou(self, outputs, targets, indices, num_objs, loss_extra_param): \n tgt_boxes = targets['boxes'] # list[n tl 4], b\n has_ann = targets['has_ann'] # b t\n\n src_boxes = outputs['pred_boxes'].sigmoid().permute(0, 2, 1, 3).contiguous() # b nq t 4\n\n src_boxes = torch.cat([t[J][:, haosen] for t, (J, _), haosen in zip(src_boxes, indices, has_ann)], dim=0) # n_sum t' 4\n tgt_boxes = torch.cat([t[J] for t, (_, J) in zip(tgt_boxes, indices)], dim=0) # n_sum t' 4\n \n nf = tgt_boxes.shape[1]\n loss_l1 = F.l1_loss(src_boxes, tgt_boxes, reduction='none').flatten(1) # n_sum t'4\n\n loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(\n box_ops.box_cxcywh_to_xyxy(src_boxes.flatten(0, 1)),\n box_ops.box_cxcywh_to_xyxy(tgt_boxes.flatten(0, 1)))) # n_sumt'\n \n loss_giou = loss_giou.view(-1, nf).contiguous()\n return {\n 'box_l1': loss_l1.sum(-1).sum() / num_objs,\n 'box_giou': loss_giou.sum(-1).sum() / num_objs\n }\n\n def _get_src_permutation_idx(self, indices):\n # permute predictions following indices\n batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])\n src_idx = torch.cat([src for (src, _) in indices])\n return batch_idx, src_idx\n\n def _get_tgt_permutation_idx(self, indices):\n # permute targets following indices\n batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])\n tgt_idx = torch.cat([tgt for (_, tgt) in indices])\n return batch_idx, tgt_idx\n\n@META_ARCH_REGISTRY.register()\nclass Video_MaskedAttn_MultiscaleMaskDecoder_v3(nn.Module):\n def __init__(self,\n configs,\n multiscale_shapes):\n super().__init__()\n d_model = configs['d_model']\n attn_configs = configs['attn']\n self.video_nqueries = configs['video_nqueries']\n self.nlayers = configs['nlayers']\n self.memory_scales = configs['memory_scales']\n self.mask_scale = configs['mask_scale']\n self.mask_spatial_stride = multiscale_shapes[self.mask_scale].spatial_stride\n num_classes = configs['num_classes']\n\n inputs_projs = configs['inputs_projs']\n self.inputs_projs = nn.Sequential()\n if inputs_projs is not None:\n self.inputs_projs = META_ARCH_REGISTRY.get(inputs_projs['name'])(inputs_projs, \n multiscale_shapes=multiscale_shapes,\n out_dim=d_model)\n self.level_embeds = nn.Embedding(len(self.memory_scales), d_model)\n assert self.nlayers % len(self.memory_scales) == 0\n self.cross_layers = _get_clones(CrossAttentionLayer(d_model=d_model,\n nhead=attn_configs['nheads'],\n dropout=0.0,\n normalize_before=attn_configs['normalize_before']),\n self.nlayers)\n self.self_layers = _get_clones(SelfAttentionLayer(d_model=d_model,\n nhead=attn_configs['nheads'],\n dropout=0.0,\n normalize_before=attn_configs['normalize_before']),\n self.nlayers) \n self.ffn_layers = _get_clones(FFNLayer(d_model=d_model,\n dim_feedforward=attn_configs['dim_feedforward'],\n dropout=0.0,\n normalize_before=attn_configs['normalize_before']),\n self.nlayers) \n \n self.nheads = attn_configs['nheads']\n self.temporal_query_poses = nn.Embedding(self.video_nqueries, d_model)\n self.temporal_query_feats = nn.Embedding(self.video_nqueries, d_model)\n self.temporal_query_norm = nn.LayerNorm(d_model)\n self.pos_3d = build_position_encoding(hidden_dim=d_model, position_embedding_name='3d') # b t c h w\n \n self.head_outputs = configs['head_outputs']\n assert 'mask' in self.head_outputs\n self.query_mask = MLP(d_model, d_model, d_model, 3)\n if 'class' in self.head_outputs:\n self.query_class = nn.Linear(d_model, num_classes + 1) \n\n self.loss_module = Video_SetMatchingLoss(loss_config=configs['loss'], num_classes=num_classes)\n \n @property\n def device(self,):\n return self.temporal_query_feats.weight.device\n\n def get_memories_and_mask_features(self, multiscales):\n # b c t h w \n memories = [multiscales[scale] for scale in self.memory_scales]\n size_list = [mem_feat.shape[-2:] for mem_feat in memories]\n memories_poses = [self.pos_3d(mem.permute(0, 2, 1,3, 4)).permute(0, 2, 1, 3, 4) for mem in memories] # b c t h w\n memories = [rearrange(mem, 'b c t h w -> (t h w) b c').contiguous() for mem in memories]\n memories_poses = [rearrange(mem_pos, 'b c t h w -> (t h w) b c').contiguous() for mem_pos in memories_poses]\n mask_features = multiscales[self.mask_scale] # b c t h w\n return memories, memories_poses, mask_features, size_list\n\n def forward(self, \n multiscales, # b c t h w\n video_aux_dict=None\n ):\n multiscales = self.inputs_projs(multiscales[0])\n # thw b c; b c t h w\n memories, memories_poses, mask_features, size_list = self.get_memories_and_mask_features(multiscales)\n memories = [mem_feat + self.level_embeds.weight[i][None, None, :] for i, mem_feat in enumerate(memories)]\n batch_size, _, nf, *_ = mask_features.shape\n \n # nq b c\n temporal_query_poses = self.temporal_query_poses.weight.unsqueeze(1).repeat(1, batch_size, 1)\n temporal_query_feats = self.temporal_query_feats.weight.unsqueeze(1).repeat(1, batch_size, 1)\n\n vid_ret = []\n # b nq class, b nq t h w; b*head nq thw\n vid_class, vid_mask, attn_mask = \\\n self.forward_heads(temporal_query_feats=temporal_query_feats,\n mask_features=mask_features, attn_mask_target_size=size_list[0]) # first sight you re not human\n vid_ret.append({'pred_class':vid_class, 'pred_masks': vid_mask})\n\n for i in range(self.nlayers):\n level_index = i % len(self.memory_scales)\n attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False # 全masked掉的 全注意, 比如有padding\n\n temporal_query_feats = self.cross_layers[i](\n tgt=temporal_query_feats, # nq b c\n memory=memories[level_index], # thw b c\n memory_mask=attn_mask, # b*h nq thw\n memory_key_padding_mask=None,\n pos=memories_poses[level_index], # thw b c\n query_pos=temporal_query_poses, # nq b c\n )\n temporal_query_feats = self.self_layers[i](\n temporal_query_feats,\n tgt_mask=None,\n tgt_key_padding_mask=None,\n query_pos=temporal_query_poses,\n )\n temporal_query_feats = self.ffn_layers[i](\n temporal_query_feats \n )\n # b nq class, b nq t h w\n vid_class, vid_mask, attn_mask = \\\n self.forward_heads(temporal_query_feats=temporal_query_feats,\n mask_features=mask_features, attn_mask_target_size=size_list[(i + 1) % len(self.memory_scales)]) # first sight you re not human\n vid_ret.append({'pred_class':vid_class, 'pred_masks': vid_mask})\n \n return vid_ret\n\n def forward_heads(self, temporal_query_feats, mask_features, attn_mask_target_size): # nq b c; b c t h w\n batch_size, _, nf, *_ = mask_features.shape\n\n temporal_query_feats = self.temporal_query_norm(temporal_query_feats) # nq b c\n temporal_query_feats = temporal_query_feats.transpose(0, 1).contiguous() # b nq c\n\n class_logits = self.query_class(temporal_query_feats) if 'class' in self.head_outputs else None # b n class+1\n mask_embeds = self.query_mask(temporal_query_feats) # b n c\n mask_logits = torch.einsum(\"bqc,bcthw->bqthw\", mask_embeds, mask_features) \n batch_size, nq, nf = mask_logits.shape[:3]\n mask_logits = F.interpolate(mask_logits.flatten(0, 1), scale_factor=self.mask_spatial_stride, mode='bilinear', align_corners=False)\n mask_logits = rearrange(mask_logits, '(b n) t h w -> b t n h w',b=batch_size, n=nq)\n\n # bt nq h w\n attn_mask = mask_logits.detach().clone().flatten(0, 1)\n attn_mask = (F.interpolate(attn_mask, size=attn_mask_target_size, mode=\"bilinear\", align_corners=False) < 0.5).bool()\n attn_mask = repeat(attn_mask, '(b t) nq h w -> (b head) nq (t h w)', b=batch_size, t=nf, head=self.nheads)\n \n if self.training:\n return class_logits, mask_logits, attn_mask\n else:\n return class_logits.softmax(-1).unsqueeze(1).repeat(1, nf, 1, 1) if class_logits is not None else None, mask_logits, attn_mask\n \n def compute_loss(self, outputs, targets, video_aux_dict, **kwargs):\n assert self.training\n return self.loss_module.compute_loss(model_outs=outputs,\n targets=targets,\n video_aux_dict=video_aux_dict)\n"
},
{
"path": "models/encoder/__init__.py",
"content": "\nfrom . import localGlobal\nfrom . import input_projs\nfrom . import neighborhood_qk\n"
},
{
"path": "models/encoder/input_projs.py",
"content": "\nimport torch.nn as nn\nfrom detectron2.modeling import META_ARCH_REGISTRY\nfrom models.layers.anyc_trans import Linear_NormAct\n\nfrom models.layers.anyc_trans import Conv3d_NormAct, Conv2d_NormAct\nfrom einops import rearrange\n\n@META_ARCH_REGISTRY.register()\nclass VideoConv3d_TextLinear(nn.Module):\n \"\"\"\n 如果multiscale_shapes是None, 那么每个multiscale_shape的input_dim都是out_dim\n 如果multiscale_shapes给出了, 那么按照multiscale shapes里的dim\n \"\"\"\n def __init__(self,\n configs,\n out_dim,\n text_dim=None, # 如果是none的话, 那么假设等于out_dim\n multiscale_shapes=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n text_dim = out_dim if text_dim is None else text_dim\n\n multiscale_projs_config = configs['video_multiscale_projs']\n proj_names = multiscale_projs_config.keys() # list[str]\n\n in_dims = {}\n if multiscale_shapes is not None:\n assert set(proj_names).issubset(set(list(multiscale_shapes.keys())))\n for name in proj_names:\n in_dims[name] = multiscale_shapes[name].dim\n else:\n for name in proj_names:\n in_dims[name] = out_dim\n\n projs = {}\n for name, config in multiscale_projs_config.items():\n projs[name] = Conv3d_NormAct(in_channels=in_dims[name],\n out_channels=out_dim, \n **config)\n \n self.video_multiscale_projs = nn.ModuleDict(projs)\n\n text_proj_config = configs['text_proj']\n if text_proj_config is None:\n self.text_proj = nn.Identity()\n else:\n self.text_proj = Linear_NormAct(in_features=text_dim, out_features=out_dim, **text_proj_config)\n\n def forward(self, multiscales, text_dict):\n ret = {}\n for scale_name, scale_feat in multiscales.items(): # b c t h w\n if scale_name in self.video_multiscale_projs:\n scale_feat = self.video_multiscale_projs[scale_name](scale_feat)\n ret[scale_name] = scale_feat\n else:\n ret[scale_name] = scale_feat\n \n if isinstance(text_dict, AMRData):\n text_dict.amr_feats = self.text_proj(text_dict.amr_feats)\n text_dict.text_feats = self.text_proj(text_dict.text_feats)\n else:\n raise ValueError()\n \n return ret, text_dict\n\n\n\n@META_ARCH_REGISTRY.register()\nclass VideoConv2d_TextLinear(nn.Module):\n \"\"\"\n 如果multiscale_shapes是None, 那么每个multiscale_shape的input_dim都是out_dim\n 如果multiscale_shapes给出了, 那么按照multiscale shapes里的dim\n \"\"\"\n def __init__(self,\n configs,\n out_dim,\n text_dim=None, # 如果是none的话, 那么假设等于out_dim\n multiscale_shapes=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n text_dim = out_dim if text_dim is None else text_dim\n multiscale_projs_config = configs['video_multiscale_projs']\n proj_names = multiscale_projs_config.keys() # list[str]\n\n in_dims = {}\n if multiscale_shapes is not None:\n assert set(proj_names).issubset(set(list(multiscale_shapes.keys())))\n for name in proj_names:\n in_dims[name] = multiscale_shapes[name].dim\n else:\n for name in proj_names:\n in_dims[name] = out_dim\n\n projs = {}\n for name, config in multiscale_projs_config.items():\n projs[name] = Conv2d_NormAct(in_channels=in_dims[name],\n out_channels=out_dim, \n **config)\n \n self.video_multiscale_projs = nn.ModuleDict(projs)\n\n text_proj_config = configs['text_proj']\n if text_proj_config is None:\n self.text_proj = nn.Identity()\n else:\n self.text_proj = Linear_NormAct(in_features=text_dim, out_features=out_dim, **text_proj_config)\n\n def forward(self, multiscales, text_dict):\n ret = {}\n for scale_name, scale_feat in multiscales.items(): # b c t h w\n if scale_name in self.video_multiscale_projs:\n batch_size, _, nf = scale_feat.shape[:3]\n scale_feat = rearrange(scale_feat, 'b c t h w -> (b t) c h w')\n scale_feat = self.video_multiscale_projs[scale_name](scale_feat)\n scale_feat = rearrange(scale_feat, '(b t) c h w -> b c t h w', b=batch_size, t=nf)\n ret[scale_name] = scale_feat\n else:\n ret[scale_name] = scale_feat\n \n if isinstance(text_dict, AMRData):\n text_dict.amr_feats = self.text_proj(text_dict.amr_feats)\n text_dict.text_feats = self.text_proj(text_dict.text_feats)\n else:\n raise ValueError()\n \n return ret, text_dict\n\n\n@META_ARCH_REGISTRY.register()\nclass ImageConv_MultiscaleProj(nn.Module):\n def __init__(self,\n configs,\n out_dim,\n multiscale_shapes=None,\n ) -> None:\n \"\"\"\n 如果multiscale_shape是空, 那么输入的dim = out_dim\n \"\"\"\n super().__init__()\n projs_configs = configs['projs']\n proj_names = list(projs_configs.keys()) # list[str]\n\n in_dims = {}\n if multiscale_shapes is not None:\n assert set(proj_names).issubset(set(list(multiscale_shapes.keys())))\n for name in proj_names:\n in_dims[name] = multiscale_shapes[name].dim\n else:\n for name in proj_names:\n in_dims[name] = out_dim\n\n projs = {}\n for name, config in projs_configs.items():\n projs[name] = Conv2d_NormAct(in_channels=in_dims[name], out_channels=out_dim, \n **config)\n self.multiscale_projs = nn.ModuleDict(projs)\n\n def forward(self, multiscales):\n ret = {}\n for scale_name, scale_feat in multiscales.items():\n if scale_name in self.multiscale_projs:\n scale_feat = self.multiscale_projs[scale_name](scale_feat)\n ret[scale_name] = scale_feat\n else:\n ret[scale_name] = scale_feat\n return ret\n\n\n@META_ARCH_REGISTRY.register()\nclass Video2D_ImageConv_MultiscaleProj(nn.Module):\n def __init__(self,\n configs,\n out_dim,\n multiscale_shapes=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n self.image_homo = ImageConv_MultiscaleProj(configs=configs, out_dim=out_dim, multiscale_shapes=multiscale_shapes)\n \n def forward(self, multiscales):\n batch_sisze, _, nf = multiscales[list(multiscales.keys())[0]].shape[:3]\n # b c t h w -> bt c h w\n multiscales = {key: value.permute(0, 2, 1, 3, 4).flatten(0, 1).contiguous() for key,value in multiscales.items()}\n multiscales = self.image_homo(multiscales)\n multiscales = {key: rearrange(value, '(b t) c h w -> b c t h w',b=batch_sisze, t=nf).contiguous()\\\n for key,value in multiscales.items()}\n return multiscales \n\n@META_ARCH_REGISTRY.register()\nclass VideoConv_MultiscaleProj(nn.Module):\n def __init__(self,\n configs,\n out_dim,\n multiscale_shapes=None,\n ) -> None:\n \"\"\"\n 如果multiscale_shape是空, 那么输入的dim = out_dim\n \"\"\"\n super().__init__()\n projs_configs = configs['projs']\n proj_names = list(projs_configs.keys()) # list[str]\n\n in_dims = {}\n if multiscale_shapes is not None:\n assert set(proj_names).issubset(set(list(multiscale_shapes.keys())))\n for name in proj_names:\n in_dims[name] = multiscale_shapes[name].dim\n else:\n for name in proj_names:\n in_dims[name] = out_dim\n\n projs = {}\n for name, config in projs_configs.items():\n projs[name] = Conv3d_NormAct(in_channels=in_dims[name], out_channels=out_dim, \n **config)\n self.multiscale_projs = nn.ModuleDict(projs)\n\n def forward(self, multiscales):\n ret = {}\n for scale_name, scale_feat in multiscales.items():\n if scale_name in self.multiscale_projs:\n scale_feat = self.multiscale_projs[scale_name](scale_feat)\n ret[scale_name] = scale_feat\n else:\n ret[scale_name] = scale_feat\n return ret\n\n \n\n\n\n\n@META_ARCH_REGISTRY.register()\nclass FrameQueryLinear_TextLinear(nn.Module):\n def __init__(self,\n configs,\n out_dim,\n text_dim=None, # int\n query_dim=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n query_proj_config = configs['query_proj']\n query_dim = out_dim if query_dim is None else query_dim\n text_dim = out_dim if text_dim is None else text_dim\n\n self.query_proj = Linear_NormAct(in_features=query_dim, out_features=out_dim, **query_proj_config)\n text_proj_config = configs['text_proj']\n if text_proj_config is None:\n self.text_proj = nn.Identity()\n else:\n self.text_proj = Linear_NormAct(in_features=text_dim, out_features=out_dim, **text_proj_config)\n\n def forward(self, frame_query, text_dict):\n # b T nqf c\n # text_dict\n frame_query = self.query_proj(frame_query)\n\n if isinstance(text_dict, AMRData):\n text_dict.amr_feats = self.text_proj(text_dict.amr_feats)\n text_dict.text_feats = self.text_proj(text_dict.text_feats)\n else:\n raise ValueError()\n \n return frame_query, text_dict\n\n\n@META_ARCH_REGISTRY.register()\nclass VideoConv3d_FrameQueryLinear_TextLinear(nn.Module):\n def __init__(self,\n configs,\n out_dim,\n feat_dim=None,\n text_dim=None, # int\n query_dim=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n query_proj_config = configs['query_proj']\n feat_proj_config = configs['feat_proj']\n text_proj_config = configs['text_proj']\n\n feat_dim = out_dim if feat_dim is None else feat_dim\n query_dim = out_dim if query_dim is None else query_dim\n text_dim = out_dim if text_dim is None else text_dim\n\n self.query_proj = Linear_NormAct(in_features=query_dim, out_features=out_dim, **query_proj_config) if query_proj_config is not None else nn.Identity()\n self.text_proj = Linear_NormAct(in_features=text_dim, out_features=out_dim, **text_proj_config) if text_proj_config is not None else nn.Identity()\n self.feat_proj = Conv3d_NormAct(in_channels=feat_dim, out_channels=out_dim, **feat_proj_config)\n\n def forward(self, mask_feat, frame_query, text_dict):\n mask_feat = self.feat_proj(mask_feat)\n frame_query = self.query_proj(frame_query)\n if isinstance(text_dict, AMRData):\n text_dict.amr_feats = self.text_proj(text_dict.amr_feats)\n text_dict.text_feats = self.text_proj(text_dict.text_feats)\n else:\n raise ValueError()\n\n return mask_feat, frame_query, text_dict\n\n# 每一个module应该都把input进行一边proj, proj到自己的空间里\n@META_ARCH_REGISTRY.register()\nclass VideoConv3d_FrameQueryLinear(nn.Module):\n \"\"\"\n 如果multiscale_shapes是None, 那么每个multiscale_shape的input_dim都是out_dim\n 如果multiscale_shapes给出了, 那么按照multiscale shapes里的dim\n \"\"\"\n def __init__(self,\n configs,\n out_dim,\n query_dim=None, # 如果是none的话, 那么假设等于out_dim\n multiscale_shapes=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n query_dim = out_dim if query_dim is None else query_dim\n\n multiscale_projs_config = configs['video_multiscale_projs']\n proj_names = multiscale_projs_config.keys() # list[str]\n\n in_dims = {}\n if multiscale_shapes is not None:\n assert set(proj_names).issubset(set(list(multiscale_shapes.keys())))\n for name in proj_names:\n in_dims[name] = multiscale_shapes[name].dim\n else:\n for name in proj_names:\n in_dims[name] = out_dim\n\n projs = {}\n for name, config in multiscale_projs_config.items():\n projs[name] = Conv3d_NormAct(in_channels=in_dims[name],\n out_channels=out_dim, \n **config)\n \n self.video_multiscale_projs = nn.ModuleDict(projs)\n\n query_proj_config = configs['query_proj']\n if query_proj_config is None:\n self.query_proj = nn.Identity()\n else:\n self.query_proj = Linear_NormAct(in_features=query_dim, out_features=out_dim, **query_proj_config)\n\n def forward(self, multiscales, frame_queries):\n # b t nq c\n ret = {}\n for scale_name, scale_feat in multiscales.items(): # b c t h w\n if scale_name in self.video_multiscale_projs:\n scale_feat = self.video_multiscale_projs[scale_name](scale_feat)\n ret[scale_name] = scale_feat\n else:\n ret[scale_name] = scale_feat\n \n frame_queries = self.query_proj(frame_queries)\n return ret, frame_queries\n\n\n\n@META_ARCH_REGISTRY.register()\nclass FrameQueryLinear(nn.Module):\n def __init__(self,\n configs,\n out_dim,\n query_dim=None, # scale_name: (dim, [temporal_scale, spatial_scale])\n ) -> None:\n super().__init__()\n query_proj_config = configs['query_proj']\n query_dim = out_dim if query_dim is None else query_dim\n\n self.query_proj = Linear_NormAct(in_features=query_dim, out_features=out_dim, **query_proj_config)\n\n def forward(self, frame_query):\n # b T nqf c\n frame_query = self.query_proj(frame_query)\n \n return frame_query\n\n"
},
{
"path": "models/encoder/localGlobal.py",
"content": "# Copyright (c) Facebook, Inc. and its affiliates.\nimport logging\nimport numpy as np\nfrom typing import Callable, Dict, List, Optional, Tuple, Union\n\nimport fvcore.nn.weight_init as weight_init\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.nn.init import xavier_uniform_, constant_, uniform_, normal_\nfrom torch.cuda.amp import autocast\n\nfrom detectron2.config import configurable\nfrom detectron2.layers import Conv2d, ShapeSpec, get_norm\nfrom detectron2.modeling import META_ARCH_REGISTRY\nfrom models.layers.position_encoding import PositionEmbeddingSine\nfrom models.layers.utils import _get_clones, _get_activation_fn\nfrom .ops.modules import MSDeformAttn\n\n# MSDeformAttn Transformer encoder in deformable detr\nclass MSDeformAttnTransformerEncoderOnly(nn.Module):\n def __init__(self, d_model=256, nhead=8,\n num_encoder_layers=6, dim_feedforward=1024, dropout=0.1,\n activation=\"relu\",\n num_feature_levels=4, enc_n_points=4,\n add_local = False,\n add_global=False,\n local_configs=None,\n global_configs=None,\n frame_nqueries=None,\n ):\n super().__init__()\n\n self.d_model = d_model\n self.nhead = nhead\n\n encoder_layer = MSDeformAttnTransformerEncoderLayer(d_model = d_model,\n d_ffn = dim_feedforward,\n dropout = dropout,\n activation = activation,\n n_levels = num_feature_levels,\n n_heads = nhead,\n n_points = enc_n_points,\n add_local = add_local,\n add_global = add_global,\n local_configs = local_configs,\n global_configs = global_configs\n )\n self.encoder = MSDeformAttnTransformerEncoder(encoder_layer, num_encoder_layers,\n d_model=d_model,\n frame_nqueries=frame_nqueries)\n\n self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))\n\n self._reset_parameters()\n\n def _reset_parameters(self):\n for p in self.parameters():\n if p.dim() > 1:\n nn.init.xavier_uniform_(p)\n for m in self.modules():\n if isinstance(m, MSDeformAttn):\n m._reset_parameters()\n normal_(self.level_embed)\n\n def get_valid_ratio(self, mask):\n _, H, W = mask.shape # b h w\n valid_H = torch.sum(~mask[:, :, 0], 1) # b \n valid_W = torch.sum(~mask[:, 0, :], 1) # b\n valid_ratio_h = valid_H.float() / H \n valid_ratio_w = valid_W.float() / W\n valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1) # b 2\n return valid_ratio\n\n def forward(self, \n srcs=None, \n pos_embeds=None,\n video_aux_dict=None,\n **kwargs):\n masks = [torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) for x in srcs]\n # prepare input for encoder\n src_flatten = []\n mask_flatten = []\n lvl_pos_embed_flatten = []\n spatial_shapes = []\n for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):\n bs, c, h, w = src.shape\n spatial_shape = (h, w)\n spatial_shapes.append(spatial_shape)\n src = src.flatten(2).transpose(1, 2)\n mask = mask.flatten(1)\n pos_embed = pos_embed.flatten(2).transpose(1, 2)\n lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)\n lvl_pos_embed_flatten.append(lvl_pos_embed)\n src_flatten.append(src)\n mask_flatten.append(mask)\n src_flatten = torch.cat(src_flatten, 1)\n mask_flatten = torch.cat(mask_flatten, 1)\n lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)\n spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)\n level_start_index = torch.cat((spatial_shapes.new_zeros((1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))\n valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1) # b #scale 2\n\n # encoder\n memory, frame_feats, frame_poses = self.encoder(src=src_flatten, # bt hw_sigma c\n spatial_shapes=spatial_shapes, \n level_start_index=level_start_index, \n valid_ratios=valid_ratios, \n pos=lvl_pos_embed_flatten, \n padding_mask=mask_flatten,\n video_aux_dict=video_aux_dict)\n\n return memory, spatial_shapes, level_start_index, frame_feats, frame_poses\n\nclass MSDeformAttnTransformerEncoderLayer(nn.Module):\n def __init__(self,\n d_model=256, d_ffn=1024,\n dropout=0.1, activation=\"relu\",\n n_levels=4, n_heads=8, n_points=4,\n add_local=False,\n add_global=False,\n local_configs=None,\n global_configs=None):\n super().__init__()\n # deform2d\n self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)\n self.dropout1 = nn.Dropout(dropout)\n self.norm1 = nn.LayerNorm(d_model)\n\n self.add_local = add_local\n if self.add_local:\n from .neighborhood_qk import NA_qk_Layer\n # self\n self.local_cnp = NA_qk_Layer(d_model=d_model, configs=local_configs)\n\n\n self.add_global = add_global\n if self.add_global:\n from models.layers.decoder_layers import CrossAttentionLayer \n # cross\n self.frame_query_cross_multiscale = CrossAttentionLayer(d_model=d_model, nhead=8, dropout=0.0,\n activation=\"relu\", normalize_before=False)\n self.cross_num_heads = 8\n self.global_add_attn_mask = global_configs['add_attn_mask'] if 'add_attn_mask' in global_configs else False\n # self+ffn\n from models.encoder.ops.modules.frame_query_ss2d import FrameQuery_SS2DLayer_hilbert \n self.global_hiss = FrameQuery_SS2DLayer_hilbert(global_configs)\n\n self.multiscale_cross_query = CrossAttentionLayer(d_model=d_model, nhead=8, dropout=0.0,\n activation=\"relu\", normalize_before=False)\n # ffn\n self.linear1 = nn.Linear(d_model, d_ffn)\n self.activation = _get_activation_fn(activation)\n self.dropout2 = nn.Dropout(dropout)\n self.linear2 = nn.Linear(d_ffn, d_model)\n self.dropout3 = nn.Dropout(dropout)\n self.norm2 = nn.LayerNorm(d_model)\n\n @staticmethod\n def with_pos_embed(tensor, pos):\n return tensor if pos is None else tensor + pos\n\n def forward_ffn(self, src):\n src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))\n src = src + self.dropout3(src2)\n src = self.norm2(src)\n return src\n\n @torch.no_grad()\n def get_attn_mask(self, frame_query_feats, src, spatial_shapes, level_start_index,):\n # nq bt c\n # bt hw_sigma c\n assert len(spatial_shapes) == 3\n frame_query_feats = frame_query_feats.permute(1, 0, 2) # bt nq c\n feat = src[:, level_start_index[-1]: (level_start_index[-1] + spatial_shapes[-1][0] * spatial_shapes[-1][1])]\n feat = rearrange(feat, 'b (h w) c -> b c h w',h=spatial_shapes[-1][0],w=spatial_shapes[-1][1])\n mask = torch.einsum('bnc, bchw -> b n h w',frame_query_feats, feat)\n mask_2 = F.interpolate(mask, size=spatial_shapes[0].tolist(), mode='bilinear',align_corners=False)\n mask_3 = F.interpolate(mask, size=spatial_shapes[1].tolist(), mode='bilinear', align_corners=False)\n attn_mask = torch.cat([mask_2.flatten(2), mask_3.flatten(2), mask.flatten(2)], dim=-1) #bt n hw_sigma\n attn_mask = (attn_mask.unsqueeze(1).repeat(1, self.cross_num_heads, 1, 1).flatten(0, 1).sigmoid() < 0.5).bool()\n return attn_mask\n\n \n def forward(self, \n src=None, pos=None, \n reference_points=None, spatial_shapes=None, level_start_index=None, padding_mask=None,\n video_aux_dict=None,\n frame_query_feats=None, # nq bt c\n frame_query_poses=None):\n \n if self.add_local:\n # local_self\n src = self.local_cnp(tgt=src, \n scale_shapes=spatial_shapes,\n level_start_idxs=level_start_index,\n nf=video_aux_dict['nf']) \n if self.add_global:\n if self.global_add_attn_mask:\n attn_mask = self.get_attn_mask(frame_query_feats, src, spatial_shapes, level_start_index,) # bthead nq hw\n attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False # 全masked掉的 全注意, 比如有padding\n else:\n attn_mask = None\n # cross\n frame_query_feats = self.frame_query_cross_multiscale(\n tgt=frame_query_feats, # nq bt c\n memory=src.permute(1, 0, 2), # hw_sigma bt c\n memory_mask=attn_mask,\n memory_key_padding_mask=None,\n pos= pos.permute(1,0,2),\n query_pos=frame_query_poses,\n )\n # self+ffn\n frame_query_feats = self.global_hiss(frame_query_feats=frame_query_feats,\n frame_query_poses=frame_query_poses,\n video_aux_dict=video_aux_dict)\n\n # self\n src = self.multiscale_cross_query(\n tgt=src.permute(1, 0, 2), # hw_sigma bt c\n memory=frame_query_feats, # nq bt c\n memory_mask=None,\n memory_key_padding_mask=None,\n pos= frame_query_poses,\n query_pos=pos.permute(1,0,2), \n ).permute(1, 0, 2)\n\n # self attention\n src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)[0]\n src = src + self.dropout1(src2)\n src = self.norm1(src)\n\n # ffn\n src = self.forward_ffn(src)\n return src, frame_query_feats\n\n\nclass MSDeformAttnTransformerEncoder(nn.Module):\n def __init__(self,\n encoder_layer=None, \n num_layers=None, \n d_model=None, frame_nqueries=None):\n super().__init__()\n self.layers = _get_clones(encoder_layer, num_layers)\n self.num_layers = num_layers\n\n self.frame_nqueries = frame_nqueries # 10\n self.frame_query_feats = nn.Embedding(self.frame_nqueries, d_model)\n self.frame_query_poses = nn.Embedding(self.frame_nqueries, d_model)\n\n @staticmethod\n def get_reference_points(spatial_shapes, valid_ratios, device):\n # b #scale 2, valid_w(0-1), valid_h(0-1), 整个feature map有多少是非padding的\n # list[h w] #scale\n reference_points_list = []\n for lvl, (H_, W_) in enumerate(spatial_shapes):\n ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),\n torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))\n ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_) # 1 hw / b 1 -> b hw(0-1), y的绝对坐标\n ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_) # 1 hw / b 1 -> b hw(0-1), x的绝对坐标\n ref = torch.stack((ref_x, ref_y), -1) # b hw 2\n reference_points_list.append(ref)\n reference_points = torch.cat(reference_points_list, 1) # b hw_sigma 2, 每个点的相对坐标(0-1)\n reference_points = reference_points[:, :, None] * valid_ratios[:, None] # b hw_sigma 1 2 * b 1 #scale 2\n return reference_points # b hw_sigma #scale 2\n\n def forward(self, \n src, # bt hw_sigma c\n spatial_shapes, \n level_start_index, \n valid_ratios, \n pos=None, \n padding_mask=None,\n video_aux_dict=None):\n \n output = src # bt hw_sigma c\n batch_size_nf = output.shape[0]\n frame_query_feats = self.frame_query_feats.weight.unsqueeze(1).repeat(1,batch_size_nf, 1)\n frame_query_poses = self.frame_query_poses.weight.unsqueeze(1).repeat(1,batch_size_nf,1) # n bt c\n \n reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)\n frame_feats = []\n for _, layer in enumerate(self.layers):\n output, frame_query_feats = layer(src=output, \n pos=pos, \n reference_points=reference_points, \n spatial_shapes=spatial_shapes, \n level_start_index=level_start_index, \n padding_mask=padding_mask,\n video_aux_dict=video_aux_dict,\n frame_query_feats=frame_query_feats,\n frame_query_poses=frame_query_poses)\n frame_feats.append(frame_query_feats)\n\n return output, frame_feats, frame_query_poses\n\n\nimport copy\nfrom einops import rearrange\nfrom models.layers.utils import _get_clones\nfrom models.layers.position_encoding import build_position_encoding\n# video multiscale, text_dict\n\n@META_ARCH_REGISTRY.register()\nclass Video_Deform2D_DividedTemporal_MultiscaleEncoder_localGlobal(nn.Module):\n def __init__(\n self,\n configs,\n multiscale_shapes, # {'res2': .temporal_stride, .spatial_stride, .dim}\n ):\n super().__init__()\n d_model = configs['d_model']\n fpn_norm = configs['fpn_norm'] # fpn的norm\n nlayers = configs['nlayers']\n\n # 4, 8, 16, 32\n self.multiscale_shapes = dict(sorted(copy.deepcopy(multiscale_shapes).items(), key=lambda x: x[1].spatial_stride))\n self.encoded_scales = sorted(configs['encoded_scales'], \n key=lambda x:self.multiscale_shapes[x].spatial_stride) # res3, res4, res5\n \n # 4 -> 8 -> 16 -> 32 \n self.scale_dims = [val.dim for val in multiscale_shapes.values()]\n self.video_projs = META_ARCH_REGISTRY.get(configs['video_projs']['name'])(configs=configs['video_projs'],\n multiscale_shapes=multiscale_shapes, out_dim=d_model)\n\n self.pos_2d = build_position_encoding(position_embedding_name='2d')\n\n deform_attn = configs['deform_attn']\n self.transformer = MSDeformAttnTransformerEncoderOnly(\n d_model=d_model,\n dropout=deform_attn['dropout'],\n nhead=deform_attn['nheads'],\n dim_feedforward=deform_attn['dim_feedforward'],\n activation=deform_attn['activation'],\n num_encoder_layers=nlayers,\n num_feature_levels=len(self.encoded_scales),\n enc_n_points=deform_attn['enc_n_points'],\n add_local = configs['add_local'],\n add_global = configs['add_global'],\n local_configs = configs['local_configs'],\n global_configs = configs['global_configs'],\n frame_nqueries=configs['frame_nqueries']\n \n )\n\n min_encode_stride = self.multiscale_shapes[self.encoded_scales[0]].spatial_stride # 8\n min_stride = list(self.multiscale_shapes.values())[0].spatial_stride # 4\n self.num_fpn_levels = int(np.log2(min_encode_stride) - np.log2(min_stride))\n lateral_convs = [] \n output_convs = []\n use_bias = fpn_norm == \"\"\n for idx, in_channels in enumerate(self.scale_dims[:self.num_fpn_levels]):\n lateral_norm = get_norm(fpn_norm, d_model)\n output_norm = get_norm(fpn_norm, d_model)\n\n lateral_conv = Conv2d(in_channels, d_model, kernel_size=1, bias=use_bias, norm=lateral_norm)\n output_conv = Conv2d(d_model, d_model, kernel_size=3, padding=1, bias=use_bias, norm=output_norm, activation=F.relu)\n\n self.add_module(\"adapter_{}\".format(idx + 1), lateral_conv)\n self.add_module(\"layer_{}\".format(idx + 1), output_conv)\n\n lateral_convs.append(lateral_conv)\n output_convs.append(output_conv)\n # Place convs into top-down order (from low to high resolution)\n # to make the top-down computation in forward clearer.\n self.lateral_convs = lateral_convs[::-1] # 8 4\n self.output_convs = output_convs[::-1] # 8 4\n\n def forward(self, \n multiscales=None, # b c t h w\n video_aux_dict=None, # dict{}\n **kwargs):\n batch_size, _, nf = multiscales[list(multiscales.keys())[0]].shape[:3]\n video_aux_dict['nf'] = nf\n multiscales = self.video_projs(multiscales) \n assert set(list(multiscales.keys())).issubset(set(list(self.multiscale_shapes.keys())))\n assert set(list(self.multiscale_shapes.keys())).issubset(set(list(multiscales.keys())))\n\n srcs = []\n poses = [] # 32, 16, 8\n for idx, scale_name in enumerate(self.encoded_scales[::-1]):\n x = multiscales[scale_name].permute(0, 2, 1, 3, 4).flatten(0,1).contiguous() # bt c h w\n srcs.append(x)\n poses.append(self.pos_2d(torch.zeros_like(x)[:, 0, :, :].bool(), hidden_dim=x.shape[1]))\n\n memory, spatial_shapes, level_start_index, frame_feats, frame_poses = self.transformer(srcs=srcs, \n pos_embeds=poses,\n video_aux_dict=video_aux_dict)\n bs = memory.shape[0]\n spatial_index = 0\n memory_features = [] # 32 16 8\n for lvl in range(len(self.encoded_scales)):\n h, w = spatial_shapes[lvl]\n memory_lvl = memory[:, spatial_index : spatial_index + h * w, :].reshape(bs, h, w, -1).permute(0, 3, 1, 2).contiguous() \n memory_features.append(memory_lvl)\n spatial_index += h * w\n\n for idx, f in enumerate(list(self.multiscale_shapes.keys())[:self.num_fpn_levels][::-1]):\n x = multiscales[f].permute(0, 2, 1, 3, 4).flatten(0,1).contiguous() # bt c h w\n cur_fpn = self.lateral_convs[idx](x)\n y = cur_fpn + F.interpolate(memory_features[-1], size=cur_fpn.shape[-2:], mode=\"bilinear\", align_corners=False)\n y = self.output_convs[idx](y)\n memory_features.append(y)\n\n assert len(memory_features) == len(list(self.multiscale_shapes.keys()))\n\n ret = {}\n for key, out_feat in zip(list(self.multiscale_shapes.keys()), memory_features[::-1]):\n ret[key] = rearrange(out_feat, '(b t) c h w -> b c t h w', b=batch_size, t=nf)\n return ret, frame_feats[::-1], frame_poses # 32, 16, 8\n\n"
},
{
"path": "models/encoder/neighborhood_qk.py",
"content": "from typing import Optional\n\nimport torch\nfrom torch import nn, Tensor\nfrom torch.nn.functional import pad\nfrom torch.nn.init import trunc_normal_\n\nfrom natten.functional import na2d_av, na2d_qk_with_bias\nfrom einops import rearrange\nfrom natten import NeighborhoodAttention2D\nfrom detectron2.modeling import META_ARCH_REGISTRY\n\nclass NeighborhoodAttention2D_qk(nn.Module):\n \"\"\"\n Neighborhood Attention 2D Module\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n num_heads: int,\n kernel_size: int,\n dilation: int = 1,\n bias: bool = True,\n qkv_bias: bool = True,\n qk_scale: Optional[float] = None,\n attn_drop: float = 0.0,\n proj_drop: float = 0.0,\n ):\n super().__init__()\n self.num_heads = num_heads\n self.head_dim = dim // self.num_heads\n self.scale = qk_scale or self.head_dim**-0.5\n assert (\n kernel_size > 1 and kernel_size % 2 == 1\n ), f\"Kernel size must be an odd number greater than 1, got {kernel_size}.\"\n self.kernel_size = kernel_size\n assert (\n dilation is None or dilation >= 1\n ), f\"Dilation must be greater than or equal to 1, got {dilation}.\"\n self.dilation = dilation or 1\n self.window_size = self.kernel_size * self.dilation\n\n self.q_linear = nn.Linear(dim, dim, bias=qkv_bias)\n self.kv_linear = nn.Linear(dim, dim * 2, bias=qkv_bias)\n if bias:\n self.rpb = nn.Parameter(\n torch.zeros(num_heads, (2 * kernel_size - 1), (2 * kernel_size - 1))\n )\n trunc_normal_(self.rpb, std=0.02, mean=0.0, a=-2.0, b=2.0)\n else:\n self.register_parameter(\"rpb\", None)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, \n x_q: Tensor, \n x_k: Tensor) -> Tensor:\n # bt h w c; bt h w c, 前一帧\n if x_q.dim() != 4:\n raise ValueError(\n f\"NeighborhoodAttention2D expected a rank-4 input tensor; got {x.dim()=}.\"\n )\n\n B, H, W, C = x_q.shape\n # Pad if the input is small than the minimum supported size\n H_padded, W_padded = H, W\n padding_h = padding_w = 0\n if H < self.window_size or W < self.window_size:\n padding_h = max(0, self.window_size - H_padded)\n padding_w = max(0, self.window_size - W_padded)\n x_q = pad(x_q, (0, 0, 0, padding_w, 0, padding_h))\n x_k = pad(x_k, (0, 0, 0, padding_w, 0, padding_h))\n _, H_padded, W_padded, _ = x_q.shape\n\n # b h w c -> b h w h c_h\n q = self.q_linear(x_q).reshape(B, H_padded, W_padded, self.num_heads, self.head_dim)\n q = q.permute(0, 3, 1, 2, 4) # b head h w c_h\n kv = (\n self.kv_linear(x_k)\n .reshape(B, H_padded, W_padded, 2, self.num_heads, self.head_dim)\n .permute(3, 0, 4, 1, 2, 5)\n ) # b \n k, v = kv[0], kv[1]\n q = q * self.scale\n attn = na2d_qk_with_bias(q, k, self.rpb, self.kernel_size, self.dilation)\n attn = attn.softmax(dim=-1)\n attn = self.attn_drop(attn)\n x_q = na2d_av(attn, v, self.kernel_size, self.dilation) # b head h w c_h\n x_q = x_q.permute(0, 2, 3, 1, 4).reshape(B, H_padded, W_padded, C) # b h w head c_h\n\n # Remove padding, if added any\n if padding_h or padding_w:\n x_q = x_q[:, :H, :W, :].contiguous()\n\n return self.proj_drop(self.proj(x_q))\n\n def extra_repr(self) -> str:\n return (\n f\"head_dim={self.head_dim}, num_heads={self.num_heads}, \"\n + f\"kernel_size={self.kernel_size}, dilation={self.dilation}, \"\n + f\"has_bias={self.rpb is not None}\"\n )\n\nfrom models.layers.utils import _get_clones\nclass NA_qk_Layer(nn.Module):\n\n def __init__(self, d_model, configs):\n super().__init__()\n self.self_attn = NeighborhoodAttention2D_qk(dim=configs['d_model'],\n num_heads=configs['num_heads'],\n kernel_size=configs['kernel_size'],\n dilation=configs['dilation'],\n bias=False,\n qkv_bias=False,)\n self.num_steps = configs['num_steps'] if 'num_steps' in configs else 1\n self.norm = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(configs['dropout'])\n \n\n def forward(self, tgt=None, scale_shapes=None, level_start_idxs=None, nf=None):\n # bt hw_sigma c -> list[b t h w c], 3\n video_feats = [tgt[:, start_idx:(start_idx + haosen[0]*haosen[1])].contiguous() for start_idx, haosen in zip(level_start_idxs, scale_shapes)]\n video_feats = [rearrange(haosen, '(b t) (h w) c -> b t h w c', t=nf, h=scale_shapes[idx][0], w=scale_shapes[idx][1]).contiguous() for idx, haosen in enumerate(video_feats)]\n\n video_key_feats = []\n for haosen in video_feats:\n scale_feats = torch.stack([torch.roll(haosen, shifts=k, dims=1) for k in range(1, self.num_steps+1)], dim=0) # s b t h w c\n video_key_feats.append(scale_feats.flatten(0, 2)) #sbt h w c\n \n # sbt h w c\n video_feats = [haosen.unsqueeze(0).repeat(self.num_steps, 1,1,1,1,1).flatten(0, 2) for haosen in video_feats]\n \n local_feats = [] # list[sbt h w c]\n for idx, (q_feat, k_feat) in enumerate(zip(video_feats, video_key_feats)):\n local_feats.append(self.self_attn(q_feat, k_feat))\n local_feats = [rearrange(haosen, '(s bt) h w c -> s bt h w c',s=self.num_steps) for haosen in local_feats]\n local_feats = [haosen.sum(dim=0) for haosen in local_feats] # bt h w c\n local_feats = torch.cat([haosen.flatten(1, 2) for haosen in local_feats], dim=1) # bt hw_sigma c\n tgt = tgt + self.dropout(local_feats)\n tgt = self.norm(tgt) \n return tgt\n\n@META_ARCH_REGISTRY.register()\nclass NA_qk_Layer_v2(nn.Module):\n\n def __init__(self, configs):\n super().__init__()\n self.self_attn = NeighborhoodAttention2D_qk(dim=configs['d_model'],\n num_heads=configs['num_heads'],\n kernel_size=configs['kernel_size'],\n dilation=configs['dilation'],\n bias=False,\n qkv_bias=False,)\n \n\n def forward(self,\n query=None, \n spatial_shapes=None,\n level_start_index=None,\n video_aux_dict=None,):\n \n # bt hw_sigma c -> list[b t h w c], 3\n video_feat = [query[:, start_idx:(start_idx + haosen[0]*haosen[1])].contiguous() for start_idx, haosen in zip(level_start_index, spatial_shapes)]\n video_feat = [rearrange(haosen, '(b t) (h w) c -> b t h w c',t=video_aux_dict['nf'], h=spatial_shapes[idx][0], w=spatial_shapes[idx][1]).contiguous() for idx, haosen in enumerate(video_feat)]\n video_key_feats = [torch.roll(haosen, shifts=1, dims=1).contiguous() for haosen in video_feat]\n\n local_feats = [] # list[bt h w c]\n for idx, (q_feat, k_feat) in enumerate(zip(video_feat, video_key_feats)):\n local_feats.append(self.self_attn(q_feat.flatten(0, 1), k_feat.flatten(0, 1)))\n \n local_feats = torch.cat([haosen.flatten(1, 2) for haosen in local_feats], dim=1) # bt hw_sigma c\n return local_feats, None\n"
},
{
"path": "models/encoder/ops/MultiScaleDeformableAttention.egg-info/PKG-INFO",
"content": "Metadata-Version: 2.1\nName: MultiScaleDeformableAttention\nVersion: 1.0\nSummary: PyTorch Wrapper for CUDA Functions of Multi-Scale Deformable Attention\nHome-page: https://github.com/fundamentalvision/Deformable-DETR\nAuthor: Weijie Su\n"
},
{
"path": "models/encoder/ops/attention.py",
"content": "from inspect import isfunction\nimport math\nimport torch\nfrom torch.nn.init import xavier_uniform_, constant_\nimport torch.nn.functional as F\nfrom torch import nn, einsum\nfrom einops import rearrange, repeat\nfrom .functions import MSDeformAttnFunction\nimport warnings\n\ndef exists(val):\n return val is not None\n\n\ndef uniq(arr):\n return{el: True for el in arr}.keys()\n\n\ndef default(val, d):\n if exists(val):\n return val\n return d() if isfunction(d) else d\n\n\ndef max_neg_value(t):\n return -torch.finfo(t.dtype).max\n\n\ndef init_(tensor):\n dim = tensor.shape[-1]\n std = 1 / math.sqrt(dim)\n tensor.uniform_(-std, std)\n return tensor\n\n\n# feedforward\nclass GEGLU(nn.Module):\n def __init__(self, dim_in, dim_out):\n super().__init__()\n self.proj = nn.Linear(dim_in, dim_out * 2)\n\n def forward(self, x):\n x, gate = self.proj(x).chunk(2, dim=-1)\n return x * F.gelu(gate)\n\n\nclass FeedForward(nn.Module):\n def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):\n super().__init__()\n inner_dim = int(dim * mult)\n dim_out = default(dim_out, dim)\n project_in = nn.Sequential(\n nn.Linear(dim, inner_dim),\n nn.GELU()\n ) if not glu else GEGLU(dim, inner_dim)\n\n self.net = nn.Sequential(\n project_in,\n nn.Dropout(dropout),\n nn.Linear(inner_dim, dim_out)\n )\n\n def forward(self, x):\n return self.net(x)\n\n\ndef zero_module(module):\n \"\"\"\n Zero out the parameters of a module and return it.\n \"\"\"\n for p in module.parameters():\n p.detach().zero_()\n return module\n\n\ndef Normalize(in_channels):\n return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)\n\n\nclass LinearAttention(nn.Module):\n def __init__(self, dim, heads=4, dim_head=32):\n super().__init__()\n self.heads = heads\n hidden_dim = dim_head * heads\n self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)\n self.to_out = nn.Conv2d(hidden_dim, dim, 1)\n\n def forward(self, x):\n b, c, h, w = x.shape\n qkv = self.to_qkv(x)\n q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)\n k = k.softmax(dim=-1) \n context = torch.einsum('bhdn,bhen->bhde', k, v)\n out = torch.einsum('bhde,bhdn->bhen', context, q)\n out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)\n return self.to_out(out)\n\n\nclass SpatialSelfAttention(nn.Module):\n def __init__(self, in_channels):\n super().__init__()\n self.in_channels = in_channels\n\n self.norm = Normalize(in_channels)\n self.q = torch.nn.Conv2d(in_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0)\n self.k = torch.nn.Conv2d(in_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0)\n self.v = torch.nn.Conv2d(in_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0)\n self.proj_out = torch.nn.Conv2d(in_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0)\n\n def forward(self, x):\n h_ = x\n h_ = self.norm(h_)\n q = self.q(h_)\n k = self.k(h_)\n v = self.v(h_)\n\n # compute attention\n b,c,h,w = q.shape\n q = rearrange(q, 'b c h w -> b (h w) c')\n k = rearrange(k, 'b c h w -> b c (h w)')\n w_ = torch.einsum('bij,bjk->bik', q, k)\n\n w_ = w_ * (int(c)**(-0.5))\n w_ = torch.nn.functional.softmax(w_, dim=2)\n\n # attend to values\n v = rearrange(v, 'b c h w -> b c (h w)')\n w_ = rearrange(w_, 'b i j -> b j i')\n h_ = torch.einsum('bij,bjk->bik', v, w_)\n h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)\n h_ = self.proj_out(h_)\n\n return x+h_\n\n\nclass CrossAttention(nn.Module):\n def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):\n super().__init__()\n inner_dim = dim_head * heads\n context_dim = default(context_dim, query_dim)\n\n self.scale = dim_head ** -0.5\n self.heads = heads\n\n self.to_q = nn.Linear(query_dim, inner_dim, bias=False)\n self.to_k = nn.Linear(context_dim, inner_dim, bias=False)\n self.to_v = nn.Linear(context_dim, inner_dim, bias=False)\n\n self.to_out = nn.Sequential(\n nn.Linear(inner_dim, query_dim),\n nn.Dropout(dropout)\n )\n\n def forward(self, x, context=None, mask=None):\n h = self.heads\n\n q = self.to_q(x)\n context = default(context, x)\n k = self.to_k(context)\n v = self.to_v(context)\n\n q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))\n\n sim = einsum('b i d, b j d -> b i j', q, k) * self.scale\n\n if exists(mask):\n mask = rearrange(mask, 'b ... -> b (...)')\n max_neg_value = -torch.finfo(sim.dtype).max\n mask = repeat(mask, 'b j -> (b h) () j', h=h)\n sim.masked_fill_(~mask, max_neg_value)\n\n # attention, what we cannot get enough of\n attn = sim.softmax(dim=-1)\n\n out = einsum('b i j, b j d -> b i d', attn, v)\n out = rearrange(out, '(b h) n d -> b n (h d)', h=h)\n return self.to_out(out)\n\nclass BasicTransformerBlock(nn.Module):\n def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True):\n super().__init__()\n self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout) # is a self-attention\n self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)\n self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,\n heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none\n self.norm1 = nn.LayerNorm(dim)\n self.norm2 = nn.LayerNorm(dim)\n self.norm3 = nn.LayerNorm(dim)\n\n def forward(self, x, context=None):\n x = self.attn1(self.norm1(x)) + x\n x = self.attn2(self.norm2(x), context=context) + x\n x = self.ff(self.norm3(x)) + x\n return x\n\n\ndef _is_power_of_2(n):\n if (not isinstance(n, int)) or (n < 0):\n raise ValueError(\"invalid input for _is_power_of_2: {} (type: {})\".format(n, type(n)))\n return (n & (n-1) == 0) and n != 0\n\nclass DeformAttn(nn.Module):\n def __init__(self, \n d_model=256,\n nheads=8,\n npoints=4,\n nlevels=4,\n key_dim=None):\n super().__init__()\n query_dim = d_model\n key_dim = d_model\n head_dim = d_model // nheads\n \n if d_model % nheads != 0:\n raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, nheads))\n # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation\n if not _is_power_of_2(head_dim):\n warnings.warn(\"You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 \"\n \"which is more efficient in our CUDA implementation.\")\n self.im2col_step = 64\n \n self.d_model = nheads * head_dim\n key_dim = default(key_dim, query_dim)\n \n self.value_proj = nn.Linear(key_dim, nheads * head_dim)\n self.sampling_offsets = nn.Linear(query_dim, nheads * nlevels * npoints * 2)\n \n self.attention_weights = nn.Linear(query_dim, nheads * nlevels * npoints)\n self.output_proj = nn.Linear(nheads * head_dim, query_dim)\n \n self.n_heads = nheads\n self.n_levels = nlevels\n self.head_dim = head_dim\n self.n_points = npoints\n self._reset_parameters()\n \n def _reset_parameters(self):\n constant_(self.sampling_offsets.weight.data, 0.)\n thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)\n grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)\n grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)\n for i in range(self.n_points):\n grid_init[:, :, i, :] *= i + 1\n with torch.no_grad():\n self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))\n constant_(self.attention_weights.weight.data, 0.)\n constant_(self.attention_weights.bias.data, 0.)\n xavier_uniform_(self.value_proj.weight.data)\n constant_(self.value_proj.bias.data, 0.)\n xavier_uniform_(self.output_proj.weight.data)\n constant_(self.output_proj.bias.data, 0.) \n\n def forward(self, query, reference_points, \n input_flatten, input_spatial_shapes, input_level_start_index,input_padding_mask=None):\n \"\"\"\n multi-scale deformable attention, self attention if query == input_flatten\n Input: \n - query: \n T(b n c)\n - reference_points: center or reference boxes, normalized, [0, 1], including padding area,\n add additional (w, h) to form reference boxes\n T(b n level 2) or T(b n level 4)\n - input_flatten: multi-scale特征\n T(b (h_\\sigma w_\\sigma) c)\n - input_spatial_shapes: 每个level的大小\n T(level 2)\n - input_level_start_index: [0, level1_start, level2_start]\n - input_padding_mask: True/False\n T(b), (h_\\sigma w_\\sigma))\n Output: \n - query results:\n T(b, n c)\n - sampling_locations: normalized\n T(b, n, m*l*k, 2)\n - attention_weights: after softmax\n T(b, n, m*l*k)\n \"\"\"\n batch_size, Nq, _ = query.shape\n _, Nk, _ = input_flatten.shape\n assert Nk == (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum()\n \n # B (h w) M * V\n value = self.value_proj(input_flatten)\n if input_padding_mask is not None:\n value = value.masked_fill(input_padding_mask[...,None], float(0))\n value = value.view(batch_size, Nk, self.n_heads, self.head_dim)\n sampling_offesets = self.sampling_offsets(query).view(batch_size, Nq, self.n_heads, self.n_levels, self.n_points, 2)\n attention_weights= self.attention_weights(query).view(batch_size, Nq, self.n_heads, self.n_levels * self.n_points)\n attention_weights = F.softmax(attention_weights, dim=-1).view(batch_size, Nq, self.n_heads, self.n_levels, self.n_points)\n # b, n ,head, level, point, 2\n if reference_points.shape[-1] == 2:\n # T(2 level)\n offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)\n sampling_locations = reference_points[:, :, None, :, None, :] + \\\n sampling_offesets / offset_normalizer[None, None, None, :, None,:]\n elif reference_points.shape[-1] == 4:\n sampling_locations = reference_points[:, :, None, :, None, :2] + \\\n sampling_offesets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5\n else:\n raise NotImplementedError\n output = MSDeformAttnFunction.apply(value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)\n output = self.output_proj(output)\n \n return output, sampling_locations, attention_weights\n \n \nclass ContextuallSelfAttention(nn.Module):\n def __init__(self, \n d_model, \n n_points,\n n_heads,\n context_dim=None):\n super().__init__()\n context_dim = default(context_dim, d_model)\n query_dim = key_dim = d_model\n \n if d_model % n_heads != 0:\n raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))\n head_dim = d_model // n_heads\n # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation\n if not _is_power_of_2(head_dim):\n warnings.warn(\"You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 \"\n \"which is more efficient in our CUDA implementation.\")\n self.im2col_step = 64\n \n self.d_model = d_model\n self.nheads = n_heads\n self.head_dim = head_dim\n self.npoints = n_points\n \n self.nlevels = 1\n \n self.value_proj = nn.Linear(key_dim, n_heads * head_dim)\n self.sampling_offsets = nn.Linear(query_dim, n_heads * n_points * 2)\n self.attention_weights = nn.Linear(query_dim, n_heads * n_points)\n \n self.output_proj = nn.Linear(n_heads * head_dim, query_dim)\n \n def forward(self,\n context, context_mask,\n query, reference_points, query_padding_mask = None,):\n \"\"\"\n contextual deformable attention\n Input: \n - context:\n T(b n c)\n - context_mask:\n T(b n)\n - query: \n T(b (h w) c)\n - reference_points: center or reference boxes, normalized, [0, 1], including padding area,\n T(b (h w) 2/4)\n - query_padding_mask:\n T(b (h w))\n Output: \n - query results:\n T(b, n c)\n - sampling_locations: normalized\n T(b, n, m*l*k, 2)\n - attention_weights: after softmax\n T(b, n, m*l*k)\n \"\"\"\n key = query\n key_padding_mask = query_padding_mask\n batch_size, Nq, _ = query.shape\n Nk = Nq\n \n input_spatial_shapes = torch.tensor(query.shape[-2:]).unsqueeze(0) # T(1, 2)\n input_level_start_index = [0, ]\n \n # B (h w) M * V\n value = self.value_proj(key)\n if key_padding_mask is not None:\n value = value.masked_fill(key_padding_mask[..., None], float(0))\n value = value.view(batch_size, Nk, self.nheads, self.head_dim)\n \n sampling_offesets = self.sampling_offsets(query).view(batch_size, Nq, self.n_heads, self.n_levels, self.n_points, 2)\n attention_weights= self.attention_weights(query).view(batch_size, Nq, self.n_heads, self.n_levels * self.n_points)\n attention_weights = F.softmax(attention_weights, dim=-1).view(batch_size, Nq, self.n_heads, self.n_levels, self.n_points)\n # b, n ,head, level, point, 2\n if reference_points.shape[-1] == 2:\n # T(2 level)\n offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)\n sampling_locations = reference_points[:, :, None, :, None, :] + \\\n sampling_offesets / offset_normalizer[None, None, None, :, None,:]\n elif reference_points.shape[-1] == 4:\n sampling_locations = reference_points[:, :, None, :, None, :2] + \\\n sampling_offesets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5\n else:\n raise NotImplementedError\n output = MSDeformAttnFunction.apply(value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)\n output = self.output_proj(output)\n \n return output, sampling_locations, attention_weights\n\nclass BasicTransformerBlock_v2(nn.Module):\n def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True):\n super().__init__()\n self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)\n \n self.norm1 = nn.LayerNorm(dim)\n self.norm2 = nn.LayerNorm(dim)\n\n def forward(self, x, context=None):\n x = self.norm1(self.attn1(x, context=context) + x)\n x = self.norm2(self.ff(x) + x)\n return x\n\n\n\nclass SpatialTransformer(nn.Module):\n \"\"\"\n Transformer block for image-like data.\n First, project the input (aka embedding)\n and reshape to b, t, d.\n Then apply standard transformer action.\n Finally, reshape to image\n \"\"\"\n def __init__(self, in_channels, n_heads, d_head,\n depth=1, dropout=0., context_dim=None):\n super().__init__()\n self.in_channels = in_channels\n inner_dim = n_heads * d_head\n self.norm = Normalize(in_channels)\n\n self.proj_in = nn.Conv2d(in_channels,\n inner_dim,\n kernel_size=1,\n stride=1,\n padding=0)\n\n self.transformer_blocks = nn.ModuleList(\n [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)\n for d in range(depth)]\n )\n\n self.proj_out = zero_module(nn.Conv2d(inner_dim,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0))\n\n def forward(self, x, context=None):\n # note: if no context is given, cross-attention defaults to self-attention\n b, c, h, w = x.shape\n x_in = x\n x = self.norm(x)\n x = self.proj_in(x)\n x = rearrange(x, 'b c h w -> b (h w) c')\n for block in self.transformer_blocks:\n x = block(x, context=context)\n x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)\n x = self.proj_out(x)\n return x + x_in\n "
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-311/functions/__init__.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom .ms_deform_attn_func import MSDeformAttnFunction\n\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-311/functions/ms_deform_attn_func.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.autograd import Function\nfrom torch.autograd.function import once_differentiable\n\nimport MultiScaleDeformableAttention as MSDA\n\n\nclass MSDeformAttnFunction(Function):\n @staticmethod\n def forward(ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step):\n ctx.im2col_step = im2col_step\n output = MSDA.ms_deform_attn_forward(\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, ctx.im2col_step)\n ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights)\n return output\n\n @staticmethod\n @once_differentiable\n def backward(ctx, grad_output):\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors\n grad_value, grad_sampling_loc, grad_attn_weight = \\\n MSDA.ms_deform_attn_backward(\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, grad_output, ctx.im2col_step)\n\n return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None\n\n\ndef ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):\n # for debug and test only,\n # need to use cuda version instead\n N_, S_, M_, D_ = value.shape\n _, Lq_, M_, L_, P_, _ = sampling_locations.shape\n value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)\n sampling_grids = 2 * sampling_locations - 1\n sampling_value_list = []\n for lid_, (H_, W_) in enumerate(value_spatial_shapes):\n # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_\n value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_*M_, D_, H_, W_)\n # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2\n sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)\n # N_*M_, D_, Lq_, P_\n sampling_value_l_ = F.grid_sample(value_l_, sampling_grid_l_,\n mode='bilinear', padding_mode='zeros', align_corners=False)\n sampling_value_list.append(sampling_value_l_)\n # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_)\n attention_weights = attention_weights.transpose(1, 2).reshape(N_*M_, 1, Lq_, L_*P_)\n output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_*D_, Lq_)\n return output.transpose(1, 2).contiguous()\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-311/modules/__init__.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom .ms_deform_attn import MSDeformAttn\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-311/modules/ms_deform_attn.py",
"content": "# Modify for sample points visualization\n# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport warnings\nimport math\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.nn.init import xavier_uniform_, constant_\n\nfrom ..functions import MSDeformAttnFunction\n\n\ndef _is_power_of_2(n):\n if (not isinstance(n, int)) or (n < 0):\n raise ValueError(\"invalid input for _is_power_of_2: {} (type: {})\".format(n, type(n)))\n return (n & (n-1) == 0) and n != 0\n\n\nclass MSDeformAttn(nn.Module):\n def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):\n \"\"\"\n Multi-Scale Deformable Attention Module\n :param d_model hidden dimension\n :param n_levels number of feature levels\n :param n_heads number of attention heads\n :param n_points number of sampling points per attention head per feature level\n \"\"\"\n super().__init__()\n if d_model % n_heads != 0:\n raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))\n _d_per_head = d_model // n_heads\n # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation\n if not _is_power_of_2(_d_per_head):\n warnings.warn(\"You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 \"\n \"which is more efficient in our CUDA implementation.\")\n\n self.im2col_step = 64\n\n self.d_model = d_model\n self.n_levels = n_levels\n self.n_heads = n_heads\n self.n_points = n_points\n\n self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)\n self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)\n self.value_proj = nn.Linear(d_model, d_model)\n self.output_proj = nn.Linear(d_model, d_model)\n\n self._reset_parameters()\n\n def _reset_parameters(self):\n constant_(self.sampling_offsets.weight.data, 0.)\n thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)\n grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)\n grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)\n for i in range(self.n_points):\n grid_init[:, :, i, :] *= i + 1\n with torch.no_grad():\n self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))\n constant_(self.attention_weights.weight.data, 0.)\n constant_(self.attention_weights.bias.data, 0.)\n xavier_uniform_(self.value_proj.weight.data)\n constant_(self.value_proj.bias.data, 0.)\n xavier_uniform_(self.output_proj.weight.data)\n constant_(self.output_proj.bias.data, 0.)\n\n def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None):\n \"\"\"\n :param query (N, Length_{query}, C)\n :param reference_points (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area\n or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes\n :param input_flatten (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l, C)\n :param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]\n :param input_level_start_index (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]\n :param input_padding_mask (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l), True for padding elements, False for non-padding elements\n\n :return output (N, Length_{query}, C)\n \"\"\"\n N, Len_q, _ = query.shape\n N, Len_in, _ = input_flatten.shape\n assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in\n\n value = self.value_proj(input_flatten)\n if input_padding_mask is not None:\n value = value.masked_fill(input_padding_mask[..., None], float(0))\n value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)\n sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)\n attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)\n attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)\n # N, Len_q, n_heads, n_levels, n_points, 2\n if reference_points.shape[-1] == 2:\n offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)\n sampling_locations = reference_points[:, :, None, :, None, :] \\\n + sampling_offsets / offset_normalizer[None, None, None, :, None, :]\n elif reference_points.shape[-1] == 4:\n sampling_locations = reference_points[:, :, None, :, None, :2] \\\n + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5\n else:\n raise ValueError(\n 'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))\n output = MSDeformAttnFunction.apply(\n value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)\n output = self.output_proj(output)\n\n return output, sampling_locations, attention_weights\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-38/functions/__init__.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom .ms_deform_attn_func import MSDeformAttnFunction\n\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-38/functions/ms_deform_attn_func.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.autograd import Function\nfrom torch.autograd.function import once_differentiable\n\nimport MultiScaleDeformableAttention as MSDA\n\n\nclass MSDeformAttnFunction(Function):\n @staticmethod\n def forward(ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step):\n ctx.im2col_step = im2col_step\n output = MSDA.ms_deform_attn_forward(\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, ctx.im2col_step)\n ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights)\n return output\n\n @staticmethod\n @once_differentiable\n def backward(ctx, grad_output):\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors\n grad_value, grad_sampling_loc, grad_attn_weight = \\\n MSDA.ms_deform_attn_backward(\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, grad_output, ctx.im2col_step)\n\n return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None\n\n\ndef ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):\n # for debug and test only,\n # need to use cuda version instead\n N_, S_, M_, D_ = value.shape\n _, Lq_, M_, L_, P_, _ = sampling_locations.shape\n value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)\n sampling_grids = 2 * sampling_locations - 1\n sampling_value_list = []\n for lid_, (H_, W_) in enumerate(value_spatial_shapes):\n # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_\n value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_*M_, D_, H_, W_)\n # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2\n sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)\n # N_*M_, D_, Lq_, P_\n sampling_value_l_ = F.grid_sample(value_l_, sampling_grid_l_,\n mode='bilinear', padding_mode='zeros', align_corners=False)\n sampling_value_list.append(sampling_value_l_)\n # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_)\n attention_weights = attention_weights.transpose(1, 2).reshape(N_*M_, 1, Lq_, L_*P_)\n output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_*D_, Lq_)\n return output.transpose(1, 2).contiguous()\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-38/modules/__init__.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom .ms_deform_attn import MSDeformAttn\n"
},
{
"path": "models/encoder/ops/build/lib.linux-x86_64-cpython-38/modules/ms_deform_attn.py",
"content": "# Modify for sample points visualization\n# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport warnings\nimport math\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.nn.init import xavier_uniform_, constant_\n\nfrom ..functions import MSDeformAttnFunction\n\n\ndef _is_power_of_2(n):\n if (not isinstance(n, int)) or (n < 0):\n raise ValueError(\"invalid input for _is_power_of_2: {} (type: {})\".format(n, type(n)))\n return (n & (n-1) == 0) and n != 0\n\n\nclass MSDeformAttn(nn.Module):\n def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):\n \"\"\"\n Multi-Scale Deformable Attention Module\n :param d_model hidden dimension\n :param n_levels number of feature levels\n :param n_heads number of attention heads\n :param n_points number of sampling points per attention head per feature level\n \"\"\"\n super().__init__()\n if d_model % n_heads != 0:\n raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))\n _d_per_head = d_model // n_heads\n # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation\n if not _is_power_of_2(_d_per_head):\n warnings.warn(\"You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 \"\n \"which is more efficient in our CUDA implementation.\")\n\n self.im2col_step = 64\n\n self.d_model = d_model\n self.n_levels = n_levels\n self.n_heads = n_heads\n self.n_points = n_points\n\n self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)\n self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)\n self.value_proj = nn.Linear(d_model, d_model)\n self.output_proj = nn.Linear(d_model, d_model)\n\n self._reset_parameters()\n\n def _reset_parameters(self):\n constant_(self.sampling_offsets.weight.data, 0.)\n thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)\n grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)\n grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)\n for i in range(self.n_points):\n grid_init[:, :, i, :] *= i + 1\n with torch.no_grad():\n self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))\n constant_(self.attention_weights.weight.data, 0.)\n constant_(self.attention_weights.bias.data, 0.)\n xavier_uniform_(self.value_proj.weight.data)\n constant_(self.value_proj.bias.data, 0.)\n xavier_uniform_(self.output_proj.weight.data)\n constant_(self.output_proj.bias.data, 0.)\n\n def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None):\n \"\"\"\n :param query (N, Length_{query}, C)\n :param reference_points (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area\n or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes\n :param input_flatten (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l, C)\n :param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]\n :param input_level_start_index (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]\n :param input_padding_mask (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l), True for padding elements, False for non-padding elements\n\n :return output (N, Length_{query}, C)\n \"\"\"\n N, Len_q, _ = query.shape\n N, Len_in, _ = input_flatten.shape\n assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in\n\n value = self.value_proj(input_flatten)\n if input_padding_mask is not None:\n value = value.masked_fill(input_padding_mask[..., None], float(0))\n value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)\n sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)\n attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)\n attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)\n # N, Len_q, n_heads, n_levels, n_points, 2\n if reference_points.shape[-1] == 2:\n offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)\n sampling_locations = reference_points[:, :, None, :, None, :] \\\n + sampling_offsets / offset_normalizer[None, None, None, :, None, :]\n elif reference_points.shape[-1] == 4:\n sampling_locations = reference_points[:, :, None, :, None, :2] \\\n + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5\n else:\n raise ValueError(\n 'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))\n output = MSDeformAttnFunction.apply(\n value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)\n output = self.output_proj(output)\n\n return output, sampling_locations, attention_weights\n"
},
{
"path": "models/encoder/ops/build/temp.linux-x86_64-cpython-311/.ninja_log",
"content": "# ninja log v5\n0\t5344\t1685604027\t/home/xhh/workspace/rvos_encoder/models/ops/build/temp.linux-x86_64-cpython-311/home/xhh/workspace/rvos_encoder/models/ops/src/cpu/ms_deform_attn_cpu.o\t1eaabdd4515aceab\n1\t20910\t1685604042\t/home/xhh/workspace/rvos_encoder/models/ops/build/temp.linux-x86_64-cpython-311/home/xhh/workspace/rvos_encoder/models/ops/src/vision.o\tb8641c4a4f7766f9\n0\t21063\t1685604042\t/home/xhh/workspace/rvos_encoder/models/ops/build/temp.linux-x86_64-cpython-311/home/xhh/workspace/rvos_encoder/models/ops/src/cuda/ms_deform_attn_cuda.o\td77fcd8ae1c377bb\n"
},
{
"path": "models/encoder/ops/build/temp.linux-x86_64-cpython-311/build.ninja",
"content": "ninja_required_version = 1.3\ncxx = c++\nnvcc = /usr/local/cuda/bin/nvcc\n\ncflags = -pthread -B /home/xhh/anaconda3/envs/natten/compiler_compat -DNDEBUG -fwrapv -O2 -Wall -fPIC -O2 -isystem /home/xhh/anaconda3/envs/natten/include -fPIC -O2 -isystem /home/xhh/anaconda3/envs/natten/include -fPIC -DWITH_CUDA -I/home/xhh/workspace/rvos_encoder/models/ops/src -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include/TH -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include/THC -I/usr/local/cuda/include -I/home/xhh/anaconda3/envs/natten/include/python3.11 -c\npost_cflags = -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE=\"_gcc\"' '-DPYBIND11_STDLIB=\"_libstdcpp\"' '-DPYBIND11_BUILD_ABI=\"_cxxabi1011\"' -DTORCH_EXTENSION_NAME=MultiScaleDeformableAttention -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++17\ncuda_cflags = -DWITH_CUDA -I/home/xhh/workspace/rvos_encoder/models/ops/src -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include/TH -I/home/xhh/anaconda3/envs/natten/lib/python3.11/site-packages/torch/include/THC -I/usr/local/cuda/include -I/home/xhh/anaconda3/envs/natten/include/python3.11 -c\ncuda_post_cflags = -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''\"'\"'-fPIC'\"'\"'' -DCUDA_HAS_FP16=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE=\"_gcc\"' '-DPYBIND11_STDLIB=\"_libstdcpp\"' '-DPYBIND11_BUILD_ABI=\"_cxxabi1011\"' -DTORCH_EXTENSION_NAME=MultiScaleDeformableAttention -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 -std=c++17\ncuda_dlink_post_cflags = \nldflags = \n\nrule compile\n command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags\n depfile = $out.d\n deps = gcc\n\nrule cuda_compile\n depfile = $out.d\n deps = gcc\n command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags\n\n\n\n\n\nbuild /home/xhh/workspace/rvos_encoder/models/ops/build/temp.linux-x86_64-cpython-311/home/xhh/workspace/rvos_encoder/models/ops/src/cpu/ms_deform_attn_cpu.o: compile /home/xhh/workspace/rvos_encoder/models/ops/src/cpu/ms_deform_attn_cpu.cpp\nbuild /home/xhh/workspace/rvos_encoder/models/ops/build/temp.linux-x86_64-cpython-311/home/xhh/workspace/rvos_encoder/models/ops/src/cuda/ms_deform_attn_cuda.o: cuda_compile /home/xhh/workspace/rvos_encoder/models/ops/src/cuda/ms_deform_attn_cuda.cu\nbuild /home/xhh/workspace/rvos_encoder/models/ops/build/temp.linux-x86_64-cpython-311/home/xhh/workspace/rvos_encoder/models/ops/src/vision.o: compile /home/xhh/workspace/rvos_encoder/models/ops/src/vision.cpp\n\n\n\n\n\n\n\n"
},
{
"path": "models/encoder/ops/functions/__init__.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom .ms_deform_attn_func import MSDeformAttnFunction\n\n"
},
{
"path": "models/encoder/ops/functions/ms_deform_attn_func.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.autograd import Function\nfrom torch.autograd.function import once_differentiable\n\nimport MultiScaleDeformableAttention as MSDA\n\n\nclass MSDeformAttnFunction(Function):\n @staticmethod\n def forward(ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step):\n ctx.im2col_step = im2col_step\n output = MSDA.ms_deform_attn_forward(\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, ctx.im2col_step)\n ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights)\n return output\n\n @staticmethod\n @once_differentiable\n def backward(ctx, grad_output):\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors\n grad_value, grad_sampling_loc, grad_attn_weight = \\\n MSDA.ms_deform_attn_backward(\n value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, grad_output, ctx.im2col_step)\n\n return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None\n\n\ndef ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):\n # for debug and test only,\n # need to use cuda version instead\n N_, S_, M_, D_ = value.shape\n _, Lq_, M_, L_, P_, _ = sampling_locations.shape\n value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)\n sampling_grids = 2 * sampling_locations - 1\n sampling_value_list = []\n for lid_, (H_, W_) in enumerate(value_spatial_shapes):\n # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_\n value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_*M_, D_, H_, W_)\n # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2\n sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)\n # N_*M_, D_, Lq_, P_\n sampling_value_l_ = F.grid_sample(value_l_, sampling_grid_l_,\n mode='bilinear', padding_mode='zeros', align_corners=False)\n sampling_value_list.append(sampling_value_l_)\n # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_)\n attention_weights = attention_weights.transpose(1, 2).reshape(N_*M_, 1, Lq_, L_*P_)\n output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_*D_, Lq_)\n return output.transpose(1, 2).contiguous()\n"
},
{
"path": "models/encoder/ops/make.sh",
"content": "#!/usr/bin/env bash\n# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\npython setup.py build install\n"
},
{
"path": "models/encoder/ops/modules/__init__.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom .ms_deform_attn import MSDeformAttn\nfrom . import frame_query_ss2d\n"
},
{
"path": "models/encoder/ops/modules/frame_query_ss2d.py",
"content": "\nfrom models.layers.position_encoding import build_position_encoding\nfrom mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref\nimport math\nimport warnings\nimport math\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.nn.init import xavier_uniform_, constant_\n\nfrom ..functions import MSDeformAttnFunction\n\nfrom mamba_ssm import Mamba\nfrom einops import rearrange, reduce, repeat\nfrom detectron2.modeling import META_ARCH_REGISTRY\n\n# v1\nclass SS2D(nn.Module):\n def __init__(\n self,\n d_model,\n d_state=16,\n # d_state=\"auto\", # 20240109\n d_conv=3,\n expand=2,\n dt_rank=\"auto\",\n dt_min=0.001,\n dt_max=0.1,\n dt_init=\"random\",\n dt_scale=1.0,\n dt_init_floor=1e-4,\n dropout=0.,\n conv_bias=True,\n bias=False,\n device=None,\n dtype=None,\n **kwargs,\n ):\n factory_kwargs = {\"device\": device, \"dtype\": dtype}\n super().__init__()\n self.d_model = d_model\n self.d_state = d_state\n # self.d_state = math.ceil(self.d_model / 6) if d_state == \"auto\" else d_model # 20240109\n self.d_conv = d_conv\n self.expand = expand\n self.d_inner = int(self.expand * self.d_model)\n self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == \"auto\" else dt_rank\n\n self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)\n self.conv2d = nn.Conv2d(\n in_channels=self.d_inner,\n out_channels=self.d_inner,\n groups=self.d_inner,\n bias=conv_bias,\n kernel_size=d_conv,\n padding=(d_conv - 1) // 2,\n **factory_kwargs,\n )\n self.act = nn.SiLU()\n\n self.x_proj = (\n nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs), \n nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs), \n nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs), \n nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs), \n )\n self.x_proj_weight = nn.Parameter(torch.stack([t.weight for t in self.x_proj], dim=0)) # (K=4, N, inner)\n del self.x_proj\n\n self.dt_projs = (\n self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs),\n self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs),\n self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs),\n self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs),\n )\n self.dt_projs_weight = nn.Parameter(torch.stack([t.weight for t in self.dt_projs], dim=0)) # (K=4, inner, rank)\n self.dt_projs_bias = nn.Parameter(torch.stack([t.bias for t in self.dt_projs], dim=0)) # (K=4, inner)\n del self.dt_projs\n \n self.A_logs = self.A_log_init(self.d_state, self.d_inner, copies=4, merge=True) # (K=4, D, N)\n self.Ds = self.D_init(self.d_inner, copies=4, merge=True) # (K=4, D, N)\n\n # self.selective_scan = selective_scan_fn\n self.forward_core = self.forward_corev0\n\n self.out_norm = nn.LayerNorm(self.d_inner)\n self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)\n self.dropout = nn.Dropout(dropout) if dropout > 0. else None\n\n @staticmethod\n def dt_init(dt_rank, d_inner, dt_scale=1.0, dt_init=\"random\", dt_min=0.001, dt_max=0.1, dt_init_floor=1e-4, **factory_kwargs):\n dt_proj = nn.Linear(dt_rank, d_inner, bias=True, **factory_kwargs)\n\n # Initialize special dt projection to preserve variance at initialization\n dt_init_std = dt_rank**-0.5 * dt_scale\n if dt_init == \"constant\":\n nn.init.constant_(dt_proj.weight, dt_init_std)\n elif dt_init == \"random\":\n nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)\n else:\n raise NotImplementedError\n\n # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max\n dt = torch.exp(\n torch.rand(d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))\n + math.log(dt_min)\n ).clamp(min=dt_init_floor)\n # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759\n inv_dt = dt + torch.log(-torch.expm1(-dt))\n with torch.no_grad():\n dt_proj.bias.copy_(inv_dt)\n # Our initialization would set all Linear.bias to zero, need to mark this one as _no_reinit\n dt_proj.bias._no_reinit = True\n \n return dt_proj\n\n @staticmethod\n def A_log_init(d_state, d_inner, copies=1, device=None, merge=True):\n # S4D real initialization\n A = repeat(\n torch.arange(1, d_state + 1, dtype=torch.float32, device=device),\n \"n -> d n\",\n d=d_inner,\n ).contiguous()\n A_log = torch.log(A) # Keep A_log in fp32\n if copies > 1:\n A_log = repeat(A_log, \"d n -> r d n\", r=copies)\n if merge:\n A_log = A_log.flatten(0, 1)\n A_log = nn.Parameter(A_log)\n A_log._no_weight_decay = True\n return A_log\n\n @staticmethod\n def D_init(d_inner, copies=1, device=None, merge=True):\n # D \"skip\" parameter\n D = torch.ones(d_inner, device=device)\n if copies > 1:\n D = repeat(D, \"n1 -> r n1\", r=copies)\n if merge:\n D = D.flatten(0, 1)\n D = nn.Parameter(D) # Keep in fp32\n D._no_weight_decay = True\n return D\n\n def forward_corev0(self, x: torch.Tensor):\n self.selective_scan = selective_scan_fn\n \n B, C, H, W = x.shape\n L = H * W\n K = 4\n\n x_hwwh = torch.stack([x.view(B, -1, L), torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, -1, L)], dim=1).view(B, 2, -1, L)\n xs = torch.cat([x_hwwh, torch.flip(x_hwwh, dims=[-1])], dim=1) # (b, k, d, l)\n\n x_dbl = torch.einsum(\"b k d l, k c d -> b k c l\", xs.view(B, K, -1, L), self.x_proj_weight)\n # x_dbl = x_dbl + self.x_proj_bias.view(1, K, -1, 1)\n dts, Bs, Cs = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=2)\n dts = torch.einsum(\"b k r l, k d r -> b k d l\", dts.view(B, K, -1, L), self.dt_projs_weight)\n # dts = dts + self.dt_projs_bias.view(1, K, -1, 1)\n\n xs = xs.float().view(B, -1, L) # (b, k * d, l)\n dts = dts.contiguous().float().view(B, -1, L) # (b, k * d, l)\n Bs = Bs.float().view(B, K, -1, L) # (b, k, d_state, l)\n Cs = Cs.float().view(B, K, -1, L) # (b, k, d_state, l)\n Ds = self.Ds.float().view(-1) # (k * d)\n As = -torch.exp(self.A_logs.float()).view(-1, self.d_state) # (k * d, d_state)\n dt_projs_bias = self.dt_projs_bias.float().view(-1) # (k * d)\n\n out_y = self.selective_scan(\n xs, dts, \n As, Bs, Cs, Ds, z=None,\n delta_bias=dt_projs_bias,\n delta_softplus=True,\n return_last_state=False,\n ).view(B, K, -1, L)\n assert out_y.dtype == torch.float\n\n inv_y = torch.flip(out_y[:, 2:4], dims=[-1]).view(B, 2, -1, L)\n wh_y = torch.transpose(out_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)\n invwh_y = torch.transpose(inv_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)\n\n return out_y[:, 0], inv_y[:, 0], wh_y, invwh_y\n\n def forward(self, x: torch.Tensor, **kwargs):\n B, H, W, C = x.shape\n xz = self.in_proj(x)\n x, z = xz.chunk(2, dim=-1) # (b, h, w, d)\n\n x = x.permute(0, 3, 1, 2).contiguous()\n x = self.act(self.conv2d(x)) # (b, d, h, w)\n y1, y2, y3, y4 = self.forward_core(x) # B C hw\n assert y1.dtype == torch.float32\n y = y1 + y2 + y3 + y4\n y = torch.transpose(y, dim0=1, dim1=2).contiguous().view(B, H, W, -1)\n y = self.out_norm(y)\n y = y * F.silu(z)\n out = self.out_proj(y)\n if self.dropout is not None:\n out = self.dropout(out)\n return out\n\nclass SS2D_FrameQuery(nn.Module):\n def __init__(self, configs,):\n super().__init__()\n d_model = configs['d_model']\n self.homo = SS2D(d_model=configs['d_model'],\n d_state=configs['d_state'] if 'd_state' in configs else 16,\n d_conv=configs['d_conv'] if 'd_conv' in configs else 3,\n expand=configs['expand'] if 'expand' in configs else 2,\n dt_rank=configs['dt_rank'] if 'dt_rank' in configs else 'auto',\n dt_min=configs['dt_min'] if 'dt_min' in configs else 0.001,\n dt_max=configs['dt_max'] if 'dt_max' in configs else 0.1,\n dt_init=configs['dt_init'] if 'dt_init' in configs else 'random',\n dt_scale=configs['dt_scale'] if 'dt_scale' in configs else 1.0,\n dt_init_floor=configs['dt_init_floor'] if 'dt_init_floor' in configs else 1e-4,\n dropout=configs['dropout'] if 'dropout' in configs else 0,\n conv_bias=configs['conv_bias'] if 'conv_bias' in configs else True,\n bias=configs['bias'] if 'bias' in configs else False,\n )\n self.pos_1d = build_position_encoding(position_embedding_name='1d') # t上的position embedding\n\n def forward(self, \n frame_query_feats=None, # n bt c \n frame_query_poses=None, # n bt c # nq上的Position embedding\n nf=None,\n **kwargs\n ):\n batch_size = frame_query_feats.shape[1] // nf # b\n frame_query_feats += frame_query_poses \n frame_query_feats = rearrange(frame_query_feats, 'n (b t) c -> b t n c',b=batch_size,t=nf).contiguous()\n\n sin_poses = self.pos_1d(torch.zeros_like(frame_query_feats[..., 0].permute(0, 2, 1).flatten(0, 1)).bool(), \n hidden_dim=frame_query_feats.shape[-1]) # bn c t\n sin_poses = rearrange(sin_poses, '(b n) c t -> b t n c', b=batch_size)\n frame_query_feats += sin_poses\n\n frame_query_feats = self.homo(frame_query_feats) # b t n c\n \n frame_query_feats = frame_query_feats.permute(2, 0, 1, 3).flatten(1, 2).contiguous() # n bt c\n\n return frame_query_feats, None\n\n@META_ARCH_REGISTRY.register()\nclass FrameQuery_SS2DLayer(nn.Module):\n def __init__(self, \n configs, \n dropout=0.0,\n activation=\"relu\", normalize_before=False):\n super().__init__()\n d_model = configs['d_model']\n dropout = configs['dropout']\n self.self_attn = SS2D_FrameQuery(configs)\n self.norm = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(dropout)\n\n from models.layers.decoder_layers import FFNLayer\n self.ffn = FFNLayer(d_model=d_model,\n dim_feedforward=configs['dim_feedforward'],\n dropout=dropout,)\n\n def forward(self, \n frame_query_feats, # n bt c \n frame_query_poses, # n bt c # nq上的Position embedding\n nf=None,\n **kwargs):\n \n tgt2 = self.self_attn(frame_query_feats=frame_query_feats, # n bt c \n frame_query_poses=frame_query_poses, # n bt c # nq上的Position embedding\n nf=nf,)[0]\n frame_query_feats = frame_query_feats + self.dropout(tgt2)\n frame_query_feats = self.norm(frame_query_feats)\n\n frame_query_feats = self.ffn(frame_query_feats)\n return frame_query_feats\n\nfrom models.layers.decoder_layers import CrossAttentionLayer, SelfAttentionLayer, FFNLayer\n@META_ARCH_REGISTRY.register()\nclass TemporalQuery_CrossSelf(nn.Module):\n def __init__(self, configs) -> None:\n super().__init__()\n d_model = configs['d_model']\n attn_configs = configs['attn']\n self.cross_layers = CrossAttentionLayer(d_model=d_model,\n nhead=attn_configs['nheads'],\n dropout=0.0,\n normalize_before=attn_configs['normalize_before'])\n self.self_layers = SelfAttentionLayer(d_model=d_model,\n nhead=attn_configs['nheads'],\n dropout=0.0,\n normalize_before=attn_configs['normalize_before'])\n self.ffn_layers = FFNLayer(d_model=d_model,\n dim_feedforward=attn_configs['dim_feedforward'],\n dropout=0.0,\n normalize_before=attn_configs['normalize_before'])\n \n def forward(self, \n temporal_query_feats, \n temporal_query_poses,\n frame_query_feats, frame_query_poses,\n video_aux_dict=None, **kwargs):\n # nq b c; nq bt c\n nq, batch_size, _ = temporal_query_feats.shape\n nf = frame_query_feats.shape[1] // batch_size\n nqf = frame_query_feats.shape[0]\n frame_query_feats = rearrange(frame_query_feats, 'nq (b t) c -> (t nq) b c',b=batch_size, t=nf)\n frame_query_poses = rearrange(frame_query_poses, 'nq (b t) c -> (t nq) b c',b=batch_size, t=nf)\n temporal_query_feats = self.cross_layers(\n tgt=temporal_query_feats, # n b c\n memory=frame_query_feats, # t nqf b c\n pos=frame_query_poses, \n query_pos=temporal_query_poses,\n )\n temporal_query_feats = self.self_layers(\n temporal_query_feats,\n query_pos=temporal_query_poses,\n )\n temporal_query_feats = self.ffn_layers(\n temporal_query_feats \n )\n return temporal_query_feats\n \n# v2 多层\nclass SS2D_FrameQuery_v2(nn.Module):\n def __init__(self, configs,):\n super().__init__()\n d_model = configs['d_model']\n self.homo = SS2D(d_model=configs['d_model'],\n d_state=configs['d_state'] if 'd_state' in configs else 16,\n d_conv=configs['d_conv'] if 'd_conv' in configs else 3,\n expand=configs['expand'] if 'expand' in configs else 2,\n dt_rank=configs['dt_rank'] if 'dt_rank' in configs else 'auto',\n dt_min=configs['dt_min'] if 'dt_min' in configs else 0.001,\n dt_max=configs['dt_max'] if 'dt_max' in configs else 0.1,\n dt_init=configs['dt_init'] if 'dt_init' in configs else 'random',\n dt_scale=configs['dt_scale'] if 'dt_scale' in configs else 1.0,\n dt_init_floor=configs['dt_init_floor'] if 'dt_init_floor' in configs else 1e-4,\n dropout=configs['dropout'] if 'dropout' in configs else 0,\n conv_bias=configs['conv_bias'] if 'conv_bias' in configs else True,\n bias=configs['bias'] if 'bias' in configs else False,\n )\n self.pos_1d = build_position_encoding(position_embedding_name='1d') # t上的position embedding\n self.norm = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(configs['dropout'])\n\n def forward(self, \n frame_query_feats=None, # n bt c \n frame_query_poses=None, # n bt c # nq上的Position embedding\n nf=None,\n **kwargs\n ):\n batch_size = frame_query_feats.shape[1] // nf # b\n tgt2 = frame_query_feats + frame_query_poses \n tgt2 = rearrange(tgt2, 'n (b t) c -> b t n c',b=batch_size,t=nf).contiguous()\n\n sin_poses = self.pos_1d(torch.zeros_like(tgt2[..., 0].permute(0, 2, 1).flatten(0, 1)).bool(), \n hidden_dim=tgt2.shape[-1]) # bn c t\n sin_poses = rearrange(sin_poses, '(b n) c t -> b t n c', b=batch_size)\n tgt2 += sin_poses\n\n tgt2 = self.homo(tgt2) # b t n c\n \n tgt2 = tgt2.permute(2, 0, 1, 3).flatten(1, 2).contiguous() # n bt c\n\n\n frame_query_feats = frame_query_feats + self.dropout(tgt2)\n frame_query_feats = self.norm(frame_query_feats)\n\n return frame_query_feats, None\n \nfrom models.layers.utils import _get_clones\n@META_ARCH_REGISTRY.register()\nclass FrameQuery_SS2DLayer_v2(nn.Module):\n def __init__(self, \n configs, \n dropout=0.0):\n super().__init__()\n d_model = configs['d_model']\n n_layers = configs['nlayers'] if 'nlayers' in configs else 1\n self.nlayers = n_layers\n self.self_attn = _get_clones(SS2D_FrameQuery_v2(configs), n_layers)\n\n from models.layers.decoder_layers import FFNLayer\n self.ffn = FFNLayer(d_model=d_model,\n dim_feedforward=configs['dim_feedforward'],\n dropout=configs['dropout'],)\n\n def forward(self, \n frame_query_feats, # n bt c \n frame_query_poses, # n bt c # nq上的Position embedding\n nf=None,\n **kwargs):\n \n for i in range(self.nlayers):\n frame_query_feats = self.self_attn[i](frame_query_feats=frame_query_feats, # n bt c \n frame_query_poses=frame_query_poses, # n bt c # nq上的Position embedding\n nf=nf,)[0]\n \n frame_query_feats = self.ffn(frame_query_feats)\n return frame_query_feats\n\n\nclass Hilbert_2DSelectiveScan(nn.Module):\n def __init__(\n self,\n d_model,\n d_state=16,\n # d_state=\"auto\", # 20240109\n d_conv=3,\n expand=2,\n dt_rank=\"auto\",\n dt_min=0.001,\n dt_max=0.1,\n dt_init=\"random\",\n dt_scale=1.0,\n dt_init_floor=1e-4,\n dropout=0.,\n conv_bias=True,\n bias=False,\n device=None,\n dtype=None,\n scan_order=None,\n **kwargs,\n ):\n factory_kwargs = {\"device\": device, \"dtype\": dtype}\n super().__init__()\n self.d_model = d_model\n self.d_state = d_state\n # self.d_state = math.ceil(self.d_model / 6) if d_state == \"auto\" else d_model # 20240109\n self.d_conv = d_conv\n self.expand = expand\n self.d_inner = int(self.expand * self.d_model)\n self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == \"auto\" else dt_rank\n\n self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)\n self.conv2d = nn.Conv2d(\n in_channels=self.d_inner,\n out_channels=self.d_inner,\n groups=self.d_inner,\n bias=conv_bias,\n kernel_size=d_conv,\n padding=(d_conv - 1) // 2,\n **factory_kwargs,\n )\n self.act = nn.SiLU()\n\n self.x_proj = (\n nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs), \n nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs), \n )\n self.x_proj_weight = nn.Parameter(torch.stack([t.weight for t in self.x_proj], dim=0)) # (K=2, N, inner)\n del self.x_proj\n\n self.dt_projs = (\n self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs),\n self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs),\n )\n self.dt_projs_weight = nn.Parameter(torch.stack([t.weight for t in self.dt_projs], dim=0)) # (K=2, inner, rank)\n self.dt_projs_bias = nn.Parameter(torch.stack([t.bias for t in self.dt_projs], dim=0)) # (K=2, inner)\n del self.dt_projs\n \n self.A_logs = self.A_log_init(self.d_state, self.d_inner, copies=2, merge=True) # (K=2, D, N)\n self.Ds = self.D_init(self.d_inner, copies=2, merge=True) # (K=2, D, N)\n\n # self.selective_scan = selective_scan_fn\n self.forward_core = self.forward_corev0\n\n self.out_norm = nn.LayerNorm(self.d_inner)\n self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)\n self.dropout = nn.Dropout(dropout) if dropout > 0. else None\n self.scan_order = scan_order\n\n @staticmethod\n def dt_init(dt_rank, d_inner, dt_scale=1.0, dt_init=\"random\", dt_min=0.001, dt_max=0.1, dt_init_floor=1e-4, **factory_kwargs):\n dt_proj = nn.Linear(dt_rank, d_inner, bias=True, **factory_kwargs)\n\n # Initialize special dt projection to preserve variance at initialization\n dt_init_std = dt_rank**-0.5 * dt_scale\n if dt_init == \"constant\":\n nn.init.constant_(dt_proj.weight, dt_init_std)\n elif dt_init == \"random\":\n nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)\n else:\n raise NotImplementedError\n\n # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max\n dt = torch.exp(\n torch.rand(d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))\n + math.log(dt_min)\n ).clamp(min=dt_init_floor)\n # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759\n inv_dt = dt + torch.log(-torch.expm1(-dt))\n with torch.no_grad():\n dt_proj.bias.copy_(inv_dt)\n # Our initialization would set all Linear.bias to zero, need to mark this one as _no_reinit\n dt_proj.bias._no_reinit = True\n \n return dt_proj\n\n @staticmethod\n def A_log_init(d_state, d_inner, copies=1, device=None, merge=True):\n # S4D real initialization\n A = repeat(\n torch.arange(1, d_state + 1, dtype=torch.float32, device=device),\n \"n -> d n\",\n d=d_inner,\n ).contiguous()\n A_log = torch.log(A) # Keep A_log in fp32\n if copies > 1:\n A_log = repeat(A_log, \"d n -> r d n\", r=copies)\n if merge:\n A_log = A_log.flatten(0, 1)\n A_log = nn.Parameter(A_log)\n A_log._no_weight_decay = True\n return A_log\n\n @staticmethod\n def D_init(d_inner, copies=1, device=None, merge=True):\n # D \"skip\" parameter\n D = torch.ones(d_inner, device=device)\n if copies > 1:\n D = repeat(D, \"n1 -> r n1\", r=copies)\n if merge:\n D = D.flatten(0, 1)\n D = nn.Parameter(D) # Keep in fp32\n D._no_weight_decay = True\n return D\n\n def forward_corev0(self, x: torch.Tensor, hilbert_curve):\n # LongTensor[int] 按照hw进行flatten之后的hilbert排序\n self.selective_scan = selective_scan_fn\n \n B, C, H, W = x.shape\n L = H * W\n K = 2\n\n if self.scan_order == 'zigzag':\n x_hw = x.view(B, -1, L).contiguous() # b c hw\n xs = torch.stack([x_hw, torch.flip(x_hw, dims=[-1])], dim=1) # (b, k, d, l)\n elif self.scan_order == 'hilbert':\n x_hw = x.flatten(2).contiguous() # b c hw\n x_hil = x_hw.index_select(dim=-1, index=hilbert_curve)\n xs = torch.stack([x_hil, torch.flip(x_hil, dims=[-1])], dim=1) # (b, k, d, l)\n else:\n raise ValueError()\n\n x_dbl = torch.einsum(\"b k d l, k c d -> b k c l\", xs.view(B, K, -1, L), self.x_proj_weight)\n # x_dbl = x_dbl + self.x_proj_bias.view(1, K, -1, 1)\n dts, Bs, Cs = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=2)\n dts = torch.einsum(\"b k r l, k d r -> b k d l\", dts.view(B, K, -1, L), self.dt_projs_weight)\n # dts = dts + self.dt_projs_bias.view(1, K, -1, 1)\n\n xs = xs.float().view(B, -1, L) # (b, k * d, l)\n dts = dts.contiguous().float().view(B, -1, L) # (b, k * d, l)\n Bs = Bs.float().view(B, K, -1, L) # (b, k, d_state, l)\n Cs = Cs.float().view(B, K, -1, L) # (b, k, d_state, l)\n Ds = self.Ds.float().view(-1) # (k * d)\n As = -torch.exp(self.A_logs.float()).view(-1, self.d_state) # (k * d, d_state)\n dt_projs_bias = self.dt_projs_bias.float().view(-1) # (k * d)\n\n out_y = self.selective_scan(\n xs, dts, \n As, Bs, Cs, Ds, z=None,\n delta_bias=dt_projs_bias,\n delta_softplus=True,\n return_last_state=False,\n ).view(B, K, -1, L)\n\n assert out_y.dtype == torch.float\n\n if self.scan_order == 'zigzag':\n hw_order = out_y[:, 0].contiguous().view(B, -1, H, W).contiguous()\n rhw_order = torch.flip(out_y[:, 1].contiguous(), dims=[-1]).contiguous()\n rhw_order = rhw_order.view(B, -1, H, W,).contiguous()\n return hw_order + rhw_order\n \n elif self.scan_order == 'hilbert':\n hil_order = out_y[:, 0].contiguous() # b c hw\n rhil_order = torch.flip(out_y[:, 1].contiguous(), dims=[-1]).contiguous() # b c hw\n\n sum_out = torch.zeros_like(hil_order)\n hilbert_curve = repeat(hilbert_curve, 'hw -> b c hw', b=hil_order.shape[0], c=hil_order.shape[1])\n assert hil_order.shape == hilbert_curve.shape\n sum_out.scatter_add_(dim=-1, index=hilbert_curve, src=hil_order)\n sum_out.scatter_add_(dim=-1, index=hilbert_curve, src=rhil_order)\n sum_out = sum_out.view(B, -1, H, W).contiguous()\n return sum_out\n\n\n # def forward_corev0(self, x: torch.Tensor, hilbert_curve):\n # # LongTensor[int] 按照hw进行flatten之后的hilbert排序\n # self.selective_scan = selective_scan_fn\n \n # B, C, H, W, T = x.shape\n # L = H * W * T\n # K = 2\n\n # if self.scan_order == 'zigzag':\n # x_hw = x.view(B, -1, L).contiguous() # b c hwt\n # xs = torch.stack([x_hw, torch.flip(x_hw, dims=[-1])], dim=1) # (b, k, d, l)\n # elif self.scan_order == 'hilbert':\n # x_hw = x.flatten(2).contiguous() # b c hwt\n # x_hil = x_hw.index_select(dim=-1, index=hilbert_curve)\n # xs = torch.stack([x_hil, torch.flip(x_hil, dims=[-1])], dim=1) # (b, k, d, l)\n # else:\n # raise ValueError()\n\n # x_dbl = torch.einsum(\"b k d l, k c d -> b k c l\", xs.view(B, K, -1, L), self.x_proj_weight)\n # # x_dbl = x_dbl + self.x_proj_bias.view(1, K, -1, 1)\n # dts, Bs, Cs = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=2)\n # dts = torch.einsum(\"b k r l, k d r -> b k d l\", dts.view(B, K, -1, L), self.dt_projs_weight)\n # # dts = dts + self.dt_projs_bias.view(1, K, -1, 1)\n\n # xs = xs.float().view(B, -1, L) # (b, k * d, l)\n # dts = dts.contiguous().float().view(B, -1, L) # (b, k * d, l)\n # Bs = Bs.float().view(B, K, -1, L) # (b, k, d_state, l)\n # Cs = Cs.float().view(B, K, -1, L) # (b, k, d_state, l)\n # Ds = self.Ds.float().view(-1) # (k * d)\n # As = -torch.exp(self.A_logs.float()).view(-1, self.d_state) # (k * d, d_state)\n # dt_projs_bias = self.dt_projs_bias.float().view(-1) # (k * d)\n\n # out_y = self.selective_scan(\n # xs, dts, \n # As, Bs, Cs, Ds, z=None,\n # delta_bias=dt_projs_bias,\n # delta_softplus=True,\n # return_last_state=False,\n # ).view(B, K, -1, L)\n\n # assert out_y.dtype == torch.float\n\n # if self.scan_order == 'zigzag':\n # hw_order = out_y[:, 0].contiguous().view(B, -1, H, W).contiguous()\n # rhw_order = torch.flip(out_y[:, 1].contiguous(), dims=[-1]).contiguous()\n # rhw_order = rhw_order.view(B, -1, H, W,).contiguous()\n # return hw_order + rhw_order\n \n # elif self.scan_order == 'hilbert':\n # hil_order = out_y[:, 0].contiguous() # b c hw\n # rhil_order = torch.flip(out_y[:, 1].contiguous(), dims=[-1]).contiguous() # b c hw\n\n # sum_out = torch.zeros_like(hil_order)\n # hilbert_curve = repeat(hilbert_curve, 'hwt -> b c hwt', b=hil_order.shape[0], c=hil_order.shape[1])\n # assert hil_order.shape == hilbert_curve.shape\n # sum_out.scatter_add_(dim=-1, index=hilbert_curve, src=hil_order)\n # sum_out.scatter_add_(dim=-1, index=hilbert_curve, src=rhil_order)\n # sum_out = sum_out.view(B, -1, H, W).contiguous()\n # return sum_out\n\n\n def forward(self, x: torch.Tensor, hilbert_curve, **kwargs):\n\n B, H, W, C = x.shape\n xz = self.in_proj(x)\n x, z = xz.chunk(2, dim=-1) # (b, h, w, d)\n\n x = x.permute(0, 3, 1, 2).contiguous()\n x = self.act(self.conv2d(x)) # (b, d, h, w)\n \n y = self.forward_core(x, hilbert_curve=hilbert_curve) # B C h w\n y = y.permute(0, 2, 3, 1).contiguous() # b h w c\n y = self.out_norm(y)\n y = y * F.silu(z)\n out = self.out_proj(y)\n if self.dropout is not None:\n out = self.dropout(out)\n return out\n\n\nclass SS2D_FrameQuery_hilbert(nn.Module):\n def __init__(self, configs,):\n super().__init__()\n d_model = configs['d_model']\n self.homo = Hilbert_2DSelectiveScan(d_model=configs['d_model'],\n d_state=configs['d_state'] if 'd_state' in configs else 16,\n d_conv=configs['d_conv'] if 'd_conv' in configs else 3,\n expand=configs['expand'] if 'expand' in configs else 2,\n dt_rank=configs['dt_rank'] if 'dt_rank' in configs else 'auto',\n dt_min=configs['dt_min'] if 'dt_min' in configs else 0.001,\n dt_max=configs['dt_max'] if 'dt_max' in configs else 0.1,\n dt_init=configs['dt_init'] if 'dt_init' in configs else 'random',\n dt_scale=configs['dt_scale'] if 'dt_scale' in configs else 1.0,\n dt_init_floor=configs['dt_init_floor'] if 'dt_init_floor' in configs else 1e-4,\n dropout=configs['dropout'] if 'dropout' in configs else 0,\n conv_bias=configs['conv_bias'] if 'conv_bias' in configs else True,\n bias=configs['bias'] if 'bias' in configs else False,\n scan_order=configs['scan_order']\n )\n self.pos_1d = build_position_encoding(position_embedding_name='1d') # t上的position embedding\n self.norm = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(configs['dropout'])\n\n def forward(self, \n frame_query_feats=None, # n bt c \n frame_query_poses=None, # n bt c # nq上的Position embedding\n hilbert_curve=None,\n nf=None,\n **kwargs\n ):\n batch_size = frame_query_feats.shape[1] // nf # b\n tgt2 = frame_query_feats + frame_query_poses \n tgt2 = rearrange(tgt2, 'n (b t) c -> b t n c',b=batch_size,t=nf).contiguous()\n\n sin_poses = self.pos_1d(torch.zeros_like(tgt2[..., 0].permute(0, 2, 1).flatten(0, 1)).bool(), \n hidden_dim=tgt2.shape[-1]) # bn c t\n sin_poses = rearrange(sin_poses, '(b n) c t -> b t n c', b=batch_size)\n tgt2 += sin_poses\n\n tgt2 = self.homo(tgt2, hilbert_curve=hilbert_curve) # b t n c\n \n tgt2 = tgt2.permute(2, 0, 1, 3).flatten(1, 2).contiguous() # n bt c\n\n\n frame_query_feats = frame_query_feats + self.dropout(tgt2)\n frame_query_feats = self.norm(frame_query_feats)\n\n return frame_query_feats, None\n \nfrom models.layers.utils import _get_clones\n\n@META_ARCH_REGISTRY.register()\nclass FrameQuery_SS2DLayer_hilbert(nn.Module):\n def __init__(self, \n configs, \n dropout=0.0):\n super().__init__()\n d_model = configs['d_model']\n n_layers = configs['nlayers'] if 'nlayers' in configs else 1\n self.nlayers = n_layers\n self.self_attn = _get_clones(SS2D_FrameQuery_hilbert(configs), n_layers)\n from models.layers.decoder_layers import FFNLayer\n self.ffn = FFNLayer(d_model=d_model,\n dim_feedforward=configs['dim_feedforward'],\n dropout=configs['dropout'],)\n\n def forward(self, \n frame_query_feats, # n bt c \n frame_query_poses, # n bt c # nq上的Position embedding\n video_aux_dict=None,\n **kwargs):\n\n for i in range(self.nlayers):\n frame_query_feats = self.self_attn[i](frame_query_feats=frame_query_feats, # n bt c \n frame_query_poses=frame_query_poses, # n bt c # nq上的Position embedding\n nf=video_aux_dict['nf'],\n hilbert_curve=video_aux_dict['hilbert_curve'])[0]\n frame_query_feats = self.ffn(frame_query_feats)\n return frame_query_feats\n\n\n\n"
},
{
"path": "models/encoder/ops/modules/ms_deform_attn.py",
"content": "# Modify for sample points visualization\n# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport warnings\nimport math\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.nn.init import xavier_uniform_, constant_\n\nfrom ..functions import MSDeformAttnFunction\n\n\ndef _is_power_of_2(n):\n if (not isinstance(n, int)) or (n < 0):\n raise ValueError(\"invalid input for _is_power_of_2: {} (type: {})\".format(n, type(n)))\n return (n & (n-1) == 0) and n != 0\n\n\nclass MSDeformAttn(nn.Module):\n def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):\n \"\"\"\n Multi-Scale Deformable Attention Module\n :param d_model hidden dimension\n :param n_levels number of feature levels\n :param n_heads number of attention heads\n :param n_points number of sampling points per attention head per feature level\n \"\"\"\n super().__init__()\n if d_model % n_heads != 0:\n raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))\n _d_per_head = d_model // n_heads\n # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation\n if not _is_power_of_2(_d_per_head):\n warnings.warn(\"You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 \"\n \"which is more efficient in our CUDA implementation.\")\n\n self.im2col_step = 512\n\n self.d_model = d_model\n self.n_levels = n_levels\n self.n_heads = n_heads\n self.n_points = n_points\n\n self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)\n self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)\n self.value_proj = nn.Linear(d_model, d_model)\n self.output_proj = nn.Linear(d_model, d_model)\n\n self._reset_parameters()\n\n def _reset_parameters(self):\n constant_(self.sampling_offsets.weight.data, 0.)\n thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)\n grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)\n grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)\n for i in range(self.n_points):\n grid_init[:, :, i, :] *= i + 1\n with torch.no_grad():\n self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))\n constant_(self.attention_weights.weight.data, 0.)\n constant_(self.attention_weights.bias.data, 0.)\n xavier_uniform_(self.value_proj.weight.data)\n constant_(self.value_proj.bias.data, 0.)\n xavier_uniform_(self.output_proj.weight.data)\n constant_(self.output_proj.bias.data, 0.)\n\n def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None):\n \"\"\"\n :param query (N, Length_{query}, C)\n :param reference_points (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area\n or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes\n :param input_flatten (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l, C)\n :param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]\n :param input_level_start_index (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]\n :param input_padding_mask (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l), True for padding elements, False for non-padding elements\n\n :return output (N, Length_{query}, C)\n \"\"\"\n N, Len_q, _ = query.shape\n N, Len_in, _ = input_flatten.shape\n assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in\n\n value = self.value_proj(input_flatten)\n if input_padding_mask is not None:\n value = value.masked_fill(input_padding_mask[..., None], float(0))\n value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)\n sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)\n attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)\n attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)\n # N, Len_q, n_heads, n_levels, n_points, 2\n if reference_points.shape[-1] == 2:\n offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)\n sampling_locations = reference_points[:, :, None, :, None, :] \\\n + sampling_offsets / offset_normalizer[None, None, None, :, None, :]\n elif reference_points.shape[-1] == 4:\n sampling_locations = reference_points[:, :, None, :, None, :2] \\\n + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5\n else:\n raise ValueError(\n 'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))\n output = MSDeformAttnFunction.apply(\n value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)\n output = self.output_proj(output)\n\n return output, sampling_locations, attention_weights\n"
},
{
"path": "models/encoder/ops/setup.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nimport os\nimport glob\n\nimport torch\n\nfrom torch.utils.cpp_extension import CUDA_HOME\nfrom torch.utils.cpp_extension import CppExtension\nfrom torch.utils.cpp_extension import CUDAExtension\n\nfrom setuptools import find_packages\nfrom setuptools import setup\n\nrequirements = [\"torch\", \"torchvision\"]\n\ndef get_extensions():\n this_dir = os.path.dirname(os.path.abspath(__file__))\n extensions_dir = os.path.join(this_dir, \"src\")\n\n main_file = glob.glob(os.path.join(extensions_dir, \"*.cpp\"))\n source_cpu = glob.glob(os.path.join(extensions_dir, \"cpu\", \"*.cpp\"))\n source_cuda = glob.glob(os.path.join(extensions_dir, \"cuda\", \"*.cu\"))\n\n sources = main_file + source_cpu\n extension = CppExtension\n extra_compile_args = {\"cxx\": []}\n define_macros = []\n\n if torch.cuda.is_available() and CUDA_HOME is not None:\n extension = CUDAExtension\n sources += source_cuda\n define_macros += [(\"WITH_CUDA\", None)]\n extra_compile_args[\"nvcc\"] = [\n \"-DCUDA_HAS_FP16=1\",\n \"-D__CUDA_NO_HALF_OPERATORS__\",\n \"-D__CUDA_NO_HALF_CONVERSIONS__\",\n \"-D__CUDA_NO_HALF2_OPERATORS__\",\n ]\n else:\n raise NotImplementedError('Cuda is not availabel')\n\n sources = [os.path.join(extensions_dir, s) for s in sources]\n include_dirs = [extensions_dir]\n ext_modules = [\n extension(\n \"MultiScaleDeformableAttention\",\n sources,\n include_dirs=include_dirs,\n define_macros=define_macros,\n extra_compile_args=extra_compile_args,\n )\n ]\n return ext_modules\n\nsetup(\n name=\"MultiScaleDeformableAttention\",\n version=\"1.0\",\n author=\"Weijie Su\",\n url=\"https://github.com/fundamentalvision/Deformable-DETR\",\n description=\"PyTorch Wrapper for CUDA Functions of Multi-Scale Deformable Attention\",\n packages=find_packages(exclude=(\"configs\", \"tests\",)),\n ext_modules=get_extensions(),\n cmdclass={\"build_ext\": torch.utils.cpp_extension.BuildExtension},\n)\n"
},
{
"path": "models/encoder/ops/src/cpu/ms_deform_attn_cpu.cpp",
"content": "/*!\n**************************************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************************************\n* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n**************************************************************************************************\n*/\n\n#include \n\n#include \n#include \n\n\nat::Tensor\nms_deform_attn_cpu_forward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const int im2col_step)\n{\n AT_ERROR(\"Not implement on cpu\");\n}\n\nstd::vector\nms_deform_attn_cpu_backward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const at::Tensor &grad_output,\n const int im2col_step)\n{\n AT_ERROR(\"Not implement on cpu\");\n}\n\n"
},
{
"path": "models/encoder/ops/src/cpu/ms_deform_attn_cpu.h",
"content": "/*!\n**************************************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************************************\n* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n**************************************************************************************************\n*/\n\n#pragma once\n#include \n\nat::Tensor\nms_deform_attn_cpu_forward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const int im2col_step);\n\nstd::vector\nms_deform_attn_cpu_backward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const at::Tensor &grad_output,\n const int im2col_step);\n\n\n"
},
{
"path": "models/encoder/ops/src/cuda/ms_deform_attn_cuda.cu",
"content": "/*!\n**************************************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************************************\n* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n**************************************************************************************************\n*/\n\n#include \n#include \"cuda/ms_deform_im2col_cuda.cuh\"\n\n#include \n#include \n#include \n#include \n\n\nat::Tensor ms_deform_attn_cuda_forward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const int im2col_step)\n{\n AT_ASSERTM(value.is_contiguous(), \"value tensor has to be contiguous\");\n AT_ASSERTM(spatial_shapes.is_contiguous(), \"spatial_shapes tensor has to be contiguous\");\n AT_ASSERTM(level_start_index.is_contiguous(), \"level_start_index tensor has to be contiguous\");\n AT_ASSERTM(sampling_loc.is_contiguous(), \"sampling_loc tensor has to be contiguous\");\n AT_ASSERTM(attn_weight.is_contiguous(), \"attn_weight tensor has to be contiguous\");\n\n AT_ASSERTM(value.type().is_cuda(), \"value must be a CUDA tensor\");\n AT_ASSERTM(spatial_shapes.type().is_cuda(), \"spatial_shapes must be a CUDA tensor\");\n AT_ASSERTM(level_start_index.type().is_cuda(), \"level_start_index must be a CUDA tensor\");\n AT_ASSERTM(sampling_loc.type().is_cuda(), \"sampling_loc must be a CUDA tensor\");\n AT_ASSERTM(attn_weight.type().is_cuda(), \"attn_weight must be a CUDA tensor\");\n\n const int batch = value.size(0);\n const int spatial_size = value.size(1);\n const int num_heads = value.size(2);\n const int channels = value.size(3);\n\n const int num_levels = spatial_shapes.size(0);\n\n const int num_query = sampling_loc.size(1);\n const int num_point = sampling_loc.size(4);\n\n const int im2col_step_ = std::min(batch, im2col_step);\n\n AT_ASSERTM(batch % im2col_step_ == 0, \"batch(%d) must divide im2col_step(%d)\", batch, im2col_step_);\n \n auto output = at::zeros({batch, num_query, num_heads, channels}, value.options());\n\n const int batch_n = im2col_step_;\n auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels});\n auto per_value_size = spatial_size * num_heads * channels;\n auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;\n auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;\n for (int n = 0; n < batch/im2col_step_; ++n)\n {\n auto columns = output_n.select(0, n);\n AT_DISPATCH_FLOATING_TYPES(value.type(), \"ms_deform_attn_forward_cuda\", ([&] {\n ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(),\n value.data() + n * im2col_step_ * per_value_size,\n spatial_shapes.data(),\n level_start_index.data(),\n sampling_loc.data() + n * im2col_step_ * per_sample_loc_size,\n attn_weight.data() + n * im2col_step_ * per_attn_weight_size,\n batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,\n columns.data());\n\n }));\n }\n\n output = output.view({batch, num_query, num_heads*channels});\n\n return output;\n}\n\n\nstd::vector ms_deform_attn_cuda_backward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const at::Tensor &grad_output,\n const int im2col_step)\n{\n\n AT_ASSERTM(value.is_contiguous(), \"value tensor has to be contiguous\");\n AT_ASSERTM(spatial_shapes.is_contiguous(), \"spatial_shapes tensor has to be contiguous\");\n AT_ASSERTM(level_start_index.is_contiguous(), \"level_start_index tensor has to be contiguous\");\n AT_ASSERTM(sampling_loc.is_contiguous(), \"sampling_loc tensor has to be contiguous\");\n AT_ASSERTM(attn_weight.is_contiguous(), \"attn_weight tensor has to be contiguous\");\n AT_ASSERTM(grad_output.is_contiguous(), \"grad_output tensor has to be contiguous\");\n\n AT_ASSERTM(value.type().is_cuda(), \"value must be a CUDA tensor\");\n AT_ASSERTM(spatial_shapes.type().is_cuda(), \"spatial_shapes must be a CUDA tensor\");\n AT_ASSERTM(level_start_index.type().is_cuda(), \"level_start_index must be a CUDA tensor\");\n AT_ASSERTM(sampling_loc.type().is_cuda(), \"sampling_loc must be a CUDA tensor\");\n AT_ASSERTM(attn_weight.type().is_cuda(), \"attn_weight must be a CUDA tensor\");\n AT_ASSERTM(grad_output.type().is_cuda(), \"grad_output must be a CUDA tensor\");\n\n const int batch = value.size(0);\n const int spatial_size = value.size(1);\n const int num_heads = value.size(2);\n const int channels = value.size(3);\n\n const int num_levels = spatial_shapes.size(0);\n\n const int num_query = sampling_loc.size(1);\n const int num_point = sampling_loc.size(4);\n\n const int im2col_step_ = std::min(batch, im2col_step);\n\n AT_ASSERTM(batch % im2col_step_ == 0, \"batch(%d) must divide im2col_step(%d)\", batch, im2col_step_);\n\n auto grad_value = at::zeros_like(value);\n auto grad_sampling_loc = at::zeros_like(sampling_loc);\n auto grad_attn_weight = at::zeros_like(attn_weight);\n\n const int batch_n = im2col_step_;\n auto per_value_size = spatial_size * num_heads * channels;\n auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;\n auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;\n auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels});\n \n for (int n = 0; n < batch/im2col_step_; ++n)\n {\n auto grad_output_g = grad_output_n.select(0, n);\n AT_DISPATCH_FLOATING_TYPES(value.type(), \"ms_deform_attn_backward_cuda\", ([&] {\n ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(),\n grad_output_g.data(),\n value.data() + n * im2col_step_ * per_value_size,\n spatial_shapes.data(),\n level_start_index.data(),\n sampling_loc.data() + n * im2col_step_ * per_sample_loc_size,\n attn_weight.data() + n * im2col_step_ * per_attn_weight_size,\n batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,\n grad_value.data() + n * im2col_step_ * per_value_size,\n grad_sampling_loc.data() + n * im2col_step_ * per_sample_loc_size,\n grad_attn_weight.data() + n * im2col_step_ * per_attn_weight_size);\n\n }));\n }\n\n return {\n grad_value, grad_sampling_loc, grad_attn_weight\n };\n}"
},
{
"path": "models/encoder/ops/src/cuda/ms_deform_attn_cuda.h",
"content": "/*!\n**************************************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************************************\n* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n**************************************************************************************************\n*/\n\n#pragma once\n#include \n\nat::Tensor ms_deform_attn_cuda_forward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const int im2col_step);\n\nstd::vector ms_deform_attn_cuda_backward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const at::Tensor &grad_output,\n const int im2col_step);\n\n"
},
{
"path": "models/encoder/ops/src/cuda/ms_deform_im2col_cuda.cuh",
"content": "/*!\n**************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************\n* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)\n* Copyright (c) 2018 Microsoft\n**************************************************************************\n*/\n\n#include \n#include \n#include \n\n#include \n#include \n\n#include \n\n#define CUDA_KERNEL_LOOP(i, n) \\\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; \\\n i < (n); \\\n i += blockDim.x * gridDim.x)\n\nconst int CUDA_NUM_THREADS = 1024;\ninline int GET_BLOCKS(const int N, const int num_threads)\n{\n return (N + num_threads - 1) / num_threads;\n}\n\n\ntemplate \n__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, \n const int &height, const int &width, const int &nheads, const int &channels,\n const scalar_t &h, const scalar_t &w, const int &m, const int &c)\n{\n const int h_low = floor(h);\n const int w_low = floor(w);\n const int h_high = h_low + 1;\n const int w_high = w_low + 1;\n\n const scalar_t lh = h - h_low;\n const scalar_t lw = w - w_low;\n const scalar_t hh = 1 - lh, hw = 1 - lw;\n\n const int w_stride = nheads * channels;\n const int h_stride = width * w_stride;\n const int h_low_ptr_offset = h_low * h_stride;\n const int h_high_ptr_offset = h_low_ptr_offset + h_stride;\n const int w_low_ptr_offset = w_low * w_stride;\n const int w_high_ptr_offset = w_low_ptr_offset + w_stride;\n const int base_ptr = m * channels + c;\n\n scalar_t v1 = 0;\n if (h_low >= 0 && w_low >= 0)\n {\n const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;\n v1 = bottom_data[ptr1];\n }\n scalar_t v2 = 0;\n if (h_low >= 0 && w_high <= width - 1)\n {\n const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;\n v2 = bottom_data[ptr2];\n }\n scalar_t v3 = 0;\n if (h_high <= height - 1 && w_low >= 0)\n {\n const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;\n v3 = bottom_data[ptr3];\n }\n scalar_t v4 = 0;\n if (h_high <= height - 1 && w_high <= width - 1)\n {\n const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;\n v4 = bottom_data[ptr4];\n }\n\n const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;\n\n const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);\n return val;\n}\n\n\ntemplate \n__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, \n const int &height, const int &width, const int &nheads, const int &channels,\n const scalar_t &h, const scalar_t &w, const int &m, const int &c,\n const scalar_t &top_grad,\n const scalar_t &attn_weight,\n scalar_t* &grad_value, \n scalar_t* grad_sampling_loc,\n scalar_t* grad_attn_weight)\n{\n const int h_low = floor(h);\n const int w_low = floor(w);\n const int h_high = h_low + 1;\n const int w_high = w_low + 1;\n\n const scalar_t lh = h - h_low;\n const scalar_t lw = w - w_low;\n const scalar_t hh = 1 - lh, hw = 1 - lw;\n\n const int w_stride = nheads * channels;\n const int h_stride = width * w_stride;\n const int h_low_ptr_offset = h_low * h_stride;\n const int h_high_ptr_offset = h_low_ptr_offset + h_stride;\n const int w_low_ptr_offset = w_low * w_stride;\n const int w_high_ptr_offset = w_low_ptr_offset + w_stride;\n const int base_ptr = m * channels + c;\n\n const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;\n const scalar_t top_grad_value = top_grad * attn_weight;\n scalar_t grad_h_weight = 0, grad_w_weight = 0;\n\n scalar_t v1 = 0;\n if (h_low >= 0 && w_low >= 0)\n {\n const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;\n v1 = bottom_data[ptr1];\n grad_h_weight -= hw * v1;\n grad_w_weight -= hh * v1;\n atomicAdd(grad_value+ptr1, w1*top_grad_value);\n }\n scalar_t v2 = 0;\n if (h_low >= 0 && w_high <= width - 1)\n {\n const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;\n v2 = bottom_data[ptr2];\n grad_h_weight -= lw * v2;\n grad_w_weight += hh * v2;\n atomicAdd(grad_value+ptr2, w2*top_grad_value);\n }\n scalar_t v3 = 0;\n if (h_high <= height - 1 && w_low >= 0)\n {\n const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;\n v3 = bottom_data[ptr3];\n grad_h_weight += hw * v3;\n grad_w_weight -= lh * v3;\n atomicAdd(grad_value+ptr3, w3*top_grad_value); \n }\n scalar_t v4 = 0;\n if (h_high <= height - 1 && w_high <= width - 1)\n {\n const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;\n v4 = bottom_data[ptr4];\n grad_h_weight += lw * v4;\n grad_w_weight += lh * v4;\n atomicAdd(grad_value+ptr4, w4*top_grad_value);\n }\n\n const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);\n *grad_attn_weight = top_grad * val;\n *grad_sampling_loc = width * grad_w_weight * top_grad_value;\n *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value;\n}\n\n\ntemplate \n__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, \n const int &height, const int &width, const int &nheads, const int &channels,\n const scalar_t &h, const scalar_t &w, const int &m, const int &c,\n const scalar_t &top_grad,\n const scalar_t &attn_weight,\n scalar_t* &grad_value, \n scalar_t* grad_sampling_loc,\n scalar_t* grad_attn_weight)\n{\n const int h_low = floor(h);\n const int w_low = floor(w);\n const int h_high = h_low + 1;\n const int w_high = w_low + 1;\n\n const scalar_t lh = h - h_low;\n const scalar_t lw = w - w_low;\n const scalar_t hh = 1 - lh, hw = 1 - lw;\n\n const int w_stride = nheads * channels;\n const int h_stride = width * w_stride;\n const int h_low_ptr_offset = h_low * h_stride;\n const int h_high_ptr_offset = h_low_ptr_offset + h_stride;\n const int w_low_ptr_offset = w_low * w_stride;\n const int w_high_ptr_offset = w_low_ptr_offset + w_stride;\n const int base_ptr = m * channels + c;\n\n const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;\n const scalar_t top_grad_value = top_grad * attn_weight;\n scalar_t grad_h_weight = 0, grad_w_weight = 0;\n\n scalar_t v1 = 0;\n if (h_low >= 0 && w_low >= 0)\n {\n const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;\n v1 = bottom_data[ptr1];\n grad_h_weight -= hw * v1;\n grad_w_weight -= hh * v1;\n atomicAdd(grad_value+ptr1, w1*top_grad_value);\n }\n scalar_t v2 = 0;\n if (h_low >= 0 && w_high <= width - 1)\n {\n const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;\n v2 = bottom_data[ptr2];\n grad_h_weight -= lw * v2;\n grad_w_weight += hh * v2;\n atomicAdd(grad_value+ptr2, w2*top_grad_value);\n }\n scalar_t v3 = 0;\n if (h_high <= height - 1 && w_low >= 0)\n {\n const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;\n v3 = bottom_data[ptr3];\n grad_h_weight += hw * v3;\n grad_w_weight -= lh * v3;\n atomicAdd(grad_value+ptr3, w3*top_grad_value); \n }\n scalar_t v4 = 0;\n if (h_high <= height - 1 && w_high <= width - 1)\n {\n const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;\n v4 = bottom_data[ptr4];\n grad_h_weight += lw * v4;\n grad_w_weight += lh * v4;\n atomicAdd(grad_value+ptr4, w4*top_grad_value);\n }\n\n const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);\n atomicAdd(grad_attn_weight, top_grad * val); \n atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value);\n atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value);\n}\n\n\ntemplate \n__global__ void ms_deformable_im2col_gpu_kernel(const int n,\n const scalar_t *data_value, \n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *data_col)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n scalar_t *data_col_ptr = data_col + index;\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n scalar_t col = 0;\n \n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight;\n }\n\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n }\n }\n *data_col_ptr = col;\n }\n}\n\ntemplate \n__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n,\n const scalar_t *grad_col,\n const scalar_t *data_value,\n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *grad_value,\n scalar_t *grad_sampling_loc,\n scalar_t *grad_attn_weight)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];\n __shared__ scalar_t cache_grad_attn_weight[blockSize];\n unsigned int tid = threadIdx.x;\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n const scalar_t top_grad = grad_col[index];\n\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int grad_sampling_ptr = data_weight_ptr;\n grad_sampling_loc += grad_sampling_ptr << 1;\n grad_attn_weight += grad_sampling_ptr;\n const int grad_weight_stride = 1;\n const int grad_loc_stride = 2;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n\n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;\n const scalar_t *data_value_ptr = data_value + value_ptr_offset;\n scalar_t *grad_value_ptr = grad_value + value_ptr_offset;\n\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;\n *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;\n *(cache_grad_attn_weight+threadIdx.x)=0;\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n ms_deform_attn_col2im_bilinear(\n data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,\n top_grad, weight, grad_value_ptr, \n cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);\n }\n \n __syncthreads();\n if (tid == 0)\n {\n scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0];\n int sid=2;\n for (unsigned int tid = 1; tid < blockSize; ++tid)\n {\n _grad_w += cache_grad_sampling_loc[sid];\n _grad_h += cache_grad_sampling_loc[sid + 1];\n _grad_a += cache_grad_attn_weight[tid];\n sid += 2;\n }\n \n \n *grad_sampling_loc = _grad_w;\n *(grad_sampling_loc + 1) = _grad_h;\n *grad_attn_weight = _grad_a;\n }\n __syncthreads();\n\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n grad_attn_weight += grad_weight_stride;\n grad_sampling_loc += grad_loc_stride;\n }\n }\n }\n}\n\n\ntemplate \n__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n,\n const scalar_t *grad_col,\n const scalar_t *data_value,\n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *grad_value,\n scalar_t *grad_sampling_loc,\n scalar_t *grad_attn_weight)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];\n __shared__ scalar_t cache_grad_attn_weight[blockSize];\n unsigned int tid = threadIdx.x;\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n const scalar_t top_grad = grad_col[index];\n\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int grad_sampling_ptr = data_weight_ptr;\n grad_sampling_loc += grad_sampling_ptr << 1;\n grad_attn_weight += grad_sampling_ptr;\n const int grad_weight_stride = 1;\n const int grad_loc_stride = 2;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n\n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;\n const scalar_t *data_value_ptr = data_value + value_ptr_offset;\n scalar_t *grad_value_ptr = grad_value + value_ptr_offset;\n\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;\n *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;\n *(cache_grad_attn_weight+threadIdx.x)=0;\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n ms_deform_attn_col2im_bilinear(\n data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,\n top_grad, weight, grad_value_ptr, \n cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);\n }\n \n __syncthreads();\n\n for (unsigned int s=blockSize/2; s>0; s>>=1)\n {\n if (tid < s) {\n const unsigned int xid1 = tid << 1;\n const unsigned int xid2 = (tid + s) << 1;\n cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];\n cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];\n cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];\n }\n __syncthreads();\n }\n\n if (tid == 0)\n { \n *grad_sampling_loc = cache_grad_sampling_loc[0];\n *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];\n *grad_attn_weight = cache_grad_attn_weight[0];\n }\n __syncthreads();\n\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n grad_attn_weight += grad_weight_stride;\n grad_sampling_loc += grad_loc_stride;\n }\n }\n }\n}\n\n\ntemplate \n__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n,\n const scalar_t *grad_col,\n const scalar_t *data_value,\n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *grad_value,\n scalar_t *grad_sampling_loc,\n scalar_t *grad_attn_weight)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n extern __shared__ int _s[];\n scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;\n scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;\n unsigned int tid = threadIdx.x;\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n const scalar_t top_grad = grad_col[index];\n\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int grad_sampling_ptr = data_weight_ptr;\n grad_sampling_loc += grad_sampling_ptr << 1;\n grad_attn_weight += grad_sampling_ptr;\n const int grad_weight_stride = 1;\n const int grad_loc_stride = 2;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n\n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;\n const scalar_t *data_value_ptr = data_value + value_ptr_offset;\n scalar_t *grad_value_ptr = grad_value + value_ptr_offset;\n\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;\n *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;\n *(cache_grad_attn_weight+threadIdx.x)=0;\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n ms_deform_attn_col2im_bilinear(\n data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,\n top_grad, weight, grad_value_ptr, \n cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);\n }\n \n __syncthreads();\n if (tid == 0)\n {\n scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0];\n int sid=2;\n for (unsigned int tid = 1; tid < blockDim.x; ++tid)\n {\n _grad_w += cache_grad_sampling_loc[sid];\n _grad_h += cache_grad_sampling_loc[sid + 1];\n _grad_a += cache_grad_attn_weight[tid];\n sid += 2;\n }\n \n \n *grad_sampling_loc = _grad_w;\n *(grad_sampling_loc + 1) = _grad_h;\n *grad_attn_weight = _grad_a;\n }\n __syncthreads();\n\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n grad_attn_weight += grad_weight_stride;\n grad_sampling_loc += grad_loc_stride;\n }\n }\n }\n}\n\ntemplate \n__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n,\n const scalar_t *grad_col,\n const scalar_t *data_value,\n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *grad_value,\n scalar_t *grad_sampling_loc,\n scalar_t *grad_attn_weight)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n extern __shared__ int _s[];\n scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;\n scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;\n unsigned int tid = threadIdx.x;\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n const scalar_t top_grad = grad_col[index];\n\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int grad_sampling_ptr = data_weight_ptr;\n grad_sampling_loc += grad_sampling_ptr << 1;\n grad_attn_weight += grad_sampling_ptr;\n const int grad_weight_stride = 1;\n const int grad_loc_stride = 2;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n\n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;\n const scalar_t *data_value_ptr = data_value + value_ptr_offset;\n scalar_t *grad_value_ptr = grad_value + value_ptr_offset;\n\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;\n *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;\n *(cache_grad_attn_weight+threadIdx.x)=0;\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n ms_deform_attn_col2im_bilinear(\n data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,\n top_grad, weight, grad_value_ptr, \n cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);\n }\n \n __syncthreads();\n\n for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1)\n {\n if (tid < s) {\n const unsigned int xid1 = tid << 1;\n const unsigned int xid2 = (tid + s) << 1;\n cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];\n cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];\n cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];\n if (tid + (s << 1) < spre)\n {\n cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];\n cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];\n cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];\n } \n }\n __syncthreads();\n }\n\n if (tid == 0)\n {\n *grad_sampling_loc = cache_grad_sampling_loc[0];\n *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];\n *grad_attn_weight = cache_grad_attn_weight[0];\n }\n __syncthreads();\n\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n grad_attn_weight += grad_weight_stride;\n grad_sampling_loc += grad_loc_stride;\n }\n }\n }\n}\n\ntemplate \n__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n,\n const scalar_t *grad_col,\n const scalar_t *data_value,\n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *grad_value,\n scalar_t *grad_sampling_loc,\n scalar_t *grad_attn_weight)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n extern __shared__ int _s[];\n scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;\n scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;\n unsigned int tid = threadIdx.x;\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n const scalar_t top_grad = grad_col[index];\n\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int grad_sampling_ptr = data_weight_ptr;\n grad_sampling_loc += grad_sampling_ptr << 1;\n grad_attn_weight += grad_sampling_ptr;\n const int grad_weight_stride = 1;\n const int grad_loc_stride = 2;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n\n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;\n const scalar_t *data_value_ptr = data_value + value_ptr_offset;\n scalar_t *grad_value_ptr = grad_value + value_ptr_offset;\n\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;\n *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;\n *(cache_grad_attn_weight+threadIdx.x)=0;\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n ms_deform_attn_col2im_bilinear(\n data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,\n top_grad, weight, grad_value_ptr, \n cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);\n }\n \n __syncthreads();\n\n for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1)\n {\n if (tid < s) {\n const unsigned int xid1 = tid << 1;\n const unsigned int xid2 = (tid + s) << 1;\n cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];\n cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];\n cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];\n if (tid + (s << 1) < spre)\n {\n cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];\n cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];\n cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];\n }\n }\n __syncthreads();\n }\n\n if (tid == 0)\n {\n atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]);\n atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]);\n atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]);\n }\n __syncthreads();\n\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n grad_attn_weight += grad_weight_stride;\n grad_sampling_loc += grad_loc_stride;\n }\n }\n }\n}\n\n\ntemplate \n__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n,\n const scalar_t *grad_col,\n const scalar_t *data_value,\n const int64_t *data_spatial_shapes,\n const int64_t *data_level_start_index, \n const scalar_t *data_sampling_loc,\n const scalar_t *data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point,\n scalar_t *grad_value,\n scalar_t *grad_sampling_loc,\n scalar_t *grad_attn_weight)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n int _temp = index;\n const int c_col = _temp % channels;\n _temp /= channels;\n const int sampling_index = _temp; \n const int m_col = _temp % num_heads;\n _temp /= num_heads;\n const int q_col = _temp % num_query;\n _temp /= num_query;\n const int b_col = _temp;\n\n const scalar_t top_grad = grad_col[index];\n\n int data_weight_ptr = sampling_index * num_levels * num_point;\n int data_loc_w_ptr = data_weight_ptr << 1;\n const int grad_sampling_ptr = data_weight_ptr;\n grad_sampling_loc += grad_sampling_ptr << 1;\n grad_attn_weight += grad_sampling_ptr;\n const int grad_weight_stride = 1;\n const int grad_loc_stride = 2;\n const int qid_stride = num_heads * channels;\n const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;\n\n for (int l_col=0; l_col < num_levels; ++l_col)\n {\n const int level_start_id = data_level_start_index[l_col];\n const int spatial_h_ptr = l_col << 1;\n const int spatial_h = data_spatial_shapes[spatial_h_ptr];\n const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];\n const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;\n const scalar_t *data_value_ptr = data_value + value_ptr_offset;\n scalar_t *grad_value_ptr = grad_value + value_ptr_offset;\n\n for (int p_col=0; p_col < num_point; ++p_col)\n {\n const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];\n const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];\n const scalar_t weight = data_attn_weight[data_weight_ptr];\n\n const scalar_t h_im = loc_h * spatial_h - 0.5;\n const scalar_t w_im = loc_w * spatial_w - 0.5;\n if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)\n {\n ms_deform_attn_col2im_bilinear_gm(\n data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,\n top_grad, weight, grad_value_ptr, \n grad_sampling_loc, grad_attn_weight);\n }\n data_weight_ptr += 1;\n data_loc_w_ptr += 2;\n grad_attn_weight += grad_weight_stride;\n grad_sampling_loc += grad_loc_stride;\n }\n }\n }\n}\n\n\ntemplate \nvoid ms_deformable_im2col_cuda(cudaStream_t stream,\n const scalar_t* data_value,\n const int64_t* data_spatial_shapes, \n const int64_t* data_level_start_index, \n const scalar_t* data_sampling_loc,\n const scalar_t* data_attn_weight,\n const int batch_size,\n const int spatial_size, \n const int num_heads, \n const int channels, \n const int num_levels, \n const int num_query,\n const int num_point,\n scalar_t* data_col)\n{\n const int num_kernels = batch_size * num_query * num_heads * channels;\n const int num_actual_kernels = batch_size * num_query * num_heads * channels;\n const int num_threads = CUDA_NUM_THREADS;\n ms_deformable_im2col_gpu_kernel\n <<>>(\n num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, \n batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);\n \n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in ms_deformable_im2col_cuda: %s\\n\", cudaGetErrorString(err));\n }\n\n}\n\ntemplate \nvoid ms_deformable_col2im_cuda(cudaStream_t stream,\n const scalar_t* grad_col,\n const scalar_t* data_value,\n const int64_t * data_spatial_shapes,\n const int64_t * data_level_start_index,\n const scalar_t * data_sampling_loc,\n const scalar_t * data_attn_weight,\n const int batch_size, \n const int spatial_size, \n const int num_heads,\n const int channels, \n const int num_levels,\n const int num_query,\n const int num_point, \n scalar_t* grad_value,\n scalar_t* grad_sampling_loc,\n scalar_t* grad_attn_weight)\n{\n const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels;\n const int num_kernels = batch_size * num_query * num_heads * channels;\n const int num_actual_kernels = batch_size * num_query * num_heads * channels;\n if (channels > 1024)\n {\n if ((channels & 1023) == 0)\n {\n ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n }\n else\n {\n ms_deformable_col2im_gpu_kernel_gm\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n }\n }\n else{\n switch(channels)\n {\n case 1:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 2:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 4:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 8:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 16:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 32:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 64:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 128:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 256:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 512:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n case 1024:\n ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n break;\n default:\n if (channels < 64)\n {\n ms_deformable_col2im_gpu_kernel_shm_reduce_v1\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n }\n else\n {\n ms_deformable_col2im_gpu_kernel_shm_reduce_v2\n <<>>(\n num_kernels, \n grad_col,\n data_value,\n data_spatial_shapes,\n data_level_start_index, \n data_sampling_loc,\n data_attn_weight,\n batch_size, \n spatial_size, \n num_heads,\n channels, \n num_levels,\n num_query,\n num_point,\n grad_value,\n grad_sampling_loc,\n grad_attn_weight);\n }\n }\n }\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in ms_deformable_col2im_cuda: %s\\n\", cudaGetErrorString(err));\n }\n\n}"
},
{
"path": "models/encoder/ops/src/ms_deform_attn.h",
"content": "/*!\n**************************************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************************************\n* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n**************************************************************************************************\n*/\n\n#pragma once\n\n#include \"cpu/ms_deform_attn_cpu.h\"\n\n#ifdef WITH_CUDA\n#include \"cuda/ms_deform_attn_cuda.h\"\n#endif\n\n\nat::Tensor\nms_deform_attn_forward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const int im2col_step)\n{\n if (value.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return ms_deform_attn_cuda_forward(\n value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n AT_ERROR(\"Not implemented on the CPU\");\n}\n\nstd::vector\nms_deform_attn_backward(\n const at::Tensor &value, \n const at::Tensor &spatial_shapes,\n const at::Tensor &level_start_index,\n const at::Tensor &sampling_loc,\n const at::Tensor &attn_weight,\n const at::Tensor &grad_output,\n const int im2col_step)\n{\n if (value.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return ms_deform_attn_cuda_backward(\n value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n AT_ERROR(\"Not implemented on the CPU\");\n}\n\n"
},
{
"path": "models/encoder/ops/src/vision.cpp",
"content": "/*!\n**************************************************************************************************\n* Deformable DETR\n* Copyright (c) 2020 SenseTime. All Rights Reserved.\n* Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n**************************************************************************************************\n* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n**************************************************************************************************\n*/\n\n#include \"ms_deform_attn.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"ms_deform_attn_forward\", &ms_deform_attn_forward, \"ms_deform_attn_forward\");\n m.def(\"ms_deform_attn_backward\", &ms_deform_attn_backward, \"ms_deform_attn_backward\");\n}\n"
},
{
"path": "models/encoder/ops/test.py",
"content": "# ------------------------------------------------------------------------------------------------\n# Deformable DETR\n# Copyright (c) 2020 SenseTime. All Rights Reserved.\n# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n# ------------------------------------------------------------------------------------------------\n# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0\n# ------------------------------------------------------------------------------------------------\n\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport time\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import gradcheck\n\nfrom functions.ms_deform_attn_func import MSDeformAttnFunction, ms_deform_attn_core_pytorch\n\n\nN, M, D = 1, 2, 2\nLq, L, P = 2, 2, 2\nshapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()\nlevel_start_index = torch.cat((shapes.new_zeros((1, )), shapes.prod(1).cumsum(0)[:-1]))\nS = sum([(H*W).item() for H, W in shapes])\n\n\ntorch.manual_seed(3)\n\n\n@torch.no_grad()\ndef check_forward_equal_with_pytorch_double():\n value = torch.rand(N, S, M, D).cuda() * 0.01\n sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()\n attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5\n attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)\n im2col_step = 2\n output_pytorch = ms_deform_attn_core_pytorch(value.double(), shapes, sampling_locations.double(), attention_weights.double()).detach().cpu()\n output_cuda = MSDeformAttnFunction.apply(value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step).detach().cpu()\n fwdok = torch.allclose(output_cuda, output_pytorch)\n max_abs_err = (output_cuda - output_pytorch).abs().max()\n max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()\n\n print(f'* {fwdok} check_forward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')\n\n\n@torch.no_grad()\ndef check_forward_equal_with_pytorch_float():\n value = torch.rand(N, S, M, D).cuda() * 0.01\n sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()\n attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5\n attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)\n im2col_step = 2\n output_pytorch = ms_deform_attn_core_pytorch(value, shapes, sampling_locations, attention_weights).detach().cpu()\n output_cuda = MSDeformAttnFunction.apply(value, shapes, level_start_index, sampling_locations, attention_weights, im2col_step).detach().cpu()\n fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)\n max_abs_err = (output_cuda - output_pytorch).abs().max()\n max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()\n\n print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')\n\n\ndef check_gradient_numerical(channels=4, grad_value=True, grad_sampling_loc=True, grad_attn_weight=True):\n\n value = torch.rand(N, S, M, channels).cuda() * 0.01\n sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()\n attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5\n attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)\n im2col_step = 2\n func = MSDeformAttnFunction.apply\n\n value.requires_grad = grad_value\n sampling_locations.requires_grad = grad_sampling_loc\n attention_weights.requires_grad = grad_attn_weight\n\n gradok = gradcheck(func, (value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step))\n\n print(f'* {gradok} check_gradient_numerical(D={channels})')\n\n\nif __name__ == '__main__':\n check_forward_equal_with_pytorch_double()\n check_forward_equal_with_pytorch_float()\n\n for channels in [30, 32, 64, 71, 1025, 2048, 3096]:\n check_gradient_numerical(channels, True, True, True)\n\n\n\n"
},
{
"path": "models/layers/anyc_trans.py",
"content": "\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import repeat, rearrange, reduce\nfrom typing import Any, Optional\nfrom torch import Tensor\nfrom .utils import _get_activation_fn\n\nclass MLP(nn.Module):\n \"\"\" Very simple multi-layer perceptron (also called FFN)\"\"\"\n def __init__(self, input_dim, hidden_dim, output_dim, num_layers):\n super().__init__()\n self.num_layers = num_layers\n h = [hidden_dim] * (num_layers - 1)\n self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))\n\n def forward(self, x):\n for i, layer in enumerate(self.layers):\n x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)\n return x\n\nclass Linear_NormAct(nn.Linear):\n def __init__(self, *args, **kwargs):\n norm = kwargs.pop(\"norm\", None)\n activation = kwargs.pop(\"activation\", None)\n super().__init__(*args, **kwargs)\n out_features = kwargs['out_features']\n if norm == None:\n self.norm = None\n elif norm == 'ln':\n self.norm = nn.LayerNorm(out_features)\n else:\n raise ValueError()\n \n self.activation = _get_activation_fn(activation)\n\n def forward(self, x):\n x = F.linear(x, self.weight, self.bias)\n if self.norm is not None:\n x = self.norm(x)\n if self.activation is not None:\n x = self.activation(x)\n return x\n\nclass Conv2d_NormAct(torch.nn.Conv2d):\n def __init__(self, *args, **kwargs):\n norm = kwargs.pop(\"norm\", None)\n activation = kwargs.pop(\"activation\", None)\n super().__init__(*args, **kwargs)\n out_dim = kwargs['out_channels']\n if norm is None:\n self.norm = None\n elif norm == 'bn2d':\n # b c h w\n self.norm = nn.BatchNorm2d(out_dim)\n elif 'gn' in norm:\n # b c ..\n n_groups = int(norm.split('_')[-1])\n self.norm = nn.GroupNorm(n_groups, out_dim)\n else:\n raise ValueError()\n self.activation = _get_activation_fn(activation)\n\n def forward(self, x):\n # b c h w\n x = self._conv_forward(x, self.weight, self.bias)\n if self.norm is not None:\n x = self.norm(x)\n if self.activation is not None:\n x = self.activation(x)\n return x\n\nclass Conv3d_NormAct(torch.nn.Conv3d):\n def __init__(self, *args, **kwargs):\n norm = kwargs.pop(\"norm\", None)\n activation = kwargs.pop(\"activation\", None)\n super().__init__(*args, **kwargs)\n out_dim = kwargs['out_channels']\n if norm == None:\n self.norm = None\n elif 'gn' in norm:\n n_groups = int(norm.split('_')[-1])\n self.norm = nn.GroupNorm(n_groups, out_dim)\n else:\n raise ValueError()\n self.activation = _get_activation_fn(activation)\n\n def forward(self, x):\n x = self._conv_forward(x, self.weight, self.bias)\n if self.norm is not None:\n x = self.norm(x)\n if self.activation is not None:\n x = self.activation(x)\n return x\n\n"
},
{
"path": "models/layers/decoder_layers.py",
"content": "\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import repeat, rearrange, reduce\nfrom typing import Any, Optional\nfrom torch import Tensor\nfrom .utils import _get_activation_fn\n\nclass SelfAttentionLayer(nn.Module):\n\n def __init__(self, d_model, nhead, dropout=0.0,\n activation=\"relu\", normalize_before=False):\n super().__init__()\n self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)\n\n self.norm = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(dropout)\n\n self.activation = _get_activation_fn(activation)\n self.normalize_before = normalize_before\n\n self._reset_parameters()\n \n def _reset_parameters(self):\n for p in self.parameters():\n if p.dim() > 1:\n nn.init.xavier_uniform_(p)\n\n def with_pos_embed(self, tensor, pos: Optional[Tensor]):\n return tensor if pos is None else tensor + pos\n\n def forward_post(self, tgt,\n tgt_mask: Optional[Tensor] = None,\n tgt_key_padding_mask: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,\n ):\n q = k = self.with_pos_embed(tgt, query_pos)\n tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,\n key_padding_mask=tgt_key_padding_mask)[0]\n tgt = tgt + self.dropout(tgt2)\n tgt = self.norm(tgt)\n \n\n\n return tgt\n\n def forward_pre(self, tgt,\n tgt_mask: Optional[Tensor] = None,\n tgt_key_padding_mask: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,\n ):\n tgt2 = self.norm(tgt)\n q = k = self.with_pos_embed(tgt2, query_pos)\n tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,\n key_padding_mask=tgt_key_padding_mask)[0]\n tgt = tgt + self.dropout(tgt2)\n \n\n \n return tgt\n\n def forward(self, tgt,\n tgt_mask: Optional[Tensor] = None,\n tgt_key_padding_mask: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,):\n if self.normalize_before:\n return self.forward_pre(tgt, tgt_mask,\n tgt_key_padding_mask, query_pos)\n return self.forward_post(tgt, tgt_mask,\n tgt_key_padding_mask, query_pos)\n\nclass CrossAttentionLayer(nn.Module):\n\n def __init__(self, d_model, nhead, dropout=0.0,\n activation=\"relu\", normalize_before=False):\n super().__init__()\n self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)\n\n self.norm = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(dropout)\n\n self.activation = _get_activation_fn(activation)\n self.normalize_before = normalize_before\n \n self._reset_parameters()\n \n def _reset_parameters(self):\n for p in self.parameters():\n if p.dim() > 1:\n nn.init.xavier_uniform_(p)\n\n def with_pos_embed(self, tensor, pos: Optional[Tensor]):\n return tensor if pos is None else tensor + pos\n\n def forward_post(self, tgt, memory,\n memory_mask: Optional[Tensor] = None,\n memory_key_padding_mask: Optional[Tensor] = None,\n pos: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,\n ):\n tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),\n key=self.with_pos_embed(memory, pos),\n value=memory, attn_mask=memory_mask,\n key_padding_mask=memory_key_padding_mask)[0]\n tgt = tgt + self.dropout(tgt2)\n tgt = self.norm(tgt) # n b d\n\n \n return tgt\n\n def forward_pre(self, tgt, memory,\n memory_mask: Optional[Tensor] = None,\n memory_key_padding_mask: Optional[Tensor] = None,\n pos: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,\n ):\n tgt2 = self.norm(tgt)\n tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),\n key=self.with_pos_embed(memory, pos),\n value=memory, attn_mask=memory_mask,\n key_padding_mask=memory_key_padding_mask)[0]\n tgt = tgt + self.dropout(tgt2)\n return tgt\n\n def forward(self, tgt, memory,\n memory_mask: Optional[Tensor] = None,\n memory_key_padding_mask: Optional[Tensor] = None,\n pos: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,\n ):\n if self.normalize_before:\n return self.forward_pre(tgt, memory, memory_mask,\n memory_key_padding_mask, pos, query_pos)\n return self.forward_post(tgt, memory, memory_mask,\n memory_key_padding_mask, pos, query_pos)\n\nclass FFNLayer(nn.Module):\n\n def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,\n activation=\"relu\", normalize_before=False):\n super().__init__()\n # Implementation of Feedforward model\n self.linear1 = nn.Linear(d_model, dim_feedforward)\n self.dropout = nn.Dropout(dropout)\n self.linear2 = nn.Linear(dim_feedforward, d_model)\n\n self.norm = nn.LayerNorm(d_model)\n\n self.activation = _get_activation_fn(activation)\n self.normalize_before = normalize_before\n\n self._reset_parameters()\n \n def _reset_parameters(self):\n for p in self.parameters():\n if p.dim() > 1:\n nn.init.xavier_uniform_(p)\n\n def with_pos_embed(self, tensor, pos: Optional[Tensor]):\n return tensor if pos is None else tensor + pos\n\n def forward_post(self, tgt):\n tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))\n tgt = tgt + self.dropout(tgt2)\n tgt = self.norm(tgt)\n return tgt\n\n def forward_pre(self, tgt):\n tgt2 = self.norm(tgt)\n tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))\n tgt = tgt + self.dropout(tgt2)\n return tgt\n\n def forward(self, tgt):\n if self.normalize_before:\n return self.forward_pre(tgt)\n return self.forward_post(tgt)\n \n"
},
{
"path": "models/layers/gilbert/demo/index.html",
"content": "\n\n\n \n Gilbert Curve\n\n \n\n \n \n\n\n\n \n \n\n\n\n \n
![](\"assets/images/pipeline.png\")
\n