├── mmdet ├── VERSION ├── ops │ ├── iou │ │ └── __init__.py │ ├── box_iou_rotated │ │ ├── __init__.py │ │ ├── src │ │ │ ├── box_iou_rotated_ext.cpp │ │ │ └── box_iou_rotated_cpu.cpp │ │ └── box_iou_rotated_wrapper.py │ ├── convex │ │ ├── __init__.py │ │ ├── convex_wrapper.py │ │ └── src │ │ │ └── convex_ext.cpp │ ├── corner_pool │ │ ├── __init__.py │ │ └── corner_pool.py │ ├── roi_pool │ │ ├── __init__.py │ │ ├── gradcheck.py │ │ └── roi_pool.py │ ├── roi_align │ │ ├── __init__.py │ │ └── gradcheck.py │ ├── masked_conv │ │ ├── __init__.py │ │ └── src │ │ │ ├── masked_conv2d_ext.cpp │ │ │ └── cuda │ │ │ └── masked_conv2d_cuda.cpp │ ├── nms │ │ ├── __init__.py │ │ └── src │ │ │ ├── cuda │ │ │ └── nms_cuda.cpp │ │ │ └── nms_ext.cpp │ ├── roi_align_rotated │ │ ├── __init__.py │ │ └── roi_align_rotated.py │ ├── sigmoid_focal_loss │ │ ├── __init__.py │ │ ├── sigmoid_focal_loss.py │ │ └── src │ │ │ └── sigmoid_focal_loss_ext.cpp │ ├── nms_rotated │ │ ├── __init__.py │ │ └── src │ │ │ ├── poly_nms_cpu.cpp │ │ │ ├── nms_rotated_ext.cpp │ │ │ └── nms_rotated_cpu.cpp │ ├── carafe │ │ ├── __init__.py │ │ ├── setup.py │ │ ├── src │ │ │ ├── carafe_naive_ext.cpp │ │ │ ├── carafe_ext.cpp │ │ │ └── cuda │ │ │ │ └── carafe_naive_cuda.cpp │ │ └── grad_check.py │ ├── dcn2 │ │ └── __init__.py │ ├── utils │ │ ├── __init__.py │ │ └── src │ │ │ └── compiling_info.cpp │ ├── dcn │ │ ├── __init__.py │ │ └── src │ │ │ └── deform_pool_ext.cpp │ ├── plugin.py │ └── __init__.py ├── datasets │ ├── obb │ │ └── __init__.py │ ├── pipelines │ │ ├── obb │ │ │ ├── __init__.py │ │ │ └── .base.py.swp │ │ ├── __init__.py │ │ ├── compose.py │ │ └── auto_augment.py │ ├── samplers │ │ ├── __init__.py │ │ └── distributed_sampler.py │ ├── deepfashion.py │ ├── __init__.py │ └── wider_face.py ├── models │ ├── losses │ │ ├── obb │ │ │ └── _init_.py │ │ ├── __init__.py │ │ ├── mse_loss.py │ │ └── accuracy.py │ ├── detectors │ │ ├── obb │ │ │ ├── __init__.py │ │ │ ├── retinanet_obb.py │ │ │ ├── gliding_vertex.py │ │ │ ├── faster_rcnn_obb.py │ │ │ ├── roi_transformer.py │ │ │ ├── obb_test_mixins.py │ │ │ └── obb_base.py │ │ ├── htc.py │ │ ├── gfl.py │ │ ├── atss.py │ │ ├── fcos.py │ │ ├── fsaf.py │ │ ├── fovea.py │ │ ├── retinanet.py │ │ ├── nasfcos.py │ │ ├── mask_rcnn.py │ │ ├── faster_rcnn.py │ │ ├── mask_scoring_rcnn.py │ │ ├── point_rend.py │ │ ├── grid_rcnn.py │ │ ├── cascade_rcnn.py │ │ ├── __init__.py │ │ └── fast_rcnn.py │ ├── dense_heads │ │ ├── obb │ │ │ └── __init__.py │ │ ├── __init__.py │ │ ├── base_dense_head.py │ │ ├── rpn_test_mixin.py │ │ └── nasfcos_head.py │ ├── roi_heads │ │ ├── obb │ │ │ ├── __init__.py │ │ │ └── obb_base_roi_head.py │ │ ├── bbox_heads │ │ │ ├── obb │ │ │ │ └── __init__.py │ │ │ └── __init__.py │ │ ├── roi_extractors │ │ │ ├── obb │ │ │ │ ├── __init__.py │ │ │ │ └── obb_base_roi_extractor.py │ │ │ ├── __init__.py │ │ │ └── base_roi_extractor.py │ │ ├── shared_heads │ │ │ ├── __init__.py │ │ │ └── res_layer.py │ │ ├── mask_heads │ │ │ ├── __init__.py │ │ │ └── htc_mask_head.py │ │ ├── double_roi_head.py │ │ └── __init__.py │ ├── utils │ │ └── __init__.py │ ├── necks │ │ └── __init__.py │ ├── backbones │ │ └── __init__.py │ ├── __init__.py │ ├── show.py │ └── builder.py ├── core │ ├── bbox │ │ ├── coder │ │ │ ├── obb │ │ │ │ └── __init__.py │ │ │ ├── base_bbox_coder.py │ │ │ ├── pseudo_bbox_coder.py │ │ │ └── __init__.py │ │ ├── iou_calculators │ │ │ ├── obb │ │ │ │ ├── __init__.py │ │ │ │ └── obbiou_calculator.py │ │ │ ├── builder.py │ │ │ └── __init__.py │ │ ├── samplers │ │ │ ├── obb │ │ │ │ ├── __init__.py │ │ │ │ └── obb_random_sampler.py │ │ │ ├── combined_sampler.py │ │ │ ├── __init__.py │ │ │ ├── pseudo_sampler.py │ │ │ ├── instance_balanced_pos_sampler.py │ │ │ └── random_sampler.py │ │ ├── assigners │ │ │ ├── base_assigner.py │ │ │ └── __init__.py │ │ ├── transforms_obb │ │ │ └── __init__.py │ │ ├── builder.py │ │ ├── demodata.py │ │ └── __init__.py │ ├── evaluation │ │ ├── obb │ │ │ └── __init__.py │ │ ├── __init__.py │ │ ├── bbox_overlaps.py │ │ └── eval_hooks.py │ ├── fp16 │ │ ├── __init__.py │ │ └── utils.py │ ├── utils │ │ ├── __init__.py │ │ ├── dist_utils.py │ │ └── misc.py │ ├── anchor │ │ ├── builder.py │ │ ├── __init__.py │ │ ├── point_generator.py │ │ └── utils.py │ ├── post_processing │ │ ├── obb │ │ │ ├── __init__.py │ │ │ └── obb_nms.py │ │ ├── __init__.py │ │ └── bbox_nms.py │ ├── mask │ │ ├── __init__.py │ │ ├── utils.py │ │ └── mask_target.py │ └── __init__.py ├── __init__.py ├── utils │ ├── __init__.py │ ├── logger.py │ ├── profiling.py │ └── collect_env.py └── apis │ └── __init__.py ├── requirements ├── readthedocs.txt ├── docs.txt ├── optional.txt ├── build.txt ├── runtime.txt └── tests.txt ├── architecture.jpg ├── BboxToolkit ├── BboxToolkit │ ├── vis │ │ ├── __init__.py │ │ ├── vis_xie.py │ │ └── base.py │ ├── datasets │ │ ├── __init__.py │ │ └── io.py │ ├── __init__.py │ ├── utils.py │ └── move.py └── setup.py ├── requirements.txt ├── tools ├── dist_train.sh ├── dist_test.sh ├── slurm_test.sh └── slurm_train.sh ├── configs ├── schedule_1x.py └── default_runtime.py └── .gitignore /mmdet/VERSION: -------------------------------------------------------------------------------- 1 | 2.2.0 2 | -------------------------------------------------------------------------------- /mmdet/ops/iou/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/datasets/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/models/losses/obb/_init_.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/core/bbox/coder/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/core/evaluation/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/datasets/pipelines/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/models/dense_heads/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/obb/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /mmdet/core/bbox/iou_calculators/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/bbox_heads/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/roi_extractors/obb/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requirements/readthedocs.txt: -------------------------------------------------------------------------------- 1 | mmcv 2 | torch 3 | torchvision 4 | -------------------------------------------------------------------------------- /architecture.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yanqingyao1994/QPDET/HEAD/architecture.jpg -------------------------------------------------------------------------------- /mmdet/ops/box_iou_rotated/__init__.py: -------------------------------------------------------------------------------- 1 | from .box_iou_rotated_wrapper import obb_overlaps 2 | -------------------------------------------------------------------------------- /mmdet/models/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .res_layer import ResLayer 2 | 3 | __all__ = ['ResLayer'] 4 | -------------------------------------------------------------------------------- /requirements/docs.txt: -------------------------------------------------------------------------------- 1 | recommonmark 2 | sphinx 3 | sphinx_markdown_tables 4 | sphinx_rtd_theme 5 | -------------------------------------------------------------------------------- /requirements/optional.txt: -------------------------------------------------------------------------------- 1 | albumentations>=0.3.2 2 | cityscapesscripts 3 | imagecorruptions 4 | -------------------------------------------------------------------------------- /mmdet/ops/convex/__init__.py: -------------------------------------------------------------------------------- 1 | from .convex_wrapper import convex_sort 2 | 3 | __all = ['convex_sort'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/corner_pool/__init__.py: -------------------------------------------------------------------------------- 1 | from .corner_pool import CornerPool 2 | 3 | __all__ = ['CornerPool'] 4 | -------------------------------------------------------------------------------- /requirements/build.txt: -------------------------------------------------------------------------------- 1 | # These must be installed before building mmdetection 2 | numpy 3 | torch>=1.3 4 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/shared_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .res_layer import ResLayer 2 | 3 | __all__ = ['ResLayer'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/roi_pool/__init__.py: -------------------------------------------------------------------------------- 1 | from .roi_pool import RoIPool, roi_pool 2 | 3 | __all__ = ['roi_pool', 'RoIPool'] 4 | -------------------------------------------------------------------------------- /mmdet/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import __version__, short_version 2 | 3 | __all__ = ['__version__', 'short_version'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/roi_align/__init__.py: -------------------------------------------------------------------------------- 1 | from .roi_align import RoIAlign, roi_align 2 | 3 | __all__ = ['roi_align', 'RoIAlign'] 4 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/vis/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import color_val, imshow 2 | from .bbox import imshow_bboxes, imshow_det_bboxes 3 | -------------------------------------------------------------------------------- /mmdet/datasets/pipelines/obb/.base.py.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yanqingyao1994/QPDET/HEAD/mmdet/datasets/pipelines/obb/.base.py.swp -------------------------------------------------------------------------------- /mmdet/ops/masked_conv/__init__.py: -------------------------------------------------------------------------------- 1 | from .masked_conv import MaskedConv2d, masked_conv2d 2 | 3 | __all__ = ['masked_conv2d', 'MaskedConv2d'] 4 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | -r requirements/build.txt 2 | -r requirements/optional.txt 3 | -r requirements/runtime.txt 4 | -r requirements/tests.txt 5 | -------------------------------------------------------------------------------- /mmdet/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .collect_env import collect_env 2 | from .logger import get_root_logger 3 | 4 | __all__ = ['get_root_logger', 'collect_env'] 5 | -------------------------------------------------------------------------------- /mmdet/ops/nms/__init__.py: -------------------------------------------------------------------------------- 1 | from .nms_wrapper import batched_nms, nms, nms_match, soft_nms 2 | 3 | __all__ = ['nms', 'soft_nms', 'batched_nms', 'nms_match'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/roi_align_rotated/__init__.py: -------------------------------------------------------------------------------- 1 | from .roi_align_rotated import RoIAlignRotated, roi_align_rotated 2 | 3 | __all__ = ['roi_align_rotated', 'RoIAlignRotated'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/sigmoid_focal_loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .sigmoid_focal_loss import SigmoidFocalLoss, sigmoid_focal_loss 2 | 3 | __all__ = ['SigmoidFocalLoss', 'sigmoid_focal_loss'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/nms_rotated/__init__.py: -------------------------------------------------------------------------------- 1 | from .nms_rotated_wrapper import obb_nms, poly_nms, BT_nms, arb_batched_nms 2 | 3 | __all__ = ['obb_nms', 'poly_nms', 'BT_nms', 'arb_batched_nms'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/carafe/__init__.py: -------------------------------------------------------------------------------- 1 | from .carafe import CARAFE, CARAFENaive, CARAFEPack, carafe, carafe_naive 2 | 3 | __all__ = ['carafe', 'carafe_naive', 'CARAFE', 'CARAFENaive', 'CARAFEPack'] 4 | -------------------------------------------------------------------------------- /mmdet/ops/nms_rotated/src/poly_nms_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | template 4 | at::Tensor poly_nms_cpu_kernel(const at::Tensor& dets, const float threshold) { 5 | 6 | -------------------------------------------------------------------------------- /requirements/runtime.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | mmcv>=0.6.2 3 | numpy 4 | # need older pillow until torchvision is fixed 5 | Pillow<=6.2.2 6 | six 7 | terminaltables 8 | torch>=1.3 9 | torchvision 10 | -------------------------------------------------------------------------------- /mmdet/core/fp16/__init__.py: -------------------------------------------------------------------------------- 1 | from .decorators import auto_fp16, force_fp32 2 | from .hooks import Fp16OptimizerHook, wrap_fp16_model 3 | 4 | __all__ = ['auto_fp16', 'force_fp32', 'Fp16OptimizerHook', 'wrap_fp16_model'] 5 | -------------------------------------------------------------------------------- /mmdet/datasets/samplers/__init__.py: -------------------------------------------------------------------------------- 1 | from .distributed_sampler import DistributedSampler 2 | from .group_sampler import DistributedGroupSampler, GroupSampler 3 | 4 | __all__ = ['DistributedSampler', 'DistributedGroupSampler', 'GroupSampler'] 5 | -------------------------------------------------------------------------------- /mmdet/core/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .dist_utils import DistOptimizerHook, allreduce_grads 2 | from .misc import multi_apply, tensor2imgs, unmap 3 | 4 | __all__ = [ 5 | 'allreduce_grads', 'DistOptimizerHook', 'tensor2imgs', 'multi_apply', 6 | 'unmap' 7 | ] 8 | -------------------------------------------------------------------------------- /mmdet/core/anchor/builder.py: -------------------------------------------------------------------------------- 1 | from mmcv.utils import Registry, build_from_cfg 2 | 3 | ANCHOR_GENERATORS = Registry('Anchor generator') 4 | 5 | 6 | def build_anchor_generator(cfg, default_args=None): 7 | return build_from_cfg(cfg, ANCHOR_GENERATORS, default_args) 8 | -------------------------------------------------------------------------------- /mmdet/ops/dcn2/__init__.py: -------------------------------------------------------------------------------- 1 | from .deform_pool import (DeformRoIPooling2, DeformRoIPoolingPack2, 2 | ModulatedDeformRoIPoolingPack2) 3 | 4 | __all__ = [ 5 | 'DeformRoIPooling2', 'DeformRoIPoolingPack2', 'ModulatedDeformRoIPoolingPack2' 6 | ] 7 | -------------------------------------------------------------------------------- /requirements/tests.txt: -------------------------------------------------------------------------------- 1 | asynctest 2 | codecov 3 | flake8 4 | interrogate 5 | isort 6 | # Note: used for kwarray.group_items, this may be ported to mmcv in the future. 7 | kwarray 8 | pytest 9 | pytest-cov 10 | pytest-runner 11 | ubelt 12 | xdoctest >= 0.10.0 13 | yapf 14 | -------------------------------------------------------------------------------- /mmdet/core/post_processing/obb/__init__.py: -------------------------------------------------------------------------------- 1 | from .obb_nms import multiclass_arb_nms 2 | from .obb_merge_augs import (merge_rotate_aug_proposals, merge_rotate_aug_hbb, 3 | merge_rotate_aug_obb, merge_rotate_aug_poly, 4 | merge_rotate_aug_arb) 5 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/obb/__init__.py: -------------------------------------------------------------------------------- 1 | from .obb_sampling_result import OBBSamplingResult 2 | from .obb_base_sampler import OBBBaseSampler 3 | from .obb_random_sampler import OBBRandomSampler 4 | from .obb_ohem_sampler import OBBOHEMSampler 5 | from .obb_random_sampler_xyrab import OBBRandomSamplerXYRAB 6 | -------------------------------------------------------------------------------- /tools/dist_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CONFIG=$1 4 | GPUS=$2 5 | PORT=${PORT:-29500} 6 | 7 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 8 | python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 9 | $(dirname "$0")/train.py $CONFIG --launcher pytorch ${@:3} 10 | -------------------------------------------------------------------------------- /mmdet/core/bbox/iou_calculators/builder.py: -------------------------------------------------------------------------------- 1 | from mmcv.utils import Registry, build_from_cfg 2 | 3 | IOU_CALCULATORS = Registry('IoU calculator') 4 | 5 | 6 | def build_iou_calculator(cfg, default_args=None): 7 | """Builder of IoU calculator""" 8 | return build_from_cfg(cfg, IOU_CALCULATORS, default_args) 9 | -------------------------------------------------------------------------------- /mmdet/core/mask/__init__.py: -------------------------------------------------------------------------------- 1 | from .mask_target import mask_target 2 | from .structures import BitmapMasks, PolygonMasks 3 | from .utils import encode_mask_results, split_combined_polys 4 | 5 | __all__ = [ 6 | 'split_combined_polys', 'mask_target', 'BitmapMasks', 'PolygonMasks', 7 | 'encode_mask_results' 8 | ] 9 | -------------------------------------------------------------------------------- /tools/dist_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CONFIG=$1 4 | CHECKPOINT=$2 5 | GPUS=$3 6 | PORT=${PORT:-29500} 7 | 8 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 9 | python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \ 10 | $(dirname "$0")/test.py $CONFIG $CHECKPOINT --launcher pytorch ${@:4} 11 | -------------------------------------------------------------------------------- /mmdet/core/__init__.py: -------------------------------------------------------------------------------- 1 | from .anchor import * # noqa: F401, F403 2 | from .bbox import * # noqa: F401, F403 3 | from .evaluation import * # noqa: F401, F403 4 | from .fp16 import * # noqa: F401, F403 5 | from .mask import * # noqa: F401, F403 6 | from .post_processing import * # noqa: F401, F403 7 | from .utils import * # noqa: F401, F403 8 | -------------------------------------------------------------------------------- /mmdet/core/bbox/iou_calculators/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_iou_calculator 2 | from .iou2d_calculator import BboxOverlaps2D, bbox_overlaps 3 | 4 | from .obb.obbiou_calculator import OBBOverlaps, PolyOverlaps 5 | 6 | __all__ = ['build_iou_calculator', 'BboxOverlaps2D', 'bbox_overlaps', 7 | 'OBBOverlaps', 'PolyOverlaps' 8 | ] 9 | -------------------------------------------------------------------------------- /mmdet/ops/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # from . import compiling_info 2 | from .compiling_info import get_compiler_version, get_compiling_cuda_version 3 | 4 | # get_compiler_version = compiling_info.get_compiler_version 5 | # get_compiling_cuda_version = compiling_info.get_compiling_cuda_version 6 | 7 | __all__ = ['get_compiler_version', 'get_compiling_cuda_version'] 8 | -------------------------------------------------------------------------------- /configs/schedule_1x.py: -------------------------------------------------------------------------------- 1 | # optimizer 2 | optimizer = dict(type='SGD', lr=0.005, momentum=0.9, weight_decay=0.0001) 3 | optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2)) 4 | # learning policy 5 | lr_config = dict( 6 | policy='step', 7 | warmup='linear', 8 | warmup_iters=500, 9 | warmup_ratio=0.001, 10 | step=[8, 11]) 11 | total_epochs = 12 12 | -------------------------------------------------------------------------------- /mmdet/datasets/deepfashion.py: -------------------------------------------------------------------------------- 1 | from .builder import DATASETS 2 | from .coco import CocoDataset 3 | 4 | 5 | @DATASETS.register_module() 6 | class DeepFashionDataset(CocoDataset): 7 | 8 | CLASSES = ('top', 'skirt', 'leggings', 'dress', 'outer', 'pants', 'bag', 9 | 'neckwear', 'headwear', 'eyeglass', 'belt', 'footwear', 'hair', 10 | 'skin', 'face') 11 | -------------------------------------------------------------------------------- /mmdet/models/necks/__init__.py: -------------------------------------------------------------------------------- 1 | from .bfp import BFP 2 | from .fpn import FPN 3 | from .fpn_carafe import FPN_CARAFE 4 | from .hrfpn import HRFPN 5 | from .nas_fpn import NASFPN 6 | from .nasfcos_fpn import NASFCOS_FPN 7 | from .pafpn import PAFPN 8 | from .rfp import RFP 9 | 10 | __all__ = [ 11 | 'FPN', 'BFP', 'HRFPN', 'NASFPN', 'FPN_CARAFE', 'PAFPN', 'NASFCOS_FPN', 12 | 'RFP' 13 | ] 14 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .mean_ap import eval_map 2 | from .misc import (get_classes, change_cls_order, merge_prior_contents, 3 | split_imgset) 4 | from .io import load_imgs, load_pkl, save_pkl 5 | from .DOTAio import load_dota, load_dota_submission, save_dota_submission 6 | from .DIORio import load_dior_hbb, load_dior_obb, load_dior 7 | from .HRSCio import load_hrsc 8 | -------------------------------------------------------------------------------- /configs/default_runtime.py: -------------------------------------------------------------------------------- 1 | checkpoint_config = dict(interval=1) 2 | # yapf:disable 3 | log_config = dict( 4 | interval=50, 5 | hooks=[ 6 | dict(type='TextLoggerHook'), 7 | # dict(type='TensorboardLoggerHook') 8 | ]) 9 | # yapf:enable 10 | dist_params = dict(backend='nccl') 11 | log_level = 'INFO' 12 | load_from = None 13 | resume_from = None 14 | workflow = [('train', 1)] 15 | -------------------------------------------------------------------------------- /BboxToolkit/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup(name='BboxToolkit', 4 | version='1.0', 5 | description='a tiny toolkit for bounding boxes', 6 | author='jbwang', 7 | packages=find_packages(), 8 | install_requires=[ 9 | 'opencv-python', 10 | 'terminaltables', 11 | 'pillow', 12 | 'shapely', 13 | 'numpy']) 14 | -------------------------------------------------------------------------------- /mmdet/core/bbox/assigners/base_assigner.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | 4 | class BaseAssigner(metaclass=ABCMeta): 5 | """Base assigner that assigns boxes to ground truth boxes""" 6 | 7 | @abstractmethod 8 | def assign(self, bboxes, gt_bboxes, gt_bboxes_ignore=None, gt_labels=None): 9 | """Assign boxes to either a ground truth boxe or a negative boxes""" 10 | pass 11 | -------------------------------------------------------------------------------- /mmdet/core/anchor/__init__.py: -------------------------------------------------------------------------------- 1 | from .anchor_generator import AnchorGenerator, LegacyAnchorGenerator 2 | from .builder import ANCHOR_GENERATORS, build_anchor_generator 3 | from .point_generator import PointGenerator 4 | from .utils import anchor_inside_flags, calc_region, images_to_levels 5 | 6 | __all__ = [ 7 | 'AnchorGenerator', 'LegacyAnchorGenerator', 'anchor_inside_flags', 8 | 'PointGenerator', 'images_to_levels', 'calc_region', 9 | 'build_anchor_generator', 'ANCHOR_GENERATORS' 10 | ] 11 | -------------------------------------------------------------------------------- /mmdet/apis/__init__.py: -------------------------------------------------------------------------------- 1 | from .inference import (async_inference_detector, inference_detector, 2 | init_detector, show_result_pyplot) 3 | from .test import multi_gpu_test, single_gpu_test 4 | from .train import get_root_logger, set_random_seed, train_detector 5 | 6 | __all__ = [ 7 | 'get_root_logger', 'set_random_seed', 'train_detector', 'init_detector', 8 | 'async_inference_detector', 'inference_detector', 'show_result_pyplot', 9 | 'multi_gpu_test', 'single_gpu_test' 10 | ] 11 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/mask_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .coarse_mask_head import CoarseMaskHead 2 | from .fcn_mask_head import FCNMaskHead 3 | from .fused_semantic_head import FusedSemanticHead 4 | from .grid_head import GridHead 5 | from .htc_mask_head import HTCMaskHead 6 | from .mask_point_head import MaskPointHead 7 | from .maskiou_head import MaskIoUHead 8 | 9 | __all__ = [ 10 | 'FCNMaskHead', 'HTCMaskHead', 'FusedSemanticHead', 'GridHead', 11 | 'MaskIoUHead', 'CoarseMaskHead', 'MaskPointHead' 12 | ] 13 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/__init__.py: -------------------------------------------------------------------------------- 1 | pi = 3.141592 2 | from .datasets import * 3 | import BboxToolkit.datasets 4 | 5 | from .vis import * 6 | import BboxToolkit.vis 7 | 8 | from .utils import (get_bbox_type, get_bbox_dim, choice_by_type, 9 | regular_theta, regular_obb) 10 | from .transforms import (poly2hbb, poly2obb, rectpoly2obb, obb2poly, obb2hbb, 11 | hbb2poly, hbb2obb, bbox2type) 12 | from .geometry import bbox_overlaps, bbox_areas, bbox_nms 13 | from .move import translate, flip, warp 14 | -------------------------------------------------------------------------------- /mmdet/models/detectors/htc.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .cascade_rcnn import CascadeRCNN 3 | 4 | 5 | @DETECTORS.register_module() 6 | class HybridTaskCascade(CascadeRCNN): 7 | """Implementation of `HTC `_""" 8 | 9 | def __init__(self, **kwargs): 10 | super(HybridTaskCascade, self).__init__(**kwargs) 11 | 12 | @property 13 | def with_semantic(self): 14 | """bool: whether the detector has a semantic head""" 15 | return self.roi_head.with_semantic 16 | -------------------------------------------------------------------------------- /mmdet/models/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | from .detectors_resnet import DetectoRS_ResNet 2 | from .detectors_resnext import DetectoRS_ResNeXt 3 | from .hourglass import HourglassNet 4 | from .hrnet import HRNet 5 | from .regnet import RegNet 6 | from .res2net import Res2Net 7 | from .resnet import ResNet, ResNetV1d 8 | from .resnext import ResNeXt 9 | from .ssd_vgg import SSDVGG 10 | 11 | __all__ = [ 12 | 'RegNet', 'ResNet', 'ResNetV1d', 'ResNeXt', 'SSDVGG', 'HRNet', 'Res2Net', 13 | 'HourglassNet', 'DetectoRS_ResNet', 'DetectoRS_ResNeXt' 14 | ] 15 | -------------------------------------------------------------------------------- /mmdet/models/detectors/gfl.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class GFL(SingleStageDetector): 7 | 8 | def __init__(self, 9 | backbone, 10 | neck, 11 | bbox_head, 12 | train_cfg=None, 13 | test_cfg=None, 14 | pretrained=None): 15 | super(GFL, self).__init__(backbone, neck, bbox_head, train_cfg, 16 | test_cfg, pretrained) 17 | -------------------------------------------------------------------------------- /mmdet/models/detectors/atss.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class ATSS(SingleStageDetector): 7 | 8 | def __init__(self, 9 | backbone, 10 | neck, 11 | bbox_head, 12 | train_cfg=None, 13 | test_cfg=None, 14 | pretrained=None): 15 | super(ATSS, self).__init__(backbone, neck, bbox_head, train_cfg, 16 | test_cfg, pretrained) 17 | -------------------------------------------------------------------------------- /mmdet/utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from mmcv.utils import get_logger 4 | 5 | 6 | def get_root_logger(log_file=None, log_level=logging.INFO): 7 | """Get root logger 8 | 9 | Args: 10 | log_file (str, optional): File path of log. Defaults to None. 11 | log_level (int, optional): The level of logger. 12 | Defaults to logging.INFO. 13 | 14 | Returns: 15 | :obj:`logging.Logger`: The obtained logger 16 | """ 17 | logger = get_logger(name='mmdet', log_file=log_file, log_level=log_level) 18 | 19 | return logger 20 | -------------------------------------------------------------------------------- /mmdet/ops/nms/src/cuda/nms_cuda.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #include 3 | 4 | #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ") 5 | 6 | at::Tensor nms_cuda_forward(const at::Tensor boxes, float nms_overlap_thresh); 7 | 8 | at::Tensor nms_cuda(const at::Tensor& dets, const float threshold) { 9 | CHECK_CUDA(dets); 10 | if (dets.numel() == 0) 11 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 12 | return nms_cuda_forward(dets, threshold); 13 | } 14 | -------------------------------------------------------------------------------- /mmdet/core/bbox/coder/base_bbox_coder.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | 4 | class BaseBBoxCoder(metaclass=ABCMeta): 5 | """Base bounding box coder""" 6 | 7 | def __init__(self, **kwargs): 8 | pass 9 | 10 | @abstractmethod 11 | def encode(self, bboxes, gt_bboxes): 12 | """Encode deltas between bboxes and ground truth boxes""" 13 | pass 14 | 15 | @abstractmethod 16 | def decode(self, bboxes, bboxes_pred): 17 | """ 18 | Decode the predicted bboxes according to prediction and base boxes 19 | """ 20 | pass 21 | -------------------------------------------------------------------------------- /mmdet/ops/roi_pool/gradcheck.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import torch 5 | from torch.autograd import gradcheck 6 | 7 | sys.path.append(osp.abspath(osp.join(__file__, '../../'))) 8 | from roi_pool import RoIPool # noqa: E402, isort:skip 9 | 10 | feat = torch.randn(4, 16, 15, 15, requires_grad=True).cuda() 11 | rois = torch.Tensor([[0, 0, 0, 50, 50], [0, 10, 30, 43, 55], 12 | [1, 67, 40, 110, 120]]).cuda() 13 | inputs = (feat, rois) 14 | print('Gradcheck for roi pooling...') 15 | test = gradcheck(RoIPool(4, 1.0 / 8), inputs, eps=1e-5, atol=1e-3) 16 | print(test) 17 | -------------------------------------------------------------------------------- /mmdet/core/bbox/transforms_obb/__init__.py: -------------------------------------------------------------------------------- 1 | from .form import (poly2obb, rectpoly2obb, poly2hbb, obb2poly, obb2hbb, 2 | hbb2poly, hbb2obb, bbox2type) 3 | from .mapping import (hbb_flip, obb_flip, poly_flip, hbb_warp, obb_warp, 4 | poly_warp, hbb_mapping, obb_mapping, poly_mapping, 5 | hbb_mapping_back, obb_mapping_back, poly_mapping_back, 6 | arb_mapping, arb_mapping_back) 7 | from .misc import (get_bbox_type, get_bbox_dim, get_bbox_areas, choice_by_type, arb2result, 8 | arb2roi, regular_theta, regular_obb, mintheta_obb) 9 | -------------------------------------------------------------------------------- /mmdet/core/bbox/assigners/__init__.py: -------------------------------------------------------------------------------- 1 | from .approx_max_iou_assigner import ApproxMaxIoUAssigner 2 | from .assign_result import AssignResult 3 | from .atss_assigner import ATSSAssigner 4 | from .base_assigner import BaseAssigner 5 | from .center_region_assigner import CenterRegionAssigner 6 | from .max_iou_assigner import MaxIoUAssigner 7 | from .point_assigner import PointAssigner 8 | 9 | from .obb2hbb_max_iou_assigner import OBB2HBBMaxIoUAssigner 10 | 11 | __all__ = [ 12 | 'BaseAssigner', 'MaxIoUAssigner', 'ApproxMaxIoUAssigner', 'AssignResult', 13 | 'PointAssigner', 'ATSSAssigner', 'CenterRegionAssigner' 14 | ] 15 | -------------------------------------------------------------------------------- /mmdet/ops/convex/convex_wrapper.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Function 2 | from . import convex_ext 3 | 4 | 5 | class ConvexSortFunction(Function): 6 | 7 | @staticmethod 8 | def forward(ctx, pts, masks, circular): 9 | idx = convex_ext.convex_sort(pts, masks, circular) 10 | ctx.mark_non_differentiable(idx) 11 | return idx 12 | 13 | @staticmethod 14 | def backward(ctx, grad_output): 15 | return () 16 | 17 | convex_sort_func = ConvexSortFunction.apply 18 | 19 | 20 | def convex_sort(pts, masks, circular=True): 21 | return convex_sort_func(pts, masks, circular) 22 | 23 | 24 | -------------------------------------------------------------------------------- /mmdet/core/bbox/coder/pseudo_bbox_coder.py: -------------------------------------------------------------------------------- 1 | from ..builder import BBOX_CODERS 2 | from .base_bbox_coder import BaseBBoxCoder 3 | 4 | 5 | @BBOX_CODERS.register_module() 6 | class PseudoBBoxCoder(BaseBBoxCoder): 7 | """Pseudo bounding box coder""" 8 | 9 | def __init__(self, **kwargs): 10 | super(BaseBBoxCoder, self).__init__(**kwargs) 11 | 12 | def encode(self, bboxes, gt_bboxes): 13 | """torch.Tensor: return the given ``bboxes``""" 14 | return gt_bboxes 15 | 16 | def decode(self, bboxes, pred_bboxes): 17 | """torch.Tensor: return the given ``pred_bboxes``""" 18 | return pred_bboxes 19 | -------------------------------------------------------------------------------- /mmdet/models/detectors/fcos.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class FCOS(SingleStageDetector): 7 | """Implementation of `FCOS `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | neck, 12 | bbox_head, 13 | train_cfg=None, 14 | test_cfg=None, 15 | pretrained=None): 16 | super(FCOS, self).__init__(backbone, neck, bbox_head, train_cfg, 17 | test_cfg, pretrained) 18 | -------------------------------------------------------------------------------- /mmdet/models/detectors/fsaf.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class FSAF(SingleStageDetector): 7 | """Implementation of `FSAF `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | neck, 12 | bbox_head, 13 | train_cfg=None, 14 | test_cfg=None, 15 | pretrained=None): 16 | super(FSAF, self).__init__(backbone, neck, bbox_head, train_cfg, 17 | test_cfg, pretrained) 18 | -------------------------------------------------------------------------------- /mmdet/ops/dcn/__init__.py: -------------------------------------------------------------------------------- 1 | from .deform_conv import (DeformConv, DeformConvPack, ModulatedDeformConv, 2 | ModulatedDeformConvPack, deform_conv, 3 | modulated_deform_conv) 4 | from .deform_pool import (DeformRoIPooling, DeformRoIPoolingPack, 5 | ModulatedDeformRoIPoolingPack, deform_roi_pooling) 6 | 7 | __all__ = [ 8 | 'DeformConv', 'DeformConvPack', 'ModulatedDeformConv', 9 | 'ModulatedDeformConvPack', 'DeformRoIPooling', 'DeformRoIPoolingPack', 10 | 'ModulatedDeformRoIPoolingPack', 'deform_conv', 'modulated_deform_conv', 11 | 'deform_roi_pooling' 12 | ] 13 | -------------------------------------------------------------------------------- /mmdet/core/post_processing/__init__.py: -------------------------------------------------------------------------------- 1 | from .bbox_nms import multiclass_nms 2 | from .merge_augs import (merge_aug_bboxes, merge_aug_masks, 3 | merge_aug_proposals, merge_aug_scores) 4 | 5 | from .obb import (multiclass_arb_nms, merge_rotate_aug_proposals, merge_rotate_aug_hbb, 6 | merge_rotate_aug_obb, merge_rotate_aug_arb) 7 | 8 | __all__ = [ 9 | 'multiclass_nms', 'merge_aug_proposals', 'merge_aug_bboxes', 10 | 'merge_aug_scores', 'merge_aug_masks', 11 | 'multiclass_arb_nms', 'merge_rotate_aug_proposals', 'merge_rotate_aug_hbb', 12 | 'merge_rotate_aug_obb', 'merge_rotate_aug_arb' 13 | ] 14 | -------------------------------------------------------------------------------- /mmdet/models/detectors/fovea.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class FOVEA(SingleStageDetector): 7 | """Implementation of `FoveaBox `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | neck, 12 | bbox_head, 13 | train_cfg=None, 14 | test_cfg=None, 15 | pretrained=None): 16 | super(FOVEA, self).__init__(backbone, neck, bbox_head, train_cfg, 17 | test_cfg, pretrained) 18 | -------------------------------------------------------------------------------- /mmdet/models/detectors/retinanet.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class RetinaNet(SingleStageDetector): 7 | """Implementation of `RetinaNet `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | neck, 12 | bbox_head, 13 | train_cfg=None, 14 | test_cfg=None, 15 | pretrained=None): 16 | super(RetinaNet, self).__init__(backbone, neck, bbox_head, train_cfg, 17 | test_cfg, pretrained) 18 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/roi_extractors/__init__.py: -------------------------------------------------------------------------------- 1 | from .generic_roi_extractor import GenericRoIExtractor 2 | from .single_level_roi_extractor import SingleRoIExtractor 3 | from .single_level_roi_extractor_more import SingleRoIExtractor_more 4 | 5 | from .obb.obb_single_level_roi_extractor import OBBSingleRoIExtractor 6 | from .obb.obb_single_level_roi_extractor_more import OBBSingleRoIExtractor_more 7 | from .obb.hbb_select_level_roi_extractor import HBBSelectLVLRoIExtractor 8 | 9 | __all__ = [ 10 | 'SingleRoIExtractor', 11 | 'GenericRoIExtractor', 12 | 'SingleRoIExtractor', 13 | 14 | 'OBBSingleRoIExtractor', 15 | 'OBBSingleRoIExtractor_more', 16 | ] 17 | -------------------------------------------------------------------------------- /mmdet/core/bbox/builder.py: -------------------------------------------------------------------------------- 1 | from mmcv.utils import Registry, build_from_cfg 2 | 3 | BBOX_ASSIGNERS = Registry('bbox_assigner') 4 | BBOX_SAMPLERS = Registry('bbox_sampler') 5 | BBOX_CODERS = Registry('bbox_coder') 6 | 7 | 8 | def build_assigner(cfg, **default_args): 9 | """Builder of box assigner""" 10 | return build_from_cfg(cfg, BBOX_ASSIGNERS, default_args) 11 | 12 | 13 | def build_sampler(cfg, **default_args): 14 | """Builder of box sampler""" 15 | return build_from_cfg(cfg, BBOX_SAMPLERS, default_args) 16 | 17 | 18 | def build_bbox_coder(cfg, **default_args): 19 | """Builder of box coder""" 20 | return build_from_cfg(cfg, BBOX_CODERS, default_args) 21 | -------------------------------------------------------------------------------- /tools/slurm_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | JOB_NAME=$2 7 | CONFIG=$3 8 | CHECKPOINT=$4 9 | GPUS=${GPUS:-8} 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | PY_ARGS=${@:5} 13 | SRUN_ARGS=${SRUN_ARGS:-""} 14 | 15 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 16 | srun -p ${PARTITION} \ 17 | --job-name=${JOB_NAME} \ 18 | --gres=gpu:${GPUS_PER_NODE} \ 19 | --ntasks=${GPUS} \ 20 | --ntasks-per-node=${GPUS_PER_NODE} \ 21 | --cpus-per-task=${CPUS_PER_TASK} \ 22 | --kill-on-bad-exit=1 \ 23 | ${SRUN_ARGS} \ 24 | python -u tools/test.py ${CONFIG} ${CHECKPOINT} --launcher="slurm" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /tools/slurm_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | JOB_NAME=$2 7 | CONFIG=$3 8 | WORK_DIR=$4 9 | GPUS=${GPUS:-8} 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | SRUN_ARGS=${SRUN_ARGS:-""} 13 | PY_ARGS=${@:5} 14 | 15 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 16 | srun -p ${PARTITION} \ 17 | --job-name=${JOB_NAME} \ 18 | --gres=gpu:${GPUS_PER_NODE} \ 19 | --ntasks=${GPUS} \ 20 | --ntasks-per-node=${GPUS_PER_NODE} \ 21 | --cpus-per-task=${CPUS_PER_TASK} \ 22 | --kill-on-bad-exit=1 \ 23 | ${SRUN_ARGS} \ 24 | python -u tools/train.py ${CONFIG} --work-dir=${WORK_DIR} --launcher="slurm" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/retinanet_obb.py: -------------------------------------------------------------------------------- 1 | from mmdet.models.builder import DETECTORS 2 | from .obb_single_stage import OBBSingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class RetinaNetOBB(OBBSingleStageDetector): 7 | """Implementation of `RetinaNet `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | neck, 12 | bbox_head, 13 | train_cfg=None, 14 | test_cfg=None, 15 | pretrained=None): 16 | super(RetinaNetOBB, self).__init__(backbone, neck, bbox_head, train_cfg, 17 | test_cfg, pretrained) 18 | -------------------------------------------------------------------------------- /mmdet/models/detectors/nasfcos.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .single_stage import SingleStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class NASFCOS(SingleStageDetector): 7 | """NAS-FCOS: Fast Neural Architecture Search for Object Detection. 8 | 9 | https://arxiv.org/abs/1906.0442 10 | """ 11 | 12 | def __init__(self, 13 | backbone, 14 | neck, 15 | bbox_head, 16 | train_cfg=None, 17 | test_cfg=None, 18 | pretrained=None): 19 | super(NASFCOS, self).__init__(backbone, neck, bbox_head, train_cfg, 20 | test_cfg, pretrained) 21 | -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/gliding_vertex.py: -------------------------------------------------------------------------------- 1 | from mmdet.models.builder import DETECTORS 2 | from .obb_two_stage import OBBTwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class GlidingVertex(OBBTwoStageDetector): 7 | 8 | def __init__(self, 9 | backbone, 10 | rpn_head, 11 | roi_head, 12 | train_cfg, 13 | test_cfg, 14 | neck=None, 15 | pretrained=None): 16 | super(GlidingVertex, self).__init__( 17 | backbone=backbone, 18 | neck=neck, 19 | rpn_head=rpn_head, 20 | roi_head=roi_head, 21 | train_cfg=train_cfg, 22 | test_cfg=test_cfg, 23 | pretrained=pretrained) 24 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/combined_sampler.py: -------------------------------------------------------------------------------- 1 | from ..builder import BBOX_SAMPLERS, build_sampler 2 | from .base_sampler import BaseSampler 3 | 4 | 5 | @BBOX_SAMPLERS.register_module() 6 | class CombinedSampler(BaseSampler): 7 | """A sampler that combines positive sampler and negative sampler""" 8 | 9 | def __init__(self, pos_sampler, neg_sampler, **kwargs): 10 | super(CombinedSampler, self).__init__(**kwargs) 11 | self.pos_sampler = build_sampler(pos_sampler, **kwargs) 12 | self.neg_sampler = build_sampler(neg_sampler, **kwargs) 13 | 14 | def _sample_pos(self, **kwargs): 15 | """Sample positive samples""" 16 | raise NotImplementedError 17 | 18 | def _sample_neg(self, **kwargs): 19 | """Sample negative samples""" 20 | raise NotImplementedError 21 | -------------------------------------------------------------------------------- /mmdet/models/detectors/mask_rcnn.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class MaskRCNN(TwoStageDetector): 7 | """Implementation of `Mask R-CNN `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | rpn_head, 12 | roi_head, 13 | train_cfg, 14 | test_cfg, 15 | neck=None, 16 | pretrained=None): 17 | super(MaskRCNN, self).__init__( 18 | backbone=backbone, 19 | neck=neck, 20 | rpn_head=rpn_head, 21 | roi_head=roi_head, 22 | train_cfg=train_cfg, 23 | test_cfg=test_cfg, 24 | pretrained=pretrained) 25 | -------------------------------------------------------------------------------- /mmdet/models/detectors/faster_rcnn.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class FasterRCNN(TwoStageDetector): 7 | """Implementation of `Faster R-CNN `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | rpn_head, 12 | roi_head, 13 | train_cfg, 14 | test_cfg, 15 | neck=None, 16 | pretrained=None): 17 | super(FasterRCNN, self).__init__( 18 | backbone=backbone, 19 | neck=neck, 20 | rpn_head=rpn_head, 21 | roi_head=roi_head, 22 | train_cfg=train_cfg, 23 | test_cfg=test_cfg, 24 | pretrained=pretrained) 25 | -------------------------------------------------------------------------------- /mmdet/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .backbones import * # noqa: F401,F403 2 | from .builder import (BACKBONES, DETECTORS, HEADS, LOSSES, NECKS, 3 | ROI_EXTRACTORS, SHARED_HEADS, build_backbone, 4 | build_detector, build_head, build_loss, build_neck, 5 | build_roi_extractor, build_shared_head) 6 | from .dense_heads import * # noqa: F401,F403 7 | from .detectors import * # noqa: F401,F403 8 | from .losses import * # noqa: F401,F403 9 | from .necks import * # noqa: F401,F403 10 | from .roi_heads import * # noqa: F401,F403 11 | 12 | __all__ = [ 13 | 'BACKBONES', 'NECKS', 'ROI_EXTRACTORS', 'SHARED_HEADS', 'HEADS', 'LOSSES', 14 | 'DETECTORS', 'build_backbone', 'build_neck', 'build_roi_extractor', 15 | 'build_shared_head', 'build_head', 'build_loss', 'build_detector' 16 | ] 17 | -------------------------------------------------------------------------------- /mmdet/models/detectors/mask_scoring_rcnn.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class MaskScoringRCNN(TwoStageDetector): 7 | """Mask Scoring RCNN. 8 | 9 | https://arxiv.org/abs/1903.00241 10 | """ 11 | 12 | def __init__(self, 13 | backbone, 14 | rpn_head, 15 | roi_head, 16 | train_cfg, 17 | test_cfg, 18 | neck=None, 19 | pretrained=None): 20 | super(MaskScoringRCNN, self).__init__( 21 | backbone=backbone, 22 | neck=neck, 23 | rpn_head=rpn_head, 24 | roi_head=roi_head, 25 | train_cfg=train_cfg, 26 | test_cfg=test_cfg, 27 | pretrained=pretrained) 28 | -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/faster_rcnn_obb.py: -------------------------------------------------------------------------------- 1 | from mmdet.models.builder import DETECTORS 2 | from .obb_two_stage import OBBTwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class FasterRCNNOBB(OBBTwoStageDetector): 7 | """Implementation of `Faster R-CNN `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | rpn_head, 12 | roi_head, 13 | train_cfg, 14 | test_cfg, 15 | neck=None, 16 | pretrained=None): 17 | super(FasterRCNNOBB, self).__init__( 18 | backbone=backbone, 19 | neck=neck, 20 | rpn_head=rpn_head, 21 | roi_head=roi_head, 22 | train_cfg=train_cfg, 23 | test_cfg=test_cfg, 24 | pretrained=pretrained) 25 | -------------------------------------------------------------------------------- /mmdet/ops/convex/src/convex_ext.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #ifdef WITH_CUDA 5 | at::Tensor convex_sort_cuda( 6 | const at::Tensor& pts, const at::Tensor& masks, const bool circular); 7 | #endif 8 | 9 | at::Tensor convex_sort_cpu( 10 | const at::Tensor& pts, const at::Tensor& masks, const bool circular); 11 | 12 | 13 | at::Tensor convex_sort( 14 | const at::Tensor& pts, const at::Tensor& masks, const bool circular) { 15 | if (pts.device().is_cuda()) { 16 | #ifdef WITH_CUDA 17 | return convex_sort_cuda(pts, masks, circular); 18 | #else 19 | AT_ERROR("sort_vert is not compiled with GPU support"); 20 | #endif 21 | } 22 | return convex_sort_cpu(pts, masks, circular); 23 | } 24 | 25 | 26 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 27 | m.def("convex_sort", &convex_sort, "select the convex points and sort them"); 28 | } 29 | -------------------------------------------------------------------------------- /mmdet/models/show.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy 3 | 4 | img = cv2.imread('/disk0/evannnnnnn/home/Datasets/DOTA/train/images/P0896.png') 5 | # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) 6 | # img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) 7 | lbl = open('/disk0/evannnnnnn/home/Datasets/DOTA/train/labelTxt/P0896.txt') 8 | lbl.readline() 9 | lbl.readline() 10 | 11 | for l in lbl: 12 | x1, y1, x2, y2, x3, y3, x4, y4, c, _ = l.split() 13 | x1, y1, x2, y2, x3, y3, x4, y4 =\ 14 | int(x1), int(y1), int(x2), int(y2), int(x3), int(y3), int(x4), int(y4) 15 | b = [[x1, y1], [x2, y2], [x3, y3], [x4, y4]] 16 | if c == 'harbor': color = (255, 255, 0) 17 | if c == 'ship': color = (0, 255, 255) 18 | cv2.polylines(img, [numpy.array(b)], isClosed=True, color=color, thickness=2) 19 | 20 | cv2.imwrite('/disk0/evannnnnnn/home/P0896.jpg', img) 21 | cv2.imshow('img', img) 22 | # cv2.waitKey() -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/roi_transformer.py: -------------------------------------------------------------------------------- 1 | import mmcv 2 | import warnings 3 | import numpy as np 4 | import BboxToolkit as bt 5 | 6 | from mmdet.models.builder import DETECTORS 7 | from .obb_two_stage import OBBTwoStageDetector 8 | 9 | 10 | @DETECTORS.register_module() 11 | class RoITransformer(OBBTwoStageDetector): 12 | 13 | def __init__(self, 14 | backbone, 15 | neck=None, 16 | rpn_head=None, 17 | roi_head=None, 18 | train_cfg=None, 19 | test_cfg=None, 20 | pretrained=None): 21 | super(RoITransformer, self).__init__( 22 | backbone=backbone, 23 | neck=neck, 24 | rpn_head=rpn_head, 25 | roi_head=roi_head, 26 | train_cfg=train_cfg, 27 | test_cfg=test_cfg, 28 | pretrained=pretrained) 29 | -------------------------------------------------------------------------------- /mmdet/models/detectors/point_rend.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class PointRend(TwoStageDetector): 7 | """PointRend: Image Segmentation as Rendering 8 | 9 | This detector is the implementation of 10 | `PointRend `_. 11 | 12 | """ 13 | 14 | def __init__(self, 15 | backbone, 16 | rpn_head, 17 | roi_head, 18 | train_cfg, 19 | test_cfg, 20 | neck=None, 21 | pretrained=None): 22 | super(PointRend, self).__init__( 23 | backbone=backbone, 24 | neck=neck, 25 | rpn_head=rpn_head, 26 | roi_head=roi_head, 27 | train_cfg=train_cfg, 28 | test_cfg=test_cfg, 29 | pretrained=pretrained) 30 | -------------------------------------------------------------------------------- /mmdet/models/detectors/grid_rcnn.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class GridRCNN(TwoStageDetector): 7 | """Grid R-CNN. 8 | 9 | This detector is the implementation of: 10 | - Grid R-CNN (https://arxiv.org/abs/1811.12030) 11 | - Grid R-CNN Plus: Faster and Better (https://arxiv.org/abs/1906.05688) 12 | """ 13 | 14 | def __init__(self, 15 | backbone, 16 | rpn_head, 17 | roi_head, 18 | train_cfg, 19 | test_cfg, 20 | neck=None, 21 | pretrained=None): 22 | super(GridRCNN, self).__init__( 23 | backbone=backbone, 24 | neck=neck, 25 | rpn_head=rpn_head, 26 | roi_head=roi_head, 27 | train_cfg=train_cfg, 28 | test_cfg=test_cfg, 29 | pretrained=pretrained) 30 | -------------------------------------------------------------------------------- /mmdet/core/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | from .class_names import (cityscapes_classes, coco_classes, dataset_aliases, 2 | get_classes, imagenet_det_classes, 3 | imagenet_vid_classes, voc_classes) 4 | from .eval_hooks import DistEvalHook, EvalHook 5 | from .mean_ap import average_precision, eval_map, print_map_summary 6 | from .recall import (eval_recalls, plot_iou_recall, plot_num_recall, 7 | print_recall_summary) 8 | from .obb.obb_mean_ap import eval_arb_map 9 | from .obb.obb_recall import eval_arb_recalls 10 | 11 | __all__ = [ 12 | 'voc_classes', 'imagenet_det_classes', 'imagenet_vid_classes', 13 | 'coco_classes', 'cityscapes_classes', 'dataset_aliases', 'get_classes', 14 | 'DistEvalHook', 'EvalHook', 'average_precision', 'eval_map', 15 | 'print_map_summary', 'eval_recalls', 'print_recall_summary', 16 | 'plot_num_recall', 'plot_iou_recall', 17 | 'eval_arb_map', 'eval_arb_recalls' 18 | ] 19 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_sampler import BaseSampler 2 | from .combined_sampler import CombinedSampler 3 | from .instance_balanced_pos_sampler import InstanceBalancedPosSampler 4 | from .iou_balanced_neg_sampler import IoUBalancedNegSampler 5 | from .ohem_sampler import OHEMSampler 6 | from .pseudo_sampler import PseudoSampler 7 | from .random_sampler import RandomSampler 8 | from .sampling_result import SamplingResult 9 | from .score_hlr_sampler import ScoreHLRSampler 10 | 11 | from .obb import (OBBSamplingResult, OBBBaseSampler, OBBRandomSampler, 12 | OBBOHEMSampler, OBBRandomSamplerXYRAB) 13 | 14 | __all__ = [ 15 | 'BaseSampler', 'PseudoSampler', 'RandomSampler', 16 | 'InstanceBalancedPosSampler', 'IoUBalancedNegSampler', 'CombinedSampler', 17 | 'OHEMSampler', 'SamplingResult', 'ScoreHLRSampler', 18 | 19 | 'OBBSamplingResult', 'OBBBaseSampler', 'OBBRandomSampler', 20 | 'OBBOHEMSampler', 'OBBRandomSamplerXYRAB' 21 | ] 22 | -------------------------------------------------------------------------------- /mmdet/core/bbox/coder/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_bbox_coder import BaseBBoxCoder 2 | from .delta_xywh_bbox_coder import DeltaXYWHBBoxCoder 3 | from .legacy_delta_xywh_bbox_coder import LegacyDeltaXYWHBBoxCoder 4 | from .pseudo_bbox_coder import PseudoBBoxCoder 5 | from .tblr_bbox_coder import TBLRBBoxCoder 6 | 7 | from .obb.obb2obb_delta_xywht_coder import OBB2OBBDeltaXYWHTCoder 8 | from .obb.hbb2obb_delta_xywht_coder import HBB2OBBDeltaXYWHTCoder 9 | from .obb.hbb_qp_coder import HBBQPCoder 10 | from .obb.obb_qp_coder import OBBQPCoder 11 | 12 | from .obb.hbb_slide_point_coder_xyrab import HBBSlidePointCoderXYRAB 13 | from .obb.obb2obb_delta_xywht_coder_xyrab import OBB2OBBDeltaXYWHTCoderXYRAB 14 | 15 | __all__ = [ 16 | 'BaseBBoxCoder', 'PseudoBBoxCoder', 'DeltaXYWHBBoxCoder', 17 | 'LegacyDeltaXYWHBBoxCoder', 'TBLRBBoxCoder', 18 | 'OBB2OBBDeltaXYWHTCoder', 'HBB2OBBDeltaXYWHTCoder', 19 | 'HBBQPCoder', 'OBBQPCoder', 'HBBSlidePointCoderXYRAB', 'OBB2OBBDeltaXYWHTCoderXYRAB' 20 | ] 21 | -------------------------------------------------------------------------------- /mmdet/ops/roi_align/gradcheck.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import numpy as np 5 | import torch 6 | from torch.autograd import gradcheck 7 | 8 | sys.path.append(osp.abspath(osp.join(__file__, '../../'))) 9 | from roi_align import RoIAlign # noqa: E402, isort:skip 10 | 11 | feat_size = 15 12 | spatial_scale = 1.0 / 8 13 | img_size = feat_size / spatial_scale 14 | num_imgs = 2 15 | num_rois = 20 16 | 17 | batch_ind = np.random.randint(num_imgs, size=(num_rois, 1)) 18 | rois = np.random.rand(num_rois, 4) * img_size * 0.5 19 | rois[:, 2:] += img_size * 0.5 20 | rois = np.hstack((batch_ind, rois)) 21 | 22 | feat = torch.randn( 23 | num_imgs, 16, feat_size, feat_size, requires_grad=True, device='cuda:0') 24 | rois = torch.from_numpy(rois).float().cuda() 25 | inputs = (feat, rois) 26 | print('Gradcheck for roi align...') 27 | test = gradcheck(RoIAlign(3, spatial_scale), inputs, atol=1e-3, eps=1e-3) 28 | print(test) 29 | test = gradcheck(RoIAlign(3, spatial_scale, 2), inputs, atol=1e-3, eps=1e-3) 30 | print(test) 31 | -------------------------------------------------------------------------------- /mmdet/datasets/samplers/distributed_sampler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DistributedSampler as _DistributedSampler 3 | 4 | 5 | class DistributedSampler(_DistributedSampler): 6 | 7 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): 8 | super().__init__(dataset, num_replicas=num_replicas, rank=rank) 9 | self.shuffle = shuffle 10 | 11 | def __iter__(self): 12 | # deterministically shuffle based on epoch 13 | if self.shuffle: 14 | g = torch.Generator() 15 | g.manual_seed(self.epoch) 16 | indices = torch.randperm(len(self.dataset), generator=g).tolist() 17 | else: 18 | indices = torch.arange(len(self.dataset)).tolist() 19 | 20 | # add extra samples to make it evenly divisible 21 | indices += indices[:(self.total_size - len(indices))] 22 | assert len(indices) == self.total_size 23 | 24 | # subsample 25 | indices = indices[self.rank:self.total_size:self.num_replicas] 26 | assert len(indices) == self.num_samples 27 | 28 | return iter(indices) 29 | -------------------------------------------------------------------------------- /mmdet/core/fp16/utils.py: -------------------------------------------------------------------------------- 1 | from collections import abc 2 | 3 | import numpy as np 4 | import torch 5 | 6 | 7 | def cast_tensor_type(inputs, src_type, dst_type): 8 | """Recursively convert Tensor in inputs from src_type to dst_type. 9 | 10 | Args: 11 | inputs: Inputs that to be casted. 12 | src_type (torch.dtype): Source type.. 13 | dst_type (torch.dtype): Destination type. 14 | 15 | Returns: 16 | The same type with inputs, but all contained Tensors have been cast. 17 | """ 18 | if isinstance(inputs, torch.Tensor): 19 | return inputs.to(dst_type) 20 | elif isinstance(inputs, str): 21 | return inputs 22 | elif isinstance(inputs, np.ndarray): 23 | return inputs 24 | elif isinstance(inputs, abc.Mapping): 25 | return type(inputs)({ 26 | k: cast_tensor_type(v, src_type, dst_type) 27 | for k, v in inputs.items() 28 | }) 29 | elif isinstance(inputs, abc.Iterable): 30 | return type(inputs)( 31 | cast_tensor_type(item, src_type, dst_type) for item in inputs) 32 | else: 33 | return inputs 34 | -------------------------------------------------------------------------------- /mmdet/ops/box_iou_rotated/src/box_iou_rotated_ext.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #ifdef WITH_CUDA 5 | at::Tensor box_iou_rotated_cuda( 6 | const at::Tensor& boxes1, 7 | const at::Tensor& boxes2, 8 | const bool iou_or_iof); 9 | #endif 10 | 11 | at::Tensor box_iou_rotated_cpu( 12 | const at::Tensor& boxes1, 13 | const at::Tensor& boxes2, 14 | const bool iou_or_iof); 15 | 16 | 17 | inline at::Tensor box_iou_rotated( 18 | const at::Tensor& boxes1, 19 | const at::Tensor& boxes2, 20 | const bool iou_or_iof) { 21 | assert(boxes1.device().is_cuda() == boxes2.device().is_cuda()); 22 | if (boxes1.device().is_cuda()) { 23 | #ifdef WITH_CUDA 24 | return box_iou_rotated_cuda( 25 | boxes1.contiguous(), 26 | boxes2.contiguous(), 27 | iou_or_iof); 28 | #else 29 | AT_ERROR("Not compiled with GPU support"); 30 | #endif 31 | } 32 | 33 | return box_iou_rotated_cpu( 34 | boxes1.contiguous(), 35 | boxes2.contiguous(), 36 | iou_or_iof); 37 | } 38 | 39 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 40 | m.def("overlaps", box_iou_rotated, "calculate iou or iof of two group boxes"); 41 | } 42 | -------------------------------------------------------------------------------- /mmdet/ops/carafe/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | 5 | NVCC_ARGS = [ 6 | '-D__CUDA_NO_HALF_OPERATORS__', 7 | '-D__CUDA_NO_HALF_CONVERSIONS__', 8 | '-D__CUDA_NO_HALF2_OPERATORS__', 9 | ] 10 | 11 | setup( 12 | name='carafe', 13 | ext_modules=[ 14 | CUDAExtension( 15 | 'carafe_ext', [ 16 | 'src/cuda/carafe_cuda.cpp', 'src/cuda/carafe_cuda_kernel.cu', 17 | 'src/carafe_ext.cpp' 18 | ], 19 | define_macros=[('WITH_CUDA', None)], 20 | extra_compile_args={ 21 | 'cxx': [], 22 | 'nvcc': NVCC_ARGS 23 | }), 24 | CUDAExtension( 25 | 'carafe_naive_ext', [ 26 | 'src/cuda/carafe_naive_cuda.cpp', 27 | 'src/cuda/carafe_naive_cuda_kernel.cu', 28 | 'src/carafe_naive_ext.cpp' 29 | ], 30 | define_macros=[('WITH_CUDA', None)], 31 | extra_compile_args={ 32 | 'cxx': [], 33 | 'nvcc': NVCC_ARGS 34 | }) 35 | ], 36 | cmdclass={'build_ext': BuildExtension}) 37 | -------------------------------------------------------------------------------- /mmdet/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset 2 | from .cityscapes import CityscapesDataset 3 | from .coco import CocoDataset 4 | from .custom import CustomDataset 5 | from .dataset_wrappers import (ClassBalancedDataset, ConcatDataset, 6 | RepeatDataset) 7 | from .deepfashion import DeepFashionDataset 8 | from .lvis import LVISDataset 9 | from .samplers import DistributedGroupSampler, DistributedSampler, GroupSampler 10 | from .voc import VOCDataset 11 | from .wider_face import WIDERFaceDataset 12 | from .xml_style import XMLDataset 13 | 14 | from .obb.dota import DOTADataset 15 | from .obb.dior import DIORDataset 16 | from .obb.hrsc import HRSCDataset 17 | 18 | __all__ = [ 19 | 'CustomDataset', 'XMLDataset', 'CocoDataset', 'DeepFashionDataset', 20 | 'VOCDataset', 'CityscapesDataset', 'LVISDataset', 'GroupSampler', 21 | 'CustomDataset', 'XMLDataset', 'CocoDataset', 'VOCDataset', 22 | 'CityscapesDataset', 'LVISDataset', 'DeepFashionDataset', 'GroupSampler', 23 | 'DistributedGroupSampler', 'DistributedSampler', 'build_dataloader', 24 | 'ConcatDataset', 'RepeatDataset', 'ClassBalancedDataset', 25 | 'WIDERFaceDataset', 'DATASETS', 'PIPELINES', 'build_dataset' 26 | ] 27 | -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/obb_test_mixins.py: -------------------------------------------------------------------------------- 1 | from mmdet.core import merge_rotate_aug_proposals 2 | 3 | 4 | class RotateAugRPNTestMixin(object): 5 | 6 | def rotate_aug_test_rpn(self, feats, img_metas): 7 | samples_per_gpu = len(img_metas[0]) 8 | aug_proposals = [[] for _ in range(samples_per_gpu)] 9 | for x, img_meta in zip(feats, img_metas): 10 | proposal_list = self.rpn_head.simple_test_rpn(x, img_meta) 11 | for i, proposals in enumerate(proposal_list): 12 | aug_proposals[i].append(proposals) 13 | # reorganize the order of 'img_metas' to match the dimensions 14 | # of 'aug_proposals' 15 | aug_img_metas = [] 16 | for i in range(samples_per_gpu): 17 | aug_img_meta = [] 18 | for j in range(len(img_metas)): 19 | aug_img_meta.append(img_metas[j][i]) 20 | aug_img_metas.append(aug_img_meta) 21 | # after merging, proposals will be rescaled to the original image size 22 | merged_proposals = [ 23 | merge_rotate_aug_proposals(proposals, aug_img_meta, 24 | self.rpn_head.test_cfg) 25 | for proposals, aug_img_meta in zip(aug_proposals, aug_img_metas) 26 | ] 27 | return merged_proposals 28 | -------------------------------------------------------------------------------- /mmdet/ops/box_iou_rotated/src/box_iou_rotated_cpu.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 2 | #include 3 | #include "box_iou_rotated_utils.h" 4 | 5 | 6 | template 7 | void box_iou_rotated_cpu_kernel( 8 | const at::Tensor& boxes1, 9 | const at::Tensor& boxes2, 10 | const bool iou_or_iof, 11 | at::Tensor& ious) { 12 | auto num_boxes1 = boxes1.size(0); 13 | auto num_boxes2 = boxes2.size(0); 14 | 15 | for (int i = 0; i < num_boxes1; i++) { 16 | for (int j = 0; j < num_boxes2; j++) { 17 | ious[i * num_boxes2 + j] = single_box_iou_rotated( 18 | boxes1[i].data_ptr(), boxes2[j].data_ptr(), iou_or_iof); 19 | } 20 | } 21 | } 22 | 23 | at::Tensor box_iou_rotated_cpu( 24 | // input must be contiguous: 25 | const at::Tensor& boxes1, 26 | const at::Tensor& boxes2, 27 | const bool iou_or_iof) { 28 | auto num_boxes1 = boxes1.size(0); 29 | auto num_boxes2 = boxes2.size(0); 30 | at::Tensor ious = 31 | at::empty({num_boxes1 * num_boxes2}, boxes1.options().dtype(at::kFloat)); 32 | 33 | box_iou_rotated_cpu_kernel(boxes1, boxes2, iou_or_iof, ious); 34 | 35 | // reshape from 1d array to 2d array 36 | auto shape = std::vector{num_boxes1, num_boxes2}; 37 | return ious.reshape(shape); 38 | } 39 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/double_roi_head.py: -------------------------------------------------------------------------------- 1 | from ..builder import HEADS 2 | from .standard_roi_head import StandardRoIHead 3 | 4 | 5 | @HEADS.register_module() 6 | class DoubleHeadRoIHead(StandardRoIHead): 7 | """RoI head for Double Head RCNN 8 | 9 | https://arxiv.org/abs/1904.06493 10 | """ 11 | 12 | def __init__(self, reg_roi_scale_factor, **kwargs): 13 | super(DoubleHeadRoIHead, self).__init__(**kwargs) 14 | self.reg_roi_scale_factor = reg_roi_scale_factor 15 | 16 | def _bbox_forward(self, x, rois): 17 | """Box head forward function used in both training and testing time""" 18 | bbox_cls_feats = self.bbox_roi_extractor( 19 | x[:self.bbox_roi_extractor.num_inputs], rois) 20 | bbox_reg_feats = self.bbox_roi_extractor( 21 | x[:self.bbox_roi_extractor.num_inputs], 22 | rois, 23 | roi_scale_factor=self.reg_roi_scale_factor) 24 | if self.with_shared_head: 25 | bbox_cls_feats = self.shared_head(bbox_cls_feats) 26 | bbox_reg_feats = self.shared_head(bbox_reg_feats) 27 | cls_score, bbox_pred = self.bbox_head(bbox_cls_feats, bbox_reg_feats) 28 | 29 | bbox_results = dict( 30 | cls_score=cls_score, 31 | bbox_pred=bbox_pred, 32 | bbox_feats=bbox_cls_feats) 33 | return bbox_results 34 | -------------------------------------------------------------------------------- /mmdet/models/detectors/cascade_rcnn.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class CascadeRCNN(TwoStageDetector): 7 | """Implementation of `Cascade R-CNN and Cascade Mask R-CNN 8 | `_""" 9 | 10 | def __init__(self, 11 | backbone, 12 | neck=None, 13 | rpn_head=None, 14 | roi_head=None, 15 | train_cfg=None, 16 | test_cfg=None, 17 | pretrained=None): 18 | super(CascadeRCNN, self).__init__( 19 | backbone=backbone, 20 | neck=neck, 21 | rpn_head=rpn_head, 22 | roi_head=roi_head, 23 | train_cfg=train_cfg, 24 | test_cfg=test_cfg, 25 | pretrained=pretrained) 26 | 27 | def show_result(self, data, result, **kwargs): 28 | """Show prediction results of the detector""" 29 | if self.with_mask: 30 | ms_bbox_result, ms_segm_result = result 31 | if isinstance(ms_bbox_result, dict): 32 | result = (ms_bbox_result['ensemble'], 33 | ms_segm_result['ensemble']) 34 | else: 35 | if isinstance(result, dict): 36 | result = result['ensemble'] 37 | return super(CascadeRCNN, self).show_result(data, result, **kwargs) 38 | -------------------------------------------------------------------------------- /mmdet/utils/profiling.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import sys 3 | import time 4 | 5 | import torch 6 | 7 | if sys.version_info >= (3, 7): 8 | 9 | @contextlib.contextmanager 10 | def profile_time(trace_name, 11 | name, 12 | enabled=True, 13 | stream=None, 14 | end_stream=None): 15 | """Print time spent by CPU and GPU. 16 | 17 | Useful as a temporary context manager to find sweet spots of 18 | code suitable for async implementation. 19 | 20 | """ 21 | if (not enabled) or not torch.cuda.is_available(): 22 | yield 23 | return 24 | stream = stream if stream else torch.cuda.current_stream() 25 | end_stream = end_stream if end_stream else stream 26 | start = torch.cuda.Event(enable_timing=True) 27 | end = torch.cuda.Event(enable_timing=True) 28 | stream.record_event(start) 29 | try: 30 | cpu_start = time.monotonic() 31 | yield 32 | finally: 33 | cpu_end = time.monotonic() 34 | end_stream.record_event(end) 35 | end.synchronize() 36 | cpu_time = (cpu_end - cpu_start) * 1000 37 | gpu_time = start.elapsed_time(end) 38 | msg = f'{trace_name} {name} cpu_time {cpu_time:.2f} ms ' 39 | msg += f'gpu_time {gpu_time:.2f} ms stream {stream}' 40 | print(msg, end_stream) 41 | -------------------------------------------------------------------------------- /mmdet/ops/plugin.py: -------------------------------------------------------------------------------- 1 | from mmcv.cnn import ConvModule 2 | 3 | from .context_block import ContextBlock 4 | from .generalized_attention import GeneralizedAttention 5 | from .non_local import NonLocal2D 6 | 7 | plugin_cfg = { 8 | # format: layer_type: (abbreviation, module) 9 | 'ContextBlock': ('context_block', ContextBlock), 10 | 'GeneralizedAttention': ('gen_attention_block', GeneralizedAttention), 11 | 'NonLocal2D': ('nonlocal_block', NonLocal2D), 12 | 'ConvModule': ('conv_block', ConvModule), 13 | } 14 | 15 | 16 | def build_plugin_layer(cfg, postfix='', **kwargs): 17 | """ Build plugin layer 18 | 19 | Args: 20 | cfg (None or dict): cfg should contain: 21 | type (str): identify plugin layer type. 22 | layer args: args needed to instantiate a plugin layer. 23 | postfix (int, str): appended into norm abbreviation to 24 | create named layer. 25 | 26 | Returns: 27 | name (str): abbreviation + postfix 28 | layer (nn.Module): created plugin layer 29 | """ 30 | assert isinstance(cfg, dict) and 'type' in cfg 31 | cfg_ = cfg.copy() 32 | 33 | layer_type = cfg_.pop('type') 34 | if layer_type not in plugin_cfg: 35 | raise KeyError(f'Unrecognized plugin type {layer_type}') 36 | else: 37 | abbr, plugin_layer = plugin_cfg[layer_type] 38 | 39 | assert isinstance(postfix, (int, str)) 40 | name = abbr + str(postfix) 41 | 42 | layer = plugin_layer(**kwargs, **cfg_) 43 | 44 | return name, layer 45 | -------------------------------------------------------------------------------- /mmdet/core/anchor/point_generator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .builder import ANCHOR_GENERATORS 4 | 5 | 6 | @ANCHOR_GENERATORS.register_module() 7 | class PointGenerator(object): 8 | 9 | def _meshgrid(self, x, y, row_major=True): 10 | xx = x.repeat(len(y)) 11 | yy = y.view(-1, 1).repeat(1, len(x)).view(-1) 12 | if row_major: 13 | return xx, yy 14 | else: 15 | return yy, xx 16 | 17 | def grid_points(self, featmap_size, stride=16, device='cuda'): 18 | feat_h, feat_w = featmap_size 19 | shift_x = torch.arange(0., feat_w, device=device) * stride 20 | shift_y = torch.arange(0., feat_h, device=device) * stride 21 | shift_xx, shift_yy = self._meshgrid(shift_x, shift_y) 22 | stride = shift_x.new_full((shift_xx.shape[0], ), stride) 23 | shifts = torch.stack([shift_xx, shift_yy, stride], dim=-1) 24 | all_points = shifts.to(device) 25 | return all_points 26 | 27 | def valid_flags(self, featmap_size, valid_size, device='cuda'): 28 | feat_h, feat_w = featmap_size 29 | valid_h, valid_w = valid_size 30 | assert valid_h <= feat_h and valid_w <= feat_w 31 | valid_x = torch.zeros(feat_w, dtype=torch.bool, device=device) 32 | valid_y = torch.zeros(feat_h, dtype=torch.bool, device=device) 33 | valid_x[:valid_w] = 1 34 | valid_y[:valid_h] = 1 35 | valid_xx, valid_yy = self._meshgrid(valid_x, valid_y) 36 | valid = valid_xx & valid_yy 37 | return valid 38 | -------------------------------------------------------------------------------- /mmdet/core/post_processing/obb/obb_nms.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from mmdet.ops.nms_rotated import arb_batched_nms 4 | from mmdet.core.bbox.transforms_obb import get_bbox_dim 5 | 6 | 7 | def multiclass_arb_nms(multi_bboxes, 8 | multi_scores, 9 | score_thr, 10 | nms_cfg, 11 | max_num=-1, 12 | score_factors=None, 13 | bbox_type='hbb'): 14 | bbox_dim = get_bbox_dim(bbox_type) 15 | num_classes = multi_scores.size(1) - 1 16 | # exclude background category 17 | if multi_bboxes.shape[1] > bbox_dim: 18 | bboxes = multi_bboxes.view(multi_scores.size(0), -1, bbox_dim) 19 | else: 20 | bboxes = multi_bboxes[:, None].expand(-1, num_classes, bbox_dim) 21 | scores = multi_scores[:, :-1] 22 | 23 | # filter out boxes with low scores 24 | valid_mask = scores > score_thr 25 | bboxes = bboxes[valid_mask] 26 | if score_factors is not None: 27 | scores = scores * score_factors[:, None] 28 | scores = scores[valid_mask] 29 | labels = valid_mask.nonzero()[:, 1] 30 | 31 | if bboxes.numel() == 0: 32 | bboxes = multi_bboxes.new_zeros((0, bbox_dim+1)) 33 | labels = multi_bboxes.new_zeros((0, ), dtype=torch.long) 34 | return bboxes, labels 35 | 36 | dets, keep = arb_batched_nms(bboxes, scores, labels, nms_cfg) 37 | 38 | if max_num > 0: 39 | dets = dets[:max_num] 40 | keep = keep[:max_num] 41 | 42 | return dets, labels[keep] 43 | -------------------------------------------------------------------------------- /mmdet/models/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from .accuracy import Accuracy, accuracy 2 | from .ae_loss import AssociativeEmbeddingLoss 3 | from .balanced_l1_loss import BalancedL1Loss, balanced_l1_loss 4 | from .cross_entropy_loss import (CrossEntropyLoss, binary_cross_entropy, 5 | cross_entropy, mask_cross_entropy) 6 | from .focal_loss import FocalLoss, sigmoid_focal_loss 7 | from .gaussian_focal_loss import GaussianFocalLoss 8 | from .gfocal_loss import DistributionFocalLoss, QualityFocalLoss 9 | from .ghm_loss import GHMC, GHMR 10 | from .iou_loss import (BoundedIoULoss, GIoULoss, IoULoss, bounded_iou_loss, 11 | iou_loss) 12 | from .mse_loss import MSELoss, mse_loss 13 | from .pisa_loss import carl_loss, isr_p 14 | from .smooth_l1_loss import L1Loss, SmoothL1Loss, l1_loss, smooth_l1_loss 15 | from .utils import reduce_loss, weight_reduce_loss, weighted_loss 16 | from .obb.poly_iou_loss import PolyIoULoss, PolyGIoULoss 17 | __all__ = [ 18 | 'accuracy', 'Accuracy', 'cross_entropy', 'binary_cross_entropy', 19 | 'mask_cross_entropy', 'CrossEntropyLoss', 'sigmoid_focal_loss', 20 | 'FocalLoss', 'smooth_l1_loss', 'SmoothL1Loss', 'balanced_l1_loss', 21 | 'BalancedL1Loss', 'mse_loss', 'MSELoss', 'iou_loss', 'bounded_iou_loss', 22 | 'IoULoss', 'BoundedIoULoss', 'GIoULoss', 'GHMC', 'GHMR', 'reduce_loss', 23 | 'weight_reduce_loss', 'weighted_loss', 'L1Loss', 'l1_loss', 'isr_p', 24 | 'carl_loss', 'AssociativeEmbeddingLoss', 'GaussianFocalLoss', 25 | 'QualityFocalLoss', 'DistributionFocalLoss','PolyIoULoss', 'PolyGIoULoss' 26 | ] 27 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/mask_heads/htc_mask_head.py: -------------------------------------------------------------------------------- 1 | from mmcv.cnn import ConvModule 2 | 3 | from mmdet.models.builder import HEADS 4 | from .fcn_mask_head import FCNMaskHead 5 | 6 | 7 | @HEADS.register_module() 8 | class HTCMaskHead(FCNMaskHead): 9 | 10 | def __init__(self, with_conv_res=True, *args, **kwargs): 11 | super(HTCMaskHead, self).__init__(*args, **kwargs) 12 | self.with_conv_res = with_conv_res 13 | if self.with_conv_res: 14 | self.conv_res = ConvModule( 15 | self.conv_out_channels, 16 | self.conv_out_channels, 17 | 1, 18 | conv_cfg=self.conv_cfg, 19 | norm_cfg=self.norm_cfg) 20 | 21 | def init_weights(self): 22 | super(HTCMaskHead, self).init_weights() 23 | if self.with_conv_res: 24 | self.conv_res.init_weights() 25 | 26 | def forward(self, x, res_feat=None, return_logits=True, return_feat=True): 27 | if res_feat is not None: 28 | assert self.with_conv_res 29 | res_feat = self.conv_res(res_feat) 30 | x = x + res_feat 31 | for conv in self.convs: 32 | x = conv(x) 33 | res_feat = x 34 | outs = [] 35 | if return_logits: 36 | x = self.upsample(x) 37 | if self.upsample_method == 'deconv': 38 | x = self.relu(x) 39 | mask_pred = self.conv_logits(x) 40 | outs.append(mask_pred) 41 | if return_feat: 42 | outs.append(res_feat) 43 | return outs if len(outs) > 1 else outs[0] 44 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/pseudo_sampler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ..builder import BBOX_SAMPLERS 4 | from .base_sampler import BaseSampler 5 | from .sampling_result import SamplingResult 6 | 7 | 8 | @BBOX_SAMPLERS.register_module() 9 | class PseudoSampler(BaseSampler): 10 | """A pseudo sampler that does not do sampling actually.""" 11 | 12 | def __init__(self, **kwargs): 13 | pass 14 | 15 | def _sample_pos(self, **kwargs): 16 | """Sample positive samples""" 17 | raise NotImplementedError 18 | 19 | def _sample_neg(self, **kwargs): 20 | """Sample negative samples""" 21 | raise NotImplementedError 22 | 23 | def sample(self, assign_result, bboxes, gt_bboxes, **kwargs): 24 | """Directly returns the positive and negative indices of samples 25 | 26 | Args: 27 | assign_result (:obj:`AssignResult`): Assigned results 28 | bboxes (torch.Tensor): Bounding boxes 29 | gt_bboxes (torch.Tensor): Ground truth boxes 30 | 31 | Returns: 32 | :obj:`SamplingResult`: sampler results 33 | """ 34 | pos_inds = torch.nonzero( 35 | assign_result.gt_inds > 0, as_tuple=False).squeeze(-1).unique() 36 | neg_inds = torch.nonzero( 37 | assign_result.gt_inds == 0, as_tuple=False).squeeze(-1).unique() 38 | gt_flags = bboxes.new_zeros(bboxes.shape[0], dtype=torch.uint8) 39 | sampling_result = SamplingResult(pos_inds, neg_inds, bboxes, gt_bboxes, 40 | assign_result, gt_flags) 41 | return sampling_result 42 | -------------------------------------------------------------------------------- /mmdet/models/detectors/__init__.py: -------------------------------------------------------------------------------- 1 | from .atss import ATSS 2 | from .base import BaseDetector 3 | from .cascade_rcnn import CascadeRCNN 4 | from .fast_rcnn import FastRCNN 5 | from .faster_rcnn import FasterRCNN 6 | from .fcos import FCOS 7 | from .fovea import FOVEA 8 | from .fsaf import FSAF 9 | from .gfl import GFL 10 | from .grid_rcnn import GridRCNN 11 | from .htc import HybridTaskCascade 12 | from .mask_rcnn import MaskRCNN 13 | from .mask_scoring_rcnn import MaskScoringRCNN 14 | from .nasfcos import NASFCOS 15 | from .point_rend import PointRend 16 | from .reppoints_detector import RepPointsDetector 17 | from .retinanet import RetinaNet 18 | from .rpn import RPN 19 | from .single_stage import SingleStageDetector 20 | from .two_stage import TwoStageDetector 21 | 22 | from .obb.obb_base import OBBBaseDetector 23 | from .obb.obb_two_stage import OBBTwoStageDetector 24 | from .obb.obb_single_stage import OBBSingleStageDetector 25 | from .obb.faster_rcnn_obb import FasterRCNNOBB 26 | from .obb.roi_transformer import RoITransformer 27 | from .obb.retinanet_obb import RetinaNetOBB 28 | from .obb.gliding_vertex import GlidingVertex 29 | from .obb.obb_rpn import OBBRPN 30 | 31 | __all__ = [ 32 | 'ATSS', 'BaseDetector', 'SingleStageDetector', 'TwoStageDetector', 'RPN', 33 | 'FastRCNN', 'FasterRCNN', 'MaskRCNN', 'CascadeRCNN', 'HybridTaskCascade', 34 | 'RetinaNet', 'FCOS', 'GridRCNN', 'MaskScoringRCNN', 'RepPointsDetector', 35 | 'FOVEA', 'FSAF', 'NASFCOS', 'PointRend', 'GFL', 36 | 37 | 'OBBBaseDetector', 'OBBTwoStageDetector', 'OBBSingleStageDetector', 38 | 'FasterRCNNOBB', 'RetinaNetOBB', 'RoITransformer' 39 | ] 40 | -------------------------------------------------------------------------------- /mmdet/datasets/pipelines/__init__.py: -------------------------------------------------------------------------------- 1 | from .auto_augment import AutoAugment 2 | from .compose import Compose 3 | from .formating import (Collect, ImageToTensor, ToDataContainer, ToTensor, 4 | Transpose, to_tensor) 5 | from .instaboost import InstaBoost 6 | from .loading import (LoadAnnotations, LoadImageFromFile, 7 | LoadMultiChannelImageFromFiles, LoadProposals) 8 | from .test_time_aug import MultiScaleFlipAug 9 | from .transforms import (Albu, Expand, MinIoURandomCrop, Normalize, Pad, 10 | PhotoMetricDistortion, RandomCenterCropPad, 11 | RandomCrop, RandomFlip, Resize, SegRescale) 12 | 13 | from .obb.base import mask2obb, mask2poly, poly2mask 14 | from .obb.base import (LoadOBBAnnotations, Mask2OBB, OBBDefaultFormatBundle, 15 | OBBRandomFlip, RandomOBBRotate, MultiScaleFlipRotateAug, 16 | FliterEmpty) 17 | from .obb.dota import LoadDOTASpecialInfo, DOTASpecialIgnore 18 | 19 | __all__ = [ 20 | 'Compose', 'to_tensor', 'ToTensor', 'ImageToTensor', 'ToDataContainer', 21 | 'Transpose', 'Collect', 'LoadAnnotations', 'LoadImageFromFile', 22 | 'LoadMultiChannelImageFromFiles', 'LoadProposals', 'MultiScaleFlipAug', 23 | 'Resize', 'RandomFlip', 'Pad', 'RandomCrop', 'Normalize', 'SegRescale', 24 | 'MinIoURandomCrop', 'Expand', 'PhotoMetricDistortion', 'Albu', 25 | 'InstaBoost', 'RandomCenterCropPad', 'AutoAugment', 26 | 'LoadOBBAnnotations', 'Mask2OBB', 'OBBDefaultFormatBundle', 'OBBRandomFlip', 27 | 'RandomOBBRotate', 'LoadDOTASpecialInfo', 'DOTASpecialIgnore', 'FliterEmpty' 28 | ] 29 | -------------------------------------------------------------------------------- /mmdet/models/losses/mse_loss.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | 4 | from ..builder import LOSSES 5 | from .utils import weighted_loss 6 | 7 | 8 | @weighted_loss 9 | def mse_loss(pred, target): 10 | """Warpper of mse loss""" 11 | return F.mse_loss(pred, target, reduction='none') 12 | 13 | 14 | @LOSSES.register_module() 15 | class MSELoss(nn.Module): 16 | """MSELoss 17 | 18 | Args: 19 | reduction (str, optional): The method that reduces the loss to a 20 | scalar. Options are "none", "mean" and "sum". 21 | loss_weight (float, optional): The weight of the loss. Defaults to 1.0 22 | """ 23 | 24 | def __init__(self, reduction='mean', loss_weight=1.0): 25 | super().__init__() 26 | self.reduction = reduction 27 | self.loss_weight = loss_weight 28 | 29 | def forward(self, pred, target, weight=None, avg_factor=None): 30 | """Forward function of loss 31 | 32 | Args: 33 | pred (torch.Tensor): The prediction. 34 | target (torch.Tensor): The learning target of the prediction. 35 | weight (torch.Tensor, optional): Weight of the loss for each 36 | prediction. Defaults to None. 37 | avg_factor (int, optional): Average factor that is used to average 38 | the loss. Defaults to None. 39 | 40 | Returns: 41 | torch.Tensor: The calculated loss 42 | """ 43 | loss = self.loss_weight * mse_loss( 44 | pred, 45 | target, 46 | weight, 47 | reduction=self.reduction, 48 | avg_factor=avg_factor) 49 | return loss 50 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_roi_head import BaseRoIHead 2 | from .bbox_heads import (BBoxHead, ConvFCBBoxHead, DoubleConvFCBBoxHead, 3 | Shared2FCBBoxHead, Shared4Conv1FCBBoxHead) 4 | from .cascade_roi_head import CascadeRoIHead 5 | from .double_roi_head import DoubleHeadRoIHead 6 | from .dynamic_roi_head import DynamicRoIHead 7 | from .grid_roi_head import GridRoIHead 8 | from .htc_roi_head import HybridTaskCascadeRoIHead 9 | from .mask_heads import (CoarseMaskHead, FCNMaskHead, FusedSemanticHead, 10 | GridHead, HTCMaskHead, MaskIoUHead, MaskPointHead) 11 | from .mask_scoring_roi_head import MaskScoringRoIHead 12 | from .pisa_roi_head import PISARoIHead 13 | from .point_rend_roi_head import PointRendRoIHead 14 | from .roi_extractors import SingleRoIExtractor 15 | from .shared_heads import ResLayer 16 | 17 | from .obb.obb_base_roi_head import OBBBaseRoIHead 18 | from .obb.roitrans_roi_head import RoITransRoIHead 19 | from .obb.obb_standard_roi_head import OBBStandardRoIHead 20 | from .obb.obb_standard_roi_head_xyrab import OBBStandardRoIHeadXYRAB 21 | from .obb.gv_ratio_roi_head import GVRatioRoIHead 22 | 23 | __all__ = [ 24 | 'BaseRoIHead', 'CascadeRoIHead', 'DoubleHeadRoIHead', 'MaskScoringRoIHead', 25 | 'HybridTaskCascadeRoIHead', 'GridRoIHead', 'ResLayer', 'BBoxHead', 26 | 'ConvFCBBoxHead', 'Shared2FCBBoxHead', 'Shared4Conv1FCBBoxHead', 27 | 'DoubleConvFCBBoxHead', 'FCNMaskHead', 'HTCMaskHead', 'FusedSemanticHead', 28 | 'GridHead', 'MaskIoUHead', 'SingleRoIExtractor', 'PISARoIHead', 29 | 'PointRendRoIHead', 'MaskPointHead', 'CoarseMaskHead', 'DynamicRoIHead', 30 | 31 | 'RoITransRoIHead' 32 | ] 33 | -------------------------------------------------------------------------------- /mmdet/datasets/pipelines/compose.py: -------------------------------------------------------------------------------- 1 | import collections 2 | 3 | from mmcv.utils import build_from_cfg 4 | 5 | from ..builder import PIPELINES 6 | 7 | 8 | @PIPELINES.register_module() 9 | class Compose(object): 10 | """Compose multiple transforms sequentially. 11 | 12 | Args: 13 | transforms (Sequence[dict | callable]): Sequence of transform object or 14 | config dict to be composed. 15 | """ 16 | 17 | def __init__(self, transforms): 18 | assert isinstance(transforms, collections.abc.Sequence) 19 | self.transforms = [] 20 | for transform in transforms: 21 | if isinstance(transform, dict): 22 | transform = build_from_cfg(transform, PIPELINES) 23 | self.transforms.append(transform) 24 | elif callable(transform): 25 | self.transforms.append(transform) 26 | else: 27 | raise TypeError('transform must be callable or a dict') 28 | 29 | def __call__(self, data): 30 | """Call function to apply transforms sequentially. 31 | 32 | Args: 33 | data (dict): A result dict contains the data to transform. 34 | 35 | Returns: 36 | dict: Transformed data. 37 | """ 38 | 39 | for t in self.transforms: 40 | data = t(data) 41 | if data is None: 42 | return None 43 | return data 44 | 45 | def __repr__(self): 46 | format_string = self.__class__.__name__ + '(' 47 | for t in self.transforms: 48 | format_string += '\n' 49 | format_string += f' {t}' 50 | format_string += '\n)' 51 | return format_string 52 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/bbox_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .bbox_head import BBoxHead 2 | from .convfc_bbox_head import (ConvFCBBoxHead, Shared2FCBBoxHead, 3 | Shared4Conv1FCBBoxHead) 4 | from .double_bbox_head import DoubleConvFCBBoxHead 5 | 6 | from .obb.obbox_head import OBBoxHead 7 | from .obb.obb_convfc_bbox_head import (OBBConvFCBBoxHead, OBBShared2FCBBoxHead, 8 | OBBShared4Conv1FCBBoxHead) 9 | from .obb.obb_convfc_bbox_head_refine import (OBBConvFCBBoxHeadRefine, OBBFCBBoxHeadRefine) 10 | from .obb.obb_convfc_bbox_head2 import (OBBConvFCBBoxHead2, OBBShared2FCBBoxHead2, 11 | OBBShared4Conv1FCBBoxHead2) 12 | from .obb.obb_convfc_bbox_head2_dcn import (OBBConvFCBBoxHead2DCN, OBBShared2FCBBoxHead2DCN, 13 | OBBShared4Conv1FCBBoxHead2DCN) 14 | from .obb.obb_convfc_bbox_head2_dcn_hbb import (OBBConvFCBBoxHead2DCNHBB, OBBShared2FCBBoxHead2DCNHBB, 15 | OBBShared4Conv1FCBBoxHead2DCNHBB) 16 | from .obb.gv_bbox_head import GVBBoxHead 17 | 18 | __all__ = [ 19 | 'BBoxHead', 'ConvFCBBoxHead', 'Shared2FCBBoxHead', 20 | 'Shared4Conv1FCBBoxHead', 'DoubleConvFCBBoxHead', 21 | 22 | 'OBBoxHead', 'OBBConvFCBBoxHead', 'OBBShared2FCBBoxHead', 23 | 'OBBShared4Conv1FCBBoxHead', 24 | 'OBBConvFCBBoxHeadRefine', 'OBBFCBBoxHeadRefine', 25 | 'OBBConvFCBBoxHead2', 'OBBShared2FCBBoxHead2', 26 | 'OBBShared4Conv1FCBBoxHead2', 27 | 'OBBConvFCBBoxHead2DCN', 'OBBShared2FCBBoxHead2DCN', 28 | 'OBBShared4Conv1FCBBoxHead2DCN', 29 | 'OBBConvFCBBoxHead2DCNHBB', 'OBBShared2FCBBoxHead2DCNHBB', 30 | 'OBBShared4Conv1FCBBoxHead2DCNHBB', 31 | ] 32 | -------------------------------------------------------------------------------- /mmdet/ops/utils/src/compiling_info.cpp: -------------------------------------------------------------------------------- 1 | // modified from 2 | // https://github.com/facebookresearch/detectron2/blob/master/detectron2/layers/csrc/vision.cpp 3 | #include 4 | 5 | #ifdef WITH_CUDA 6 | #include 7 | int get_cudart_version() { return CUDART_VERSION; } 8 | #endif 9 | 10 | std::string get_compiling_cuda_version() { 11 | #ifdef WITH_CUDA 12 | std::ostringstream oss; 13 | 14 | // copied from 15 | // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/detail/CUDAHooks.cpp#L231 16 | auto printCudaStyleVersion = [&](int v) { 17 | oss << (v / 1000) << "." << (v / 10 % 100); 18 | if (v % 10 != 0) { 19 | oss << "." << (v % 10); 20 | } 21 | }; 22 | printCudaStyleVersion(get_cudart_version()); 23 | return oss.str(); 24 | #else 25 | return std::string("not available"); 26 | #endif 27 | } 28 | 29 | // similar to 30 | // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Version.cpp 31 | std::string get_compiler_version() { 32 | std::ostringstream ss; 33 | #if defined(__GNUC__) 34 | #ifndef __clang__ 35 | { ss << "GCC " << __GNUC__ << "." << __GNUC_MINOR__; } 36 | #endif 37 | #endif 38 | 39 | #if defined(__clang_major__) 40 | { 41 | ss << "clang " << __clang_major__ << "." << __clang_minor__ << "." 42 | << __clang_patchlevel__; 43 | } 44 | #endif 45 | 46 | #if defined(_MSC_VER) 47 | { ss << "MSVC " << _MSC_FULL_VER; } 48 | #endif 49 | return ss.str(); 50 | } 51 | 52 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 53 | m.def("get_compiler_version", &get_compiler_version, "get_compiler_version"); 54 | m.def("get_compiling_cuda_version", &get_compiling_cuda_version, 55 | "get_compiling_cuda_version"); 56 | } 57 | -------------------------------------------------------------------------------- /mmdet/datasets/wider_face.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import xml.etree.ElementTree as ET 3 | 4 | import mmcv 5 | 6 | from .builder import DATASETS 7 | from .xml_style import XMLDataset 8 | 9 | 10 | @DATASETS.register_module() 11 | class WIDERFaceDataset(XMLDataset): 12 | """ 13 | Reader for the WIDER Face dataset in PASCAL VOC format. 14 | Conversion scripts can be found in 15 | https://github.com/sovrasov/wider-face-pascal-voc-annotations 16 | """ 17 | CLASSES = ('face', ) 18 | 19 | def __init__(self, **kwargs): 20 | super(WIDERFaceDataset, self).__init__(**kwargs) 21 | 22 | def load_annotations(self, ann_file): 23 | """Load annotation from WIDERFace XML style annotation file. 24 | 25 | Args: 26 | ann_file (str): Path of XML file. 27 | 28 | Returns: 29 | list[dict]: Annotation info from XML file. 30 | """ 31 | 32 | data_infos = [] 33 | img_ids = mmcv.list_from_file(ann_file) 34 | for img_id in img_ids: 35 | filename = f'{img_id}.jpg' 36 | xml_path = osp.join(self.img_prefix, 'Annotations', 37 | f'{img_id}.xml') 38 | tree = ET.parse(xml_path) 39 | root = tree.getroot() 40 | size = root.find('size') 41 | width = int(size.find('width').text) 42 | height = int(size.find('height').text) 43 | folder = root.find('folder').text 44 | data_infos.append( 45 | dict( 46 | id=img_id, 47 | filename=osp.join(folder, filename), 48 | width=width, 49 | height=height)) 50 | 51 | return data_infos 52 | -------------------------------------------------------------------------------- /mmdet/core/evaluation/bbox_overlaps.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def bbox_overlaps(bboxes1, bboxes2, mode='iou', eps=1e-6): 5 | """Calculate the ious between each bbox of bboxes1 and bboxes2. 6 | 7 | Args: 8 | bboxes1(ndarray): shape (n, 4) 9 | bboxes2(ndarray): shape (k, 4) 10 | mode(str): iou (intersection over union) or iof (intersection 11 | over foreground) 12 | 13 | Returns: 14 | ious(ndarray): shape (n, k) 15 | """ 16 | 17 | assert mode in ['iou', 'iof'] 18 | 19 | bboxes1 = bboxes1.astype(np.float32) 20 | bboxes2 = bboxes2.astype(np.float32) 21 | rows = bboxes1.shape[0] 22 | cols = bboxes2.shape[0] 23 | ious = np.zeros((rows, cols), dtype=np.float32) 24 | if rows * cols == 0: 25 | return ious 26 | exchange = False 27 | if bboxes1.shape[0] > bboxes2.shape[0]: 28 | bboxes1, bboxes2 = bboxes2, bboxes1 29 | ious = np.zeros((cols, rows), dtype=np.float32) 30 | exchange = True 31 | area1 = (bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1]) 32 | area2 = (bboxes2[:, 2] - bboxes2[:, 0]) * (bboxes2[:, 3] - bboxes2[:, 1]) 33 | for i in range(bboxes1.shape[0]): 34 | x_start = np.maximum(bboxes1[i, 0], bboxes2[:, 0]) 35 | y_start = np.maximum(bboxes1[i, 1], bboxes2[:, 1]) 36 | x_end = np.minimum(bboxes1[i, 2], bboxes2[:, 2]) 37 | y_end = np.minimum(bboxes1[i, 3], bboxes2[:, 3]) 38 | overlap = np.maximum(x_end - x_start, 0) * np.maximum( 39 | y_end - y_start, 0) 40 | if mode == 'iou': 41 | union = area1[i] + area2 - overlap 42 | else: 43 | union = area1[i] if not exchange else area2 44 | union = np.maximum(union, eps) 45 | ious[i, :] = overlap / union 46 | if exchange: 47 | ious = ious.T 48 | return ious 49 | -------------------------------------------------------------------------------- /mmdet/models/dense_heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .anchor_free_head import AnchorFreeHead 2 | from .anchor_head import AnchorHead 3 | from .atss_head import ATSSHead 4 | from .fcos_head import FCOSHead 5 | from .fovea_head import FoveaHead 6 | from .free_anchor_retina_head import FreeAnchorRetinaHead 7 | from .fsaf_head import FSAFHead 8 | from .ga_retina_head import GARetinaHead 9 | from .ga_rpn_head import GARPNHead 10 | from .gfl_head import GFLHead 11 | from .guided_anchor_head import FeatureAdaption, GuidedAnchorHead 12 | from .nasfcos_head import NASFCOSHead 13 | from .pisa_retinanet_head import PISARetinaHead 14 | from .pisa_ssd_head import PISASSDHead 15 | from .reppoints_head import RepPointsHead 16 | from .retina_head import RetinaHead 17 | from .retina_sepbn_head import RetinaSepBNHead 18 | from .rpn_head import RPNHead 19 | from .ssd_head import SSDHead 20 | 21 | from .obb.obb_anchor_head import OBBAnchorHead 22 | from .obb.sp_orpn_head import SPORPNHead 23 | from .obb.sp_orpn_head2 import SPORPNHead2 24 | from .obb.obb_retina_head import OBBRetinaHead 25 | from .obb.aog_rpn_head import AOGRPNHead 26 | from .obb.aopg_head import AOPGHead 27 | from .obb.aopg_wo_arm_head import AOPGWOARMHead 28 | from .obb.aopg_obbnms_head import AOPGOBBNMSHead 29 | from .obb.obb_rpn_head import OBBRPNHead 30 | 31 | from .obb.obb_rpn_head_xyrab import OBBRPNHeadXYRAB 32 | 33 | __all__ = [ 34 | 'AnchorFreeHead', 'AnchorHead', 'GuidedAnchorHead', 'FeatureAdaption', 35 | 'RPNHead', 'GARPNHead', 'RetinaHead', 'RetinaSepBNHead', 'GARetinaHead', 36 | 'SSDHead', 'FCOSHead', 'RepPointsHead', 'FoveaHead', 37 | 'FreeAnchorRetinaHead', 'ATSSHead', 'FSAFHead', 'NASFCOSHead', 38 | 'PISARetinaHead', 'PISASSDHead', 'GFLHead', 'SPORPNHead2', 'OBBAnchorHead', 39 | 'SPORPNHead', 'OBBRetinaHead', 'AOGRPNHead', 'AOPGHead', 'AOPGWOARMHead', 40 | 'AOPGOBBNMSHead', 'OBBRPNHead', 'OBBRPNHeadXYRAB' 41 | ] 42 | -------------------------------------------------------------------------------- /mmdet/ops/sigmoid_focal_loss/sigmoid_focal_loss.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.autograd import Function 3 | from torch.autograd.function import once_differentiable 4 | 5 | from . import sigmoid_focal_loss_ext 6 | 7 | 8 | class SigmoidFocalLossFunction(Function): 9 | 10 | @staticmethod 11 | def forward(ctx, input, target, gamma=2.0, alpha=0.25): 12 | ctx.save_for_backward(input, target) 13 | num_classes = input.shape[1] 14 | ctx.num_classes = num_classes 15 | ctx.gamma = gamma 16 | ctx.alpha = alpha 17 | 18 | loss = sigmoid_focal_loss_ext.forward(input, target, num_classes, 19 | gamma, alpha) 20 | return loss 21 | 22 | @staticmethod 23 | @once_differentiable 24 | def backward(ctx, d_loss): 25 | input, target = ctx.saved_tensors 26 | num_classes = ctx.num_classes 27 | gamma = ctx.gamma 28 | alpha = ctx.alpha 29 | d_loss = d_loss.contiguous() 30 | d_input = sigmoid_focal_loss_ext.backward(input, target, d_loss, 31 | num_classes, gamma, alpha) 32 | return d_input, None, None, None, None 33 | 34 | 35 | sigmoid_focal_loss = SigmoidFocalLossFunction.apply 36 | 37 | 38 | # TODO: remove this module 39 | class SigmoidFocalLoss(nn.Module): 40 | 41 | def __init__(self, gamma, alpha): 42 | super(SigmoidFocalLoss, self).__init__() 43 | self.gamma = gamma 44 | self.alpha = alpha 45 | 46 | def forward(self, logits, targets): 47 | assert logits.is_cuda 48 | loss = sigmoid_focal_loss(logits, targets, self.gamma, self.alpha) 49 | return loss.sum() 50 | 51 | def __repr__(self): 52 | tmpstr = self.__class__.__name__ 53 | tmpstr += f'(gamma={self.gamma}, alpha={self.alpha})' 54 | return tmpstr 55 | -------------------------------------------------------------------------------- /mmdet/models/builder.py: -------------------------------------------------------------------------------- 1 | from mmcv.utils import Registry, build_from_cfg 2 | from torch import nn 3 | 4 | BACKBONES = Registry('backbone') 5 | NECKS = Registry('neck') 6 | ROI_EXTRACTORS = Registry('roi_extractor') 7 | SHARED_HEADS = Registry('shared_head') 8 | HEADS = Registry('head') 9 | LOSSES = Registry('loss') 10 | DETECTORS = Registry('detector') 11 | 12 | 13 | def build(cfg, registry, default_args=None): 14 | """Build a module 15 | 16 | Args: 17 | cfg (dict, list[dict]): The config of modules, is is either a dict 18 | or a list of configs. 19 | registry (:obj:`Registry`): A registry the module belongs to. 20 | default_args (dict, optional): Default arguments to build the module. 21 | Defaults to None. 22 | 23 | Returns: 24 | nn.Module: A built nn module. 25 | """ 26 | if isinstance(cfg, list): 27 | modules = [ 28 | build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg 29 | ] 30 | return nn.Sequential(*modules) 31 | else: 32 | return build_from_cfg(cfg, registry, default_args) 33 | 34 | 35 | def build_backbone(cfg): 36 | """Build backbone""" 37 | return build(cfg, BACKBONES) 38 | 39 | 40 | def build_neck(cfg): 41 | """Build neck""" 42 | return build(cfg, NECKS) 43 | 44 | 45 | def build_roi_extractor(cfg): 46 | """Build roi extractor""" 47 | return build(cfg, ROI_EXTRACTORS) 48 | 49 | 50 | def build_shared_head(cfg): 51 | """Build shared head""" 52 | return build(cfg, SHARED_HEADS) 53 | 54 | 55 | def build_head(cfg): 56 | """Build head""" 57 | return build(cfg, HEADS) 58 | 59 | 60 | def build_loss(cfg): 61 | """Build loss""" 62 | return build(cfg, LOSSES) 63 | 64 | 65 | def build_detector(cfg, train_cfg=None, test_cfg=None): 66 | """Build detector""" 67 | return build(cfg, DETECTORS, dict(train_cfg=train_cfg, test_cfg=test_cfg)) 68 | -------------------------------------------------------------------------------- /mmdet/ops/nms/src/nms_ext.cpp: -------------------------------------------------------------------------------- 1 | // Modified from https://github.com/bharatsingh430/soft-nms/blob/master/lib/nms/cpu_nms.pyx, Soft-NMS is added 2 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 3 | #include 4 | 5 | at::Tensor nms_cpu(const at::Tensor& dets, const float threshold); 6 | 7 | at::Tensor soft_nms_cpu(const at::Tensor& dets, const float threshold, 8 | const unsigned char method, const float sigma, const 9 | float min_score); 10 | 11 | std::vector > nms_match_cpu(const at::Tensor& dets, const float threshold); 12 | 13 | 14 | #ifdef WITH_CUDA 15 | at::Tensor nms_cuda(const at::Tensor& dets, const float threshold); 16 | #endif 17 | 18 | at::Tensor nms(const at::Tensor& dets, const float threshold){ 19 | if (dets.device().is_cuda()) { 20 | #ifdef WITH_CUDA 21 | return nms_cuda(dets, threshold); 22 | #else 23 | AT_ERROR("nms is not compiled with GPU support"); 24 | #endif 25 | } 26 | return nms_cpu(dets, threshold); 27 | } 28 | 29 | at::Tensor soft_nms(const at::Tensor& dets, const float threshold, 30 | const unsigned char method, const float sigma, const 31 | float min_score) { 32 | if (dets.device().is_cuda()) { 33 | AT_ERROR("soft_nms is not implemented on GPU"); 34 | } 35 | return soft_nms_cpu(dets, threshold, method, sigma, min_score); 36 | } 37 | 38 | std::vector > nms_match(const at::Tensor& dets, const float threshold) { 39 | if (dets.type().is_cuda()) { 40 | AT_ERROR("nms_match is not implemented on GPU"); 41 | } 42 | return nms_match_cpu(dets, threshold); 43 | } 44 | 45 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 46 | m.def("nms", &nms, "non-maximum suppression"); 47 | m.def("soft_nms", &soft_nms, "soft non-maximum suppression"); 48 | m.def("nms_match", &nms_match, "non-maximum suppression match"); 49 | } 50 | -------------------------------------------------------------------------------- /mmdet/ops/nms_rotated/src/nms_rotated_ext.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/facebookresearch/detectron2/tree/master/detectron2/layers/csrc/nms_rotated 3 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 4 | #include 5 | #include 6 | 7 | 8 | #ifdef WITH_CUDA 9 | at::Tensor nms_rotated_cuda( 10 | const at::Tensor& dets, 11 | const at::Tensor& scores, 12 | const float iou_threshold); 13 | 14 | at::Tensor poly_nms_cuda( 15 | const at::Tensor boxes, 16 | float nms_overlap_thresh); 17 | #endif 18 | 19 | at::Tensor nms_rotated_cpu( 20 | const at::Tensor& dets, 21 | const at::Tensor& scores, 22 | const float iou_threshold); 23 | 24 | 25 | inline at::Tensor nms_rotated( 26 | const at::Tensor& dets, 27 | const at::Tensor& scores, 28 | const float iou_threshold) { 29 | assert(dets.device().is_cuda() == scores.device().is_cuda()); 30 | if (dets.device().is_cuda()) { 31 | #ifdef WITH_CUDA 32 | return nms_rotated_cuda( 33 | dets.contiguous(), scores.contiguous(), iou_threshold); 34 | #else 35 | AT_ERROR("Not compiled with GPU support"); 36 | #endif 37 | } 38 | return nms_rotated_cpu(dets.contiguous(), scores.contiguous(), iou_threshold); 39 | } 40 | 41 | 42 | inline at::Tensor nms_poly( 43 | const at::Tensor& dets, 44 | const float iou_threshold) { 45 | if (dets.device().is_cuda()) { 46 | #ifdef WITH_CUDA 47 | if (dets.numel() == 0) 48 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 49 | return poly_nms_cuda(dets, iou_threshold); 50 | #else 51 | AT_ERROR("POLY_NMS is not compiled with GPU support"); 52 | #endif 53 | } 54 | AT_ERROR("POLY_NMS is not implemented on CPU"); 55 | } 56 | 57 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 58 | m.def("nms_rotated", &nms_rotated, "nms for rotated bboxes"); 59 | m.def("nms_poly", &nms_poly, "nms for poly bboxes"); 60 | } 61 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | mmdet/version.py 107 | data/ 108 | .vscode 109 | .idea 110 | .DS_Store 111 | 112 | # custom 113 | *.pkl 114 | *.pkl.json 115 | *.log.json 116 | work_dirs/ 117 | 118 | # Pytorch 119 | *.pth 120 | *.py~ 121 | *.sh~ 122 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from . import pi 3 | 4 | 5 | def get_bbox_type(bboxes, with_score=False): 6 | dim = bboxes.shape[-1] 7 | if with_score: 8 | dim -= 1 9 | 10 | if dim == 4: 11 | return 'hbb' 12 | if dim == 5: 13 | return 'obb' 14 | if dim == 8: 15 | return 'poly' 16 | return 'notype' 17 | 18 | 19 | def get_bbox_dim(bbox_type, with_score=False): 20 | if bbox_type == 'hbb': 21 | dim = 4 22 | elif bbox_type == 'obb': 23 | dim = 5 24 | elif bbox_type == 'poly': 25 | dim = 8 26 | else: 27 | raise ValueError(f"don't know {bbox_type} bbox dim") 28 | 29 | if with_score: 30 | dim += 1 31 | return dim 32 | 33 | 34 | def choice_by_type(hbb_op, obb_op, poly_op, bboxes_or_type, 35 | with_score=False): 36 | if isinstance(bboxes_or_type, np.ndarray): 37 | bbox_type = get_bbox_type(bboxes_or_type, with_score) 38 | elif isinstance(bboxes_or_type, str): 39 | bbox_type = bboxes_or_type 40 | else: 41 | raise TypeError(f'need np.ndarray or str,', 42 | f'but get {type(bboxes_or_type)}') 43 | 44 | if bbox_type == 'hbb': 45 | return hbb_op 46 | elif bbox_type == 'obb': 47 | return obb_op 48 | elif bbox_type == 'poly': 49 | return poly_op 50 | else: 51 | raise ValueError('notype bboxes is not suppert') 52 | 53 | 54 | def regular_theta(theta, mode='180', start=-pi/2): 55 | assert mode in ['360', '180'] 56 | cycle = 2 * pi if mode == '360' else pi 57 | 58 | theta = theta - start 59 | theta = theta % cycle 60 | return theta + start 61 | 62 | 63 | def regular_obb(obboxes): 64 | x, y, w, h, theta = [obboxes[..., i] for i in range(5)] 65 | w_regular = np.where(w > h, w, h) 66 | h_regular = np.where(w > h, h, w) 67 | theta_regular = np.where(w > h, theta, theta+pi/2) 68 | theta_regular = regular_theta(theta_regular) 69 | return np.stack([x, y, w_regular, h_regular, theta_regular], axis=-1) 70 | -------------------------------------------------------------------------------- /mmdet/ops/box_iou_rotated/box_iou_rotated_wrapper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from . import box_iou_rotated_ext 5 | 6 | 7 | def obb_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, device_id=None): 8 | assert mode in ['iou', 'iof'] 9 | assert type(bboxes1) is type(bboxes2) 10 | if is_aligned: 11 | assert bboxes1.shape[0] == bboxes2.shape[0] 12 | 13 | if isinstance(bboxes1, torch.Tensor): 14 | is_numpy = False 15 | bboxes1_th = bboxes1 16 | bboxes2_th = bboxes2 17 | elif isinstance(bboxes1, np.ndarray): 18 | is_numpy = True 19 | device = 'cpu' if device_id is None else f'cuda:{device_id}' 20 | bboxes1_th = torch.from_numpy(bboxes1).float().to(device) 21 | bboxes2_th = torch.from_numpy(bboxes2).float().to(device) 22 | else: 23 | raise TypeError('bboxes must be either a Tensor or numpy array, ' 24 | f'but got {type(bboxes1)}') 25 | 26 | if bboxes1_th.numel() == 0 or bboxes2_th.numel() == 0: 27 | rows, cols = bboxes1_th.size(0), bboxes2_th.size(0) 28 | outputs = bboxes1_th.new_zeros(rows, 1) if is_aligned else \ 29 | bboxes1_th.new_zeros(rows, cols) 30 | else: 31 | outputs = box_iou_rotated_ext.overlaps( 32 | bboxes1_th, 33 | bboxes2_th, 34 | mode == 'iou') 35 | 36 | # same bug will happen when bbox size is to small 37 | too_small1 = bboxes1_th[:, [2, 3]].min(1)[0] < 0.001 38 | too_small2 = bboxes2_th[:, [2, 3]].min(1)[0] < 0.001 39 | if too_small1.any() or too_small2.any(): 40 | inds1 = torch.nonzero(too_small1, as_tuple=False) 41 | inds2 = torch.nonzero(too_small2, as_tuple=False) 42 | outputs[inds1, :] = 0. 43 | outputs[:, inds2] = 0. 44 | 45 | if is_aligned: 46 | eye_index = torch.arange(bboxes1.shape[0])[..., None] 47 | outputs = torch.gather(outputs, dim=1, index=eye_index) 48 | if is_numpy: 49 | outputs = outputs.cpu().numpy() 50 | return outputs 51 | -------------------------------------------------------------------------------- /mmdet/ops/carafe/src/carafe_naive_ext.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | #include 6 | 7 | #ifdef WITH_CUDA 8 | int carafe_naive_forward_cuda(at::Tensor features, at::Tensor masks, 9 | int kernel_size, int group_size, int scale_factor, 10 | at::Tensor output); 11 | 12 | int carafe_naive_backward_cuda(at::Tensor top_grad, at::Tensor features, 13 | at::Tensor masks, int kernel_size, 14 | int group_size, int scale_factor, 15 | at::Tensor bottom_grad, at::Tensor mask_grad); 16 | #endif 17 | 18 | int carafe_naive_forward(at::Tensor features, at::Tensor masks, 19 | int kernel_size, int group_size, int scale_factor, 20 | at::Tensor output) { 21 | if (features.device().is_cuda()) { 22 | #ifdef WITH_CUDA 23 | return carafe_naive_forward_cuda(features, masks, kernel_size, 24 | group_size, scale_factor, output); 25 | #else 26 | AT_ERROR("carafe naive is not compiled with GPU support"); 27 | #endif 28 | } 29 | AT_ERROR("carafe naive is not implemented on CPU"); 30 | } 31 | 32 | int carafe_naive_backward(at::Tensor top_grad, at::Tensor features, 33 | at::Tensor masks, int kernel_size, 34 | int group_size, int scale_factor, 35 | at::Tensor bottom_grad, at::Tensor mask_grad) { 36 | if (top_grad.device().is_cuda()) { 37 | #ifdef WITH_CUDA 38 | return carafe_naive_backward_cuda(top_grad, features, masks, kernel_size, 39 | group_size, scale_factor, bottom_grad, mask_grad); 40 | #else 41 | AT_ERROR("carafe naive is not compiled with GPU support"); 42 | #endif 43 | } 44 | AT_ERROR("carafe naive is not implemented on CPU"); 45 | 46 | } 47 | 48 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 49 | m.def("forward", &carafe_naive_forward, "carafe_naive forward"); 50 | m.def("backward", &carafe_naive_backward, "carafe_naive backward"); 51 | } 52 | -------------------------------------------------------------------------------- /mmdet/core/bbox/demodata.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | 5 | def ensure_rng(rng=None): 6 | """ 7 | Simple version of the ``kwarray.ensure_rng`` 8 | 9 | Args: 10 | rng (int | numpy.random.RandomState | None): 11 | if None, then defaults to the global rng. Otherwise this can be an 12 | integer or a RandomState class 13 | Returns: 14 | (numpy.random.RandomState) : rng - 15 | a numpy random number generator 16 | 17 | References: 18 | https://gitlab.kitware.com/computer-vision/kwarray/blob/master/kwarray/util_random.py#L270 19 | """ 20 | 21 | if rng is None: 22 | rng = np.random.mtrand._rand 23 | elif isinstance(rng, int): 24 | rng = np.random.RandomState(rng) 25 | else: 26 | rng = rng 27 | return rng 28 | 29 | 30 | def random_boxes(num=1, scale=1, rng=None): 31 | """ 32 | Simple version of ``kwimage.Boxes.random`` 33 | 34 | Returns: 35 | Tensor: shape (n, 4) in x1, y1, x2, y2 format. 36 | 37 | References: 38 | https://gitlab.kitware.com/computer-vision/kwimage/blob/master/kwimage/structs/boxes.py#L1390 39 | 40 | Example: 41 | >>> num = 3 42 | >>> scale = 512 43 | >>> rng = 0 44 | >>> boxes = random_boxes(num, scale, rng) 45 | >>> print(boxes) 46 | tensor([[280.9925, 278.9802, 308.6148, 366.1769], 47 | [216.9113, 330.6978, 224.0446, 456.5878], 48 | [405.3632, 196.3221, 493.3953, 270.7942]]) 49 | """ 50 | rng = ensure_rng(rng) 51 | 52 | tlbr = rng.rand(num, 4).astype(np.float32) 53 | 54 | tl_x = np.minimum(tlbr[:, 0], tlbr[:, 2]) 55 | tl_y = np.minimum(tlbr[:, 1], tlbr[:, 3]) 56 | br_x = np.maximum(tlbr[:, 0], tlbr[:, 2]) 57 | br_y = np.maximum(tlbr[:, 1], tlbr[:, 3]) 58 | 59 | tlbr[:, 0] = tl_x * scale 60 | tlbr[:, 1] = tl_y * scale 61 | tlbr[:, 2] = br_x * scale 62 | tlbr[:, 3] = br_y * scale 63 | 64 | boxes = torch.from_numpy(tlbr) 65 | return boxes 66 | 67 | 68 | def random_obboxes(num=1, scale=1, rng=None): 69 | raise NotImplementedError 70 | -------------------------------------------------------------------------------- /mmdet/core/post_processing/bbox_nms.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from mmdet.ops.nms import batched_nms 4 | 5 | 6 | def multiclass_nms(multi_bboxes, 7 | multi_scores, 8 | score_thr, 9 | nms_cfg, 10 | max_num=-1, 11 | score_factors=None): 12 | """NMS for multi-class bboxes. 13 | 14 | Args: 15 | multi_bboxes (Tensor): shape (n, #class*4) or (n, 4) 16 | multi_scores (Tensor): shape (n, #class), where the last column 17 | contains scores of the background class, but this will be ignored. 18 | score_thr (float): bbox threshold, bboxes with scores lower than it 19 | will not be considered. 20 | nms_thr (float): NMS IoU threshold 21 | max_num (int): if there are more than max_num bboxes after NMS, 22 | only top max_num will be kept. 23 | score_factors (Tensor): The factors multiplied to scores before 24 | applying NMS 25 | 26 | Returns: 27 | tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels 28 | are 0-based. 29 | """ 30 | num_classes = multi_scores.size(1) - 1 31 | # exclude background category 32 | if multi_bboxes.shape[1] > 4: 33 | bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4) 34 | else: 35 | bboxes = multi_bboxes[:, None].expand(-1, num_classes, 4) 36 | scores = multi_scores[:, :-1] 37 | 38 | # filter out boxes with low scores 39 | valid_mask = scores > score_thr 40 | bboxes = bboxes[valid_mask] 41 | if score_factors is not None: 42 | scores = scores * score_factors[:, None] 43 | scores = scores[valid_mask] 44 | labels = valid_mask.nonzero()[:, 1] 45 | 46 | if bboxes.numel() == 0: 47 | bboxes = multi_bboxes.new_zeros((0, 5)) 48 | labels = multi_bboxes.new_zeros((0, ), dtype=torch.long) 49 | return bboxes, labels 50 | 51 | dets, keep = batched_nms(bboxes, scores, labels, nms_cfg) 52 | 53 | if max_num > 0: 54 | dets = dets[:max_num] 55 | keep = keep[:max_num] 56 | 57 | return dets, labels[keep] 58 | -------------------------------------------------------------------------------- /mmdet/models/losses/accuracy.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | def accuracy(pred, target, topk=1): 5 | """Calculate accuracy according to the prediction and target 6 | 7 | Args: 8 | pred (torch.Tensor): The model prediction. 9 | target (torch.Tensor): The target of each prediction 10 | topk (int | tuple[int], optional): If the predictions in ``topk`` 11 | matches the target, the predictions will be regarded as 12 | correct ones. Defaults to 1. 13 | 14 | Returns: 15 | float | tuple[float]: If the input ``topk`` is a single integer, 16 | the function will return a single float as accuracy. If 17 | ``topk`` is a tuple containing multiple integers, the 18 | function will return a tuple containing accuracies of 19 | each ``topk`` number. 20 | """ 21 | assert isinstance(topk, (int, tuple)) 22 | if isinstance(topk, int): 23 | topk = (topk, ) 24 | return_single = True 25 | else: 26 | return_single = False 27 | 28 | maxk = max(topk) 29 | _, pred_label = pred.topk(maxk, dim=1) 30 | pred_label = pred_label.t() 31 | correct = pred_label.eq(target.view(1, -1).expand_as(pred_label)) 32 | 33 | res = [] 34 | for k in topk: 35 | correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) 36 | res.append(correct_k.mul_(100.0 / pred.size(0))) 37 | return res[0] if return_single else res 38 | 39 | 40 | class Accuracy(nn.Module): 41 | 42 | def __init__(self, topk=(1, )): 43 | """Module to calculate the accuracy 44 | 45 | Args: 46 | topk (tuple, optional): The criterion used to calculate the 47 | accuracy. Defaults to (1,). 48 | """ 49 | super().__init__() 50 | self.topk = topk 51 | 52 | def forward(self, pred, target): 53 | """Forward function to calculate accuracy 54 | 55 | Args: 56 | pred (torch.Tensor): Prediction of models. 57 | target (torch.Tensor): Target for each prediction. 58 | 59 | Returns: 60 | tuple[float]: The accuracies under different topk criterions. 61 | """ 62 | return accuracy(pred, target, self.topk) 63 | -------------------------------------------------------------------------------- /mmdet/ops/masked_conv/src/masked_conv2d_ext.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | #ifdef WITH_CUDA 7 | int masked_im2col_forward_cuda(const at::Tensor im, const at::Tensor mask_h_idx, 8 | const at::Tensor mask_w_idx, const int kernel_h, 9 | const int kernel_w, const int pad_h, 10 | const int pad_w, at::Tensor col); 11 | 12 | int masked_col2im_forward_cuda(const at::Tensor col, 13 | const at::Tensor mask_h_idx, 14 | const at::Tensor mask_w_idx, int height, 15 | int width, int channels, at::Tensor im); 16 | #endif 17 | 18 | int masked_im2col_forward(const at::Tensor im, const at::Tensor mask_h_idx, 19 | const at::Tensor mask_w_idx, const int kernel_h, 20 | const int kernel_w, const int pad_h, 21 | const int pad_w, at::Tensor col) { 22 | if (im.device().is_cuda()) { 23 | #ifdef WITH_CUDA 24 | return masked_im2col_forward_cuda(im, mask_h_idx, mask_w_idx, kernel_h, 25 | kernel_w, pad_h, pad_w, col); 26 | #else 27 | AT_ERROR("masked_im2col is not compiled with GPU support"); 28 | #endif 29 | } 30 | AT_ERROR("masked_im2col is not implemented on CPU"); 31 | } 32 | 33 | int masked_col2im_forward(const at::Tensor col, 34 | const at::Tensor mask_h_idx, 35 | const at::Tensor mask_w_idx, int height, 36 | int width, int channels, at::Tensor im) { 37 | if (col.device().is_cuda()) { 38 | #ifdef WITH_CUDA 39 | return masked_col2im_forward_cuda(col, mask_h_idx, mask_w_idx, height, 40 | width, channels, im); 41 | #else 42 | AT_ERROR("masked_col2im is not compiled with GPU support"); 43 | #endif 44 | } 45 | AT_ERROR("masked_col2im is not implemented on CPU"); 46 | } 47 | 48 | 49 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 50 | m.def("masked_im2col_forward", &masked_im2col_forward, 51 | "masked_im2col forward"); 52 | m.def("masked_col2im_forward", &masked_col2im_forward, 53 | "masked_col2im forward"); 54 | } 55 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/move.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .utils import get_bbox_type, regular_obb 4 | from .transforms import bbox2type 5 | 6 | 7 | def translate(bboxes, x, y): 8 | assert get_bbox_type(bboxes) != 'notype' 9 | 10 | if get_bbox_type(bboxes) == 'obb': 11 | translated = bboxes.copy() 12 | translated[..., :2] = translated[..., :2] + np.array([x, y]) 13 | else: 14 | dim = bboxes.shape[-1] 15 | translated = bboxes + np.array([x, y]*int(dim/2)) 16 | return translated 17 | 18 | 19 | def flip(bboxes, W, H, direction='horizontal'): 20 | assert get_bbox_type(bboxes) != 'notype' 21 | assert direction in ['horizontal', 'vertical'] 22 | 23 | flipped = bboxes.copy() 24 | if get_bbox_type(bboxes) == 'poly': 25 | if direction == 'horizontal': 26 | flipped[..., 0::2] = W - bboxes[..., 0::2] 27 | else: 28 | flipped[..., 1::2] = H - bboxes[..., 1::2] 29 | 30 | if get_bbox_type(bboxes) == 'obb': 31 | if direction == 'horizontal': 32 | flipped[..., 0] = W - bboxes[..., 0] 33 | else: 34 | flipped[..., 1] = H - bboxes[..., 1] 35 | flipped[..., 4] = -flipped[..., 4] 36 | flipped = regular_obb(flipped) 37 | 38 | if get_bbox_type(bboxes) == 'hbb': 39 | if direction == 'horizontal': 40 | flipped[..., 0::4] = W - bboxes[..., 2::4] 41 | flipped[..., 2::4] = W - bboxes[..., 0::4] 42 | else: 43 | flipped[..., 1::4] = H - bboxes[..., 3::4] 44 | flipped[..., 3::4] = H - bboxes[..., 1::4] 45 | return flipped 46 | 47 | 48 | def warp(bboxes, M, keep_type=False): 49 | ori_type = get_bbox_type(bboxes) 50 | assert ori_type != 'notype' 51 | assert M.ndim == 2 52 | 53 | polys = bbox2type(bboxes, 'poly') 54 | shape = polys.shape 55 | group_pts = polys.reshape(*shape[:-1], shape[-1]//2, 2) 56 | group_pts = np.insert(group_pts, 2, 1, axis=-1) 57 | warped_pts = np.matmul(group_pts, M.T) 58 | 59 | if M.shape[0] == 3: 60 | warped_pts = (warped_pts / warped_pts[..., -1:])[..., :-1] 61 | warped_pts = warped_pts.reshape(*shape) 62 | if keep_type: 63 | warped_pts = bbox2type(warped_pts, ori_type) 64 | return warped_pts 65 | -------------------------------------------------------------------------------- /mmdet/utils/collect_env.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import subprocess 3 | import sys 4 | from collections import defaultdict 5 | 6 | import cv2 7 | import mmcv 8 | import torch 9 | import torchvision 10 | 11 | import mmdet 12 | 13 | 14 | def collect_env(): 15 | """Collect the information of the running environments.""" 16 | env_info = {} 17 | env_info['sys.platform'] = sys.platform 18 | env_info['Python'] = sys.version.replace('\n', '') 19 | 20 | cuda_available = torch.cuda.is_available() 21 | env_info['CUDA available'] = cuda_available 22 | 23 | if cuda_available: 24 | from torch.utils.cpp_extension import CUDA_HOME 25 | env_info['CUDA_HOME'] = CUDA_HOME 26 | 27 | if CUDA_HOME is not None and osp.isdir(CUDA_HOME): 28 | try: 29 | nvcc = osp.join(CUDA_HOME, 'bin/nvcc') 30 | nvcc = subprocess.check_output( 31 | f'"{nvcc}" -V | tail -n1', shell=True) 32 | nvcc = nvcc.decode('utf-8').strip() 33 | except subprocess.SubprocessError: 34 | nvcc = 'Not Available' 35 | env_info['NVCC'] = nvcc 36 | 37 | devices = defaultdict(list) 38 | for k in range(torch.cuda.device_count()): 39 | devices[torch.cuda.get_device_name(k)].append(str(k)) 40 | for name, devids in devices.items(): 41 | env_info['GPU ' + ','.join(devids)] = name 42 | 43 | gcc = subprocess.check_output('gcc --version | head -n1', shell=True) 44 | gcc = gcc.decode('utf-8').strip() 45 | env_info['GCC'] = gcc 46 | 47 | env_info['PyTorch'] = torch.__version__ 48 | env_info['PyTorch compiling details'] = torch.__config__.show() 49 | 50 | env_info['TorchVision'] = torchvision.__version__ 51 | 52 | env_info['OpenCV'] = cv2.__version__ 53 | 54 | env_info['MMCV'] = mmcv.__version__ 55 | env_info['MMDetection'] = mmdet.__version__ 56 | from mmdet.ops import get_compiler_version, get_compiling_cuda_version 57 | env_info['MMDetection Compiler'] = get_compiler_version() 58 | env_info['MMDetection CUDA Compiler'] = get_compiling_cuda_version() 59 | return env_info 60 | 61 | 62 | if __name__ == '__main__': 63 | for name, val in collect_env().items(): 64 | print(f'{name}: {val}') 65 | -------------------------------------------------------------------------------- /mmdet/core/utils/dist_utils.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | from collections import OrderedDict 3 | 4 | import torch.distributed as dist 5 | from mmcv.runner import OptimizerHook 6 | from torch._utils import (_flatten_dense_tensors, _take_tensors, 7 | _unflatten_dense_tensors) 8 | 9 | 10 | def _allreduce_coalesced(tensors, world_size, bucket_size_mb=-1): 11 | if bucket_size_mb > 0: 12 | bucket_size_bytes = bucket_size_mb * 1024 * 1024 13 | buckets = _take_tensors(tensors, bucket_size_bytes) 14 | else: 15 | buckets = OrderedDict() 16 | for tensor in tensors: 17 | tp = tensor.type() 18 | if tp not in buckets: 19 | buckets[tp] = [] 20 | buckets[tp].append(tensor) 21 | buckets = buckets.values() 22 | 23 | for bucket in buckets: 24 | flat_tensors = _flatten_dense_tensors(bucket) 25 | dist.all_reduce(flat_tensors) 26 | flat_tensors.div_(world_size) 27 | for tensor, synced in zip( 28 | bucket, _unflatten_dense_tensors(flat_tensors, bucket)): 29 | tensor.copy_(synced) 30 | 31 | 32 | def allreduce_grads(params, coalesce=True, bucket_size_mb=-1): 33 | """Allreduce gradients 34 | 35 | Args: 36 | params (list[torch.Parameters]): List of parameters of a model 37 | coalesce (bool, optional): Whether allreduce parameters as a whole. 38 | Defaults to True. 39 | bucket_size_mb (int, optional): Size of bucket, the unit is MB. 40 | Defaults to -1. 41 | """ 42 | grads = [ 43 | param.grad.data for param in params 44 | if param.requires_grad and param.grad is not None 45 | ] 46 | world_size = dist.get_world_size() 47 | if coalesce: 48 | _allreduce_coalesced(grads, world_size, bucket_size_mb) 49 | else: 50 | for tensor in grads: 51 | dist.all_reduce(tensor.div_(world_size)) 52 | 53 | 54 | class DistOptimizerHook(OptimizerHook): 55 | """Deprecated optimizer hook for distributed training""" 56 | 57 | def __init__(self, *args, **kwargs): 58 | warnings.warn('"DistOptimizerHook" is deprecated, please switch to' 59 | '"mmcv.runner.OptimizerHook".') 60 | super().__init__(*args, **kwargs) 61 | -------------------------------------------------------------------------------- /mmdet/models/dense_heads/base_dense_head.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | import torch.nn as nn 4 | 5 | 6 | class BaseDenseHead(nn.Module, metaclass=ABCMeta): 7 | """Base class for DenseHeads""" 8 | 9 | def __init__(self): 10 | super(BaseDenseHead, self).__init__() 11 | 12 | @abstractmethod 13 | def loss(self, **kwargs): 14 | """Compute losses of the head.""" 15 | pass 16 | 17 | @abstractmethod 18 | def get_bboxes(self, **kwargs): 19 | """Transform network output for a batch into bbox predictions.""" 20 | pass 21 | 22 | def forward_train(self, 23 | x, 24 | img_metas, 25 | gt_bboxes, 26 | gt_labels=None, 27 | gt_bboxes_ignore=None, 28 | proposal_cfg=None, 29 | **kwargs): 30 | """ 31 | Args: 32 | x (list[Tensor]): Features from FPN. 33 | img_metas (list[dict]): Meta information of each image, e.g., 34 | image size, scaling factor, etc. 35 | gt_bboxes (Tensor): Ground truth bboxes of the image, 36 | shape (num_gts, 4). 37 | gt_labels (Tensor): Ground truth labels of each box, 38 | shape (num_gts,). 39 | gt_bboxes_ignore (Tensor): Ground truth bboxes to be 40 | ignored, shape (num_ignored_gts, 4). 41 | proposal_cfg (mmcv.Config): Test / postprocessing configuration, 42 | if None, test_cfg would be used 43 | 44 | Returns: 45 | tuple: 46 | losses: (dict[str, Tensor]): A dictionary of loss components. 47 | proposal_list (list[Tensor]): Proposals of each image. 48 | """ 49 | outs = self(x) 50 | if gt_labels is None: 51 | loss_inputs = outs + (gt_bboxes, img_metas) 52 | else: 53 | loss_inputs = outs + (gt_bboxes, gt_labels, img_metas) 54 | losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore) 55 | if proposal_cfg is None: 56 | return losses 57 | else: 58 | proposal_list = self.get_bboxes(*outs, img_metas, cfg=proposal_cfg) 59 | return losses, proposal_list 60 | -------------------------------------------------------------------------------- /mmdet/models/dense_heads/rpn_test_mixin.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | from mmdet.core import merge_aug_proposals 4 | 5 | if sys.version_info >= (3, 7): 6 | from mmdet.utils.contextmanagers import completed 7 | 8 | 9 | class RPNTestMixin(object): 10 | """Test methods of RPN.""" 11 | 12 | if sys.version_info >= (3, 7): 13 | 14 | async def async_simple_test_rpn(self, x, img_metas): 15 | sleep_interval = self.rpn_head.test_cfg.pop( 16 | 'async_sleep_interval', 0.025) 17 | async with completed( 18 | __name__, 'rpn_head_forward', 19 | sleep_interval=sleep_interval): 20 | rpn_outs = self(x) 21 | 22 | proposal_list = self.get_bboxes(*rpn_outs, img_metas) 23 | return proposal_list 24 | 25 | def simple_test_rpn(self, x, img_metas): 26 | """Test without augmentation. 27 | 28 | Args: 29 | x (tuple[Tensor]): Features from the upstream network, each is 30 | a 4D-tensor. 31 | img_metas (list[dict]): Meta info of each image. 32 | 33 | Returns: 34 | list[Tensor]: Proposals of each image. 35 | """ 36 | rpn_outs = self(x) 37 | proposal_list = self.get_bboxes(*rpn_outs, img_metas) 38 | return proposal_list 39 | 40 | def aug_test_rpn(self, feats, img_metas): 41 | samples_per_gpu = len(img_metas[0]) 42 | aug_proposals = [[] for _ in range(samples_per_gpu)] 43 | for x, img_meta in zip(feats, img_metas): 44 | proposal_list = self.simple_test_rpn(x, img_meta) 45 | for i, proposals in enumerate(proposal_list): 46 | aug_proposals[i].append(proposals) 47 | # reorganize the order of 'img_metas' to match the dimensions 48 | # of 'aug_proposals' 49 | aug_img_metas = [] 50 | for i in range(samples_per_gpu): 51 | aug_img_meta = [] 52 | for j in range(len(img_metas)): 53 | aug_img_meta.append(img_metas[j][i]) 54 | aug_img_metas.append(aug_img_meta) 55 | # after merging, proposals will be rescaled to the original image size 56 | merged_proposals = [ 57 | merge_aug_proposals(proposals, aug_img_meta, self.test_cfg) 58 | for proposals, aug_img_meta in zip(aug_proposals, aug_img_metas) 59 | ] 60 | return merged_proposals 61 | -------------------------------------------------------------------------------- /mmdet/ops/carafe/grad_check.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import sys 3 | 4 | import mmcv 5 | import torch 6 | from torch.autograd import gradcheck 7 | 8 | sys.path.append(osp.abspath(osp.join(__file__, '../../'))) 9 | from mmdet.ops.carafe import CARAFE, CARAFENaive # noqa: E402, isort:skip 10 | from mmdet.ops.carafe import carafe, carafe_naive # noqa: E402, isort:skip 11 | 12 | feat = torch.randn(2, 64, 3, 3, requires_grad=True, device='cuda:0').double() 13 | mask = torch.randn( 14 | 2, 100, 6, 6, requires_grad=True, device='cuda:0').sigmoid().double() 15 | 16 | print('Gradcheck for carafe...') 17 | test = gradcheck(CARAFE(5, 4, 2), (feat, mask), atol=1e-4, eps=1e-4) 18 | print(test) 19 | 20 | print('Gradcheck for carafe naive...') 21 | test = gradcheck(CARAFENaive(5, 4, 2), (feat, mask), atol=1e-4, eps=1e-4) 22 | print(test) 23 | 24 | feat = torch.randn( 25 | 2, 1024, 100, 100, requires_grad=True, device='cuda:0').float() 26 | mask = torch.randn( 27 | 2, 25, 200, 200, requires_grad=True, device='cuda:0').sigmoid().float() 28 | loop_num = 500 29 | 30 | time_forward = 0 31 | time_backward = 0 32 | bar = mmcv.ProgressBar(loop_num) 33 | timer = mmcv.Timer() 34 | for i in range(loop_num): 35 | x = carafe(feat.clone(), mask.clone(), 5, 1, 2) 36 | torch.cuda.synchronize() 37 | time_forward += timer.since_last_check() 38 | x.sum().backward(retain_graph=True) 39 | torch.cuda.synchronize() 40 | time_backward += timer.since_last_check() 41 | bar.update() 42 | forward_speed = (time_forward + 1e-3) * 1e3 / loop_num 43 | backward_speed = (time_backward + 1e-3) * 1e3 / loop_num 44 | print(f'\nCARAFE time forward: {forward_speed} ' 45 | f'ms/iter | time backward: {backward_speed} ms/iter') 46 | 47 | time_naive_forward = 0 48 | time_naive_backward = 0 49 | bar = mmcv.ProgressBar(loop_num) 50 | timer = mmcv.Timer() 51 | for i in range(loop_num): 52 | x = carafe_naive(feat.clone(), mask.clone(), 5, 1, 2) 53 | torch.cuda.synchronize() 54 | time_naive_forward += timer.since_last_check() 55 | x.sum().backward(retain_graph=True) 56 | torch.cuda.synchronize() 57 | time_naive_backward += timer.since_last_check() 58 | bar.update() 59 | forward_speed = (time_naive_forward + 1e-3) * 1e3 / loop_num 60 | backward_speed = (time_naive_backward + 1e-3) * 1e3 / loop_num 61 | print('\nCARAFE naive time forward: ' 62 | f'{forward_speed} ms/iter | time backward: {backward_speed} ms/iter') 63 | -------------------------------------------------------------------------------- /mmdet/models/detectors/fast_rcnn.py: -------------------------------------------------------------------------------- 1 | from ..builder import DETECTORS 2 | from .two_stage import TwoStageDetector 3 | 4 | 5 | @DETECTORS.register_module() 6 | class FastRCNN(TwoStageDetector): 7 | """Implementation of `Fast R-CNN `_""" 8 | 9 | def __init__(self, 10 | backbone, 11 | roi_head, 12 | train_cfg, 13 | test_cfg, 14 | neck=None, 15 | pretrained=None): 16 | super(FastRCNN, self).__init__( 17 | backbone=backbone, 18 | neck=neck, 19 | roi_head=roi_head, 20 | train_cfg=train_cfg, 21 | test_cfg=test_cfg, 22 | pretrained=pretrained) 23 | 24 | def forward_test(self, imgs, img_metas, proposals, **kwargs): 25 | """ 26 | Args: 27 | imgs (List[Tensor]): the outer list indicates test-time 28 | augmentations and inner Tensor should have a shape NxCxHxW, 29 | which contains all images in the batch. 30 | img_metas (List[List[dict]]): the outer list indicates test-time 31 | augs (multiscale, flip, etc.) and the inner list indicates 32 | images in a batch. 33 | proposals (List[List[Tensor]]): the outer list indicates test-time 34 | augs (multiscale, flip, etc.) and the inner list indicates 35 | images in a batch. The Tensor should have a shape Px4, where 36 | P is the number of proposals. 37 | """ 38 | for var, name in [(imgs, 'imgs'), (img_metas, 'img_metas')]: 39 | if not isinstance(var, list): 40 | raise TypeError(f'{name} must be a list, but got {type(var)}') 41 | 42 | num_augs = len(imgs) 43 | if num_augs != len(img_metas): 44 | raise ValueError(f'num of augmentations ({len(imgs)}) ' 45 | f'!= num of image meta ({len(img_metas)})') 46 | # TODO: remove the restriction of samples_per_gpu == 1 when prepared 47 | samples_per_gpu = imgs[0].size(0) 48 | assert samples_per_gpu == 1 49 | 50 | if num_augs == 1: 51 | return self.simple_test(imgs[0], img_metas[0], proposals[0], 52 | **kwargs) 53 | else: 54 | # TODO: support test-time augmentation 55 | assert NotImplementedError 56 | -------------------------------------------------------------------------------- /mmdet/ops/__init__.py: -------------------------------------------------------------------------------- 1 | from .context_block import ContextBlock 2 | from .conv_ws import ConvWS2d, conv_ws_2d 3 | from .corner_pool import CornerPool 4 | from .dcn import (DeformConv, DeformConvPack, DeformRoIPooling, 5 | DeformRoIPoolingPack, ModulatedDeformConv, 6 | ModulatedDeformConvPack, ModulatedDeformRoIPoolingPack, 7 | deform_conv, deform_roi_pooling, modulated_deform_conv) 8 | from .dcn2 import (DeformRoIPooling2, DeformRoIPoolingPack2, ModulatedDeformRoIPoolingPack2) 9 | from .generalized_attention import GeneralizedAttention 10 | from .masked_conv import MaskedConv2d 11 | from .nms import batched_nms, nms, nms_match, soft_nms 12 | from .non_local import NonLocal2D 13 | from .plugin import build_plugin_layer 14 | from .point_sample import (SimpleRoIAlign, point_sample, 15 | rel_roi_point_to_rel_img_point) 16 | from .roi_align import RoIAlign, roi_align 17 | from .roi_pool import RoIPool, roi_pool 18 | from .saconv import SAConv2d 19 | from .sigmoid_focal_loss import SigmoidFocalLoss, sigmoid_focal_loss 20 | from .utils import get_compiler_version, get_compiling_cuda_version 21 | from .wrappers import Conv2d, ConvTranspose2d, Linear, MaxPool2d 22 | 23 | from .roi_align_rotated import roi_align_rotated, RoIAlignRotated 24 | from .nms_rotated import obb_nms, poly_nms, BT_nms, arb_batched_nms 25 | from .box_iou_rotated import obb_overlaps 26 | from .convex import convex_sort 27 | 28 | __all__ = [ 29 | 'nms', 'soft_nms', 'RoIAlign', 'roi_align', 'RoIPool', 'roi_pool', 30 | 'DeformConv', 'DeformConvPack', 'DeformRoIPooling', 'DeformRoIPoolingPack', 31 | 'ModulatedDeformRoIPoolingPack', 'ModulatedDeformConv', 32 | 'ModulatedDeformConvPack', 'deform_conv', 'modulated_deform_conv', 33 | 'deform_roi_pooling', 'SigmoidFocalLoss', 'sigmoid_focal_loss', 34 | 'MaskedConv2d', 'ContextBlock', 'GeneralizedAttention', 'NonLocal2D', 35 | 'get_compiler_version', 'get_compiling_cuda_version', 'ConvWS2d', 36 | 'conv_ws_2d', 'build_plugin_layer', 'batched_nms', 'Conv2d', 37 | 'ConvTranspose2d', 'MaxPool2d', 'Linear', 'nms_match', 'CornerPool', 38 | 'point_sample', 'rel_roi_point_to_rel_img_point', 'SimpleRoIAlign', 39 | 'SAConv2d', 40 | 41 | 'roi_align_rotated', 'RoIAlignRotated', 'obb_nms', 'BT_nms', 42 | 'arb_batched_nms', 'obb_overlaps', 'convex_sort', 43 | 44 | 'DeformRoIPooling2', 'DeformRoIPoolingPack2', 'ModulatedDeformRoIPoolingPack2' 45 | ] 46 | -------------------------------------------------------------------------------- /mmdet/core/utils/misc.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import mmcv 4 | import numpy as np 5 | import torch 6 | from six.moves import map, zip 7 | 8 | 9 | def tensor2imgs(tensor, mean=(0, 0, 0), std=(1, 1, 1), to_rgb=True): 10 | """Convert tensor to images 11 | 12 | Args: 13 | tensor (torch.Tensor): Tensor that contains multiple images 14 | mean (tuple[float], optional): Mean of images. Defaults to (0, 0, 0). 15 | std (tuple[float], optional): Standard deviation of images. 16 | Defaults to (1, 1, 1). 17 | to_rgb (bool, optional): Whether convert the images to RGB format. 18 | Defaults to True. 19 | 20 | Returns: 21 | list[np.ndarray]: A list that contains multiple images. 22 | """ 23 | num_imgs = tensor.size(0) 24 | mean = np.array(mean, dtype=np.float32) 25 | std = np.array(std, dtype=np.float32) 26 | imgs = [] 27 | for img_id in range(num_imgs): 28 | img = tensor[img_id, ...].cpu().numpy().transpose(1, 2, 0) 29 | img = mmcv.imdenormalize( 30 | img, mean, std, to_bgr=to_rgb).astype(np.uint8) 31 | imgs.append(np.ascontiguousarray(img)) 32 | return imgs 33 | 34 | 35 | def multi_apply(func, *args, **kwargs): 36 | """Apply function to a list of arguments 37 | 38 | Note: 39 | This function applies the ``func`` to multiple inputs and 40 | map the multiple outputs of the ``func`` into different 41 | list. Each list contains the same type of outputs corresponding 42 | to different inputs. 43 | 44 | Args: 45 | func (Function): A function that will be applied to a list of 46 | arguments 47 | 48 | Returns: 49 | tuple(list): A tuple containing multiple list, each list contains 50 | a kind of returned results by the function 51 | """ 52 | pfunc = partial(func, **kwargs) if kwargs else func 53 | map_results = map(pfunc, *args) 54 | return tuple(map(list, zip(*map_results))) 55 | 56 | 57 | def unmap(data, count, inds, fill=0): 58 | """ Unmap a subset of item (data) back to the original set of items (of 59 | size count) """ 60 | if data.dim() == 1: 61 | ret = data.new_full((count, ), fill) 62 | ret[inds.type(torch.bool)] = data 63 | else: 64 | new_size = (count, ) + data.size()[1:] 65 | ret = data.new_full(new_size, fill) 66 | ret[inds.type(torch.bool), :] = data 67 | return ret 68 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/instance_balanced_pos_sampler.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from ..builder import BBOX_SAMPLERS 5 | from .random_sampler import RandomSampler 6 | 7 | 8 | @BBOX_SAMPLERS.register_module() 9 | class InstanceBalancedPosSampler(RandomSampler): 10 | """Instance balanced sampler that samples equal number of positive samples 11 | for each instance.""" 12 | 13 | def _sample_pos(self, assign_result, num_expected, **kwargs): 14 | """Sample positive boxes 15 | 16 | Args: 17 | assign_result (:obj:`AssignResult`): The assigned results of boxes. 18 | num_expected (int): The number of expected positive samples 19 | 20 | Returns: 21 | Tensor or ndarray: sampled indices. 22 | """ 23 | pos_inds = torch.nonzero(assign_result.gt_inds > 0, as_tuple=False) 24 | if pos_inds.numel() != 0: 25 | pos_inds = pos_inds.squeeze(1) 26 | if pos_inds.numel() <= num_expected: 27 | return pos_inds 28 | else: 29 | unique_gt_inds = assign_result.gt_inds[pos_inds].unique() 30 | num_gts = len(unique_gt_inds) 31 | num_per_gt = int(round(num_expected / float(num_gts)) + 1) 32 | sampled_inds = [] 33 | for i in unique_gt_inds: 34 | inds = torch.nonzero( 35 | assign_result.gt_inds == i.item(), as_tuple=False) 36 | if inds.numel() != 0: 37 | inds = inds.squeeze(1) 38 | else: 39 | continue 40 | if len(inds) > num_per_gt: 41 | inds = self.random_choice(inds, num_per_gt) 42 | sampled_inds.append(inds) 43 | sampled_inds = torch.cat(sampled_inds) 44 | if len(sampled_inds) < num_expected: 45 | num_extra = num_expected - len(sampled_inds) 46 | extra_inds = np.array( 47 | list(set(pos_inds.cpu()) - set(sampled_inds.cpu()))) 48 | if len(extra_inds) > num_extra: 49 | extra_inds = self.random_choice(extra_inds, num_extra) 50 | extra_inds = torch.from_numpy(extra_inds).to( 51 | assign_result.gt_inds.device).long() 52 | sampled_inds = torch.cat([sampled_inds, extra_inds]) 53 | elif len(sampled_inds) > num_expected: 54 | sampled_inds = self.random_choice(sampled_inds, num_expected) 55 | return sampled_inds 56 | -------------------------------------------------------------------------------- /mmdet/ops/carafe/src/carafe_ext.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | #include 6 | 7 | #ifdef WITH_CUDA 8 | int carafe_forward_cuda(at::Tensor features, at::Tensor rfeatures, 9 | at::Tensor masks, at::Tensor rmasks, int kernel_size, 10 | int group_size, int scale_factor, at::Tensor routput, 11 | at::Tensor output); 12 | 13 | int carafe_backward_cuda(at::Tensor top_grad, at::Tensor rfeatures, 14 | at::Tensor masks, int kernel_size, int group_size, 15 | int scale_factor, at::Tensor rtop_grad, 16 | at::Tensor rbottom_grad_hs, at::Tensor rbottom_grad, 17 | at::Tensor rmask_grad, at::Tensor bottom_grad, 18 | at::Tensor mask_grad); 19 | #endif 20 | 21 | int carafe_forward(at::Tensor features, at::Tensor rfeatures, 22 | at::Tensor masks, at::Tensor rmasks, int kernel_size, 23 | int group_size, int scale_factor, at::Tensor routput, 24 | at::Tensor output) { 25 | if (features.device().is_cuda()) { 26 | #ifdef WITH_CUDA 27 | return carafe_forward_cuda(features, rfeatures, masks, rmasks, kernel_size, 28 | group_size, scale_factor, routput, output); 29 | #else 30 | AT_ERROR("carafe is not compiled with GPU support"); 31 | #endif 32 | } 33 | AT_ERROR("carafe is not implemented on CPU"); 34 | } 35 | 36 | int carafe_backward(at::Tensor top_grad, at::Tensor rfeatures, 37 | at::Tensor masks, int kernel_size, int group_size, 38 | int scale_factor, at::Tensor rtop_grad, 39 | at::Tensor rbottom_grad_hs, at::Tensor rbottom_grad, 40 | at::Tensor rmask_grad, at::Tensor bottom_grad, 41 | at::Tensor mask_grad) { 42 | if (top_grad.device().is_cuda()) { 43 | #ifdef WITH_CUDA 44 | return carafe_backward_cuda(top_grad, rfeatures, masks, kernel_size, 45 | group_size, scale_factor, rtop_grad, rbottom_grad_hs, rbottom_grad, 46 | rmask_grad, bottom_grad, mask_grad); 47 | #else 48 | AT_ERROR("carafe is not compiled with GPU support"); 49 | #endif 50 | } 51 | AT_ERROR("carafe is not implemented on CPU"); 52 | } 53 | 54 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 55 | m.def("forward", &carafe_forward, "carafe forward"); 56 | m.def("backward", &carafe_backward, "carafe backward"); 57 | } 58 | -------------------------------------------------------------------------------- /mmdet/core/mask/utils.py: -------------------------------------------------------------------------------- 1 | import mmcv 2 | import numpy as np 3 | import pycocotools.mask as mask_util 4 | 5 | 6 | def split_combined_polys(polys, poly_lens, polys_per_mask): 7 | """Split the combined 1-D polys into masks. 8 | 9 | A mask is represented as a list of polys, and a poly is represented as 10 | a 1-D array. In dataset, all masks are concatenated into a single 1-D 11 | tensor. Here we need to split the tensor into original representations. 12 | 13 | Args: 14 | polys (list): a list (length = image num) of 1-D tensors 15 | poly_lens (list): a list (length = image num) of poly length 16 | polys_per_mask (list): a list (length = image num) of poly number 17 | of each mask 18 | 19 | Returns: 20 | list: a list (length = image num) of list (length = mask num) of 21 | list (length = poly num) of numpy array 22 | """ 23 | mask_polys_list = [] 24 | for img_id in range(len(polys)): 25 | polys_single = polys[img_id] 26 | polys_lens_single = poly_lens[img_id].tolist() 27 | polys_per_mask_single = polys_per_mask[img_id].tolist() 28 | 29 | split_polys = mmcv.slice_list(polys_single, polys_lens_single) 30 | mask_polys = mmcv.slice_list(split_polys, polys_per_mask_single) 31 | mask_polys_list.append(mask_polys) 32 | return mask_polys_list 33 | 34 | 35 | # TODO: move this function to more proper place 36 | def encode_mask_results(mask_results): 37 | """Encode bitmap mask to RLE code. 38 | 39 | Args: 40 | mask_results (list | tuple[list]): bitmap mask results. 41 | In mask scoring rcnn, mask_results is a tuple of (segm_results, 42 | segm_cls_score). 43 | 44 | Returns: 45 | list | tuple: RLE encoded mask. 46 | """ 47 | if isinstance(mask_results, tuple): # mask scoring 48 | cls_segms, cls_mask_scores = mask_results 49 | else: 50 | cls_segms = mask_results 51 | num_classes = len(cls_segms) 52 | encoded_mask_results = [[] for _ in range(num_classes)] 53 | for i in range(len(cls_segms)): 54 | for cls_segm in cls_segms[i]: 55 | encoded_mask_results[i].append( 56 | mask_util.encode( 57 | np.array( 58 | cls_segm[:, :, np.newaxis], order='F', 59 | dtype='uint8'))[0]) # encoded with RLE 60 | if isinstance(mask_results, tuple): 61 | return encoded_mask_results, cls_mask_scores 62 | else: 63 | return encoded_mask_results 64 | -------------------------------------------------------------------------------- /mmdet/core/mask/mask_target.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch.nn.modules.utils import _pair 4 | 5 | 6 | def mask_target(pos_proposals_list, pos_assigned_gt_inds_list, gt_masks_list, 7 | cfg): 8 | """ Compute mask target for positive proposals in multiple images. 9 | 10 | Args: 11 | pos_proposals_list (list[Tensor]): Positive proposals in multiple 12 | images. 13 | pos_assigned_gt_inds_list (list[Tensor]): Assigned GT indices for each 14 | positive proposals. 15 | gt_masks_list (list[:obj:`BaseInstanceMasks`]): Ground truth masks of 16 | each image. 17 | cfg (dict): Config dict that specifies the mask size. 18 | 19 | Returns: 20 | list[Tensor]: Mask target of each image. 21 | """ 22 | cfg_list = [cfg for _ in range(len(pos_proposals_list))] 23 | mask_targets = map(mask_target_single, pos_proposals_list, 24 | pos_assigned_gt_inds_list, gt_masks_list, cfg_list) 25 | mask_targets = list(mask_targets) 26 | if len(mask_targets) > 0: 27 | mask_targets = torch.cat(mask_targets) 28 | return mask_targets 29 | 30 | 31 | def mask_target_single(pos_proposals, pos_assigned_gt_inds, gt_masks, cfg): 32 | """Compute mask target for each positive proposal in the image. 33 | 34 | Args: 35 | pos_proposals (Tensor): Positive proposals. 36 | pos_assigned_gt_inds (Tensor): Assigned GT inds of positive proposals. 37 | gt_masks (:obj:`BaseInstanceMasks`): GT masks in the format of Bitmap 38 | or Polygon. 39 | cfg (dict): Config dict that indicate the mask size. 40 | 41 | Returns: 42 | Tensor: Mask target of each positive proposals in the image. 43 | """ 44 | device = pos_proposals.device 45 | mask_size = _pair(cfg.mask_size) 46 | num_pos = pos_proposals.size(0) 47 | if num_pos > 0: 48 | proposals_np = pos_proposals.cpu().numpy() 49 | maxh, maxw = gt_masks.height, gt_masks.width 50 | proposals_np[:, [0, 2]] = np.clip(proposals_np[:, [0, 2]], 0, maxw) 51 | proposals_np[:, [1, 3]] = np.clip(proposals_np[:, [1, 3]], 0, maxh) 52 | pos_assigned_gt_inds = pos_assigned_gt_inds.cpu().numpy() 53 | 54 | mask_targets = gt_masks.crop_and_resize( 55 | proposals_np, mask_size, device=device, 56 | inds=pos_assigned_gt_inds).to_ndarray() 57 | 58 | mask_targets = torch.from_numpy(mask_targets).float().to(device) 59 | else: 60 | mask_targets = pos_proposals.new_zeros((0, ) + mask_size) 61 | 62 | return mask_targets 63 | -------------------------------------------------------------------------------- /mmdet/ops/sigmoid_focal_loss/src/sigmoid_focal_loss_ext.cpp: -------------------------------------------------------------------------------- 1 | // modify from 2 | // https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/csrc/SigmoidFocalLoss.h 3 | #include 4 | 5 | #ifdef WITH_CUDA 6 | at::Tensor SigmoidFocalLoss_forward_cuda(const at::Tensor &logits, 7 | const at::Tensor &targets, 8 | const int num_classes, 9 | const float gamma, const float alpha); 10 | 11 | at::Tensor SigmoidFocalLoss_backward_cuda(const at::Tensor &logits, 12 | const at::Tensor &targets, 13 | const at::Tensor &d_losses, 14 | const int num_classes, 15 | const float gamma, const float alpha); 16 | #endif 17 | 18 | // Interface for Python 19 | at::Tensor SigmoidFocalLoss_forward(const at::Tensor &logits, 20 | const at::Tensor &targets, 21 | const int num_classes, const float gamma, 22 | const float alpha) { 23 | if (logits.device().is_cuda()) { 24 | #ifdef WITH_CUDA 25 | at::DeviceGuard guard(logits.device()); 26 | return SigmoidFocalLoss_forward_cuda(logits, targets, num_classes, gamma, 27 | alpha); 28 | #else 29 | AT_ERROR("SigmoidFocalLoss is not compiled with GPU support"); 30 | #endif 31 | } 32 | AT_ERROR("SigmoidFocalLoss is not implemented on the CPU"); 33 | } 34 | 35 | at::Tensor SigmoidFocalLoss_backward(const at::Tensor &logits, 36 | const at::Tensor &targets, 37 | const at::Tensor &d_losses, 38 | const int num_classes, const float gamma, 39 | const float alpha) { 40 | if (logits.device().is_cuda()) { 41 | #ifdef WITH_CUDA 42 | at::DeviceGuard guard(logits.device()); 43 | return SigmoidFocalLoss_backward_cuda(logits, targets, d_losses, 44 | num_classes, gamma, alpha); 45 | #else 46 | AT_ERROR("SigmoidFocalLoss is not compiled with GPU support"); 47 | #endif 48 | } 49 | AT_ERROR("SigmoidFocalLoss is not implemented on the CPU"); 50 | } 51 | 52 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 53 | m.def("forward", &SigmoidFocalLoss_forward, 54 | "SigmoidFocalLoss forward"); 55 | m.def("backward", &SigmoidFocalLoss_backward, 56 | "SigmoidFocalLoss backward"); 57 | } 58 | -------------------------------------------------------------------------------- /mmdet/ops/nms_rotated/src/nms_rotated_cpu.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/facebookresearch/detectron2/tree/master/detectron2/layers/csrc/nms_rotated 3 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 4 | #include 5 | #include "box_iou_rotated_utils.h" 6 | 7 | 8 | template 9 | at::Tensor nms_rotated_cpu_kernel( 10 | const at::Tensor& dets, 11 | const at::Tensor& scores, 12 | const float iou_threshold) { 13 | // nms_rotated_cpu_kernel is modified from torchvision's nms_cpu_kernel, 14 | // however, the code in this function is much shorter because 15 | // we delegate the IoU computation for rotated boxes to 16 | // the single_box_iou_rotated function in box_iou_rotated_utils.h 17 | AT_ASSERTM(dets.device().is_cpu(), "dets must be a CPU tensor"); 18 | AT_ASSERTM(scores.device().is_cpu(), "scores must be a CPU tensor"); 19 | AT_ASSERTM( 20 | dets.scalar_type() == scores.scalar_type(), 21 | "dets should have the same type as scores"); 22 | 23 | if (dets.numel() == 0) { 24 | return at::empty({0}, dets.options().dtype(at::kLong)); 25 | } 26 | 27 | auto order_t = std::get<1>(scores.sort(0, /* descending=*/true)); 28 | 29 | auto ndets = dets.size(0); 30 | at::Tensor suppressed_t = at::zeros({ndets}, dets.options().dtype(at::kByte)); 31 | at::Tensor keep_t = at::zeros({ndets}, dets.options().dtype(at::kLong)); 32 | 33 | auto suppressed = suppressed_t.data_ptr(); 34 | auto keep = keep_t.data_ptr(); 35 | auto order = order_t.data_ptr(); 36 | 37 | int64_t num_to_keep = 0; 38 | 39 | for (int64_t _i = 0; _i < ndets; _i++) { 40 | auto i = order[_i]; 41 | if (suppressed[i] == 1) { 42 | continue; 43 | } 44 | 45 | keep[num_to_keep++] = i; 46 | 47 | for (int64_t _j = _i + 1; _j < ndets; _j++) { 48 | auto j = order[_j]; 49 | if (suppressed[j] == 1) { 50 | continue; 51 | } 52 | 53 | auto ovr = single_box_iou_rotated( 54 | dets[i].data_ptr(), dets[j].data_ptr()); 55 | if (ovr >= iou_threshold) { 56 | suppressed[j] = 1; 57 | } 58 | } 59 | } 60 | return keep_t.narrow(/*dim=*/0, /*start=*/0, /*length=*/num_to_keep); 61 | } 62 | 63 | at::Tensor nms_rotated_cpu( 64 | // input must be contiguous 65 | const at::Tensor& dets, 66 | const at::Tensor& scores, 67 | const float iou_threshold) { 68 | auto result = at::empty({0}, dets.options()); 69 | 70 | AT_DISPATCH_FLOATING_TYPES(dets.scalar_type(), "nms_rotated", [&] { 71 | result = nms_rotated_cpu_kernel(dets, scores, iou_threshold); 72 | }); 73 | return result; 74 | } 75 | -------------------------------------------------------------------------------- /mmdet/datasets/pipelines/auto_augment.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | import numpy as np 4 | 5 | from ..builder import PIPELINES 6 | from .compose import Compose 7 | 8 | 9 | @PIPELINES.register_module() 10 | class AutoAugment(object): 11 | """Auto augmentation. 12 | 13 | This data augmentation is proposed in 14 | `Learning Data Augmentation Strategies for Object Detection `_ # noqa: E501 15 | 16 | Args: 17 | policies (list[list[dict]]): The policies of auto augmentation. Each 18 | policy in ``policies`` is a specific augmentation policy, and is 19 | composed by several augmentations (dict). When AutoAugment is 20 | called, a random policy in ``policies`` will be selected to 21 | augment images. 22 | 23 | Examples: 24 | >>> replace = (104, 116, 124) 25 | >>> policies = [ 26 | >>> [ 27 | >>> dict(type='Sharpness', prob=0.0, level=8), 28 | >>> dict( 29 | >>> type='Shear', 30 | >>> prob=0.4, 31 | >>> level=0, 32 | >>> replace=replace, 33 | >>> axis='x') 34 | >>> ], 35 | >>> [ 36 | >>> dict( 37 | >>> type='Rotate', 38 | >>> prob=0.6, 39 | >>> level=10, 40 | >>> replace=replace), 41 | >>> dict(type='Color', prob=1.0, level=6) 42 | >>> ] 43 | >>> ] 44 | >>> augmentation = AutoAugment(policies) 45 | >>> img = np.ones(100, 100, 3) 46 | >>> gt_bboxes = np.ones(10, 4) 47 | >>> results = dict(img=img, gt_bboxes=gt_bboxes) 48 | >>> results = augmentation(results) 49 | """ 50 | 51 | def __init__(self, policies): 52 | assert isinstance(policies, list) and len(policies) > 0, \ 53 | 'Policies must be a non-empty list.' 54 | for policy in policies: 55 | assert isinstance(policy, list) and len(policy) > 0, \ 56 | 'Each policy in policies must be a non-empty list.' 57 | for augment in policy: 58 | assert isinstance(augment, dict) and 'type' in augment, \ 59 | 'Each specific augmentation must be a dict with key' \ 60 | ' "type".' 61 | 62 | self.policies = copy.deepcopy(policies) 63 | self.transforms = [Compose(policy) for policy in self.policies] 64 | 65 | def __call__(self, results): 66 | transform = np.random.choice(self.transforms) 67 | return transform(results) 68 | 69 | def __repr__(self): 70 | return f'{self.__class__.__name__}(policies={self.policies}' 71 | -------------------------------------------------------------------------------- /mmdet/core/anchor/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def images_to_levels(target, num_levels): 5 | """Convert targets by image to targets by feature level. 6 | 7 | [target_img0, target_img1] -> [target_level0, target_level1, ...] 8 | """ 9 | target = torch.stack(target, 0) 10 | level_targets = [] 11 | start = 0 12 | for n in num_levels: 13 | end = start + n 14 | # level_targets.append(target[:, start:end].squeeze(0)) 15 | level_targets.append(target[:, start:end]) 16 | start = end 17 | return level_targets 18 | 19 | 20 | def anchor_inside_flags(flat_anchors, 21 | valid_flags, 22 | img_shape, 23 | allowed_border=0): 24 | """Check whether the anchors are inside the border 25 | 26 | Args: 27 | flat_anchors (torch.Tensor): Flatten anchors, shape (n, 4). 28 | valid_flags (torch.Tensor): An existing valid flags of anchors. 29 | img_shape (tuple(int)): Shape of current image. 30 | allowed_border (int, optional): The border to allow the valid anchor. 31 | Defaults to 0. 32 | 33 | Returns: 34 | torch.Tensor: Flags indicating whether the anchors are inside a 35 | valid range. 36 | """ 37 | img_h, img_w = img_shape[:2] 38 | if allowed_border >= 0: 39 | inside_flags = valid_flags & \ 40 | (flat_anchors[:, 0] >= -allowed_border) & \ 41 | (flat_anchors[:, 1] >= -allowed_border) & \ 42 | (flat_anchors[:, 2] < img_w + allowed_border) & \ 43 | (flat_anchors[:, 3] < img_h + allowed_border) 44 | else: 45 | inside_flags = valid_flags 46 | return inside_flags 47 | 48 | 49 | def calc_region(bbox, ratio, featmap_size=None): 50 | """Calculate a proportional bbox region. 51 | 52 | The bbox center are fixed and the new h' and w' is h * ratio and w * ratio. 53 | 54 | Args: 55 | bbox (Tensor): Bboxes to calculate regions, shape (n, 4). 56 | ratio (float): Ratio of the output region. 57 | featmap_size (tuple): Feature map size used for clipping the boundary. 58 | 59 | Returns: 60 | tuple: x1, y1, x2, y2 61 | """ 62 | x1 = torch.round((1 - ratio) * bbox[0] + ratio * bbox[2]).long() 63 | y1 = torch.round((1 - ratio) * bbox[1] + ratio * bbox[3]).long() 64 | x2 = torch.round(ratio * bbox[0] + (1 - ratio) * bbox[2]).long() 65 | y2 = torch.round(ratio * bbox[1] + (1 - ratio) * bbox[3]).long() 66 | if featmap_size is not None: 67 | x1 = x1.clamp(min=0, max=featmap_size[1]) 68 | y1 = y1.clamp(min=0, max=featmap_size[0]) 69 | x2 = x2.clamp(min=0, max=featmap_size[1]) 70 | y2 = y2.clamp(min=0, max=featmap_size[0]) 71 | return (x1, y1, x2, y2) 72 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/vis/vis_xie.py: -------------------------------------------------------------------------------- 1 | from multiprocessing import Pool 2 | import os.path as osp 3 | import BboxToolkit as bt 4 | 5 | 6 | 7 | # DOTA1': ('large-vehicle (LV)', 'swimming-pool (SP)', 'helicopter (HC)', 'bridge (BR)', 8 | # 'plane (PL)', 'ship (SH)', 'soccer-ball-field (SBF)', 'basketball-court (BC)', 9 | # 'ground-track-field (GTF)', 'small-vehicle (SV)', 'baseball-diamond (BD)','tennis-court (TC)', 10 | # 'roundabout (RA)', 'storage-tank (ST)', 'harbor (HA)') 11 | 12 | 13 | img_dir = '/disk0/evannnnnnn/home/Datasets/DOTA/train/images/p0775.png' 14 | ann_dir = '/disk0/evannnnnnn/home/Datasets/DOTA/train/labelTxt/p0775.txt' 15 | save_dir = '/disk2/xiexingxing/wjb/vis/base_wo_merge_xie2' 16 | colors = [(0,128,255),(116,115,147),(255,0,0),(0,255,0), 17 | (42,42,165), (255,0,255),(205,250,255),(193,193,255), 18 | (0,0,255),(226,43,138), (107,183,189), (255,255,0), 19 | (139,139,0),(153,51,0),(0,255,255) ] 20 | 21 | infos, classes = bt.load_dota_submission(ann_dir) 22 | colors = colors[:len(classes)] 23 | 24 | # img_dir = '/disk2/xiexingxing/wjb/HRSC2016/FullDataSet/AllImages/' 25 | # ann_dir = '/disk2/xiexingxing/wjb/results/sp_orpn/hrsc.pkl' 26 | # save_dir = '/disk2/xiexingxing/wjb/vis/hrsc' 27 | # infos, classes = bt.load_pkl(ann_dir) 28 | # colors = 'green' 29 | 30 | def vis(info): 31 | print(info['id']) 32 | bboxes = info['ann']['bboxes'] 33 | labels = info['ann']['labels'] 34 | scores = info['ann']['scores'] 35 | 36 | bboxes = bboxes[scores > 0.3] 37 | labels = labels[scores > 0.3] 38 | scores = scores[scores > 0.3] 39 | if len(bboxes) == 0: 40 | return 41 | bboxes = [bboxes[labels == i] for i in range(len(classes))] 42 | 43 | bt.imshow_bboxes( 44 | osp.join(img_dir, info['id'] + '.png'), 45 | bboxes, 46 | colors=colors, 47 | show=False, 48 | thickness=3, 49 | out_file=osp.join(save_dir, info['id'] + '.png') 50 | ) 51 | 52 | pool = Pool(5) 53 | pool.map(vis, infos) 54 | pool.close() 55 | # #def vis(info): 56 | # info=infos[0] 57 | # print(info['id']) 58 | # bboxes = info['ann']['bboxes'] 59 | # labels = info['ann']['labels'] 60 | # scores = info['ann']['scores'] 61 | 62 | # bboxes = bboxes[scores > 0.3] 63 | # labels = labels[scores > 0.3] 64 | # scores = scores[scores > 0.3] 65 | # # if len(bboxes) == 0: 66 | # # return 67 | # bboxes = [bboxes[labels == i] for i in range(len(classes))] 68 | 69 | # bt.imshow_bboxes( 70 | # osp.join(img_dir, info['id'] + '.png'), 71 | # bboxes, 72 | # colors=colors, 73 | # show=False, 74 | # thickness=3, 75 | # out_file=osp.join(save_dir, info['id'] + '.png') 76 | # ) 77 | 78 | # # pool = Pool(5) 79 | # # pool.map(vis, infos) 80 | # # pool.close() 81 | 82 | 83 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/shared_heads/res_layer.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from mmcv.cnn import constant_init, kaiming_init 3 | from mmcv.runner import load_checkpoint 4 | 5 | from mmdet.core import auto_fp16 6 | from mmdet.models.backbones import ResNet 7 | from mmdet.models.builder import SHARED_HEADS 8 | from mmdet.models.utils import ResLayer as _ResLayer 9 | from mmdet.utils import get_root_logger 10 | 11 | 12 | @SHARED_HEADS.register_module() 13 | class ResLayer(nn.Module): 14 | 15 | def __init__(self, 16 | depth, 17 | stage=3, 18 | stride=2, 19 | dilation=1, 20 | style='pytorch', 21 | norm_cfg=dict(type='BN', requires_grad=True), 22 | norm_eval=True, 23 | with_cp=False, 24 | dcn=None): 25 | super(ResLayer, self).__init__() 26 | self.norm_eval = norm_eval 27 | self.norm_cfg = norm_cfg 28 | self.stage = stage 29 | self.fp16_enabled = False 30 | block, stage_blocks = ResNet.arch_settings[depth] 31 | stage_block = stage_blocks[stage] 32 | planes = 64 * 2**stage 33 | inplanes = 64 * 2**(stage - 1) * block.expansion 34 | 35 | res_layer = _ResLayer( 36 | block, 37 | inplanes, 38 | planes, 39 | stage_block, 40 | stride=stride, 41 | dilation=dilation, 42 | style=style, 43 | with_cp=with_cp, 44 | norm_cfg=self.norm_cfg, 45 | dcn=dcn) 46 | self.add_module(f'layer{stage + 1}', res_layer) 47 | 48 | def init_weights(self, pretrained=None): 49 | """Initialize the weights in the module 50 | 51 | Args: 52 | pretrained (str, optional): Path to pre-trained weights. 53 | Defaults to None. 54 | """ 55 | if isinstance(pretrained, str): 56 | logger = get_root_logger() 57 | load_checkpoint(self, pretrained, strict=False, logger=logger) 58 | elif pretrained is None: 59 | for m in self.modules(): 60 | if isinstance(m, nn.Conv2d): 61 | kaiming_init(m) 62 | elif isinstance(m, nn.BatchNorm2d): 63 | constant_init(m, 1) 64 | else: 65 | raise TypeError('pretrained must be a str or None') 66 | 67 | @auto_fp16() 68 | def forward(self, x): 69 | res_layer = getattr(self, f'layer{self.stage + 1}') 70 | out = res_layer(x) 71 | return out 72 | 73 | def train(self, mode=True): 74 | super(ResLayer, self).train(mode) 75 | if self.norm_eval: 76 | for m in self.modules(): 77 | if isinstance(m, nn.BatchNorm2d): 78 | m.eval() 79 | -------------------------------------------------------------------------------- /mmdet/models/detectors/obb/obb_base.py: -------------------------------------------------------------------------------- 1 | import mmcv 2 | import warnings 3 | import numpy as np 4 | import BboxToolkit as bt 5 | 6 | from ..base import BaseDetector 7 | 8 | 9 | class OBBBaseDetector(BaseDetector): 10 | 11 | def show_result(self, 12 | img, 13 | result, 14 | score_thr=0.3, 15 | bbox_color='green', 16 | text_color='green', 17 | thickness=1, 18 | font_scale=0.5, 19 | win_name='', 20 | show=False, 21 | wait_time=0, 22 | out_file=None): 23 | 24 | img = mmcv.imread(img) 25 | img = img.copy() 26 | if isinstance(result, tuple): 27 | bbox_result, segm_result = result 28 | if isinstance(segm_result, tuple): 29 | segm_result = segm_result[0] # ms rcnn 30 | else: 31 | bbox_result, segm_result = result, None 32 | bboxes = np.vstack(bbox_result) 33 | labels = [ 34 | np.full(bbox.shape[0], i, dtype=np.int32) 35 | for i, bbox in enumerate(bbox_result) 36 | ] 37 | labels = np.concatenate(labels) 38 | # draw segmentation masks 39 | if segm_result is not None and len(labels) > 0: # non empty 40 | segms = mmcv.concat_list(segm_result) 41 | inds = np.where(bboxes[:, -1] > score_thr)[0] 42 | np.random.seed(42) 43 | color_masks = [ 44 | np.random.randint(0, 256, (1, 3), dtype=np.uint8) 45 | for _ in range(max(labels) + 1) 46 | ] 47 | for i in inds: 48 | i = int(i) 49 | color_mask = color_masks[labels[i]] 50 | mask = segms[i] 51 | img[mask] = img[mask] * 0.5 + color_mask * 0.5 52 | # if out_file specified, do not show image in window 53 | if out_file is not None: 54 | show = False 55 | # draw bounding boxes 56 | bboxes, scores = bboxes[:, :-1], bboxes[:, -1] 57 | bt.imshow_det_bboxes( 58 | img, 59 | bboxes, 60 | labels, 61 | scores=scores, 62 | class_names=self.CLASSES, 63 | score_thr=score_thr, 64 | bbox_color=bbox_color, 65 | text_color=text_color, 66 | thickness=thickness, 67 | font_scale=font_scale, 68 | win_name=win_name, 69 | show=show, 70 | wait_time=wait_time, 71 | out_file=out_file) 72 | 73 | if not (show or out_file): 74 | warnings.warn('show==False and out_file is not specified, only ' 75 | 'result image will be returned') 76 | return img 77 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/vis/base.py: -------------------------------------------------------------------------------- 1 | # Modified from mmcv.visualization 2 | # Copyright (c) Open-MMLab. All rights reserved. 3 | import cv2 4 | import numpy as np 5 | 6 | from enum import Enum 7 | 8 | 9 | class Color(Enum): 10 | """An enum that defines common colors. 11 | 12 | Contains red, green, blue, cyan, yellow, magenta, white and black. 13 | """ 14 | red = (0, 0, 255) 15 | green = (0, 255, 0) 16 | blue = (255, 0, 0) 17 | cyan = (255, 255, 0) 18 | yellow = (0, 255, 255) 19 | magenta = (255, 0, 255) 20 | white = (255, 255, 255) 21 | black = (0, 0, 0) 22 | 23 | 24 | def color_val(color): 25 | """Convert various input to color tuples. 26 | 27 | Args: 28 | color (:obj:`Color`/str/tuple/int/ndarray): Color inputs 29 | 30 | Returns: 31 | tuple[int]: A tuple of 3 integers indicating BGR channels. 32 | """ 33 | if isinstance(color, str): 34 | return Color[color].value 35 | elif isinstance(color, Color): 36 | return color.value 37 | elif isinstance(color, tuple): 38 | assert len(color) == 3 39 | for channel in color: 40 | assert channel >= 0 and channel <= 255 41 | return color 42 | elif isinstance(color, int): 43 | assert color >= 0 and color <= 255 44 | return color, color, color 45 | elif isinstance(color, np.ndarray): 46 | assert color.ndim == 1 and color.size == 3 47 | assert np.all((color >= 0) & (color <= 255)) 48 | color = color.astype(np.uint8) 49 | return tuple(color) 50 | else: 51 | raise TypeError('Invalid type for color: {}'.format(type(color))) 52 | 53 | 54 | def imshow(img, win_name='', wait_time=0, max_size=1000): 55 | """Show an image. 56 | 57 | Args: 58 | img (str or ndarray): The image to be displayed. 59 | win_name (str): The window name. 60 | wait_time (int): Value of waitKey param. 61 | max_size (int): Max size of window 62 | """ 63 | height, width = img.shape[:2] 64 | if max(height, width) > max_size: 65 | win_height = max_size if height > width else int(max_size*height/width) 66 | win_width = max_size if width >= height else int(max_size*width/height) 67 | else: 68 | win_height, win_width = height, width 69 | cv2.namedWindow(win_name, 0) 70 | cv2.resizeWindow(win_name, win_width, win_height) 71 | cv2.imshow(win_name, img) 72 | if wait_time == 0: # prevent from hangning if windows was closed 73 | while True: 74 | ret = cv2.waitKey(1) 75 | 76 | closed = cv2.getWindowProperty(win_name, cv2.WND_PROP_VISIBLE) < 1 77 | # if user closed window or if some key pressed 78 | if closed or ret != -1: 79 | break 80 | else: 81 | ret = cv2.waitKey(wait_time) 82 | 83 | -------------------------------------------------------------------------------- /mmdet/ops/roi_pool/roi_pool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Function 4 | from torch.autograd.function import once_differentiable 5 | from torch.nn.modules.utils import _pair 6 | 7 | from . import roi_pool_ext 8 | 9 | 10 | class RoIPoolFunction(Function): 11 | 12 | @staticmethod 13 | def forward(ctx, features, rois, out_size, spatial_scale): 14 | assert features.is_cuda 15 | out_h, out_w = _pair(out_size) 16 | assert isinstance(out_h, int) and isinstance(out_w, int) 17 | ctx.save_for_backward(rois) 18 | num_channels = features.size(1) 19 | num_rois = rois.size(0) 20 | out_size = (num_rois, num_channels, out_h, out_w) 21 | output = features.new_zeros(out_size) 22 | argmax = features.new_zeros(out_size, dtype=torch.int) 23 | roi_pool_ext.forward(features, rois, out_h, out_w, spatial_scale, 24 | output, argmax) 25 | ctx.spatial_scale = spatial_scale 26 | ctx.feature_size = features.size() 27 | ctx.argmax = argmax 28 | 29 | return output 30 | 31 | @staticmethod 32 | @once_differentiable 33 | def backward(ctx, grad_output): 34 | assert grad_output.is_cuda 35 | spatial_scale = ctx.spatial_scale 36 | feature_size = ctx.feature_size 37 | argmax = ctx.argmax 38 | rois = ctx.saved_tensors[0] 39 | assert feature_size is not None 40 | 41 | grad_input = grad_rois = None 42 | if ctx.needs_input_grad[0]: 43 | grad_input = grad_output.new_zeros(feature_size) 44 | roi_pool_ext.backward(grad_output.contiguous(), rois, argmax, 45 | spatial_scale, grad_input) 46 | 47 | return grad_input, grad_rois, None, None 48 | 49 | 50 | roi_pool = RoIPoolFunction.apply 51 | 52 | 53 | class RoIPool(nn.Module): 54 | 55 | def __init__(self, out_size, spatial_scale, use_torchvision=False): 56 | super(RoIPool, self).__init__() 57 | 58 | self.out_size = _pair(out_size) 59 | self.spatial_scale = float(spatial_scale) 60 | self.use_torchvision = use_torchvision 61 | 62 | def forward(self, features, rois): 63 | if self.use_torchvision: 64 | from torchvision.ops import roi_pool as tv_roi_pool 65 | return tv_roi_pool(features, rois, self.out_size, 66 | self.spatial_scale) 67 | else: 68 | return roi_pool(features, rois, self.out_size, self.spatial_scale) 69 | 70 | def __repr__(self): 71 | format_str = self.__class__.__name__ 72 | format_str += f'(out_size={self.out_size}, ' 73 | format_str += f'spatial_scale={self.spatial_scale}, ' 74 | format_str += f'use_torchvision={self.use_torchvision})' 75 | return format_str 76 | -------------------------------------------------------------------------------- /mmdet/core/evaluation/eval_hooks.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | 3 | from mmcv.runner import Hook 4 | from torch.utils.data import DataLoader 5 | 6 | 7 | class EvalHook(Hook): 8 | """Evaluation hook. 9 | 10 | Attributes: 11 | dataloader (DataLoader): A PyTorch dataloader. 12 | interval (int): Evaluation interval (by epochs). Default: 1. 13 | """ 14 | 15 | def __init__(self, dataloader, interval=1, **eval_kwargs): 16 | if not isinstance(dataloader, DataLoader): 17 | raise TypeError('dataloader must be a pytorch DataLoader, but got' 18 | f' {type(dataloader)}') 19 | self.dataloader = dataloader 20 | self.interval = interval 21 | self.eval_kwargs = eval_kwargs 22 | 23 | def after_train_epoch(self, runner): 24 | if not self.every_n_epochs(runner, self.interval): 25 | return 26 | from mmdet.apis import single_gpu_test 27 | results = single_gpu_test(runner.model, self.dataloader, show=False) 28 | self.evaluate(runner, results) 29 | 30 | def evaluate(self, runner, results): 31 | eval_res = self.dataloader.dataset.evaluate( 32 | results, logger=runner.logger, **self.eval_kwargs) 33 | for name, val in eval_res.items(): 34 | runner.log_buffer.output[name] = val 35 | runner.log_buffer.ready = True 36 | 37 | 38 | class DistEvalHook(EvalHook): 39 | """Distributed evaluation hook. 40 | 41 | Attributes: 42 | dataloader (DataLoader): A PyTorch dataloader. 43 | interval (int): Evaluation interval (by epochs). Default: 1. 44 | tmpdir (str | None): Temporary directory to save the results of all 45 | processes. Default: None. 46 | gpu_collect (bool): Whether to use gpu or cpu to collect results. 47 | Default: False. 48 | """ 49 | 50 | def __init__(self, 51 | dataloader, 52 | interval=1, 53 | gpu_collect=False, 54 | **eval_kwargs): 55 | if not isinstance(dataloader, DataLoader): 56 | raise TypeError('dataloader must be a pytorch DataLoader, but got ' 57 | f'{type(dataloader)}') 58 | self.dataloader = dataloader 59 | self.interval = interval 60 | self.gpu_collect = gpu_collect 61 | self.eval_kwargs = eval_kwargs 62 | 63 | def after_train_epoch(self, runner): 64 | if not self.every_n_epochs(runner, self.interval): 65 | return 66 | from mmdet.apis import multi_gpu_test 67 | results = multi_gpu_test( 68 | runner.model, 69 | self.dataloader, 70 | tmpdir=osp.join(runner.work_dir, '.eval_hook'), 71 | gpu_collect=self.gpu_collect) 72 | if runner.rank == 0: 73 | print('\n') 74 | self.evaluate(runner, results) 75 | -------------------------------------------------------------------------------- /mmdet/core/bbox/__init__.py: -------------------------------------------------------------------------------- 1 | from .assigners import (AssignResult, BaseAssigner, CenterRegionAssigner, 2 | MaxIoUAssigner) 3 | from .builder import build_assigner, build_bbox_coder, build_sampler 4 | from .coder import (BaseBBoxCoder, DeltaXYWHBBoxCoder, PseudoBBoxCoder, 5 | TBLRBBoxCoder) 6 | from .iou_calculators import BboxOverlaps2D, bbox_overlaps 7 | from .samplers import (BaseSampler, CombinedSampler, 8 | InstanceBalancedPosSampler, IoUBalancedNegSampler, 9 | PseudoSampler, RandomSampler, SamplingResult) 10 | from .transforms import (bbox2distance, bbox2result, bbox2roi, bbox_flip, 11 | bbox_mapping, bbox_mapping_back, distance2bbox, 12 | roi2bbox) 13 | 14 | from .transforms_obb import (poly2obb, rectpoly2obb, poly2hbb, obb2poly, obb2hbb, 15 | hbb2poly, hbb2obb, bbox2type, hbb_flip, obb_flip, poly_flip, 16 | hbb_warp, obb_warp, poly_warp, hbb_mapping, obb_mapping, 17 | poly_mapping, hbb_mapping_back, obb_mapping_back, 18 | poly_mapping_back, arb_mapping, arb_mapping_back, 19 | get_bbox_type, get_bbox_dim, get_bbox_areas, choice_by_type, arb2result, 20 | arb2roi, regular_theta, regular_obb, mintheta_obb) 21 | from .iou_calculators import OBBOverlaps, PolyOverlaps 22 | from .samplers import (OBBSamplingResult, OBBBaseSampler, OBBRandomSampler, 23 | OBBOHEMSampler) 24 | from .coder import OBB2OBBDeltaXYWHTCoder, HBB2OBBDeltaXYWHTCoder 25 | 26 | __all__ = [ 27 | 'bbox_overlaps', 'BboxOverlaps2D', 'BaseAssigner', 'MaxIoUAssigner', 28 | 'AssignResult', 'BaseSampler', 'PseudoSampler', 'RandomSampler', 29 | 'InstanceBalancedPosSampler', 'IoUBalancedNegSampler', 'CombinedSampler', 30 | 'SamplingResult', 'build_assigner', 'build_sampler', 'bbox_flip', 31 | 'bbox_mapping', 'bbox_mapping_back', 'bbox2roi', 'roi2bbox', 'bbox2result', 32 | 'distance2bbox', 'bbox2distance', 'build_bbox_coder', 'BaseBBoxCoder', 33 | 'PseudoBBoxCoder', 'DeltaXYWHBBoxCoder', 'TBLRBBoxCoder', 34 | 'CenterRegionAssigner', 35 | 36 | 'poly2obb', 'rectpoly2obb', 'poly2hbb', 'obb2poly', 'obb2hbb', 'hbb2poly', 37 | 'hbb2obb', 'bbox2type', 'hbb_flip', 'obb_flip', 'poly_flip', 'hbb_warp', 'obb_warp', 38 | 'poly_warp', 'hbb_mapping', 'obb_mapping', 'poly_mapping', 'hbb_mapping_back', 39 | 'obb_mapping_back', 'poly_mapping_back', 'get_bbox_type', 'get_bbox_dim', 'get_bbox_areas', 40 | 'choice_by_type', 'arb2roi', 'arb2result', 'arb_mapping', 'arb_mapping_back', 41 | 'OBBOverlaps', 'PolyOverlaps', 'OBBSamplingResult', 'OBBBaseSampler', 'OBBRandomSampler', 42 | 'OBBOHEMSampler', 'OBB2OBBDeltaXYWHTCoder', 'HBB2OBBDeltaXYWHTCoder', 'regular_theta', 43 | 'regular_obb', 'mintheta_obb' 44 | ] 45 | -------------------------------------------------------------------------------- /mmdet/ops/masked_conv/src/cuda/masked_conv2d_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | int MaskedIm2colForwardLaucher(const at::Tensor im, const int height, 7 | const int width, const int channels, 8 | const int kernel_h, const int kernel_w, 9 | const int pad_h, const int pad_w, 10 | const at::Tensor mask_h_idx, 11 | const at::Tensor mask_w_idx, const int mask_cnt, 12 | at::Tensor col); 13 | 14 | int MaskedCol2imForwardLaucher(const at::Tensor col, const int height, 15 | const int width, const int channels, 16 | const at::Tensor mask_h_idx, 17 | const at::Tensor mask_w_idx, const int mask_cnt, 18 | at::Tensor im); 19 | 20 | #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ") 21 | #define CHECK_CONTIGUOUS(x) \ 22 | TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 23 | #define CHECK_INPUT(x) \ 24 | CHECK_CUDA(x); \ 25 | CHECK_CONTIGUOUS(x) 26 | 27 | int masked_im2col_forward_cuda(const at::Tensor im, const at::Tensor mask_h_idx, 28 | const at::Tensor mask_w_idx, const int kernel_h, 29 | const int kernel_w, const int pad_h, 30 | const int pad_w, at::Tensor col) { 31 | CHECK_INPUT(im); 32 | CHECK_INPUT(mask_h_idx); 33 | CHECK_INPUT(mask_w_idx); 34 | CHECK_INPUT(col); 35 | // im: (n, ic, h, w), kernel size (kh, kw) 36 | // kernel: (oc, ic * kh * kw), col: (kh * kw * ic, ow * oh) 37 | at::DeviceGuard guard(im.device()); 38 | 39 | int channels = im.size(1); 40 | int height = im.size(2); 41 | int width = im.size(3); 42 | int mask_cnt = mask_h_idx.size(0); 43 | 44 | MaskedIm2colForwardLaucher(im, height, width, channels, kernel_h, kernel_w, 45 | pad_h, pad_w, mask_h_idx, mask_w_idx, mask_cnt, 46 | col); 47 | 48 | return 1; 49 | } 50 | 51 | int masked_col2im_forward_cuda(const at::Tensor col, 52 | const at::Tensor mask_h_idx, 53 | const at::Tensor mask_w_idx, int height, 54 | int width, int channels, at::Tensor im) { 55 | CHECK_INPUT(col); 56 | CHECK_INPUT(mask_h_idx); 57 | CHECK_INPUT(mask_w_idx); 58 | CHECK_INPUT(im); 59 | // im: (n, ic, h, w), kernel size (kh, kw) 60 | // kernel: (oc, ic * kh * kh), col: (kh * kw * ic, ow * oh) 61 | at::DeviceGuard guard(col.device()); 62 | 63 | int mask_cnt = mask_h_idx.size(0); 64 | 65 | MaskedCol2imForwardLaucher(col, height, width, channels, mask_h_idx, 66 | mask_w_idx, mask_cnt, im); 67 | 68 | return 1; 69 | } 70 | -------------------------------------------------------------------------------- /BboxToolkit/BboxToolkit/datasets/io.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import pickle 4 | import time 5 | import numpy as np 6 | 7 | from multiprocessing import Pool 8 | 9 | from ..utils import get_bbox_dim 10 | from .misc import read_img_info, change_cls_order, get_classes 11 | 12 | 13 | def load_imgs(img_dir, ann_dir=None, classes=None, nproc=10, 14 | def_bbox_type='poly'): 15 | assert def_bbox_type in ['hbb', 'obb', 'poly'] 16 | if ann_dir is not None: 17 | print('ann_dir is no use in load_pseudo function') 18 | 19 | print('Starting loading images information') 20 | start_time = time.time() 21 | imgpaths = [osp.join(img_dir, imgfile) 22 | for imgfile in os.listdir(img_dir)] 23 | if nproc > 1: 24 | pool = Pool(nproc) 25 | infos = pool.map(read_img_info, imgpaths) 26 | pool.close() 27 | else: 28 | infos = list(map(read_img_info, imgpaths)) 29 | 30 | contents = [] 31 | for info in infos: 32 | if info is None: 33 | continue 34 | bbox_dim = get_bbox_dim(def_bbox_type) 35 | bboxes = np.zeros((0, bbox_dim), dtype=np.float) 36 | labels = np.zeros((0, ), dtype=np.int) 37 | info['ann'] = dict(bboxes=bboxes, labels=labels) 38 | contents.append(info) 39 | classes = [] if classes is None else classes 40 | end_time = time.time() 41 | print(f'Finishing loading images, get {len(contents)} iamges,', 42 | f'using {end_time-start_time:.3f}s.') 43 | return contents, classes 44 | 45 | 46 | def load_pkl(ann_dir, img_dir=None, classes=None, nproc=10): 47 | print('Starting loading pkl information') 48 | start_time = time.time() 49 | data = pickle.load(open(ann_dir, 'rb')) 50 | old_classes, contents = data['cls'], data['content'] 51 | 52 | if img_dir is not None: 53 | imgpaths = [osp.join(img_dir, content['filename']) 54 | for content in contents] 55 | if nproc > 1: 56 | pool = Pool(nproc) 57 | infos = pool.map(read_img_info, imgpaths) 58 | pool.close() 59 | else: 60 | infos = list(map(read_img_info, imgpaths)) 61 | 62 | for info, content in zip(infos, contents): 63 | content.update(info) 64 | 65 | if classes is None: 66 | classes = old_classes 67 | else: 68 | classes = get_classes(classes) 69 | change_cls_order(contents, old_classes, classes) 70 | end_time = time.time() 71 | print(f'Finishing loading pkl, get {len(contents)} iamges,', 72 | f'using {end_time-start_time:.3f}s.') 73 | return contents, classes 74 | 75 | 76 | def save_pkl(save_dir, contents, classes): 77 | assert save_dir.endswith('.pkl') 78 | filepath = osp.split(save_dir)[0] 79 | if not osp.exists(filepath): 80 | os.makedirs(filepath) 81 | 82 | data = dict(cls=classes, content=contents) 83 | pickle.dump(data, open(save_dir, 'wb')) 84 | -------------------------------------------------------------------------------- /mmdet/core/bbox/iou_calculators/obb/obbiou_calculator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import BboxToolkit as bt 3 | 4 | from mmdet.ops import obb_overlaps 5 | from ..builder import IOU_CALCULATORS 6 | 7 | 8 | @IOU_CALCULATORS.register_module() 9 | class OBBOverlaps(object): 10 | """2D IoU Calculator""" 11 | 12 | def __call__(self, bboxes1, bboxes2, mode='iou', is_aligned=False): 13 | """Calculate IoU between 2D bboxes 14 | 15 | Args: 16 | bboxes1 (Tensor): bboxes have shape (m, 4) in 17 | format, or shape (m, 5) in format. 18 | bboxes2 (Tensor): bboxes have shape (m, 4) in 19 | format, shape (m, 5) in format, or be 20 | empty. If is_aligned is ``True``, then m and n must be equal. 21 | mode (str): "iou" (intersection over union) or iof (intersection 22 | over foreground). 23 | 24 | Returns: 25 | ious(Tensor): shape (m, n) if is_aligned == False else shape (m, 1) 26 | """ 27 | assert bboxes1.size(-1) in [0, 5, 6, 11] 28 | assert bboxes2.size(-1) in [0, 5, 6, 11] 29 | if bboxes2.size(-1) == 6: 30 | bboxes2 = bboxes2[..., :5] 31 | if bboxes1.size(-1) == 6: 32 | bboxes1 = bboxes1[..., :5] 33 | return obb_overlaps(bboxes1, bboxes2, mode, is_aligned) 34 | 35 | def __repr__(self): 36 | """str: a string describing the module""" 37 | repr_str = self.__class__.__name__ + '()' 38 | return repr_str 39 | 40 | 41 | @IOU_CALCULATORS.register_module() 42 | class PolyOverlaps(object): 43 | """2D IoU Calculator""" 44 | 45 | def __call__(self, bboxes1, bboxes2, mode='iou', is_aligned=False): 46 | """Calculate IoU between 2D bboxes 47 | 48 | Args: 49 | bboxes1 (Tensor): bboxes have shape (m, 4) in 50 | format, or shape (m, 5) in format. 51 | bboxes2 (Tensor): bboxes have shape (m, 4) in 52 | format, shape (m, 5) in format, or be 53 | empty. If is_aligned is ``True``, then m and n must be equal. 54 | mode (str): "iou" (intersection over union) or iof (intersection 55 | over foreground). 56 | 57 | Returns: 58 | ious(Tensor): shape (m, n) if is_aligned == False else shape (m, 1) 59 | """ 60 | assert bboxes1.size(-1) in [0, 8, 9] 61 | assert bboxes2.size(-1) in [0, 8, 9] 62 | if bboxes2.size(-1) == 9: 63 | bboxes2 = bboxes2[..., :8] 64 | if bboxes1.size(-1) == 9: 65 | bboxes1 = bboxes1[..., :8] 66 | return bt.bbox_overlaps(bboxes1, bboxes2, mode, is_aligned) 67 | 68 | def __repr__(self): 69 | """str: a string describing the module""" 70 | repr_str = self.__class__.__name__ + '()' 71 | return repr_str 72 | -------------------------------------------------------------------------------- /mmdet/ops/carafe/src/cuda/carafe_naive_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | #include 6 | 7 | int CARAFENAIVEForwardLaucher(const at::Tensor features, const at::Tensor masks, 8 | const int kernel_size, const int group_size, 9 | const int scale_factor, const int batch_size, 10 | const int channels, const int height, 11 | const int width, at::Tensor output); 12 | 13 | int CARAFENAIVEBackwardLaucher(const at::Tensor top_grad, 14 | const at::Tensor features, 15 | const at::Tensor masks, const int kernel_size, 16 | const int group_size, const int scale_factor, 17 | const int batch_size, const int channels, 18 | const int height, const int width, 19 | at::Tensor bottom_grad, at::Tensor mask_grad); 20 | 21 | #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ") 22 | #define CHECK_CONTIGUOUS(x) \ 23 | TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ") 24 | #define CHECK_INPUT(x) \ 25 | CHECK_CUDA(x); \ 26 | CHECK_CONTIGUOUS(x) 27 | 28 | int carafe_naive_forward_cuda(at::Tensor features, at::Tensor masks, 29 | int kernel_size, int group_size, int scale_factor, 30 | at::Tensor output) { 31 | CHECK_INPUT(features); 32 | CHECK_INPUT(masks); 33 | CHECK_INPUT(output); 34 | at::DeviceGuard guard(features.device()); 35 | 36 | int batch_size = output.size(0); 37 | int num_channels = output.size(1); 38 | int data_height = output.size(2); 39 | int data_width = output.size(3); 40 | 41 | CARAFENAIVEForwardLaucher(features, masks, kernel_size, group_size, 42 | scale_factor, batch_size, num_channels, data_height, 43 | data_width, output); 44 | 45 | return 1; 46 | } 47 | 48 | int carafe_naive_backward_cuda(at::Tensor top_grad, at::Tensor features, 49 | at::Tensor masks, int kernel_size, 50 | int group_size, int scale_factor, 51 | at::Tensor bottom_grad, at::Tensor mask_grad) { 52 | CHECK_INPUT(top_grad); 53 | CHECK_INPUT(features); 54 | CHECK_INPUT(masks); 55 | CHECK_INPUT(bottom_grad); 56 | CHECK_INPUT(mask_grad); 57 | at::DeviceGuard guard(top_grad.device()); 58 | 59 | int batch_size = top_grad.size(0); 60 | int num_channels = top_grad.size(1); 61 | int data_height = top_grad.size(2); 62 | int data_width = top_grad.size(3); 63 | 64 | CARAFENAIVEBackwardLaucher(top_grad, features, masks, kernel_size, group_size, 65 | scale_factor, batch_size, num_channels, 66 | data_height, data_width, bottom_grad, mask_grad); 67 | 68 | return 1; 69 | } 70 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/roi_extractors/obb/obb_base_roi_extractor.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | from torch.nn.modules.utils import _pair 7 | from mmdet import ops 8 | 9 | 10 | class OBBBaseRoIExtractor(nn.Module, metaclass=ABCMeta): 11 | """Base class for RoI extractor. 12 | 13 | Args: 14 | roi_layer (dict): Specify RoI layer type and arguments. 15 | out_channels (int): Output channels of RoI layers. 16 | featmap_strides (int): Strides of input feature maps. 17 | """ 18 | 19 | def __init__(self, roi_layer, out_channels, featmap_strides): 20 | super(OBBBaseRoIExtractor, self).__init__() 21 | self.roi_layers = self.build_roi_layers(roi_layer, featmap_strides) 22 | self.out_channels = out_channels 23 | self.featmap_strides = featmap_strides 24 | self.fp16_enabled = False 25 | 26 | @property 27 | def num_inputs(self): 28 | """int: Number of input feature maps.""" 29 | return len(self.featmap_strides) 30 | 31 | def init_weights(self): 32 | pass 33 | 34 | def build_roi_layers(self, layer_cfg, featmap_strides): 35 | """Build RoI operator to extract feature from each level feature map. 36 | 37 | Args: 38 | layer_cfg (dict): Dictionary to construct and config RoI layer 39 | operation. Options are modules under ``mmdet/ops`` such as 40 | ``RoIAlign``. 41 | featmap_strides (int): The stride of input feature map w.r.t to the 42 | original image size, which would be used to scale RoI 43 | coordinate (original image coordinate system) to feature 44 | coordinate system. 45 | 46 | Returns: 47 | nn.ModuleList: The RoI extractor modules for each level feature 48 | map. 49 | """ 50 | 51 | cfg = layer_cfg.copy() 52 | layer_type = cfg.pop('type') 53 | assert hasattr(ops, layer_type) 54 | layer_cls = getattr(ops, layer_type) 55 | roi_layers = nn.ModuleList( 56 | [layer_cls(spatial_scale=1 / s, **cfg) for s in featmap_strides]) 57 | return roi_layers 58 | 59 | def roi_rescale(self, rois, scale_factor): 60 | """Scale RoI coordinates by scale factor. 61 | 62 | Args: 63 | rois (torch.Tensor): RoI (Region of Interest), shape (n, 6) 64 | scale_factor (float): Scale factor that RoI will be multiplied by. 65 | 66 | Returns: 67 | torch.Tensor: Scaled RoI. 68 | """ 69 | if scale_factor is None: 70 | return rois 71 | h_scale_factor, w_scale_factor = _pair(scale_factor) 72 | new_rois = rois.clone() 73 | new_rois[:, 3] = w_scale_factor * new_rois[:, 3] 74 | new_rois[:, 4] = h_scale_factor * new_rois[:, 4] 75 | return new_rois 76 | 77 | @abstractmethod 78 | def forward(self, feats, rois, roi_scale_factor=None): 79 | pass 80 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/random_sampler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from ..builder import BBOX_SAMPLERS 4 | from .base_sampler import BaseSampler 5 | 6 | 7 | @BBOX_SAMPLERS.register_module() 8 | class RandomSampler(BaseSampler): 9 | """Random sampler 10 | 11 | Args: 12 | num (int): Number of samples 13 | pos_fraction (float): Fraction of positive samples 14 | neg_pos_up (int, optional): Upper bound number of negative and 15 | positive samples. Defaults to -1. 16 | add_gt_as_proposals (bool, optional): Whether to add ground truth 17 | boxes as proposals. Defaults to True. 18 | """ 19 | 20 | def __init__(self, 21 | num, 22 | pos_fraction, 23 | neg_pos_ub=-1, 24 | add_gt_as_proposals=True, 25 | **kwargs): 26 | from mmdet.core.bbox import demodata 27 | super(RandomSampler, self).__init__(num, pos_fraction, neg_pos_ub, 28 | add_gt_as_proposals) 29 | self.rng = demodata.ensure_rng(kwargs.get('rng', None)) 30 | 31 | def random_choice(self, gallery, num): 32 | """Random select some elements from the gallery. 33 | 34 | If `gallery` is a Tensor, the returned indices will be a Tensor; 35 | If `gallery` is a ndarray or list, the returned indices will be a 36 | ndarray. 37 | 38 | Args: 39 | gallery (Tensor | ndarray | list): indices pool. 40 | num (int): expected sample num. 41 | 42 | Returns: 43 | Tensor or ndarray: sampled indices. 44 | """ 45 | assert len(gallery) >= num 46 | 47 | is_tensor = isinstance(gallery, torch.Tensor) 48 | if not is_tensor: 49 | gallery = torch.tensor( 50 | gallery, dtype=torch.long, device=torch.cuda.current_device()) 51 | perm = torch.randperm(gallery.numel(), device=gallery.device)[:num] 52 | rand_inds = gallery[perm] 53 | if not is_tensor: 54 | rand_inds = rand_inds.cpu().numpy() 55 | return rand_inds 56 | 57 | def _sample_pos(self, assign_result, num_expected, **kwargs): 58 | """Randomly sample some positive samples.""" 59 | pos_inds = torch.nonzero(assign_result.gt_inds > 0, as_tuple=False) 60 | if pos_inds.numel() != 0: 61 | pos_inds = pos_inds.squeeze(1) 62 | if pos_inds.numel() <= num_expected: 63 | return pos_inds 64 | else: 65 | return self.random_choice(pos_inds, num_expected) 66 | 67 | def _sample_neg(self, assign_result, num_expected, **kwargs): 68 | """Randomly sample some negative samples.""" 69 | neg_inds = torch.nonzero(assign_result.gt_inds == 0, as_tuple=False) 70 | if neg_inds.numel() != 0: 71 | neg_inds = neg_inds.squeeze(1) 72 | if len(neg_inds) <= num_expected: 73 | return neg_inds 74 | else: 75 | return self.random_choice(neg_inds, num_expected) 76 | -------------------------------------------------------------------------------- /mmdet/core/bbox/samplers/obb/obb_random_sampler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from mmdet.core.bbox.builder import BBOX_SAMPLERS 4 | from .obb_base_sampler import OBBBaseSampler 5 | 6 | 7 | @BBOX_SAMPLERS.register_module() 8 | class OBBRandomSampler(OBBBaseSampler): 9 | """Random sampler 10 | 11 | Args: 12 | num (int): Number of samples 13 | pos_fraction (float): Fraction of positive samples 14 | neg_pos_up (int, optional): Upper bound number of negative and 15 | positive samples. Defaults to -1. 16 | add_gt_as_proposals (bool, optional): Whether to add ground truth 17 | boxes as proposals. Defaults to True. 18 | """ 19 | 20 | def __init__(self, 21 | num, 22 | pos_fraction, 23 | neg_pos_ub=-1, 24 | add_gt_as_proposals=True, 25 | **kwargs): 26 | from mmdet.core.bbox import demodata 27 | super(OBBRandomSampler, self).__init__(num, pos_fraction, neg_pos_ub, 28 | add_gt_as_proposals) 29 | self.rng = demodata.ensure_rng(kwargs.get('rng', None)) 30 | 31 | def random_choice(self, gallery, num): 32 | """Random select some elements from the gallery. 33 | 34 | If `gallery` is a Tensor, the returned indices will be a Tensor; 35 | If `gallery` is a ndarray or list, the returned indices will be a 36 | ndarray. 37 | 38 | Args: 39 | gallery (Tensor | ndarray | list): indices pool. 40 | num (int): expected sample num. 41 | 42 | Returns: 43 | Tensor or ndarray: sampled indices. 44 | """ 45 | assert len(gallery) >= num 46 | 47 | is_tensor = isinstance(gallery, torch.Tensor) 48 | if not is_tensor: 49 | gallery = torch.tensor( 50 | gallery, dtype=torch.long, device=torch.cuda.current_device()) 51 | perm = torch.randperm(gallery.numel(), device=gallery.device)[:num] 52 | rand_inds = gallery[perm] 53 | if not is_tensor: 54 | rand_inds = rand_inds.cpu().numpy() 55 | return rand_inds 56 | 57 | def _sample_pos(self, assign_result, num_expected, **kwargs): 58 | """Randomly sample some positive samples.""" 59 | pos_inds = torch.nonzero(assign_result.gt_inds > 0, as_tuple=False) 60 | if pos_inds.numel() != 0: 61 | pos_inds = pos_inds.squeeze(1) 62 | if pos_inds.numel() <= num_expected: 63 | return pos_inds 64 | else: 65 | return self.random_choice(pos_inds, num_expected) 66 | 67 | def _sample_neg(self, assign_result, num_expected, **kwargs): 68 | """Randomly sample some negative samples.""" 69 | neg_inds = torch.nonzero(assign_result.gt_inds == 0, as_tuple=False) 70 | if neg_inds.numel() != 0: 71 | neg_inds = neg_inds.squeeze(1) 72 | if len(neg_inds) <= num_expected: 73 | return neg_inds 74 | else: 75 | return self.random_choice(neg_inds, num_expected) 76 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/obb/obb_base_roi_head.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | import torch.nn as nn 4 | 5 | from mmdet.models.builder import build_shared_head 6 | 7 | 8 | class OBBBaseRoIHead(nn.Module, metaclass=ABCMeta): 9 | """Base class for RoIHeads""" 10 | 11 | def __init__(self, 12 | bbox_roi_extractor=None, 13 | bbox_head=None, 14 | shared_head=None, 15 | train_cfg=None, 16 | test_cfg=None): 17 | super(OBBBaseRoIHead, self).__init__() 18 | self.train_cfg = train_cfg 19 | self.test_cfg = test_cfg 20 | if shared_head is not None: 21 | self.shared_head = build_shared_head(shared_head) 22 | 23 | if bbox_head is not None: 24 | self.init_bbox_head(bbox_roi_extractor, bbox_head) 25 | 26 | self.init_assigner_sampler() 27 | 28 | @property 29 | def with_bbox(self): 30 | """bool: whether the RoI head contains a `bbox_head`""" 31 | return hasattr(self, 'bbox_head') and self.bbox_head is not None 32 | 33 | @property 34 | def with_shared_head(self): 35 | """bool: whether the RoI head contains a `shared_head`""" 36 | return hasattr(self, 'shared_head') and self.shared_head is not None 37 | 38 | @abstractmethod 39 | def init_weights(self, pretrained): 40 | """Initialize the weights in head 41 | 42 | Args: 43 | pretrained (str, optional): Path to pre-trained weights. 44 | Defaults to None. 45 | """ 46 | pass 47 | 48 | @abstractmethod 49 | def init_bbox_head(self): 50 | """Initialize ``bbox_head``""" 51 | pass 52 | 53 | @abstractmethod 54 | def init_assigner_sampler(self): 55 | """Initialize assigner and sampler""" 56 | pass 57 | 58 | @abstractmethod 59 | def forward_train(self, 60 | x, 61 | img_meta, 62 | proposal_list, 63 | gt_bboxes, 64 | gt_obboxes, 65 | gt_labels, 66 | gt_bboxes_ignore=None, 67 | gt_obboxes_ignore=None, 68 | **kwargs): 69 | """Forward function during training""" 70 | pass 71 | 72 | async def async_simple_test(self, x, img_meta, **kwargs): 73 | """Asynchronized test function""" 74 | raise NotImplementedError 75 | 76 | def simple_test(self, 77 | x, 78 | proposal_list, 79 | img_meta, 80 | proposals=None, 81 | rescale=False, 82 | **kwargs): 83 | """Test without augmentation.""" 84 | pass 85 | 86 | def aug_test(self, x, proposal_list, img_metas, rescale=False, **kwargs): 87 | """Test with augmentations. 88 | 89 | If rescale is False, then returned bboxes and masks will fit the scale 90 | of imgs[0]. 91 | """ 92 | pass 93 | -------------------------------------------------------------------------------- /mmdet/models/dense_heads/nasfcos_head.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | import torch.nn as nn 4 | from mmcv.cnn import (ConvModule, Scale, bias_init_with_prob, 5 | caffe2_xavier_init, normal_init) 6 | 7 | from mmdet.models.dense_heads.fcos_head import FCOSHead 8 | from ..builder import HEADS 9 | 10 | 11 | @HEADS.register_module() 12 | class NASFCOSHead(FCOSHead): 13 | """Anchor-free head used in `NASFCOS `_. 14 | 15 | It is quite similar with FCOS head, except for the searched structure 16 | of classification branch and bbox regression branch, where a structure 17 | of "dconv3x3, conv3x3, dconv3x3, conv1x1" is utilized instead. 18 | 19 | """ 20 | 21 | def _init_layers(self): 22 | """Initialize layers of the head.""" 23 | dconv3x3_config = dict( 24 | type='DCNv2', 25 | kernel_size=3, 26 | use_bias=True, 27 | deformable_groups=2, 28 | padding=1) 29 | conv3x3_config = dict(type='Conv', kernel_size=3, padding=1) 30 | conv1x1_config = dict(type='Conv', kernel_size=1) 31 | 32 | self.arch_config = [ 33 | dconv3x3_config, conv3x3_config, dconv3x3_config, conv1x1_config 34 | ] 35 | self.cls_convs = nn.ModuleList() 36 | self.reg_convs = nn.ModuleList() 37 | for i, op_ in enumerate(self.arch_config): 38 | op = copy.deepcopy(op_) 39 | chn = self.in_channels if i == 0 else self.feat_channels 40 | assert isinstance(op, dict) 41 | use_bias = op.pop('use_bias', False) 42 | padding = op.pop('padding', 0) 43 | kernel_size = op.pop('kernel_size') 44 | module = ConvModule( 45 | chn, 46 | self.feat_channels, 47 | kernel_size, 48 | stride=1, 49 | padding=padding, 50 | norm_cfg=self.norm_cfg, 51 | bias=use_bias, 52 | conv_cfg=op) 53 | 54 | self.cls_convs.append(copy.deepcopy(module)) 55 | self.reg_convs.append(copy.deepcopy(module)) 56 | 57 | self.fcos_cls = nn.Conv2d( 58 | self.feat_channels, self.cls_out_channels, 3, padding=1) 59 | self.fcos_reg = nn.Conv2d(self.feat_channels, 4, 3, padding=1) 60 | self.fcos_centerness = nn.Conv2d(self.feat_channels, 1, 3, padding=1) 61 | 62 | self.scales = nn.ModuleList([Scale(1.0) for _ in self.strides]) 63 | 64 | def init_weights(self): 65 | """Initialize weights of the head.""" 66 | # retinanet_bias_init 67 | bias_cls = bias_init_with_prob(0.01) 68 | normal_init(self.fcos_reg, std=0.01) 69 | normal_init(self.fcos_centerness, std=0.01) 70 | normal_init(self.fcos_cls, std=0.01, bias=bias_cls) 71 | 72 | for branch in [self.cls_convs, self.reg_convs]: 73 | for module in branch.modules(): 74 | if isinstance(module, ConvModule) \ 75 | and isinstance(module.conv, nn.Conv2d): 76 | caffe2_xavier_init(module.conv) 77 | -------------------------------------------------------------------------------- /mmdet/models/roi_heads/roi_extractors/base_roi_extractor.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | from mmdet import ops 7 | 8 | 9 | class BaseRoIExtractor(nn.Module, metaclass=ABCMeta): 10 | """Base class for RoI extractor. 11 | 12 | Args: 13 | roi_layer (dict): Specify RoI layer type and arguments. 14 | out_channels (int): Output channels of RoI layers. 15 | featmap_strides (int): Strides of input feature maps. 16 | """ 17 | 18 | def __init__(self, roi_layer, out_channels, featmap_strides): 19 | super(BaseRoIExtractor, self).__init__() 20 | self.roi_layers = self.build_roi_layers(roi_layer, featmap_strides) 21 | self.out_channels = out_channels 22 | self.featmap_strides = featmap_strides 23 | self.fp16_enabled = False 24 | 25 | @property 26 | def num_inputs(self): 27 | """int: Number of input feature maps.""" 28 | return len(self.featmap_strides) 29 | 30 | def init_weights(self): 31 | pass 32 | 33 | def build_roi_layers(self, layer_cfg, featmap_strides): 34 | """Build RoI operator to extract feature from each level feature map. 35 | 36 | Args: 37 | layer_cfg (dict): Dictionary to construct and config RoI layer 38 | operation. Options are modules under ``mmdet/ops`` such as 39 | ``RoIAlign``. 40 | featmap_strides (int): The stride of input feature map w.r.t to the 41 | original image size, which would be used to scale RoI 42 | coordinate (original image coordinate system) to feature 43 | coordinate system. 44 | 45 | Returns: 46 | nn.ModuleList: The RoI extractor modules for each level feature 47 | map. 48 | """ 49 | 50 | cfg = layer_cfg.copy() 51 | layer_type = cfg.pop('type') 52 | assert hasattr(ops, layer_type) 53 | layer_cls = getattr(ops, layer_type) 54 | roi_layers = nn.ModuleList( 55 | [layer_cls(spatial_scale=1 / s, **cfg) for s in featmap_strides]) 56 | return roi_layers 57 | 58 | def roi_rescale(self, rois, scale_factor): 59 | """Scale RoI coordinates by scale factor. 60 | 61 | Args: 62 | rois (torch.Tensor): RoI (Region of Interest), shape (n, 5) 63 | scale_factor (float): Scale factor that RoI will be multiplied by. 64 | 65 | Returns: 66 | torch.Tensor: Scaled RoI. 67 | """ 68 | 69 | cx = (rois[:, 1] + rois[:, 3]) * 0.5 70 | cy = (rois[:, 2] + rois[:, 4]) * 0.5 71 | w = rois[:, 3] - rois[:, 1] 72 | h = rois[:, 4] - rois[:, 2] 73 | new_w = w * scale_factor 74 | new_h = h * scale_factor 75 | x1 = cx - new_w * 0.5 76 | x2 = cx + new_w * 0.5 77 | y1 = cy - new_h * 0.5 78 | y2 = cy + new_h * 0.5 79 | new_rois = torch.stack((rois[:, 0], x1, y1, x2, y2), dim=-1) 80 | return new_rois 81 | 82 | @abstractmethod 83 | def forward(self, feats, rois, roi_scale_factor=None): 84 | pass 85 | -------------------------------------------------------------------------------- /mmdet/ops/roi_align_rotated/roi_align_rotated.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch import nn 3 | from torch.autograd import Function 4 | from torch.autograd.function import once_differentiable 5 | from torch.nn.modules.utils import _pair 6 | 7 | from . import roi_align_rotated_ext 8 | 9 | 10 | class RoIAlignRotatedFunction(Function): 11 | 12 | @staticmethod 13 | def forward(ctx, 14 | features, 15 | rois, 16 | out_size, 17 | spatial_scale, 18 | sample_num=0, 19 | aligned=True): 20 | out_h, out_w = _pair(out_size) 21 | assert isinstance(out_h, int) and isinstance(out_w, int) 22 | ctx.spatial_scale = spatial_scale 23 | ctx.sample_num = sample_num 24 | ctx.save_for_backward(rois) 25 | ctx.feature_size = features.size() 26 | ctx.aligned = aligned 27 | 28 | output = roi_align_rotated_ext.forward( 29 | features, rois, spatial_scale, out_h, out_w, sample_num, aligned) 30 | 31 | return output 32 | 33 | @staticmethod 34 | @once_differentiable 35 | def backward(ctx, grad_output): 36 | feature_size = ctx.feature_size 37 | spatial_scale = ctx.spatial_scale 38 | sample_num = ctx.sample_num 39 | rois = ctx.saved_tensors[0] 40 | aligned = ctx.aligned 41 | assert feature_size is not None 42 | 43 | batch_size, num_channels, data_height, data_width = feature_size 44 | out_w = grad_output.size(3) 45 | out_h = grad_output.size(2) 46 | 47 | grad_input = grad_rois = None 48 | grad_input = roi_align_rotated_ext.backward( 49 | grad_output, rois, spatial_scale, out_h, out_w, 50 | batch_size, num_channels, data_height, data_width, 51 | sample_num, aligned) 52 | 53 | return grad_input, grad_rois, None, None, None, None 54 | 55 | 56 | roi_align_rotated = RoIAlignRotatedFunction.apply 57 | 58 | 59 | class RoIAlignRotated(nn.Module): 60 | 61 | def __init__(self, 62 | out_size, 63 | spatial_scale, 64 | sample_num=0, 65 | aligned=True): 66 | super(RoIAlignRotated, self).__init__() 67 | self.out_size = _pair(out_size) 68 | self.spatial_scale = float(spatial_scale) 69 | self.sample_num = int(sample_num) 70 | self.aligned = aligned 71 | 72 | def forward(self, features, rois): 73 | assert rois.dim() == 2 and rois.size(1) == 6 74 | return roi_align_rotated(features, rois, self.out_size, 75 | self.spatial_scale, self.sample_num, 76 | self.aligned) 77 | 78 | def __repr__(self): 79 | indent_str = '\n ' 80 | format_str = self.__class__.__name__ 81 | format_str += f'({indent_str}out_size={self.out_size},' 82 | format_str += f'{indent_str}spatial_scale={self.spatial_scale},' 83 | format_str += f'{indent_str}sample_num={self.sample_num},' 84 | format_str += f'{indent_str}aligned={self.aligned},' 85 | return format_str 86 | -------------------------------------------------------------------------------- /mmdet/ops/corner_pool/corner_pool.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.autograd import Function 3 | 4 | from . import corner_pool_ext 5 | 6 | 7 | class TopPoolFunction(Function): 8 | 9 | @staticmethod 10 | def forward(ctx, input): 11 | output = corner_pool_ext.top_pool_forward(input) 12 | ctx.save_for_backward(input) 13 | return output 14 | 15 | @staticmethod 16 | def backward(ctx, grad_output): 17 | input = ctx.saved_variables[0] 18 | output = corner_pool_ext.top_pool_backward(input, grad_output) 19 | return output 20 | 21 | 22 | class BottomPoolFunction(Function): 23 | 24 | @staticmethod 25 | def forward(ctx, input): 26 | output = corner_pool_ext.bottom_pool_forward(input) 27 | ctx.save_for_backward(input) 28 | return output 29 | 30 | @staticmethod 31 | def backward(ctx, grad_output): 32 | input = ctx.saved_variables[0] 33 | output = corner_pool_ext.bottom_pool_backward(input, grad_output) 34 | return output 35 | 36 | 37 | class LeftPoolFunction(Function): 38 | 39 | @staticmethod 40 | def forward(ctx, input): 41 | output = corner_pool_ext.left_pool_forward(input) 42 | ctx.save_for_backward(input) 43 | return output 44 | 45 | @staticmethod 46 | def backward(ctx, grad_output): 47 | input = ctx.saved_variables[0] 48 | output = corner_pool_ext.left_pool_backward(input, grad_output) 49 | return output 50 | 51 | 52 | class RightPoolFunction(Function): 53 | 54 | @staticmethod 55 | def forward(ctx, input): 56 | output = corner_pool_ext.right_pool_forward(input) 57 | ctx.save_for_backward(input) 58 | return output 59 | 60 | @staticmethod 61 | def backward(ctx, grad_output): 62 | input = ctx.saved_variables[0] 63 | output = corner_pool_ext.right_pool_backward(input, grad_output) 64 | return output 65 | 66 | 67 | class CornerPool(nn.Module): 68 | """Corner Pooling. 69 | 70 | Corner Pooling is a new type of pooling layer that helps a 71 | convolutional network better localize corners of bounding boxes. 72 | 73 | Please refer to https://arxiv.org/abs/1808.01244 for more details. 74 | Code is modified from https://github.com/princeton-vl/CornerNet-Lite. 75 | 76 | Args: 77 | mode(str): Pooling orientation for the pooling layer 78 | 79 | - 'bottom': Bottom Pooling 80 | - 'left': Left Pooling 81 | - 'right': Right Pooling 82 | - 'top': Top Pooling 83 | 84 | Returns: 85 | Feature map after pooling. 86 | """ 87 | 88 | pool_functions = { 89 | 'bottom': BottomPoolFunction, 90 | 'left': LeftPoolFunction, 91 | 'right': RightPoolFunction, 92 | 'top': TopPoolFunction, 93 | } 94 | 95 | def __init__(self, mode): 96 | super(CornerPool, self).__init__() 97 | assert mode in self.pool_functions 98 | self.corner_pool = self.pool_functions[mode] 99 | 100 | def forward(self, x): 101 | return self.corner_pool.apply(x) 102 | -------------------------------------------------------------------------------- /mmdet/ops/dcn/src/deform_pool_ext.cpp: -------------------------------------------------------------------------------- 1 | // modify from 2 | // https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/blob/mmdetection/mmdet/ops/dcn/src/modulated_dcn_cuda.c 3 | 4 | // based on 5 | // author: Charles Shang 6 | // https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | #include 13 | 14 | #ifdef WITH_CUDA 15 | void deform_psroi_pooling_cuda_forward( 16 | at::Tensor input, at::Tensor bbox, at::Tensor trans, at::Tensor out, 17 | at::Tensor top_count, const int no_trans, const float spatial_scale, 18 | const int output_dim, const int group_size, const int pooled_size, 19 | const int part_size, const int sample_per_part, const float trans_std); 20 | 21 | void deform_psroi_pooling_cuda_backward( 22 | at::Tensor out_grad, at::Tensor input, at::Tensor bbox, at::Tensor trans, 23 | at::Tensor top_count, at::Tensor input_grad, at::Tensor trans_grad, 24 | const int no_trans, const float spatial_scale, const int output_dim, 25 | const int group_size, const int pooled_size, const int part_size, 26 | const int sample_per_part, const float trans_std); 27 | #endif 28 | 29 | void deform_psroi_pooling_forward( 30 | at::Tensor input, at::Tensor bbox, at::Tensor trans, at::Tensor out, 31 | at::Tensor top_count, const int no_trans, const float spatial_scale, 32 | const int output_dim, const int group_size, const int pooled_size, 33 | const int part_size, const int sample_per_part, const float trans_std) { 34 | if (input.device().is_cuda()) { 35 | #ifdef WITH_CUDA 36 | return deform_psroi_pooling_cuda_forward(input, bbox, trans, out, top_count, 37 | no_trans, spatial_scale, output_dim, group_size, pooled_size, 38 | part_size, sample_per_part, trans_std); 39 | #else 40 | AT_ERROR("deform psroi pooling is not compiled with GPU support"); 41 | #endif 42 | } 43 | AT_ERROR("deform psroi pooling is not implemented on CPU"); 44 | } 45 | 46 | void deform_psroi_pooling_backward( 47 | at::Tensor out_grad, at::Tensor input, at::Tensor bbox, at::Tensor trans, 48 | at::Tensor top_count, at::Tensor input_grad, at::Tensor trans_grad, 49 | const int no_trans, const float spatial_scale, const int output_dim, 50 | const int group_size, const int pooled_size, const int part_size, 51 | const int sample_per_part, const float trans_std) { 52 | if (input.device().is_cuda()) { 53 | #ifdef WITH_CUDA 54 | return deform_psroi_pooling_cuda_backward(out_grad, input, bbox, trans, 55 | top_count, input_grad, trans_grad, no_trans, spatial_scale, 56 | output_dim, group_size, pooled_size, part_size, sample_per_part, 57 | trans_std); 58 | #else 59 | AT_ERROR("deform psroi pooling is not compiled with GPU support"); 60 | #endif 61 | } 62 | AT_ERROR("deform psroi pooling is not implemented on CPU"); 63 | } 64 | 65 | 66 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 67 | m.def("deform_psroi_pooling_forward", &deform_psroi_pooling_forward, 68 | "deform psroi pooling forward"); 69 | m.def("deform_psroi_pooling_backward", &deform_psroi_pooling_backward, 70 | "deform psroi pooling backward"); 71 | } 72 | --------------------------------------------------------------------------------