diff --git a/README.md b/README.md index c795afe6..85fe7b23 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,41 @@ +# Cycle-CenterNet +Table structure parsing (TSP), Wired Table in the Wild (WTW). + +The unofficial Cycle-CenterNet architecture repository, based on the MMDetection fork. +![](demo/schema.jpg) +> [**Parsing table structures in the wild**]([http://openaccess.thecvf.com/content/ICCV2021/papers/Long_Parsing_Table_Structures_in_the_Wild_ICCV_2021_paper.pdf]), +> Long, R., Wang, W., Xue, N., Gao, F., Yang, Z., Wang, Y., & Xia, G. S. + +Contact: [arkhipov.ai@phystech.edu](mailto:arkhipov.ai@phystech.edu). Any questions or discussions are welcome! + +## Abstract +This paper tackles the problem of table structure parsing (TSP) from images in the wild. In contrast to existing studies that mainly focus on parsing well-aligned tabular images with simple layouts from scanned PDF documents, we aim to establish a practical table structure parsing system for real-world scenarios where tabular input images are taken or scanned with severe deformation, bending or occlusions. For designing such a system, we propose an approach named Cycle-CenterNet on the top of CenterNet with a novel cycle-pairing module to simultaneously detect and group tabular cells into structured tables. In the cycle-pairing module, a new pairing loss function is proposed for the network training. Alongside with our Cycle-CenterNet, we also present a large-scale dataset, named Wired Table in the Wild (WTW), which includes well-annotated structure parsing of multiple style tables in several scenes like photo, scanning files, web pages, etc.. In experiments, we demonstrate that our Cycle-CenterNet consistently achieves the best accuracy of table structure parsing on the new WTW dataset by 24.6% absolute improvement evaluated by the TEDS metric. A more comprehensive experimental analysis also validates the advantages of our proposed methods for the TSP task. + +## Models + +[**Link**](https://drive.google.com/file/d/1ZQZycYwWXMlfZnvC9hFUCtlqd4eOwP8v/view?usp=share_link) to download model CenterNet(ResNet backbone). [**Config**](https://github.com/ArchieAlexArkhipov/Cycle-CenterNet/blob/master/configs/centernet/wtw_centernet_resnet34_dcnv2_150e_coco.py) + +[**Link**](https://drive.google.com/file/d/1aZ2IF0tQq1Ino4QAsmrk8jcJ85urqsFE/view?usp=share_link) to download model CenterNet(DLA backbone). [**Config**](https://github.com/ArchieAlexArkhipov/Cycle-CenterNet/blob/master/configs/centernet/wtw_centernet_dla34_dcnv2_150e_coco.py) + +[**Link**](https://drive.google.com/file/d/1oeqtA84eF_KJi953f8qX8cggDBjOVS3A/view?usp=share_link) to download model Cycle-CenterNet(DLA backbone) trained on bounding boxes. [**Config**](https://github.com/ArchieAlexArkhipov/Cycle-CenterNet/blob/master/configs/centernet/wtw_bbox_cycle_centernet_dla34_dcnv2_150e_coco.py) + +[**Link**](https://drive.google.com/file/d/1taFK_co-9ofbL1pivhKrqHf8-icb1ZaK/view?usp=share_link) **to download model Cycle-CenterNet(DLA backbone) trained on bounding quadrangles**. [**Config**](https://github.com/ArchieAlexArkhipov/Cycle-CenterNet/blob/master/configs/centernet/wtw_quad_cycle_centernet_dla34_dcnv2_150e_coco.py) + +## Installation + +```conda create --name openmmlab python=3.8 -y``` + +```conda activate openmmlab``` + +```pip install torch==1.8.1+cu101 torchvision==0.9.1+cu101 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html``` + +```pip install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.8/index.html``` + +```pip install -e .``` + +```conda install -n openmmlab ipykernel --update-deps --force-reinstall``` + +# MMDetection README
 
diff --git a/configs/centernet/centernet_resnet18_dcnv2_140e_coco.py b/configs/centernet/centernet_resnet18_dcnv2_140e_coco.py index b8a0bb10..898e1db3 100644 --- a/configs/centernet/centernet_resnet18_dcnv2_140e_coco.py +++ b/configs/centernet/centernet_resnet18_dcnv2_140e_coco.py @@ -1,92 +1,110 @@ _base_ = [ - '../_base_/datasets/coco_detection.py', - '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py' + "../_base_/datasets/coco_detection.py", + "../_base_/schedules/schedule_1x.py", + "../_base_/default_runtime.py", ] model = dict( - type='CenterNet', + type="CenterNet", backbone=dict( - type='ResNet', + type="ResNet", depth=18, norm_eval=False, - norm_cfg=dict(type='BN'), - init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet18')), + norm_cfg=dict(type="BN"), + init_cfg=dict(type="Pretrained", checkpoint="torchvision://resnet18"), + ), neck=dict( - type='CTResNetNeck', + type="CTResNetNeck", in_channel=512, num_deconv_filters=(256, 128, 64), num_deconv_kernels=(4, 4, 4), - use_dcn=True), + use_dcn=True, + ), bbox_head=dict( - type='CenterNetHead', + type="CenterNetHead", num_classes=80, in_channel=64, feat_channel=64, - loss_center_heatmap=dict(type='GaussianFocalLoss', loss_weight=1.0), - loss_wh=dict(type='L1Loss', loss_weight=0.1), - loss_offset=dict(type='L1Loss', loss_weight=1.0)), + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_wh=dict(type="L1Loss", loss_weight=0.1), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + ), train_cfg=None, - test_cfg=dict(topk=100, local_maximum_kernel=3, max_per_img=100)) + test_cfg=dict(topk=100, local_maximum_kernel=3, max_per_img=100), +) # We fixed the incorrect img_norm_cfg problem in the source code. -img_norm_cfg = dict( - mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) +img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) train_pipeline = [ - dict(type='LoadImageFromFile', to_float32=True, color_type='color'), - dict(type='LoadAnnotations', with_bbox=True), + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), dict( - type='PhotoMetricDistortion', + type="PhotoMetricDistortion", brightness_delta=32, contrast_range=(0.5, 1.5), saturation_range=(0.5, 1.5), - hue_delta=18), + hue_delta=18, + ), dict( - type='RandomCenterCropPad', + type="RandomCenterCropPad", crop_size=(512, 512), ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), mean=[0, 0, 0], std=[1, 1, 1], to_rgb=True, - test_pad_mode=None), - dict(type='Resize', img_scale=(512, 512), keep_ratio=True), - dict(type='RandomFlip', flip_ratio=0.5), - dict(type='Normalize', **img_norm_cfg), - dict(type='DefaultFormatBundle'), - dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']) + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(512, 512), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict(type="Normalize", **img_norm_cfg), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), ] test_pipeline = [ - dict(type='LoadImageFromFile', to_float32=True), + dict(type="LoadImageFromFile", to_float32=True), dict( - type='MultiScaleFlipAug', + type="MultiScaleFlipAug", scale_factor=1.0, flip=False, transforms=[ - dict(type='Resize', keep_ratio=True), + dict(type="Resize", keep_ratio=True), dict( - type='RandomCenterCropPad', + type="RandomCenterCropPad", ratios=None, border=None, mean=[0, 0, 0], std=[1, 1, 1], to_rgb=True, test_mode=True, - test_pad_mode=['logical_or', 31], - test_pad_add_pix=1), - dict(type='RandomFlip'), - dict(type='Normalize', **img_norm_cfg), - dict(type='DefaultFormatBundle'), + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict(type="Normalize", **img_norm_cfg), + dict(type="DefaultFormatBundle"), dict( - type='Collect', - meta_keys=('filename', 'ori_filename', 'ori_shape', - 'img_shape', 'pad_shape', 'scale_factor', 'flip', - 'flip_direction', 'img_norm_cfg', 'border'), - keys=['img']) - ]) + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), ] -dataset_type = 'CocoDataset' -data_root = 'data/coco/' +dataset_type = "CocoDataset" +data_root = "data/coco/" # Use RepeatDataset to speed up training data = dict( @@ -94,31 +112,34 @@ workers_per_gpu=4, train=dict( _delete_=True, - type='RepeatDataset', + type="RepeatDataset", times=5, dataset=dict( type=dataset_type, - ann_file=data_root + 'annotations/instances_train2017.json', - img_prefix=data_root + 'train2017/', - pipeline=train_pipeline)), + ann_file=data_root + "annotations/instances_train2017.json", + img_prefix=data_root + "train2017/", + pipeline=train_pipeline, + ), + ), val=dict(pipeline=test_pipeline), - test=dict(pipeline=test_pipeline)) + test=dict(pipeline=test_pipeline), +) # optimizer # Based on the default settings of modern detectors, the SGD effect is better # than the Adam in the source code, so we use SGD default settings and # if you use adam+lr5e-4, the map is 29.1. -optimizer_config = dict( - _delete_=True, grad_clip=dict(max_norm=35, norm_type=2)) +optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=35, norm_type=2)) # learning policy # Based on the default settings of modern detectors, we added warmup settings. lr_config = dict( - policy='step', - warmup='linear', + policy="step", + warmup="linear", warmup_iters=1000, warmup_ratio=1.0 / 1000, - step=[18, 24]) # the real step is [18*5, 24*5] + step=[18, 24], +) # the real step is [18*5, 24*5] runner = dict(max_epochs=28) # the real epoch is 28*5=140 # NOTE: `auto_scale_lr` is for automatically scaling LR, diff --git a/configs/centernet/wtw_bbox_cycle_centernet_dla34_dcnv2_150e_coco.py b/configs/centernet/wtw_bbox_cycle_centernet_dla34_dcnv2_150e_coco.py new file mode 100644 index 00000000..5188607a --- /dev/null +++ b/configs/centernet/wtw_bbox_cycle_centernet_dla34_dcnv2_150e_coco.py @@ -0,0 +1,326 @@ +TEST_NAME = "31_long_check" +EVAL_LAG = 30 +CHECKPOINT_LAG = 5 +EPOCHS = 150 +LR = 0.00125 +BACKBONE = "DLANetMMDet3D" +BATCH = 8 +local_maximum_kernel = 1 +LOG_LAG = 100 +TAGS = [ + f"local_maximum_kernel={local_maximum_kernel}", + "1hm", + f"{EPOCHS}_epoch", + f"{LR}_lr", + f"{EVAL_LAG}_eval_lag", + f"{CHECKPOINT_LAG}_chkpt_lag", + f"{BACKBONE}_backbone", + f"{BATCH}_batch", +] +# DATA AND AUG +dataset_type = "CocoDataset" +data_root = "/home/aiarhipov/datasets/WTW-dataset/" + +img_norm_cfg = dict(mean=[103.53, 116.28, 123.675], std=[1.0, 1.0, 1.0], to_rgb=False) + +data = dict( + samples_per_gpu=BATCH, + workers_per_gpu=2, + train=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/train/train.json", + img_prefix="train/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val_loss=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + test=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), +) + + +# MODEL CycleCenterNet(dcnv2) DLANetMMDet3D, +load_from = None +resume_from = None + +model = dict( + type="CenterNet", + backbone=dict( + type=BACKBONE, + depth=34, + norm_cfg=dict(type="BN"), + init_cfg=dict( + type="Pretrained", + checkpoint="http://dl.yf.io/dla/models/imagenet/dla34%2Btricks-24a49e58.pth", + ), + ), + neck=dict( + type="CTResNetNeck", + in_channel=512, + num_deconv_filters=(256, 128, 64), + num_deconv_kernels=(4, 4, 4), + use_dcn=True, + ), + bbox_head=dict( + type="CycleCenterNetHead", + in_channel=64, + feat_channel=64, + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + loss_c2v=dict(type="L1Loss", loss_weight=1.0), + loss_v2c=dict(type="L1Loss", loss_weight=0.5), + ), + train_cfg=None, + test_cfg=dict( + topk=3000, + local_maximum_kernel=local_maximum_kernel, + max_per_img=3000, + # nms=dict(type="nms", iou_threshold=0.475, split_thr=2000), + ), +) + + +# GPU +gpu_ids = [5] +device = "cuda" + + +# OPTIMIZATION +optimizer = dict(type="SGD", lr=LR, momentum=0.9, weight_decay=0.0001) +# Based on the default settings of modern detectors, the SGD effect is better +# than the Adam in the source code, so we use SGD default settings and +# if you use adam+lr5e-4, the map is 29.1. +optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2)) + + +# LEARNING POLICY +runner = dict(type="EpochBasedRunner", max_epochs=EPOCHS) # the real epoch is 28*5=140 + +# Based on the default settings of modern detectors, we added warmup settings. +lr_config = dict( + policy="step", + warmup="linear", + warmup_iters=1000, + warmup_ratio=0.001, + step=[90, 120], # the real step is [18*5, 24*5] +) + +# NOTE: `auto_scale_lr` is for automatically scaling LR, +# USER SHOULD NOT CHANGE ITS VALUES. +# base_batch_size = (8 GPUs) x (16 samples per GPU) +auto_scale_lr = dict(enable=False, base_batch_size=16) + + +# LOGGING +work_dir = f"/home/aiarhipov/centernet/exps/{TEST_NAME}" +INTERVAL = int(10976 / (LOG_LAG * BATCH)) + (10976 % (LOG_LAG * BATCH) > 0) +log_config = dict( + interval=INTERVAL, + hooks=[ + dict(type="TextLoggerHook"), + dict(type="TensorboardLoggerHook"), + dict( + type="MMDetWandbHook", + init_kwargs={ + "project": "CenterNet", + "entity": "centernet", + "name": TEST_NAME, + "dir": "/home/aiarhipov/centernet/exps/wandb", + "tags": TAGS, + }, + interval=INTERVAL, + log_checkpoint=True, + log_checkpoint_metadata=False, + num_eval_images=5, + ), + ], +) +log_level = "INFO" + + +# EVALUATION +evaluation = dict(interval=EVAL_LAG, metric="bbox") +checkpoint_config = dict(interval=CHECKPOINT_LAG, max_keep_ckpts=5) + + +# RUNTIME +seed = 0 + +custom_hooks = [dict(type="NumClassCheckHook")] +dist_params = dict(backend="nccl") +workflow = [("train", 1), ("val", 1)] + +# disable opencv multithreading to avoid system being overloaded +opencv_num_threads = 0 +# set multi-process start method as `fork` to speed up the training +mp_start_method = "fork" diff --git a/configs/centernet/wtw_centernet_dla34_dcnv2_150e_coco.py b/configs/centernet/wtw_centernet_dla34_dcnv2_150e_coco.py new file mode 100644 index 00000000..4798d795 --- /dev/null +++ b/configs/centernet/wtw_centernet_dla34_dcnv2_150e_coco.py @@ -0,0 +1,302 @@ +TEST_NAME = "16_paper_params_dla34_batch8" + +# DATA AND AUG +dataset_type = "CocoDataset" +data_root = "/home/aiarhipov/datasets/WTW-dataset/" + +img_norm_cfg = dict(mean=[103.53, 116.28, 123.675], std=[1.0, 1.0, 1.0], to_rgb=False) + +data = dict( + samples_per_gpu=8, + workers_per_gpu=2, + train=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/train/train.json", + img_prefix="train/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val_loss=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + test=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), +) + + +# MODEL CenterNet(dcnv2) without changes +load_from = None +resume_from = None + +model = dict( + type="CenterNet", + backbone=dict( + type="DLANetMMDet3D", + depth=34, + norm_cfg=dict(type="BN"), + init_cfg=dict( + type="Pretrained", + checkpoint="http://dl.yf.io/dla/models/imagenet/dla34%2Btricks-24a49e58.pth", + ), + ), + neck=dict( + type="CTResNetNeck", + in_channel=512, + num_deconv_filters=(256, 128, 64), + num_deconv_kernels=(4, 4, 4), + use_dcn=True, + ), + bbox_head=dict( + type="CenterNetHead", + num_classes=1, + in_channel=64, + feat_channel=64, + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_wh=dict(type="L1Loss", loss_weight=0.1), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + ), + train_cfg=None, + test_cfg=dict(topk=3000, local_maximum_kernel=1, max_per_img=3000), +) + + +# GPU +gpu_ids = [6] +device = "cuda" + + +# OPTIMIZATION +optimizer = dict(type="SGD", lr=0.00125, momentum=0.9, weight_decay=0.0001) +# Based on the default settings of modern detectors, the SGD effect is better +# than the Adam in the source code, so we use SGD default settings and +# if you use adam+lr5e-4, the map is 29.1. +optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2)) + + +# LEARNING POLICY +runner = dict(type="EpochBasedRunner", max_epochs=150) # the real epoch is 28*5=140 + +# Based on the default settings of modern detectors, we added warmup settings. +lr_config = dict( + policy="step", + warmup="linear", + warmup_iters=1000, + warmup_ratio=0.001, + step=[90, 120], # the real step is [18*5, 24*5] +) + +# NOTE: `auto_scale_lr` is for automatically scaling LR, +# USER SHOULD NOT CHANGE ITS VALUES. +# base_batch_size = (8 GPUs) x (16 samples per GPU) +auto_scale_lr = dict(enable=False, base_batch_size=16) + + +# LOGGING +work_dir = f"/home/aiarhipov/centernet/exps/{TEST_NAME}" + +log_config = dict( + interval=1000, + hooks=[ + dict(type="TextLoggerHook"), + dict(type="TensorboardLoggerHook"), + dict( + type="MMDetWandbHook", + init_kwargs={ + "project": "CenterNet", + "entity": "centernet", + "name": TEST_NAME, + }, + interval=1000, + log_checkpoint=True, + log_checkpoint_metadata=True, + num_eval_images=15, + ), + ], +) +log_level = "INFO" + + +# EVALUATION +evaluation = dict(interval=30, metric="bbox") +checkpoint_config = dict(interval=30) + + +# RUNTIME +seed = 0 + +custom_hooks = [dict(type="NumClassCheckHook")] +dist_params = dict(backend="nccl") +workflow = [("train", 1), ("val", 1)] + +# disable opencv multithreading to avoid system being overloaded +opencv_num_threads = 0 +# set multi-process start method as `fork` to speed up the training +mp_start_method = "fork" diff --git a/configs/centernet/wtw_centernet_resnet34_dcnv2_150e_coco.py b/configs/centernet/wtw_centernet_resnet34_dcnv2_150e_coco.py new file mode 100644 index 00000000..e670ac06 --- /dev/null +++ b/configs/centernet/wtw_centernet_resnet34_dcnv2_150e_coco.py @@ -0,0 +1,300 @@ +TEST_NAME = "15_paper_params_resnet34_batch8" + +# DATA AND AUG +dataset_type = "CocoDataset" +data_root = "/home/aiarhipov/datasets/WTW-dataset/" + +img_norm_cfg = dict(mean=[103.53, 116.28, 123.675], std=[1.0, 1.0, 1.0], to_rgb=False) + +data = dict( + samples_per_gpu=8, + workers_per_gpu=2, + train=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/train/train.json", + img_prefix="train/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val_loss=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + test=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), +) + + +# MODEL CenterNet(dcnv2) without changes +load_from = None +resume_from = None + +model = dict( + type="CenterNet", + backbone=dict( + type="ResNet", + depth=34, + norm_eval=False, + norm_cfg=dict(type="BN"), + init_cfg=dict(type="Pretrained", checkpoint="torchvision://resnet34"), + ), + neck=dict( + type="CTResNetNeck", + in_channel=512, + num_deconv_filters=(256, 128, 64), + num_deconv_kernels=(4, 4, 4), + use_dcn=True, + ), + bbox_head=dict( + type="CenterNetHead", + num_classes=1, + in_channel=64, + feat_channel=64, + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_wh=dict(type="L1Loss", loss_weight=0.1), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + ), + train_cfg=None, + test_cfg=dict(topk=3000, local_maximum_kernel=1, max_per_img=3000), +) + + +# GPU +gpu_ids = [6] +device = "cuda" + + +# OPTIMIZATION +optimizer = dict(type="SGD", lr=0.00125, momentum=0.9, weight_decay=0.0001) +# Based on the default settings of modern detectors, the SGD effect is better +# than the Adam in the source code, so we use SGD default settings and +# if you use adam+lr5e-4, the map is 29.1. +optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2)) + + +# LEARNING POLICY +runner = dict(type="EpochBasedRunner", max_epochs=150) # the real epoch is 28*5=140 + +# Based on the default settings of modern detectors, we added warmup settings. +lr_config = dict( + policy="step", + warmup="linear", + warmup_iters=1000, + warmup_ratio=0.001, + step=[90, 120], # the real step is [18*5, 24*5] +) + +# NOTE: `auto_scale_lr` is for automatically scaling LR, +# USER SHOULD NOT CHANGE ITS VALUES. +# base_batch_size = (8 GPUs) x (16 samples per GPU) +auto_scale_lr = dict(enable=False, base_batch_size=16) + + +# LOGGING +work_dir = f"/home/aiarhipov/centernet/exps/{TEST_NAME}" + +log_config = dict( + interval=1000, + hooks=[ + dict(type="TextLoggerHook"), + dict(type="TensorboardLoggerHook"), + dict( + type="MMDetWandbHook", + init_kwargs={ + "project": "CenterNet", + "entity": "centernet", + "name": TEST_NAME, + }, + interval=1000, + log_checkpoint=True, + log_checkpoint_metadata=True, + num_eval_images=15, + ), + ], +) +log_level = "INFO" + + +# EVALUATION +evaluation = dict(interval=30, metric="bbox") +checkpoint_config = dict(interval=30) + + +# RUNTIME +seed = 0 + +custom_hooks = [dict(type="NumClassCheckHook")] +dist_params = dict(backend="nccl") +workflow = [("train", 1), ("val", 1)] + +# disable opencv multithreading to avoid system being overloaded +opencv_num_threads = 0 +# set multi-process start method as `fork` to speed up the training +mp_start_method = "fork" diff --git a/configs/centernet/wtw_quad_cycle_centernet_dla34_dcnv2_150e_coco.py b/configs/centernet/wtw_quad_cycle_centernet_dla34_dcnv2_150e_coco.py new file mode 100644 index 00000000..8401b89b --- /dev/null +++ b/configs/centernet/wtw_quad_cycle_centernet_dla34_dcnv2_150e_coco.py @@ -0,0 +1,327 @@ +TEST_NAME = "32_quad_long" +EVAL_LAG = 150 +CHECKPOINT_LAG = 1 +EPOCHS = 150 +LR = 0.00125 +BACKBONE = "DLANetMMDet3D" +BATCH = 8 +local_maximum_kernel = 1 +ITER_PERIOD = 2 +TAGS = [ + f"local_maximum_kernel={local_maximum_kernel}", + "1hm", + f"{EPOCHS}_epoch", + f"{LR}_lr", + f"{EVAL_LAG}_eval_lag", + f"{CHECKPOINT_LAG}_chkpt_lag", + f"{BACKBONE}_backbone", + f"{BATCH}_batch", +] +# DATA AND AUG +dataset_type = "CocoDataset" +data_root = "/home/aiarhipov/datasets/WTW-dataset/" + +img_norm_cfg = dict(mean=[103.53, 116.28, 123.675], std=[1.0, 1.0, 1.0], to_rgb=False) + +data = dict( + # train_dataloader=dict(shuffle=), + samples_per_gpu=BATCH, + workers_per_gpu=2, + train=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/train/train_4x10.json", + img_prefix="train/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val_loss=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test_4x10.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True, color_type="color"), + dict(type="LoadAnnotations", with_bbox=True), + dict( + type="PhotoMetricDistortion", + brightness_delta=32, + contrast_range=(0.5, 1.5), + saturation_range=(0.5, 1.5), + hue_delta=18, + ), + dict( + type="RandomCenterCropPad", + crop_size=(1024, 1024), + ratios=(0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3), + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_pad_mode=None, + ), + dict(type="Resize", img_scale=(1024, 1024), keep_ratio=True), + dict(type="RandomFlip", flip_ratio=0.5), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict(type="Collect", keys=["img", "gt_bboxes", "gt_labels"]), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + val=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test_4x10.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), + test=dict( + type=dataset_type, + ann_file="/home/aiarhipov/datasets/WTW-dataset/test/test_4x10.json", + img_prefix="test/images/", + pipeline=[ + dict(type="LoadImageFromFile", to_float32=True), + dict( + type="MultiScaleFlipAug", + scale_factor=1.0, + flip=False, + transforms=[ + dict(type="Resize", keep_ratio=True), + dict( + type="RandomCenterCropPad", + ratios=None, + border=None, + mean=[0, 0, 0], + std=[1, 1, 1], + to_rgb=True, + test_mode=True, + test_pad_mode=["logical_or", 31], + test_pad_add_pix=1, + ), + dict(type="RandomFlip"), + dict( + type="Normalize", + mean=[103.53, 116.28, 123.675], + std=[1.0, 1.0, 1.0], + to_rgb=False, + ), + dict(type="DefaultFormatBundle"), + dict( + type="Collect", + meta_keys=( + "filename", + "ori_filename", + "ori_shape", + "img_shape", + "pad_shape", + "scale_factor", + "flip", + "flip_direction", + "img_norm_cfg", + "border", + ), + keys=["img"], + ), + ], + ), + ], + data_root="/home/aiarhipov/datasets/WTW-dataset/", + classes=("box",), + ), +) + + +# MODEL CycleCenterNet(dcnv2) DLANetMMDet3D, +load_from = None +resume_from = "/home/aiarhipov/centernet/exps/32_quad_long/epoch_85.pth" + +model = dict( + type="CenterNet", + backbone=dict( + type=BACKBONE, + depth=34, + norm_cfg=dict(type="BN"), + init_cfg=dict( + type="Pretrained", + checkpoint="http://dl.yf.io/dla/models/imagenet/dla34%2Btricks-24a49e58.pth", + ), + ), + neck=dict( + type="CTResNetNeck", + in_channel=512, + num_deconv_filters=(256, 128, 64), + num_deconv_kernels=(4, 4, 4), + use_dcn=True, + ), + bbox_head=dict( + type="CycleCenterNetHead", + in_channel=64, + feat_channel=64, + num_classes=1, + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + loss_c2v=dict(type="L1Loss", loss_weight=1.0), + loss_v2c=dict(type="L1Loss", loss_weight=0.5), + ), + train_cfg=None, + test_cfg=dict( + topk=3000, + local_maximum_kernel=local_maximum_kernel, + max_per_img=3000, + ), +) + + +# GPU +gpu_ids = [5] +device = "cuda" + + +# OPTIMIZATION +optimizer = dict(type="SGD", lr=LR, momentum=0.9, weight_decay=0.0001) +# Based on the default settings of modern detectors, the SGD effect is better +# than the Adam in the source code, so we use SGD default settings and +# if you use adam+lr5e-4, the map is 29.1. +optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2)) + + +# LEARNING POLICY +runner = dict(type="EpochBasedRunner", max_epochs=EPOCHS) # the real epoch is 28*5=140 + +# Based on the default settings of modern detectors, we added warmup settings. +lr_config = dict( + policy="step", + warmup="linear", + warmup_iters=1000, + warmup_ratio=0.001, + step=[90, 120], # the real step is [18*5, 24*5] +) + +# NOTE: `auto_scale_lr` is for automatically scaling LR, +# USER SHOULD NOT CHANGE ITS VALUES. +# base_batch_size = (8 GPUs) x (16 samples per GPU) +auto_scale_lr = dict(enable=False, base_batch_size=16) + + +# LOGGING +work_dir = f"/home/aiarhipov/centernet/exps/{TEST_NAME}" +INTERVAL = int(10976 / (ITER_PERIOD * BATCH)) + (10976 % (ITER_PERIOD * BATCH) > 0) +log_config = dict( + interval=INTERVAL, + hooks=[ + dict(type="TextLoggerHook"), + dict(type="TensorboardLoggerHook"), + dict( + type="MMDetWandbHook", + init_kwargs={ + "project": "CenterNet", + "entity": "centernet", + "name": TEST_NAME, + "dir": "/home/aiarhipov/centernet/exps/wandb", + "tags": TAGS, + }, + interval=INTERVAL, + log_checkpoint=True, + log_checkpoint_metadata=False, + num_eval_images=5, + ), + ], +) +log_level = "INFO" + + +# EVALUATION +evaluation = dict(interval=EVAL_LAG, metric="bbox") +checkpoint_config = dict(interval=CHECKPOINT_LAG, max_keep_ckpts=1) + + +# RUNTIME +seed = 0 + +custom_hooks = [dict(type="NumClassCheckHook")] +dist_params = dict(backend="nccl") +workflow = [("train", 1), ("val", 1)] + +# disable opencv multithreading to avoid system being overloaded +opencv_num_threads = 0 +# set multi-process start method as `fork` to speed up the training +mp_start_method = "fork" diff --git a/demo/schema.jpg b/demo/schema.jpg new file mode 100644 index 00000000..833313a0 Binary files /dev/null and b/demo/schema.jpg differ diff --git a/mmdet/models/backbones/__init__.py b/mmdet/models/backbones/__init__.py index 91b50d25..a66650a3 100644 --- a/mmdet/models/backbones/__init__.py +++ b/mmdet/models/backbones/__init__.py @@ -3,6 +3,7 @@ from .darknet import Darknet from .detectors_resnet import DetectoRS_ResNet from .detectors_resnext import DetectoRS_ResNeXt +from .dla_mmdet3d import DLANetMMDet3D from .efficientnet import EfficientNet from .hourglass import HourglassNet from .hrnet import HRNet @@ -18,9 +19,24 @@ from .trident_resnet import TridentResNet __all__ = [ - 'RegNet', 'ResNet', 'ResNetV1d', 'ResNeXt', 'SSDVGG', 'HRNet', - 'MobileNetV2', 'Res2Net', 'HourglassNet', 'DetectoRS_ResNet', - 'DetectoRS_ResNeXt', 'Darknet', 'ResNeSt', 'TridentResNet', 'CSPDarknet', - 'SwinTransformer', 'PyramidVisionTransformer', - 'PyramidVisionTransformerV2', 'EfficientNet' + "RegNet", + "ResNet", + "ResNetV1d", + "ResNeXt", + "SSDVGG", + "HRNet", + "MobileNetV2", + "Res2Net", + "HourglassNet", + "DetectoRS_ResNet", + "DetectoRS_ResNeXt", + "Darknet", + "ResNeSt", + "TridentResNet", + "CSPDarknet", + "SwinTransformer", + "PyramidVisionTransformer", + "PyramidVisionTransformerV2", + "EfficientNet", + "DLANetMMDet3D", ] diff --git a/mmdet/models/backbones/dla_mmdet3d.py b/mmdet/models/backbones/dla_mmdet3d.py new file mode 100644 index 00000000..2f0b0965 --- /dev/null +++ b/mmdet/models/backbones/dla_mmdet3d.py @@ -0,0 +1,477 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import warnings + +import torch +from mmcv.cnn import build_conv_layer, build_norm_layer +from mmcv.runner import BaseModule +from torch import nn + +from ..builder import BACKBONES + + +def dla_build_norm_layer(cfg, num_features): + """Build normalization layer specially designed for DLANet. + + Args: + cfg (dict): The norm layer config, which should contain: + + - type (str): Layer type. + - layer args: Args needed to instantiate a norm layer. + - requires_grad (bool, optional): Whether stop gradient updates. + num_features (int): Number of input channels. + + + Returns: + Function: Build normalization layer in mmcv. + """ + cfg_ = cfg.copy() + if cfg_["type"] == "GN": + if num_features % 32 == 0: + return build_norm_layer(cfg_, num_features) + else: + assert "num_groups" in cfg_ + cfg_["num_groups"] = cfg_["num_groups"] // 2 + return build_norm_layer(cfg_, num_features) + else: + return build_norm_layer(cfg_, num_features) + + +class BasicBlock(BaseModule): + """BasicBlock in DLANet. + + Args: + in_channels (int): Input feature channel. + out_channels (int): Output feature channel. + norm_cfg (dict): Dictionary to construct and config + norm layer. + conv_cfg (dict): Dictionary to construct and config + conv layer. + stride (int, optional): Conv stride. Default: 1. + dilation (int, optional): Conv dilation. Default: 1. + init_cfg (dict, optional): Initialization config. + Default: None. + """ + + def __init__( + self, + in_channels, + out_channels, + norm_cfg, + conv_cfg, + stride=1, + dilation=1, + init_cfg=None, + ): + super(BasicBlock, self).__init__(init_cfg) + self.conv1 = build_conv_layer( + conv_cfg, + in_channels, + out_channels, + 3, + stride=stride, + padding=dilation, + dilation=dilation, + bias=False, + ) + self.norm1 = dla_build_norm_layer(norm_cfg, out_channels)[1] + self.relu = nn.ReLU(inplace=True) + self.conv2 = build_conv_layer( + conv_cfg, + out_channels, + out_channels, + 3, + stride=1, + padding=dilation, + dilation=dilation, + bias=False, + ) + self.norm2 = dla_build_norm_layer(norm_cfg, out_channels)[1] + self.stride = stride + + def forward(self, x, identity=None): + """Forward function.""" + + if identity is None: + identity = x + out = self.conv1(x) + out = self.norm1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.norm2(out) + out += identity + out = self.relu(out) + + return out + + +class Root(BaseModule): + """Root in DLANet. + + Args: + in_channels (int): Input feature channel. + out_channels (int): Output feature channel. + norm_cfg (dict): Dictionary to construct and config + norm layer. + conv_cfg (dict): Dictionary to construct and config + conv layer. + kernel_size (int): Size of convolution kernel. + add_identity (bool): Whether to add identity in root. + init_cfg (dict, optional): Initialization config. + Default: None. + """ + + def __init__( + self, + in_channels, + out_channels, + norm_cfg, + conv_cfg, + kernel_size, + add_identity, + init_cfg=None, + ): + super(Root, self).__init__(init_cfg) + self.conv = build_conv_layer( + conv_cfg, + in_channels, + out_channels, + 1, + stride=1, + padding=(kernel_size - 1) // 2, + bias=False, + ) + self.norm = dla_build_norm_layer(norm_cfg, out_channels)[1] + self.relu = nn.ReLU(inplace=True) + self.add_identity = add_identity + + def forward(self, feat_list): + """Forward function. + + Args: + feat_list (list[torch.Tensor]): Output features from + multiple layers. + """ + children = feat_list + x = self.conv(torch.cat(feat_list, 1)) + x = self.norm(x) + if self.add_identity: + x += children[0] + x = self.relu(x) + + return x + + +class Tree(BaseModule): + """Tree in DLANet. + + Args: + levels (int): The level of the tree. + block (nn.Module): The block module in tree. + in_channels: Input feature channel. + out_channels: Output feature channel. + norm_cfg (dict): Dictionary to construct and config + norm layer. + conv_cfg (dict): Dictionary to construct and config + conv layer. + stride (int, optional): Convolution stride. + Default: 1. + level_root (bool, optional): whether belongs to the + root layer. + root_dim (int, optional): Root input feature channel. + root_kernel_size (int, optional): Size of root + convolution kernel. Default: 1. + dilation (int, optional): Conv dilation. Default: 1. + add_identity (bool, optional): Whether to add + identity in root. Default: False. + init_cfg (dict, optional): Initialization config. + Default: None. + """ + + def __init__( + self, + levels, + block, + in_channels, + out_channels, + norm_cfg, + conv_cfg, + stride=1, + level_root=False, + root_dim=None, + root_kernel_size=1, + dilation=1, + add_identity=False, + init_cfg=None, + ): + super(Tree, self).__init__(init_cfg) + if root_dim is None: + root_dim = 2 * out_channels + if level_root: + root_dim += in_channels + if levels == 1: + self.root = Root( + root_dim, + out_channels, + norm_cfg, + conv_cfg, + root_kernel_size, + add_identity, + ) + self.tree1 = block( + in_channels, + out_channels, + norm_cfg, + conv_cfg, + stride, + dilation=dilation, + ) + self.tree2 = block( + out_channels, + out_channels, + norm_cfg, + conv_cfg, + 1, + dilation=dilation, + ) + else: + self.tree1 = Tree( + levels - 1, + block, + in_channels, + out_channels, + norm_cfg, + conv_cfg, + stride, + root_dim=None, + root_kernel_size=root_kernel_size, + dilation=dilation, + add_identity=add_identity, + ) + self.tree2 = Tree( + levels - 1, + block, + out_channels, + out_channels, + norm_cfg, + conv_cfg, + root_dim=root_dim + out_channels, + root_kernel_size=root_kernel_size, + dilation=dilation, + add_identity=add_identity, + ) + self.level_root = level_root + self.root_dim = root_dim + self.downsample = None + self.project = None + self.levels = levels + if stride > 1: + self.downsample = nn.MaxPool2d(stride, stride=stride) + if in_channels != out_channels: + self.project = nn.Sequential( + build_conv_layer( + conv_cfg, in_channels, out_channels, 1, stride=1, bias=False + ), + dla_build_norm_layer(norm_cfg, out_channels)[1], + ) + + def forward(self, x, identity=None, children=None): + children = [] if children is None else children + bottom = self.downsample(x) if self.downsample else x + identity = self.project(bottom) if self.project else bottom + if self.level_root: + children.append(bottom) + x1 = self.tree1(x, identity) + if self.levels == 1: + x2 = self.tree2(x1) + feat_list = [x2, x1] + children + x = self.root(feat_list) + else: + children.append(x1) + x = self.tree2(x1, children=children) + return x + + +@BACKBONES.register_module() +class DLANetMMDet3D(BaseModule): + r"""`DLA backbone `_. + + Args: + depth (int): Depth of DLA. Default: 34. + in_channels (int, optional): Number of input image channels. + Default: 3. + norm_cfg (dict, optional): Dictionary to construct and config + norm layer. Default: None. + conv_cfg (dict, optional): Dictionary to construct and config + conv layer. Default: None. + layer_with_level_root (list[bool], optional): Whether to apply + level_root in each DLA layer, this is only used for + tree levels. Default: (False, True, True, True). + with_identity_root (bool, optional): Whether to add identity + in root layer. Default: False. + pretrained (str, optional): model pretrained path. + Default: None. + init_cfg (dict or list[dict], optional): Initialization + config dict. Default: None + """ + arch_settings = { + 34: (BasicBlock, (1, 1, 1, 2, 2, 1), (16, 32, 64, 128, 256, 512)), + } + + def __init__( + self, + depth, + in_channels=3, + out_indices=(0, 1, 2, 3, 4, 5), + frozen_stages=-1, + norm_cfg=None, + conv_cfg=None, + layer_with_level_root=(False, True, True, True), + with_identity_root=False, + pretrained=None, + init_cfg=None, + ): + super(DLANetMMDet3D, self).__init__(init_cfg) + if depth not in self.arch_settings: + raise KeyError(f"invalida depth {depth} for DLA") + + assert not ( + init_cfg and pretrained + ), "init_cfg and pretrained cannot be setting at the same time" + if isinstance(pretrained, str): + warnings.warn( + "DeprecationWarning: pretrained is a deprecated, " + 'please use "init_cfg" instead' + ) + self.init_cfg = dict(type="Pretrained", checkpoint=pretrained) + elif pretrained is None: + if init_cfg is None: + self.init_cfg = [ + dict(type="Kaiming", layer="Conv2d"), + dict( + type="Constant", + val=1, + layer=["_BatchNorm", "GroupNorm"], + ), + ] + + block, levels, channels = self.arch_settings[depth] + self.channels = channels + self.num_levels = len(levels) + self.frozen_stages = frozen_stages + self.out_indices = out_indices + assert max(out_indices) < self.num_levels + self.base_layer = nn.Sequential( + build_conv_layer( + conv_cfg, + in_channels, + channels[0], + 7, + stride=1, + padding=3, + bias=False, + ), + dla_build_norm_layer(norm_cfg, channels[0])[1], + nn.ReLU(inplace=True), + ) + + # DLANet first uses two conv layers then uses several + # Tree layers + for i in range(2): + level_layer = self._make_conv_level( + channels[0], + channels[i], + levels[i], + norm_cfg, + conv_cfg, + stride=i + 1, + ) + layer_name = f"level{i}" + self.add_module(layer_name, level_layer) + + for i in range(2, self.num_levels): + dla_layer = Tree( + levels[i], + block, + channels[i - 1], + channels[i], + norm_cfg, + conv_cfg, + 2, + level_root=layer_with_level_root[i - 2], + add_identity=with_identity_root, + ) + layer_name = f"level{i}" + self.add_module(layer_name, dla_layer) + + self._freeze_stages() + + def _make_conv_level( + self, + in_channels, + out_channels, + num_convs, + norm_cfg, + conv_cfg, + stride=1, + dilation=1, + ): + """Conv modules. + + Args: + in_channels (int): Input feature channel. + out_channels (int): Output feature channel. + num_convs (int): Number of Conv module. + norm_cfg (dict): Dictionary to construct and config + norm layer. + conv_cfg (dict): Dictionary to construct and config + conv layer. + stride (int, optional): Conv stride. Default: 1. + dilation (int, optional): Conv dilation. Default: 1. + """ + modules = [] + for i in range(num_convs): + modules.extend( + [ + build_conv_layer( + conv_cfg, + in_channels, + out_channels, + 3, + stride=stride if i == 0 else 1, + padding=dilation, + bias=False, + dilation=dilation, + ), + dla_build_norm_layer(norm_cfg, out_channels)[1], + nn.ReLU(inplace=True), + ] + ) + in_channels = out_channels + return nn.Sequential(*modules) + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.base_layer.eval() + for param in self.base_layer.parameters(): + param.requires_grad = False + + for i in range(2): + m = getattr(self, f"level{i}") + m.eval() + for param in m.parameters(): + param.requires_grad = False + + for i in range(1, self.frozen_stages + 1): + m = getattr(self, f"level{i+1}") + m.eval() + for param in m.parameters(): + param.requires_grad = False + + def forward(self, x): + outs = [] + x = self.base_layer(x) + for i in range(self.num_levels): + x = getattr(self, "level{}".format(i))(x) + if i in self.out_indices: + outs.append(x) + return tuple(outs) diff --git a/mmdet/models/dense_heads/__init__.py b/mmdet/models/dense_heads/__init__.py index 1c228699..409b11a5 100644 --- a/mmdet/models/dense_heads/__init__.py +++ b/mmdet/models/dense_heads/__init__.py @@ -41,18 +41,56 @@ from .yolo_head import YOLOV3Head from .yolof_head import YOLOFHead from .yolox_head import YOLOXHead +from .cycle_centernet_head import CycleCenterNetHead __all__ = [ - 'AnchorFreeHead', 'AnchorHead', 'GuidedAnchorHead', 'FeatureAdaption', - 'RPNHead', 'GARPNHead', 'RetinaHead', 'RetinaSepBNHead', 'GARetinaHead', - 'SSDHead', 'FCOSHead', 'RepPointsHead', 'FoveaHead', - 'FreeAnchorRetinaHead', 'ATSSHead', 'FSAFHead', 'NASFCOSHead', - 'PISARetinaHead', 'PISASSDHead', 'GFLHead', 'CornerHead', 'YOLACTHead', - 'YOLACTSegmHead', 'YOLACTProtonet', 'YOLOV3Head', 'PAAHead', - 'SABLRetinaHead', 'CentripetalHead', 'VFNetHead', 'StageCascadeRPNHead', - 'CascadeRPNHead', 'EmbeddingRPNHead', 'LDHead', 'AutoAssignHead', - 'DETRHead', 'YOLOFHead', 'DeformableDETRHead', 'SOLOHead', - 'DecoupledSOLOHead', 'CenterNetHead', 'YOLOXHead', - 'DecoupledSOLOLightHead', 'LADHead', 'TOODHead', 'MaskFormerHead', - 'Mask2FormerHead', 'SOLOV2Head', 'DDODHead' + "AnchorFreeHead", + "AnchorHead", + "GuidedAnchorHead", + "FeatureAdaption", + "RPNHead", + "GARPNHead", + "RetinaHead", + "RetinaSepBNHead", + "GARetinaHead", + "SSDHead", + "FCOSHead", + "RepPointsHead", + "FoveaHead", + "FreeAnchorRetinaHead", + "ATSSHead", + "FSAFHead", + "NASFCOSHead", + "PISARetinaHead", + "PISASSDHead", + "GFLHead", + "CornerHead", + "YOLACTHead", + "YOLACTSegmHead", + "YOLACTProtonet", + "YOLOV3Head", + "PAAHead", + "SABLRetinaHead", + "CentripetalHead", + "VFNetHead", + "StageCascadeRPNHead", + "CascadeRPNHead", + "EmbeddingRPNHead", + "LDHead", + "AutoAssignHead", + "DETRHead", + "YOLOFHead", + "DeformableDETRHead", + "SOLOHead", + "DecoupledSOLOHead", + "CenterNetHead", + "YOLOXHead", + "DecoupledSOLOLightHead", + "LADHead", + "TOODHead", + "MaskFormerHead", + "Mask2FormerHead", + "SOLOV2Head", + "DDODHead", + "CycleCenterNetHead", ] diff --git a/mmdet/models/dense_heads/base_dense_head.py b/mmdet/models/dense_heads/base_dense_head.py index 0c7abb7b..31726870 100644 --- a/mmdet/models/dense_heads/base_dense_head.py +++ b/mmdet/models/dense_heads/base_dense_head.py @@ -20,7 +20,7 @@ def init_weights(self): # avoid init_cfg overwrite the initialization of `conv_offset` for m in self.modules(): # DeformConv2dPack, ModulatedDeformConv2dPack - if hasattr(m, 'conv_offset'): + if hasattr(m, "conv_offset"): constant_init(m.conv_offset, 0) @abstractmethod @@ -28,16 +28,18 @@ def loss(self, **kwargs): """Compute losses of the head.""" pass - @force_fp32(apply_to=('cls_scores', 'bbox_preds')) - def get_bboxes(self, - cls_scores, - bbox_preds, - score_factors=None, - img_metas=None, - cfg=None, - rescale=False, - with_nms=True, - **kwargs): + @force_fp32(apply_to=("cls_scores", "bbox_preds")) + def get_bboxes( + self, + cls_scores, + bbox_preds, + score_factors=None, + img_metas=None, + cfg=None, + rescale=False, + with_nms=True, + **kwargs + ): """Transform network outputs of a batch into bbox results. Note: When score_factors is not None, the cls_scores are @@ -84,9 +86,8 @@ def get_bboxes(self, featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)] mlvl_priors = self.prior_generator.grid_priors( - featmap_sizes, - dtype=cls_scores[0].dtype, - device=cls_scores[0].device) + featmap_sizes, dtype=cls_scores[0].dtype, device=cls_scores[0].device + ) result_list = [] @@ -99,23 +100,32 @@ def get_bboxes(self, else: score_factor_list = [None for _ in range(num_levels)] - results = self._get_bboxes_single(cls_score_list, bbox_pred_list, - score_factor_list, mlvl_priors, - img_meta, cfg, rescale, with_nms, - **kwargs) + results = self._get_bboxes_single( + cls_score_list, + bbox_pred_list, + score_factor_list, + mlvl_priors, + img_meta, + cfg, + rescale, + with_nms, + **kwargs + ) result_list.append(results) return result_list - def _get_bboxes_single(self, - cls_score_list, - bbox_pred_list, - score_factor_list, - mlvl_priors, - img_meta, - cfg, - rescale=False, - with_nms=True, - **kwargs): + def _get_bboxes_single( + self, + cls_score_list, + bbox_pred_list, + score_factor_list, + mlvl_priors, + img_meta, + cfg, + rescale=False, + with_nms=True, + **kwargs + ): """Transform outputs of a single image into bbox predictions. Args: @@ -164,8 +174,8 @@ def _get_bboxes_single(self, with_score_factors = True cfg = self.test_cfg if cfg is None else cfg - img_shape = img_meta['img_shape'] - nms_pre = cfg.get('nms_pre', -1) + img_shape = img_meta["img_shape"] + nms_pre = cfg.get("nms_pre", -1) mlvl_bboxes = [] mlvl_scores = [] @@ -174,18 +184,15 @@ def _get_bboxes_single(self, mlvl_score_factors = [] else: mlvl_score_factors = None - for level_idx, (cls_score, bbox_pred, score_factor, priors) in \ - enumerate(zip(cls_score_list, bbox_pred_list, - score_factor_list, mlvl_priors)): - + for level_idx, (cls_score, bbox_pred, score_factor, priors) in enumerate( + zip(cls_score_list, bbox_pred_list, score_factor_list, mlvl_priors) + ): assert cls_score.size()[-2:] == bbox_pred.size()[-2:] bbox_pred = bbox_pred.permute(1, 2, 0).reshape(-1, 4) if with_score_factors: - score_factor = score_factor.permute(1, 2, - 0).reshape(-1).sigmoid() - cls_score = cls_score.permute(1, 2, - 0).reshape(-1, self.cls_out_channels) + score_factor = score_factor.permute(1, 2, 0).reshape(-1).sigmoid() + cls_score = cls_score.permute(1, 2, 0).reshape(-1, self.cls_out_channels) if self.use_sigmoid_cls: scores = cls_score.sigmoid() else: @@ -200,18 +207,17 @@ def _get_bboxes_single(self, # find a slight drop in performance, you can set a larger # `nms_pre` than before. results = filter_scores_and_topk( - scores, cfg.score_thr, nms_pre, - dict(bbox_pred=bbox_pred, priors=priors)) + scores, cfg.score_thr, nms_pre, dict(bbox_pred=bbox_pred, priors=priors) + ) scores, labels, keep_idxs, filtered_results = results - bbox_pred = filtered_results['bbox_pred'] - priors = filtered_results['priors'] + bbox_pred = filtered_results["bbox_pred"] + priors = filtered_results["priors"] if with_score_factors: score_factor = score_factor[keep_idxs] - bboxes = self.bbox_coder.decode( - priors, bbox_pred, max_shape=img_shape) + bboxes = self.bbox_coder.decode(priors, bbox_pred, max_shape=img_shape) mlvl_bboxes.append(bboxes) mlvl_scores.append(scores) @@ -219,20 +225,30 @@ def _get_bboxes_single(self, if with_score_factors: mlvl_score_factors.append(score_factor) - return self._bbox_post_process(mlvl_scores, mlvl_labels, mlvl_bboxes, - img_meta['scale_factor'], cfg, rescale, - with_nms, mlvl_score_factors, **kwargs) - - def _bbox_post_process(self, - mlvl_scores, - mlvl_labels, - mlvl_bboxes, - scale_factor, - cfg, - rescale=False, - with_nms=True, - mlvl_score_factors=None, - **kwargs): + return self._bbox_post_process( + mlvl_scores, + mlvl_labels, + mlvl_bboxes, + img_meta["scale_factor"], + cfg, + rescale, + with_nms, + mlvl_score_factors, + **kwargs + ) + + def _bbox_post_process( + self, + mlvl_scores, + mlvl_labels, + mlvl_bboxes, + scale_factor, + cfg, + rescale=False, + with_nms=True, + mlvl_score_factors=None, + **kwargs + ): """bbox post-processing method. The boxes would be rescaled to the original image scale and do @@ -292,22 +308,24 @@ def _bbox_post_process(self, det_bboxes = torch.cat([mlvl_bboxes, mlvl_scores[:, None]], -1) return det_bboxes, mlvl_labels - det_bboxes, keep_idxs = batched_nms(mlvl_bboxes, mlvl_scores, - mlvl_labels, cfg.nms) - det_bboxes = det_bboxes[:cfg.max_per_img] - det_labels = mlvl_labels[keep_idxs][:cfg.max_per_img] + det_bboxes, keep_idxs = batched_nms(mlvl_bboxes, mlvl_scores, mlvl_labels, cfg.nms) + det_bboxes = det_bboxes[: cfg.max_per_img] + det_labels = mlvl_labels[keep_idxs][: cfg.max_per_img] return det_bboxes, det_labels else: return mlvl_bboxes, mlvl_scores, mlvl_labels - def forward_train(self, - x, - img_metas, - gt_bboxes, - gt_labels=None, - gt_bboxes_ignore=None, - proposal_cfg=None, - **kwargs): + def forward_train( + self, + x, + img_metas, + gt_bboxes, + gt_labels=None, + gt_masks=None, + gt_bboxes_ignore=None, + proposal_cfg=None, + **kwargs + ): """ Args: x (list[Tensor]): Features from FPN. @@ -332,12 +350,11 @@ def forward_train(self, loss_inputs = outs + (gt_bboxes, img_metas) else: loss_inputs = outs + (gt_bboxes, gt_labels, img_metas) - losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore) + losses = self.loss(*loss_inputs, gt_masks=gt_masks, gt_bboxes_ignore=gt_bboxes_ignore) if proposal_cfg is None: return losses else: - proposal_list = self.get_bboxes( - *outs, img_metas=img_metas, cfg=proposal_cfg) + proposal_list = self.get_bboxes(*outs, img_metas=img_metas, cfg=proposal_cfg) return losses, proposal_list def simple_test(self, feats, img_metas, rescale=False): @@ -359,13 +376,8 @@ def simple_test(self, feats, img_metas, rescale=False): """ return self.simple_test_bboxes(feats, img_metas, rescale=rescale) - @force_fp32(apply_to=('cls_scores', 'bbox_preds')) - def onnx_export(self, - cls_scores, - bbox_preds, - score_factors=None, - img_metas=None, - with_nms=True): + @force_fp32(apply_to=("cls_scores", "bbox_preds")) + def onnx_export(self, cls_scores, bbox_preds, score_factors=None, img_metas=None, with_nms=True): """Transform network output for a batch into bbox predictions. Args: @@ -395,25 +407,21 @@ def onnx_export(self, featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores] mlvl_priors = self.prior_generator.grid_priors( - featmap_sizes, - dtype=bbox_preds[0].dtype, - device=bbox_preds[0].device) + featmap_sizes, dtype=bbox_preds[0].dtype, device=bbox_preds[0].device + ) mlvl_cls_scores = [cls_scores[i].detach() for i in range(num_levels)] mlvl_bbox_preds = [bbox_preds[i].detach() for i in range(num_levels)] - assert len( - img_metas - ) == 1, 'Only support one input image while in exporting to ONNX' - img_shape = img_metas[0]['img_shape_for_onnx'] + assert len(img_metas) == 1, "Only support one input image while in exporting to ONNX" + img_shape = img_metas[0]["img_shape_for_onnx"] cfg = self.test_cfg assert len(cls_scores) == len(bbox_preds) == len(mlvl_priors) device = cls_scores[0].device batch_size = cls_scores[0].shape[0] # convert to tensor to keep tracing - nms_pre_tensor = torch.tensor( - cfg.get('nms_pre', -1), device=device, dtype=torch.long) + nms_pre_tensor = torch.tensor(cfg.get("nms_pre", -1), device=device, dtype=torch.long) # e.g. Retina, FreeAnchor, etc. if score_factors is None: @@ -422,22 +430,18 @@ def onnx_export(self, else: # e.g. FCOS, PAA, ATSS, etc. with_score_factors = True - mlvl_score_factor = [ - score_factors[i].detach() for i in range(num_levels) - ] + mlvl_score_factor = [score_factors[i].detach() for i in range(num_levels)] mlvl_score_factors = [] mlvl_batch_bboxes = [] mlvl_scores = [] for cls_score, bbox_pred, score_factors, priors in zip( - mlvl_cls_scores, mlvl_bbox_preds, mlvl_score_factor, - mlvl_priors): + mlvl_cls_scores, mlvl_bbox_preds, mlvl_score_factor, mlvl_priors + ): assert cls_score.size()[-2:] == bbox_pred.size()[-2:] - scores = cls_score.permute(0, 2, 3, - 1).reshape(batch_size, -1, - self.cls_out_channels) + scores = cls_score.permute(0, 2, 3, 1).reshape(batch_size, -1, self.cls_out_channels) if self.use_sigmoid_cls: scores = scores.sigmoid() nms_pre_score = scores @@ -446,18 +450,16 @@ def onnx_export(self, nms_pre_score = scores if with_score_factors: - score_factors = score_factors.permute(0, 2, 3, 1).reshape( - batch_size, -1).sigmoid() - bbox_pred = bbox_pred.permute(0, 2, 3, - 1).reshape(batch_size, -1, 4) + score_factors = score_factors.permute(0, 2, 3, 1).reshape(batch_size, -1).sigmoid() + bbox_pred = bbox_pred.permute(0, 2, 3, 1).reshape(batch_size, -1, 4) priors = priors.expand(batch_size, -1, priors.size(-1)) # Get top-k predictions from mmdet.core.export import get_k_for_topk + nms_pre = get_k_for_topk(nms_pre_tensor, bbox_pred.shape[1]) if nms_pre > 0: - if with_score_factors: - nms_pre_score = (nms_pre_score * score_factors[..., None]) + nms_pre_score = nms_pre_score * score_factors[..., None] else: nms_pre_score = nms_pre_score @@ -471,26 +473,22 @@ def onnx_export(self, max_scores, _ = nms_pre_score[..., :-1].max(-1) _, topk_inds = max_scores.topk(nms_pre) - batch_inds = torch.arange( - batch_size, device=bbox_pred.device).view( - -1, 1).expand_as(topk_inds).long() + batch_inds = ( + torch.arange(batch_size, device=bbox_pred.device).view(-1, 1).expand_as(topk_inds).long() + ) # Avoid onnx2tensorrt issue in https://github.com/NVIDIA/TensorRT/issues/1134 # noqa: E501 transformed_inds = bbox_pred.shape[1] * batch_inds + topk_inds - priors = priors.reshape( - -1, priors.size(-1))[transformed_inds, :].reshape( - batch_size, -1, priors.size(-1)) - bbox_pred = bbox_pred.reshape(-1, - 4)[transformed_inds, :].reshape( - batch_size, -1, 4) - scores = scores.reshape( - -1, self.cls_out_channels)[transformed_inds, :].reshape( - batch_size, -1, self.cls_out_channels) + priors = priors.reshape(-1, priors.size(-1))[transformed_inds, :].reshape( + batch_size, -1, priors.size(-1) + ) + bbox_pred = bbox_pred.reshape(-1, 4)[transformed_inds, :].reshape(batch_size, -1, 4) + scores = scores.reshape(-1, self.cls_out_channels)[transformed_inds, :].reshape( + batch_size, -1, self.cls_out_channels + ) if with_score_factors: - score_factors = score_factors.reshape( - -1, 1)[transformed_inds].reshape(batch_size, -1) + score_factors = score_factors.reshape(-1, 1)[transformed_inds].reshape(batch_size, -1) - bboxes = self.bbox_coder.decode( - priors, bbox_pred, max_shape=img_shape) + bboxes = self.bbox_coder.decode(priors, bbox_pred, max_shape=img_shape) mlvl_batch_bboxes.append(bboxes) mlvl_scores.append(scores) @@ -507,20 +505,24 @@ def onnx_export(self, from mmdet.core.export import add_dummy_nms_for_onnx if not self.use_sigmoid_cls: - batch_scores = batch_scores[..., :self.num_classes] + batch_scores = batch_scores[..., : self.num_classes] if with_score_factors: batch_scores = batch_scores * (batch_score_factors.unsqueeze(2)) if with_nms: - max_output_boxes_per_class = cfg.nms.get( - 'max_output_boxes_per_class', 200) - iou_threshold = cfg.nms.get('iou_threshold', 0.5) + max_output_boxes_per_class = cfg.nms.get("max_output_boxes_per_class", 200) + iou_threshold = cfg.nms.get("iou_threshold", 0.5) score_threshold = cfg.score_thr - nms_pre = cfg.get('deploy_nms_pre', -1) - return add_dummy_nms_for_onnx(batch_bboxes, batch_scores, - max_output_boxes_per_class, - iou_threshold, score_threshold, - nms_pre, cfg.max_per_img) + nms_pre = cfg.get("deploy_nms_pre", -1) + return add_dummy_nms_for_onnx( + batch_bboxes, + batch_scores, + max_output_boxes_per_class, + iou_threshold, + score_threshold, + nms_pre, + cfg.max_per_img, + ) else: return batch_bboxes, batch_scores diff --git a/mmdet/models/dense_heads/centernet_head.py b/mmdet/models/dense_heads/centernet_head.py index b9d5d2f0..26cde2e9 100644 --- a/mmdet/models/dense_heads/centernet_head.py +++ b/mmdet/models/dense_heads/centernet_head.py @@ -8,8 +8,7 @@ from mmdet.core import multi_apply from mmdet.models import HEADS, build_loss from mmdet.models.utils import gaussian_radius, gen_gaussian_target -from ..utils.gaussian_target import (get_local_maximum, get_topk_from_heatmap, - transpose_and_gather_feat) +from ..utils.gaussian_target import get_local_maximum, get_topk_from_heatmap, transpose_and_gather_feat from .base_dense_head import BaseDenseHead from .dense_test_mixins import BBoxTestMixin @@ -35,21 +34,21 @@ class CenterNetHead(BaseDenseHead, BBoxTestMixin): Default: None """ - def __init__(self, - in_channel, - feat_channel, - num_classes, - loss_center_heatmap=dict( - type='GaussianFocalLoss', loss_weight=1.0), - loss_wh=dict(type='L1Loss', loss_weight=0.1), - loss_offset=dict(type='L1Loss', loss_weight=1.0), - train_cfg=None, - test_cfg=None, - init_cfg=None): + def __init__( + self, + in_channel, + feat_channel, + num_classes, + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_wh=dict(type="L1Loss", loss_weight=0.1), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + train_cfg=None, + test_cfg=None, + init_cfg=None, + ): super(CenterNetHead, self).__init__(init_cfg) self.num_classes = num_classes - self.heatmap_head = self._build_head(in_channel, feat_channel, - num_classes) + self.heatmap_head = self._build_head(in_channel, feat_channel, num_classes) self.wh_head = self._build_head(in_channel, feat_channel, 2) self.offset_head = self._build_head(in_channel, feat_channel, 2) @@ -66,7 +65,8 @@ def _build_head(self, in_channel, feat_channel, out_channel): layer = nn.Sequential( nn.Conv2d(in_channel, feat_channel, kernel_size=3, padding=1), nn.ReLU(inplace=True), - nn.Conv2d(feat_channel, out_channel, kernel_size=1)) + nn.Conv2d(feat_channel, out_channel, kernel_size=1), + ) return layer def init_weights(self): @@ -112,15 +112,10 @@ def forward_single(self, feat): offset_pred = self.offset_head(feat) return center_heatmap_pred, wh_pred, offset_pred - @force_fp32(apply_to=('center_heatmap_preds', 'wh_preds', 'offset_preds')) - def loss(self, - center_heatmap_preds, - wh_preds, - offset_preds, - gt_bboxes, - gt_labels, - img_metas, - gt_bboxes_ignore=None): + @force_fp32(apply_to=("center_heatmap_preds", "wh_preds", "offset_preds")) + def loss( + self, center_heatmap_preds, wh_preds, offset_preds, gt_bboxes, gt_labels, img_metas, gt_bboxes_ignore=None + ): """Compute losses of the head. Args: @@ -144,39 +139,30 @@ def loss(self, - loss_wh (Tensor): loss of hw heatmap - loss_offset (Tensor): loss of offset heatmap. """ - assert len(center_heatmap_preds) == len(wh_preds) == len( - offset_preds) == 1 + assert len(center_heatmap_preds) == len(wh_preds) == len(offset_preds) == 1 center_heatmap_pred = center_heatmap_preds[0] wh_pred = wh_preds[0] offset_pred = offset_preds[0] - target_result, avg_factor = self.get_targets(gt_bboxes, gt_labels, - center_heatmap_pred.shape, - img_metas[0]['pad_shape']) + target_result, avg_factor = self.get_targets( + gt_bboxes, gt_labels, center_heatmap_pred.shape, img_metas[0]["pad_shape"] + ) - center_heatmap_target = target_result['center_heatmap_target'] - wh_target = target_result['wh_target'] - offset_target = target_result['offset_target'] - wh_offset_target_weight = target_result['wh_offset_target_weight'] + center_heatmap_target = target_result["center_heatmap_target"] + wh_target = target_result["wh_target"] + offset_target = target_result["offset_target"] + wh_offset_target_weight = target_result["wh_offset_target_weight"] # Since the channel of wh_target and offset_target is 2, the avg_factor # of loss_center_heatmap is always 1/2 of loss_wh and loss_offset. loss_center_heatmap = self.loss_center_heatmap( - center_heatmap_pred, center_heatmap_target, avg_factor=avg_factor) - loss_wh = self.loss_wh( - wh_pred, - wh_target, - wh_offset_target_weight, - avg_factor=avg_factor * 2) + center_heatmap_pred, center_heatmap_target, avg_factor=avg_factor + ) + loss_wh = self.loss_wh(wh_pred, wh_target, wh_offset_target_weight, avg_factor=avg_factor * 2) loss_offset = self.loss_offset( - offset_pred, - offset_target, - wh_offset_target_weight, - avg_factor=avg_factor * 2) - return dict( - loss_center_heatmap=loss_center_heatmap, - loss_wh=loss_wh, - loss_offset=loss_offset) + offset_pred, offset_target, wh_offset_target_weight, avg_factor=avg_factor * 2 + ) + return dict(loss_center_heatmap=loss_center_heatmap, loss_wh=loss_wh, loss_offset=loss_offset) def get_targets(self, gt_bboxes, gt_labels, feat_shape, img_shape): """Compute regression and classification targets in multiple images. @@ -206,12 +192,10 @@ def get_targets(self, gt_bboxes, gt_labels, feat_shape, img_shape): width_ratio = float(feat_w / img_w) height_ratio = float(feat_h / img_h) - center_heatmap_target = gt_bboxes[-1].new_zeros( - [bs, self.num_classes, feat_h, feat_w]) + center_heatmap_target = gt_bboxes[-1].new_zeros([bs, self.num_classes, feat_h, feat_w]) wh_target = gt_bboxes[-1].new_zeros([bs, 2, feat_h, feat_w]) offset_target = gt_bboxes[-1].new_zeros([bs, 2, feat_h, feat_w]) - wh_offset_target_weight = gt_bboxes[-1].new_zeros( - [bs, 2, feat_h, feat_w]) + wh_offset_target_weight = gt_bboxes[-1].new_zeros([bs, 2, feat_h, feat_w]) for batch_id in range(bs): gt_bbox = gt_bboxes[batch_id] @@ -225,12 +209,10 @@ def get_targets(self, gt_bboxes, gt_labels, feat_shape, img_shape): ctx, cty = ct scale_box_h = (gt_bbox[j][3] - gt_bbox[j][1]) * height_ratio scale_box_w = (gt_bbox[j][2] - gt_bbox[j][0]) * width_ratio - radius = gaussian_radius([scale_box_h, scale_box_w], - min_overlap=0.3) + radius = gaussian_radius([scale_box_h, scale_box_w], min_overlap=0.3) radius = max(0, int(radius)) ind = gt_label[j] - gen_gaussian_target(center_heatmap_target[batch_id, ind], - [ctx_int, cty_int], radius) + gen_gaussian_target(center_heatmap_target[batch_id, ind], [ctx_int, cty_int], radius) wh_target[batch_id, 0, cty_int, ctx_int] = scale_box_w wh_target[batch_id, 1, cty_int, ctx_int] = scale_box_h @@ -245,17 +227,12 @@ def get_targets(self, gt_bboxes, gt_labels, feat_shape, img_shape): center_heatmap_target=center_heatmap_target, wh_target=wh_target, offset_target=offset_target, - wh_offset_target_weight=wh_offset_target_weight) + wh_offset_target_weight=wh_offset_target_weight, + ) return target_result, avg_factor - @force_fp32(apply_to=('center_heatmap_preds', 'wh_preds', 'offset_preds')) - def get_bboxes(self, - center_heatmap_preds, - wh_preds, - offset_preds, - img_metas, - rescale=True, - with_nms=False): + @force_fp32(apply_to=("center_heatmap_preds", "wh_preds", "offset_preds")) + def get_bboxes(self, center_heatmap_preds, wh_preds, offset_preds, img_metas, rescale=True, with_nms=False): """Transform network output for a batch into bbox predictions. Args: @@ -280,27 +257,24 @@ def get_bboxes(self, each element represents the class label of the corresponding box. """ - assert len(center_heatmap_preds) == len(wh_preds) == len( - offset_preds) == 1 + assert len(center_heatmap_preds) == len(wh_preds) == len(offset_preds) == 1 result_list = [] for img_id in range(len(img_metas)): result_list.append( self._get_bboxes_single( - center_heatmap_preds[0][img_id:img_id + 1, ...], - wh_preds[0][img_id:img_id + 1, ...], - offset_preds[0][img_id:img_id + 1, ...], + center_heatmap_preds[0][img_id : img_id + 1, ...], + wh_preds[0][img_id : img_id + 1, ...], + offset_preds[0][img_id : img_id + 1, ...], img_metas[img_id], rescale=rescale, - with_nms=with_nms)) + with_nms=with_nms, + ) + ) return result_list - def _get_bboxes_single(self, - center_heatmap_pred, - wh_pred, - offset_pred, - img_meta, - rescale=False, - with_nms=True): + def _get_bboxes_single( + self, center_heatmap_pred, wh_pred, offset_pred, img_meta, rescale=False, with_nms=True + ): """Transform outputs of a single image into bbox results. Args: @@ -328,33 +302,25 @@ def _get_bboxes_single(self, center_heatmap_pred, wh_pred, offset_pred, - img_meta['batch_input_shape'], + img_meta["batch_input_shape"], k=self.test_cfg.topk, - kernel=self.test_cfg.local_maximum_kernel) + kernel=self.test_cfg.local_maximum_kernel, + ) det_bboxes = batch_det_bboxes.view([-1, 5]) det_labels = batch_labels.view(-1) - batch_border = det_bboxes.new_tensor(img_meta['border'])[..., - [2, 0, 2, 0]] + batch_border = det_bboxes.new_tensor(img_meta["border"])[..., [2, 0, 2, 0]] det_bboxes[..., :4] -= batch_border if rescale: - det_bboxes[..., :4] /= det_bboxes.new_tensor( - img_meta['scale_factor']) + det_bboxes[..., :4] /= det_bboxes.new_tensor(img_meta["scale_factor"]) if with_nms: - det_bboxes, det_labels = self._bboxes_nms(det_bboxes, det_labels, - self.test_cfg) + det_bboxes, det_labels = self._bboxes_nms(det_bboxes, det_labels, self.test_cfg) return det_bboxes, det_labels - def decode_heatmap(self, - center_heatmap_pred, - wh_pred, - offset_pred, - img_shape, - k=100, - kernel=3): + def decode_heatmap(self, center_heatmap_pred, wh_pred, offset_pred, img_shape, k=100, kernel=3): """Transform outputs into detections raw bbox prediction. Args: @@ -378,11 +344,9 @@ def decode_heatmap(self, height, width = center_heatmap_pred.shape[2:] inp_h, inp_w = img_shape - center_heatmap_pred = get_local_maximum( - center_heatmap_pred, kernel=kernel) + center_heatmap_pred = get_local_maximum(center_heatmap_pred, kernel=kernel) - *batch_dets, topk_ys, topk_xs = get_topk_from_heatmap( - center_heatmap_pred, k=k) + *batch_dets, topk_ys, topk_xs = get_topk_from_heatmap(center_heatmap_pred, k=k) batch_scores, batch_index, batch_topk_labels = batch_dets wh = transpose_and_gather_feat(wh_pred, batch_index) @@ -395,16 +359,13 @@ def decode_heatmap(self, br_y = (topk_ys + wh[..., 1] / 2) * (inp_h / height) batch_bboxes = torch.stack([tl_x, tl_y, br_x, br_y], dim=2) - batch_bboxes = torch.cat((batch_bboxes, batch_scores[..., None]), - dim=-1) + batch_bboxes = torch.cat((batch_bboxes, batch_scores[..., None]), dim=-1) return batch_bboxes, batch_topk_labels def _bboxes_nms(self, bboxes, labels, cfg): if labels.numel() > 0: max_num = cfg.max_per_img - bboxes, keep = batched_nms(bboxes[:, :4], bboxes[:, - -1].contiguous(), - labels, cfg.nms) + bboxes, keep = batched_nms(bboxes[:, :4], bboxes[:, -1].contiguous(), labels, cfg.nms) if max_num > 0: bboxes = bboxes[:max_num] labels = labels[keep][:max_num] diff --git a/mmdet/models/dense_heads/cycle_centernet_head.py b/mmdet/models/dense_heads/cycle_centernet_head.py new file mode 100644 index 00000000..463144ed --- /dev/null +++ b/mmdet/models/dense_heads/cycle_centernet_head.py @@ -0,0 +1,590 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import torch +import torch.nn as nn +from numpy import pi +from mmcv.cnn import bias_init_with_prob, normal_init +from mmcv.ops import batched_nms +from mmcv.runner import force_fp32 + +from mmdet.core import multi_apply +from mmdet.models import HEADS, build_loss +from mmdet.models.utils import gaussian_radius, gen_gaussian_target +from ..utils.gaussian_target import get_local_maximum, get_topk_from_heatmap, transpose_and_gather_feat +from .base_dense_head import BaseDenseHead +from .dense_test_mixins import BBoxTestMixin + + +@HEADS.register_module() +class CycleCenterNetHead(BaseDenseHead, BBoxTestMixin): + """Parsing Table Structures in the Wild Head. CycleCenterHead use + center point to indicate cell's position in a table. + Paper link + + Args: + in_channel (int): Number of channel in the input feature map. + feat_channel (int): Number of channel in the intermediate feature map. + loss_center_heatmap (dict | None): Config of center heatmap loss. + Default: GaussianFocalLoss. + loss_wh (dict | None): Config of wh loss. Default: L1Loss. + loss_offset (dict | None): Config of offset loss. Default: L1Loss. + train_cfg (dict | None): Training config. Useless in CenterNet, + but we keep this variable for SingleStageDetector. Default: None. + test_cfg (dict | None): Testing config of CenterNet. Default: None. + init_cfg (dict or list[dict], optional): Initialization config dict. + Default: None + """ + + def __init__( + self, + in_channel, + feat_channel, + loss_center_heatmap=dict(type="GaussianFocalLoss", loss_weight=1.0), + loss_offset=dict(type="L1Loss", loss_weight=1.0), + loss_c2v=dict(type="L1Loss", loss_weight=1.0), + loss_v2c=dict(type="L1Loss", loss_weight=0.5), + train_cfg=None, + test_cfg=None, + init_cfg=None, + ): + super(CycleCenterNetHead, self).__init__(init_cfg) + self.heatmap_head = self._build_head(in_channel, feat_channel, 2) + self.offset_head = self._build_head(in_channel, feat_channel, 2) + self.center2vertex_head = self._build_head(in_channel, feat_channel, 8) + self.vertex2center_head = self._build_head(in_channel, feat_channel, 8) + + self.loss_center_heatmap = build_loss(loss_center_heatmap) + self.loss_offset = build_loss(loss_offset) + self.loss_c2v = build_loss(loss_c2v) + self.loss_v2c = build_loss(loss_v2c) + + self.train_cfg = train_cfg + self.test_cfg = test_cfg + self.fp16_enabled = False + + def _build_head(self, in_channel, feat_channel, out_channel): + """Build head for each branch.""" + layer = nn.Sequential( + nn.Conv2d(in_channel, feat_channel, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(feat_channel, out_channel, kernel_size=1), + ) + return layer + + def init_weights(self): + """Initialize weights of the head.""" + bias_init = bias_init_with_prob(0.1) + self.heatmap_head[-1].bias.data.fill_(bias_init) + for head in [ + self.offset_head, + self.center2vertex_head, + self.vertex2center_head, + ]: + for m in head.modules(): + if isinstance(m, nn.Conv2d): + normal_init(m, std=0.001) + + def forward(self, feats): + """Forward features. Notice CenterNet head does not use FPN. + + Args: + feats (tuple[Tensor]): Features from the upstream network, each is + a 4D-tensor. + + Returns: + center_heatmap_preds (List[Tensor]): center predict heatmaps for + all levels, the channels number is 2. + wh_preds (List[Tensor]): wh predicts for all levels, the channels + number is 2. + offset_preds (List[Tensor]): offset predicts for all levels, the + channels number is 2. + """ + return multi_apply(self.forward_single, feats) + + def forward_single(self, feat): + """Forward feature of a single level. + + Args: + feat (Tensor): Feature of a single level. + + Returns: + center_heatmap_pred (Tensor): center predict heatmaps, the + channels number is 2. + wh_pred (Tensor): wh predicts, the channels number is 2. + offset_pred (Tensor): offset predicts, the channels number is 2. + """ + center_heatmap_pred = self.heatmap_head(feat).sigmoid() + offset_pred = self.offset_head(feat) + center2vertex_pred = self.center2vertex_head(feat) + vertex2center_pred = self.vertex2center_head(feat) + return ( + center_heatmap_pred, + offset_pred, + center2vertex_pred, + vertex2center_pred, + ) + + @force_fp32( + apply_to=( + "center_heatmap_preds", + "offset_preds", + "center2vertex_pred", + "vertex2center_pred", + ) + ) + def loss( + self, + center_heatmap_preds, + offset_preds, + center2vertex_pred, + vertex2center_pred, + gt_bboxes, + gt_labels, + img_metas, + gt_masks=None, + gt_bboxes_ignore=None, + ): + """Compute losses of the head. + + Args: + center_heatmap_preds (list[Tensor]): center predict heatmaps for + all levels with shape (B, 2, H, W). + offset_preds (list[Tensor]): offset predicts for all levels + with shape (B, 2, H, W). + center2vertex_pred (list[Tensor]): center2vertex predicts for all levels + with shape (B, 8, H, W). + vertex2center_pred (list[Tensor]): vertex2center predicts for all levels + with shape (B, 8, H, W). + gt_bboxes (list[Tensor]): Ground truth bboxes for each image with + shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format. + gt_labels (list[Tensor]): class indices corresponding to each box. + img_metas (list[dict]): Meta information of each image, e.g., + image size, scaling factor, etc. + + Returns: + dict[str, Tensor]: which has components below: + - loss_center_heatmap (Tensor): loss of center heatmap. + - loss_wh (Tensor): loss of hw heatmap + - loss_offset (Tensor): loss of offset heatmap. + """ + assert ( + len(center_heatmap_preds) + == len(offset_preds) + == len(center2vertex_pred) + == len(vertex2center_pred) + == 1 + ) + center_heatmap_pred = center_heatmap_preds[0] + offset_pred = offset_preds[0] + center2vertex_pred = center2vertex_pred[0] + vertex2center_pred = vertex2center_pred[0] + + target_result, avg_factor = self.get_targets( + gt_bboxes, + gt_labels, + gt_masks, + center_heatmap_pred.shape, + img_metas[0]["pad_shape"], + center2vertex_pred, + vertex2center_pred, + ) + + center_heatmap_target = target_result["center_heatmap_target"] + offset_target = target_result["offset_target"] + offset_target_weight = target_result["offset_target_weight"] + center2vertex_target = target_result["center2vertex_target"] + vertex2center_target = target_result["vertex2center_target"] + pairing_weight = target_result["pairing_weight"] + + # Since the channel of wh_target and offset_target is 2, the avg_factor + # of loss_center_heatmap is always 1/2 of loss_wh and loss_offset. + loss_center_heatmap = self.loss_center_heatmap( + center_heatmap_pred, center_heatmap_target, avg_factor=avg_factor + ) + loss_offset = self.loss_offset( + offset_pred, + offset_target, + offset_target_weight, + avg_factor=avg_factor * 2, + ) + loss_c2v = self.loss_c2v( + center2vertex_pred, + center2vertex_target, + pairing_weight, + avg_factor=avg_factor * 8, + ) + loss_v2c = self.loss_v2c( + vertex2center_pred, + vertex2center_target, + pairing_weight, + avg_factor=avg_factor * 8, + ) + + return dict( + loss_center_heatmap=loss_center_heatmap, + loss_offset=loss_offset, + loss_c2v=loss_c2v, + loss_v2c=loss_v2c, + ) + + def get_targets(self, gt_bboxes, gt_labels, gt_masks, feat_shape, img_shape, c2v_pred, v2c_pred): + """Compute regression and classification targets in multiple images and + compute weight for Pairing Loss. + + Args: + gt_bboxes (list[Tensor]): Ground truth bboxes for each image with + shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format. + gt_labels (list[Tensor]): class indices corresponding to each box. + feat_shape (list[int]): feature map shape with value [B, _, H, W] + img_shape (list[int]): image shape in [h, w] format. + c2v_pred (list[Tensor]): center2vertex branch shape with value [B, 8, H, W] + v2c_pred (list[Tensor]): vertex2center branch shape with value [B, 8, H, W] + + Returns: + tuple[dict,float]: The float value is mean avg_factor, the dict has + components below: + - center_heatmap_target (Tensor): targets of center heatmap, \ + shape (B, 2, H, W). + - wh_target (Tensor): targets of wh predict, shape \ + (B, 2, H, W). + - offset_target (Tensor): targets of offset predict, shape \ + (B, 2, H, W). + - wh_offset_target_weight (Tensor): weights of wh and offset \ + predict, shape (B, 2, H, W). + - center2vertex_target (Tensor): targets of center2vertex predict, shape \ + (B, 8, H, W). + - vertex2center_target (Tensor): targets of vertex2center predict, shape \ + (B, 8, H, W). + """ + img_h, img_w = img_shape[:2] + bs, _, feat_h, feat_w = feat_shape + + width_ratio = float(feat_w / img_w) + height_ratio = float(feat_h / img_h) + + center_heatmap_target = gt_bboxes[-1].new_zeros([bs, 2, feat_h, feat_w]) + offset_target = gt_bboxes[-1].new_zeros([bs, 2, feat_h, feat_w]) + offset_target_weight = gt_bboxes[-1].new_zeros([bs, 2, feat_h, feat_w]) + c2v_target = gt_bboxes[-1].new_zeros([bs, 8, feat_h, feat_w]) + v2c_target = gt_bboxes[-1].new_zeros([bs, 8, feat_h, feat_w]) + pairing_weight = gt_bboxes[-1].new_zeros([bs, 8, feat_h, feat_w]) + + radius = gaussian_radius([feat_h / 8, feat_w / 8], min_overlap=0.3) + radius = max(0, int(radius)) + + for batch_id in range(bs): + for mask in gt_masks[batch_id].masks: + mask = mask[0] + ctx, cty = (mask[0] + mask[2] + mask[4] + mask[6]) / 4, (mask[1] + mask[3] + mask[5] + mask[7]) / 4 + ctx_int, cty_int = int(ctx * width_ratio), int(cty * height_ratio) + + gen_gaussian_target( + heatmap=center_heatmap_target[batch_id, 0], + center=[ctx_int, cty_int], + radius=radius, + ) + + tl_x, tr_x, br_x, bl_x = map(lambda x: x * width_ratio, (mask[0], mask[2], mask[4], mask[6])) + tl_y, tr_y, br_y, bl_y = map(lambda x: x * height_ratio, (mask[1], mask[3], mask[5], mask[7])) + tl_x_int, tl_y_int, tr_x_int, tr_y_int, br_x_int, br_y_int, bl_x_int, bl_y_int = map( + int, (tl_x, tl_y, tr_x, tr_y, br_x, br_y, bl_x, bl_y) + ) + tl_x_int, tr_x_int, br_x_int, bl_x_int = map( + lambda x: feat_w - 1 if x >= feat_w else x if x >= 0 else 0, + (tl_x_int, tr_x_int, br_x_int, bl_x_int), + ) + tl_y_int, tr_y_int, br_y_int, bl_y_int = map( + lambda x: feat_h - 1 if x >= feat_h else x if x >= 0 else 0, + (tl_y_int, tr_y_int, br_y_int, bl_y_int), + ) + for x, y in ( + (tl_x_int, tl_y_int), + (tr_x_int, tr_y_int), + (br_x_int, br_y_int), + (bl_x_int, bl_y_int), + ): + if center_heatmap_target[batch_id, 1, y, x] != 1.0: + gen_gaussian_target( + heatmap=center_heatmap_target[batch_id, 1], + center=[x, y], + radius=radius, + ) + + offset_target[batch_id, 0, cty_int, ctx_int] = ctx - ctx_int + offset_target[batch_id, 0, tl_y_int, tl_x_int] = tl_x - tl_x_int + offset_target[batch_id, 0, tr_y_int, tr_x_int] = tr_x - tr_x_int + offset_target[batch_id, 0, br_y_int, br_x_int] = br_x - br_x_int + offset_target[batch_id, 0, bl_y_int, bl_x_int] = bl_x - bl_x_int + + offset_target[batch_id, 1, cty_int, ctx_int] = cty - cty_int + offset_target[batch_id, 1, tl_y_int, tl_x_int] = tl_y - tl_y_int + offset_target[batch_id, 1, tr_y_int, tr_x_int] = tr_y - tr_y_int + offset_target[batch_id, 1, br_y_int, br_x_int] = br_y - br_y_int + offset_target[batch_id, 1, bl_y_int, bl_x_int] = bl_y - bl_y_int + + offset_target_weight[batch_id, :, cty_int, ctx_int] = 1 + offset_target_weight[batch_id, :, tl_y_int, tl_x_int] = 1 + offset_target_weight[batch_id, :, tr_y_int, tr_x_int] = 1 + offset_target_weight[batch_id, :, br_y_int, br_x_int] = 1 + offset_target_weight[batch_id, :, bl_y_int, bl_x_int] = 1 + + c2v_target[batch_id, 0, cty_int, ctx_int] = tl_x - ctx + c2v_target[batch_id, 1, cty_int, ctx_int] = tl_y - cty + c2v_target[batch_id, 2, cty_int, ctx_int] = tr_x - ctx + c2v_target[batch_id, 3, cty_int, ctx_int] = tr_y - cty + c2v_target[batch_id, 4, cty_int, ctx_int] = br_x - ctx + c2v_target[batch_id, 5, cty_int, ctx_int] = br_y - cty + c2v_target[batch_id, 6, cty_int, ctx_int] = bl_x - ctx + c2v_target[batch_id, 7, cty_int, ctx_int] = bl_y - cty + + v2c_target[batch_id, 0, tl_y_int, tl_x_int] = ctx - tl_x + v2c_target[batch_id, 1, tl_y_int, tl_x_int] = cty - tl_y + v2c_target[batch_id, 2, tr_y_int, tr_x_int] = ctx - tr_x + v2c_target[batch_id, 3, tr_y_int, tr_x_int] = cty - tr_y + v2c_target[batch_id, 4, br_y_int, br_x_int] = ctx - br_x + v2c_target[batch_id, 5, br_y_int, br_x_int] = cty - br_y + v2c_target[batch_id, 6, bl_y_int, bl_x_int] = ctx - bl_x + v2c_target[batch_id, 7, bl_y_int, bl_x_int] = cty - bl_y + + # Pairing loss + for idx, v_x_int, v_y_int in ( + (0, tl_x_int, tl_y_int), + (1, tr_x_int, tr_y_int), + (2, br_x_int, br_y_int), + (3, bl_x_int, bl_y_int), + ): + for k in range(2): + D_cv = min( + torch.tensor(1.0), + ( + torch.abs( + c2v_pred[batch_id, 2 * idx + k, cty_int, ctx_int] + - c2v_target[batch_id, 2 * idx + k, cty_int, ctx_int] + ) + + torch.abs( + v2c_pred[batch_id, 2 * idx + k, v_y_int, v_x_int] + - v2c_target[batch_id, 2 * idx + k, v_y_int, v_x_int] + ) + ) + / torch.abs(c2v_target[batch_id, 2 * idx + k, cty_int, ctx_int]), + ) + w = 1 - torch.exp(-pi * D_cv) + pairing_weight[batch_id, 2 * idx + k, cty_int, ctx_int] += w + pairing_weight[batch_id, 2 * idx + k, v_y_int, v_x_int] += w + + avg_factor = max(1, center_heatmap_target.eq(1).sum()) + target_result = dict( + center_heatmap_target=center_heatmap_target, + offset_target=offset_target, + offset_target_weight=offset_target_weight, + center2vertex_target=c2v_target, + vertex2center_target=v2c_target, + pairing_weight=pairing_weight, + ) + return target_result, avg_factor + + @force_fp32( + apply_to=( + "center_heatmap_preds", + "offset_preds", + "center2vertex_preds", + "vertex2center_preds", + ) + ) + def get_bboxes( + self, + center_heatmap_preds, + offset_preds, + center2vertex_preds, + vertex2center_preds, + img_metas, + rescale=True, + with_nms=False, + ): + """Transform network output for a batch into bbox predictions. + + Args: + center_heatmap_preds (list[Tensor]): Center predict heatmaps for + all levels with shape (B, 2, H, W). + wh_preds (list[Tensor]): WH predicts for all levels with + shape (B, 2, H, W). + offset_preds (list[Tensor]): Offset predicts for all levels + with shape (B, 2, H, W). + center2vertex_pred (list[Tensor]): center2vertex predicts for all levels + with shape (B, 8, H, W). + vertex2center_pred (list[Tensor]): vertex2center predicts for all levels + with shape (B, 8, H, W). + img_metas (list[dict]): Meta information of each image, e.g., + image size, scaling factor, etc. + rescale (bool): If True, return boxes in original image space. + Default: True. + with_nms (bool): If True, do nms before return boxes. + Default: False. + + Returns: + list[tuple[Tensor, Tensor]]: Each item in result_list is 2-tuple. + The first item is an (n, 5) tensor, where 5 represent + (tl_x, tl_y, br_x, br_y, score) and the score between 0 and 1. + The shape of the second tensor in the tuple is (n,), and + each element represents the class label of the corresponding + box. + """ + assert ( + len(center_heatmap_preds) + == len(offset_preds) + == len(center2vertex_preds) + == len(vertex2center_preds) + == 1 + ) + wh_preds = [torch.zeros_like(offset_preds[0])] + # wh_preds_vc = [torch.ones_like(offset_preds[0]) * 5] + wh_preds[0][:, 0, ...] = ( + -center2vertex_preds[0][:, 0, ...] + + center2vertex_preds[0][:, 2, ...] + + center2vertex_preds[0][:, 4, ...] + - center2vertex_preds[0][:, 6, ...] + ) / 2 + wh_preds[0][:, 1, ...] = ( + -center2vertex_preds[0][:, 1, ...] + - center2vertex_preds[0][:, 3, ...] + + center2vertex_preds[0][:, 5, ...] + + center2vertex_preds[0][:, 7, ...] + ) / 2 + result_list = [] + for img_id in range(len(img_metas)): + result_list.append( + self._get_bboxes_single( + center_heatmap_preds[0][img_id : img_id + 1, 0:1, ...], + wh_preds[0][img_id : img_id + 1, ...], + offset_preds[0][img_id : img_id + 1, ...], + img_metas[img_id], + rescale=rescale, + with_nms=with_nms, + ) + ) + # for img_id in range(len(img_metas)): + # result_list.append( + # self._get_bboxes_single( + # center_heatmap_preds[0][img_id : img_id + 1, 1:2, ...], + # wh_preds_vc[0][img_id : img_id + 1, ...], + # offset_preds[0][img_id : img_id + 1, ...], + # img_metas[img_id], + # rescale=rescale, + # with_nms=with_nms, + # ) + # ) + # ) + return result_list + + def _get_bboxes_single( + self, + center_heatmap_pred, + wh_pred, + offset_pred, + img_meta, + rescale=False, + with_nms=True, + ): + """Transform outputs of a single image into bbox results. + + Args: + center_heatmap_pred (Tensor): Center heatmap for current level with + shape (1, 2, H, W). + wh_pred (Tensor): WH heatmap for current level with shape + (1, 2, H, W). + offset_pred (Tensor): Offset for current level with shape + (1, corner_offset_channels, H, W). + img_meta (dict): Meta information of current image, e.g., + image size, scaling factor, etc. + rescale (bool): If True, return boxes in original image space. + Default: False. + with_nms (bool): If True, do nms before return boxes. + Default: True. + + Returns: + tuple[Tensor, Tensor]: The first item is an (n, 5) tensor, where + 5 represent (tl_x, tl_y, br_x, br_y, score) and the score + between 0 and 1. The shape of the second tensor in the tuple + is (n,), and each element represents the class label of the + corresponding box. + """ + batch_det_bboxes, batch_labels = self.decode_heatmap( + center_heatmap_pred, + wh_pred, + offset_pred, + img_meta["batch_input_shape"], + k=self.test_cfg.topk, + kernel=self.test_cfg.local_maximum_kernel, + ) + + det_bboxes = batch_det_bboxes.view([-1, 5]) + det_labels = batch_labels.view(-1) + + batch_border = det_bboxes.new_tensor(img_meta["border"])[..., [2, 0, 2, 0]] + det_bboxes[..., :4] -= batch_border + + if rescale: + det_bboxes[..., :4] /= det_bboxes.new_tensor(img_meta["scale_factor"]) + + if with_nms: + det_bboxes, det_labels = self._bboxes_nms(det_bboxes, det_labels, self.test_cfg) + return det_bboxes, det_labels + + def decode_heatmap( + self, + center_heatmap_pred, + wh_pred, + offset_pred, + img_shape, + k=100, + kernel=3, + ): + """Transform outputs into detections raw bbox prediction. + + Args: + center_heatmap_pred (Tensor): center predict heatmap, + shape (B, 2, H, W). + wh_pred (Tensor): wh predict, shape (B, 2, H, W). + offset_pred (Tensor): offset predict, shape (B, 2, H, W). + img_shape (list[int]): image shape in [h, w] format. + k (int): Get top k center keypoints from heatmap. Default 100. + kernel (int): Max pooling kernel for extract local maximum pixels. + Default 3. + + Returns: + tuple[torch.Tensor]: Decoded output of CenterNetHead, containing + the following Tensors: + + - batch_bboxes (Tensor): Coords of each box with shape (B, k, 5) + - batch_topk_labels (Tensor): Categories of each box with \ + shape (B, k) + """ + height, width = center_heatmap_pred.shape[2:] + inp_h, inp_w = img_shape + + center_heatmap_pred = get_local_maximum(center_heatmap_pred, kernel=kernel) + + *batch_dets, topk_ys, topk_xs = get_topk_from_heatmap(center_heatmap_pred, k=k) + batch_scores, batch_index, batch_topk_labels = batch_dets + + wh = transpose_and_gather_feat(wh_pred, batch_index) + offset = transpose_and_gather_feat(offset_pred, batch_index) + topk_xs = topk_xs + offset[..., 0] + topk_ys = topk_ys + offset[..., 1] + tl_x = (topk_xs - wh[..., 0] / 2) * (inp_w / width) + tl_y = (topk_ys - wh[..., 1] / 2) * (inp_h / height) + br_x = (topk_xs + wh[..., 0] / 2) * (inp_w / width) + br_y = (topk_ys + wh[..., 1] / 2) * (inp_h / height) + + batch_bboxes = torch.stack([tl_x, tl_y, br_x, br_y], dim=2) + batch_bboxes = torch.cat((batch_bboxes, batch_scores[..., None]), dim=-1) + return batch_bboxes, batch_topk_labels + + def _bboxes_nms(self, bboxes, labels, cfg): + if labels.numel() > 0: + max_num = cfg.max_per_img + bboxes, keep = batched_nms(bboxes[:, :4], bboxes[:, -1].contiguous(), labels, cfg.nms) + if max_num > 0: + bboxes = bboxes[:max_num] + labels = labels[keep][:max_num] + + return bboxes, labels diff --git a/mmdet/models/detectors/single_stage.py b/mmdet/models/detectors/single_stage.py index c375c72d..aeeef0d0 100644 --- a/mmdet/models/detectors/single_stage.py +++ b/mmdet/models/detectors/single_stage.py @@ -16,18 +16,12 @@ class SingleStageDetector(BaseDetector): output features of the backbone+neck. """ - def __init__(self, - backbone, - neck=None, - bbox_head=None, - train_cfg=None, - test_cfg=None, - pretrained=None, - init_cfg=None): + def __init__( + self, backbone, neck=None, bbox_head=None, train_cfg=None, test_cfg=None, pretrained=None, init_cfg=None + ): super(SingleStageDetector, self).__init__(init_cfg) if pretrained: - warnings.warn('DeprecationWarning: pretrained is deprecated, ' - 'please use "init_cfg" instead') + warnings.warn("DeprecationWarning: pretrained is deprecated, " 'please use "init_cfg" instead') backbone.pretrained = pretrained self.backbone = build_backbone(backbone) if neck is not None: @@ -54,12 +48,7 @@ def forward_dummy(self, img): outs = self.bbox_head(x) return outs - def forward_train(self, - img, - img_metas, - gt_bboxes, - gt_labels, - gt_bboxes_ignore=None): + def forward_train(self, img, img_metas, gt_bboxes, gt_labels, gt_masks=None, gt_bboxes_ignore=None): """ Args: img (Tensor): Input images of shape (N, C, H, W). @@ -80,8 +69,7 @@ def forward_train(self, """ super(SingleStageDetector, self).forward_train(img, img_metas) x = self.extract_feat(img) - losses = self.bbox_head.forward_train(x, img_metas, gt_bboxes, - gt_labels, gt_bboxes_ignore) + losses = self.bbox_head.forward_train(x, img_metas, gt_bboxes, gt_labels, gt_masks, gt_bboxes_ignore) return losses def simple_test(self, img, img_metas, rescale=False): @@ -99,8 +87,7 @@ def simple_test(self, img, img_metas, rescale=False): corresponds to each class. """ feat = self.extract_feat(img) - results_list = self.bbox_head.simple_test( - feat, img_metas, rescale=rescale) + results_list = self.bbox_head.simple_test(feat, img_metas, rescale=rescale) bbox_results = [ bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes) for det_bboxes, det_labels in results_list @@ -125,13 +112,12 @@ def aug_test(self, imgs, img_metas, rescale=False): The outer list corresponds to each image. The inner list corresponds to each class. """ - assert hasattr(self.bbox_head, 'aug_test'), \ - f'{self.bbox_head.__class__.__name__}' \ - ' does not support test-time augmentation' + assert hasattr(self.bbox_head, "aug_test"), ( + f"{self.bbox_head.__class__.__name__}" " does not support test-time augmentation" + ) feats = self.extract_feats(imgs) - results_list = self.bbox_head.aug_test( - feats, img_metas, rescale=rescale) + results_list = self.bbox_head.aug_test(feats, img_metas, rescale=rescale) bbox_results = [ bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes) for det_bboxes, det_labels in results_list @@ -155,17 +141,16 @@ def onnx_export(self, img, img_metas, with_nms=True): # get shape as tensor img_shape = torch._shape_as_tensor(img)[2:] - img_metas[0]['img_shape_for_onnx'] = img_shape + img_metas[0]["img_shape_for_onnx"] = img_shape # get pad input shape to support onnx dynamic shape for exporting # `CornerNet` and `CentripetalNet`, which 'pad_shape' is used # for inference - img_metas[0]['pad_shape_for_onnx'] = img_shape + img_metas[0]["pad_shape_for_onnx"] = img_shape if len(outs) == 2: # add dummy score_factor outs = (*outs, None) # TODO Can we change to `get_bboxes` when `onnx_export` fail - det_bboxes, det_labels = self.bbox_head.onnx_export( - *outs, img_metas, with_nms=with_nms) + det_bboxes, det_labels = self.bbox_head.onnx_export(*outs, img_metas, with_nms=with_nms) return det_bboxes, det_labels diff --git a/mmdet/models/necks/dla_neck.py b/mmdet/models/necks/dla_neck.py new file mode 100644 index 00000000..ad520033 --- /dev/null +++ b/mmdet/models/necks/dla_neck.py @@ -0,0 +1,250 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import math + +import numpy as np +from mmcv.cnn import ConvModule, build_conv_layer +from mmcv.runner import BaseModule +from torch import nn as nn + +from ..builder import NECKS + + +def fill_up_weights(up): + """Simulated bilinear upsampling kernel. + + Args: + up (nn.Module): ConvTranspose2d module. + """ + w = up.weight.data + f = math.ceil(w.size(2) / 2) + c = (2 * f - 1 - f % 2) / (2.0 * f) + for i in range(w.size(2)): + for j in range(w.size(3)): + w[0, 0, i, j] = (1 - math.fabs(i / f - c)) * ( + 1 - math.fabs(j / f - c) + ) + for c in range(1, w.size(0)): + w[c, 0, :, :] = w[0, 0, :, :] + + +class IDAUpsample(BaseModule): + """Iterative Deep Aggregation (IDA) Upsampling module to upsample features + of different scales to a similar scale. + + Args: + out_channels (int): Number of output channels for DeformConv. + in_channels (List[int]): List of input channels of multi-scale + feature maps. + kernel_sizes (List[int]): List of size of the convolving + kernel of different scales. + norm_cfg (dict, optional): Config dict for normalization layer. + Default: None. + use_dcn (bool, optional): If True, use DCNv2. Default: True. + """ + + def __init__( + self, + out_channels, + in_channels, + kernel_sizes, + norm_cfg=None, + use_dcn=True, + init_cfg=None, + ): + super(IDAUpsample, self).__init__(init_cfg) + self.use_dcn = use_dcn + self.projs = nn.ModuleList() + self.ups = nn.ModuleList() + self.nodes = nn.ModuleList() + + for i in range(1, len(in_channels)): + in_channel = in_channels[i] + up_kernel_size = int(kernel_sizes[i]) + proj = ConvModule( + in_channel, + out_channels, + 3, + padding=1, + bias=True, + conv_cfg=dict(type="DCNv2") if self.use_dcn else None, + norm_cfg=norm_cfg, + ) + node = ConvModule( + out_channels, + out_channels, + 3, + padding=1, + bias=True, + conv_cfg=dict(type="DCNv2") if self.use_dcn else None, + norm_cfg=norm_cfg, + ) + up = build_conv_layer( + dict(type="deconv"), + out_channels, + out_channels, + up_kernel_size * 2, + stride=up_kernel_size, + padding=up_kernel_size // 2, + output_padding=0, + groups=out_channels, + bias=False, + ) + + self.projs.append(proj) + self.ups.append(up) + self.nodes.append(node) + + def forward(self, mlvl_features, start_level, end_level): + """Forward function. + + Args: + mlvl_features (list[torch.Tensor]): Features from multiple layers. + start_level (int): Start layer for feature upsampling. + end_level (int): End layer for feature upsampling. + """ + for i in range(start_level, end_level - 1): + upsample = self.ups[i - start_level] + project = self.projs[i - start_level] + mlvl_features[i + 1] = upsample(project(mlvl_features[i + 1])) + node = self.nodes[i - start_level] + mlvl_features[i + 1] = node(mlvl_features[i + 1] + mlvl_features[i]) + + +class DLAUpsample(BaseModule): + """Deep Layer Aggregation (DLA) Upsampling module for different scales + feature extraction, upsampling and fusion, It consists of groups of + IDAupsample modules. + + Args: + start_level (int): The start layer. + channels (List[int]): List of input channels of multi-scale + feature maps. + scales(List[int]): List of scale of different layers' feature. + in_channels (NoneType, optional): List of input channels of + different scales. Default: None. + norm_cfg (dict, optional): Config dict for normalization layer. + Default: None. + use_dcn (bool, optional): Whether to use dcn in IDAup module. + Default: True. + """ + + def __init__( + self, + start_level, + channels, + scales, + in_channels=None, + norm_cfg=None, + use_dcn=True, + init_cfg=None, + ): + super(DLAUpsample, self).__init__(init_cfg) + self.start_level = start_level + if in_channels is None: + in_channels = channels + self.channels = channels + channels = list(channels) + scales = np.array(scales, dtype=int) + for i in range(len(channels) - 1): + j = -i - 2 + setattr( + self, + "ida_{}".format(i), + IDAUpsample( + channels[j], + in_channels[j:], + scales[j:] // scales[j], + norm_cfg, + use_dcn, + ), + ) + scales[j + 1 :] = scales[j] + in_channels[j + 1 :] = [channels[j] for _ in channels[j + 1 :]] + + def forward(self, mlvl_features): + """Forward function. + + Args: + mlvl_features(list[torch.Tensor]): Features from multi-scale + layers. + + Returns: + tuple[torch.Tensor]: Up-sampled features of different layers. + """ + outs = [mlvl_features[-1]] + for i in range(len(mlvl_features) - self.start_level - 1): + ida = getattr(self, "ida_{}".format(i)) + ida(mlvl_features, len(mlvl_features) - i - 2, len(mlvl_features)) + outs.insert(0, mlvl_features[-1]) + return outs + + +@NECKS.register_module() +class DLANeck(BaseModule): + """DLA Neck. + + Args: + in_channels (list[int], optional): List of input channels + of multi-scale feature map. + start_level (int, optional): The scale level where upsampling + starts. Default: 2. + end_level (int, optional): The scale level where upsampling + ends. Default: 5. + norm_cfg (dict, optional): Config dict for normalization + layer. Default: None. + use_dcn (bool, optional): Whether to use dcn in IDAup module. + Default: True. + """ + + def __init__( + self, + in_channels=[16, 32, 64, 128, 256, 512], + start_level=2, + end_level=5, + norm_cfg=None, + use_dcn=True, + init_cfg=None, + ): + super(DLANeck, self).__init__(init_cfg) + self.start_level = start_level + self.end_level = end_level + scales = [2**i for i in range(len(in_channels[self.start_level :]))] + self.dla_up = DLAUpsample( + start_level=self.start_level, + channels=in_channels[self.start_level :], + scales=scales, + norm_cfg=norm_cfg, + use_dcn=use_dcn, + ) + self.ida_up = IDAUpsample( + in_channels[self.start_level], + in_channels[self.start_level : self.end_level], + [2**i for i in range(self.end_level - self.start_level)], + norm_cfg, + use_dcn, + ) + + def forward(self, x): + mlvl_features = [x[i] for i in range(len(x))] + mlvl_features = self.dla_up(mlvl_features) + outs = [] + for i in range(self.end_level - self.start_level): + outs.append(mlvl_features[i].clone()) + self.ida_up(outs, 0, len(outs)) + return [outs[-1]] + + def init_weights(self): + for m in self.modules(): + if isinstance(m, nn.ConvTranspose2d): + # In order to be consistent with the source code, + # reset the ConvTranspose2d initialization parameters + m.reset_parameters() + # Simulated bilinear upsampling kernel + fill_up_weights(m) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Conv2d): + # In order to be consistent with the source code, + # reset the Conv2d initialization parameters + m.reset_parameters() diff --git a/tools/analysis_tools/analyze_results.py b/tools/analysis_tools/analyze_results.py index 4d8b60c9..dd106671 100644 --- a/tools/analysis_tools/analyze_results.py +++ b/tools/analysis_tools/analyze_results.py @@ -6,7 +6,6 @@ import mmcv import numpy as np from mmcv import Config, DictAction - from mmdet.core.evaluation import eval_map from mmdet.core.visualization import imshow_gt_det_bboxes from mmdet.datasets import build_dataset, get_loading_pipeline @@ -42,17 +41,16 @@ def bbox_map_eval(det_result, annotation, nproc=4): else: bbox_det_result = [det_result] # mAP - iou_thrs = np.linspace( - .5, 0.95, int(np.round((0.95 - .5) / .05)) + 1, endpoint=True) + iou_thrs = np.linspace(0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True) processes = [] workers = Pool(processes=nproc) for thr in iou_thrs: - p = workers.apply_async(eval_map, (bbox_det_result, [annotation]), { - 'iou_thr': thr, - 'logger': 'silent', - 'nproc': 1 - }) + p = workers.apply_async( + eval_map, + (bbox_det_result, [annotation]), + {"iou_thr": thr, "logger": "silent", "nproc": 1}, + ) processes.append(p) workers.close() @@ -80,21 +78,13 @@ class ResultVisualizer: with the prediction result. Default: False. """ - def __init__(self, - show=False, - wait_time=0, - score_thr=0, - overlay_gt_pred=False): + def __init__(self, show=False, wait_time=0, score_thr=0, overlay_gt_pred=False): self.show = show self.wait_time = wait_time self.score_thr = score_thr self.overlay_gt_pred = overlay_gt_pred - def _save_image_gts_results(self, - dataset, - results, - performances, - out_dir=None): + def _save_image_gts_results(self, dataset, results, performances, out_dir=None): """Display or save image with groung truths and predictions from a model. @@ -114,16 +104,16 @@ def _save_image_gts_results(self, data_info = dataset.prepare_train_img(index) # calc save file path - filename = data_info['filename'] - if data_info['img_prefix'] is not None: - filename = osp.join(data_info['img_prefix'], filename) + filename = data_info["filename"] + if data_info["img_prefix"] is not None: + filename = osp.join(data_info["img_prefix"], filename) else: - filename = data_info['filename'] + filename = data_info["filename"] fname, name = osp.splitext(osp.basename(filename)) - save_filename = fname + '_' + str(round(performance, 3)) + name + save_filename = fname + "_" + str(round(performance, 3)) + name out_file = osp.join(out_dir, save_filename) imshow_gt_det_bboxes( - data_info['img'], + data_info["img"], data_info, results[index], dataset.CLASSES, @@ -137,13 +127,10 @@ def _save_image_gts_results(self, score_thr=self.score_thr, wait_time=self.wait_time, out_file=out_file, - overlay_gt_pred=self.overlay_gt_pred) + overlay_gt_pred=self.overlay_gt_pred, + ) - def evaluate_and_show(self, - dataset, - results, - topk=20, - show_dir='work_dir'): + def evaluate_and_show(self, dataset, results, topk=20, show_dir="work_dir"): """Evaluate and show results. Args: @@ -162,22 +149,17 @@ def evaluate_and_show(self, topk = len(dataset) // 2 if isinstance(results[0], dict): - good_samples, bad_samples = self.panoptic_evaluate( - dataset, results, topk=topk) + good_samples, bad_samples = self.panoptic_evaluate(dataset, results, topk=topk) elif isinstance(results[0], list): - good_samples, bad_samples = self.detection_evaluate( - dataset, results, topk=topk) + good_samples, bad_samples = self.detection_evaluate(dataset, results, topk=topk) elif isinstance(results[0], tuple): results_ = [result[0] for result in results] - good_samples, bad_samples = self.detection_evaluate( - dataset, results_, topk=topk) + good_samples, bad_samples = self.detection_evaluate(dataset, results_, topk=topk) else: - raise 'The format of result is not supported yet. ' \ - 'Current dict for panoptic segmentation and list ' \ - 'or tuple for object detection are supported.' + raise "The format of result is not supported yet. " "Current dict for panoptic segmentation and list " "or tuple for object detection are supported." - good_dir = osp.abspath(osp.join(show_dir, 'good')) - bad_dir = osp.abspath(osp.join(show_dir, 'bad')) + good_dir = osp.abspath(osp.join(show_dir, "good")) + bad_dir = osp.abspath(osp.join(show_dir, "bad")) self._save_image_gts_results(dataset, results, good_samples, good_dir) self._save_image_gts_results(dataset, results, bad_samples, bad_dir) @@ -208,11 +190,11 @@ def detection_evaluate(self, dataset, results, topk=20, eval_fn=None): prog_bar = mmcv.ProgressBar(len(results)) _mAPs = {} - for i, (result, ) in enumerate(zip(results)): + for i, (result,) in enumerate(zip(results)): # self.dataset[i] should not call directly # because there is a risk of mismatch data_info = dataset.prepare_train_img(i) - mAP = eval_fn(result, data_info['ann_info']) + mAP = eval_fn(result, data_info["ann_info"]) _mAPs[i] = mAP prog_bar.update() # descending select topk image @@ -245,31 +227,32 @@ def panoptic_evaluate(self, dataset, results, topk=20): gt_json = dataset.coco.img_ann_map result_files, tmp_dir = dataset.format_results(results) - pred_json = mmcv.load(result_files['panoptic'])['annotations'] - pred_folder = osp.join(tmp_dir.name, 'panoptic') + pred_json = mmcv.load(result_files["panoptic"])["annotations"] + pred_folder = osp.join(tmp_dir.name, "panoptic") gt_folder = dataset.seg_prefix pqs = {} prog_bar = mmcv.ProgressBar(len(results)) for i in range(len(results)): data_info = dataset.prepare_train_img(i) - image_id = data_info['img_info']['id'] + image_id = data_info["img_info"]["id"] gt_ann = { - 'image_id': image_id, - 'segments_info': gt_json[image_id], - 'file_name': data_info['img_info']['segm_file'] + "image_id": image_id, + "segments_info": gt_json[image_id], + "file_name": data_info["img_info"]["segm_file"], } pred_ann = pred_json[i] pq_stat = pq_compute_single_core( - i, [(gt_ann, pred_ann)], + i, + [(gt_ann, pred_ann)], gt_folder, pred_folder, dataset.categories, dataset.file_client, - print_log=False) - pq_results, classwise_results = pq_stat.pq_average( - dataset.categories, isthing=None) - pqs[i] = pq_results['pq'] + print_log=False, + ) + pq_results, classwise_results = pq_stat.pq_average(dataset.categories, isthing=None) + pqs[i] = pq_results["pq"] prog_bar.update() if tmp_dir is not None: @@ -284,48 +267,49 @@ def panoptic_evaluate(self, dataset, results, topk=20): def parse_args(): - parser = argparse.ArgumentParser( - description='MMDet eval image prediction result for each') - parser.add_argument('config', help='test config file path') - parser.add_argument( - 'prediction_path', help='prediction path where test pkl result') - parser.add_argument( - 'show_dir', help='directory where painted images will be saved') - parser.add_argument('--show', action='store_true', help='show results') + parser = argparse.ArgumentParser(description="MMDet eval image prediction result for each") + parser.add_argument("config", help="test config file path") + parser.add_argument("prediction_path", help="prediction path where test pkl result") + parser.add_argument("show_dir", help="directory where painted images will be saved") + parser.add_argument("--show", action="store_true", help="show results") parser.add_argument( - '--wait-time', + "--wait-time", type=float, default=0, - help='the interval of show (s), 0 is block') + help="the interval of show (s), 0 is block", + ) parser.add_argument( - '--topk', + "--topk", default=20, type=int, - help='saved Number of the highest topk ' - 'and lowest topk after index sorting') + help="saved Number of the highest topk " "and lowest topk after index sorting", + ) parser.add_argument( - '--show-score-thr', + "--show-score-thr", type=float, default=0, - help='score threshold (default: 0.)') + help="score threshold (default: 0.)", + ) parser.add_argument( - '--overlay-gt-pred', - action='store_true', - help='whether to plot gts and predictions on the same image.' - 'If False, predictions and gts will be plotted on two same' - 'image which will be concatenated in vertical direction.' - 'The image above is drawn with gt, and the image below is' - 'drawn with the prediction result.') + "--overlay-gt-pred", + action="store_true", + help="whether to plot gts and predictions on the same image." + "If False, predictions and gts will be plotted on two same" + "image which will be concatenated in vertical direction." + "The image above is drawn with gt, and the image below is" + "drawn with the prediction result.", + ) parser.add_argument( - '--cfg-options', - nargs='+', + "--cfg-options", + nargs="+", action=DictAction, - help='override some settings in the used config, the key-value pair ' - 'in xxx=yyy format will be merged into config file. If the value to ' + help="override some settings in the used config, the key-value pair " + "in xxx=yyy format will be merged into config file. If the value to " 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' - 'Note that the quotation marks are necessary and that no white space ' - 'is allowed.') + "Note that the quotation marks are necessary and that no white space " + "is allowed.", + ) args = parser.parse_args() return args @@ -347,23 +331,23 @@ def main(): cfg.merge_from_dict(args.cfg_options) cfg.data.test.test_mode = True - cfg.data.test.pop('samples_per_gpu', 0) - if cfg.data.train.type in ('MultiImageMixDataset', 'ClassBalancedDataset', - 'RepeatDataset', 'ConcatDataset'): - cfg.data.test.pipeline = get_loading_pipeline( - cfg.data.train.dataset.pipeline) + cfg.data.test.pop("samples_per_gpu", 0) + if cfg.data.train.type in ( + "MultiImageMixDataset", + "ClassBalancedDataset", + "RepeatDataset", + "ConcatDataset", + ): + cfg.data.test.pipeline = get_loading_pipeline(cfg.data.train.dataset.pipeline) else: cfg.data.test.pipeline = get_loading_pipeline(cfg.data.train.pipeline) dataset = build_dataset(cfg.data.test) outputs = mmcv.load(args.prediction_path) - result_visualizer = ResultVisualizer(args.show, args.wait_time, - args.show_score_thr, - args.overlay_gt_pred) - result_visualizer.evaluate_and_show( - dataset, outputs, topk=args.topk, show_dir=args.show_dir) + result_visualizer = ResultVisualizer(args.show, args.wait_time, args.show_score_thr, args.overlay_gt_pred) + result_visualizer.evaluate_and_show(dataset, outputs, topk=args.topk, show_dir=args.show_dir) -if __name__ == '__main__': +if __name__ == "__main__": main()