birth
This commit is contained in:
Yuwen Xiong
96
segmentation/README.md
Normal file
96
segmentation/README.md
Normal file
@@ -0,0 +1,96 @@
|
||||
# FlashInternImage for Semantic Segmentation
|
||||
|
||||
This folder contains the implementation of the InternImage for semantic segmentation.
|
||||
|
||||
Our segmentation code is developed on top of [MMSegmentation v0.27.0](https://github.com/open-mmlab/mmsegmentation/tree/v0.27.0).
|
||||
|
||||
## Usage
|
||||
|
||||
### Install
|
||||
|
||||
- Clone this repo:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/OpenGVLab/DCNv4.git
|
||||
cd DCNv4
|
||||
```
|
||||
|
||||
- Create a conda virtual environment and activate it:
|
||||
|
||||
```bash
|
||||
conda create -n dcnv4 python=3.7 -y
|
||||
conda activate dcnv4
|
||||
```
|
||||
|
||||
- Install `CUDA>=10.2` with `cudnn>=7` following
|
||||
the [official installation instructions](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html)
|
||||
- Install `PyTorch>=1.10.0` and `torchvision>=0.9.0` with `CUDA>=10.2`:
|
||||
|
||||
For examples, to install torch==1.11 with CUDA==11.3 and nvcc:
|
||||
```bash
|
||||
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch -y
|
||||
conda install -c conda-forge cudatoolkit-dev=11.3 -y # to install nvcc
|
||||
```
|
||||
|
||||
- Install other requirements:
|
||||
|
||||
note: conda opencv will break torchvision as not to support GPU, so we need to install opencv using pip.
|
||||
|
||||
```bash
|
||||
conda install -c conda-forge termcolor yacs pyyaml scipy pip -y
|
||||
pip install opencv-python
|
||||
```
|
||||
|
||||
- Install `timm` and `mmcv-full` and `mmsegmentation':
|
||||
|
||||
```bash
|
||||
pip install -U openmim
|
||||
mim install mmcv-full==1.5.0
|
||||
mim install mmsegmentation==0.27.0
|
||||
pip install timm==0.6.11 mmdet==2.28.1
|
||||
```
|
||||
|
||||
- Install DCNv4
|
||||
```bash
|
||||
pip install DCNv4
|
||||
```
|
||||
|
||||
### Data Preparation
|
||||
|
||||
Prepare datasets according to the [guidelines](https://github.com/open-mmlab/mmsegmentation/blob/master/docs/en/dataset_prepare.md#prepare-datasets) in MMSegmentation.
|
||||
|
||||
|
||||
### Evaluation
|
||||
|
||||
To evaluate our `FlashInternImage` on ADE20K val, run:
|
||||
|
||||
```bash
|
||||
sh dist_test.sh <config-file> <checkpoint> <gpu-num> --eval mIoU
|
||||
```
|
||||
You can download checkpoint files from [here](https://huggingface.co/OpenGVLab/DCNv4). Then place it to segmentation/checkpoint_dir/seg.
|
||||
|
||||
For example, to evaluate the `FlashInternImage-T` with a single GPU:
|
||||
|
||||
```bash
|
||||
python test.py configs/ade20k/upernet_flash_internimage_t_512_160k_ade20k.py checkpoint_dir/seg/upernet_flash_internimage_t_512_160k_ade20k.pth --eval mIoU
|
||||
```
|
||||
|
||||
For example, to evaluate the `FlashInternImage-B` with a single node with 8 GPUs:
|
||||
|
||||
```bash
|
||||
sh dist_test.sh configs/ade20k/upernet_flash_internimage_b_512_160k_ade20k.py checkpoint_dir/seg/upernet_flash_internimage_b_512_160k_ade20k.pth 8 --eval mIoU
|
||||
```
|
||||
|
||||
### Training
|
||||
|
||||
To train an `FlashInternImage` on ADE20K, run:
|
||||
|
||||
```bash
|
||||
sh dist_train.sh <config-file> <gpu-num>
|
||||
```
|
||||
|
||||
For example, to train `FlashInternImage-T` with 8 GPU on 1 node (total batch size 16), run:
|
||||
|
||||
```bash
|
||||
sh dist_train.sh configs/ade20k/upernet_flash_internimage_t_512_160k_ade20k.py 8
|
||||
```
|
||||
54
segmentation/configs/_base_/datasets/ade20k.py
Normal file
54
segmentation/configs/_base_/datasets/ade20k.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# dataset settings
|
||||
dataset_type = 'ADE20KDataset'
|
||||
data_root = 'data/ADEChallengeData2016'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 512)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/training',
|
||||
ann_dir='annotations/training',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline))
|
||||
54
segmentation/configs/_base_/datasets/ade20k_640x640.py
Normal file
54
segmentation/configs/_base_/datasets/ade20k_640x640.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# dataset settings
|
||||
dataset_type = 'ADE20KDataset'
|
||||
data_root = 'data/ADEChallengeData2016'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (640, 640)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2560, 640), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2560, 640),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/training',
|
||||
ann_dir='annotations/training',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline))
|
||||
59
segmentation/configs/_base_/datasets/chase_db1.py
Normal file
59
segmentation/configs/_base_/datasets/chase_db1.py
Normal file
@@ -0,0 +1,59 @@
|
||||
# dataset settings
|
||||
dataset_type = 'ChaseDB1Dataset'
|
||||
data_root = 'data/CHASE_DB1'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
img_scale = (960, 999)
|
||||
crop_size = (128, 128)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=img_scale, ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=img_scale,
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img'])
|
||||
])
|
||||
]
|
||||
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type='RepeatDataset',
|
||||
times=40000,
|
||||
dataset=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/training',
|
||||
ann_dir='annotations/training',
|
||||
pipeline=train_pipeline)),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline))
|
||||
54
segmentation/configs/_base_/datasets/cityscapes.py
Normal file
54
segmentation/configs/_base_/datasets/cityscapes.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# dataset settings
|
||||
dataset_type = 'CityscapesDataset'
|
||||
data_root = 'data/cityscapes/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 1024)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 1024),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=2,
|
||||
workers_per_gpu=2,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='leftImg8bit/train',
|
||||
ann_dir='gtFine/train',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline))
|
||||
35
segmentation/configs/_base_/datasets/cityscapes_1024x1024.py
Normal file
35
segmentation/configs/_base_/datasets/cityscapes_1024x1024.py
Normal file
@@ -0,0 +1,35 @@
|
||||
_base_ = './cityscapes.py'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (1024, 1024)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 1024),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
54
segmentation/configs/_base_/datasets/cityscapes_extra.py
Normal file
54
segmentation/configs/_base_/datasets/cityscapes_extra.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# dataset settings
|
||||
dataset_type = 'CityscapesDataset'
|
||||
data_root = 'data/cityscapes/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 1024)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 1024),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=2,
|
||||
workers_per_gpu=2,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir=['leftImg8bit/train', 'leftImg8bit/train_extra'],
|
||||
ann_dir=['gtFine/train', 'refinement_final/train_extra'],
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline))
|
||||
57
segmentation/configs/_base_/datasets/coco-stuff10k.py
Normal file
57
segmentation/configs/_base_/datasets/coco-stuff10k.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# dataset settings
|
||||
dataset_type = 'COCOStuffDataset'
|
||||
data_root = 'data/coco_stuff10k'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 512)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
reduce_zero_label=True,
|
||||
img_dir='images/train2014',
|
||||
ann_dir='annotations/train2014',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
reduce_zero_label=True,
|
||||
img_dir='images/test2014',
|
||||
ann_dir='annotations/test2014',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
reduce_zero_label=True,
|
||||
img_dir='images/test2014',
|
||||
ann_dir='annotations/test2014',
|
||||
pipeline=test_pipeline))
|
||||
54
segmentation/configs/_base_/datasets/coco-stuff164k.py
Normal file
54
segmentation/configs/_base_/datasets/coco-stuff164k.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# dataset settings
|
||||
dataset_type = 'COCOStuffDataset'
|
||||
data_root = 'data/coco_stuff164k'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 512)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/train2017',
|
||||
ann_dir='annotations/train2017',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/val2017',
|
||||
ann_dir='annotations/val2017',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/val2017',
|
||||
ann_dir='annotations/val2017',
|
||||
pipeline=test_pipeline))
|
||||
59
segmentation/configs/_base_/datasets/drive.py
Normal file
59
segmentation/configs/_base_/datasets/drive.py
Normal file
@@ -0,0 +1,59 @@
|
||||
# dataset settings
|
||||
dataset_type = 'DRIVEDataset'
|
||||
data_root = 'data/DRIVE'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
img_scale = (584, 565)
|
||||
crop_size = (64, 64)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=img_scale, ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=img_scale,
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img'])
|
||||
])
|
||||
]
|
||||
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type='RepeatDataset',
|
||||
times=40000,
|
||||
dataset=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/training',
|
||||
ann_dir='annotations/training',
|
||||
pipeline=train_pipeline)),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline))
|
||||
59
segmentation/configs/_base_/datasets/hrf.py
Normal file
59
segmentation/configs/_base_/datasets/hrf.py
Normal file
@@ -0,0 +1,59 @@
|
||||
# dataset settings
|
||||
dataset_type = 'HRFDataset'
|
||||
data_root = 'data/HRF'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
img_scale = (2336, 3504)
|
||||
crop_size = (256, 256)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=img_scale, ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=img_scale,
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img'])
|
||||
])
|
||||
]
|
||||
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type='RepeatDataset',
|
||||
times=40000,
|
||||
dataset=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/training',
|
||||
ann_dir='annotations/training',
|
||||
pipeline=train_pipeline)),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline))
|
||||
54
segmentation/configs/_base_/datasets/loveda.py
Normal file
54
segmentation/configs/_base_/datasets/loveda.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# dataset settings
|
||||
dataset_type = 'LoveDADataset'
|
||||
data_root = 'data/loveDA'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 512)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(1024, 1024),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='img_dir/train',
|
||||
ann_dir='ann_dir/train',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='img_dir/val',
|
||||
ann_dir='ann_dir/val',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='img_dir/val',
|
||||
ann_dir='ann_dir/val',
|
||||
pipeline=test_pipeline))
|
||||
55
segmentation/configs/_base_/datasets/mapillary.py
Normal file
55
segmentation/configs/_base_/datasets/mapillary.py
Normal file
@@ -0,0 +1,55 @@
|
||||
# dataset settings
|
||||
dataset_type = 'MapillaryDataset'
|
||||
data_root = 'data/Mapillary/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 1024)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='MapillaryHack'),
|
||||
dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 1.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 1024),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=2,
|
||||
workers_per_gpu=2,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root='data/Mapillary/',
|
||||
img_dir=['training/images', 'validation/images'],
|
||||
ann_dir=['training/labels', 'validation/labels'],
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type='CityscapesDataset',
|
||||
data_root='data/cityscapes/',
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type='CityscapesDataset',
|
||||
data_root='data/cityscapes/',
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline))
|
||||
55
segmentation/configs/_base_/datasets/mapillary_1024x1024.py
Normal file
55
segmentation/configs/_base_/datasets/mapillary_1024x1024.py
Normal file
@@ -0,0 +1,55 @@
|
||||
# dataset settings
|
||||
dataset_type = 'MapillaryDataset'
|
||||
data_root = 'data/Mapillary/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (1024, 1024)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='MapillaryHack'),
|
||||
dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 1.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 1024),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=2,
|
||||
workers_per_gpu=2,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root='data/Mapillary/',
|
||||
img_dir=['training/images', 'validation/images'],
|
||||
ann_dir=['training/labels', 'validation/labels'],
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type='CityscapesDataset',
|
||||
data_root='data/cityscapes/',
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type='CityscapesDataset',
|
||||
data_root='data/cityscapes/',
|
||||
img_dir='leftImg8bit/val',
|
||||
ann_dir='gtFine/val',
|
||||
pipeline=test_pipeline))
|
||||
59
segmentation/configs/_base_/datasets/nyu_depth_v2.py
Normal file
59
segmentation/configs/_base_/datasets/nyu_depth_v2.py
Normal file
@@ -0,0 +1,59 @@
|
||||
# dataset settings
|
||||
dataset_type = 'NYUDepthV2Dataset'
|
||||
data_root = 'data/nyu_depth_v2/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
|
||||
crop_size = (480, 480)
|
||||
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(640, 480), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(640, 480),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='image',
|
||||
ann_dir='label40',
|
||||
split='train.txt',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='image',
|
||||
ann_dir='label40',
|
||||
split='test.txt',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='image',
|
||||
ann_dir='label40',
|
||||
split='test.txt',
|
||||
pipeline=test_pipeline))
|
||||
60
segmentation/configs/_base_/datasets/pascal_context.py
Normal file
60
segmentation/configs/_base_/datasets/pascal_context.py
Normal file
@@ -0,0 +1,60 @@
|
||||
# dataset settings
|
||||
dataset_type = 'PascalContextDataset'
|
||||
data_root = 'data/VOCdevkit/VOC2010/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
|
||||
img_scale = (520, 520)
|
||||
crop_size = (480, 480)
|
||||
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=img_scale, ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=img_scale,
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClassContext',
|
||||
split='ImageSets/SegmentationContext/train.txt',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClassContext',
|
||||
split='ImageSets/SegmentationContext/val.txt',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClassContext',
|
||||
split='ImageSets/SegmentationContext/val.txt',
|
||||
pipeline=test_pipeline))
|
||||
60
segmentation/configs/_base_/datasets/pascal_context_59.py
Normal file
60
segmentation/configs/_base_/datasets/pascal_context_59.py
Normal file
@@ -0,0 +1,60 @@
|
||||
# dataset settings
|
||||
dataset_type = 'PascalContextDataset59'
|
||||
data_root = 'data/VOCdevkit/VOC2010/'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
|
||||
img_scale = (520, 520)
|
||||
crop_size = (480, 480)
|
||||
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=img_scale, ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=img_scale,
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClassContext',
|
||||
split='ImageSets/SegmentationContext/train.txt',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClassContext',
|
||||
split='ImageSets/SegmentationContext/val.txt',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClassContext',
|
||||
split='ImageSets/SegmentationContext/val.txt',
|
||||
pipeline=test_pipeline))
|
||||
57
segmentation/configs/_base_/datasets/pascal_voc12.py
Normal file
57
segmentation/configs/_base_/datasets/pascal_voc12.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# dataset settings
|
||||
dataset_type = 'PascalVOCDataset'
|
||||
data_root = 'data/VOCdevkit/VOC2012'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 512)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClass',
|
||||
split='ImageSets/Segmentation/train.txt',
|
||||
pipeline=train_pipeline),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClass',
|
||||
split='ImageSets/Segmentation/val.txt',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='JPEGImages',
|
||||
ann_dir='SegmentationClass',
|
||||
split='ImageSets/Segmentation/val.txt',
|
||||
pipeline=test_pipeline))
|
||||
9
segmentation/configs/_base_/datasets/pascal_voc12_aug.py
Normal file
9
segmentation/configs/_base_/datasets/pascal_voc12_aug.py
Normal file
@@ -0,0 +1,9 @@
|
||||
_base_ = './pascal_voc12.py'
|
||||
# dataset settings
|
||||
data = dict(
|
||||
train=dict(
|
||||
ann_dir=['SegmentationClass', 'SegmentationClassAug'],
|
||||
split=[
|
||||
'ImageSets/Segmentation/train.txt',
|
||||
'ImageSets/Segmentation/aug.txt'
|
||||
]))
|
||||
0
segmentation/configs/_base_/datasets/potsdam.py
Normal file
0
segmentation/configs/_base_/datasets/potsdam.py
Normal file
59
segmentation/configs/_base_/datasets/stare.py
Normal file
59
segmentation/configs/_base_/datasets/stare.py
Normal file
@@ -0,0 +1,59 @@
|
||||
# dataset settings
|
||||
dataset_type = 'STAREDataset'
|
||||
data_root = 'data/STARE'
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
img_scale = (605, 700)
|
||||
crop_size = (128, 128)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations'),
|
||||
dict(type='Resize', img_scale=img_scale, ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=img_scale,
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img'])
|
||||
])
|
||||
]
|
||||
|
||||
data = dict(
|
||||
samples_per_gpu=4,
|
||||
workers_per_gpu=4,
|
||||
train=dict(
|
||||
type='RepeatDataset',
|
||||
times=40000,
|
||||
dataset=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/training',
|
||||
ann_dir='annotations/training',
|
||||
pipeline=train_pipeline)),
|
||||
val=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline),
|
||||
test=dict(
|
||||
type=dataset_type,
|
||||
data_root=data_root,
|
||||
img_dir='images/validation',
|
||||
ann_dir='annotations/validation',
|
||||
pipeline=test_pipeline))
|
||||
14
segmentation/configs/_base_/default_runtime.py
Normal file
14
segmentation/configs/_base_/default_runtime.py
Normal file
@@ -0,0 +1,14 @@
|
||||
# yapf:disable
|
||||
log_config = dict(
|
||||
interval=50,
|
||||
hooks=[
|
||||
dict(type='TextLoggerHook', by_epoch=False),
|
||||
# dict(type='TensorboardLoggerHook')
|
||||
])
|
||||
# yapf:enable
|
||||
dist_params = dict(backend='nccl')
|
||||
log_level = 'INFO'
|
||||
load_from = None
|
||||
resume_from = None
|
||||
workflow = [('train', 1)]
|
||||
cudnn_benchmark = True
|
||||
138
segmentation/configs/_base_/models/mask2former_beit.py
Normal file
138
segmentation/configs/_base_/models/mask2former_beit.py
Normal file
@@ -0,0 +1,138 @@
|
||||
# model_cfg
|
||||
num_things_classes = 100
|
||||
num_stuff_classes = 50
|
||||
num_classes = num_things_classes + num_stuff_classes
|
||||
norm_cfg = dict(type='SyncBN', requires_grad=True)
|
||||
model = dict(
|
||||
type='EncoderDecoderMask2Former',
|
||||
pretrained=None,
|
||||
backbone=dict(
|
||||
type='XCiT',
|
||||
patch_size=16,
|
||||
embed_dim=384,
|
||||
depth=12,
|
||||
num_heads=8,
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
use_abs_pos_emb=True,
|
||||
use_rel_pos_bias=False,
|
||||
),
|
||||
decode_head=dict(
|
||||
type='Mask2FormerHead',
|
||||
in_channels=[256, 512, 1024, 2048], # pass to pixel_decoder inside
|
||||
# strides=[4, 8, 16, 32],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
in_index=[0, 1, 2, 3],
|
||||
num_things_classes=num_things_classes,
|
||||
num_stuff_classes=num_stuff_classes,
|
||||
num_queries=100,
|
||||
num_transformer_feat_level=3,
|
||||
pixel_decoder=dict(
|
||||
type='MSDeformAttnPixelDecoder',
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiScaleDeformableAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
ffn_cfgs=dict(
|
||||
type='FFN',
|
||||
embed_dims=256,
|
||||
feedforward_channels=1024,
|
||||
num_fcs=2,
|
||||
ffn_drop=0.0,
|
||||
act_cfg=dict(type='ReLU', inplace=True)),
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
init_cfg=None),
|
||||
enforce_decoder_input_project=False,
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
transformer_decoder=dict(
|
||||
type='DetrTransformerDecoder',
|
||||
return_intermediate=True,
|
||||
num_layers=9,
|
||||
transformerlayers=dict(
|
||||
type='DetrTransformerDecoderLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiheadAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
attn_drop=0.0,
|
||||
proj_drop=0.0,
|
||||
dropout_layer=None,
|
||||
batch_first=False),
|
||||
ffn_cfgs=dict(
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
act_cfg=dict(type='ReLU', inplace=True),
|
||||
ffn_drop=0.0,
|
||||
dropout_layer=None,
|
||||
add_identity=True),
|
||||
feedforward_channels=2048,
|
||||
operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
|
||||
'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=False,
|
||||
loss_weight=2.0,
|
||||
reduction='mean',
|
||||
class_weight=[1.0] * num_classes + [0.1]),
|
||||
loss_mask=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=True,
|
||||
reduction='mean',
|
||||
loss_weight=5.0),
|
||||
loss_dice=dict(
|
||||
type='DiceLoss',
|
||||
use_sigmoid=True,
|
||||
activate=True,
|
||||
reduction='mean',
|
||||
naive_dice=True,
|
||||
eps=1.0,
|
||||
loss_weight=5.0)),
|
||||
train_cfg=dict(
|
||||
num_points=12544,
|
||||
oversample_ratio=3.0,
|
||||
importance_sample_ratio=0.75,
|
||||
assigner=dict(
|
||||
type='MaskHungarianAssigner',
|
||||
cls_cost=dict(type='ClassificationCost', weight=2.0),
|
||||
mask_cost=dict(
|
||||
type='CrossEntropyLossCost', weight=5.0, use_sigmoid=True),
|
||||
dice_cost=dict(
|
||||
type='DiceCost', weight=5.0, pred_act=True, eps=1.0)),
|
||||
sampler=dict(type='MaskPseudoSampler')),
|
||||
test_cfg=dict(
|
||||
panoptic_on=True,
|
||||
# For now, the dataset does not support
|
||||
# evaluating semantic segmentation metric.
|
||||
semantic_on=False,
|
||||
instance_on=True,
|
||||
# max_per_image is for instance segmentation.
|
||||
max_per_image=100,
|
||||
iou_thr=0.8,
|
||||
# In Mask2Former's panoptic postprocessing,
|
||||
# it will filter mask area where score is less than 0.5 .
|
||||
filter_low_score=True),
|
||||
init_cfg=None)
|
||||
|
||||
# find_unused_parameters = True
|
||||
34
segmentation/configs/_base_/models/segformer_mit-b0.py
Normal file
34
segmentation/configs/_base_/models/segformer_mit-b0.py
Normal file
@@ -0,0 +1,34 @@
|
||||
# model settings
|
||||
norm_cfg = dict(type='SyncBN', requires_grad=True)
|
||||
model = dict(
|
||||
type='EncoderDecoder',
|
||||
pretrained=None,
|
||||
backbone=dict(
|
||||
type='MixVisionTransformer',
|
||||
in_channels=3,
|
||||
embed_dims=32,
|
||||
num_stages=4,
|
||||
num_layers=[2, 2, 2, 2],
|
||||
num_heads=[1, 2, 5, 8],
|
||||
patch_sizes=[7, 3, 3, 3],
|
||||
sr_ratios=[8, 4, 2, 1],
|
||||
out_indices=(0, 1, 2, 3),
|
||||
mlp_ratio=4,
|
||||
qkv_bias=True,
|
||||
drop_rate=0.0,
|
||||
attn_drop_rate=0.0,
|
||||
drop_path_rate=0.1),
|
||||
decode_head=dict(
|
||||
type='SegformerHead',
|
||||
in_channels=[32, 64, 160, 256],
|
||||
in_index=[0, 1, 2, 3],
|
||||
channels=256,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
|
||||
# model training and testing settings
|
||||
train_cfg=dict(),
|
||||
test_cfg=dict(mode='whole'))
|
||||
46
segmentation/configs/_base_/models/upernet_convnext.py
Normal file
46
segmentation/configs/_base_/models/upernet_convnext.py
Normal file
@@ -0,0 +1,46 @@
|
||||
norm_cfg = dict(type='SyncBN', requires_grad=True)
|
||||
custom_imports = dict(imports='mmcls.models', allow_failed_imports=False)
|
||||
# checkpoint_file = 'https://download.openmmlab.com/mmclassification/v0/convnext/downstream/convnext-base_3rdparty_32xb128-noema_in1k_20220301-2a0ee547.pth' # noqa
|
||||
model = dict(
|
||||
type='EncoderDecoder',
|
||||
pretrained=None,
|
||||
backbone=dict(
|
||||
type='mmcls.ConvNeXt',
|
||||
arch='base',
|
||||
norm_cfg=dict(type='LN2dv2', eps=1e-6),
|
||||
out_indices=[0, 1, 2, 3],
|
||||
drop_path_rate=0.4,
|
||||
layer_scale_init_value=1.0,
|
||||
gap_before_final_norm=False,
|
||||
# init_cfg=dict(
|
||||
# type='Pretrained', checkpoint=checkpoint_file,
|
||||
# prefix='backbone.')
|
||||
),
|
||||
decode_head=dict(
|
||||
type='UPerHead',
|
||||
in_channels=[128, 256, 512, 1024],
|
||||
in_index=[0, 1, 2, 3],
|
||||
pool_scales=(1, 2, 3, 6),
|
||||
channels=512,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
|
||||
auxiliary_head=dict(
|
||||
type='FCNHead',
|
||||
in_channels=384,
|
||||
in_index=2,
|
||||
channels=256,
|
||||
num_convs=1,
|
||||
concat_input=False,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
|
||||
# model training and testing settings
|
||||
train_cfg=dict(),
|
||||
test_cfg=dict(mode='whole'))
|
||||
44
segmentation/configs/_base_/models/upernet_r50.py
Normal file
44
segmentation/configs/_base_/models/upernet_r50.py
Normal file
@@ -0,0 +1,44 @@
|
||||
# model settings
|
||||
norm_cfg = dict(type='SyncBN', requires_grad=True)
|
||||
model = dict(
|
||||
type='EncoderDecoder',
|
||||
pretrained='open-mmlab://resnet50_v1c',
|
||||
backbone=dict(
|
||||
type='ResNetV1c',
|
||||
depth=50,
|
||||
num_stages=4,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
dilations=(1, 1, 1, 1),
|
||||
strides=(1, 2, 2, 2),
|
||||
norm_cfg=norm_cfg,
|
||||
norm_eval=False,
|
||||
style='pytorch',
|
||||
contract_dilation=True),
|
||||
decode_head=dict(
|
||||
type='UPerHead',
|
||||
in_channels=[256, 512, 1024, 2048],
|
||||
in_index=[0, 1, 2, 3],
|
||||
pool_scales=(1, 2, 3, 6),
|
||||
channels=512,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
|
||||
auxiliary_head=dict(
|
||||
type='FCNHead',
|
||||
in_channels=1024,
|
||||
in_index=2,
|
||||
channels=256,
|
||||
num_convs=1,
|
||||
concat_input=False,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
|
||||
# model training and testing settings
|
||||
train_cfg=dict(),
|
||||
test_cfg=dict(mode='whole'))
|
||||
54
segmentation/configs/_base_/models/upernet_swin.py
Normal file
54
segmentation/configs/_base_/models/upernet_swin.py
Normal file
@@ -0,0 +1,54 @@
|
||||
# model settings
|
||||
norm_cfg = dict(type='SyncBN', requires_grad=True)
|
||||
backbone_norm_cfg = dict(type='LN', requires_grad=True)
|
||||
model = dict(
|
||||
type='EncoderDecoder',
|
||||
pretrained=None,
|
||||
backbone=dict(
|
||||
type='SwinTransformer',
|
||||
pretrain_img_size=224,
|
||||
embed_dims=96,
|
||||
patch_size=4,
|
||||
window_size=7,
|
||||
mlp_ratio=4,
|
||||
depths=[2, 2, 6, 2],
|
||||
num_heads=[3, 6, 12, 24],
|
||||
strides=(4, 2, 2, 2),
|
||||
out_indices=(0, 1, 2, 3),
|
||||
qkv_bias=True,
|
||||
qk_scale=None,
|
||||
patch_norm=True,
|
||||
drop_rate=0.,
|
||||
attn_drop_rate=0.,
|
||||
drop_path_rate=0.3,
|
||||
use_abs_pos_embed=False,
|
||||
act_cfg=dict(type='GELU'),
|
||||
norm_cfg=backbone_norm_cfg),
|
||||
decode_head=dict(
|
||||
type='UPerHead',
|
||||
in_channels=[96, 192, 384, 768],
|
||||
in_index=[0, 1, 2, 3],
|
||||
pool_scales=(1, 2, 3, 6),
|
||||
channels=512,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
|
||||
auxiliary_head=dict(
|
||||
type='FCNHead',
|
||||
in_channels=384,
|
||||
in_index=2,
|
||||
channels=256,
|
||||
num_convs=1,
|
||||
concat_input=False,
|
||||
dropout_ratio=0.1,
|
||||
num_classes=19,
|
||||
norm_cfg=norm_cfg,
|
||||
align_corners=False,
|
||||
loss_decode=dict(
|
||||
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
|
||||
# model training and testing settings
|
||||
train_cfg=dict(),
|
||||
test_cfg=dict(mode='whole'))
|
||||
9
segmentation/configs/_base_/schedules/schedule_160k.py
Normal file
9
segmentation/configs/_base_/schedules/schedule_160k.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# optimizer
|
||||
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer_config = dict()
|
||||
# learning policy
|
||||
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
|
||||
# runtime settings
|
||||
runner = dict(type='IterBasedRunner', max_iters=160000)
|
||||
checkpoint_config = dict(by_epoch=False, interval=16000)
|
||||
evaluation = dict(interval=16000, metric='mIoU', pre_eval=True)
|
||||
9
segmentation/configs/_base_/schedules/schedule_20k.py
Normal file
9
segmentation/configs/_base_/schedules/schedule_20k.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# optimizer
|
||||
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer_config = dict()
|
||||
# learning policy
|
||||
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
|
||||
# runtime settings
|
||||
runner = dict(type='IterBasedRunner', max_iters=20000)
|
||||
checkpoint_config = dict(by_epoch=False, interval=2000)
|
||||
evaluation = dict(interval=2000, metric='mIoU', pre_eval=True)
|
||||
9
segmentation/configs/_base_/schedules/schedule_320k.py
Normal file
9
segmentation/configs/_base_/schedules/schedule_320k.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# optimizer
|
||||
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer_config = dict()
|
||||
# learning policy
|
||||
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
|
||||
# runtime settings
|
||||
runner = dict(type='IterBasedRunner', max_iters=320000)
|
||||
checkpoint_config = dict(by_epoch=False, interval=32000)
|
||||
evaluation = dict(interval=32000, metric='mIoU')
|
||||
9
segmentation/configs/_base_/schedules/schedule_40k.py
Normal file
9
segmentation/configs/_base_/schedules/schedule_40k.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# optimizer
|
||||
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer_config = dict()
|
||||
# learning policy
|
||||
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
|
||||
# runtime settings
|
||||
runner = dict(type='IterBasedRunner', max_iters=40000)
|
||||
checkpoint_config = dict(by_epoch=False, interval=4000)
|
||||
evaluation = dict(interval=4000, metric='mIoU', pre_eval=True)
|
||||
9
segmentation/configs/_base_/schedules/schedule_80k.py
Normal file
9
segmentation/configs/_base_/schedules/schedule_80k.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# optimizer
|
||||
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer_config = dict()
|
||||
# learning policy
|
||||
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
|
||||
# runtime settings
|
||||
runner = dict(type='IterBasedRunner', max_iters=80000)
|
||||
checkpoint_config = dict(by_epoch=False, interval=8000)
|
||||
evaluation = dict(interval=8000, metric='mIoU', pre_eval=True)
|
||||
33
segmentation/configs/ade20k/README.md
Normal file
33
segmentation/configs/ade20k/README.md
Normal file
@@ -0,0 +1,33 @@
|
||||
# ADE20K
|
||||
|
||||
Introduced by Zhou et al. in [Scene Parsing Through ADE20K Dataset](https://paperswithcode.com/paper/scene-parsing-through-ade20k-dataset).
|
||||
|
||||
The ADE20K semantic segmentation dataset contains more than 20K scene-centric images exhaustively annotated with pixel-level objects and object parts labels. There are totally 150 semantic categories, which include stuffs like sky, road, grass, and discrete objects like person, car, bed.
|
||||
|
||||
|
||||
## Model Zoo
|
||||
|
||||
### UperNet + InternImage
|
||||
|
||||
|
||||
| backbone | resolution | mIoU (ss/ms) | Config | Download |
|
||||
|:--------------:|:----------:|:-----------:|:-----------:|:----------:
|
||||
| FlashInternImage-T | 512x512 | 49.3 / 50.3 | [config](./upernet_flash_internimage_t_512_160k_ade20k.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_t_512_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_t_512_160k_ade20k.log) |
|
||||
| FlashInternImage-S | 512x512 | 50.6 / 51.6 | [config](./upernet_flash_internimage_s_512_160k_ade20k.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_s_512_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_s_512_160k_ade20k.log) |
|
||||
| FlashInternImage-B | 512x512 | 52.0 / 52.6 | [config](./upernet_flash_internimage_b_512_160k_ade20k.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_b_512_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_s_512_160k_ade20k.log) |
|
||||
| FlashInternImage-L | 640x640 | 55.6 / 56.0 | [config](./upernet_flash_internimage_l_640_160k_ade20k.py)| [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_l_640_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_l_640_160k_ade20k.log) |
|
||||
|
||||
- Training speed is measured with A100 GPU.
|
||||
- Please set `with_cp=True` to save memory if you meet `out-of-memory` issues.
|
||||
- The logs are our recent newly trained ones. There are slight differences between the results in logs and our paper.
|
||||
|
||||
|
||||
### Mask2Former + InternImage
|
||||
|
||||
| backbone | resolution | mIoU (ss) | Config | Download |
|
||||
|:--------------:|:----------:|:-----------:|:-----------:|:----------:
|
||||
| FlashInternImage-T | 512x512 | 51.2 | [config](./mask2former_flash_internimage_t_512_160k_ade20k_ss.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_t_512_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_t_512_160k_ade20k_ss.log) |
|
||||
| FlashInternImage-S | 640x640 | 52.2 | [config](./mask2former_flash_internimage_s_640_160k_ade20k_ss.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_s_640_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_s_640_160k_ade20k_ss.log) |
|
||||
| FlashInternImage-B | 640x640 | 53.4 | [config](./mask2former_flash_internimage_b_640_160k_ade20k_ss.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_b_640_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_b_640_160k_ade20k_ss.log) |
|
||||
| FlashInternImage-L | 640x640 | 56.7 | [config](./mask2former_flash_internimage_l_640_160k_ade20k_ss.py)| [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_l_640_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_l_640_160k_ade20k_ss.log) |
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/mask2former_beit.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
num_classes = 150
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_b_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=112,
|
||||
depths=[4, 4, 21, 4],
|
||||
groups=[7, 14, 28, 56],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.4,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=0.5,
|
||||
post_norm=True,
|
||||
with_cp=False,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(
|
||||
in_channels=[112, 224, 448, 896],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
num_classes=num_classes,
|
||||
num_queries=200,
|
||||
pixel_decoder=dict(
|
||||
type='MSDeformAttnPixelDecoder',
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiScaleDeformableAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
ffn_cfgs=dict(
|
||||
type='FFN',
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
ffn_drop=0.0,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
act_cfg=dict(type='ReLU', inplace=True)),
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
transformer_decoder=dict(
|
||||
type='DetrTransformerDecoder',
|
||||
return_intermediate=True,
|
||||
num_layers=9,
|
||||
transformerlayers=dict(
|
||||
type='DetrTransformerDecoderLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiheadAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
attn_drop=0.0,
|
||||
proj_drop=0.0,
|
||||
dropout_layer=None,
|
||||
batch_first=False),
|
||||
ffn_cfgs=dict(
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
act_cfg=dict(type='ReLU', inplace=True),
|
||||
ffn_drop=0.0,
|
||||
dropout_layer=None,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
add_identity=True),
|
||||
feedforward_channels=2048,
|
||||
operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
|
||||
'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=False,
|
||||
loss_weight=2.0,
|
||||
reduction='mean',
|
||||
class_weight=[1.0] * num_classes + [0.1])
|
||||
),
|
||||
test_cfg=dict(mode='whole'))
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (640, 640)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2560, 640), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='ToMask'),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2560, 640),
|
||||
# img_ratios=[768./896., 832./896., 1.0, 960./896., 1024./896.],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=33, layer_decay_rate=1.0,
|
||||
depths=[4, 4, 21, 4]))
|
||||
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data = dict(samples_per_gpu=2,
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=0.1, norm_type=2))
|
||||
checkpoint_config = dict(by_epoch=False, interval=5000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=5000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,160 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/mask2former_beit.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
num_classes = 150
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_l_22k_384.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=160,
|
||||
depths=[5, 5, 22, 5],
|
||||
groups=[10, 20, 40, 80],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.5,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=2.0,
|
||||
post_norm=True,
|
||||
with_cp=True,
|
||||
dcn_output_bias=True,
|
||||
mlp_fc2_bias=True,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(
|
||||
in_channels=[160, 320, 640, 1280],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
num_classes=num_classes,
|
||||
num_queries=200,
|
||||
pixel_decoder=dict(
|
||||
type='MSDeformAttnPixelDecoder',
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiScaleDeformableAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
ffn_cfgs=dict(
|
||||
type='FFN',
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
ffn_drop=0.0,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
act_cfg=dict(type='ReLU', inplace=True)),
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
transformer_decoder=dict(
|
||||
type='DetrTransformerDecoder',
|
||||
return_intermediate=True,
|
||||
num_layers=9,
|
||||
transformerlayers=dict(
|
||||
type='DetrTransformerDecoderLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiheadAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
attn_drop=0.0,
|
||||
proj_drop=0.0,
|
||||
dropout_layer=None,
|
||||
batch_first=False),
|
||||
ffn_cfgs=dict(
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
act_cfg=dict(type='ReLU', inplace=True),
|
||||
ffn_drop=0.0,
|
||||
dropout_layer=None,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
add_identity=True),
|
||||
feedforward_channels=2048,
|
||||
operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
|
||||
'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=False,
|
||||
loss_weight=2.0,
|
||||
reduction='mean',
|
||||
class_weight=[1.0] * num_classes + [0.1])
|
||||
),
|
||||
test_cfg=dict(mode='whole'))
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (640, 640)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2560, 640), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='ToMask'),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2560, 640),
|
||||
# img_ratios=[768./896., 832./896., 1.0, 960./896., 1024./896.],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=37, layer_decay_rate=0.94,
|
||||
depths=[5, 5, 22, 5], offset_lr_scale=1.0))
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data = dict(samples_per_gpu=2,
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=0.1, norm_type=2))
|
||||
checkpoint_config = dict(by_epoch=False, interval=2000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=2000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,159 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/mask2former_beit.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
num_classes = 150
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_s_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=80,
|
||||
depths=[4, 4, 21, 4],
|
||||
groups=[5, 10, 20, 40],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.3,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=1.0,
|
||||
post_norm=True,
|
||||
with_cp=False,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(
|
||||
in_channels=[80, 160, 320, 640],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
num_classes=num_classes,
|
||||
num_queries=200,
|
||||
pixel_decoder=dict(
|
||||
type='MSDeformAttnPixelDecoder',
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiScaleDeformableAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
ffn_cfgs=dict(
|
||||
type='FFN',
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
ffn_drop=0.0,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
act_cfg=dict(type='ReLU', inplace=True)),
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
transformer_decoder=dict(
|
||||
type='DetrTransformerDecoder',
|
||||
return_intermediate=True,
|
||||
num_layers=9,
|
||||
transformerlayers=dict(
|
||||
type='DetrTransformerDecoderLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiheadAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
attn_drop=0.0,
|
||||
proj_drop=0.0,
|
||||
dropout_layer=None,
|
||||
batch_first=False),
|
||||
ffn_cfgs=dict(
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
act_cfg=dict(type='ReLU', inplace=True),
|
||||
ffn_drop=0.0,
|
||||
dropout_layer=None,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
add_identity=True),
|
||||
feedforward_channels=2048,
|
||||
operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
|
||||
'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=False,
|
||||
loss_weight=2.0,
|
||||
reduction='mean',
|
||||
class_weight=[1.0] * num_classes + [0.1])
|
||||
),
|
||||
test_cfg=dict(mode='whole'))
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (640, 640)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2560, 640), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='ToMask'),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2560, 640),
|
||||
# img_ratios=[768./896., 832./896., 1.0, 960./896., 1024./896.],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=33, layer_decay_rate=1.0,
|
||||
depths=[4, 4, 21, 4]))
|
||||
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data = dict(samples_per_gpu=2,
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=0.1, norm_type=2))
|
||||
checkpoint_config = dict(by_epoch=False, interval=5000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=5000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,160 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/mask2former_beit.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
num_classes = 150
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_s_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=80,
|
||||
depths=[4, 4, 21, 4],
|
||||
groups=[5, 10, 20, 40],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.3,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=1.0,
|
||||
post_norm=True,
|
||||
with_cp=False,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(
|
||||
in_channels=[80, 160, 320, 640],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
num_classes=num_classes,
|
||||
num_queries=200,
|
||||
pixel_decoder=dict(
|
||||
type='MSDeformAttnPixelDecoder',
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='CustomMultiScaleDeformableAttention',
|
||||
use_softmax=False,
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
ffn_cfgs=dict(
|
||||
type='FFN',
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
ffn_drop=0.0,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
act_cfg=dict(type='ReLU', inplace=True)),
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
transformer_decoder=dict(
|
||||
type='DetrTransformerDecoder',
|
||||
return_intermediate=True,
|
||||
num_layers=9,
|
||||
transformerlayers=dict(
|
||||
type='DetrTransformerDecoderLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiheadAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
attn_drop=0.0,
|
||||
proj_drop=0.0,
|
||||
dropout_layer=None,
|
||||
batch_first=False),
|
||||
ffn_cfgs=dict(
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
act_cfg=dict(type='ReLU', inplace=True),
|
||||
ffn_drop=0.0,
|
||||
dropout_layer=None,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
add_identity=True),
|
||||
feedforward_channels=2048,
|
||||
operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
|
||||
'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=False,
|
||||
loss_weight=2.0,
|
||||
reduction='mean',
|
||||
class_weight=[1.0] * num_classes + [0.1])
|
||||
),
|
||||
test_cfg=dict(mode='whole'))
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (640, 640)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2560, 640), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='ToMask'),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2560, 640),
|
||||
# img_ratios=[768./896., 832./896., 1.0, 960./896., 1024./896.],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=33, layer_decay_rate=1.0,
|
||||
depths=[4, 4, 21, 4]))
|
||||
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data = dict(samples_per_gpu=2,
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=0.1, norm_type=2))
|
||||
checkpoint_config = dict(by_epoch=False, interval=5000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=5000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,157 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2023 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/mask2former_beit.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
num_classes = 150
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_t_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=64,
|
||||
depths=[4, 4, 18, 4],
|
||||
groups=[4, 8, 16, 32],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.2,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=1.0,
|
||||
post_norm=False,
|
||||
with_cp=False,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(
|
||||
in_channels=[64, 128, 256, 512],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
num_classes=num_classes,
|
||||
num_queries=100,
|
||||
pixel_decoder=dict(
|
||||
type='MSDeformAttnPixelDecoder',
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiScaleDeformableAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
ffn_cfgs=dict(
|
||||
type='FFN',
|
||||
embed_dims=256,
|
||||
feedforward_channels=1024,
|
||||
num_fcs=2,
|
||||
ffn_drop=0.0,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
act_cfg=dict(type='ReLU', inplace=True)),
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding', num_feats=128, normalize=True),
|
||||
transformer_decoder=dict(
|
||||
type='DetrTransformerDecoder',
|
||||
return_intermediate=True,
|
||||
num_layers=9,
|
||||
transformerlayers=dict(
|
||||
type='DetrTransformerDecoderLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiheadAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
attn_drop=0.0,
|
||||
proj_drop=0.0,
|
||||
dropout_layer=None,
|
||||
batch_first=False),
|
||||
ffn_cfgs=dict(
|
||||
embed_dims=256,
|
||||
feedforward_channels=2048,
|
||||
num_fcs=2,
|
||||
act_cfg=dict(type='ReLU', inplace=True),
|
||||
ffn_drop=0.0,
|
||||
dropout_layer=None,
|
||||
with_cp=False, # set with_cp=True to save memory
|
||||
add_identity=True),
|
||||
feedforward_channels=2048,
|
||||
operation_order=('cross_attn', 'norm', 'self_attn', 'norm',
|
||||
'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
use_sigmoid=False,
|
||||
loss_weight=2.0,
|
||||
reduction='mean',
|
||||
class_weight=[1.0] * num_classes + [0.1])
|
||||
),
|
||||
test_cfg=dict(mode='whole'))
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (512, 512)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='ToMask'),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'])
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[768./896., 832./896., 1.0, 960./896., 1024./896.],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=30, layer_decay_rate=0.9,
|
||||
depths=[4, 4, 18, 4]))
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data = dict(samples_per_gpu=2,
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=0.01, norm_type=2))
|
||||
checkpoint_config = dict(by_epoch=False, interval=5000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=5000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,68 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/upernet_r50.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_b_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=112,
|
||||
depths=[4, 4, 21, 4],
|
||||
groups=[7, 14, 28, 56],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.3,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=0.5,
|
||||
post_norm=True,
|
||||
with_cp=False,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(num_classes=150, in_channels=[112, 224, 448, 896]),
|
||||
auxiliary_head=dict(num_classes=150, in_channels=448),
|
||||
test_cfg=dict(mode='whole')
|
||||
)
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=33, layer_decay_rate=1.0,
|
||||
depths=[4, 4, 21, 4]))
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data=dict(samples_per_gpu=2,
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=16000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,84 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/upernet_r50.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_l_22k_384.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=160,
|
||||
depths=[5, 5, 22, 5],
|
||||
groups=[10, 20, 40, 80],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.4,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=2.0,
|
||||
post_norm=True,
|
||||
with_cp=False,
|
||||
dcn_output_bias=True,
|
||||
mlp_fc2_bias=True,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(num_classes=150, in_channels=[160, 320, 640, 1280]),
|
||||
auxiliary_head=dict(num_classes=150, in_channels=640),
|
||||
test_cfg=dict(mode='whole'))
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
crop_size = (640, 640)
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(type='LoadAnnotations', reduce_zero_label=True),
|
||||
dict(type='Resize', img_scale=(2560, 640), ratio_range=(0.5, 2.0)),
|
||||
dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
|
||||
dict(type='RandomFlip', prob=0.5),
|
||||
dict(type='PhotoMetricDistortion'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
|
||||
dict(type='DefaultFormatBundle'),
|
||||
dict(type='Collect', keys=['img', 'gt_semantic_seg']),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2560, 640),
|
||||
# img_ratios=[0.75, 1.0, 1.25],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.00002, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=37, layer_decay_rate=0.94,
|
||||
depths=[5, 5, 22, 5], offset_lr_scale=1.0))
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data = dict(samples_per_gpu=2,
|
||||
train=dict(pipeline=train_pipeline),
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=0.1, norm_type=2))
|
||||
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=16000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,68 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/upernet_r50.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_s_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=80,
|
||||
depths=[4, 4, 21, 4],
|
||||
groups=[5, 10, 20, 40],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.3,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=1.0,
|
||||
post_norm=True,
|
||||
with_cp=True,
|
||||
dw_kernel_size=3,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(num_classes=150, in_channels=[80, 160, 320, 640]),
|
||||
auxiliary_head=dict(num_classes=150, in_channels=320),
|
||||
test_cfg=dict(mode='whole')
|
||||
)
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=33, layer_decay_rate=1.0,
|
||||
depths=[4, 4, 21, 4]))
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data=dict(samples_per_gpu=2,
|
||||
val=dict(pipeline=test_pipeline),
|
||||
test=dict(pipeline=test_pipeline))
|
||||
runner = dict(type='IterBasedRunner')
|
||||
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=16000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
@@ -0,0 +1,68 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
_base_ = [
|
||||
'../_base_/models/upernet_r50.py', '../_base_/datasets/ade20k.py',
|
||||
'../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
|
||||
]
|
||||
pretrained = 'https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_t_1k_224.pth'
|
||||
model = dict(
|
||||
backbone=dict(
|
||||
_delete_=True,
|
||||
type='FlashInternImage',
|
||||
core_op='DCNv4',
|
||||
channels=64,
|
||||
depths=[4, 4, 18, 4],
|
||||
groups=[4, 8, 16, 32],
|
||||
mlp_ratio=4.,
|
||||
drop_path_rate=0.2,
|
||||
norm_layer='LN',
|
||||
layer_scale=1.0,
|
||||
offset_scale=1.0,
|
||||
post_norm=False,
|
||||
with_cp=True,
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
|
||||
decode_head=dict(num_classes=150, in_channels=[64, 128, 256, 512]),
|
||||
auxiliary_head=dict(num_classes=150, in_channels=256),
|
||||
test_cfg=dict(mode='whole')
|
||||
)
|
||||
img_norm_cfg = dict(
|
||||
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiScaleFlipAug',
|
||||
img_scale=(2048, 512),
|
||||
# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
|
||||
flip=False,
|
||||
transforms=[
|
||||
dict(type='Resize', keep_ratio=True),
|
||||
dict(type='ResizeToMultiple', size_divisor=32),
|
||||
dict(type='RandomFlip'),
|
||||
dict(type='Normalize', **img_norm_cfg),
|
||||
dict(type='ImageToTensor', keys=['img']),
|
||||
dict(type='Collect', keys=['img']),
|
||||
])
|
||||
]
|
||||
optimizer = dict(
|
||||
_delete_=True, type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.05,
|
||||
constructor='CustomLayerDecayOptimizerConstructor',
|
||||
paramwise_cfg=dict(num_layers=30, layer_decay_rate=1.0,
|
||||
depths=[4, 4, 18, 4]))
|
||||
lr_config = dict(_delete_=True, policy='poly',
|
||||
warmup='linear',
|
||||
warmup_iters=1500,
|
||||
warmup_ratio=1e-6,
|
||||
power=1.0, min_lr=0.0, by_epoch=False)
|
||||
# By default, models are trained on 8 GPUs with 2 images per GPU
|
||||
data=dict(samples_per_gpu=2,
|
||||
# val=dict(pipeline=test_pipeline),
|
||||
# test=dict(pipeline=test_pipeline)
|
||||
)
|
||||
runner = dict(type='IterBasedRunner')
|
||||
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=1)
|
||||
evaluation = dict(interval=16000, metric='mIoU', save_best='mIoU')
|
||||
# fp16 = dict(loss_scale=dict(init_scale=512))
|
||||
2
segmentation/deploy/configs/_base_/backends/tensorrt.py
Normal file
2
segmentation/deploy/configs/_base_/backends/tensorrt.py
Normal file
@@ -0,0 +1,2 @@
|
||||
backend_config = dict(
|
||||
type='tensorrt', common_config=dict(fp16_mode=False, max_workspace_size=0))
|
||||
10
segmentation/deploy/configs/_base_/onnx_config.py
Normal file
10
segmentation/deploy/configs/_base_/onnx_config.py
Normal file
@@ -0,0 +1,10 @@
|
||||
onnx_config = dict(
|
||||
type='onnx',
|
||||
export_params=True,
|
||||
keep_initializers_as_inputs=False,
|
||||
opset_version=11,
|
||||
save_file='end2end.onnx',
|
||||
input_names=['input'],
|
||||
output_names=['output'],
|
||||
input_shape=None,
|
||||
optimize=True)
|
||||
2
segmentation/deploy/configs/mmseg/segmentation_static.py
Normal file
2
segmentation/deploy/configs/mmseg/segmentation_static.py
Normal file
@@ -0,0 +1,2 @@
|
||||
_base_ = ['../_base_/onnx_config.py']
|
||||
codebase_config = dict(type='mmseg', task='Segmentation', with_argmax=True)
|
||||
@@ -0,0 +1,13 @@
|
||||
_base_ = ['./segmentation_static.py', '../_base_/backends/tensorrt.py']
|
||||
|
||||
onnx_config = dict(input_shape=[512, 512])
|
||||
backend_config = dict(
|
||||
common_config=dict(max_workspace_size=1 << 30),
|
||||
model_inputs=[
|
||||
dict(
|
||||
input_shapes=dict(
|
||||
input=dict(
|
||||
min_shape=[1, 3, 512, 512],
|
||||
opt_shape=[1, 3, 512, 512],
|
||||
max_shape=[1, 3, 512, 512])))
|
||||
])
|
||||
BIN
segmentation/deploy/demo.png
Normal file
BIN
segmentation/deploy/demo.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 301 KiB |
9
segmentation/dist_test.sh
Executable file
9
segmentation/dist_test.sh
Executable file
@@ -0,0 +1,9 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
CONFIG=$1
|
||||
CHECKPOINT=$2
|
||||
GPUS=$3
|
||||
PORT=${PORT:-29510}
|
||||
PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
|
||||
python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \
|
||||
$(dirname "$0")/test.py $CONFIG $CHECKPOINT --launcher pytorch ${@:4}
|
||||
9
segmentation/dist_train.sh
Executable file
9
segmentation/dist_train.sh
Executable file
@@ -0,0 +1,9 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
CONFIG=$1
|
||||
GPUS=$2
|
||||
PORT=${PORT:-29300}
|
||||
|
||||
PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
|
||||
python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \
|
||||
$(dirname "$0")/train.py $CONFIG --launcher pytorch ${@:3}
|
||||
119
segmentation/get_flops.py
Normal file
119
segmentation/get_flops.py
Normal file
@@ -0,0 +1,119 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from mmcv import Config, DictAction
|
||||
|
||||
from mmseg.models import build_segmentor
|
||||
import mmcv_custom # noqa: F401,F403
|
||||
import mmseg_custom # noqa: F401,F403
|
||||
|
||||
try:
|
||||
from mmcv.cnn.utils.flops_counter import flops_to_string, params_to_string
|
||||
from mmcv.cnn import get_model_complexity_info
|
||||
except ImportError:
|
||||
raise ImportError('Please upgrade mmcv to >0.6.2')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Train a detector')
|
||||
parser.add_argument('config', help='train config file path')
|
||||
parser.add_argument(
|
||||
'--shape',
|
||||
type=int,
|
||||
nargs='+',
|
||||
default=[512, 2048],
|
||||
help='input image size')
|
||||
parser.add_argument(
|
||||
'--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 '
|
||||
'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.')
|
||||
parser.add_argument(
|
||||
'--size-divisor',
|
||||
type=int,
|
||||
default=32,
|
||||
help='Pad the input image, the minimum size that is divisible '
|
||||
'by size_divisor, -1 means do not pad the image.')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
# dcnv3_flops(n=h*w, k=3*3, c=channels)
|
||||
def dcnv3_flops(n, k, c):
|
||||
return 5 * n * k * c
|
||||
|
||||
def get_flops(model, input_shape):
|
||||
flops, params = get_model_complexity_info(model, input_shape, as_strings=False)
|
||||
|
||||
backbone = model.backbone
|
||||
backbone_name = type(backbone).__name__
|
||||
_, H, W = input_shape
|
||||
|
||||
temp = 0
|
||||
if 'Intern' in backbone_name:
|
||||
depths = backbone.depths # [4, 4, 18, 4]
|
||||
for idx, depth in enumerate(depths):
|
||||
channels = backbone.channels * (2 ** idx)
|
||||
h = H / (4 * (2 ** idx))
|
||||
w = W / (4 * (2 ** idx))
|
||||
temp += depth * dcnv3_flops(n=h*w, k=3*3, c=channels)
|
||||
|
||||
flops = flops + temp
|
||||
return flops_to_string(flops), params_to_string(params)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parse_args()
|
||||
|
||||
if len(args.shape) == 1:
|
||||
h = w = args.shape[0]
|
||||
elif len(args.shape) == 2:
|
||||
h, w = args.shape
|
||||
else:
|
||||
raise ValueError('invalid input shape')
|
||||
orig_shape = (3, h, w)
|
||||
divisor = args.size_divisor
|
||||
if divisor > 0:
|
||||
h = int(np.ceil(h / divisor)) * divisor
|
||||
w = int(np.ceil(w / divisor)) * divisor
|
||||
|
||||
input_shape = (3, h, w)
|
||||
|
||||
cfg = Config.fromfile(args.config)
|
||||
if args.cfg_options is not None:
|
||||
cfg.merge_from_dict(args.cfg_options)
|
||||
|
||||
model = build_segmentor(
|
||||
cfg.model,
|
||||
train_cfg=cfg.get('train_cfg'),
|
||||
test_cfg=cfg.get('test_cfg'))
|
||||
|
||||
if torch.cuda.is_available():
|
||||
model.cuda()
|
||||
model.eval()
|
||||
if hasattr(model, 'forward_dummy'):
|
||||
model.forward = model.forward_dummy
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
'FLOPs counter is currently not currently supported with {}'.
|
||||
format(model.__class__.__name__))
|
||||
|
||||
flops, params = get_flops(model, input_shape)
|
||||
split_line = '=' * 30
|
||||
|
||||
if divisor > 0 and \
|
||||
input_shape != orig_shape:
|
||||
print(f'{split_line}\nUse size divisor set input shape '
|
||||
f'from {orig_shape} to {input_shape}\n')
|
||||
print(f'{split_line}\nInput shape: {input_shape}\n'
|
||||
f'Flops: {flops}\nParams: {params}\n{split_line}')
|
||||
print('!!!Please be cautious if you use the results in papers. '
|
||||
'You may need to check if all ops are supported and verify that the '
|
||||
'flops computation is correct.')
|
||||
70
segmentation/image_demo.py
Normal file
70
segmentation/image_demo.py
Normal file
@@ -0,0 +1,70 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from argparse import ArgumentParser
|
||||
|
||||
import mmcv
|
||||
|
||||
import mmcv_custom # noqa: F401,F403
|
||||
import mmseg_custom # noqa: F401,F403
|
||||
from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
|
||||
from mmseg.core.evaluation import get_palette
|
||||
from mmcv.runner import load_checkpoint
|
||||
from mmseg.core import get_classes
|
||||
import cv2
|
||||
import os.path as osp
|
||||
import os
|
||||
|
||||
|
||||
def test_single_image(model, img_name, out_dir, color_palette, opacity):
|
||||
result = inference_segmentor(model, img_name)
|
||||
|
||||
# show the results
|
||||
if hasattr(model, 'module'):
|
||||
model = model.module
|
||||
img = model.show_result(img_name, result,
|
||||
palette=color_palette,
|
||||
show=False, opacity=opacity)
|
||||
|
||||
# save the results
|
||||
mmcv.mkdir_or_exist(out_dir)
|
||||
out_path = osp.join(out_dir, osp.basename(img_name))
|
||||
cv2.imwrite(out_path, img)
|
||||
print(f"Result is save at {out_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument('img', help='Image file or a directory contains images')
|
||||
parser.add_argument('config', help='Config file')
|
||||
parser.add_argument('checkpoint', help='Checkpoint file')
|
||||
parser.add_argument('--out', type=str, default="demo", help='out dir')
|
||||
parser.add_argument(
|
||||
'--device', default='cuda:0', help='Device used for inference')
|
||||
parser.add_argument(
|
||||
'--palette',
|
||||
default='ade20k',
|
||||
choices=['ade20k', 'cityscapes', 'cocostuff'],
|
||||
help='Color palette used for segmentation map')
|
||||
parser.add_argument(
|
||||
'--opacity',
|
||||
type=float,
|
||||
default=0.5,
|
||||
help='Opacity of painted segmentation map. In (0, 1] range.')
|
||||
args = parser.parse_args()
|
||||
|
||||
# build the model from a config file and a checkpoint file
|
||||
model = init_segmentor(args.config, checkpoint=None, device=args.device)
|
||||
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
|
||||
if 'CLASSES' in checkpoint.get('meta', {}):
|
||||
model.CLASSES = checkpoint['meta']['CLASSES']
|
||||
else:
|
||||
model.CLASSES = get_classes(args.palette)
|
||||
|
||||
# check arg.img is directory of a single image.
|
||||
if osp.isdir(args.img):
|
||||
for img in os.listdir(args.img):
|
||||
test_single_image(model, osp.join(args.img, img), args.out, get_palette(args.palette), args.opacity)
|
||||
else:
|
||||
test_single_image(model, args.img, args.out, get_palette(args.palette), args.opacity)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
11
segmentation/mmcv_custom/__init__.py
Normal file
11
segmentation/mmcv_custom/__init__.py
Normal file
@@ -0,0 +1,11 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
# -*- coding: utf-8 -*-
|
||||
from .custom_layer_decay_optimizer_constructor import CustomLayerDecayOptimizerConstructor
|
||||
from .layer_decay import LearningRateDecayOptimizerConstructor
|
||||
from .layer_decay_vit import LayerDecayOptimizerConstructor_vit
|
||||
__all__ = ['CustomLayerDecayOptimizerConstructor',]
|
||||
@@ -0,0 +1,157 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
"""
|
||||
Mostly copy-paste from BEiT library:
|
||||
https://github.com/microsoft/unilm/blob/master/beit/semantic_segmentation/mmcv_custom/layer_decay_optimizer_constructor.py
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from mmcv.runner import OPTIMIZER_BUILDERS, DefaultOptimizerConstructor
|
||||
from mmcv.runner import get_dist_info
|
||||
from mmseg.utils import get_root_logger
|
||||
|
||||
|
||||
def get_num_layer_for_swin(var_name, num_max_layer, depths):
|
||||
if var_name.startswith("backbone.patch_embed"):
|
||||
return 0
|
||||
elif var_name.startswith('decode_head.mask_embed'):
|
||||
return 0
|
||||
elif var_name.startswith('decode_head.cls_embed'):
|
||||
return 0
|
||||
elif var_name.startswith('decode_head.level_embed'):
|
||||
return 0
|
||||
elif var_name.startswith('decode_head.query_embed'):
|
||||
return 0
|
||||
elif var_name.startswith('decode_head.query_feat'):
|
||||
return 0
|
||||
if var_name.startswith("backbone.cb_modules.0.patch_embed"):
|
||||
return 0
|
||||
elif "level_embeds" in var_name:
|
||||
return 0
|
||||
elif var_name.startswith("backbone.layers") or var_name.startswith(
|
||||
"backbone.levels"):
|
||||
if var_name.split('.')[3] not in ['downsample', 'norm']:
|
||||
stage_id = int(var_name.split('.')[2])
|
||||
layer_id = int(var_name.split('.')[4])
|
||||
# layers for Swin-Large: [2, 2, 18, 2]
|
||||
if stage_id == 0:
|
||||
return layer_id + 1
|
||||
elif stage_id == 1:
|
||||
return layer_id + 1 + depths[0]
|
||||
elif stage_id == 2:
|
||||
return layer_id + 1 + depths[0] + depths[1]
|
||||
else:
|
||||
return layer_id + 1 + depths[0] + depths[1] + depths[2]
|
||||
else:
|
||||
stage_id = int(var_name.split('.')[2])
|
||||
if stage_id == 0:
|
||||
return 1 + depths[0]
|
||||
elif stage_id == 1:
|
||||
return 1 + depths[0] + depths[1]
|
||||
elif stage_id == 2:
|
||||
return 1 + depths[0] + depths[1] + depths[2]
|
||||
else:
|
||||
return 1 + depths[0] + depths[1] + depths[2]
|
||||
else:
|
||||
return num_max_layer - 1
|
||||
|
||||
|
||||
@OPTIMIZER_BUILDERS.register_module()
|
||||
class CustomLayerDecayOptimizerConstructor(DefaultOptimizerConstructor):
|
||||
|
||||
def add_params(self, params, module, prefix='', is_dcn_module=None):
|
||||
"""Add all parameters of module to the params list.
|
||||
The parameters of the given module will be added to the list of param
|
||||
groups, with specific rules defined by paramwise_cfg.
|
||||
Args:
|
||||
params (list[dict]): A list of param groups, it will be modified
|
||||
in place.
|
||||
module (nn.Module): The module to be added.
|
||||
prefix (str): The prefix of the module
|
||||
is_dcn_module (int|float|None): If the current module is a
|
||||
submodule of DCN, `is_dcn_module` will be passed to
|
||||
control conv_offset layer's learning rate. Defaults to None.
|
||||
"""
|
||||
parameter_groups = {}
|
||||
logger = get_root_logger()
|
||||
logger.info(self.paramwise_cfg)
|
||||
backbone_small_lr = self.paramwise_cfg.get('backbone_small_lr', False)
|
||||
dino_head = self.paramwise_cfg.get('dino_head', False)
|
||||
num_layers = self.paramwise_cfg.get('num_layers') + 2
|
||||
layer_decay_rate = self.paramwise_cfg.get('layer_decay_rate')
|
||||
depths = self.paramwise_cfg.get('depths')
|
||||
offset_lr_scale = self.paramwise_cfg.get('offset_lr_scale', 1.0)
|
||||
|
||||
logger.info("Build CustomLayerDecayOptimizerConstructor %f - %d" %
|
||||
(layer_decay_rate, num_layers))
|
||||
weight_decay = self.base_wd
|
||||
|
||||
for name, param in module.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue # frozen weights
|
||||
if len(param.shape) == 1 or name.endswith(".bias") or \
|
||||
"relative_position" in name or \
|
||||
"norm" in name or\
|
||||
"sampling_offsets" in name:
|
||||
group_name = "no_decay"
|
||||
this_weight_decay = 0.
|
||||
else:
|
||||
group_name = "decay"
|
||||
this_weight_decay = weight_decay
|
||||
|
||||
layer_id = get_num_layer_for_swin(name, num_layers, depths)
|
||||
if layer_id == num_layers - 1 and dino_head and \
|
||||
("sampling_offsets" in name or "reference_points" in name):
|
||||
group_name = "layer_%d_%s_0.1x" % (layer_id, group_name)
|
||||
elif ("sampling_offsets" in name or "reference_points" in name) and "backbone" in name:
|
||||
group_name = "layer_%d_%s_offset_lr_scale" % (layer_id,
|
||||
group_name)
|
||||
elif "offset_mask" in name and "offset_mask_dw" not in name:
|
||||
group_name = "layer_%d_%s_offset_lr_scale" % (layer_id,
|
||||
group_name)
|
||||
else:
|
||||
group_name = "layer_%d_%s" % (layer_id, group_name)
|
||||
|
||||
if group_name not in parameter_groups:
|
||||
scale = layer_decay_rate ** (num_layers - layer_id - 1)
|
||||
if scale < 1 and backbone_small_lr == True:
|
||||
scale = scale * 0.1
|
||||
if "0.1x" in group_name:
|
||||
scale = scale * 0.1
|
||||
if "offset_lr_scale" in group_name:
|
||||
scale = scale * offset_lr_scale
|
||||
|
||||
parameter_groups[group_name] = {
|
||||
"weight_decay": this_weight_decay,
|
||||
"params": [],
|
||||
"param_names": [],
|
||||
"lr_scale": scale,
|
||||
"group_name": group_name,
|
||||
"lr": scale * self.base_lr,
|
||||
}
|
||||
|
||||
parameter_groups[group_name]["params"].append(param)
|
||||
parameter_groups[group_name]["param_names"].append(name)
|
||||
rank, _ = get_dist_info()
|
||||
if rank == 0:
|
||||
to_display = {}
|
||||
for key in parameter_groups:
|
||||
to_display[key] = {
|
||||
"param_names": parameter_groups[key]["param_names"],
|
||||
"lr_scale": parameter_groups[key]["lr_scale"],
|
||||
"lr": parameter_groups[key]["lr"],
|
||||
"weight_decay": parameter_groups[key]["weight_decay"],
|
||||
}
|
||||
logger.info("Param groups = %s" % json.dumps(to_display, indent=2))
|
||||
|
||||
# state_dict = module.state_dict()
|
||||
# for group_name in parameter_groups:
|
||||
# group = parameter_groups[group_name]
|
||||
# for name in group["param_names"]:
|
||||
# group["params"].append(state_dict[name])
|
||||
|
||||
params.extend(parameter_groups.values())
|
||||
123
segmentation/mmcv_custom/layer_decay.py
Normal file
123
segmentation/mmcv_custom/layer_decay.py
Normal file
@@ -0,0 +1,123 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
import json
|
||||
from mmcv.runner import OPTIMIZER_BUILDERS, DefaultOptimizerConstructor
|
||||
from mmcv.runner import get_dist_info
|
||||
|
||||
|
||||
def get_num_layer_layer_wise(var_name, num_max_layer=12):
|
||||
|
||||
if var_name in ("backbone.cls_token", "backbone.mask_token", "backbone.pos_embed"):
|
||||
return 0
|
||||
elif var_name.startswith("backbone.downsample_layers"):
|
||||
stage_id = int(var_name.split('.')[2])
|
||||
if stage_id == 0:
|
||||
layer_id = 0
|
||||
elif stage_id == 1:
|
||||
layer_id = 2
|
||||
elif stage_id == 2:
|
||||
layer_id = 3
|
||||
elif stage_id == 3:
|
||||
layer_id = num_max_layer
|
||||
return layer_id
|
||||
elif var_name.startswith("backbone.stages"):
|
||||
stage_id = int(var_name.split('.')[2])
|
||||
block_id = int(var_name.split('.')[3])
|
||||
if stage_id == 0:
|
||||
layer_id = 1
|
||||
elif stage_id == 1:
|
||||
layer_id = 2
|
||||
elif stage_id == 2:
|
||||
layer_id = 3 + block_id // 3
|
||||
elif stage_id == 3:
|
||||
layer_id = num_max_layer
|
||||
return layer_id
|
||||
else:
|
||||
return num_max_layer + 1
|
||||
|
||||
|
||||
def get_num_layer_stage_wise(var_name, num_max_layer):
|
||||
if var_name in ("backbone.cls_token", "backbone.mask_token", "backbone.pos_embed"):
|
||||
return 0
|
||||
elif var_name.startswith("backbone.downsample_layers"):
|
||||
return 0
|
||||
elif var_name.startswith("backbone.stages"):
|
||||
stage_id = int(var_name.split('.')[2])
|
||||
return stage_id + 1
|
||||
else:
|
||||
return num_max_layer - 1
|
||||
|
||||
|
||||
@OPTIMIZER_BUILDERS.register_module()
|
||||
class LearningRateDecayOptimizerConstructor(DefaultOptimizerConstructor):
|
||||
def add_params(self, params, module, prefix='', is_dcn_module=None):
|
||||
"""Add all parameters of module to the params list.
|
||||
The parameters of the given module will be added to the list of param
|
||||
groups, with specific rules defined by paramwise_cfg.
|
||||
Args:
|
||||
params (list[dict]): A list of param groups, it will be modified
|
||||
in place.
|
||||
module (nn.Module): The module to be added.
|
||||
prefix (str): The prefix of the module
|
||||
is_dcn_module (int|float|None): If the current module is a
|
||||
submodule of DCN, `is_dcn_module` will be passed to
|
||||
control conv_offset layer's learning rate. Defaults to None.
|
||||
"""
|
||||
parameter_groups = {}
|
||||
print(self.paramwise_cfg)
|
||||
num_layers = self.paramwise_cfg.get('num_layers') + 2
|
||||
decay_rate = self.paramwise_cfg.get('decay_rate')
|
||||
decay_type = self.paramwise_cfg.get('decay_type', "layer_wise")
|
||||
print("Build LearningRateDecayOptimizerConstructor %s %f - %d" % (decay_type, decay_rate, num_layers))
|
||||
weight_decay = self.base_wd
|
||||
|
||||
for name, param in module.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue # frozen weights
|
||||
if len(param.shape) == 1 or name.endswith(".bias") or name in ('pos_embed', 'cls_token'):
|
||||
group_name = "no_decay"
|
||||
this_weight_decay = 0.
|
||||
else:
|
||||
group_name = "decay"
|
||||
this_weight_decay = weight_decay
|
||||
|
||||
if decay_type == "layer_wise":
|
||||
layer_id = get_num_layer_layer_wise(name, self.paramwise_cfg.get('num_layers'))
|
||||
elif decay_type == "stage_wise":
|
||||
layer_id = get_num_layer_stage_wise(name, num_layers)
|
||||
|
||||
group_name = "layer_%d_%s" % (layer_id, group_name)
|
||||
|
||||
if group_name not in parameter_groups:
|
||||
scale = decay_rate ** (num_layers - layer_id - 1)
|
||||
|
||||
parameter_groups[group_name] = {
|
||||
"weight_decay": this_weight_decay,
|
||||
"params": [],
|
||||
"param_names": [],
|
||||
"lr_scale": scale,
|
||||
"group_name": group_name,
|
||||
"lr": scale * self.base_lr,
|
||||
}
|
||||
|
||||
parameter_groups[group_name]["params"].append(param)
|
||||
parameter_groups[group_name]["param_names"].append(name)
|
||||
rank, _ = get_dist_info()
|
||||
if rank == 0:
|
||||
to_display = {}
|
||||
for key in parameter_groups:
|
||||
to_display[key] = {
|
||||
"param_names": parameter_groups[key]["param_names"],
|
||||
"lr_scale": parameter_groups[key]["lr_scale"],
|
||||
"lr": parameter_groups[key]["lr"],
|
||||
"weight_decay": parameter_groups[key]["weight_decay"],
|
||||
}
|
||||
print("Param groups = %s" % json.dumps(to_display, indent=2))
|
||||
|
||||
params.extend(parameter_groups.values())
|
||||
105
segmentation/mmcv_custom/layer_decay_vit.py
Normal file
105
segmentation/mmcv_custom/layer_decay_vit.py
Normal file
@@ -0,0 +1,105 @@
|
||||
import json
|
||||
from mmcv.runner import OPTIMIZER_BUILDERS, DefaultOptimizerConstructor
|
||||
from mmcv.runner import get_dist_info
|
||||
|
||||
|
||||
def get_num_layer_for_vit(var_name, num_max_layer, layer_sep=None):
|
||||
if var_name in ("backbone.cls_token", "backbone.mask_token", "backbone.pos_embed"):
|
||||
return 0
|
||||
elif var_name.startswith("backbone.patch_embed"):
|
||||
return 0
|
||||
elif var_name.startswith("backbone.blocks"):
|
||||
layer_id = int(var_name.split('.')[2])
|
||||
return layer_id + 1
|
||||
elif var_name.startswith("backbone.layers"):
|
||||
assert layer_sep is not None
|
||||
split = var_name.split('.')
|
||||
start_id = layer_sep[int(split[2])]
|
||||
if split[3] == 'RC':
|
||||
return start_id
|
||||
return start_id + int(split[4]) + 1
|
||||
else:
|
||||
return num_max_layer - 1
|
||||
|
||||
|
||||
@OPTIMIZER_BUILDERS.register_module()
|
||||
class LayerDecayOptimizerConstructor_vit(DefaultOptimizerConstructor):
|
||||
def add_params(self, params, module, prefix='', is_dcn_module=None):
|
||||
"""Add all parameters of module to the params list.
|
||||
The parameters of the given module will be added to the list of param
|
||||
groups, with specific rules defined by paramwise_cfg.
|
||||
Args:
|
||||
params (list[dict]): A list of param groups, it will be modified
|
||||
in place.
|
||||
module (nn.Module): The module to be added.
|
||||
prefix (str): The prefix of the module
|
||||
is_dcn_module (int|float|None): If the current module is a
|
||||
submodule of DCN, `is_dcn_module` will be passed to
|
||||
control conv_offset layer's learning rate. Defaults to None.
|
||||
"""
|
||||
# get param-wise options
|
||||
|
||||
parameter_groups = {}
|
||||
print(self.paramwise_cfg)
|
||||
num_layers = self.paramwise_cfg.get('num_layers') + 2
|
||||
layer_sep = self.paramwise_cfg.get('layer_sep', None)
|
||||
layer_decay_rate = self.paramwise_cfg.get('layer_decay_rate')
|
||||
print("Build LayerDecayOptimizerConstructor %f - %d" % (layer_decay_rate, num_layers))
|
||||
weight_decay = self.base_wd
|
||||
|
||||
custom_keys = self.paramwise_cfg.get('custom_keys', {})
|
||||
# first sort with alphabet order and then sort with reversed len of str
|
||||
sorted_keys = sorted(custom_keys.keys())
|
||||
|
||||
for name, param in module.named_parameters():
|
||||
|
||||
if not param.requires_grad:
|
||||
continue # frozen weights
|
||||
|
||||
if len(param.shape) == 1 or name.endswith(".bias") or ('pos_embed' in name) or ('cls_token' in name) or ('rel_pos_' in name):
|
||||
group_name = "no_decay"
|
||||
this_weight_decay = 0.
|
||||
else:
|
||||
group_name = "decay"
|
||||
this_weight_decay = weight_decay
|
||||
|
||||
layer_id = get_num_layer_for_vit(name, num_layers, layer_sep)
|
||||
group_name = "layer_%d_%s" % (layer_id, group_name)
|
||||
|
||||
# if the parameter match one of the custom keys, ignore other rules
|
||||
this_lr_multi = 1.
|
||||
for key in sorted_keys:
|
||||
if key in f'{name}':
|
||||
lr_mult = custom_keys[key].get('lr_mult', 1.)
|
||||
this_lr_multi = lr_mult
|
||||
group_name = "%s_%s" % (group_name, key)
|
||||
break
|
||||
|
||||
if group_name not in parameter_groups:
|
||||
scale = layer_decay_rate ** (num_layers - layer_id - 1)
|
||||
|
||||
parameter_groups[group_name] = {
|
||||
"weight_decay": this_weight_decay,
|
||||
"params": [],
|
||||
"param_names": [],
|
||||
"lr_scale": scale,
|
||||
"group_name": group_name,
|
||||
"lr": scale * self.base_lr * this_lr_multi,
|
||||
}
|
||||
|
||||
parameter_groups[group_name]["params"].append(param)
|
||||
parameter_groups[group_name]["param_names"].append(name)
|
||||
|
||||
rank, _ = get_dist_info()
|
||||
if rank == 0:
|
||||
to_display = {}
|
||||
for key in parameter_groups:
|
||||
to_display[key] = {
|
||||
"param_names": parameter_groups[key]["param_names"],
|
||||
"lr_scale": parameter_groups[key]["lr_scale"],
|
||||
"lr": parameter_groups[key]["lr"],
|
||||
"weight_decay": parameter_groups[key]["weight_decay"],
|
||||
}
|
||||
print("Param groups = %s" % json.dumps(to_display, indent=2))
|
||||
|
||||
params.extend(parameter_groups.values())
|
||||
9
segmentation/mmseg_custom/__init__.py
Normal file
9
segmentation/mmseg_custom/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
from .models import * # noqa: F401,F403
|
||||
from .datasets import * # noqa: F401,F403
|
||||
from .core import * # noqa: F401,F403
|
||||
9
segmentation/mmseg_custom/core/__init__.py
Normal file
9
segmentation/mmseg_custom/core/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from mmseg.core.evaluation import * # noqa: F401, F403
|
||||
from mmseg.core.seg import * # noqa: F401, F403
|
||||
|
||||
from .anchor import * # noqa: F401,F403
|
||||
from .box import * # noqa: F401,F403
|
||||
from .evaluation import * # noqa: F401,F403
|
||||
from .mask import * # noqa: F401,F403
|
||||
from .utils import * # noqa: F401, F403
|
||||
2
segmentation/mmseg_custom/core/anchor/__init__.py
Normal file
2
segmentation/mmseg_custom/core/anchor/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from .point_generator import MlvlPointGenerator # noqa: F401,F403
|
||||
19
segmentation/mmseg_custom/core/anchor/builder.py
Normal file
19
segmentation/mmseg_custom/core/anchor/builder.py
Normal file
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import warnings
|
||||
|
||||
from mmcv.utils import Registry, build_from_cfg
|
||||
|
||||
PRIOR_GENERATORS = Registry('Generator for anchors and points')
|
||||
|
||||
ANCHOR_GENERATORS = PRIOR_GENERATORS
|
||||
|
||||
|
||||
def build_prior_generator(cfg, default_args=None):
|
||||
return build_from_cfg(cfg, PRIOR_GENERATORS, default_args)
|
||||
|
||||
|
||||
def build_anchor_generator(cfg, default_args=None):
|
||||
warnings.warn(
|
||||
'``build_anchor_generator`` would be deprecated soon, please use '
|
||||
'``build_prior_generator`` ')
|
||||
return build_prior_generator(cfg, default_args=default_args)
|
||||
260
segmentation/mmseg_custom/core/anchor/point_generator.py
Normal file
260
segmentation/mmseg_custom/core/anchor/point_generator.py
Normal file
@@ -0,0 +1,260 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn.modules.utils import _pair
|
||||
|
||||
from .builder import PRIOR_GENERATORS
|
||||
|
||||
|
||||
@PRIOR_GENERATORS.register_module()
|
||||
class PointGenerator:
|
||||
def _meshgrid(self, x, y, row_major=True):
|
||||
xx = x.repeat(len(y))
|
||||
yy = y.view(-1, 1).repeat(1, len(x)).view(-1)
|
||||
if row_major:
|
||||
return xx, yy
|
||||
else:
|
||||
return yy, xx
|
||||
|
||||
def grid_points(self, featmap_size, stride=16, device='cuda'):
|
||||
feat_h, feat_w = featmap_size
|
||||
shift_x = torch.arange(0., feat_w, device=device) * stride
|
||||
shift_y = torch.arange(0., feat_h, device=device) * stride
|
||||
shift_xx, shift_yy = self._meshgrid(shift_x, shift_y)
|
||||
stride = shift_x.new_full((shift_xx.shape[0], ), stride)
|
||||
shifts = torch.stack([shift_xx, shift_yy, stride], dim=-1)
|
||||
all_points = shifts.to(device)
|
||||
return all_points
|
||||
|
||||
def valid_flags(self, featmap_size, valid_size, device='cuda'):
|
||||
feat_h, feat_w = featmap_size
|
||||
valid_h, valid_w = valid_size
|
||||
assert valid_h <= feat_h and valid_w <= feat_w
|
||||
valid_x = torch.zeros(feat_w, dtype=torch.bool, device=device)
|
||||
valid_y = torch.zeros(feat_h, dtype=torch.bool, device=device)
|
||||
valid_x[:valid_w] = 1
|
||||
valid_y[:valid_h] = 1
|
||||
valid_xx, valid_yy = self._meshgrid(valid_x, valid_y)
|
||||
valid = valid_xx & valid_yy
|
||||
return valid
|
||||
|
||||
|
||||
@PRIOR_GENERATORS.register_module()
|
||||
class MlvlPointGenerator:
|
||||
"""Standard points generator for multi-level (Mlvl) feature maps in 2D
|
||||
points-based detectors.
|
||||
|
||||
Args:
|
||||
strides (list[int] | list[tuple[int, int]]): Strides of anchors
|
||||
in multiple feature levels in order (w, h).
|
||||
offset (float): The offset of points, the value is normalized with
|
||||
corresponding stride. Defaults to 0.5.
|
||||
"""
|
||||
def __init__(self, strides, offset=0.5):
|
||||
self.strides = [_pair(stride) for stride in strides]
|
||||
self.offset = offset
|
||||
|
||||
@property
|
||||
def num_levels(self):
|
||||
"""int: number of feature levels that the generator will be applied"""
|
||||
return len(self.strides)
|
||||
|
||||
@property
|
||||
def num_base_priors(self):
|
||||
"""list[int]: The number of priors (points) at a point
|
||||
on the feature grid"""
|
||||
return [1 for _ in range(len(self.strides))]
|
||||
|
||||
def _meshgrid(self, x, y, row_major=True):
|
||||
yy, xx = torch.meshgrid(y, x)
|
||||
if row_major:
|
||||
# warning .flatten() would cause error in ONNX exporting
|
||||
# have to use reshape here
|
||||
return xx.reshape(-1), yy.reshape(-1)
|
||||
|
||||
else:
|
||||
return yy.reshape(-1), xx.reshape(-1)
|
||||
|
||||
def grid_priors(self,
|
||||
featmap_sizes,
|
||||
dtype=torch.float32,
|
||||
device='cuda',
|
||||
with_stride=False):
|
||||
"""Generate grid points of multiple feature levels.
|
||||
|
||||
Args:
|
||||
featmap_sizes (list[tuple]): List of feature map sizes in
|
||||
multiple feature levels, each size arrange as
|
||||
as (h, w).
|
||||
dtype (:obj:`dtype`): Dtype of priors. Default: torch.float32.
|
||||
device (str): The device where the anchors will be put on.
|
||||
with_stride (bool): Whether to concatenate the stride to
|
||||
the last dimension of points.
|
||||
|
||||
Return:
|
||||
list[torch.Tensor]: Points of multiple feature levels.
|
||||
The sizes of each tensor should be (N, 2) when with stride is
|
||||
``False``, where N = width * height, width and height
|
||||
are the sizes of the corresponding feature level,
|
||||
and the last dimension 2 represent (coord_x, coord_y),
|
||||
otherwise the shape should be (N, 4),
|
||||
and the last dimension 4 represent
|
||||
(coord_x, coord_y, stride_w, stride_h).
|
||||
"""
|
||||
|
||||
assert self.num_levels == len(featmap_sizes)
|
||||
multi_level_priors = []
|
||||
for i in range(self.num_levels):
|
||||
priors = self.single_level_grid_priors(featmap_sizes[i],
|
||||
level_idx=i,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
with_stride=with_stride)
|
||||
multi_level_priors.append(priors)
|
||||
return multi_level_priors
|
||||
|
||||
def single_level_grid_priors(self,
|
||||
featmap_size,
|
||||
level_idx,
|
||||
dtype=torch.float32,
|
||||
device='cuda',
|
||||
with_stride=False):
|
||||
"""Generate grid Points of a single level.
|
||||
|
||||
Note:
|
||||
This function is usually called by method ``self.grid_priors``.
|
||||
|
||||
Args:
|
||||
featmap_size (tuple[int]): Size of the feature maps, arrange as
|
||||
(h, w).
|
||||
level_idx (int): The index of corresponding feature map level.
|
||||
dtype (:obj:`dtype`): Dtype of priors. Default: torch.float32.
|
||||
device (str, optional): The device the tensor will be put on.
|
||||
Defaults to 'cuda'.
|
||||
with_stride (bool): Concatenate the stride to the last dimension
|
||||
of points.
|
||||
|
||||
Return:
|
||||
Tensor: Points of single feature levels.
|
||||
The shape of tensor should be (N, 2) when with stride is
|
||||
``False``, where N = width * height, width and height
|
||||
are the sizes of the corresponding feature level,
|
||||
and the last dimension 2 represent (coord_x, coord_y),
|
||||
otherwise the shape should be (N, 4),
|
||||
and the last dimension 4 represent
|
||||
(coord_x, coord_y, stride_w, stride_h).
|
||||
"""
|
||||
feat_h, feat_w = featmap_size
|
||||
stride_w, stride_h = self.strides[level_idx]
|
||||
shift_x = (torch.arange(0, feat_w, device=device) +
|
||||
self.offset) * stride_w
|
||||
# keep featmap_size as Tensor instead of int, so that we
|
||||
# can convert to ONNX correctly
|
||||
shift_x = shift_x.to(dtype)
|
||||
|
||||
shift_y = (torch.arange(0, feat_h, device=device) +
|
||||
self.offset) * stride_h
|
||||
# keep featmap_size as Tensor instead of int, so that we
|
||||
# can convert to ONNX correctly
|
||||
shift_y = shift_y.to(dtype)
|
||||
shift_xx, shift_yy = self._meshgrid(shift_x, shift_y)
|
||||
if not with_stride:
|
||||
shifts = torch.stack([shift_xx, shift_yy], dim=-1)
|
||||
else:
|
||||
# use `shape[0]` instead of `len(shift_xx)` for ONNX export
|
||||
stride_w = shift_xx.new_full((shift_xx.shape[0], ),
|
||||
stride_w).to(dtype)
|
||||
stride_h = shift_xx.new_full((shift_yy.shape[0], ),
|
||||
stride_h).to(dtype)
|
||||
shifts = torch.stack([shift_xx, shift_yy, stride_w, stride_h],
|
||||
dim=-1)
|
||||
all_points = shifts.to(device)
|
||||
return all_points
|
||||
|
||||
def valid_flags(self, featmap_sizes, pad_shape, device='cuda'):
|
||||
"""Generate valid flags of points of multiple feature levels.
|
||||
|
||||
Args:
|
||||
featmap_sizes (list(tuple)): List of feature map sizes in
|
||||
multiple feature levels, each size arrange as
|
||||
as (h, w).
|
||||
pad_shape (tuple(int)): The padded shape of the image,
|
||||
arrange as (h, w).
|
||||
device (str): The device where the anchors will be put on.
|
||||
|
||||
Return:
|
||||
list(torch.Tensor): Valid flags of points of multiple levels.
|
||||
"""
|
||||
assert self.num_levels == len(featmap_sizes)
|
||||
multi_level_flags = []
|
||||
for i in range(self.num_levels):
|
||||
point_stride = self.strides[i]
|
||||
feat_h, feat_w = featmap_sizes[i]
|
||||
h, w = pad_shape[:2]
|
||||
valid_feat_h = min(int(np.ceil(h / point_stride[1])), feat_h)
|
||||
valid_feat_w = min(int(np.ceil(w / point_stride[0])), feat_w)
|
||||
flags = self.single_level_valid_flags((feat_h, feat_w),
|
||||
(valid_feat_h, valid_feat_w),
|
||||
device=device)
|
||||
multi_level_flags.append(flags)
|
||||
return multi_level_flags
|
||||
|
||||
def single_level_valid_flags(self,
|
||||
featmap_size,
|
||||
valid_size,
|
||||
device='cuda'):
|
||||
"""Generate the valid flags of points of a single feature map.
|
||||
|
||||
Args:
|
||||
featmap_size (tuple[int]): The size of feature maps, arrange as
|
||||
as (h, w).
|
||||
valid_size (tuple[int]): The valid size of the feature maps.
|
||||
The size arrange as as (h, w).
|
||||
device (str, optional): The device where the flags will be put on.
|
||||
Defaults to 'cuda'.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The valid flags of each points in a single level \
|
||||
feature map.
|
||||
"""
|
||||
feat_h, feat_w = featmap_size
|
||||
valid_h, valid_w = valid_size
|
||||
assert valid_h <= feat_h and valid_w <= feat_w
|
||||
valid_x = torch.zeros(feat_w, dtype=torch.bool, device=device)
|
||||
valid_y = torch.zeros(feat_h, dtype=torch.bool, device=device)
|
||||
valid_x[:valid_w] = 1
|
||||
valid_y[:valid_h] = 1
|
||||
valid_xx, valid_yy = self._meshgrid(valid_x, valid_y)
|
||||
valid = valid_xx & valid_yy
|
||||
return valid
|
||||
|
||||
def sparse_priors(self,
|
||||
prior_idxs,
|
||||
featmap_size,
|
||||
level_idx,
|
||||
dtype=torch.float32,
|
||||
device='cuda'):
|
||||
"""Generate sparse points according to the ``prior_idxs``.
|
||||
|
||||
Args:
|
||||
prior_idxs (Tensor): The index of corresponding anchors
|
||||
in the feature map.
|
||||
featmap_size (tuple[int]): feature map size arrange as (w, h).
|
||||
level_idx (int): The level index of corresponding feature
|
||||
map.
|
||||
dtype (obj:`torch.dtype`): Date type of points. Defaults to
|
||||
``torch.float32``.
|
||||
device (obj:`torch.device`): The device where the points is
|
||||
located.
|
||||
Returns:
|
||||
Tensor: Anchor with shape (N, 2), N should be equal to
|
||||
the length of ``prior_idxs``. And last dimension
|
||||
2 represent (coord_x, coord_y).
|
||||
"""
|
||||
height, width = featmap_size
|
||||
x = (prior_idxs % width + self.offset) * self.strides[level_idx][0]
|
||||
y = ((prior_idxs // width) % height +
|
||||
self.offset) * self.strides[level_idx][1]
|
||||
prioris = torch.stack([x, y], 1).to(dtype)
|
||||
prioris = prioris.to(device)
|
||||
return prioris
|
||||
3
segmentation/mmseg_custom/core/box/__init__.py
Normal file
3
segmentation/mmseg_custom/core/box/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from .builder import * # noqa: F401,F403
|
||||
from .samplers import MaskPseudoSampler # noqa: F401,F403
|
||||
15
segmentation/mmseg_custom/core/box/builder.py
Normal file
15
segmentation/mmseg_custom/core/box/builder.py
Normal file
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from mmcv.utils import Registry, build_from_cfg
|
||||
|
||||
BBOX_SAMPLERS = Registry('bbox_sampler')
|
||||
BBOX_CODERS = Registry('bbox_coder')
|
||||
|
||||
|
||||
def build_sampler(cfg, **default_args):
|
||||
"""Builder of box sampler."""
|
||||
return build_from_cfg(cfg, BBOX_SAMPLERS, default_args)
|
||||
|
||||
|
||||
def build_bbox_coder(cfg, **default_args):
|
||||
"""Builder of box coder."""
|
||||
return build_from_cfg(cfg, BBOX_CODERS, default_args)
|
||||
2
segmentation/mmseg_custom/core/box/samplers/__init__.py
Normal file
2
segmentation/mmseg_custom/core/box/samplers/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from .mask_pseudo_sampler import MaskPseudoSampler # noqa: F401,F403
|
||||
105
segmentation/mmseg_custom/core/box/samplers/base_sampler.py
Normal file
105
segmentation/mmseg_custom/core/box/samplers/base_sampler.py
Normal file
@@ -0,0 +1,105 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from abc import ABCMeta, abstractmethod
|
||||
|
||||
import torch
|
||||
|
||||
from .sampling_result import SamplingResult
|
||||
|
||||
|
||||
class BaseSampler(metaclass=ABCMeta):
|
||||
"""Base class of samplers."""
|
||||
def __init__(self,
|
||||
num,
|
||||
pos_fraction,
|
||||
neg_pos_ub=-1,
|
||||
add_gt_as_proposals=True,
|
||||
**kwargs):
|
||||
self.num = num
|
||||
self.pos_fraction = pos_fraction
|
||||
self.neg_pos_ub = neg_pos_ub
|
||||
self.add_gt_as_proposals = add_gt_as_proposals
|
||||
self.pos_sampler = self
|
||||
self.neg_sampler = self
|
||||
|
||||
@abstractmethod
|
||||
def _sample_pos(self, assign_result, num_expected, **kwargs):
|
||||
"""Sample positive samples."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _sample_neg(self, assign_result, num_expected, **kwargs):
|
||||
"""Sample negative samples."""
|
||||
pass
|
||||
|
||||
def sample(self,
|
||||
assign_result,
|
||||
bboxes,
|
||||
gt_bboxes,
|
||||
gt_labels=None,
|
||||
**kwargs):
|
||||
"""Sample positive and negative bboxes.
|
||||
|
||||
This is a simple implementation of bbox sampling given candidates,
|
||||
assigning results and ground truth bboxes.
|
||||
|
||||
Args:
|
||||
assign_result (:obj:`AssignResult`): Bbox assigning results.
|
||||
bboxes (Tensor): Boxes to be sampled from.
|
||||
gt_bboxes (Tensor): Ground truth bboxes.
|
||||
gt_labels (Tensor, optional): Class labels of ground truth bboxes.
|
||||
|
||||
Returns:
|
||||
:obj:`SamplingResult`: Sampling result.
|
||||
|
||||
Example:
|
||||
>>> from mmdet.core.bbox import RandomSampler
|
||||
>>> from mmdet.core.bbox import AssignResult
|
||||
>>> from mmdet.core.bbox.demodata import ensure_rng, random_boxes
|
||||
>>> rng = ensure_rng(None)
|
||||
>>> assign_result = AssignResult.random(rng=rng)
|
||||
>>> bboxes = random_boxes(assign_result.num_preds, rng=rng)
|
||||
>>> gt_bboxes = random_boxes(assign_result.num_gts, rng=rng)
|
||||
>>> gt_labels = None
|
||||
>>> self = RandomSampler(num=32, pos_fraction=0.5, neg_pos_ub=-1,
|
||||
>>> add_gt_as_proposals=False)
|
||||
>>> self = self.sample(assign_result, bboxes, gt_bboxes, gt_labels)
|
||||
"""
|
||||
if len(bboxes.shape) < 2:
|
||||
bboxes = bboxes[None, :]
|
||||
|
||||
bboxes = bboxes[:, :4]
|
||||
|
||||
gt_flags = bboxes.new_zeros((bboxes.shape[0], ), dtype=torch.uint8)
|
||||
if self.add_gt_as_proposals and len(gt_bboxes) > 0:
|
||||
if gt_labels is None:
|
||||
raise ValueError(
|
||||
'gt_labels must be given when add_gt_as_proposals is True')
|
||||
bboxes = torch.cat([gt_bboxes, bboxes], dim=0)
|
||||
assign_result.add_gt_(gt_labels)
|
||||
gt_ones = bboxes.new_ones(gt_bboxes.shape[0], dtype=torch.uint8)
|
||||
gt_flags = torch.cat([gt_ones, gt_flags])
|
||||
|
||||
num_expected_pos = int(self.num * self.pos_fraction)
|
||||
pos_inds = self.pos_sampler._sample_pos(assign_result,
|
||||
num_expected_pos,
|
||||
bboxes=bboxes,
|
||||
**kwargs)
|
||||
# We found that sampled indices have duplicated items occasionally.
|
||||
# (may be a bug of PyTorch)
|
||||
pos_inds = pos_inds.unique()
|
||||
num_sampled_pos = pos_inds.numel()
|
||||
num_expected_neg = self.num - num_sampled_pos
|
||||
if self.neg_pos_ub >= 0:
|
||||
_pos = max(1, num_sampled_pos)
|
||||
neg_upper_bound = int(self.neg_pos_ub * _pos)
|
||||
if num_expected_neg > neg_upper_bound:
|
||||
num_expected_neg = neg_upper_bound
|
||||
neg_inds = self.neg_sampler._sample_neg(assign_result,
|
||||
num_expected_neg,
|
||||
bboxes=bboxes,
|
||||
**kwargs)
|
||||
neg_inds = neg_inds.unique()
|
||||
|
||||
sampling_result = SamplingResult(pos_inds, neg_inds, bboxes, gt_bboxes,
|
||||
assign_result, gt_flags)
|
||||
return sampling_result
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
"""copy from
|
||||
https://github.com/ZwwWayne/K-Net/blob/main/knet/det/mask_pseudo_sampler.py."""
|
||||
|
||||
import torch
|
||||
|
||||
from ..builder import BBOX_SAMPLERS
|
||||
from .base_sampler import BaseSampler
|
||||
from .mask_sampling_result import MaskSamplingResult
|
||||
|
||||
|
||||
@BBOX_SAMPLERS.register_module()
|
||||
class MaskPseudoSampler(BaseSampler):
|
||||
"""A pseudo sampler that does not do sampling actually."""
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def _sample_pos(self, **kwargs):
|
||||
"""Sample positive samples."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _sample_neg(self, **kwargs):
|
||||
"""Sample negative samples."""
|
||||
raise NotImplementedError
|
||||
|
||||
def sample(self, assign_result, masks, gt_masks, **kwargs):
|
||||
"""Directly returns the positive and negative indices of samples.
|
||||
|
||||
Args:
|
||||
assign_result (:obj:`AssignResult`): Assigned results
|
||||
masks (torch.Tensor): Bounding boxes
|
||||
gt_masks (torch.Tensor): Ground truth boxes
|
||||
Returns:
|
||||
:obj:`SamplingResult`: sampler results
|
||||
"""
|
||||
pos_inds = torch.nonzero(assign_result.gt_inds > 0,
|
||||
as_tuple=False).squeeze(-1).unique()
|
||||
neg_inds = torch.nonzero(assign_result.gt_inds == 0,
|
||||
as_tuple=False).squeeze(-1).unique()
|
||||
gt_flags = masks.new_zeros(masks.shape[0], dtype=torch.uint8)
|
||||
sampling_result = MaskSamplingResult(pos_inds, neg_inds, masks,
|
||||
gt_masks, assign_result, gt_flags)
|
||||
return sampling_result
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
"""copy from
|
||||
https://github.com/ZwwWayne/K-Net/blob/main/knet/det/mask_pseudo_sampler.py."""
|
||||
|
||||
import torch
|
||||
|
||||
from .sampling_result import SamplingResult
|
||||
|
||||
|
||||
class MaskSamplingResult(SamplingResult):
|
||||
"""Mask sampling result."""
|
||||
def __init__(self, pos_inds, neg_inds, masks, gt_masks, assign_result,
|
||||
gt_flags):
|
||||
self.pos_inds = pos_inds
|
||||
self.neg_inds = neg_inds
|
||||
self.pos_masks = masks[pos_inds]
|
||||
self.neg_masks = masks[neg_inds]
|
||||
self.pos_is_gt = gt_flags[pos_inds]
|
||||
|
||||
self.num_gts = gt_masks.shape[0]
|
||||
self.pos_assigned_gt_inds = assign_result.gt_inds[pos_inds] - 1
|
||||
|
||||
if gt_masks.numel() == 0:
|
||||
# hack for index error case
|
||||
assert self.pos_assigned_gt_inds.numel() == 0
|
||||
self.pos_gt_masks = torch.empty_like(gt_masks)
|
||||
else:
|
||||
self.pos_gt_masks = gt_masks[self.pos_assigned_gt_inds, :]
|
||||
|
||||
if assign_result.labels is not None:
|
||||
self.pos_gt_labels = assign_result.labels[pos_inds]
|
||||
else:
|
||||
self.pos_gt_labels = None
|
||||
|
||||
@property
|
||||
def masks(self):
|
||||
"""torch.Tensor: concatenated positive and negative boxes"""
|
||||
return torch.cat([self.pos_masks, self.neg_masks])
|
||||
|
||||
def __nice__(self):
|
||||
data = self.info.copy()
|
||||
data['pos_masks'] = data.pop('pos_masks').shape
|
||||
data['neg_masks'] = data.pop('neg_masks').shape
|
||||
parts = [f"'{k}': {v!r}" for k, v in sorted(data.items())]
|
||||
body = ' ' + ',\n '.join(parts)
|
||||
return '{\n' + body + '\n}'
|
||||
|
||||
@property
|
||||
def info(self):
|
||||
"""Returns a dictionary of info about the object."""
|
||||
return {
|
||||
'pos_inds': self.pos_inds,
|
||||
'neg_inds': self.neg_inds,
|
||||
'pos_masks': self.pos_masks,
|
||||
'neg_masks': self.neg_masks,
|
||||
'pos_is_gt': self.pos_is_gt,
|
||||
'num_gts': self.num_gts,
|
||||
'pos_assigned_gt_inds': self.pos_assigned_gt_inds,
|
||||
}
|
||||
150
segmentation/mmseg_custom/core/box/samplers/sampling_result.py
Normal file
150
segmentation/mmseg_custom/core/box/samplers/sampling_result.py
Normal file
@@ -0,0 +1,150 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
from mmdet.utils import util_mixins
|
||||
|
||||
|
||||
class SamplingResult(util_mixins.NiceRepr):
|
||||
"""Bbox sampling result.
|
||||
|
||||
Example:
|
||||
>>> # xdoctest: +IGNORE_WANT
|
||||
>>> from mmdet.core.bbox.samplers.sampling_result import * # NOQA
|
||||
>>> self = SamplingResult.random(rng=10)
|
||||
>>> print(f'self = {self}')
|
||||
self = <SamplingResult({
|
||||
'neg_bboxes': torch.Size([12, 4]),
|
||||
'neg_inds': tensor([ 0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 11, 12]),
|
||||
'num_gts': 4,
|
||||
'pos_assigned_gt_inds': tensor([], dtype=torch.int64),
|
||||
'pos_bboxes': torch.Size([0, 4]),
|
||||
'pos_inds': tensor([], dtype=torch.int64),
|
||||
'pos_is_gt': tensor([], dtype=torch.uint8)
|
||||
})>
|
||||
"""
|
||||
def __init__(self, pos_inds, neg_inds, bboxes, gt_bboxes, assign_result,
|
||||
gt_flags):
|
||||
self.pos_inds = pos_inds
|
||||
self.neg_inds = neg_inds
|
||||
self.pos_bboxes = bboxes[pos_inds]
|
||||
self.neg_bboxes = bboxes[neg_inds]
|
||||
self.pos_is_gt = gt_flags[pos_inds]
|
||||
|
||||
self.num_gts = gt_bboxes.shape[0]
|
||||
self.pos_assigned_gt_inds = assign_result.gt_inds[pos_inds] - 1
|
||||
|
||||
if gt_bboxes.numel() == 0:
|
||||
# hack for index error case
|
||||
assert self.pos_assigned_gt_inds.numel() == 0
|
||||
self.pos_gt_bboxes = torch.empty_like(gt_bboxes).view(-1, 4)
|
||||
else:
|
||||
if len(gt_bboxes.shape) < 2:
|
||||
gt_bboxes = gt_bboxes.view(-1, 4)
|
||||
|
||||
self.pos_gt_bboxes = gt_bboxes[self.pos_assigned_gt_inds.long(), :]
|
||||
|
||||
if assign_result.labels is not None:
|
||||
self.pos_gt_labels = assign_result.labels[pos_inds]
|
||||
else:
|
||||
self.pos_gt_labels = None
|
||||
|
||||
@property
|
||||
def bboxes(self):
|
||||
"""torch.Tensor: concatenated positive and negative boxes"""
|
||||
return torch.cat([self.pos_bboxes, self.neg_bboxes])
|
||||
|
||||
def to(self, device):
|
||||
"""Change the device of the data inplace.
|
||||
|
||||
Example:
|
||||
>>> self = SamplingResult.random()
|
||||
>>> print(f'self = {self.to(None)}')
|
||||
>>> # xdoctest: +REQUIRES(--gpu)
|
||||
>>> print(f'self = {self.to(0)}')
|
||||
"""
|
||||
_dict = self.__dict__
|
||||
for key, value in _dict.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
_dict[key] = value.to(device)
|
||||
return self
|
||||
|
||||
def __nice__(self):
|
||||
data = self.info.copy()
|
||||
data['pos_bboxes'] = data.pop('pos_bboxes').shape
|
||||
data['neg_bboxes'] = data.pop('neg_bboxes').shape
|
||||
parts = [f"'{k}': {v!r}" for k, v in sorted(data.items())]
|
||||
body = ' ' + ',\n '.join(parts)
|
||||
return '{\n' + body + '\n}'
|
||||
|
||||
@property
|
||||
def info(self):
|
||||
"""Returns a dictionary of info about the object."""
|
||||
return {
|
||||
'pos_inds': self.pos_inds,
|
||||
'neg_inds': self.neg_inds,
|
||||
'pos_bboxes': self.pos_bboxes,
|
||||
'neg_bboxes': self.neg_bboxes,
|
||||
'pos_is_gt': self.pos_is_gt,
|
||||
'num_gts': self.num_gts,
|
||||
'pos_assigned_gt_inds': self.pos_assigned_gt_inds,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def random(cls, rng=None, **kwargs):
|
||||
"""
|
||||
Args:
|
||||
rng (None | int | numpy.random.RandomState): seed or state.
|
||||
kwargs (keyword arguments):
|
||||
- num_preds: number of predicted boxes
|
||||
- num_gts: number of true boxes
|
||||
- p_ignore (float): probability of a predicted box assigned to \
|
||||
an ignored truth.
|
||||
- p_assigned (float): probability of a predicted box not being \
|
||||
assigned.
|
||||
- p_use_label (float | bool): with labels or not.
|
||||
|
||||
Returns:
|
||||
:obj:`SamplingResult`: Randomly generated sampling result.
|
||||
|
||||
Example:
|
||||
>>> from mmdet.core.bbox.samplers.sampling_result import * # NOQA
|
||||
>>> self = SamplingResult.random()
|
||||
>>> print(self.__dict__)
|
||||
"""
|
||||
from mmdet.core.bbox import demodata
|
||||
from mmdet.core.bbox.assigners.assign_result import AssignResult
|
||||
from mmdet.core.bbox.samplers.random_sampler import RandomSampler
|
||||
rng = demodata.ensure_rng(rng)
|
||||
|
||||
# make probabalistic?
|
||||
num = 32
|
||||
pos_fraction = 0.5
|
||||
neg_pos_ub = -1
|
||||
|
||||
assign_result = AssignResult.random(rng=rng, **kwargs)
|
||||
|
||||
# Note we could just compute an assignment
|
||||
bboxes = demodata.random_boxes(assign_result.num_preds, rng=rng)
|
||||
gt_bboxes = demodata.random_boxes(assign_result.num_gts, rng=rng)
|
||||
|
||||
if rng.rand() > 0.2:
|
||||
# sometimes algorithms squeeze their data, be robust to that
|
||||
gt_bboxes = gt_bboxes.squeeze()
|
||||
bboxes = bboxes.squeeze()
|
||||
|
||||
if assign_result.labels is None:
|
||||
gt_labels = None
|
||||
else:
|
||||
gt_labels = None # todo
|
||||
|
||||
if gt_labels is None:
|
||||
add_gt_as_proposals = False
|
||||
else:
|
||||
add_gt_as_proposals = True # make probabalistic?
|
||||
|
||||
sampler = RandomSampler(num,
|
||||
pos_fraction,
|
||||
neg_pos_ub=neg_pos_ub,
|
||||
add_gt_as_proposals=add_gt_as_proposals,
|
||||
rng=rng)
|
||||
self = sampler.sample(assign_result, bboxes, gt_bboxes, gt_labels)
|
||||
return self
|
||||
2
segmentation/mmseg_custom/core/evaluation/__init__.py
Normal file
2
segmentation/mmseg_custom/core/evaluation/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from .panoptic_utils import INSTANCE_OFFSET # noqa: F401,F403
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
# A custom value to distinguish instance ID and category ID; need to
|
||||
# be greater than the number of categories.
|
||||
# For a pixel in the panoptic result map:
|
||||
# pan_id = ins_id * INSTANCE_OFFSET + cat_id
|
||||
INSTANCE_OFFSET = 1000
|
||||
2
segmentation/mmseg_custom/core/mask/__init__.py
Normal file
2
segmentation/mmseg_custom/core/mask/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from .utils import mask2bbox # noqa: F401,F403
|
||||
89
segmentation/mmseg_custom/core/mask/utils.py
Normal file
89
segmentation/mmseg_custom/core/mask/utils.py
Normal file
@@ -0,0 +1,89 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import mmcv
|
||||
import numpy as np
|
||||
import pycocotools.mask as mask_util
|
||||
import torch
|
||||
|
||||
|
||||
def split_combined_polys(polys, poly_lens, polys_per_mask):
|
||||
"""Split the combined 1-D polys into masks.
|
||||
|
||||
A mask is represented as a list of polys, and a poly is represented as
|
||||
a 1-D array. In dataset, all masks are concatenated into a single 1-D
|
||||
tensor. Here we need to split the tensor into original representations.
|
||||
|
||||
Args:
|
||||
polys (list): a list (length = image num) of 1-D tensors
|
||||
poly_lens (list): a list (length = image num) of poly length
|
||||
polys_per_mask (list): a list (length = image num) of poly number
|
||||
of each mask
|
||||
|
||||
Returns:
|
||||
list: a list (length = image num) of list (length = mask num) of \
|
||||
list (length = poly num) of numpy array.
|
||||
"""
|
||||
mask_polys_list = []
|
||||
for img_id in range(len(polys)):
|
||||
polys_single = polys[img_id]
|
||||
polys_lens_single = poly_lens[img_id].tolist()
|
||||
polys_per_mask_single = polys_per_mask[img_id].tolist()
|
||||
|
||||
split_polys = mmcv.slice_list(polys_single, polys_lens_single)
|
||||
mask_polys = mmcv.slice_list(split_polys, polys_per_mask_single)
|
||||
mask_polys_list.append(mask_polys)
|
||||
return mask_polys_list
|
||||
|
||||
|
||||
# TODO: move this function to more proper place
|
||||
def encode_mask_results(mask_results):
|
||||
"""Encode bitmap mask to RLE code.
|
||||
|
||||
Args:
|
||||
mask_results (list | tuple[list]): bitmap mask results.
|
||||
In mask scoring rcnn, mask_results is a tuple of (segm_results,
|
||||
segm_cls_score).
|
||||
|
||||
Returns:
|
||||
list | tuple: RLE encoded mask.
|
||||
"""
|
||||
if isinstance(mask_results, tuple): # mask scoring
|
||||
cls_segms, cls_mask_scores = mask_results
|
||||
else:
|
||||
cls_segms = mask_results
|
||||
num_classes = len(cls_segms)
|
||||
encoded_mask_results = [[] for _ in range(num_classes)]
|
||||
for i in range(len(cls_segms)):
|
||||
for cls_segm in cls_segms[i]:
|
||||
encoded_mask_results[i].append(
|
||||
mask_util.encode(
|
||||
np.array(
|
||||
cls_segm[:, :, np.newaxis], order='F',
|
||||
dtype='uint8'))[0]) # encoded with RLE
|
||||
if isinstance(mask_results, tuple):
|
||||
return encoded_mask_results, cls_mask_scores
|
||||
else:
|
||||
return encoded_mask_results
|
||||
|
||||
|
||||
def mask2bbox(masks):
|
||||
"""Obtain tight bounding boxes of binary masks.
|
||||
|
||||
Args:
|
||||
masks (Tensor): Binary mask of shape (n, h, w).
|
||||
|
||||
Returns:
|
||||
Tensor: Bboxe with shape (n, 4) of \
|
||||
positive region in binary mask.
|
||||
"""
|
||||
N = masks.shape[0]
|
||||
bboxes = masks.new_zeros((N, 4), dtype=torch.float32)
|
||||
x_any = torch.any(masks, dim=1)
|
||||
y_any = torch.any(masks, dim=2)
|
||||
for i in range(N):
|
||||
x = torch.where(x_any[i, :])[0]
|
||||
y = torch.where(y_any[i, :])[0]
|
||||
if len(x) > 0 and len(y) > 0:
|
||||
bboxes[i, :] = bboxes.new_tensor(
|
||||
[x[0], y[0], x[-1] + 1, y[-1] + 1])
|
||||
|
||||
return bboxes
|
||||
9
segmentation/mmseg_custom/core/utils/__init__.py
Normal file
9
segmentation/mmseg_custom/core/utils/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .dist_utils import (DistOptimizerHook, all_reduce_dict, allreduce_grads,
|
||||
reduce_mean)
|
||||
from .misc import add_prefix, multi_apply
|
||||
|
||||
__all__ = [
|
||||
'add_prefix', 'multi_apply', 'DistOptimizerHook', 'allreduce_grads',
|
||||
'all_reduce_dict', 'reduce_mean'
|
||||
]
|
||||
148
segmentation/mmseg_custom/core/utils/dist_utils.py
Normal file
148
segmentation/mmseg_custom/core/utils/dist_utils.py
Normal file
@@ -0,0 +1,148 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import functools
|
||||
import pickle
|
||||
import warnings
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from mmcv.runner import OptimizerHook, get_dist_info
|
||||
from torch._utils import (_flatten_dense_tensors, _take_tensors,
|
||||
_unflatten_dense_tensors)
|
||||
|
||||
|
||||
def _allreduce_coalesced(tensors, world_size, bucket_size_mb=-1):
|
||||
if bucket_size_mb > 0:
|
||||
bucket_size_bytes = bucket_size_mb * 1024 * 1024
|
||||
buckets = _take_tensors(tensors, bucket_size_bytes)
|
||||
else:
|
||||
buckets = OrderedDict()
|
||||
for tensor in tensors:
|
||||
tp = tensor.type()
|
||||
if tp not in buckets:
|
||||
buckets[tp] = []
|
||||
buckets[tp].append(tensor)
|
||||
buckets = buckets.values()
|
||||
|
||||
for bucket in buckets:
|
||||
flat_tensors = _flatten_dense_tensors(bucket)
|
||||
dist.all_reduce(flat_tensors)
|
||||
flat_tensors.div_(world_size)
|
||||
for tensor, synced in zip(
|
||||
bucket, _unflatten_dense_tensors(flat_tensors, bucket)):
|
||||
tensor.copy_(synced)
|
||||
|
||||
|
||||
def allreduce_grads(params, coalesce=True, bucket_size_mb=-1):
|
||||
"""Allreduce gradients.
|
||||
|
||||
Args:
|
||||
params (list[torch.Parameters]): List of parameters of a model
|
||||
coalesce (bool, optional): Whether allreduce parameters as a whole.
|
||||
Defaults to True.
|
||||
bucket_size_mb (int, optional): Size of bucket, the unit is MB.
|
||||
Defaults to -1.
|
||||
"""
|
||||
grads = [
|
||||
param.grad.data for param in params
|
||||
if param.requires_grad and param.grad is not None
|
||||
]
|
||||
world_size = dist.get_world_size()
|
||||
if coalesce:
|
||||
_allreduce_coalesced(grads, world_size, bucket_size_mb)
|
||||
else:
|
||||
for tensor in grads:
|
||||
dist.all_reduce(tensor.div_(world_size))
|
||||
|
||||
|
||||
class DistOptimizerHook(OptimizerHook):
|
||||
"""Deprecated optimizer hook for distributed training."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
warnings.warn('"DistOptimizerHook" is deprecated, please switch to'
|
||||
'"mmcv.runner.OptimizerHook".')
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
def reduce_mean(tensor):
|
||||
""""Obtain the mean of tensor on different GPUs."""
|
||||
if not (dist.is_available() and dist.is_initialized()):
|
||||
return tensor
|
||||
tensor = tensor.clone()
|
||||
dist.all_reduce(tensor.div_(dist.get_world_size()), op=dist.ReduceOp.SUM)
|
||||
return tensor
|
||||
|
||||
|
||||
def obj2tensor(pyobj, device='cuda'):
|
||||
"""Serialize picklable python object to tensor."""
|
||||
storage = torch.ByteStorage.from_buffer(pickle.dumps(pyobj))
|
||||
return torch.ByteTensor(storage).to(device=device)
|
||||
|
||||
|
||||
def tensor2obj(tensor):
|
||||
"""Deserialize tensor to picklable python object."""
|
||||
return pickle.loads(tensor.cpu().numpy().tobytes())
|
||||
|
||||
|
||||
@functools.lru_cache()
|
||||
def _get_global_gloo_group():
|
||||
"""Return a process group based on gloo backend, containing all the ranks
|
||||
The result is cached."""
|
||||
if dist.get_backend() == 'nccl':
|
||||
return dist.new_group(backend='gloo')
|
||||
else:
|
||||
return dist.group.WORLD
|
||||
|
||||
|
||||
def all_reduce_dict(py_dict, op='sum', group=None, to_float=True):
|
||||
"""Apply all reduce function for python dict object.
|
||||
|
||||
The code is modified from https://github.com/Megvii-
|
||||
BaseDetection/YOLOX/blob/main/yolox/utils/allreduce_norm.py.
|
||||
|
||||
NOTE: make sure that py_dict in different ranks has the same keys and
|
||||
the values should be in the same shape.
|
||||
|
||||
Args:
|
||||
py_dict (dict): Dict to be applied all reduce op.
|
||||
op (str): Operator, could be 'sum' or 'mean'. Default: 'sum'
|
||||
group (:obj:`torch.distributed.group`, optional): Distributed group,
|
||||
Default: None.
|
||||
to_float (bool): Whether to convert all values of dict to float.
|
||||
Default: True.
|
||||
|
||||
Returns:
|
||||
OrderedDict: reduced python dict object.
|
||||
"""
|
||||
_, world_size = get_dist_info()
|
||||
if world_size == 1:
|
||||
return py_dict
|
||||
if group is None:
|
||||
# TODO: May try not to use gloo in the future
|
||||
group = _get_global_gloo_group()
|
||||
if dist.get_world_size(group) == 1:
|
||||
return py_dict
|
||||
|
||||
# all reduce logic across different devices.
|
||||
py_key = list(py_dict.keys())
|
||||
py_key_tensor = obj2tensor(py_key)
|
||||
dist.broadcast(py_key_tensor, src=0)
|
||||
py_key = tensor2obj(py_key_tensor)
|
||||
|
||||
tensor_shapes = [py_dict[k].shape for k in py_key]
|
||||
tensor_numels = [py_dict[k].numel() for k in py_key]
|
||||
|
||||
if to_float:
|
||||
flatten_tensor = torch.cat(
|
||||
[py_dict[k].flatten().float() for k in py_key])
|
||||
else:
|
||||
flatten_tensor = torch.cat([py_dict[k].flatten() for k in py_key])
|
||||
|
||||
dist.all_reduce(flatten_tensor, op=dist.ReduceOp.SUM)
|
||||
if op == 'mean':
|
||||
flatten_tensor /= world_size
|
||||
|
||||
split_tensors = [
|
||||
x.reshape(shape) for x, shape in zip(
|
||||
torch.split(flatten_tensor, tensor_numels), tensor_shapes)
|
||||
]
|
||||
return OrderedDict({k: v for k, v in zip(py_key, split_tensors)})
|
||||
40
segmentation/mmseg_custom/core/utils/misc.py
Normal file
40
segmentation/mmseg_custom/core/utils/misc.py
Normal file
@@ -0,0 +1,40 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
def multi_apply(func, *args, **kwargs):
|
||||
"""Apply function to a list of arguments.
|
||||
|
||||
Note:
|
||||
This function applies the ``func`` to multiple inputs and
|
||||
map the multiple outputs of the ``func`` into different
|
||||
list. Each list contains the same type of outputs corresponding
|
||||
to different inputs.
|
||||
|
||||
Args:
|
||||
func (Function): A function that will be applied to a list of
|
||||
arguments
|
||||
|
||||
Returns:
|
||||
tuple(list): A tuple containing multiple list, each list contains \
|
||||
a kind of returned results by the function
|
||||
"""
|
||||
pfunc = partial(func, **kwargs) if kwargs else func
|
||||
map_results = map(pfunc, *args)
|
||||
return tuple(map(list, zip(*map_results)))
|
||||
|
||||
|
||||
def add_prefix(inputs, prefix):
|
||||
"""Add prefix for dict.
|
||||
|
||||
Args:
|
||||
inputs (dict): The input dict with str keys.
|
||||
prefix (str): The prefix to add.
|
||||
|
||||
Returns:
|
||||
|
||||
dict: The dict with keys updated with ``prefix``.
|
||||
"""
|
||||
|
||||
outputs = dict()
|
||||
for name, value in inputs.items():
|
||||
outputs[f'{prefix}.{name}'] = value
|
||||
|
||||
return outputs
|
||||
10
segmentation/mmseg_custom/datasets/__init__.py
Normal file
10
segmentation/mmseg_custom/datasets/__init__.py
Normal file
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .mapillary import MapillaryDataset # noqa: F401,F403
|
||||
from .nyu_depth_v2 import NYUDepthV2Dataset # noqa: F401,F403
|
||||
from .pipelines import * # noqa: F401,F403
|
||||
from .dataset_wrappers import ConcatDataset
|
||||
|
||||
|
||||
__all__ = [
|
||||
'MapillaryDataset', 'NYUDepthV2Dataset', 'ConcatDataset'
|
||||
]
|
||||
155
segmentation/mmseg_custom/datasets/dataset_wrappers.py
Normal file
155
segmentation/mmseg_custom/datasets/dataset_wrappers.py
Normal file
@@ -0,0 +1,155 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import bisect
|
||||
from itertools import chain
|
||||
|
||||
import mmcv
|
||||
import numpy as np
|
||||
from mmcv.utils import build_from_cfg, print_log
|
||||
from torch.utils.data.dataset import ConcatDataset as _ConcatDataset
|
||||
|
||||
from mmseg.datasets.builder import DATASETS
|
||||
|
||||
|
||||
@DATASETS.register_module(force=True)
|
||||
class ConcatDataset(_ConcatDataset):
|
||||
"""A wrapper of concatenated dataset.
|
||||
|
||||
Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but
|
||||
support evaluation and formatting results
|
||||
|
||||
Args:
|
||||
datasets (list[:obj:`Dataset`]): A list of datasets.
|
||||
separate_eval (bool): Whether to evaluate the concatenated
|
||||
dataset results separately, Defaults to True.
|
||||
"""
|
||||
|
||||
def __init__(self, datasets, separate_eval=True):
|
||||
super(ConcatDataset, self).__init__(datasets)
|
||||
self.CLASSES = datasets[0].CLASSES
|
||||
self.PALETTE = datasets[0].PALETTE
|
||||
self.separate_eval = separate_eval
|
||||
assert separate_eval in [True, False], \
|
||||
f'separate_eval can only be True or False,' \
|
||||
f'but get {separate_eval}'
|
||||
|
||||
def evaluate(self, results, logger=None, **kwargs):
|
||||
"""Evaluate the results.
|
||||
|
||||
Args:
|
||||
results (list[tuple[torch.Tensor]] | list[str]]): per image
|
||||
pre_eval results or predict segmentation map for
|
||||
computing evaluation metric.
|
||||
logger (logging.Logger | str | None): Logger used for printing
|
||||
related information during evaluation. Default: None.
|
||||
|
||||
Returns:
|
||||
dict[str: float]: evaluate results of the total dataset
|
||||
or each separate
|
||||
dataset if `self.separate_eval=True`.
|
||||
"""
|
||||
assert len(results) == self.cumulative_sizes[-1], \
|
||||
('Dataset and results have different sizes: '
|
||||
f'{self.cumulative_sizes[-1]} v.s. {len(results)}')
|
||||
|
||||
# Check whether all the datasets support evaluation
|
||||
for dataset in self.datasets:
|
||||
assert hasattr(dataset, 'evaluate'), \
|
||||
f'{type(dataset)} does not implement evaluate function'
|
||||
|
||||
if self.separate_eval:
|
||||
dataset_idx = -1
|
||||
total_eval_results = dict()
|
||||
for size, dataset in zip(self.cumulative_sizes, self.datasets):
|
||||
start_idx = 0 if dataset_idx == -1 else \
|
||||
self.cumulative_sizes[dataset_idx]
|
||||
end_idx = self.cumulative_sizes[dataset_idx + 1]
|
||||
|
||||
results_per_dataset = results[start_idx:end_idx]
|
||||
print_log(
|
||||
f'\nEvaluateing {dataset.img_dir} with '
|
||||
f'{len(results_per_dataset)} images now',
|
||||
logger=logger)
|
||||
|
||||
eval_results_per_dataset = dataset.evaluate(
|
||||
results_per_dataset, logger=logger, **kwargs)
|
||||
dataset_idx += 1
|
||||
for k, v in eval_results_per_dataset.items():
|
||||
total_eval_results.update({f'{dataset_idx}_{k}': v})
|
||||
|
||||
return total_eval_results
|
||||
|
||||
if len(set([type(ds) for ds in self.datasets])) != 1:
|
||||
raise NotImplementedError(
|
||||
'All the datasets should have same types when '
|
||||
'self.separate_eval=False')
|
||||
else:
|
||||
if mmcv.is_list_of(results, np.ndarray) or mmcv.is_list_of(
|
||||
results, str):
|
||||
# merge the generators of gt_seg_maps
|
||||
gt_seg_maps = chain(
|
||||
*[dataset.get_gt_seg_maps() for dataset in self.datasets])
|
||||
else:
|
||||
# if the results are `pre_eval` results,
|
||||
# we do not need gt_seg_maps to evaluate
|
||||
gt_seg_maps = None
|
||||
eval_results = self.datasets[0].evaluate(
|
||||
results, gt_seg_maps=gt_seg_maps, logger=logger, **kwargs)
|
||||
return eval_results
|
||||
|
||||
def get_dataset_idx_and_sample_idx(self, indice):
|
||||
"""Return dataset and sample index when given an indice of
|
||||
ConcatDataset.
|
||||
|
||||
Args:
|
||||
indice (int): indice of sample in ConcatDataset
|
||||
|
||||
Returns:
|
||||
int: the index of sub dataset the sample belong to
|
||||
int: the index of sample in its corresponding subset
|
||||
"""
|
||||
if indice < 0:
|
||||
if -indice > len(self):
|
||||
raise ValueError(
|
||||
'absolute value of index should not exceed dataset length')
|
||||
indice = len(self) + indice
|
||||
dataset_idx = bisect.bisect_right(self.cumulative_sizes, indice)
|
||||
if dataset_idx == 0:
|
||||
sample_idx = indice
|
||||
else:
|
||||
sample_idx = indice - self.cumulative_sizes[dataset_idx - 1]
|
||||
return dataset_idx, sample_idx
|
||||
|
||||
def format_results(self, results, imgfile_prefix, indices=None, **kwargs):
|
||||
"""format result for every sample of ConcatDataset."""
|
||||
if indices is None:
|
||||
indices = list(range(len(self)))
|
||||
|
||||
assert isinstance(results, list), 'results must be a list.'
|
||||
assert isinstance(indices, list), 'indices must be a list.'
|
||||
|
||||
ret_res = []
|
||||
for i, indice in enumerate(indices):
|
||||
dataset_idx, sample_idx = self.get_dataset_idx_and_sample_idx(
|
||||
indice)
|
||||
res = self.datasets[dataset_idx].format_results(
|
||||
[results[i]],
|
||||
imgfile_prefix + f'/{dataset_idx}',
|
||||
indices=[sample_idx],
|
||||
**kwargs)
|
||||
ret_res.append(res)
|
||||
return sum(ret_res, [])
|
||||
|
||||
def pre_eval(self, preds, indices):
|
||||
"""do pre eval for every sample of ConcatDataset."""
|
||||
# In order to compat with batch inference
|
||||
if not isinstance(indices, list):
|
||||
indices = [indices]
|
||||
if not isinstance(preds, list):
|
||||
preds = [preds]
|
||||
ret_res = []
|
||||
for i, indice in enumerate(indices):
|
||||
dataset_idx, sample_idx = self.get_dataset_idx_and_sample_idx(
|
||||
indice)
|
||||
res = self.datasets[dataset_idx].pre_eval(preds[i], sample_idx)
|
||||
ret_res.append(res)
|
||||
return sum(ret_res, [])
|
||||
48
segmentation/mmseg_custom/datasets/mapillary.py
Normal file
48
segmentation/mmseg_custom/datasets/mapillary.py
Normal file
@@ -0,0 +1,48 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
from mmseg.datasets.builder import DATASETS
|
||||
from mmseg.datasets.custom import CustomDataset
|
||||
|
||||
|
||||
@DATASETS.register_module()
|
||||
class MapillaryDataset(CustomDataset):
|
||||
"""Mapillary dataset.
|
||||
"""
|
||||
CLASSES = ('Bird', 'Ground Animal', 'Curb', 'Fence', 'Guard Rail', 'Barrier',
|
||||
'Wall', 'Bike Lane', 'Crosswalk - Plain', 'Curb Cut', 'Parking', 'Pedestrian Area',
|
||||
'Rail Track', 'Road', 'Service Lane', 'Sidewalk', 'Bridge', 'Building', 'Tunnel',
|
||||
'Person', 'Bicyclist', 'Motorcyclist', 'Other Rider', 'Lane Marking - Crosswalk',
|
||||
'Lane Marking - General', 'Mountain', 'Sand', 'Sky', 'Snow', 'Terrain', 'Vegetation',
|
||||
'Water', 'Banner', 'Bench', 'Bike Rack', 'Billboard', 'Catch Basin', 'CCTV Camera',
|
||||
'Fire Hydrant', 'Junction Box', 'Mailbox', 'Manhole', 'Phone Booth', 'Pothole',
|
||||
'Street Light', 'Pole', 'Traffic Sign Frame', 'Utility Pole', 'Traffic Light',
|
||||
'Traffic Sign (Back)', 'Traffic Sign (Front)', 'Trash Can', 'Bicycle', 'Boat',
|
||||
'Bus', 'Car', 'Caravan', 'Motorcycle', 'On Rails', 'Other Vehicle', 'Trailer',
|
||||
'Truck', 'Wheeled Slow', 'Car Mount', 'Ego Vehicle', 'Unlabeled')
|
||||
|
||||
PALETTE = [[165, 42, 42], [0, 192, 0], [196, 196, 196], [190, 153, 153],
|
||||
[180, 165, 180], [90, 120, 150], [102, 102, 156], [128, 64, 255],
|
||||
[140, 140, 200], [170, 170, 170], [250, 170, 160], [96, 96, 96],
|
||||
[230, 150, 140], [128, 64, 128], [110, 110, 110], [244, 35, 232],
|
||||
[150, 100, 100], [70, 70, 70], [150, 120, 90], [220, 20, 60],
|
||||
[255, 0, 0], [255, 0, 100], [255, 0, 200], [200, 128, 128],
|
||||
[255, 255, 255], [64, 170, 64], [230, 160, 50], [70, 130, 180],
|
||||
[190, 255, 255], [152, 251, 152], [107, 142, 35], [0, 170, 30],
|
||||
[255, 255, 128], [250, 0, 30], [100, 140, 180], [220, 220, 220],
|
||||
[220, 128, 128], [222, 40, 40], [100, 170, 30], [40, 40, 40],
|
||||
[33, 33, 33], [100, 128, 160], [142, 0, 0], [70, 100, 150],
|
||||
[210, 170, 100], [153, 153, 153], [128, 128, 128], [0, 0, 80],
|
||||
[250, 170, 30], [192, 192, 192], [220, 220, 0], [140, 140, 20],
|
||||
[119, 11, 32], [150, 0, 255], [0, 60, 100], [0, 0, 142], [0, 0, 90],
|
||||
[0, 0, 230], [0, 80, 100], [128, 64, 64], [0, 0, 110], [0, 0, 70],
|
||||
[0, 0, 192], [32, 32, 32], [120, 10, 10], [0, 0, 0]]
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super(MapillaryDataset, self).__init__(
|
||||
img_suffix='.jpg',
|
||||
seg_map_suffix='.png',
|
||||
reduce_zero_label=False,
|
||||
**kwargs)
|
||||
43
segmentation/mmseg_custom/datasets/nyu_depth_v2.py
Normal file
43
segmentation/mmseg_custom/datasets/nyu_depth_v2.py
Normal file
@@ -0,0 +1,43 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
from mmseg.datasets.builder import DATASETS
|
||||
from mmseg.datasets.custom import CustomDataset
|
||||
|
||||
|
||||
@DATASETS.register_module()
|
||||
class NYUDepthV2Dataset(CustomDataset):
|
||||
"""NYU Depth V2 dataset.
|
||||
"""
|
||||
|
||||
CLASSES = ('wall', 'floor', 'cabinet', 'bed', 'chair',
|
||||
'sofa', 'table', 'door', 'window', 'bookshelf',
|
||||
'picture', 'counter', 'blinds', 'desk', 'shelves',
|
||||
'curtain', 'dresser', 'pillow', 'mirror', 'floor mat',
|
||||
'clothes', 'ceiling', 'books', 'refridgerator', 'television',
|
||||
'paper', 'towel', 'shower curtain', 'box', 'whiteboard',
|
||||
'person', 'night stand', 'toilet', 'sink', 'lamp',
|
||||
'bathtub', 'bag', 'otherstructure', 'otherfurniture', 'otherprop')
|
||||
|
||||
|
||||
PALETTE = [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
|
||||
[4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
|
||||
[230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
|
||||
[150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
|
||||
[143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
|
||||
[0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
|
||||
[255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
|
||||
[255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
|
||||
[255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
|
||||
[224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],]
|
||||
|
||||
def __init__(self, split, **kwargs):
|
||||
super(NYUDepthV2Dataset, self).__init__(
|
||||
img_suffix='.png',
|
||||
seg_map_suffix='.png',
|
||||
split=split,
|
||||
reduce_zero_label=True,
|
||||
**kwargs)
|
||||
|
||||
8
segmentation/mmseg_custom/datasets/pipelines/__init__.py
Normal file
8
segmentation/mmseg_custom/datasets/pipelines/__init__.py
Normal file
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .formatting import DefaultFormatBundle, ToMask
|
||||
from .transform import MapillaryHack, PadShortSide, SETR_Resize
|
||||
|
||||
__all__ = [
|
||||
'DefaultFormatBundle', 'ToMask', 'SETR_Resize',
|
||||
'PadShortSide', 'MapillaryHack'
|
||||
]
|
||||
82
segmentation/mmseg_custom/datasets/pipelines/formatting.py
Normal file
82
segmentation/mmseg_custom/datasets/pipelines/formatting.py
Normal file
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import numpy as np
|
||||
from mmcv.parallel import DataContainer as DC
|
||||
from mmseg.datasets.builder import PIPELINES
|
||||
from mmseg.datasets.pipelines.formatting import to_tensor
|
||||
|
||||
|
||||
@PIPELINES.register_module(force=True)
|
||||
class DefaultFormatBundle(object):
|
||||
"""Default formatting bundle.
|
||||
|
||||
It simplifies the pipeline of formatting common fields, including "img"
|
||||
and "gt_semantic_seg". These fields are formatted as follows.
|
||||
|
||||
- img: (1)transpose, (2)to tensor, (3)to DataContainer (stack=True)
|
||||
- gt_semantic_seg: (1)unsqueeze dim-0 (2)to tensor,
|
||||
(3)to DataContainer (stack=True)
|
||||
"""
|
||||
def __call__(self, results):
|
||||
"""Call function to transform and format common fields in results.
|
||||
|
||||
Args:
|
||||
results (dict): Result dict contains the data to convert.
|
||||
|
||||
Returns:
|
||||
dict: The result dict contains the data that is formatted with
|
||||
default bundle.
|
||||
"""
|
||||
|
||||
if 'img' in results:
|
||||
img = results['img']
|
||||
if len(img.shape) < 3:
|
||||
img = np.expand_dims(img, -1)
|
||||
img = np.ascontiguousarray(img.transpose(2, 0, 1))
|
||||
results['img'] = DC(to_tensor(img), stack=True)
|
||||
if 'gt_semantic_seg' in results:
|
||||
# convert to long
|
||||
results['gt_semantic_seg'] = DC(to_tensor(
|
||||
results['gt_semantic_seg'][None, ...].astype(np.int64)),
|
||||
stack=True)
|
||||
if 'gt_masks' in results:
|
||||
results['gt_masks'] = DC(to_tensor(results['gt_masks']))
|
||||
if 'gt_labels' in results:
|
||||
results['gt_labels'] = DC(to_tensor(results['gt_labels']))
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__
|
||||
|
||||
|
||||
@PIPELINES.register_module()
|
||||
class ToMask(object):
|
||||
"""Transfer gt_semantic_seg to binary mask and generate gt_labels."""
|
||||
def __init__(self, ignore_index=255):
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
def __call__(self, results):
|
||||
gt_semantic_seg = results['gt_semantic_seg']
|
||||
gt_labels = np.unique(gt_semantic_seg)
|
||||
# remove ignored region
|
||||
gt_labels = gt_labels[gt_labels != self.ignore_index]
|
||||
|
||||
gt_masks = []
|
||||
for class_id in gt_labels:
|
||||
gt_masks.append(gt_semantic_seg == class_id)
|
||||
|
||||
if len(gt_masks) == 0:
|
||||
# Some image does not have annotation (all ignored)
|
||||
gt_masks = np.empty((0, ) + results['pad_shape'][:-1], dtype=np.int64)
|
||||
gt_labels = np.empty((0, ), dtype=np.int64)
|
||||
else:
|
||||
gt_masks = np.asarray(gt_masks, dtype=np.int64)
|
||||
gt_labels = np.asarray(gt_labels, dtype=np.int64)
|
||||
|
||||
results['gt_labels'] = gt_labels
|
||||
results['gt_masks'] = gt_masks
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + \
|
||||
f'(ignore_index={self.ignore_index})'
|
||||
350
segmentation/mmseg_custom/datasets/pipelines/transform.py
Normal file
350
segmentation/mmseg_custom/datasets/pipelines/transform.py
Normal file
@@ -0,0 +1,350 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import mmcv
|
||||
import numpy as np
|
||||
from mmseg.datasets.builder import PIPELINES
|
||||
|
||||
|
||||
@PIPELINES.register_module()
|
||||
class SETR_Resize(object):
|
||||
"""Resize images & seg.
|
||||
|
||||
This transform resizes the input image to some scale. If the input dict
|
||||
contains the key "scale", then the scale in the input dict is used,
|
||||
otherwise the specified scale in the init method is used.
|
||||
|
||||
``img_scale`` can either be a tuple (single-scale) or a list of tuple
|
||||
(multi-scale). There are 3 multiscale modes:
|
||||
|
||||
- ``ratio_range is not None``: randomly sample a ratio from the ratio range
|
||||
and multiply it with the image scale.
|
||||
|
||||
- ``ratio_range is None and multiscale_mode == "range"``: randomly sample a
|
||||
scale from the a range.
|
||||
|
||||
- ``ratio_range is None and multiscale_mode == "value"``: randomly sample a
|
||||
scale from multiple scales.
|
||||
|
||||
Args:
|
||||
img_scale (tuple or list[tuple]): Images scales for resizing.
|
||||
multiscale_mode (str): Either "range" or "value".
|
||||
ratio_range (tuple[float]): (min_ratio, max_ratio)
|
||||
keep_ratio (bool): Whether to keep the aspect ratio when resizing the
|
||||
image.
|
||||
"""
|
||||
def __init__(self,
|
||||
img_scale=None,
|
||||
multiscale_mode='range',
|
||||
ratio_range=None,
|
||||
keep_ratio=True,
|
||||
crop_size=None,
|
||||
setr_multi_scale=False):
|
||||
|
||||
if img_scale is None:
|
||||
self.img_scale = None
|
||||
else:
|
||||
if isinstance(img_scale, list):
|
||||
self.img_scale = img_scale
|
||||
else:
|
||||
self.img_scale = [img_scale]
|
||||
# assert mmcv.is_list_of(self.img_scale, tuple)
|
||||
|
||||
if ratio_range is not None:
|
||||
# mode 1: given a scale and a range of image ratio
|
||||
assert len(self.img_scale) == 1
|
||||
else:
|
||||
# mode 2: given multiple scales or a range of scales
|
||||
assert multiscale_mode in ['value', 'range']
|
||||
|
||||
self.multiscale_mode = multiscale_mode
|
||||
self.ratio_range = ratio_range
|
||||
self.keep_ratio = keep_ratio
|
||||
self.crop_size = crop_size
|
||||
self.setr_multi_scale = setr_multi_scale
|
||||
|
||||
@staticmethod
|
||||
def random_select(img_scales):
|
||||
"""Randomly select an img_scale from given candidates.
|
||||
|
||||
Args:
|
||||
img_scales (list[tuple]): Images scales for selection.
|
||||
|
||||
Returns:
|
||||
(tuple, int): Returns a tuple ``(img_scale, scale_dix)``,
|
||||
where ``img_scale`` is the selected image scale and
|
||||
``scale_idx`` is the selected index in the given candidates.
|
||||
"""
|
||||
|
||||
assert mmcv.is_list_of(img_scales, tuple)
|
||||
scale_idx = np.random.randint(len(img_scales))
|
||||
img_scale = img_scales[scale_idx]
|
||||
return img_scale, scale_idx
|
||||
|
||||
@staticmethod
|
||||
def random_sample(img_scales):
|
||||
"""Randomly sample an img_scale when ``multiscale_mode=='range'``.
|
||||
|
||||
Args:
|
||||
img_scales (list[tuple]): Images scale range for sampling.
|
||||
There must be two tuples in img_scales, which specify the lower
|
||||
and uper bound of image scales.
|
||||
|
||||
Returns:
|
||||
(tuple, None): Returns a tuple ``(img_scale, None)``, where
|
||||
``img_scale`` is sampled scale and None is just a placeholder
|
||||
to be consistent with :func:`random_select`.
|
||||
"""
|
||||
|
||||
assert mmcv.is_list_of(img_scales, tuple) and len(img_scales) == 2
|
||||
img_scale_long = [max(s) for s in img_scales]
|
||||
img_scale_short = [min(s) for s in img_scales]
|
||||
long_edge = np.random.randint(
|
||||
min(img_scale_long),
|
||||
max(img_scale_long) + 1)
|
||||
short_edge = np.random.randint(
|
||||
min(img_scale_short),
|
||||
max(img_scale_short) + 1)
|
||||
img_scale = (long_edge, short_edge)
|
||||
return img_scale, None
|
||||
|
||||
@staticmethod
|
||||
def random_sample_ratio(img_scale, ratio_range):
|
||||
"""Randomly sample an img_scale when ``ratio_range`` is specified.
|
||||
|
||||
A ratio will be randomly sampled from the range specified by
|
||||
``ratio_range``. Then it would be multiplied with ``img_scale`` to
|
||||
generate sampled scale.
|
||||
|
||||
Args:
|
||||
img_scale (tuple): Images scale base to multiply with ratio.
|
||||
ratio_range (tuple[float]): The minimum and maximum ratio to scale
|
||||
the ``img_scale``.
|
||||
|
||||
Returns:
|
||||
(tuple, None): Returns a tuple ``(scale, None)``, where
|
||||
``scale`` is sampled ratio multiplied with ``img_scale`` and
|
||||
None is just a placeholder to be consistent with
|
||||
:func:`random_select`.
|
||||
"""
|
||||
|
||||
assert isinstance(img_scale, tuple) and len(img_scale) == 2
|
||||
min_ratio, max_ratio = ratio_range
|
||||
assert min_ratio <= max_ratio
|
||||
ratio = np.random.random_sample() * (max_ratio - min_ratio) + min_ratio
|
||||
scale = int(img_scale[0] * ratio), int(img_scale[1] * ratio)
|
||||
return scale, None
|
||||
|
||||
def _random_scale(self, results):
|
||||
"""Randomly sample an img_scale according to ``ratio_range`` and
|
||||
``multiscale_mode``.
|
||||
|
||||
If ``ratio_range`` is specified, a ratio will be sampled and be
|
||||
multiplied with ``img_scale``.
|
||||
If multiple scales are specified by ``img_scale``, a scale will be
|
||||
sampled according to ``multiscale_mode``.
|
||||
Otherwise, single scale will be used.
|
||||
|
||||
Args:
|
||||
results (dict): Result dict from :obj:`dataset`.
|
||||
|
||||
Returns:
|
||||
dict: Two new keys 'scale` and 'scale_idx` are added into
|
||||
``results``, which would be used by subsequent pipelines.
|
||||
"""
|
||||
|
||||
if self.ratio_range is not None:
|
||||
scale, scale_idx = self.random_sample_ratio(
|
||||
self.img_scale[0], self.ratio_range)
|
||||
elif len(self.img_scale) == 1:
|
||||
scale, scale_idx = self.img_scale[0], 0
|
||||
elif self.multiscale_mode == 'range':
|
||||
scale, scale_idx = self.random_sample(self.img_scale)
|
||||
elif self.multiscale_mode == 'value':
|
||||
scale, scale_idx = self.random_select(self.img_scale)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
results['scale'] = scale
|
||||
results['scale_idx'] = scale_idx
|
||||
|
||||
def _resize_img(self, results):
|
||||
"""Resize images with ``results['scale']``."""
|
||||
|
||||
if self.keep_ratio:
|
||||
if self.setr_multi_scale:
|
||||
if min(results['scale']) < self.crop_size[0]:
|
||||
new_short = self.crop_size[0]
|
||||
else:
|
||||
new_short = min(results['scale'])
|
||||
|
||||
h, w = results['img'].shape[:2]
|
||||
if h > w:
|
||||
new_h, new_w = new_short * h / w, new_short
|
||||
else:
|
||||
new_h, new_w = new_short, new_short * w / h
|
||||
results['scale'] = (new_h, new_w)
|
||||
|
||||
img, scale_factor = mmcv.imrescale(results['img'],
|
||||
results['scale'],
|
||||
return_scale=True)
|
||||
# the w_scale and h_scale has minor difference
|
||||
# a real fix should be done in the mmcv.imrescale in the future
|
||||
new_h, new_w = img.shape[:2]
|
||||
h, w = results['img'].shape[:2]
|
||||
w_scale = new_w / w
|
||||
h_scale = new_h / h
|
||||
else:
|
||||
img, w_scale, h_scale = mmcv.imresize(results['img'],
|
||||
results['scale'],
|
||||
return_scale=True)
|
||||
scale_factor = np.array([w_scale, h_scale, w_scale, h_scale],
|
||||
dtype=np.float32)
|
||||
results['img'] = img
|
||||
results['img_shape'] = img.shape
|
||||
results['pad_shape'] = img.shape # in case that there is no padding
|
||||
results['scale_factor'] = scale_factor
|
||||
results['keep_ratio'] = self.keep_ratio
|
||||
|
||||
def _resize_seg(self, results):
|
||||
"""Resize semantic segmentation map with ``results['scale']``."""
|
||||
for key in results.get('seg_fields', []):
|
||||
if self.keep_ratio:
|
||||
gt_seg = mmcv.imrescale(results[key],
|
||||
results['scale'],
|
||||
interpolation='nearest')
|
||||
else:
|
||||
gt_seg = mmcv.imresize(results[key],
|
||||
results['scale'],
|
||||
interpolation='nearest')
|
||||
results['gt_semantic_seg'] = gt_seg
|
||||
|
||||
def __call__(self, results):
|
||||
"""Call function to resize images, bounding boxes, masks, semantic
|
||||
segmentation map.
|
||||
|
||||
Args:
|
||||
results (dict): Result dict from loading pipeline.
|
||||
|
||||
Returns:
|
||||
dict: Resized results, 'img_shape', 'pad_shape', 'scale_factor',
|
||||
'keep_ratio' keys are added into result dict.
|
||||
"""
|
||||
|
||||
if 'scale' not in results:
|
||||
self._random_scale(results)
|
||||
self._resize_img(results)
|
||||
self._resize_seg(results)
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
repr_str += (f'(img_scale={self.img_scale}, '
|
||||
f'multiscale_mode={self.multiscale_mode}, '
|
||||
f'ratio_range={self.ratio_range}, '
|
||||
f'keep_ratio={self.keep_ratio})')
|
||||
return repr_str
|
||||
|
||||
|
||||
@PIPELINES.register_module()
|
||||
class PadShortSide(object):
|
||||
"""Pad the image & mask.
|
||||
|
||||
Pad to the minimum size that is equal or larger than a number.
|
||||
Added keys are "pad_shape", "pad_fixed_size",
|
||||
|
||||
Args:
|
||||
size (int, optional): Fixed padding size.
|
||||
pad_val (float, optional): Padding value. Default: 0.
|
||||
seg_pad_val (float, optional): Padding value of segmentation map.
|
||||
Default: 255.
|
||||
"""
|
||||
def __init__(self, size=None, pad_val=0, seg_pad_val=255):
|
||||
self.size = size
|
||||
self.pad_val = pad_val
|
||||
self.seg_pad_val = seg_pad_val
|
||||
# only one of size and size_divisor should be valid
|
||||
assert size is not None
|
||||
|
||||
def _pad_img(self, results):
|
||||
"""Pad images according to ``self.size``."""
|
||||
h, w = results['img'].shape[:2]
|
||||
new_h = max(h, self.size)
|
||||
new_w = max(w, self.size)
|
||||
padded_img = mmcv.impad(results['img'],
|
||||
shape=(new_h, new_w),
|
||||
pad_val=self.pad_val)
|
||||
|
||||
results['img'] = padded_img
|
||||
results['pad_shape'] = padded_img.shape
|
||||
# results['unpad_shape'] = (h, w)
|
||||
|
||||
def _pad_seg(self, results):
|
||||
"""Pad masks according to ``results['pad_shape']``."""
|
||||
for key in results.get('seg_fields', []):
|
||||
results[key] = mmcv.impad(results[key],
|
||||
shape=results['pad_shape'][:2],
|
||||
pad_val=self.seg_pad_val)
|
||||
|
||||
def __call__(self, results):
|
||||
"""Call function to pad images, masks, semantic segmentation maps.
|
||||
|
||||
Args:
|
||||
results (dict): Result dict from loading pipeline.
|
||||
|
||||
Returns:
|
||||
dict: Updated result dict.
|
||||
"""
|
||||
h, w = results['img'].shape[:2]
|
||||
if h >= self.size and w >= self.size: # 短边比窗口大，跳过
|
||||
pass
|
||||
else:
|
||||
self._pad_img(results)
|
||||
self._pad_seg(results)
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
repr_str += f'(size={self.size}, pad_val={self.pad_val})'
|
||||
return repr_str
|
||||
|
||||
|
||||
@PIPELINES.register_module()
|
||||
class MapillaryHack(object):
|
||||
"""map MV 65 class to 19 class like Cityscapes."""
|
||||
def __init__(self):
|
||||
self.map = [[13, 24, 41], [2, 15], [17], [6], [3],
|
||||
[45, 47], [48], [50], [30], [29], [27], [19], [20, 21, 22],
|
||||
[55], [61], [54], [58], [57], [52]]
|
||||
|
||||
self.others = [i for i in range(66)]
|
||||
for i in self.map:
|
||||
for j in i:
|
||||
if j in self.others:
|
||||
self.others.remove(j)
|
||||
|
||||
def __call__(self, results):
|
||||
"""Call function to process the image with gamma correction.
|
||||
|
||||
Args:
|
||||
results (dict): Result dict from loading pipeline.
|
||||
|
||||
Returns:
|
||||
dict: Processed results.
|
||||
"""
|
||||
gt_map = results['gt_semantic_seg']
|
||||
# others -> 255
|
||||
new_gt_map = np.zeros_like(gt_map)
|
||||
|
||||
for value in self.others:
|
||||
new_gt_map[gt_map == value] = 255
|
||||
|
||||
for index, map in enumerate(self.map):
|
||||
for value in map:
|
||||
new_gt_map[gt_map == value] = index
|
||||
|
||||
results['gt_semantic_seg'] = new_gt_map
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
12
segmentation/mmseg_custom/models/__init__.py
Normal file
12
segmentation/mmseg_custom/models/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
from .backbones import * # noqa: F401,F403
|
||||
from .decode_heads import * # noqa: F401,F403
|
||||
from .losses import * # noqa: F401,F403
|
||||
from .plugins import * # noqa: F401,F403
|
||||
from .segmentors import * # noqa: F401,F403
|
||||
from .utils import * # noqa: F401,F403
|
||||
9
segmentation/mmseg_custom/models/backbones/__init__.py
Normal file
9
segmentation/mmseg_custom/models/backbones/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
# --------------------------------------------------------
|
||||
# FlashInternImage
|
||||
# Copyright (c) 2023 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
from .flash_intern_image import FlashInternImage
|
||||
|
||||
__all__ = ['FlashInternImage']
|
||||
763
segmentation/mmseg_custom/models/backbones/flash_intern_image.py
Normal file
763
segmentation/mmseg_custom/models/backbones/flash_intern_image.py
Normal file
@@ -0,0 +1,763 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from collections import OrderedDict
|
||||
import torch.utils.checkpoint as checkpoint
|
||||
from timm.models.layers import trunc_normal_, DropPath
|
||||
from mmcv.runner import _load_checkpoint
|
||||
from mmcv.cnn import constant_init, trunc_normal_init
|
||||
from mmseg.utils import get_root_logger
|
||||
from mmseg.models.builder import BACKBONES
|
||||
import torch.nn.functional as F
|
||||
import DCNv4
|
||||
|
||||
|
||||
class to_channels_first(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return x.permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
class to_channels_last(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return x.permute(0, 2, 3, 1)
|
||||
|
||||
|
||||
def build_norm_layer(dim,
|
||||
norm_layer,
|
||||
in_format='channels_last',
|
||||
out_format='channels_last',
|
||||
eps=1e-6):
|
||||
layers = []
|
||||
if norm_layer == 'BN':
|
||||
if in_format == 'channels_last':
|
||||
layers.append(to_channels_first())
|
||||
layers.append(nn.BatchNorm2d(dim))
|
||||
if out_format == 'channels_last':
|
||||
layers.append(to_channels_last())
|
||||
elif norm_layer == 'LN':
|
||||
if in_format == 'channels_first':
|
||||
layers.append(to_channels_last())
|
||||
layers.append(nn.LayerNorm(dim, eps=eps))
|
||||
if out_format == 'channels_first':
|
||||
layers.append(to_channels_first())
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'build_norm_layer does not support {norm_layer}')
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
def build_act_layer(act_layer):
|
||||
if act_layer == 'ReLU':
|
||||
return nn.ReLU(inplace=True)
|
||||
elif act_layer == 'SiLU':
|
||||
return nn.SiLU(inplace=True)
|
||||
elif act_layer == 'GELU':
|
||||
return nn.GELU()
|
||||
|
||||
raise NotImplementedError(f'build_act_layer does not support {act_layer}')
|
||||
|
||||
|
||||
class CrossAttention(nn.Module):
|
||||
r""" Cross Attention Module
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
num_heads (int): Number of attention heads. Default: 8
|
||||
qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.
|
||||
Default: False.
|
||||
qk_scale (float | None, optional): Override default qk scale of
|
||||
head_dim ** -0.5 if set. Default: None.
|
||||
attn_drop (float, optional): Dropout ratio of attention weight.
|
||||
Default: 0.0
|
||||
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
|
||||
attn_head_dim (int, optional): Dimension of attention head.
|
||||
out_dim (int, optional): Dimension of output.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
num_heads=8,
|
||||
qkv_bias=False,
|
||||
qk_scale=None,
|
||||
attn_drop=0.,
|
||||
proj_drop=0.,
|
||||
attn_head_dim=None,
|
||||
out_dim=None):
|
||||
super().__init__()
|
||||
if out_dim is None:
|
||||
out_dim = dim
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
if attn_head_dim is not None:
|
||||
head_dim = attn_head_dim
|
||||
all_head_dim = head_dim * self.num_heads
|
||||
self.scale = qk_scale or head_dim ** -0.5
|
||||
assert all_head_dim == dim
|
||||
|
||||
self.q = nn.Linear(dim, all_head_dim, bias=False)
|
||||
self.k = nn.Linear(dim, all_head_dim, bias=False)
|
||||
self.v = nn.Linear(dim, all_head_dim, bias=False)
|
||||
|
||||
if qkv_bias:
|
||||
self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
|
||||
self.k_bias = nn.Parameter(torch.zeros(all_head_dim))
|
||||
self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
|
||||
else:
|
||||
self.q_bias = None
|
||||
self.k_bias = None
|
||||
self.v_bias = None
|
||||
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.proj = nn.Linear(all_head_dim, out_dim)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
|
||||
def forward(self, x, k=None, v=None):
|
||||
B, N, C = x.shape
|
||||
N_k = k.shape[1]
|
||||
N_v = v.shape[1]
|
||||
|
||||
q_bias, k_bias, v_bias = None, None, None
|
||||
if self.q_bias is not None:
|
||||
q_bias = self.q_bias
|
||||
k_bias = self.k_bias
|
||||
v_bias = self.v_bias
|
||||
|
||||
q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
|
||||
q = q.reshape(B, N, 1, self.num_heads,
|
||||
-1).permute(2, 0, 3, 1,
|
||||
4).squeeze(0) # (B, N_head, N_q, dim)
|
||||
|
||||
k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
|
||||
k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1,
|
||||
4).squeeze(0)
|
||||
|
||||
v = F.linear(input=v, weight=self.v.weight, bias=v_bias)
|
||||
v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1,
|
||||
4).squeeze(0)
|
||||
|
||||
q = q * self.scale
|
||||
attn = (q @ k.transpose(-2, -1)) # (B, N_head, N_q, N_k)
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class AttentiveBlock(nn.Module):
|
||||
r"""Attentive Block
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
num_heads (int): Number of attention heads. Default: 8
|
||||
qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.
|
||||
Default: False.
|
||||
qk_scale (float | None, optional): Override default qk scale of
|
||||
head_dim ** -0.5 if set. Default: None.
|
||||
drop (float, optional): Dropout rate. Default: 0.0.
|
||||
attn_drop (float, optional): Attention dropout rate. Default: 0.0.
|
||||
drop_path (float | tuple[float], optional): Stochastic depth rate.
|
||||
Default: 0.0.
|
||||
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm.
|
||||
attn_head_dim (int, optional): Dimension of attention head. Default: None.
|
||||
out_dim (int, optional): Dimension of output. Default: None.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
num_heads,
|
||||
qkv_bias=False,
|
||||
qk_scale=None,
|
||||
drop=0.,
|
||||
attn_drop=0.,
|
||||
drop_path=0.,
|
||||
norm_layer="LN",
|
||||
attn_head_dim=None,
|
||||
out_dim=None):
|
||||
super().__init__()
|
||||
|
||||
self.norm1_q = build_norm_layer(dim, norm_layer, eps=1e-6)
|
||||
self.norm1_k = build_norm_layer(dim, norm_layer, eps=1e-6)
|
||||
self.norm1_v = build_norm_layer(dim, norm_layer, eps=1e-6)
|
||||
self.cross_dcn = CrossAttention(dim,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=qkv_bias,
|
||||
qk_scale=qk_scale,
|
||||
attn_drop=attn_drop,
|
||||
proj_drop=drop,
|
||||
attn_head_dim=attn_head_dim,
|
||||
out_dim=out_dim)
|
||||
|
||||
self.drop_path = DropPath(
|
||||
drop_path) if drop_path > 0. else nn.Identity()
|
||||
|
||||
def forward(self,
|
||||
x_q,
|
||||
x_kv,
|
||||
pos_q,
|
||||
pos_k,
|
||||
bool_masked_pos,
|
||||
rel_pos_bias=None):
|
||||
x_q = self.norm1_q(x_q + pos_q)
|
||||
x_k = self.norm1_k(x_kv + pos_k)
|
||||
x_v = self.norm1_v(x_kv)
|
||||
|
||||
x = self.cross_dcn(x_q, k=x_k, v=x_v)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class AttentionPoolingBlock(AttentiveBlock):
|
||||
|
||||
def forward(self, x):
|
||||
x_q = x.mean(1, keepdim=True)
|
||||
x_kv = x
|
||||
pos_q, pos_k = 0, 0
|
||||
x = super().forward(x_q, x_kv, pos_q, pos_k,
|
||||
bool_masked_pos=None,
|
||||
rel_pos_bias=None)
|
||||
x = x.squeeze(1)
|
||||
return x
|
||||
|
||||
|
||||
class StemLayer(nn.Module):
|
||||
r""" Stem layer of InternImage
|
||||
Args:
|
||||
in_chans (int): number of input channels
|
||||
out_chans (int): number of output channels
|
||||
act_layer (str): activation layer
|
||||
norm_layer (str): normalization layer
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
in_chans=3,
|
||||
out_chans=96,
|
||||
act_layer='GELU',
|
||||
norm_layer='BN'):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(in_chans,
|
||||
out_chans // 2,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1)
|
||||
self.norm1 = build_norm_layer(out_chans // 2, norm_layer,
|
||||
'channels_first', 'channels_first')
|
||||
self.act = build_act_layer(act_layer)
|
||||
self.conv2 = nn.Conv2d(out_chans // 2,
|
||||
out_chans,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1)
|
||||
self.norm2 = build_norm_layer(out_chans, norm_layer, 'channels_first',
|
||||
'channels_last')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.act(x)
|
||||
x = self.conv2(x)
|
||||
x = self.norm2(x)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class DownsampleLayer(nn.Module):
|
||||
r""" Downsample layer of InternImage
|
||||
Args:
|
||||
channels (int): number of input channels
|
||||
norm_layer (str): normalization layer
|
||||
"""
|
||||
|
||||
def __init__(self, channels, norm_layer='LN'):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(channels,
|
||||
2 * channels,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
bias=False)
|
||||
self.norm = build_norm_layer(2 * channels, norm_layer,
|
||||
'channels_first', 'channels_first')
|
||||
|
||||
|
||||
def forward(self, x, shape=None):
|
||||
H, W = shape
|
||||
N, HW, C = x.shape
|
||||
|
||||
x = x.view(N, H, W, C)
|
||||
x = self.conv(x.permute(0, 3, 1, 2))
|
||||
x = self.norm(x) # B C H W
|
||||
H, W = x.size(2), x.size(3)
|
||||
x = x.flatten(2).permute(0, 2, 1)
|
||||
|
||||
return x, (H, W)
|
||||
|
||||
|
||||
|
||||
class MLPLayer(nn.Module):
|
||||
r""" MLP layer of InternImage
|
||||
Args:
|
||||
in_features (int): number of input features
|
||||
hidden_features (int): number of hidden features
|
||||
out_features (int): number of output features
|
||||
act_layer (str): activation layer
|
||||
drop (float): dropout rate
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
in_features,
|
||||
hidden_features=None,
|
||||
out_features=None,
|
||||
act_layer='GELU',
|
||||
mlp_fc2_bias=False,
|
||||
drop=0.):
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
|
||||
self.act = build_act_layer(act_layer)
|
||||
self.fc2 = nn.Linear(hidden_features, out_features, bias=mlp_fc2_bias)
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
|
||||
def forward(self, x, shape, level_idx=0):
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class InternImageLayer(nn.Module):
|
||||
r""" Basic layer of InternImage
|
||||
Args:
|
||||
core_op (nn.Module): core operation of InternImage
|
||||
channels (int): number of input channels
|
||||
groups (list): Groups of each block.
|
||||
mlp_ratio (float): ratio of mlp hidden features to input channels
|
||||
drop (float): dropout rate
|
||||
drop_path (float): drop path rate
|
||||
act_layer (str): activation layer
|
||||
norm_layer (str): normalization layer
|
||||
post_norm (bool): whether to use post normalization
|
||||
layer_scale (float): layer scale
|
||||
offset_scale (float): offset scale
|
||||
with_cp (bool): whether to use checkpoint
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
core_op,
|
||||
channels,
|
||||
groups,
|
||||
mlp_ratio=4.,
|
||||
drop=0.,
|
||||
drop_path=0.,
|
||||
act_layer='GELU',
|
||||
norm_layer='LN',
|
||||
post_norm=False,
|
||||
layer_scale=None,
|
||||
offset_scale=1.0,
|
||||
with_cp=False,
|
||||
dcn_output_bias=False,
|
||||
mlp_fc2_bias=False,
|
||||
dw_kernel_size=None, # for InternImage-H/G
|
||||
res_post_norm=False, # for InternImage-H/G
|
||||
center_feature_scale=False): # for InternImage-H/G
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.groups = groups
|
||||
self.mlp_ratio = mlp_ratio
|
||||
self.with_cp = with_cp
|
||||
|
||||
self.norm1 = build_norm_layer(channels, 'LN')
|
||||
self.post_norm = post_norm
|
||||
self.dcn = core_op(
|
||||
channels=channels,
|
||||
group=groups,
|
||||
offset_scale=offset_scale,
|
||||
dw_kernel_size=dw_kernel_size,
|
||||
output_bias=dcn_output_bias,
|
||||
)
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0. \
|
||||
else nn.Identity()
|
||||
self.norm2 = build_norm_layer(channels, 'LN')
|
||||
self.mlp = MLPLayer(in_features=channels,
|
||||
hidden_features=int(channels * mlp_ratio),
|
||||
act_layer=act_layer,
|
||||
drop=drop,
|
||||
mlp_fc2_bias=mlp_fc2_bias
|
||||
)
|
||||
self.layer_scale = layer_scale is not None
|
||||
if self.layer_scale:
|
||||
self.gamma1 = nn.Parameter(layer_scale * torch.ones(channels),
|
||||
requires_grad=True)
|
||||
self.gamma2 = nn.Parameter(layer_scale * torch.ones(channels),
|
||||
requires_grad=True)
|
||||
self.res_post_norm = res_post_norm
|
||||
if res_post_norm:
|
||||
self.res_post_norm1 = build_norm_layer(channels, 'LN')
|
||||
self.res_post_norm2 = build_norm_layer(channels, 'LN')
|
||||
def forward(self, x, shape, level_idx=0):
|
||||
|
||||
def _inner_forward(x, shape, level_idx):
|
||||
if not self.layer_scale:
|
||||
if self.post_norm:
|
||||
x = x + self.drop_path(self.norm1(self.dcn(x, shape, level_idx)))
|
||||
x = x + self.drop_path(self.norm2(self.mlp(x, shape, level_idx)))
|
||||
elif self.res_post_norm: # for InternImage-H/G
|
||||
x = x + self.drop_path(self.res_post_norm1(self.dcn(self.norm1(x), shape, level_idx)))
|
||||
x = x + self.drop_path(self.res_post_norm2(self.mlp(self.norm2(x), shape, level_idx)))
|
||||
|
||||
else:
|
||||
x = x + self.drop_path(self.dcn(self.norm1(x), shape, level_idx))
|
||||
x = x + self.drop_path(self.mlp(self.norm2(x), shape, level_idx))
|
||||
return x
|
||||
if self.post_norm:
|
||||
x = x + self.drop_path(self.gamma1 * self.norm1(self.dcn(x, shape)))
|
||||
x = x + self.drop_path(self.gamma2 * self.norm2(self.mlp(x, shape, level_idx)))
|
||||
else:
|
||||
x = x + self.drop_path(self.gamma1 * self.dcn(self.norm1(x), shape))
|
||||
x = x + self.drop_path(self.gamma2 * self.mlp(self.norm2(x), shape, level_idx))
|
||||
return x
|
||||
|
||||
if self.with_cp and x.requires_grad:
|
||||
x = checkpoint.checkpoint(_inner_forward, x, shape, level_idx)
|
||||
else:
|
||||
x = _inner_forward(x, shape, level_idx)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class InternImageBlock(nn.Module):
|
||||
r""" Block of InternImage
|
||||
Args:
|
||||
core_op (nn.Module): core operation of InternImage
|
||||
channels (int): number of input channels
|
||||
depths (list): Depth of each block.
|
||||
groups (list): Groups of each block.
|
||||
mlp_ratio (float): ratio of mlp hidden features to input channels
|
||||
drop (float): dropout rate
|
||||
drop_path (float): drop path rate
|
||||
act_layer (str): activation layer
|
||||
norm_layer (str): normalization layer
|
||||
post_norm (bool): whether to use post normalization
|
||||
layer_scale (float): layer scale
|
||||
offset_scale (float): offset scale
|
||||
with_cp (bool): whether to use checkpoint
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
core_op,
|
||||
channels,
|
||||
depth,
|
||||
groups,
|
||||
downsample=True,
|
||||
downsample_layer=DownsampleLayer,
|
||||
mlp_ratio=4.,
|
||||
drop=0.,
|
||||
drop_path=0.,
|
||||
act_layer='GELU',
|
||||
norm_layer='LN',
|
||||
post_norm=False,
|
||||
offset_scale=0.5,
|
||||
layer_scale=None,
|
||||
with_cp=False,
|
||||
dcn_output_bias=False,
|
||||
mlp_fc2_bias=False,
|
||||
dw_kernel_size=None, # for InternImage-H/G
|
||||
post_norm_block_ids=None, # for InternImage-H/G
|
||||
res_post_norm=False, # for InternImage-H/G
|
||||
center_feature_scale=False): # for InternImage-H/G
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.depth = depth
|
||||
self.post_norm = post_norm
|
||||
self.center_feature_scale = center_feature_scale
|
||||
|
||||
self.blocks = nn.ModuleList([
|
||||
InternImageLayer(
|
||||
core_op=core_op,
|
||||
channels=channels,
|
||||
groups=groups,
|
||||
mlp_ratio=mlp_ratio,
|
||||
drop=drop,
|
||||
drop_path=drop_path[i] if isinstance(
|
||||
drop_path, list) else drop_path,
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer,
|
||||
post_norm=post_norm,
|
||||
layer_scale=layer_scale,
|
||||
offset_scale=offset_scale,
|
||||
with_cp=with_cp,
|
||||
dcn_output_bias=dcn_output_bias,
|
||||
mlp_fc2_bias=mlp_fc2_bias,
|
||||
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
|
||||
res_post_norm=res_post_norm, # for InternImage-H/G
|
||||
center_feature_scale=center_feature_scale # for InternImage-H/G
|
||||
) for i in range(depth)
|
||||
])
|
||||
if not self.post_norm or center_feature_scale:
|
||||
self.norm = build_norm_layer(channels, 'LN')
|
||||
self.post_norm_block_ids = post_norm_block_ids
|
||||
if post_norm_block_ids is not None: # for InternImage-H/G
|
||||
self.post_norms = nn.ModuleList(
|
||||
[build_norm_layer(channels, 'LN', eps=1e-6) for _ in post_norm_block_ids]
|
||||
)
|
||||
self.downsample = downsample_layer(
|
||||
channels=channels, norm_layer=norm_layer) if downsample else None
|
||||
|
||||
|
||||
def forward(self, x, return_wo_downsample=False, shape=None, level_idx=0
|
||||
):
|
||||
for i, blk in enumerate(self.blocks):
|
||||
x = blk(x, shape=shape, level_idx=level_idx)
|
||||
if (self.post_norm_block_ids is not None) and (i in self.post_norm_block_ids):
|
||||
index = self.post_norm_block_ids.index(i)
|
||||
x = self.post_norms[index](x) # for InternImage-H/G
|
||||
if not self.post_norm or self.center_feature_scale:
|
||||
x = self.norm(x)
|
||||
if return_wo_downsample:
|
||||
x_ = x.clone()
|
||||
if self.downsample is not None:
|
||||
x, shape = self.downsample(x, shape=shape)
|
||||
|
||||
if return_wo_downsample:
|
||||
return x, x_, shape
|
||||
return x, shape
|
||||
|
||||
@BACKBONES.register_module()
|
||||
class FlashInternImage(nn.Module):
|
||||
r""" FlashInternImage
|
||||
A PyTorch impl based on :
|
||||
`InternImage: Exploring Large-Scale Vision Foundation Models with Deformable Convolutions` -
|
||||
https://arxiv.org/pdf/2103.14030
|
||||
'DCNv4': TODO: add arxiv
|
||||
Args:
|
||||
core_op (str): Core operator. Default: 'DCNv4'
|
||||
channels (int): Number of the first stage. Default: 64
|
||||
depths (list): Depth of each block. Default: [3, 4, 18, 5]
|
||||
groups (list): Groups of each block. Default: [3, 6, 12, 24]
|
||||
num_classes (int): Number of classes. Default: 1000
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
|
||||
drop_rate (float): Probability of an element to be zeroed. Default: 0.
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.2
|
||||
act_layer (str): Activation layer. Default: 'GELU'
|
||||
norm_layer (str): Normalization layer. Default: 'LN'
|
||||
layer_scale (bool): Whether to use layer scale. Default: False
|
||||
cls_scale (bool): Whether to use class scale. Default: False
|
||||
with_cp (bool): Use checkpoint or not. Using checkpoint will save some
|
||||
dw_kernel_size (int): Size of the dwconv. Default: None
|
||||
use_clip_projector (bool): Whether to use clip projector. Default: False
|
||||
level2_post_norm (bool): Whether to use level2 post norm. Default: False
|
||||
level2_post_norm_block_ids (list): Indexes of post norm blocks. Default: None
|
||||
res_post_norm (bool): Whether to use res post norm. Default: False
|
||||
center_feature_scale (bool): Whether to use center feature scale. Default: False
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
core_op='DCNv4',
|
||||
channels=64,
|
||||
depths=[3, 4, 18, 5],
|
||||
groups=[3, 6, 12, 24],
|
||||
num_classes=1000,
|
||||
mlp_ratio=4.,
|
||||
drop_rate=0.,
|
||||
drop_path_rate=0.2,
|
||||
drop_path_type='linear',
|
||||
act_layer='GELU',
|
||||
norm_layer='LN',
|
||||
layer_scale=None,
|
||||
offset_scale=0.5,
|
||||
post_norm=False,
|
||||
with_cp=False,
|
||||
mlp_fc2_bias=False,
|
||||
dcn_output_bias=False,
|
||||
dw_kernel_size=None, # for InternImage-H/G
|
||||
level2_post_norm=False, # for InternImage-H/G
|
||||
level2_post_norm_block_ids=None, # for InternImage-H/G
|
||||
res_post_norm=False, # for InternImage-H/G
|
||||
center_feature_scale=False, # for InternImage-H/G
|
||||
out_indices=(0, 1, 2, 3),
|
||||
init_cfg=None,
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
self.core_op = core_op
|
||||
self.num_levels = len(depths)
|
||||
self.depths = depths
|
||||
self.channels = channels
|
||||
self.num_features = int(channels * 2**(self.num_levels - 1))
|
||||
self.post_norm = post_norm
|
||||
self.mlp_ratio = mlp_ratio
|
||||
self.init_cfg = init_cfg
|
||||
self.out_indices = out_indices
|
||||
self.level2_post_norm_block_ids = level2_post_norm_block_ids
|
||||
logger = get_root_logger()
|
||||
logger.info(f'using core type: {core_op}')
|
||||
logger.info(f'using activation layer: {act_layer}')
|
||||
logger.info(f'using main norm layer: {norm_layer}')
|
||||
logger.info(f'using dpr: {drop_path_type}, {drop_path_rate}')
|
||||
logger.info(f"level2_post_norm: {level2_post_norm}")
|
||||
logger.info(f"level2_post_norm_block_ids: {level2_post_norm_block_ids}")
|
||||
logger.info(f"res_post_norm: {res_post_norm}")
|
||||
|
||||
in_chans = 3
|
||||
self.patch_embed = StemLayer(in_chans=in_chans,
|
||||
out_chans=channels,
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer)
|
||||
self.pos_drop = nn.Dropout(p=drop_rate)
|
||||
|
||||
dpr = [
|
||||
x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
|
||||
]
|
||||
if drop_path_type == 'uniform':
|
||||
for i in range(len(dpr)):
|
||||
dpr[i] = drop_path_rate
|
||||
|
||||
self.levels = nn.ModuleList()
|
||||
for i in range(self.num_levels):
|
||||
post_norm_block_ids = level2_post_norm_block_ids if level2_post_norm and (
|
||||
i == 2) else None # for InternImage-H/G
|
||||
|
||||
level = InternImageBlock(
|
||||
core_op=getattr(DCNv4, core_op),
|
||||
channels=int(channels * 2**i),
|
||||
depth=depths[i],
|
||||
groups=groups[i],
|
||||
mlp_ratio=self.mlp_ratio,
|
||||
drop=drop_rate,
|
||||
drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer,
|
||||
post_norm=post_norm,
|
||||
downsample=(i < self.num_levels - 1),
|
||||
downsample_layer = DownsampleLayer,
|
||||
layer_scale=layer_scale,
|
||||
offset_scale=offset_scale,
|
||||
with_cp=with_cp,
|
||||
mlp_fc2_bias=mlp_fc2_bias,
|
||||
dcn_output_bias=dcn_output_bias,
|
||||
dw_kernel_size=dw_kernel_size, # for InternImage-H/G
|
||||
post_norm_block_ids=post_norm_block_ids, # for InternImage-H/G
|
||||
res_post_norm=res_post_norm, # for InternImage-H/G
|
||||
center_feature_scale=center_feature_scale # for InternImage-H/G
|
||||
)
|
||||
self.levels.append(level)
|
||||
|
||||
self.num_layers = len(depths)
|
||||
self.apply(self._init_weights)
|
||||
self.apply(self._init_deform_weights)
|
||||
|
||||
def init_weights(self):
|
||||
logger = get_root_logger()
|
||||
if self.init_cfg is None:
|
||||
logger.warn(f'No pre-trained weights for '
|
||||
f'{self.__class__.__name__}, '
|
||||
f'training start from scratch')
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
trunc_normal_init(m, std=.02, bias=0.)
|
||||
elif isinstance(m, nn.LayerNorm):
|
||||
constant_init(m, 1.0)
|
||||
else:
|
||||
assert 'checkpoint' in self.init_cfg, f'Only support ' \
|
||||
f'specify `Pretrained` in ' \
|
||||
f'`init_cfg` in ' \
|
||||
f'{self.__class__.__name__} '
|
||||
ckpt = _load_checkpoint(self.init_cfg.checkpoint,
|
||||
logger=logger,
|
||||
map_location='cpu')
|
||||
if 'state_dict' in ckpt:
|
||||
_state_dict = ckpt['state_dict']
|
||||
elif 'model_ema' in ckpt:
|
||||
_state_dict = ckpt['model_ema']
|
||||
elif 'model' in ckpt:
|
||||
_state_dict = ckpt['model']
|
||||
else:
|
||||
_state_dict = ckpt
|
||||
|
||||
state_dict = OrderedDict()
|
||||
for k, v in _state_dict.items():
|
||||
if k.startswith('backbone.'):
|
||||
state_dict[k[9:]] = v
|
||||
else:
|
||||
state_dict[k] = v
|
||||
|
||||
# strip prefix of state_dict
|
||||
if list(state_dict.keys())[0].startswith('module.'):
|
||||
state_dict = {k[7:]: v for k, v in state_dict.items()}
|
||||
|
||||
# load state_dict
|
||||
meg = self.load_state_dict(state_dict, False)
|
||||
logger.info(meg)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, nn.Linear):
|
||||
trunc_normal_(m.weight, std=.02)
|
||||
if isinstance(m, nn.Linear) and m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.LayerNorm):
|
||||
nn.init.constant_(m.bias, 0)
|
||||
nn.init.constant_(m.weight, 1.0)
|
||||
|
||||
def _init_deform_weights(self, m):
|
||||
if isinstance(m, getattr(DCNv4, self.core_op)):
|
||||
m._reset_parameters()
|
||||
|
||||
@torch.jit.ignore
|
||||
def lr_decay_keywards(self, decay_ratio=0.87):
|
||||
lr_ratios = {}
|
||||
|
||||
# blocks
|
||||
idx = 0
|
||||
for i in range(4):
|
||||
layer_num = 3 - i # 3 2 1 0
|
||||
for j in range(self.depths[layer_num]):
|
||||
block_num = self.depths[layer_num] - j - 1
|
||||
tag = 'levels.{}.blocks.{}.'.format(layer_num, block_num)
|
||||
decay = 1.0 * (decay_ratio**idx)
|
||||
lr_ratios[tag] = decay
|
||||
idx += 1
|
||||
# patch_embed (before stage-1)
|
||||
lr_ratios["patch_embed"] = lr_ratios['levels.0.blocks.0.']
|
||||
# levels.0.downsample (between stage-1 and stage-2)
|
||||
lr_ratios["levels.0.downsample"] = lr_ratios['levels.1.blocks.0.']
|
||||
lr_ratios["levels.0.norm"] = lr_ratios['levels.1.blocks.0.']
|
||||
# levels.1.downsample (between stage-2 and stage-3)
|
||||
lr_ratios["levels.1.downsample"] = lr_ratios['levels.2.blocks.0.']
|
||||
lr_ratios["levels.1.norm"] = lr_ratios['levels.2.blocks.0.']
|
||||
# levels.2.downsample (between stage-3 and stage-4)
|
||||
lr_ratios["levels.2.downsample"] = lr_ratios['levels.3.blocks.0.']
|
||||
lr_ratios["levels.2.norm"] = lr_ratios['levels.3.blocks.0.']
|
||||
return lr_ratios
|
||||
|
||||
def forward(self, x):
|
||||
x = self.patch_embed(x)
|
||||
N, H, W, C = x.shape
|
||||
x = x.view(N, H*W, C)
|
||||
|
||||
shape=(H, W)
|
||||
seq_out = []
|
||||
for level_idx, level in enumerate(self.levels):
|
||||
old_shape = shape
|
||||
x, x_ , shape = level(x, return_wo_downsample=True, shape=shape, level_idx=level_idx)
|
||||
if level_idx in self.out_indices:
|
||||
h, w= old_shape
|
||||
seq_out.append(x_.reshape(N, h, w, -1).permute(0, 3, 1, 2))
|
||||
return seq_out
|
||||
23
segmentation/mmseg_custom/models/builder.py
Normal file
23
segmentation/mmseg_custom/models/builder.py
Normal file
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import warnings # noqa: F401,F403
|
||||
|
||||
from mmcv.utils import Registry
|
||||
|
||||
TRANSFORMER = Registry('Transformer')
|
||||
MASK_ASSIGNERS = Registry('mask_assigner')
|
||||
MATCH_COST = Registry('match_cost')
|
||||
|
||||
|
||||
def build_match_cost(cfg):
|
||||
"""Build Match Cost."""
|
||||
return MATCH_COST.build(cfg)
|
||||
|
||||
|
||||
def build_assigner(cfg):
|
||||
"""Build Assigner."""
|
||||
return MASK_ASSIGNERS.build(cfg)
|
||||
|
||||
|
||||
def build_transformer(cfg):
|
||||
"""Build Transformer."""
|
||||
return TRANSFORMER.build(cfg)
|
||||
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .mask2former_head import Mask2FormerHead
|
||||
from .maskformer_head import MaskFormerHead
|
||||
from .msda import CustomMultiScaleDeformableAttention
|
||||
__all__ = [
|
||||
'MaskFormerHead',
|
||||
'Mask2FormerHead',
|
||||
]
|
||||
@@ -0,0 +1,579 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import copy
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmcv.cnn import Conv2d, build_plugin_layer, caffe2_xavier_init
|
||||
from mmcv.cnn.bricks.transformer import (build_positional_encoding,
|
||||
build_transformer_layer_sequence)
|
||||
from mmcv.ops import point_sample
|
||||
from mmcv.runner import ModuleList, force_fp32
|
||||
from mmseg.models.builder import HEADS, build_loss
|
||||
from mmseg.models.decode_heads.decode_head import BaseDecodeHead
|
||||
|
||||
from ...core import build_sampler, multi_apply, reduce_mean
|
||||
from ..builder import build_assigner
|
||||
from ..utils import get_uncertain_point_coords_with_randomness
|
||||
|
||||
|
||||
@HEADS.register_module()
|
||||
class Mask2FormerHead(BaseDecodeHead):
|
||||
"""Implements the Mask2Former head.
|
||||
|
||||
See `Masked-attention Mask Transformer for Universal Image
|
||||
Segmentation <https://arxiv.org/pdf/2112.01527>`_ for details.
|
||||
|
||||
Args:
|
||||
in_channels (list[int]): Number of channels in the input feature map.
|
||||
feat_channels (int): Number of channels for features.
|
||||
out_channels (int): Number of channels for output.
|
||||
num_classes (int): Number of classes.
|
||||
num_things_classes (int): Number of things.
|
||||
num_stuff_classes (int): Number of stuff.
|
||||
num_queries (int): Number of query in Transformer decoder.
|
||||
pixel_decoder (:obj:`mmcv.ConfigDict` | dict): Config for pixel
|
||||
decoder. Defaults to None.
|
||||
enforce_decoder_input_project (bool, optional): Whether to add
|
||||
a layer to change the embed_dim of tranformer encoder in
|
||||
pixel decoder to the embed_dim of transformer decoder.
|
||||
Defaults to False.
|
||||
transformer_decoder (:obj:`mmcv.ConfigDict` | dict): Config for
|
||||
transformer decoder. Defaults to None.
|
||||
positional_encoding (:obj:`mmcv.ConfigDict` | dict): Config for
|
||||
transformer decoder position encoding. Defaults to None.
|
||||
loss_cls (:obj:`mmcv.ConfigDict` | dict): Config of the classification
|
||||
loss. Defaults to None.
|
||||
loss_mask (:obj:`mmcv.ConfigDict` | dict): Config of the mask loss.
|
||||
Defaults to None.
|
||||
loss_dice (:obj:`mmcv.ConfigDict` | dict): Config of the dice loss.
|
||||
Defaults to None.
|
||||
train_cfg (:obj:`mmcv.ConfigDict` | dict): Training config of
|
||||
Mask2Former head.
|
||||
test_cfg (:obj:`mmcv.ConfigDict` | dict): Testing config of
|
||||
Mask2Former head.
|
||||
init_cfg (dict or list[dict], optional): Initialization config dict.
|
||||
Defaults to None.
|
||||
"""
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
feat_channels,
|
||||
out_channels,
|
||||
num_classes=80,
|
||||
num_things_classes=None,
|
||||
num_stuff_classes=None,
|
||||
num_queries=100,
|
||||
num_transformer_feat_level=3,
|
||||
pixel_decoder=None,
|
||||
enforce_decoder_input_project=False,
|
||||
transformer_decoder=None,
|
||||
positional_encoding=None,
|
||||
loss_cls=None,
|
||||
loss_mask=None,
|
||||
loss_dice=None,
|
||||
train_cfg=None,
|
||||
test_cfg=None,
|
||||
init_cfg=None,
|
||||
**kwargs):
|
||||
super(Mask2FormerHead, self).__init__(
|
||||
in_channels=in_channels,
|
||||
channels=feat_channels,
|
||||
num_classes=num_classes,
|
||||
init_cfg=init_cfg,
|
||||
input_transform='multiple_select',
|
||||
**kwargs)
|
||||
self.num_classes = num_classes
|
||||
self.num_queries = num_queries
|
||||
self.num_transformer_feat_level = num_transformer_feat_level
|
||||
self.num_heads = transformer_decoder.transformerlayers. \
|
||||
attn_cfgs.num_heads
|
||||
self.num_transformer_decoder_layers = transformer_decoder.num_layers
|
||||
assert pixel_decoder.encoder.transformerlayers. \
|
||||
attn_cfgs.num_levels == num_transformer_feat_level
|
||||
pixel_decoder_ = copy.deepcopy(pixel_decoder)
|
||||
pixel_decoder_.update(
|
||||
in_channels=in_channels,
|
||||
feat_channels=feat_channels,
|
||||
out_channels=out_channels)
|
||||
self.pixel_decoder = build_plugin_layer(pixel_decoder_)[1]
|
||||
self.transformer_decoder = build_transformer_layer_sequence(
|
||||
transformer_decoder)
|
||||
self.decoder_embed_dims = self.transformer_decoder.embed_dims
|
||||
|
||||
self.decoder_input_projs = ModuleList()
|
||||
# from low resolution to high resolution
|
||||
for _ in range(num_transformer_feat_level):
|
||||
if (self.decoder_embed_dims != feat_channels
|
||||
or enforce_decoder_input_project):
|
||||
self.decoder_input_projs.append(
|
||||
Conv2d(
|
||||
feat_channels, self.decoder_embed_dims, kernel_size=1))
|
||||
else:
|
||||
self.decoder_input_projs.append(nn.Identity())
|
||||
self.decoder_positional_encoding = build_positional_encoding(
|
||||
positional_encoding)
|
||||
self.query_embed = nn.Embedding(self.num_queries, feat_channels)
|
||||
self.query_feat = nn.Embedding(self.num_queries, feat_channels)
|
||||
# from low resolution to high resolution
|
||||
self.level_embed = nn.Embedding(self.num_transformer_feat_level,
|
||||
feat_channels)
|
||||
|
||||
self.cls_embed = nn.Linear(feat_channels, self.num_classes + 1)
|
||||
self.mask_embed = nn.Sequential(
|
||||
nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
|
||||
nn.Linear(feat_channels, feat_channels), nn.ReLU(inplace=True),
|
||||
nn.Linear(feat_channels, out_channels))
|
||||
self.conv_seg = None # fix a bug here (conv_seg is not used)
|
||||
|
||||
self.test_cfg = test_cfg
|
||||
self.train_cfg = train_cfg
|
||||
if train_cfg:
|
||||
self.assigner = build_assigner(self.train_cfg.assigner)
|
||||
self.sampler = build_sampler(self.train_cfg.sampler, context=self)
|
||||
self.num_points = self.train_cfg.get('num_points', 12544)
|
||||
self.oversample_ratio = self.train_cfg.get('oversample_ratio', 3.0)
|
||||
self.importance_sample_ratio = self.train_cfg.get(
|
||||
'importance_sample_ratio', 0.75)
|
||||
|
||||
self.class_weight = loss_cls.class_weight
|
||||
self.loss_cls = build_loss(loss_cls)
|
||||
self.loss_mask = build_loss(loss_mask)
|
||||
self.loss_dice = build_loss(loss_dice)
|
||||
|
||||
def init_weights(self):
|
||||
for m in self.decoder_input_projs:
|
||||
if isinstance(m, Conv2d):
|
||||
caffe2_xavier_init(m, bias=0)
|
||||
|
||||
self.pixel_decoder.init_weights()
|
||||
|
||||
for p in self.transformer_decoder.parameters():
|
||||
if p.dim() > 1:
|
||||
nn.init.xavier_normal_(p)
|
||||
|
||||
def get_targets(self, cls_scores_list, mask_preds_list, gt_labels_list,
|
||||
gt_masks_list, img_metas):
|
||||
"""Compute classification and mask targets for all images for a decoder
|
||||
layer.
|
||||
|
||||
Args:
|
||||
cls_scores_list (list[Tensor]): Mask score logits from a single
|
||||
decoder layer for all images. Each with shape [num_queries,
|
||||
cls_out_channels].
|
||||
mask_preds_list (list[Tensor]): Mask logits from a single decoder
|
||||
layer for all images. Each with shape [num_queries, h, w].
|
||||
gt_labels_list (list[Tensor]): Ground truth class indices for all
|
||||
images. Each with shape (n, ), n is the sum of number of stuff
|
||||
type and number of instance in a image.
|
||||
gt_masks_list (list[Tensor]): Ground truth mask for each image,
|
||||
each with shape (n, h, w).
|
||||
img_metas (list[dict]): List of image meta information.
|
||||
|
||||
Returns:
|
||||
tuple[list[Tensor]]: a tuple containing the following targets.
|
||||
|
||||
- labels_list (list[Tensor]): Labels of all images.
|
||||
Each with shape [num_queries, ].
|
||||
- label_weights_list (list[Tensor]): Label weights of all
|
||||
images.Each with shape [num_queries, ].
|
||||
- mask_targets_list (list[Tensor]): Mask targets of all images.
|
||||
Each with shape [num_queries, h, w].
|
||||
- mask_weights_list (list[Tensor]): Mask weights of all images.
|
||||
Each with shape [num_queries, ].
|
||||
- num_total_pos (int): Number of positive samples in all
|
||||
images.
|
||||
- num_total_neg (int): Number of negative samples in all
|
||||
images.
|
||||
"""
|
||||
(labels_list, label_weights_list, mask_targets_list, mask_weights_list,
|
||||
pos_inds_list,
|
||||
neg_inds_list) = multi_apply(self._get_target_single, cls_scores_list,
|
||||
mask_preds_list, gt_labels_list,
|
||||
gt_masks_list, img_metas)
|
||||
|
||||
num_total_pos = sum((inds.numel() for inds in pos_inds_list))
|
||||
num_total_neg = sum((inds.numel() for inds in neg_inds_list))
|
||||
return (labels_list, label_weights_list, mask_targets_list,
|
||||
mask_weights_list, num_total_pos, num_total_neg)
|
||||
|
||||
def _get_target_single(self, cls_score, mask_pred, gt_labels, gt_masks,
|
||||
img_metas):
|
||||
"""Compute classification and mask targets for one image.
|
||||
|
||||
Args:
|
||||
cls_score (Tensor): Mask score logits from a single decoder layer
|
||||
for one image. Shape (num_queries, cls_out_channels).
|
||||
mask_pred (Tensor): Mask logits for a single decoder layer for one
|
||||
image. Shape (num_queries, h, w).
|
||||
gt_labels (Tensor): Ground truth class indices for one image with
|
||||
shape (num_gts, ).
|
||||
gt_masks (Tensor): Ground truth mask for each image, each with
|
||||
shape (num_gts, h, w).
|
||||
img_metas (dict): Image informtation.
|
||||
|
||||
Returns:
|
||||
tuple[Tensor]: A tuple containing the following for one image.
|
||||
|
||||
- labels (Tensor): Labels of each image. \
|
||||
shape (num_queries, ).
|
||||
- label_weights (Tensor): Label weights of each image. \
|
||||
shape (num_queries, ).
|
||||
- mask_targets (Tensor): Mask targets of each image. \
|
||||
shape (num_queries, h, w).
|
||||
- mask_weights (Tensor): Mask weights of each image. \
|
||||
shape (num_queries, ).
|
||||
- pos_inds (Tensor): Sampled positive indices for each \
|
||||
image.
|
||||
- neg_inds (Tensor): Sampled negative indices for each \
|
||||
image.
|
||||
"""
|
||||
# sample points
|
||||
num_queries = cls_score.shape[0]
|
||||
num_gts = gt_labels.shape[0]
|
||||
|
||||
point_coords = torch.rand((1, self.num_points, 2),
|
||||
device=cls_score.device)
|
||||
# shape (num_queries, num_points)
|
||||
mask_points_pred = point_sample(
|
||||
mask_pred.unsqueeze(1), point_coords.repeat(num_queries, 1,
|
||||
1)).squeeze(1)
|
||||
# shape (num_gts, num_points)
|
||||
gt_points_masks = point_sample(
|
||||
gt_masks.unsqueeze(1).float(), point_coords.repeat(num_gts, 1,
|
||||
1)).squeeze(1)
|
||||
|
||||
# assign and sample
|
||||
assign_result = self.assigner.assign(cls_score, mask_points_pred,
|
||||
gt_labels, gt_points_masks,
|
||||
img_metas)
|
||||
sampling_result = self.sampler.sample(assign_result, mask_pred,
|
||||
gt_masks)
|
||||
pos_inds = sampling_result.pos_inds
|
||||
neg_inds = sampling_result.neg_inds
|
||||
|
||||
# label target
|
||||
labels = gt_labels.new_full((self.num_queries, ),
|
||||
self.num_classes,
|
||||
dtype=torch.long)
|
||||
labels[pos_inds] = gt_labels[sampling_result.pos_assigned_gt_inds]
|
||||
label_weights = gt_labels.new_ones((self.num_queries, ))
|
||||
|
||||
# mask target
|
||||
mask_targets = gt_masks[sampling_result.pos_assigned_gt_inds]
|
||||
mask_weights = mask_pred.new_zeros((self.num_queries, ))
|
||||
mask_weights[pos_inds] = 1.0
|
||||
|
||||
return (labels, label_weights, mask_targets, mask_weights, pos_inds,
|
||||
neg_inds)
|
||||
|
||||
def loss_single(self, cls_scores, mask_preds, gt_labels_list,
|
||||
gt_masks_list, img_metas):
|
||||
"""Loss function for outputs from a single decoder layer.
|
||||
|
||||
Args:
|
||||
cls_scores (Tensor): Mask score logits from a single decoder layer
|
||||
for all images. Shape (batch_size, num_queries,
|
||||
cls_out_channels). Note `cls_out_channels` should includes
|
||||
background.
|
||||
mask_preds (Tensor): Mask logits for a pixel decoder for all
|
||||
images. Shape (batch_size, num_queries, h, w).
|
||||
gt_labels_list (list[Tensor]): Ground truth class indices for each
|
||||
image, each with shape (num_gts, ).
|
||||
gt_masks_list (list[Tensor]): Ground truth mask for each image,
|
||||
each with shape (num_gts, h, w).
|
||||
img_metas (list[dict]): List of image meta information.
|
||||
|
||||
Returns:
|
||||
tuple[Tensor]: Loss components for outputs from a single \
|
||||
decoder layer.
|
||||
"""
|
||||
num_imgs = cls_scores.size(0)
|
||||
cls_scores_list = [cls_scores[i] for i in range(num_imgs)]
|
||||
mask_preds_list = [mask_preds[i] for i in range(num_imgs)]
|
||||
(labels_list, label_weights_list, mask_targets_list, mask_weights_list,
|
||||
num_total_pos,
|
||||
num_total_neg) = self.get_targets(cls_scores_list, mask_preds_list,
|
||||
gt_labels_list, gt_masks_list,
|
||||
img_metas)
|
||||
# shape (batch_size, num_queries)
|
||||
labels = torch.stack(labels_list, dim=0)
|
||||
# shape (batch_size, num_queries)
|
||||
label_weights = torch.stack(label_weights_list, dim=0)
|
||||
# shape (num_total_gts, h, w)
|
||||
mask_targets = torch.cat(mask_targets_list, dim=0)
|
||||
# shape (batch_size, num_queries)
|
||||
mask_weights = torch.stack(mask_weights_list, dim=0)
|
||||
|
||||
# classfication loss
|
||||
# shape (batch_size * num_queries, )
|
||||
cls_scores = cls_scores.flatten(0, 1)
|
||||
labels = labels.flatten(0, 1)
|
||||
label_weights = label_weights.flatten(0, 1)
|
||||
|
||||
class_weight = cls_scores.new_tensor(self.class_weight)
|
||||
loss_cls = self.loss_cls(
|
||||
cls_scores,
|
||||
labels,
|
||||
label_weights,
|
||||
avg_factor=class_weight[labels].sum())
|
||||
|
||||
num_total_masks = reduce_mean(cls_scores.new_tensor([num_total_pos]))
|
||||
num_total_masks = max(num_total_masks, 1)
|
||||
|
||||
# extract positive ones
|
||||
# shape (batch_size, num_queries, h, w) -> (num_total_gts, h, w)
|
||||
mask_preds = mask_preds[mask_weights > 0]
|
||||
|
||||
if mask_targets.shape[0] == 0:
|
||||
# zero match
|
||||
loss_dice = mask_preds.sum()
|
||||
loss_mask = mask_preds.sum()
|
||||
return loss_cls, loss_mask, loss_dice
|
||||
|
||||
with torch.no_grad():
|
||||
points_coords = get_uncertain_point_coords_with_randomness(
|
||||
mask_preds.unsqueeze(1), None, self.num_points,
|
||||
self.oversample_ratio, self.importance_sample_ratio)
|
||||
# shape (num_total_gts, h, w) -> (num_total_gts, num_points)
|
||||
mask_point_targets = point_sample(
|
||||
mask_targets.unsqueeze(1).float(), points_coords).squeeze(1)
|
||||
# shape (num_queries, h, w) -> (num_queries, num_points)
|
||||
mask_point_preds = point_sample(
|
||||
mask_preds.unsqueeze(1), points_coords).squeeze(1)
|
||||
|
||||
# dice loss
|
||||
loss_dice = self.loss_dice(
|
||||
mask_point_preds, mask_point_targets, avg_factor=num_total_masks)
|
||||
|
||||
# mask loss
|
||||
# shape (num_queries, num_points) -> (num_queries * num_points, )
|
||||
mask_point_preds = mask_point_preds.reshape(-1,1)
|
||||
# shape (num_total_gts, num_points) -> (num_total_gts * num_points, )
|
||||
mask_point_targets = mask_point_targets.reshape(-1)
|
||||
loss_mask = self.loss_mask(
|
||||
mask_point_preds,
|
||||
mask_point_targets,
|
||||
avg_factor=num_total_masks * self.num_points)
|
||||
|
||||
return loss_cls, loss_mask, loss_dice
|
||||
|
||||
@force_fp32(apply_to=('all_cls_scores', 'all_mask_preds'))
|
||||
def loss(self, all_cls_scores, all_mask_preds, gt_labels_list,
|
||||
gt_masks_list, img_metas):
|
||||
"""Loss function.
|
||||
|
||||
Args:
|
||||
all_cls_scores (Tensor): Classification scores for all decoder
|
||||
layers with shape [num_decoder, batch_size, num_queries,
|
||||
cls_out_channels].
|
||||
all_mask_preds (Tensor): Mask scores for all decoder layers with
|
||||
shape [num_decoder, batch_size, num_queries, h, w].
|
||||
gt_labels_list (list[Tensor]): Ground truth class indices for each
|
||||
image with shape (n, ). n is the sum of number of stuff type
|
||||
and number of instance in a image.
|
||||
gt_masks_list (list[Tensor]): Ground truth mask for each image with
|
||||
shape (n, h, w).
|
||||
img_metas (list[dict]): List of image meta information.
|
||||
|
||||
Returns:
|
||||
dict[str, Tensor]: A dictionary of loss components.
|
||||
"""
|
||||
num_dec_layers = len(all_cls_scores)
|
||||
all_gt_labels_list = [gt_labels_list for _ in range(num_dec_layers)]
|
||||
all_gt_masks_list = [gt_masks_list for _ in range(num_dec_layers)]
|
||||
img_metas_list = [img_metas for _ in range(num_dec_layers)]
|
||||
losses_cls, losses_mask, losses_dice = multi_apply(
|
||||
self.loss_single, all_cls_scores, all_mask_preds,
|
||||
all_gt_labels_list, all_gt_masks_list, img_metas_list)
|
||||
|
||||
loss_dict = dict()
|
||||
# loss from the last decoder layer
|
||||
loss_dict['loss_cls'] = losses_cls[-1]
|
||||
loss_dict['loss_mask'] = losses_mask[-1]
|
||||
loss_dict['loss_dice'] = losses_dice[-1]
|
||||
# loss from other decoder layers
|
||||
num_dec_layer = 0
|
||||
for loss_cls_i, loss_mask_i, loss_dice_i in zip(
|
||||
losses_cls[:-1], losses_mask[:-1], losses_dice[:-1]):
|
||||
loss_dict[f'd{num_dec_layer}.loss_cls'] = loss_cls_i
|
||||
loss_dict[f'd{num_dec_layer}.loss_mask'] = loss_mask_i
|
||||
loss_dict[f'd{num_dec_layer}.loss_dice'] = loss_dice_i
|
||||
num_dec_layer += 1
|
||||
return loss_dict
|
||||
|
||||
def forward_head(self, decoder_out, mask_feature, attn_mask_target_size):
|
||||
"""Forward for head part which is called after every decoder layer.
|
||||
|
||||
Args:
|
||||
decoder_out (Tensor): in shape (num_queries, batch_size, c).
|
||||
mask_feature (Tensor): in shape (batch_size, c, h, w).
|
||||
attn_mask_target_size (tuple[int, int]): target attention
|
||||
mask size.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple contain three elements.
|
||||
|
||||
- cls_pred (Tensor): Classification scores in shape \
|
||||
(batch_size, num_queries, cls_out_channels). \
|
||||
Note `cls_out_channels` should includes background.
|
||||
- mask_pred (Tensor): Mask scores in shape \
|
||||
(batch_size, num_queries,h, w).
|
||||
- attn_mask (Tensor): Attention mask in shape \
|
||||
(batch_size * num_heads, num_queries, h, w).
|
||||
"""
|
||||
decoder_out = self.transformer_decoder.post_norm(decoder_out)
|
||||
decoder_out = decoder_out.transpose(0, 1)
|
||||
# shape (num_queries, batch_size, c)
|
||||
cls_pred = self.cls_embed(decoder_out)
|
||||
# shape (num_queries, batch_size, c)
|
||||
mask_embed = self.mask_embed(decoder_out)
|
||||
# shape (num_queries, batch_size, h, w)
|
||||
mask_pred = torch.einsum('bqc,bchw->bqhw', mask_embed, mask_feature)
|
||||
attn_mask = F.interpolate(
|
||||
mask_pred,
|
||||
attn_mask_target_size,
|
||||
mode='bilinear',
|
||||
align_corners=False)
|
||||
# shape (num_queries, batch_size, h, w) ->
|
||||
# (batch_size * num_head, num_queries, h, w)
|
||||
attn_mask = attn_mask.flatten(2).unsqueeze(1).repeat(
|
||||
(1, self.num_heads, 1, 1)).flatten(0, 1)
|
||||
attn_mask = attn_mask.sigmoid() < 0.5
|
||||
attn_mask = attn_mask.detach()
|
||||
|
||||
return cls_pred, mask_pred, attn_mask
|
||||
|
||||
def forward(self, feats, img_metas):
|
||||
"""Forward function.
|
||||
|
||||
Args:
|
||||
feats (list[Tensor]): Multi scale Features from the
|
||||
upstream network, each is a 4D-tensor.
|
||||
img_metas (list[dict]): List of image information.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple contains two elements.
|
||||
|
||||
- cls_pred_list (list[Tensor)]: Classification logits \
|
||||
for each decoder layer. Each is a 3D-tensor with shape \
|
||||
(batch_size, num_queries, cls_out_channels). \
|
||||
Note `cls_out_channels` should includes background.
|
||||
- mask_pred_list (list[Tensor]): Mask logits for each \
|
||||
decoder layer. Each with shape (batch_size, num_queries, \
|
||||
h, w).
|
||||
"""
|
||||
batch_size = len(img_metas)
|
||||
mask_features, multi_scale_memorys = self.pixel_decoder(feats)
|
||||
# multi_scale_memorys (from low resolution to high resolution)
|
||||
decoder_inputs = []
|
||||
decoder_positional_encodings = []
|
||||
for i in range(self.num_transformer_feat_level):
|
||||
decoder_input = self.decoder_input_projs[i](multi_scale_memorys[i])
|
||||
# shape (batch_size, c, h, w) -> (h*w, batch_size, c)
|
||||
decoder_input = decoder_input.flatten(2).permute(2, 0, 1)
|
||||
level_embed = self.level_embed.weight[i].view(1, 1, -1)
|
||||
decoder_input = decoder_input + level_embed
|
||||
# shape (batch_size, c, h, w) -> (h*w, batch_size, c)
|
||||
mask = decoder_input.new_zeros(
|
||||
(batch_size, ) + multi_scale_memorys[i].shape[-2:],
|
||||
dtype=torch.bool)
|
||||
decoder_positional_encoding = self.decoder_positional_encoding(
|
||||
mask)
|
||||
decoder_positional_encoding = decoder_positional_encoding.flatten(
|
||||
2).permute(2, 0, 1)
|
||||
decoder_inputs.append(decoder_input)
|
||||
decoder_positional_encodings.append(decoder_positional_encoding)
|
||||
# shape (num_queries, c) -> (num_queries, batch_size, c)
|
||||
query_feat = self.query_feat.weight.unsqueeze(1).repeat(
|
||||
(1, batch_size, 1))
|
||||
query_embed = self.query_embed.weight.unsqueeze(1).repeat(
|
||||
(1, batch_size, 1))
|
||||
|
||||
cls_pred_list = []
|
||||
mask_pred_list = []
|
||||
cls_pred, mask_pred, attn_mask = self.forward_head(
|
||||
query_feat, mask_features, multi_scale_memorys[0].shape[-2:])
|
||||
cls_pred_list.append(cls_pred)
|
||||
mask_pred_list.append(mask_pred)
|
||||
|
||||
for i in range(self.num_transformer_decoder_layers):
|
||||
level_idx = i % self.num_transformer_feat_level
|
||||
# if a mask is all True(all background), then set it all False.
|
||||
attn_mask[torch.where(
|
||||
attn_mask.sum(-1) == attn_mask.shape[-1])] = False
|
||||
|
||||
# cross_attn + self_attn
|
||||
layer = self.transformer_decoder.layers[i]
|
||||
attn_masks = [attn_mask, None]
|
||||
query_feat = layer(
|
||||
query=query_feat,
|
||||
key=decoder_inputs[level_idx],
|
||||
value=decoder_inputs[level_idx],
|
||||
query_pos=query_embed,
|
||||
key_pos=decoder_positional_encodings[level_idx],
|
||||
attn_masks=attn_masks,
|
||||
query_key_padding_mask=None,
|
||||
# here we do not apply masking on padded region
|
||||
key_padding_mask=None)
|
||||
cls_pred, mask_pred, attn_mask = self.forward_head(
|
||||
query_feat, mask_features, multi_scale_memorys[
|
||||
(i + 1) % self.num_transformer_feat_level].shape[-2:])
|
||||
|
||||
cls_pred_list.append(cls_pred)
|
||||
mask_pred_list.append(mask_pred)
|
||||
|
||||
return cls_pred_list, mask_pred_list
|
||||
|
||||
def forward_train(self, x, img_metas, gt_semantic_seg, gt_labels,
|
||||
gt_masks):
|
||||
"""Forward function for training mode.
|
||||
|
||||
Args:
|
||||
x (list[Tensor]): Multi-level features from the upstream network,
|
||||
each is a 4D-tensor.
|
||||
img_metas (list[Dict]): List of image information.
|
||||
gt_semantic_seg (list[tensor]):Each element is the ground truth
|
||||
of semantic segmentation with the shape (N, H, W).
|
||||
train_cfg (dict): The training config, which not been used in
|
||||
maskformer.
|
||||
gt_labels (list[Tensor]): Each element is ground truth labels of
|
||||
each box, shape (num_gts,).
|
||||
gt_masks (list[BitmapMasks]): Each element is masks of instances
|
||||
of a image, shape (num_gts, h, w).
|
||||
|
||||
Returns:
|
||||
losses (dict[str, Tensor]): a dictionary of loss components
|
||||
"""
|
||||
|
||||
# forward
|
||||
all_cls_scores, all_mask_preds = self(x, img_metas)
|
||||
|
||||
# loss
|
||||
losses = self.loss(all_cls_scores, all_mask_preds, gt_labels, gt_masks,
|
||||
img_metas)
|
||||
|
||||
return losses
|
||||
|
||||
def forward_test(self, inputs, img_metas, test_cfg):
|
||||
"""Test segment without test-time aumengtation.
|
||||
|
||||
Only the output of last decoder layers was used.
|
||||
|
||||
Args:
|
||||
inputs (list[Tensor]): Multi-level features from the
|
||||
upstream network, each is a 4D-tensor.
|
||||
img_metas (list[dict]): List of image information.
|
||||
test_cfg (dict): Testing config.
|
||||
|
||||
Returns:
|
||||
seg_mask (Tensor): Predicted semantic segmentation logits.
|
||||
"""
|
||||
all_cls_scores, all_mask_preds = self(inputs, img_metas)
|
||||
cls_score, mask_pred = all_cls_scores[-1], all_mask_preds[-1]
|
||||
ori_h, ori_w, _ = img_metas[0]['ori_shape']
|
||||
|
||||
# semantic inference
|
||||
cls_score = F.softmax(cls_score, dim=-1)[..., :-1]
|
||||
mask_pred = mask_pred.sigmoid()
|
||||
seg_mask = torch.einsum('bqc,bqhw->bchw', cls_score, mask_pred)
|
||||
return seg_mask
|
||||
519
segmentation/mmseg_custom/models/decode_heads/maskformer_head.py
Normal file
519
segmentation/mmseg_custom/models/decode_heads/maskformer_head.py
Normal file
@@ -0,0 +1,519 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmcv.cnn import Conv2d, build_plugin_layer, kaiming_init
|
||||
from mmcv.cnn.bricks.transformer import (build_positional_encoding,
|
||||
build_transformer_layer_sequence)
|
||||
from mmcv.runner import force_fp32
|
||||
from mmseg.models.builder import HEADS, build_loss
|
||||
from mmseg.models.decode_heads.decode_head import BaseDecodeHead
|
||||
|
||||
from ...core import multi_apply, reduce_mean
|
||||
from ..builder import build_assigner, build_transformer
|
||||
|
||||
|
||||
@HEADS.register_module()
|
||||
class MaskFormerHead(BaseDecodeHead):
|
||||
"""Implements the MaskFormer head.
|
||||
|
||||
See `paper: Per-Pixel Classification is Not All You Need
|
||||
for Semantic Segmentation<https://arxiv.org/pdf/2107.06278>`
|
||||
for details.
|
||||
|
||||
Args:
|
||||
in_channels (list[int]): Number of channels in the input feature map.
|
||||
feat_channels (int): Number channels for feature.
|
||||
out_channels (int): Number channels for output.
|
||||
num_things_classes (int): Number of things.
|
||||
num_stuff_classes (int): Number of stuff.
|
||||
num_queries (int): Number of query in Transformer.
|
||||
pixel_decoder (obj:`mmcv.ConfigDict`|dict): Config for pixel decoder.
|
||||
Defaults to None.
|
||||
enforce_decoder_input_project (bool, optional): Whether to add a layer
|
||||
to change the embed_dim of tranformer encoder in pixel decoder to
|
||||
the embed_dim of transformer decoder. Defaults to False.
|
||||
transformer_decoder (obj:`mmcv.ConfigDict`|dict): Config for
|
||||
transformer decoder. Defaults to None.
|
||||
positional_encoding (obj:`mmcv.ConfigDict`|dict): Config for
|
||||
transformer decoder position encoding. Defaults to None.
|
||||
loss_cls (obj:`mmcv.ConfigDict`|dict): Config of the classification
|
||||
loss. Defaults to `CrossEntropyLoss`.
|
||||
loss_mask (obj:`mmcv.ConfigDict`|dict): Config of the mask loss.
|
||||
Defaults to `FocalLoss`.
|
||||
loss_dice (obj:`mmcv.ConfigDict`|dict): Config of the dice loss.
|
||||
Defaults to `DiceLoss`.
|
||||
train_cfg (obj:`mmcv.ConfigDict`|dict): Training config of Maskformer
|
||||
head.
|
||||
test_cfg (obj:`mmcv.ConfigDict`|dict): Testing config of Maskformer
|
||||
head.
|
||||
init_cfg (dict or list[dict], optional): Initialization config dict.
|
||||
Defaults to None.
|
||||
"""
|
||||
def __init__(self,
|
||||
out_channels,
|
||||
num_queries=100,
|
||||
pixel_decoder=None,
|
||||
enforce_decoder_input_project=False,
|
||||
transformer_decoder=None,
|
||||
positional_encoding=None,
|
||||
loss_cls=dict(
|
||||
type='CrossEntropyLoss',
|
||||
bg_cls_weight=0.1,
|
||||
use_sigmoid=False,
|
||||
loss_weight=1.0,
|
||||
class_weight=1.0),
|
||||
loss_mask=dict(
|
||||
type='FocalLoss',
|
||||
use_sigmoid=True,
|
||||
gamma=2.0,
|
||||
alpha=0.25,
|
||||
loss_weight=20.0),
|
||||
loss_dice=dict(
|
||||
type='DiceLoss',
|
||||
use_sigmoid=True,
|
||||
activate=True,
|
||||
naive_dice=True,
|
||||
loss_weight=1.0),
|
||||
assigner=dict(
|
||||
type='MaskHungarianAssigner',
|
||||
cls_cost=dict(type='ClassificationCost', weight=1.),
|
||||
dice_cost=dict(type='DiceCost', weight=1.0, pred_act=True,
|
||||
eps=1.0),
|
||||
mask_cost=dict(type='MaskFocalLossCost', weight=20.0)),
|
||||
**kwargs):
|
||||
super(MaskFormerHead, self).__init__(input_transform='multiple_select',
|
||||
**kwargs)
|
||||
self.num_queries = num_queries
|
||||
|
||||
pixel_decoder.update(
|
||||
in_channels=self.in_channels,
|
||||
feat_channels=self.channels,
|
||||
out_channels=out_channels)
|
||||
self.pixel_decoder = build_plugin_layer(pixel_decoder)[1]
|
||||
self.transformer_decoder = build_transformer_layer_sequence(
|
||||
transformer_decoder)
|
||||
self.decoder_embed_dims = self.transformer_decoder.embed_dims
|
||||
pixel_decoder_type = pixel_decoder.get('type')
|
||||
if pixel_decoder_type == 'PixelDecoder' and (
|
||||
self.decoder_embed_dims != self.in_channels[-1]
|
||||
or enforce_decoder_input_project):
|
||||
self.decoder_input_proj = Conv2d(
|
||||
self.in_channels[-1], self.decoder_embed_dims, kernel_size=1)
|
||||
else:
|
||||
self.decoder_input_proj = nn.Identity()
|
||||
self.decoder_pe = build_positional_encoding(positional_encoding)
|
||||
self.query_embed = nn.Embedding(self.num_queries, out_channels)
|
||||
|
||||
self.cls_embed = nn.Linear(self.channels, self.num_classes + 1)
|
||||
self.mask_embed = nn.Sequential(
|
||||
nn.Linear(self.channels, self.channels), nn.ReLU(inplace=True),
|
||||
nn.Linear(self.channels, self.channels), nn.ReLU(inplace=True),
|
||||
nn.Linear(self.channels, out_channels))
|
||||
|
||||
self.assigner = build_assigner(assigner)
|
||||
|
||||
self.bg_cls_weight = 0
|
||||
class_weight = loss_cls.get('class_weight', None)
|
||||
if class_weight is not None and (self.__class__ is MaskFormerHead):
|
||||
assert isinstance(class_weight, float), 'Expected ' \
|
||||
'class_weight to have type float. Found ' \
|
||||
f'{type(class_weight)}.'
|
||||
# NOTE following the official MaskFormerHead repo, bg_cls_weight
|
||||
# means relative classification weight of the VOID class.
|
||||
bg_cls_weight = loss_cls.get('bg_cls_weight', class_weight)
|
||||
assert isinstance(bg_cls_weight, float), 'Expected ' \
|
||||
'bg_cls_weight to have type float. Found ' \
|
||||
f'{type(bg_cls_weight)}.'
|
||||
class_weight = (self.num_classes + 1) * [class_weight]
|
||||
# set VOID class as the last indice
|
||||
class_weight[self.num_classes] = bg_cls_weight
|
||||
loss_cls.update({'class_weight': class_weight})
|
||||
if 'bg_cls_weight' in loss_cls:
|
||||
loss_cls.pop('bg_cls_weight')
|
||||
self.bg_cls_weight = bg_cls_weight
|
||||
|
||||
assert loss_cls['loss_weight'] == assigner['cls_cost']['weight'], \
|
||||
'The classification weight for loss and matcher should be' \
|
||||
'exactly the same.'
|
||||
assert loss_dice['loss_weight'] == assigner['dice_cost']['weight'], \
|
||||
f'The dice weight for loss and matcher' \
|
||||
f'should be exactly the same.'
|
||||
assert loss_mask['loss_weight'] == assigner['mask_cost']['weight'], \
|
||||
'The focal weight for loss and matcher should be' \
|
||||
'exactly the same.'
|
||||
self.loss_cls = build_loss(loss_cls)
|
||||
self.loss_mask = build_loss(loss_mask)
|
||||
self.loss_dice = build_loss(loss_dice)
|
||||
|
||||
self.init_weights()
|
||||
|
||||
def init_weights(self):
|
||||
kaiming_init(self.decoder_input_proj, a=1)
|
||||
|
||||
def get_targets(self, cls_scores_list, mask_preds_list, gt_labels_list,
|
||||
gt_masks_list, img_metas):
|
||||
"""Compute classification and mask targets for all images for a decoder
|
||||
layer.
|
||||
|
||||
Args:
|
||||
cls_scores_list (list[Tensor]): Mask score logits from a single
|
||||
decoder layer for all images. Each with shape [num_queries,
|
||||
cls_out_channels].
|
||||
mask_preds_list (list[Tensor]): Mask logits from a single decoder
|
||||
layer for all images. Each with shape [num_queries, h, w].
|
||||
gt_labels_list (list[Tensor]): Ground truth class indices for all
|
||||
images. Each with shape (n, ), n is the sum of number of stuff
|
||||
type and number of instance in a image.
|
||||
gt_masks_list (list[Tensor]): Ground truth mask for each image,
|
||||
each with shape (n, h, w).
|
||||
img_metas (list[dict]): List of image meta information.
|
||||
|
||||
Returns:
|
||||
tuple[list[Tensor]]: a tuple containing the following targets.
|
||||
|
||||
- labels_list (list[Tensor]): Labels of all images.
|
||||
Each with shape [num_queries, ].
|
||||
- label_weights_list (list[Tensor]): Label weights of all
|
||||
images.Each with shape [num_queries, ].
|
||||
- mask_targets_list (list[Tensor]): Mask targets of all images.
|
||||
Each with shape [num_queries, h, w].
|
||||
- mask_weights_list (list[Tensor]): Mask weights of all images.
|
||||
Each with shape [num_queries, ].
|
||||
- num_total_pos (int): Number of positive samples in all
|
||||
images.
|
||||
- num_total_neg (int): Number of negative samples in all
|
||||
images.
|
||||
"""
|
||||
(labels_list, label_weights_list, mask_targets_list, mask_weights_list,
|
||||
pos_inds_list,
|
||||
neg_inds_list) = multi_apply(self._get_target_single, cls_scores_list,
|
||||
mask_preds_list, gt_labels_list,
|
||||
gt_masks_list, img_metas)
|
||||
|
||||
num_total_pos = sum((inds.numel() for inds in pos_inds_list))
|
||||
num_total_neg = sum((inds.numel() for inds in neg_inds_list))
|
||||
return (labels_list, label_weights_list, mask_targets_list,
|
||||
mask_weights_list, num_total_pos, num_total_neg)
|
||||
|
||||
def _get_target_single(self, cls_score, mask_pred, gt_labels, gt_masks,
|
||||
img_metas):
|
||||
"""Compute classification and mask targets for one image.
|
||||
|
||||
Args:
|
||||
cls_score (Tensor): Mask score logits from a single decoder layer
|
||||
for one image. Shape [num_queries, cls_out_channels].
|
||||
mask_pred (Tensor): Mask logits for a single decoder layer for one
|
||||
image. Shape [num_queries, h, w].
|
||||
gt_labels (Tensor): Ground truth class indices for one image with
|
||||
shape (n, ). n is the sum of number of stuff type and number
|
||||
of instance in a image.
|
||||
gt_masks (Tensor): Ground truth mask for each image, each with
|
||||
shape (n, h, w).
|
||||
img_metas (dict): Image informtation.
|
||||
|
||||
Returns:
|
||||
tuple[Tensor]: a tuple containing the following for one image.
|
||||
|
||||
- labels (Tensor): Labels of each image.
|
||||
shape [num_queries, ].
|
||||
- label_weights (Tensor): Label weights of each image.
|
||||
shape [num_queries, ].
|
||||
- mask_targets (Tensor): Mask targets of each image.
|
||||
shape [num_queries, h, w].
|
||||
- mask_weights (Tensor): Mask weights of each image.
|
||||
shape [num_queries, ].
|
||||
- pos_inds (Tensor): Sampled positive indices for each image.
|
||||
- neg_inds (Tensor): Sampled negative indices for each image.
|
||||
"""
|
||||
target_shape = mask_pred.shape[-2:]
|
||||
gt_masks_downsampled = F.interpolate(
|
||||
gt_masks.unsqueeze(1).float(), target_shape,
|
||||
mode='nearest').squeeze(1).long()
|
||||
# assign and sample
|
||||
assign_result = self.assigner.assign(cls_score, mask_pred, gt_labels,
|
||||
gt_masks_downsampled, img_metas)
|
||||
# pos_ind: range from 1 to (self.num_classes)
|
||||
# which represents the positive index
|
||||
pos_inds = torch.nonzero(assign_result.gt_inds > 0,
|
||||
as_tuple=False).squeeze(-1).unique()
|
||||
neg_inds = torch.nonzero(assign_result.gt_inds == 0,
|
||||
as_tuple=False).squeeze(-1).unique()
|
||||
pos_assigned_gt_inds = assign_result.gt_inds[pos_inds] - 1
|
||||
|
||||
# label target
|
||||
labels = gt_labels.new_full((self.num_queries, ),
|
||||
self.num_classes,
|
||||
dtype=torch.long)
|
||||
labels[pos_inds] = gt_labels[pos_assigned_gt_inds]
|
||||
label_weights = gt_labels.new_ones(self.num_queries)
|
||||
|
||||
# mask target
|
||||
mask_targets = gt_masks[pos_assigned_gt_inds, :]
|
||||
mask_weights = mask_pred.new_zeros((self.num_queries, ))
|
||||
mask_weights[pos_inds] = 1.0
|
||||
|
||||
return (labels, label_weights, mask_targets, mask_weights, pos_inds,
|
||||
neg_inds)
|
||||
|
||||
@force_fp32(apply_to=('all_cls_scores', 'all_mask_preds'))
|
||||
def loss(self, all_cls_scores, all_mask_preds, gt_labels_list,
|
||||
gt_masks_list, img_metas):
|
||||
"""Loss function.
|
||||
|
||||
Args:
|
||||
all_cls_scores (Tensor): Classification scores for all decoder
|
||||
layers with shape [num_decoder, batch_size, num_queries,
|
||||
cls_out_channels].
|
||||
all_mask_preds (Tensor): Mask scores for all decoder layers with
|
||||
shape [num_decoder, batch_size, num_queries, h, w].
|
||||
gt_labels_list (list[Tensor]): Ground truth class indices for each
|
||||
image with shape (n, ). n is the sum of number of stuff type
|
||||
and number of instance in a image.
|
||||
gt_masks_list (list[Tensor]): Ground truth mask for each image with
|
||||
shape (n, h, w).
|
||||
img_metas (list[dict]): List of image meta information.
|
||||
|
||||
Returns:
|
||||
dict[str, Tensor]: A dictionary of loss components.
|
||||
"""
|
||||
num_dec_layers = len(all_cls_scores)
|
||||
all_gt_labels_list = [gt_labels_list for _ in range(num_dec_layers)]
|
||||
all_gt_masks_list = [gt_masks_list for _ in range(num_dec_layers)]
|
||||
img_metas_list = [img_metas for _ in range(num_dec_layers)]
|
||||
losses_cls, losses_mask, losses_dice = multi_apply(
|
||||
self.loss_single, all_cls_scores, all_mask_preds,
|
||||
all_gt_labels_list, all_gt_masks_list, img_metas_list)
|
||||
|
||||
loss_dict = dict()
|
||||
# loss from the last decoder layer
|
||||
loss_dict['loss_cls'] = losses_cls[-1]
|
||||
loss_dict['loss_mask'] = losses_mask[-1]
|
||||
loss_dict['loss_dice'] = losses_dice[-1]
|
||||
# loss from other decoder layers
|
||||
num_dec_layer = 0
|
||||
for loss_cls_i, loss_mask_i, loss_dice_i in zip(
|
||||
losses_cls[:-1], losses_mask[:-1], losses_dice[:-1]):
|
||||
loss_dict[f'd{num_dec_layer}.loss_cls'] = loss_cls_i
|
||||
loss_dict[f'd{num_dec_layer}.loss_mask'] = loss_mask_i
|
||||
loss_dict[f'd{num_dec_layer}.loss_dice'] = loss_dice_i
|
||||
num_dec_layer += 1
|
||||
return loss_dict
|
||||
|
||||
def loss_single(self, cls_scores, mask_preds, gt_labels_list,
|
||||
gt_masks_list, img_metas):
|
||||
"""Loss function for outputs from a single decoder layer.
|
||||
|
||||
Args:
|
||||
cls_scores (Tensor): Mask score logits from a single decoder layer
|
||||
for all images. Shape [batch_size, num_queries,
|
||||
cls_out_channels].
|
||||
mask_preds (Tensor): Mask logits for a pixel decoder for all
|
||||
images. Shape [batch_size, num_queries, h, w].
|
||||
gt_labels_list (list[Tensor]): Ground truth class indices for each
|
||||
image, each with shape (n, ). n is the sum of number of stuff
|
||||
types and number of instances in a image.
|
||||
gt_masks_list (list[Tensor]): Ground truth mask for each image,
|
||||
each with shape (n, h, w).
|
||||
img_metas (list[dict]): List of image meta information.
|
||||
|
||||
Returns:
|
||||
tuple[Tensor]:Loss components for outputs from a single decoder
|
||||
layer.
|
||||
"""
|
||||
num_imgs = cls_scores.size(0)
|
||||
cls_scores_list = [cls_scores[i] for i in range(num_imgs)]
|
||||
mask_preds_list = [mask_preds[i] for i in range(num_imgs)]
|
||||
|
||||
(labels_list, label_weights_list, mask_targets_list, mask_weights_list,
|
||||
num_total_pos,
|
||||
num_total_neg) = self.get_targets(cls_scores_list, mask_preds_list,
|
||||
gt_labels_list, gt_masks_list,
|
||||
img_metas)
|
||||
# shape [batch_size, num_queries]
|
||||
labels = torch.stack(labels_list, dim=0)
|
||||
# shape [batch_size, num_queries]
|
||||
label_weights = torch.stack(label_weights_list, dim=0)
|
||||
# shape [num_gts, h, w]
|
||||
mask_targets = torch.cat(mask_targets_list, dim=0)
|
||||
# shape [batch_size, num_queries]
|
||||
mask_weights = torch.stack(mask_weights_list, dim=0)
|
||||
|
||||
# classfication loss
|
||||
# shape [batch_size * num_queries, ]
|
||||
cls_scores = cls_scores.flatten(0, 1)
|
||||
# shape [batch_size * num_queries, ]
|
||||
labels = labels.flatten(0, 1)
|
||||
# shape [batch_size* num_queries, ]
|
||||
label_weights = label_weights.flatten(0, 1)
|
||||
|
||||
class_weight = cls_scores.new_ones(self.num_classes + 1)
|
||||
class_weight[-1] = self.bg_cls_weight
|
||||
|
||||
loss_cls = self.loss_cls(
|
||||
cls_scores,
|
||||
labels,
|
||||
label_weights,
|
||||
avg_factor=class_weight[labels].sum())
|
||||
|
||||
num_total_masks = reduce_mean(cls_scores.new_tensor([num_total_pos]))
|
||||
num_total_masks = max(num_total_masks, 1)
|
||||
|
||||
# extract positive ones
|
||||
mask_preds = mask_preds[mask_weights > 0]
|
||||
target_shape = mask_targets.shape[-2:]
|
||||
|
||||
if mask_targets.shape[0] == 0:
|
||||
# zero match
|
||||
loss_dice = mask_preds.sum()
|
||||
loss_mask = mask_preds.sum()
|
||||
return loss_cls, loss_mask, loss_dice
|
||||
|
||||
# upsample to shape of target
|
||||
# shape [num_gts, h, w]
|
||||
mask_preds = F.interpolate(
|
||||
mask_preds.unsqueeze(1),
|
||||
target_shape,
|
||||
mode='bilinear',
|
||||
align_corners=False).squeeze(1)
|
||||
|
||||
# dice loss
|
||||
loss_dice = self.loss_dice(
|
||||
mask_preds, mask_targets, avg_factor=num_total_masks)
|
||||
|
||||
# mask loss
|
||||
# FocalLoss support input of shape [n, num_class]
|
||||
h, w = mask_preds.shape[-2:]
|
||||
# shape [num_gts, h, w] -> [num_gts * h * w, 1]
|
||||
mask_preds = mask_preds.reshape(-1, 1)
|
||||
# shape [num_gts, h, w] -> [num_gts * h * w]
|
||||
mask_targets = mask_targets.reshape(-1)
|
||||
# target is (1 - mask_targets) !!!
|
||||
print("mask_pred:", mask_preds.shape)
|
||||
print("mask_targets:", mask_targets.shape)
|
||||
loss_mask = self.loss_mask(
|
||||
mask_preds, 1 - mask_targets, avg_factor=num_total_masks * h * w)
|
||||
|
||||
return loss_cls, loss_mask, loss_dice
|
||||
|
||||
def forward(self, feats, img_metas):
|
||||
"""Forward function.
|
||||
|
||||
Args:
|
||||
feats (list[Tensor]): Features from the upstream network, each
|
||||
is a 4D-tensor.
|
||||
img_metas (list[dict]): List of image information.
|
||||
|
||||
Returns:
|
||||
all_cls_scores (Tensor): Classification scores for each
|
||||
scale level. Each is a 4D-tensor with shape
|
||||
[num_decoder, batch_size, num_queries, cls_out_channels].
|
||||
Note `cls_out_channels` should includes background.
|
||||
all_mask_preds (Tensor): Mask scores for each decoder
|
||||
layer. Each with shape [num_decoder, batch_size,
|
||||
num_queries, h, w].
|
||||
"""
|
||||
batch_size = len(img_metas)
|
||||
input_img_h, input_img_w = img_metas[0]['pad_shape'][:-1]
|
||||
# input_img_h, input_img_w = img_metas[0]['batch_input_shape']
|
||||
padding_mask = feats[-1].new_ones(
|
||||
(batch_size, input_img_h, input_img_w), dtype=torch.float32)
|
||||
for i in range(batch_size):
|
||||
img_h, img_w, _ = img_metas[i]['img_shape']
|
||||
padding_mask[i, :img_h, :img_w] = 0
|
||||
padding_mask = F.interpolate(
|
||||
padding_mask.unsqueeze(1),
|
||||
size=feats[-1].shape[-2:],
|
||||
mode='nearest').to(torch.bool).squeeze(1)
|
||||
# when backbone is swin, memory is output of last stage of swin.
|
||||
# when backbone is r50, memory is output of tranformer encoder.
|
||||
mask_features, memory = self.pixel_decoder(feats, img_metas)
|
||||
pos_embed = self.decoder_pe(padding_mask)
|
||||
memory = self.decoder_input_proj(memory)
|
||||
# shape [batch_size, c, h, w] -> [h*w, batch_size, c]
|
||||
memory = memory.flatten(2).permute(2, 0, 1)
|
||||
pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
|
||||
# shape [batch_size, h * w]
|
||||
padding_mask = padding_mask.flatten(1)
|
||||
# shape = [num_queries, embed_dims]
|
||||
query_embed = self.query_embed.weight
|
||||
# shape = [num_queries, batch_size, embed_dims]
|
||||
query_embed = query_embed.unsqueeze(1).repeat(1, batch_size, 1)
|
||||
target = torch.zeros_like(query_embed)
|
||||
# shape [num_decoder, num_queries, batch_size, embed_dims]
|
||||
out_dec = self.transformer_decoder(
|
||||
query=target,
|
||||
key=memory,
|
||||
value=memory,
|
||||
key_pos=pos_embed,
|
||||
query_pos=query_embed,
|
||||
key_padding_mask=padding_mask)
|
||||
# shape [num_decoder, batch_size, num_queries, embed_dims]
|
||||
out_dec = out_dec.transpose(1, 2)
|
||||
|
||||
# cls_scores
|
||||
all_cls_scores = self.cls_embed(out_dec)
|
||||
|
||||
# mask_preds
|
||||
mask_embed = self.mask_embed(out_dec)
|
||||
all_mask_preds = torch.einsum('lbqc,bchw->lbqhw', mask_embed,
|
||||
mask_features)
|
||||
|
||||
return all_cls_scores, all_mask_preds
|
||||
|
||||
def forward_train(self,
|
||||
x,
|
||||
img_metas,
|
||||
gt_semantic_seg,
|
||||
gt_labels,
|
||||
gt_masks):
|
||||
"""Forward function for training mode.
|
||||
|
||||
Args:
|
||||
x (list[Tensor]): Multi-level features from the upstream network,
|
||||
each is a 4D-tensor.
|
||||
img_metas (list[Dict]): List of image information.
|
||||
gt_semantic_seg (list[tensor]):Each element is the ground truth
|
||||
of semantic segmentation with the shape (N, H, W).
|
||||
train_cfg (dict): The training config, which not been used in
|
||||
maskformer.
|
||||
gt_labels (list[Tensor]): Each element is ground truth labels of
|
||||
each box, shape (num_gts,).
|
||||
gt_masks (list[BitmapMasks]): Each element is masks of instances
|
||||
of a image, shape (num_gts, h, w).
|
||||
|
||||
Returns:
|
||||
losses (dict[str, Tensor]): a dictionary of loss components
|
||||
"""
|
||||
|
||||
# forward
|
||||
all_cls_scores, all_mask_preds = self(x, img_metas)
|
||||
|
||||
# loss
|
||||
losses = self.loss(all_cls_scores, all_mask_preds, gt_labels, gt_masks,
|
||||
img_metas)
|
||||
|
||||
return losses
|
||||
|
||||
def forward_test(self, inputs, img_metas, test_cfg):
|
||||
"""Test segment without test-time aumengtation.
|
||||
|
||||
Only the output of last decoder layers was used.
|
||||
|
||||
Args:
|
||||
inputs (list[Tensor]): Multi-level features from the
|
||||
upstream network, each is a 4D-tensor.
|
||||
img_metas (list[dict]): List of image information.
|
||||
test_cfg (dict): Testing config.
|
||||
|
||||
Returns:
|
||||
seg_mask (Tensor): Predicted semantic segmentation logits.
|
||||
"""
|
||||
all_cls_scores, all_mask_preds = self(inputs, img_metas)
|
||||
cls_score, mask_pred = all_cls_scores[-1], all_mask_preds[-1]
|
||||
ori_h, ori_w, _ = img_metas[0]['ori_shape']
|
||||
|
||||
# semantic inference
|
||||
cls_score = F.softmax(cls_score, dim=-1)[..., :-1]
|
||||
mask_pred = mask_pred.sigmoid()
|
||||
seg_mask = torch.einsum('bqc,bqhw->bchw', cls_score, mask_pred)
|
||||
return seg_mask
|
||||
355
segmentation/mmseg_custom/models/decode_heads/msda.py
Normal file
355
segmentation/mmseg_custom/models/decode_heads/msda.py
Normal file
@@ -0,0 +1,355 @@
|
||||
# --------------------------------------------------------
|
||||
# DCNv4
|
||||
# Copyright (c) 2024 OpenGVLab
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# --------------------------------------------------------
|
||||
|
||||
|
||||
import torch
|
||||
from torch.cuda.amp import custom_bwd, custom_fwd
|
||||
from torch.autograd.function import Function, once_differentiable
|
||||
from mmcv.utils import ext_loader
|
||||
ext_module = ext_loader.load_ext(
|
||||
'_ext', ['ms_deform_attn_backward', 'ms_deform_attn_forward'])
|
||||
|
||||
class MultiScaleDeformableAttnFunction_fp16(Function):
|
||||
|
||||
@staticmethod
|
||||
@custom_fwd(cast_inputs=torch.float16)
|
||||
def forward(ctx, value, value_spatial_shapes, value_level_start_index,
|
||||
sampling_locations, attention_weights, im2col_step):
|
||||
"""GPU version of multi-scale deformable attention.
|
||||
|
||||
Args:
|
||||
value (Tensor): The value has shape
|
||||
(bs, num_keys, mum_heads, embed_dims//num_heads)
|
||||
value_spatial_shapes (Tensor): Spatial shape of
|
||||
each feature map, has shape (num_levels, 2),
|
||||
last dimension 2 represent (h, w)
|
||||
sampling_locations (Tensor): The location of sampling points,
|
||||
has shape
|
||||
(bs ,num_queries, num_heads, num_levels, num_points, 2),
|
||||
the last dimension 2 represent (x, y).
|
||||
attention_weights (Tensor): The weight of sampling points used
|
||||
when calculate the attention, has shape
|
||||
(bs ,num_queries, num_heads, num_levels, num_points),
|
||||
im2col_step (Tensor): The step used in image to column.
|
||||
|
||||
Returns:
|
||||
Tensor: has shape (bs, num_queries, embed_dims)
|
||||
"""
|
||||
ctx.im2col_step = im2col_step
|
||||
output = ext_module.ms_deform_attn_forward(
|
||||
value,
|
||||
value_spatial_shapes,
|
||||
value_level_start_index,
|
||||
sampling_locations,
|
||||
attention_weights,
|
||||
im2col_step=ctx.im2col_step)
|
||||
ctx.save_for_backward(value, value_spatial_shapes,
|
||||
value_level_start_index, sampling_locations,
|
||||
attention_weights)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
@once_differentiable
|
||||
@custom_bwd
|
||||
def backward(ctx, grad_output):
|
||||
"""GPU version of backward function.
|
||||
|
||||
Args:
|
||||
grad_output (Tensor): Gradient
|
||||
of output tensor of forward.
|
||||
|
||||
Returns:
|
||||
Tuple[Tensor]: Gradient
|
||||
of input tensors in forward.
|
||||
"""
|
||||
value, value_spatial_shapes, value_level_start_index, \
|
||||
sampling_locations, attention_weights = ctx.saved_tensors
|
||||
grad_value = torch.zeros_like(value)
|
||||
grad_sampling_loc = torch.zeros_like(sampling_locations)
|
||||
grad_attn_weight = torch.zeros_like(attention_weights)
|
||||
|
||||
ext_module.ms_deform_attn_backward(
|
||||
value,
|
||||
value_spatial_shapes,
|
||||
value_level_start_index,
|
||||
sampling_locations,
|
||||
attention_weights,
|
||||
grad_output.contiguous(),
|
||||
grad_value,
|
||||
grad_sampling_loc,
|
||||
grad_attn_weight,
|
||||
im2col_step=ctx.im2col_step)
|
||||
|
||||
return grad_value, None, None, \
|
||||
grad_sampling_loc, grad_attn_weight, None
|
||||
|
||||
|
||||
|
||||
|
||||
class MultiScaleDeformableAttnFunction_fp32_old(Function):
|
||||
|
||||
@staticmethod
|
||||
@custom_fwd(cast_inputs=torch.float32)
|
||||
def forward(ctx, value, value_spatial_shapes, value_level_start_index,
|
||||
sampling_locations, attention_weights, im2col_step):
|
||||
"""GPU version of multi-scale deformable attention.
|
||||
|
||||
Args:
|
||||
value (Tensor): The value has shape
|
||||
(bs, num_keys, mum_heads, embed_dims//num_heads)
|
||||
value_spatial_shapes (Tensor): Spatial shape of
|
||||
each feature map, has shape (num_levels, 2),
|
||||
last dimension 2 represent (h, w)
|
||||
sampling_locations (Tensor): The location of sampling points,
|
||||
has shape
|
||||
(bs ,num_queries, num_heads, num_levels, num_points, 2),
|
||||
the last dimension 2 represent (x, y).
|
||||
attention_weights (Tensor): The weight of sampling points used
|
||||
when calculate the attention, has shape
|
||||
(bs ,num_queries, num_heads, num_levels, num_points),
|
||||
im2col_step (Tensor): The step used in image to column.
|
||||
|
||||
Returns:
|
||||
Tensor: has shape (bs, num_queries, embed_dims)
|
||||
"""
|
||||
|
||||
ctx.im2col_step = im2col_step
|
||||
output = ext_module.ms_deform_attn_forward(
|
||||
value.to(torch.float),
|
||||
value_spatial_shapes,
|
||||
value_level_start_index,
|
||||
sampling_locations.to(torch.float),
|
||||
attention_weights.to(torch.float),
|
||||
im2col_step=ctx.im2col_step).to(torch.float16)
|
||||
ctx.save_for_backward(value, value_spatial_shapes,
|
||||
value_level_start_index, sampling_locations,
|
||||
attention_weights)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
@once_differentiable
|
||||
@custom_bwd
|
||||
def backward(ctx, grad_output):
|
||||
"""GPU version of backward function.
|
||||
|
||||
Args:
|
||||
grad_output (Tensor): Gradient
|
||||
of output tensor of forward.
|
||||
|
||||
Returns:
|
||||
Tuple[Tensor]: Gradient
|
||||
of input tensors in forward.
|
||||
"""
|
||||
value, value_spatial_shapes, value_level_start_index, \
|
||||
sampling_locations, attention_weights = ctx.saved_tensors
|
||||
grad_value = torch.zeros_like(value)
|
||||
grad_sampling_loc = torch.zeros_like(sampling_locations)
|
||||
grad_attn_weight = torch.zeros_like(attention_weights)
|
||||
|
||||
ext_module.ms_deform_attn_backward(
|
||||
value,
|
||||
value_spatial_shapes,
|
||||
value_level_start_index,
|
||||
sampling_locations,
|
||||
attention_weights,
|
||||
grad_output.contiguous(),
|
||||
grad_value,
|
||||
grad_sampling_loc,
|
||||
grad_attn_weight,
|
||||
im2col_step=ctx.im2col_step)
|
||||
|
||||
return grad_value, None, None, \
|
||||
grad_sampling_loc, grad_attn_weight, None
|
||||
|
||||
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.autograd.function import Function, once_differentiable
|
||||
|
||||
from mmcv import deprecated_api_warning
|
||||
from mmcv.cnn import constant_init, xavier_init
|
||||
from mmcv.cnn.bricks.registry import ATTENTION
|
||||
from mmcv.runner import BaseModule
|
||||
|
||||
ext_module = ext_loader.load_ext(
|
||||
'_ext', ['ms_deform_attn_backward', 'ms_deform_attn_forward'])
|
||||
import functools
|
||||
import time
|
||||
from collections import defaultdict
|
||||
import torch
|
||||
from mmcv.ops import MultiScaleDeformableAttention
|
||||
@ATTENTION.register_module()
|
||||
class CustomMultiScaleDeformableAttention(MultiScaleDeformableAttention):
|
||||
"""An attention module used in Deformable-Detr.
|
||||
|
||||
`Deformable DETR: Deformable Transformers for End-to-End Object Detection.
|
||||
<https://arxiv.org/pdf/2010.04159.pdf>`_.
|
||||
|
||||
Args:
|
||||
embed_dims (int): The embedding dimension of Attention.
|
||||
Default: 256.
|
||||
num_heads (int): Parallel attention heads. Default: 64.
|
||||
num_levels (int): The number of feature map used in
|
||||
Attention. Default: 4.
|
||||
num_points (int): The number of sampling points for
|
||||
each query in each head. Default: 4.
|
||||
im2col_step (int): The step used in image_to_column.
|
||||
Default: 64.
|
||||
dropout (float): A Dropout layer on `inp_identity`.
|
||||
Default: 0.1.
|
||||
batch_first (bool): Key, Query and Value are shape of
|
||||
(batch, n, embed_dim)
|
||||
or (n, batch, embed_dim). Default to False.
|
||||
norm_cfg (dict): Config dict for normalization layer.
|
||||
Default: None.
|
||||
init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
|
||||
Default: None.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
use_flash=False,
|
||||
use_softmax=True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.use_flash = use_flash
|
||||
self.use_softmax = use_softmax
|
||||
|
||||
@deprecated_api_warning({'residual': 'identity'},
|
||||
cls_name='FlashMultiScaleDeformableAttention')
|
||||
# @run_time('ms_attention')
|
||||
def forward(self,
|
||||
query,
|
||||
key=None,
|
||||
value=None,
|
||||
identity=None,
|
||||
query_pos=None,
|
||||
key_padding_mask=None,
|
||||
reference_points=None,
|
||||
spatial_shapes=None,
|
||||
level_start_index=None,
|
||||
**kwargs):
|
||||
"""Forward Function of MultiScaleDeformAttention.
|
||||
|
||||
Args:
|
||||
query (torch.Tensor): Query of Transformer with shape
|
||||
(num_query, bs, embed_dims).
|
||||
key (torch.Tensor): The key tensor with shape
|
||||
`(num_key, bs, embed_dims)`.
|
||||
value (torch.Tensor): The value tensor with shape
|
||||
`(num_key, bs, embed_dims)`.
|
||||
identity (torch.Tensor): The tensor used for addition, with the
|
||||
same shape as `query`. Default None. If None,
|
||||
`query` will be used.
|
||||
query_pos (torch.Tensor): The positional encoding for `query`.
|
||||
Default: None.
|
||||
key_pos (torch.Tensor): The positional encoding for `key`. Default
|
||||
None.
|
||||
reference_points (torch.Tensor): The normalized reference
|
||||
points with shape (bs, num_query, num_levels, 2),
|
||||
all elements is range in [0, 1], top-left (0,0),
|
||||
bottom-right (1, 1), including padding area.
|
||||
or (N, Length_{query}, num_levels, 4), add
|
||||
additional two dimensions is (w, h) to
|
||||
form reference boxes.
|
||||
key_padding_mask (torch.Tensor): ByteTensor for `query`, with
|
||||
shape [bs, num_key].
|
||||
spatial_shapes (torch.Tensor): Spatial shape of features in
|
||||
different levels. With shape (num_levels, 2),
|
||||
last dimension represents (h, w).
|
||||
level_start_index (torch.Tensor): The start index of each level.
|
||||
A tensor has shape ``(num_levels, )`` and can be represented
|
||||
as [0, h_0*w_0, h_0*w_0+h_1*w_1, ...].
|
||||
|
||||
Returns:
|
||||
torch.Tensor: forwarded results with shape
|
||||
[num_query, bs, embed_dims].
|
||||
"""
|
||||
|
||||
if value is None:
|
||||
value = query
|
||||
|
||||
if identity is None:
|
||||
identity = query
|
||||
if query_pos is not None:
|
||||
query = query + query_pos
|
||||
if not self.batch_first:
|
||||
# change to (bs, num_query ,embed_dims)
|
||||
query = query.permute(1, 0, 2)
|
||||
value = value.permute(1, 0, 2)
|
||||
|
||||
bs, num_query, _ = query.shape
|
||||
bs, num_value, _ = value.shape
|
||||
assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value
|
||||
|
||||
value = self.value_proj(value)
|
||||
if key_padding_mask is not None:
|
||||
value = value.masked_fill(key_padding_mask[..., None], 0.0)
|
||||
value = value.view(bs, num_value, self.num_heads, -1)
|
||||
query = query.to(value.dtype)
|
||||
sampling_offsets = self.sampling_offsets(query).view(
|
||||
bs, num_query, self.num_heads, self.num_levels, self.num_points, 2)
|
||||
attention_weights = self.attention_weights(query).view(
|
||||
bs, num_query, self.num_heads, self.num_levels * self.num_points)
|
||||
|
||||
if not self.use_flash:
|
||||
if self.use_softmax:
|
||||
attention_weights = attention_weights.softmax(-1)
|
||||
attention_weights = attention_weights.view(bs, num_query,
|
||||
self.num_heads,
|
||||
self.num_levels,
|
||||
self.num_points)
|
||||
|
||||
else:
|
||||
attention_weights = attention_weights.view(bs, num_query,
|
||||
self.num_heads,
|
||||
self.num_levels,
|
||||
self.num_points, 1)
|
||||
|
||||
if reference_points.shape[-1] == 2:
|
||||
offset_normalizer = torch.stack(
|
||||
[spatial_shapes[..., 1], spatial_shapes[..., 0]], -1)
|
||||
sampling_locations = reference_points[:, :, None, :, None, :] \
|
||||
+ sampling_offsets \
|
||||
/ offset_normalizer[None, None, None, :, None, :]
|
||||
elif reference_points.shape[-1] == 4:
|
||||
sampling_locations = reference_points[:, :, None, :, None, :2] \
|
||||
+ sampling_offsets / self.num_points \
|
||||
* reference_points[:, :, None, :, None, 2:] \
|
||||
* 0.5
|
||||
else:
|
||||
raise ValueError(
|
||||
f'Last dim of reference_points must be'
|
||||
f' 2 or 4, but get {reference_points.shape[-1]} instead.')
|
||||
sampling_locations = sampling_locations.to(sampling_offsets.dtype)
|
||||
if torch.cuda.is_available() and value.is_cuda:
|
||||
if self.use_flash:
|
||||
assert False
|
||||
else:
|
||||
MultiScaleDeformableAttnFunction = MultiScaleDeformableAttnFunction_fp32_old
|
||||
|
||||
output = MultiScaleDeformableAttnFunction.apply(
|
||||
value, spatial_shapes, level_start_index, sampling_locations,
|
||||
attention_weights, self.im2col_step)
|
||||
|
||||
else:
|
||||
output = multi_scale_deformable_attn_pytorch(
|
||||
value, spatial_shapes, sampling_locations, attention_weights)
|
||||
|
||||
output = self.output_proj(output.to(value.dtype))
|
||||
|
||||
if not self.batch_first:
|
||||
# (num_query, bs ,embed_dims)
|
||||
output = output.permute(1, 0, 2)
|
||||
|
||||
return self.dropout(output) + identity
|
||||
|
||||
13
segmentation/mmseg_custom/models/losses/__init__.py
Normal file
13
segmentation/mmseg_custom/models/losses/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .cross_entropy_loss import (CrossEntropyLoss, binary_cross_entropy,
|
||||
cross_entropy, mask_cross_entropy)
|
||||
from .dice_loss import DiceLoss
|
||||
from .focal_loss import FocalLoss
|
||||
from .match_costs import (ClassificationCost, CrossEntropyLossCost, DiceCost,
|
||||
MaskFocalLossCost)
|
||||
|
||||
__all__ = [
|
||||
'cross_entropy', 'binary_cross_entropy', 'mask_cross_entropy',
|
||||
'CrossEntropyLoss', 'DiceLoss', 'FocalLoss', 'ClassificationCost',
|
||||
'MaskFocalLossCost', 'DiceCost', 'CrossEntropyLossCost'
|
||||
]
|
||||
291
segmentation/mmseg_custom/models/losses/cross_entropy_loss.py
Normal file
291
segmentation/mmseg_custom/models/losses/cross_entropy_loss.py
Normal file
@@ -0,0 +1,291 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmseg.models.builder import LOSSES
|
||||
from mmseg.models.losses.utils import get_class_weight, weight_reduce_loss
|
||||
|
||||
|
||||
def cross_entropy(pred,
|
||||
label,
|
||||
weight=None,
|
||||
class_weight=None,
|
||||
reduction='mean',
|
||||
avg_factor=None,
|
||||
ignore_index=-100,
|
||||
avg_non_ignore=False):
|
||||
"""cross_entropy. The wrapper function for :func:`F.cross_entropy`
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction with shape (N, 1).
|
||||
label (torch.Tensor): The learning label of the prediction.
|
||||
weight (torch.Tensor, optional): Sample-wise loss weight.
|
||||
Default: None.
|
||||
class_weight (list[float], optional): The weight for each class.
|
||||
Default: None.
|
||||
reduction (str, optional): The method used to reduce the loss.
|
||||
Options are 'none', 'mean' and 'sum'. Default: 'mean'.
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Default: None.
|
||||
ignore_index (int): Specifies a target value that is ignored and
|
||||
does not contribute to the input gradients. When
|
||||
``avg_non_ignore `` is ``True``, and the ``reduction`` is
|
||||
``''mean''``, the loss is averaged over non-ignored targets.
|
||||
Defaults: -100.
|
||||
avg_non_ignore (bool): The flag decides to whether the loss is
|
||||
only averaged over non-ignored targets. Default: False.
|
||||
`New in version 0.23.0.`
|
||||
"""
|
||||
|
||||
# class_weight is a manual rescaling weight given to each class.
|
||||
# If given, has to be a Tensor of size C element-wise losses
|
||||
loss = F.cross_entropy(
|
||||
pred,
|
||||
label,
|
||||
weight=class_weight,
|
||||
reduction='none',
|
||||
ignore_index=ignore_index)
|
||||
|
||||
# apply weights and do the reduction
|
||||
# average loss over non-ignored elements
|
||||
# pytorch's official cross_entropy average loss over non-ignored elements
|
||||
# refer to https://github.com/pytorch/pytorch/blob/56b43f4fec1f76953f15a627694d4bba34588969/torch/nn/functional.py#L2660 # noqa
|
||||
if (avg_factor is None) and avg_non_ignore and reduction == 'mean':
|
||||
avg_factor = label.numel() - (label == ignore_index).sum().item()
|
||||
if weight is not None:
|
||||
weight = weight.float()
|
||||
loss = weight_reduce_loss(
|
||||
loss, weight=weight, reduction=reduction, avg_factor=avg_factor)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def _expand_onehot_labels(labels, label_weights, target_shape, ignore_index):
|
||||
"""Expand onehot labels to match the size of prediction."""
|
||||
bin_labels = labels.new_zeros(target_shape)
|
||||
valid_mask = (labels >= 0) & (labels != ignore_index)
|
||||
inds = torch.nonzero(valid_mask, as_tuple=True)
|
||||
|
||||
if inds[0].numel() > 0:
|
||||
if labels.dim() == 3:
|
||||
bin_labels[inds[0], labels[valid_mask], inds[1], inds[2]] = 1
|
||||
else:
|
||||
bin_labels[inds[0], labels[valid_mask]] = 1
|
||||
|
||||
valid_mask = valid_mask.unsqueeze(1).expand(target_shape).float()
|
||||
|
||||
if label_weights is None:
|
||||
bin_label_weights = valid_mask
|
||||
else:
|
||||
bin_label_weights = label_weights.unsqueeze(1).expand(target_shape)
|
||||
bin_label_weights = bin_label_weights * valid_mask
|
||||
|
||||
return bin_labels, bin_label_weights, valid_mask
|
||||
|
||||
|
||||
def binary_cross_entropy(pred,
|
||||
label,
|
||||
weight=None,
|
||||
reduction='mean',
|
||||
avg_factor=None,
|
||||
class_weight=None,
|
||||
ignore_index=-100,
|
||||
avg_non_ignore=False,
|
||||
**kwargs):
|
||||
"""Calculate the binary CrossEntropy loss.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction with shape (N, 1).
|
||||
label (torch.Tensor): The learning label of the prediction.
|
||||
Note: In bce loss, label < 0 is invalid.
|
||||
weight (torch.Tensor, optional): Sample-wise loss weight.
|
||||
reduction (str, optional): The method used to reduce the loss.
|
||||
Options are "none", "mean" and "sum".
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
class_weight (list[float], optional): The weight for each class.
|
||||
ignore_index (int): The label index to be ignored. Default: -100.
|
||||
avg_non_ignore (bool): The flag decides to whether the loss is
|
||||
only averaged over non-ignored targets. Default: False.
|
||||
`New in version 0.23.0.`
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The calculated loss
|
||||
"""
|
||||
if pred.size(1) == 1:
|
||||
# For binary class segmentation, the shape of pred is
|
||||
# [N, 1, H, W] and that of label is [N, H, W].
|
||||
assert label.max() <= 1, \
|
||||
'For pred with shape [N, 1, H, W], its label must have at ' \
|
||||
'most 2 classes'
|
||||
pred = pred.squeeze()
|
||||
if pred.dim() != label.dim():
|
||||
assert (pred.dim() == 2 and label.dim() == 1) or (
|
||||
pred.dim() == 4 and label.dim() == 3), \
|
||||
'Only pred shape [N, C], label shape [N] or pred shape [N, C, ' \
|
||||
'H, W], label shape [N, H, W] are supported'
|
||||
# `weight` returned from `_expand_onehot_labels`
|
||||
# has been treated for valid (non-ignore) pixels
|
||||
label, weight, valid_mask = _expand_onehot_labels(
|
||||
label, weight, pred.shape, ignore_index)
|
||||
else:
|
||||
# should mask out the ignored elements
|
||||
valid_mask = ((label >= 0) & (label != ignore_index)).float()
|
||||
if weight is not None:
|
||||
weight = weight * valid_mask
|
||||
else:
|
||||
weight = valid_mask
|
||||
# average loss over non-ignored and valid elements
|
||||
if reduction == 'mean' and avg_factor is None and avg_non_ignore:
|
||||
avg_factor = valid_mask.sum().item()
|
||||
|
||||
loss = F.binary_cross_entropy_with_logits(
|
||||
pred, label.float(), pos_weight=class_weight, reduction='none')
|
||||
# do the reduction for the weighted loss
|
||||
loss = weight_reduce_loss(
|
||||
loss, weight, reduction=reduction, avg_factor=avg_factor)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def mask_cross_entropy(pred,
|
||||
target,
|
||||
label,
|
||||
reduction='mean',
|
||||
avg_factor=None,
|
||||
class_weight=None,
|
||||
ignore_index=None,
|
||||
**kwargs):
|
||||
"""Calculate the CrossEntropy loss for masks.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction with shape (N, C), C is the number
|
||||
of classes.
|
||||
target (torch.Tensor): The learning label of the prediction.
|
||||
label (torch.Tensor): ``label`` indicates the class label of the mask'
|
||||
corresponding object. This will be used to select the mask in the
|
||||
of the class which the object belongs to when the mask prediction
|
||||
if not class-agnostic.
|
||||
reduction (str, optional): The method used to reduce the loss.
|
||||
Options are "none", "mean" and "sum".
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
class_weight (list[float], optional): The weight for each class.
|
||||
ignore_index (None): Placeholder, to be consistent with other loss.
|
||||
Default: None.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The calculated loss
|
||||
"""
|
||||
assert ignore_index is None, 'BCE loss does not support ignore_index'
|
||||
# TODO: handle these two reserved arguments
|
||||
assert reduction == 'mean' and avg_factor is None
|
||||
num_rois = pred.size()[0]
|
||||
inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device)
|
||||
pred_slice = pred[inds, label].squeeze(1)
|
||||
return F.binary_cross_entropy_with_logits(
|
||||
pred_slice, target, weight=class_weight, reduction='mean')[None]
|
||||
|
||||
|
||||
@LOSSES.register_module(force=True)
|
||||
class CrossEntropyLoss(nn.Module):
|
||||
"""CrossEntropyLoss.
|
||||
|
||||
Args:
|
||||
use_sigmoid (bool, optional): Whether the prediction uses sigmoid
|
||||
of softmax. Defaults to False.
|
||||
use_mask (bool, optional): Whether to use mask cross entropy loss.
|
||||
Defaults to False.
|
||||
reduction (str, optional): . Defaults to 'mean'.
|
||||
Options are "none", "mean" and "sum".
|
||||
class_weight (list[float] | str, optional): Weight of each class. If in
|
||||
str format, read them from a file. Defaults to None.
|
||||
loss_weight (float, optional): Weight of the loss. Defaults to 1.0.
|
||||
loss_name (str, optional): Name of the loss item. If you want this loss
|
||||
item to be included into the backward graph, `loss_` must be the
|
||||
prefix of the name. Defaults to 'loss_ce'.
|
||||
avg_non_ignore (bool): The flag decides to whether the loss is
|
||||
only averaged over non-ignored targets. Default: False.
|
||||
`New in version 0.23.0.`
|
||||
"""
|
||||
def __init__(self,
|
||||
use_sigmoid=False,
|
||||
use_mask=False,
|
||||
reduction='mean',
|
||||
class_weight=None,
|
||||
loss_weight=1.0,
|
||||
loss_name='loss_ce',
|
||||
avg_non_ignore=False):
|
||||
super(CrossEntropyLoss, self).__init__()
|
||||
assert (use_sigmoid is False) or (use_mask is False)
|
||||
self.use_sigmoid = use_sigmoid
|
||||
self.use_mask = use_mask
|
||||
self.reduction = reduction
|
||||
self.loss_weight = loss_weight
|
||||
self.class_weight = get_class_weight(class_weight)
|
||||
self.avg_non_ignore = avg_non_ignore
|
||||
if not self.avg_non_ignore and self.reduction == 'mean':
|
||||
warnings.warn(
|
||||
'Default ``avg_non_ignore`` is False, if you would like to '
|
||||
'ignore the certain label and average loss over non-ignore '
|
||||
'labels, which is the same with PyTorch official '
|
||||
'cross_entropy, set ``avg_non_ignore=True``.')
|
||||
|
||||
if self.use_sigmoid:
|
||||
self.cls_criterion = binary_cross_entropy
|
||||
elif self.use_mask:
|
||||
self.cls_criterion = mask_cross_entropy
|
||||
else:
|
||||
self.cls_criterion = cross_entropy
|
||||
self._loss_name = loss_name
|
||||
|
||||
def extra_repr(self):
|
||||
"""Extra repr."""
|
||||
s = f'avg_non_ignore={self.avg_non_ignore}'
|
||||
return s
|
||||
|
||||
def forward(self,
|
||||
cls_score,
|
||||
label,
|
||||
weight=None,
|
||||
avg_factor=None,
|
||||
reduction_override=None,
|
||||
ignore_index=-100,
|
||||
**kwargs):
|
||||
"""Forward function."""
|
||||
assert reduction_override in (None, 'none', 'mean', 'sum')
|
||||
reduction = (reduction_override
|
||||
if reduction_override else self.reduction)
|
||||
if self.class_weight is not None:
|
||||
class_weight = cls_score.new_tensor(self.class_weight)
|
||||
else:
|
||||
class_weight = None
|
||||
# Note: for BCE loss, label < 0 is invalid.
|
||||
loss_cls = self.loss_weight * self.cls_criterion(
|
||||
cls_score,
|
||||
label,
|
||||
weight,
|
||||
class_weight=class_weight,
|
||||
reduction=reduction,
|
||||
avg_factor=avg_factor,
|
||||
avg_non_ignore=self.avg_non_ignore,
|
||||
ignore_index=ignore_index,
|
||||
**kwargs)
|
||||
return loss_cls
|
||||
|
||||
@property
|
||||
def loss_name(self):
|
||||
"""Loss Name.
|
||||
|
||||
This function must be implemented and will return the name of this
|
||||
loss function. This name will be used to combine different loss items
|
||||
by simple sum operation. In addition, if you want this loss item to be
|
||||
included into the backward graph, `loss_` must be the prefix of the
|
||||
name.
|
||||
|
||||
Returns:
|
||||
str: The name of this loss item.
|
||||
"""
|
||||
return self._loss_name
|
||||
179
segmentation/mmseg_custom/models/losses/dice_loss.py
Normal file
179
segmentation/mmseg_custom/models/losses/dice_loss.py
Normal file
@@ -0,0 +1,179 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from mmseg.models.builder import LOSSES
|
||||
from mmseg.models.losses.utils import weight_reduce_loss
|
||||
|
||||
|
||||
def dice_loss(pred,
|
||||
target,
|
||||
weight=None,
|
||||
eps=1e-3,
|
||||
reduction='mean',
|
||||
avg_factor=None):
|
||||
"""Calculate dice loss, which is proposed in
|
||||
`V-Net: Fully Convolutional Neural Networks for Volumetric
|
||||
Medical Image Segmentation <https://arxiv.org/abs/1606.04797>`_.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction, has a shape (n, *)
|
||||
target (torch.Tensor): The learning label of the prediction,
|
||||
shape (n, *), same shape of pred.
|
||||
weight (torch.Tensor, optional): The weight of loss for each
|
||||
prediction, has a shape (n,). Defaults to None.
|
||||
eps (float): Avoid dividing by zero. Default: 1e-3.
|
||||
reduction (str, optional): The method used to reduce the loss into
|
||||
a scalar. Defaults to 'mean'.
|
||||
Options are "none", "mean" and "sum".
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
"""
|
||||
|
||||
input = pred.flatten(1)
|
||||
target = target.flatten(1).float()
|
||||
|
||||
a = torch.sum(input * target, 1)
|
||||
b = torch.sum(input * input, 1) + eps
|
||||
c = torch.sum(target * target, 1) + eps
|
||||
d = (2 * a) / (b + c)
|
||||
loss = 1 - d
|
||||
if weight is not None:
|
||||
assert weight.ndim == loss.ndim
|
||||
assert len(weight) == len(pred)
|
||||
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
||||
return loss
|
||||
|
||||
|
||||
def naive_dice_loss(pred,
|
||||
target,
|
||||
weight=None,
|
||||
eps=1e-3,
|
||||
reduction='mean',
|
||||
avg_factor=None):
|
||||
"""Calculate naive dice loss, the coefficient in the denominator is the
|
||||
first power instead of the second power.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction, has a shape (n, *)
|
||||
target (torch.Tensor): The learning label of the prediction,
|
||||
shape (n, *), same shape of pred.
|
||||
weight (torch.Tensor, optional): The weight of loss for each
|
||||
prediction, has a shape (n,). Defaults to None.
|
||||
eps (float): Avoid dividing by zero. Default: 1e-3.
|
||||
reduction (str, optional): The method used to reduce the loss into
|
||||
a scalar. Defaults to 'mean'.
|
||||
Options are "none", "mean" and "sum".
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
"""
|
||||
input = pred.flatten(1)
|
||||
target = target.flatten(1).float()
|
||||
|
||||
a = torch.sum(input * target, 1)
|
||||
b = torch.sum(input, 1)
|
||||
c = torch.sum(target, 1)
|
||||
d = (2 * a + eps) / (b + c + eps)
|
||||
loss = 1 - d
|
||||
if weight is not None:
|
||||
assert weight.ndim == loss.ndim
|
||||
assert len(weight) == len(pred)
|
||||
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
||||
return loss
|
||||
|
||||
|
||||
@LOSSES.register_module(force=True)
|
||||
class DiceLoss(nn.Module):
|
||||
def __init__(self,
|
||||
use_sigmoid=True,
|
||||
activate=True,
|
||||
reduction='mean',
|
||||
naive_dice=False,
|
||||
loss_weight=1.0,
|
||||
eps=1e-3):
|
||||
"""Dice Loss, there are two forms of dice loss is supported:
|
||||
|
||||
- the one proposed in `V-Net: Fully Convolutional Neural
|
||||
Networks for Volumetric Medical Image Segmentation
|
||||
<https://arxiv.org/abs/1606.04797>`_.
|
||||
- the dice loss in which the power of the number in the
|
||||
denominator is the first power instead of the second
|
||||
power.
|
||||
|
||||
Args:
|
||||
use_sigmoid (bool, optional): Whether to the prediction is
|
||||
used for sigmoid or softmax. Defaults to True.
|
||||
activate (bool): Whether to activate the predictions inside,
|
||||
this will disable the inside sigmoid operation.
|
||||
Defaults to True.
|
||||
reduction (str, optional): The method used
|
||||
to reduce the loss. Options are "none",
|
||||
"mean" and "sum". Defaults to 'mean'.
|
||||
naive_dice (bool, optional): If false, use the dice
|
||||
loss defined in the V-Net paper, otherwise, use the
|
||||
naive dice loss in which the power of the number in the
|
||||
denominator is the first power instead of the second
|
||||
power.Defaults to False.
|
||||
loss_weight (float, optional): Weight of loss. Defaults to 1.0.
|
||||
eps (float): Avoid dividing by zero. Defaults to 1e-3.
|
||||
"""
|
||||
|
||||
super(DiceLoss, self).__init__()
|
||||
self.use_sigmoid = use_sigmoid
|
||||
self.reduction = reduction
|
||||
self.naive_dice = naive_dice
|
||||
self.loss_weight = loss_weight
|
||||
self.eps = eps
|
||||
self.activate = activate
|
||||
|
||||
def forward(self,
|
||||
pred,
|
||||
target,
|
||||
weight=None,
|
||||
reduction_override=None,
|
||||
avg_factor=None):
|
||||
"""Forward function.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction, has a shape (n, *).
|
||||
target (torch.Tensor): The label of the prediction,
|
||||
shape (n, *), same shape of pred.
|
||||
weight (torch.Tensor, optional): The weight of loss for each
|
||||
prediction, has a shape (n,). Defaults to None.
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
reduction_override (str, optional): The reduction method used to
|
||||
override the original reduction method of the loss.
|
||||
Options are "none", "mean" and "sum".
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The calculated loss
|
||||
"""
|
||||
|
||||
assert reduction_override in (None, 'none', 'mean', 'sum')
|
||||
reduction = (reduction_override
|
||||
if reduction_override else self.reduction)
|
||||
|
||||
if self.activate:
|
||||
if self.use_sigmoid:
|
||||
pred = pred.sigmoid()
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
if self.naive_dice:
|
||||
loss = self.loss_weight * naive_dice_loss(
|
||||
pred,
|
||||
target,
|
||||
weight,
|
||||
eps=self.eps,
|
||||
reduction=reduction,
|
||||
avg_factor=avg_factor)
|
||||
else:
|
||||
loss = self.loss_weight * dice_loss(
|
||||
pred,
|
||||
target,
|
||||
weight,
|
||||
eps=self.eps,
|
||||
reduction=reduction,
|
||||
avg_factor=avg_factor)
|
||||
|
||||
return loss
|
||||
180
segmentation/mmseg_custom/models/losses/focal_loss.py
Normal file
180
segmentation/mmseg_custom/models/losses/focal_loss.py
Normal file
@@ -0,0 +1,180 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmcv.ops import sigmoid_focal_loss as _sigmoid_focal_loss
|
||||
from mmseg.models.builder import LOSSES
|
||||
from mmseg.models.losses.utils import weight_reduce_loss
|
||||
|
||||
|
||||
# This method is only for debugging
|
||||
def py_sigmoid_focal_loss(pred,
|
||||
target,
|
||||
weight=None,
|
||||
gamma=2.0,
|
||||
alpha=0.25,
|
||||
reduction='mean',
|
||||
avg_factor=None):
|
||||
"""PyTorch version of `Focal Loss <https://arxiv.org/abs/1708.02002>`_.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction with shape (N, C), C is the
|
||||
number of classes
|
||||
target (torch.Tensor): The learning label of the prediction.
|
||||
weight (torch.Tensor, optional): Sample-wise loss weight.
|
||||
gamma (float, optional): The gamma for calculating the modulating
|
||||
factor. Defaults to 2.0.
|
||||
alpha (float, optional): A balanced form for Focal Loss.
|
||||
Defaults to 0.25.
|
||||
reduction (str, optional): The method used to reduce the loss into
|
||||
a scalar. Defaults to 'mean'.
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
"""
|
||||
pred_sigmoid = pred.sigmoid()
|
||||
target = target.type_as(pred)
|
||||
pt = (1 - pred_sigmoid) * target + pred_sigmoid * (1 - target)
|
||||
focal_weight = (alpha * target + (1 - alpha) *
|
||||
(1 - target)) * pt.pow(gamma)
|
||||
loss = F.binary_cross_entropy_with_logits(
|
||||
pred, target, reduction='none') * focal_weight
|
||||
if weight is not None:
|
||||
if weight.shape != loss.shape:
|
||||
if weight.size(0) == loss.size(0):
|
||||
# For most cases, weight is of shape (num_priors, ),
|
||||
# which means it does not have the second axis num_class
|
||||
weight = weight.view(-1, 1)
|
||||
else:
|
||||
# Sometimes, weight per anchor per class is also needed. e.g.
|
||||
# in FSAF. But it may be flattened of shape
|
||||
# (num_priors x num_class, ), while loss is still of shape
|
||||
# (num_priors, num_class).
|
||||
assert weight.numel() == loss.numel()
|
||||
weight = weight.view(loss.size(0), -1)
|
||||
assert weight.ndim == loss.ndim
|
||||
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
||||
return loss
|
||||
|
||||
|
||||
def sigmoid_focal_loss(pred,
|
||||
target,
|
||||
weight=None,
|
||||
gamma=2.0,
|
||||
alpha=0.25,
|
||||
reduction='mean',
|
||||
avg_factor=None):
|
||||
r"""A warpper of cuda version `Focal Loss
|
||||
<https://arxiv.org/abs/1708.02002>`_.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction with shape (N, C), C is the number
|
||||
of classes.
|
||||
target (torch.Tensor): The learning label of the prediction.
|
||||
weight (torch.Tensor, optional): Sample-wise loss weight.
|
||||
gamma (float, optional): The gamma for calculating the modulating
|
||||
factor. Defaults to 2.0.
|
||||
alpha (float, optional): A balanced form for Focal Loss.
|
||||
Defaults to 0.25.
|
||||
reduction (str, optional): The method used to reduce the loss into
|
||||
a scalar. Defaults to 'mean'. Options are "none", "mean" and "sum".
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
"""
|
||||
# Function.apply does not accept keyword arguments, so the decorator
|
||||
# "weighted_loss" is not applicable
|
||||
loss = _sigmoid_focal_loss(pred.contiguous(), target.contiguous(), gamma,
|
||||
alpha, None, 'none')
|
||||
if weight is not None:
|
||||
if weight.shape != loss.shape:
|
||||
if weight.size(0) == loss.size(0):
|
||||
# For most cases, weight is of shape (num_priors, ),
|
||||
# which means it does not have the second axis num_class
|
||||
weight = weight.view(-1, 1)
|
||||
else:
|
||||
# Sometimes, weight per anchor per class is also needed. e.g.
|
||||
# in FSAF. But it may be flattened of shape
|
||||
# (num_priors x num_class, ), while loss is still of shape
|
||||
# (num_priors, num_class).
|
||||
assert weight.numel() == loss.numel()
|
||||
weight = weight.view(loss.size(0), -1)
|
||||
assert weight.ndim == loss.ndim
|
||||
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
|
||||
return loss
|
||||
|
||||
|
||||
@LOSSES.register_module(force=True)
|
||||
class FocalLoss(nn.Module):
|
||||
def __init__(self,
|
||||
use_sigmoid=True,
|
||||
gamma=2.0,
|
||||
alpha=0.25,
|
||||
reduction='mean',
|
||||
loss_weight=1.0):
|
||||
"""`Focal Loss <https://arxiv.org/abs/1708.02002>`_
|
||||
|
||||
Args:
|
||||
use_sigmoid (bool, optional): Whether to the prediction is
|
||||
used for sigmoid or softmax. Defaults to True.
|
||||
gamma (float, optional): The gamma for calculating the modulating
|
||||
factor. Defaults to 2.0.
|
||||
alpha (float, optional): A balanced form for Focal Loss.
|
||||
Defaults to 0.25.
|
||||
reduction (str, optional): The method used to reduce the loss into
|
||||
a scalar. Defaults to 'mean'. Options are "none", "mean" and
|
||||
"sum".
|
||||
loss_weight (float, optional): Weight of loss. Defaults to 1.0.
|
||||
"""
|
||||
super(FocalLoss, self).__init__()
|
||||
assert use_sigmoid is True, 'Only sigmoid focal loss supported now.'
|
||||
self.use_sigmoid = use_sigmoid
|
||||
self.gamma = gamma
|
||||
self.alpha = alpha
|
||||
self.reduction = reduction
|
||||
self.loss_weight = loss_weight
|
||||
|
||||
def forward(self,
|
||||
pred,
|
||||
target,
|
||||
weight=None,
|
||||
avg_factor=None,
|
||||
reduction_override=None):
|
||||
"""Forward function.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The prediction.
|
||||
target (torch.Tensor): The learning label of the prediction.
|
||||
weight (torch.Tensor, optional): The weight of loss for each
|
||||
prediction. Defaults to None.
|
||||
avg_factor (int, optional): Average factor that is used to average
|
||||
the loss. Defaults to None.
|
||||
reduction_override (str, optional): The reduction method used to
|
||||
override the original reduction method of the loss.
|
||||
Options are "none", "mean" and "sum".
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The calculated loss
|
||||
"""
|
||||
assert reduction_override in (None, 'none', 'mean', 'sum')
|
||||
reduction = (
|
||||
reduction_override if reduction_override else self.reduction)
|
||||
if self.use_sigmoid:
|
||||
if torch.cuda.is_available() and pred.is_cuda:
|
||||
calculate_loss_func = sigmoid_focal_loss
|
||||
else:
|
||||
num_classes = pred.size(1)
|
||||
target = F.one_hot(target, num_classes=num_classes + 1)
|
||||
target = target[:, :num_classes]
|
||||
calculate_loss_func = py_sigmoid_focal_loss
|
||||
|
||||
loss_cls = self.loss_weight * calculate_loss_func(
|
||||
pred,
|
||||
target,
|
||||
weight,
|
||||
gamma=self.gamma,
|
||||
alpha=self.alpha,
|
||||
reduction=reduction,
|
||||
avg_factor=avg_factor)
|
||||
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return loss_cls
|
||||
233
segmentation/mmseg_custom/models/losses/match_costs.py
Normal file
233
segmentation/mmseg_custom/models/losses/match_costs.py
Normal file
@@ -0,0 +1,233 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ..builder import MATCH_COST
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class FocalLossCost:
|
||||
"""FocalLossCost.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight
|
||||
alpha (int | float, optional): focal_loss alpha
|
||||
gamma (int | float, optional): focal_loss gamma
|
||||
eps (float, optional): default 1e-12
|
||||
|
||||
Examples:
|
||||
>>> from mmdet.core.bbox.match_costs.match_cost import FocalLossCost
|
||||
>>> import torch
|
||||
>>> self = FocalLossCost()
|
||||
>>> cls_pred = torch.rand(4, 3)
|
||||
>>> gt_labels = torch.tensor([0, 1, 2])
|
||||
>>> factor = torch.tensor([10, 8, 10, 8])
|
||||
>>> self(cls_pred, gt_labels)
|
||||
tensor([[-0.3236, -0.3364, -0.2699],
|
||||
[-0.3439, -0.3209, -0.4807],
|
||||
[-0.4099, -0.3795, -0.2929],
|
||||
[-0.1950, -0.1207, -0.2626]])
|
||||
"""
|
||||
def __init__(self, weight=1., alpha=0.25, gamma=2, eps=1e-12):
|
||||
self.weight = weight
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classification logits, shape
|
||||
[num_query, num_class].
|
||||
gt_labels (Tensor): Label of `gt_bboxes`, shape (num_gt,).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: cls_cost value with weight
|
||||
"""
|
||||
cls_pred = cls_pred.sigmoid()
|
||||
neg_cost = -(1 - cls_pred + self.eps).log() * (
|
||||
1 - self.alpha) * cls_pred.pow(self.gamma)
|
||||
pos_cost = -(cls_pred + self.eps).log() * self.alpha * (
|
||||
1 - cls_pred).pow(self.gamma)
|
||||
cls_cost = pos_cost[:, gt_labels] - neg_cost[:, gt_labels]
|
||||
return cls_cost * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class MaskFocalLossCost(FocalLossCost):
|
||||
"""Cost of mask assignments based on focal losses.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight.
|
||||
alpha (int | float, optional): focal_loss alpha.
|
||||
gamma (int | float, optional): focal_loss gamma.
|
||||
eps (float, optional): default 1e-12.
|
||||
"""
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classfication logits
|
||||
in shape (N1, H, W), dtype=torch.float32.
|
||||
gt_labels (Tensor): Ground truth in shape (N2, H, W),
|
||||
dtype=torch.long.
|
||||
|
||||
Returns:
|
||||
Tensor: classification cost matrix in shape (N1, N2).
|
||||
"""
|
||||
cls_pred = cls_pred.reshape((cls_pred.shape[0], -1))
|
||||
gt_labels = gt_labels.reshape((gt_labels.shape[0], -1)).float()
|
||||
hw = cls_pred.shape[1]
|
||||
cls_pred = cls_pred.sigmoid()
|
||||
neg_cost = -(1 - cls_pred + self.eps).log() * (
|
||||
1 - self.alpha) * cls_pred.pow(self.gamma)
|
||||
pos_cost = -(cls_pred + self.eps).log() * self.alpha * (
|
||||
1 - cls_pred).pow(self.gamma)
|
||||
|
||||
cls_cost = torch.einsum('nc,mc->nm', pos_cost, gt_labels) + \
|
||||
torch.einsum('nc,mc->nm', neg_cost, (1 - gt_labels))
|
||||
return cls_cost / hw * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class ClassificationCost:
|
||||
"""ClsSoftmaxCost.Borrow from
|
||||
mmdet.core.bbox.match_costs.match_cost.ClassificationCost.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight
|
||||
|
||||
Examples:
|
||||
>>> import torch
|
||||
>>> self = ClassificationCost()
|
||||
>>> cls_pred = torch.rand(4, 3)
|
||||
>>> gt_labels = torch.tensor([0, 1, 2])
|
||||
>>> factor = torch.tensor([10, 8, 10, 8])
|
||||
>>> self(cls_pred, gt_labels)
|
||||
tensor([[-0.3430, -0.3525, -0.3045],
|
||||
[-0.3077, -0.2931, -0.3992],
|
||||
[-0.3664, -0.3455, -0.2881],
|
||||
[-0.3343, -0.2701, -0.3956]])
|
||||
"""
|
||||
def __init__(self, weight=1.):
|
||||
self.weight = weight
|
||||
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classification logits, shape
|
||||
[num_query, num_class].
|
||||
gt_labels (Tensor): Label of `gt_bboxes`, shape (num_gt,).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: cls_cost value with weight
|
||||
"""
|
||||
# Following the official DETR repo, contrary to the loss that
|
||||
# NLL is used, we approximate it in 1 - cls_score[gt_label].
|
||||
# The 1 is a constant that doesn't change the matching,
|
||||
# so it can be omitted.
|
||||
cls_score = cls_pred.softmax(-1)
|
||||
cls_cost = -cls_score[:, gt_labels]
|
||||
return cls_cost * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class DiceCost:
|
||||
"""Cost of mask assignments based on dice losses.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight. Defaults to 1.
|
||||
pred_act (bool, optional): Whether to apply sigmoid to mask_pred.
|
||||
Defaults to False.
|
||||
eps (float, optional): default 1e-12.
|
||||
"""
|
||||
def __init__(self, weight=1., pred_act=False, eps=1e-3):
|
||||
self.weight = weight
|
||||
self.pred_act = pred_act
|
||||
self.eps = eps
|
||||
|
||||
def binary_mask_dice_loss(self, mask_preds, gt_masks):
|
||||
"""
|
||||
Args:
|
||||
mask_preds (Tensor): Mask prediction in shape (N1, H, W).
|
||||
gt_masks (Tensor): Ground truth in shape (N2, H, W)
|
||||
store 0 or 1, 0 for negative class and 1 for
|
||||
positive class.
|
||||
|
||||
Returns:
|
||||
Tensor: Dice cost matrix in shape (N1, N2).
|
||||
"""
|
||||
mask_preds = mask_preds.reshape((mask_preds.shape[0], -1))
|
||||
gt_masks = gt_masks.reshape((gt_masks.shape[0], -1)).float()
|
||||
numerator = 2 * torch.einsum('nc,mc->nm', mask_preds, gt_masks)
|
||||
denominator = mask_preds.sum(-1)[:, None] + gt_masks.sum(-1)[None, :]
|
||||
loss = 1 - (numerator + self.eps) / (denominator + self.eps)
|
||||
return loss
|
||||
|
||||
def __call__(self, mask_preds, gt_masks):
|
||||
"""
|
||||
Args:
|
||||
mask_preds (Tensor): Mask prediction logits in shape (N1, H, W).
|
||||
gt_masks (Tensor): Ground truth in shape (N2, H, W).
|
||||
|
||||
Returns:
|
||||
Tensor: Dice cost matrix in shape (N1, N2).
|
||||
"""
|
||||
if self.pred_act:
|
||||
mask_preds = mask_preds.sigmoid()
|
||||
dice_cost = self.binary_mask_dice_loss(mask_preds, gt_masks)
|
||||
return dice_cost * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class CrossEntropyLossCost:
|
||||
"""CrossEntropyLossCost.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss weight. Defaults to 1.
|
||||
use_sigmoid (bool, optional): Whether the prediction uses sigmoid
|
||||
of softmax. Defaults to True.
|
||||
"""
|
||||
def __init__(self, weight=1., use_sigmoid=True):
|
||||
assert use_sigmoid, 'use_sigmoid = False is not supported yet.'
|
||||
self.weight = weight
|
||||
self.use_sigmoid = use_sigmoid
|
||||
|
||||
def _binary_cross_entropy(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): The prediction with shape (num_query, 1, *) or
|
||||
(num_query, *).
|
||||
gt_labels (Tensor): The learning label of prediction with
|
||||
shape (num_gt, *).
|
||||
Returns:
|
||||
Tensor: Cross entropy cost matrix in shape (num_query, num_gt).
|
||||
"""
|
||||
cls_pred = cls_pred.flatten(1).float()
|
||||
gt_labels = gt_labels.flatten(1).float()
|
||||
n = cls_pred.shape[1]
|
||||
pos = F.binary_cross_entropy_with_logits(
|
||||
cls_pred, torch.ones_like(cls_pred), reduction='none')
|
||||
neg = F.binary_cross_entropy_with_logits(
|
||||
cls_pred, torch.zeros_like(cls_pred), reduction='none')
|
||||
cls_cost = torch.einsum('nc,mc->nm', pos, gt_labels) + \
|
||||
torch.einsum('nc,mc->nm', neg, 1 - gt_labels)
|
||||
cls_cost = cls_cost / n
|
||||
|
||||
return cls_cost
|
||||
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classification logits.
|
||||
gt_labels (Tensor): Labels.
|
||||
Returns:
|
||||
Tensor: Cross entropy cost matrix with weight in
|
||||
shape (num_query, num_gt).
|
||||
"""
|
||||
if self.use_sigmoid:
|
||||
cls_cost = self._binary_cross_entropy(cls_pred, gt_labels)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return cls_cost * self.weight
|
||||
179
segmentation/mmseg_custom/models/losses/match_loss.py
Normal file
179
segmentation/mmseg_custom/models/losses/match_loss.py
Normal file
@@ -0,0 +1,179 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ..builder import MATCH_COST
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class FocalLossCost:
|
||||
"""FocalLossCost.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight
|
||||
alpha (int | float, optional): focal_loss alpha
|
||||
gamma (int | float, optional): focal_loss gamma
|
||||
eps (float, optional): default 1e-12
|
||||
|
||||
Examples:
|
||||
>>> from mmdet.core.bbox.match_costs.match_cost import FocalLossCost
|
||||
>>> import torch
|
||||
>>> self = FocalLossCost()
|
||||
>>> cls_pred = torch.rand(4, 3)
|
||||
>>> gt_labels = torch.tensor([0, 1, 2])
|
||||
>>> factor = torch.tensor([10, 8, 10, 8])
|
||||
>>> self(cls_pred, gt_labels)
|
||||
tensor([[-0.3236, -0.3364, -0.2699],
|
||||
[-0.3439, -0.3209, -0.4807],
|
||||
[-0.4099, -0.3795, -0.2929],
|
||||
[-0.1950, -0.1207, -0.2626]])
|
||||
"""
|
||||
def __init__(self, weight=1., alpha=0.25, gamma=2, eps=1e-12):
|
||||
self.weight = weight
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classification logits, shape
|
||||
[num_query, num_class].
|
||||
gt_labels (Tensor): Label of `gt_bboxes`, shape (num_gt,).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: cls_cost value with weight
|
||||
"""
|
||||
cls_pred = cls_pred.sigmoid()
|
||||
neg_cost = -(1 - cls_pred + self.eps).log() * (
|
||||
1 - self.alpha) * cls_pred.pow(self.gamma)
|
||||
pos_cost = -(cls_pred + self.eps).log() * self.alpha * (
|
||||
1 - cls_pred).pow(self.gamma)
|
||||
cls_cost = pos_cost[:, gt_labels] - neg_cost[:, gt_labels]
|
||||
return cls_cost * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class MaskFocalLossCost(FocalLossCost):
|
||||
"""Cost of mask assignments based on focal losses.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight.
|
||||
alpha (int | float, optional): focal_loss alpha.
|
||||
gamma (int | float, optional): focal_loss gamma.
|
||||
eps (float, optional): default 1e-12.
|
||||
"""
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classfication logits
|
||||
in shape (N1, H, W), dtype=torch.float32.
|
||||
gt_labels (Tensor): Ground truth in shape (N2, H, W),
|
||||
dtype=torch.long.
|
||||
|
||||
Returns:
|
||||
Tensor: classification cost matrix in shape (N1, N2).
|
||||
"""
|
||||
cls_pred = cls_pred.reshape((cls_pred.shape[0], -1))
|
||||
gt_labels = gt_labels.reshape((gt_labels.shape[0], -1)).float()
|
||||
hw = cls_pred.shape[1]
|
||||
cls_pred = cls_pred.sigmoid()
|
||||
neg_cost = -(1 - cls_pred + self.eps).log() * (
|
||||
1 - self.alpha) * cls_pred.pow(self.gamma)
|
||||
pos_cost = -(cls_pred + self.eps).log() * self.alpha * (
|
||||
1 - cls_pred).pow(self.gamma)
|
||||
|
||||
cls_cost = torch.einsum('nc,mc->nm', pos_cost, gt_labels) + \
|
||||
torch.einsum('nc,mc->nm', neg_cost, (1 - gt_labels))
|
||||
return cls_cost / hw * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class ClassificationCost:
|
||||
"""ClsSoftmaxCost.Borrow from
|
||||
mmdet.core.bbox.match_costs.match_cost.ClassificationCost.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight
|
||||
|
||||
Examples:
|
||||
>>> import torch
|
||||
>>> self = ClassificationCost()
|
||||
>>> cls_pred = torch.rand(4, 3)
|
||||
>>> gt_labels = torch.tensor([0, 1, 2])
|
||||
>>> factor = torch.tensor([10, 8, 10, 8])
|
||||
>>> self(cls_pred, gt_labels)
|
||||
tensor([[-0.3430, -0.3525, -0.3045],
|
||||
[-0.3077, -0.2931, -0.3992],
|
||||
[-0.3664, -0.3455, -0.2881],
|
||||
[-0.3343, -0.2701, -0.3956]])
|
||||
"""
|
||||
def __init__(self, weight=1.):
|
||||
self.weight = weight
|
||||
|
||||
def __call__(self, cls_pred, gt_labels):
|
||||
"""
|
||||
Args:
|
||||
cls_pred (Tensor): Predicted classification logits, shape
|
||||
[num_query, num_class].
|
||||
gt_labels (Tensor): Label of `gt_bboxes`, shape (num_gt,).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: cls_cost value with weight
|
||||
"""
|
||||
# Following the official DETR repo, contrary to the loss that
|
||||
# NLL is used, we approximate it in 1 - cls_score[gt_label].
|
||||
# The 1 is a constant that doesn't change the matching,
|
||||
# so it can be omitted.
|
||||
cls_score = cls_pred.softmax(-1)
|
||||
cls_cost = -cls_score[:, gt_labels]
|
||||
return cls_cost * self.weight
|
||||
|
||||
|
||||
@MATCH_COST.register_module()
|
||||
class DiceCost:
|
||||
"""Cost of mask assignments based on dice losses.
|
||||
|
||||
Args:
|
||||
weight (int | float, optional): loss_weight. Defaults to 1.
|
||||
pred_act (bool, optional): Whether to apply sigmoid to mask_pred.
|
||||
Defaults to False.
|
||||
eps (float, optional): default 1e-12.
|
||||
"""
|
||||
def __init__(self, weight=1., pred_act=False, eps=1e-3):
|
||||
self.weight = weight
|
||||
self.pred_act = pred_act
|
||||
self.eps = eps
|
||||
|
||||
def binary_mask_dice_loss(self, mask_preds, gt_masks):
|
||||
"""
|
||||
Args:
|
||||
mask_preds (Tensor): Mask prediction in shape (N1, H, W).
|
||||
gt_masks (Tensor): Ground truth in shape (N2, H, W)
|
||||
store 0 or 1, 0 for negative class and 1 for
|
||||
positive class.
|
||||
|
||||
Returns:
|
||||
Tensor: Dice cost matrix in shape (N1, N2).
|
||||
"""
|
||||
mask_preds = mask_preds.reshape((mask_preds.shape[0], -1))
|
||||
gt_masks = gt_masks.reshape((gt_masks.shape[0], -1)).float()
|
||||
numerator = 2 * torch.einsum('nc,mc->nm', mask_preds, gt_masks)
|
||||
denominator = mask_preds.sum(-1)[:, None] + gt_masks.sum(-1)[None, :]
|
||||
loss = 1 - (numerator + self.eps) / (denominator + self.eps)
|
||||
return loss
|
||||
|
||||
def __call__(self, mask_preds, gt_masks):
|
||||
"""
|
||||
Args:
|
||||
mask_preds (Tensor): Mask prediction logits in shape (N1, H, W).
|
||||
gt_masks (Tensor): Ground truth in shape (N2, H, W).
|
||||
|
||||
Returns:
|
||||
Tensor: Dice cost matrix in shape (N1, N2).
|
||||
"""
|
||||
if self.pred_act:
|
||||
mask_preds = mask_preds.sigmoid()
|
||||
dice_cost = self.binary_mask_dice_loss(mask_preds, gt_masks)
|
||||
return dice_cost * self.weight
|
||||
8
segmentation/mmseg_custom/models/plugins/__init__.py
Normal file
8
segmentation/mmseg_custom/models/plugins/__init__.py
Normal file
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Shanghai AI Lab. All rights reserved.
|
||||
from .msdeformattn_pixel_decoder import MSDeformAttnPixelDecoder
|
||||
from .pixel_decoder import PixelDecoder, TransformerEncoderPixelDecoder
|
||||
|
||||
__all__ = [
|
||||
'PixelDecoder', 'TransformerEncoderPixelDecoder',
|
||||
'MSDeformAttnPixelDecoder'
|
||||
]
|
||||
@@ -0,0 +1,268 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmcv.cnn import (PLUGIN_LAYERS, Conv2d, ConvModule, caffe2_xavier_init,
|
||||
normal_init, xavier_init)
|
||||
from mmcv.cnn.bricks.transformer import (build_positional_encoding,
|
||||
build_transformer_layer_sequence)
|
||||
from mmcv.runner import BaseModule, ModuleList
|
||||
|
||||
from ...core.anchor import MlvlPointGenerator
|
||||
from ..utils.transformer import MultiScaleDeformableAttention
|
||||
|
||||
|
||||
@PLUGIN_LAYERS.register_module()
|
||||
class MSDeformAttnPixelDecoder(BaseModule):
|
||||
"""Pixel decoder with multi-scale deformable attention.
|
||||
|
||||
Args:
|
||||
in_channels (list[int] | tuple[int]): Number of channels in the
|
||||
input feature maps.
|
||||
strides (list[int] | tuple[int]): Output strides of feature from
|
||||
backbone.
|
||||
feat_channels (int): Number of channels for feature.
|
||||
out_channels (int): Number of channels for output.
|
||||
num_outs (int): Number of output scales.
|
||||
norm_cfg (:obj:`mmcv.ConfigDict` | dict): Config for normalization.
|
||||
Defaults to dict(type='GN', num_groups=32).
|
||||
act_cfg (:obj:`mmcv.ConfigDict` | dict): Config for activation.
|
||||
Defaults to dict(type='ReLU').
|
||||
encoder (:obj:`mmcv.ConfigDict` | dict): Config for transformer
|
||||
encoder. Defaults to `DetrTransformerEncoder`.
|
||||
positional_encoding (:obj:`mmcv.ConfigDict` | dict): Config for
|
||||
transformer encoder position encoding. Defaults to
|
||||
dict(type='SinePositionalEncoding', num_feats=128,
|
||||
normalize=True).
|
||||
init_cfg (:obj:`mmcv.ConfigDict` | dict): Initialization config dict.
|
||||
"""
|
||||
def __init__(self,
|
||||
in_channels=[256, 512, 1024, 2048],
|
||||
strides=[4, 8, 16, 32],
|
||||
feat_channels=256,
|
||||
out_channels=256,
|
||||
num_outs=3,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=dict(
|
||||
type='DetrTransformerEncoder',
|
||||
num_layers=6,
|
||||
transformerlayers=dict(
|
||||
type='BaseTransformerLayer',
|
||||
attn_cfgs=dict(
|
||||
type='MultiScaleDeformableAttention',
|
||||
embed_dims=256,
|
||||
num_heads=8,
|
||||
num_levels=3,
|
||||
num_points=4,
|
||||
im2col_step=64,
|
||||
dropout=0.0,
|
||||
batch_first=False,
|
||||
norm_cfg=None,
|
||||
init_cfg=None),
|
||||
feedforward_channels=1024,
|
||||
ffn_dropout=0.0,
|
||||
operation_order=('self_attn', 'norm', 'ffn', 'norm')),
|
||||
init_cfg=None),
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding',
|
||||
num_feats=128,
|
||||
normalize=True),
|
||||
init_cfg=None):
|
||||
super().__init__(init_cfg=init_cfg)
|
||||
self.strides = strides
|
||||
self.num_input_levels = len(in_channels)
|
||||
self.num_encoder_levels = \
|
||||
encoder.transformerlayers.attn_cfgs.num_levels
|
||||
assert self.num_encoder_levels >= 1, \
|
||||
'num_levels in attn_cfgs must be at least one'
|
||||
input_conv_list = []
|
||||
# from top to down (low to high resolution)
|
||||
for i in range(self.num_input_levels - 1,
|
||||
self.num_input_levels - self.num_encoder_levels - 1,
|
||||
-1):
|
||||
input_conv = ConvModule(
|
||||
in_channels[i],
|
||||
feat_channels,
|
||||
kernel_size=1,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=None,
|
||||
bias=True)
|
||||
input_conv_list.append(input_conv)
|
||||
self.input_convs = ModuleList(input_conv_list)
|
||||
|
||||
self.encoder = build_transformer_layer_sequence(encoder)
|
||||
self.postional_encoding = build_positional_encoding(
|
||||
positional_encoding)
|
||||
# high resolution to low resolution
|
||||
self.level_encoding = nn.Embedding(self.num_encoder_levels,
|
||||
feat_channels)
|
||||
|
||||
# fpn-like structure
|
||||
self.lateral_convs = ModuleList()
|
||||
self.output_convs = ModuleList()
|
||||
self.use_bias = norm_cfg is None
|
||||
# from top to down (low to high resolution)
|
||||
# fpn for the rest features that didn't pass in encoder
|
||||
for i in range(self.num_input_levels - self.num_encoder_levels - 1, -1,
|
||||
-1):
|
||||
lateral_conv = ConvModule(
|
||||
in_channels[i],
|
||||
feat_channels,
|
||||
kernel_size=1,
|
||||
bias=self.use_bias,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=None)
|
||||
output_conv = ConvModule(
|
||||
feat_channels,
|
||||
feat_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=self.use_bias,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=act_cfg)
|
||||
self.lateral_convs.append(lateral_conv)
|
||||
self.output_convs.append(output_conv)
|
||||
|
||||
self.mask_feature = Conv2d(
|
||||
feat_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
self.num_outs = num_outs
|
||||
self.point_generator = MlvlPointGenerator(strides)
|
||||
|
||||
def init_weights(self):
|
||||
"""Initialize weights."""
|
||||
for i in range(0, self.num_encoder_levels):
|
||||
xavier_init(
|
||||
self.input_convs[i].conv,
|
||||
gain=1,
|
||||
bias=0,
|
||||
distribution='uniform')
|
||||
|
||||
for i in range(0, self.num_input_levels - self.num_encoder_levels):
|
||||
caffe2_xavier_init(self.lateral_convs[i].conv, bias=0)
|
||||
caffe2_xavier_init(self.output_convs[i].conv, bias=0)
|
||||
|
||||
caffe2_xavier_init(self.mask_feature, bias=0)
|
||||
|
||||
normal_init(self.level_encoding, mean=0, std=1)
|
||||
for p in self.encoder.parameters():
|
||||
if p.dim() > 1:
|
||||
nn.init.xavier_normal_(p)
|
||||
|
||||
# init_weights defined in MultiScaleDeformableAttention
|
||||
for layer in self.encoder.layers:
|
||||
for attn in layer.attentions:
|
||||
if isinstance(attn, MultiScaleDeformableAttention):
|
||||
attn.init_weights()
|
||||
|
||||
def forward(self, feats):
|
||||
"""
|
||||
Args:
|
||||
feats (list[Tensor]): Feature maps of each level. Each has
|
||||
shape of (batch_size, c, h, w).
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the following:
|
||||
|
||||
- mask_feature (Tensor): shape (batch_size, c, h, w).
|
||||
- multi_scale_features (list[Tensor]): Multi scale \
|
||||
features, each in shape (batch_size, c, h, w).
|
||||
"""
|
||||
# generate padding mask for each level, for each image
|
||||
batch_size = feats[0].shape[0]
|
||||
encoder_input_list = []
|
||||
padding_mask_list = []
|
||||
level_positional_encoding_list = []
|
||||
spatial_shapes = []
|
||||
reference_points_list = []
|
||||
for i in range(self.num_encoder_levels):
|
||||
level_idx = self.num_input_levels - i - 1
|
||||
feat = feats[level_idx]
|
||||
feat_projected = self.input_convs[i](feat)
|
||||
h, w = feat.shape[-2:]
|
||||
|
||||
# no padding
|
||||
padding_mask_resized = feat.new_zeros(
|
||||
(batch_size, ) + feat.shape[-2:], dtype=torch.bool)
|
||||
pos_embed = self.postional_encoding(padding_mask_resized)
|
||||
level_embed = self.level_encoding.weight[i]
|
||||
level_pos_embed = level_embed.view(1, -1, 1, 1) + pos_embed
|
||||
# (h_i * w_i, 2)
|
||||
reference_points = self.point_generator.single_level_grid_priors(
|
||||
feat.shape[-2:], level_idx, device=feat.device)
|
||||
# normalize
|
||||
factor = feat.new_tensor([[w, h]]) * self.strides[level_idx]
|
||||
reference_points = reference_points / factor
|
||||
|
||||
# shape (batch_size, c, h_i, w_i) -> (h_i * w_i, batch_size, c)
|
||||
feat_projected = feat_projected.flatten(2).permute(2, 0, 1)
|
||||
level_pos_embed = level_pos_embed.flatten(2).permute(2, 0, 1)
|
||||
padding_mask_resized = padding_mask_resized.flatten(1)
|
||||
|
||||
encoder_input_list.append(feat_projected)
|
||||
padding_mask_list.append(padding_mask_resized)
|
||||
level_positional_encoding_list.append(level_pos_embed)
|
||||
spatial_shapes.append(feat.shape[-2:])
|
||||
reference_points_list.append(reference_points)
|
||||
# shape (batch_size, total_num_query),
|
||||
# total_num_query=sum([., h_i * w_i,.])
|
||||
padding_masks = torch.cat(padding_mask_list, dim=1)
|
||||
# shape (total_num_query, batch_size, c)
|
||||
encoder_inputs = torch.cat(encoder_input_list, dim=0)
|
||||
level_positional_encodings = torch.cat(
|
||||
level_positional_encoding_list, dim=0)
|
||||
device = encoder_inputs.device
|
||||
# shape (num_encoder_levels, 2), from low
|
||||
# resolution to high resolution
|
||||
spatial_shapes = torch.as_tensor(
|
||||
spatial_shapes, dtype=torch.long, device=device)
|
||||
# shape (0, h_0*w_0, h_0*w_0+h_1*w_1, ...)
|
||||
level_start_index = torch.cat((spatial_shapes.new_zeros(
|
||||
(1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
|
||||
reference_points = torch.cat(reference_points_list, dim=0)
|
||||
reference_points = reference_points[None, :, None].repeat(
|
||||
batch_size, 1, self.num_encoder_levels, 1)
|
||||
valid_radios = reference_points.new_ones(
|
||||
(batch_size, self.num_encoder_levels, 2))
|
||||
# shape (num_total_query, batch_size, c)
|
||||
memory = self.encoder(
|
||||
query=encoder_inputs,
|
||||
key=None,
|
||||
value=None,
|
||||
query_pos=level_positional_encodings,
|
||||
key_pos=None,
|
||||
attn_masks=None,
|
||||
key_padding_mask=None,
|
||||
query_key_padding_mask=padding_masks,
|
||||
spatial_shapes=spatial_shapes,
|
||||
reference_points=reference_points,
|
||||
level_start_index=level_start_index,
|
||||
valid_radios=valid_radios)
|
||||
# (num_total_query, batch_size, c) -> (batch_size, c, num_total_query)
|
||||
memory = memory.permute(1, 2, 0)
|
||||
|
||||
# from low resolution to high resolution
|
||||
num_query_per_level = [e[0] * e[1] for e in spatial_shapes]
|
||||
outs = torch.split(memory, num_query_per_level, dim=-1)
|
||||
outs = [
|
||||
x.reshape(batch_size, -1, spatial_shapes[i][0],
|
||||
spatial_shapes[i][1]) for i, x in enumerate(outs)
|
||||
]
|
||||
|
||||
for i in range(self.num_input_levels - self.num_encoder_levels - 1, -1,
|
||||
-1):
|
||||
x = feats[i]
|
||||
cur_feat = self.lateral_convs[i](x)
|
||||
y = cur_feat + F.interpolate(
|
||||
outs[-1],
|
||||
size=cur_feat.shape[-2:],
|
||||
mode='bilinear',
|
||||
align_corners=False)
|
||||
y = self.output_convs[i](y)
|
||||
outs.append(y)
|
||||
multi_scale_features = outs[:self.num_outs]
|
||||
|
||||
mask_feature = self.mask_feature(outs[-1])
|
||||
return mask_feature, multi_scale_features
|
||||
237
segmentation/mmseg_custom/models/plugins/pixel_decoder.py
Normal file
237
segmentation/mmseg_custom/models/plugins/pixel_decoder.py
Normal file
@@ -0,0 +1,237 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from mmcv.cnn import PLUGIN_LAYERS, Conv2d, ConvModule, kaiming_init
|
||||
from mmcv.cnn.bricks.transformer import (build_positional_encoding,
|
||||
build_transformer_layer_sequence)
|
||||
from mmcv.runner import BaseModule, ModuleList
|
||||
|
||||
|
||||
@PLUGIN_LAYERS.register_module()
|
||||
class PixelDecoder(BaseModule):
|
||||
"""Pixel decoder with a structure like fpn.
|
||||
|
||||
Args:
|
||||
in_channels (list[int] | tuple[int]): Number of channels in the
|
||||
input feature maps.
|
||||
feat_channels (int): Number channels for feature.
|
||||
out_channels (int): Number channels for output.
|
||||
norm_cfg (obj:`mmcv.ConfigDict`|dict): Config for normalization.
|
||||
Defaults to dict(type='GN', num_groups=32).
|
||||
act_cfg (obj:`mmcv.ConfigDict`|dict): Config for activation.
|
||||
Defaults to dict(type='ReLU').
|
||||
encoder (obj:`mmcv.ConfigDict`|dict): Config for transorformer
|
||||
encoder.Defaults to None.
|
||||
positional_encoding (obj:`mmcv.ConfigDict`|dict): Config for
|
||||
transformer encoder position encoding. Defaults to
|
||||
dict(type='SinePositionalEncoding', num_feats=128,
|
||||
normalize=True).
|
||||
init_cfg (obj:`mmcv.ConfigDict`|dict): Initialization config dict.
|
||||
Default: None
|
||||
"""
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
feat_channels,
|
||||
out_channels,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
init_cfg=None):
|
||||
super().__init__(init_cfg=init_cfg)
|
||||
self.in_channels = in_channels
|
||||
self.num_inputs = len(in_channels)
|
||||
self.lateral_convs = ModuleList()
|
||||
self.output_convs = ModuleList()
|
||||
self.use_bias = norm_cfg is None
|
||||
for i in range(0, self.num_inputs - 1):
|
||||
l_conv = ConvModule(
|
||||
in_channels[i],
|
||||
feat_channels,
|
||||
kernel_size=1,
|
||||
bias=self.use_bias,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=None)
|
||||
o_conv = ConvModule(
|
||||
feat_channels,
|
||||
feat_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=self.use_bias,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=act_cfg)
|
||||
self.lateral_convs.append(l_conv)
|
||||
self.output_convs.append(o_conv)
|
||||
|
||||
self.last_feat_conv = ConvModule(
|
||||
in_channels[-1],
|
||||
feat_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
stride=1,
|
||||
bias=self.use_bias,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=act_cfg)
|
||||
self.mask_feature = Conv2d(
|
||||
feat_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def init_weights(self):
|
||||
"""Initialize weights."""
|
||||
for i in range(0, self.num_inputs - 2):
|
||||
kaiming_init(self.lateral_convs[i].conv, a=1)
|
||||
kaiming_init(self.output_convs[i].conv, a=1)
|
||||
|
||||
kaiming_init(self.mask_feature, a=1)
|
||||
kaiming_init(self.last_feat_conv, a=1)
|
||||
|
||||
def forward(self, feats, img_metas):
|
||||
"""
|
||||
Args:
|
||||
feats (list[Tensor]): Feature maps of each level. Each has
|
||||
shape of [bs, c, h, w].
|
||||
img_metas (list[dict]): List of image information. Pass in
|
||||
for creating more accurate padding mask. #! not used here.
|
||||
|
||||
Returns:
|
||||
tuple: a tuple containing the following:
|
||||
|
||||
- mask_feature (Tensor): Shape [bs, c, h, w].
|
||||
- memory (Tensor): Output of last stage of backbone.
|
||||
Shape [bs, c, h, w].
|
||||
"""
|
||||
y = self.last_feat_conv(feats[-1])
|
||||
for i in range(self.num_inputs - 2, -1, -1):
|
||||
x = feats[i]
|
||||
cur_fpn = self.lateral_convs[i](x)
|
||||
y = cur_fpn + \
|
||||
F.interpolate(y, size=cur_fpn.shape[-2:], mode='nearest')
|
||||
y = self.output_convs[i](y)
|
||||
|
||||
mask_feature = self.mask_feature(y)
|
||||
memory = feats[-1]
|
||||
return mask_feature, memory
|
||||
|
||||
|
||||
@PLUGIN_LAYERS.register_module()
|
||||
class TransformerEncoderPixelDecoder(PixelDecoder):
|
||||
"""Pixel decoder with transormer encoder inside.
|
||||
|
||||
Args:
|
||||
in_channels (list[int] | tuple[int]): Number of channels in the
|
||||
input feature maps.
|
||||
feat_channels (int): Number channels for feature.
|
||||
out_channels (int): Number channels for output.
|
||||
norm_cfg (obj:`mmcv.ConfigDict`|dict): Config for normalization.
|
||||
Defaults to dict(type='GN', num_groups=32).
|
||||
act_cfg (obj:`mmcv.ConfigDict`|dict): Config for activation.
|
||||
Defaults to dict(type='ReLU').
|
||||
encoder (obj:`mmcv.ConfigDict`|dict): Config for transorformer
|
||||
encoder.Defaults to None.
|
||||
positional_encoding (obj:`mmcv.ConfigDict`|dict): Config for
|
||||
transformer encoder position encoding. Defaults to
|
||||
dict(type='SinePositionalEncoding', num_feats=128,
|
||||
normalize=True).
|
||||
init_cfg (obj:`mmcv.ConfigDict`|dict): Initialization config dict.
|
||||
Default: None
|
||||
"""
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
feat_channels,
|
||||
out_channels,
|
||||
norm_cfg=dict(type='GN', num_groups=32),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
encoder=None,
|
||||
positional_encoding=dict(
|
||||
type='SinePositionalEncoding',
|
||||
num_feats=128,
|
||||
normalize=True),
|
||||
init_cfg=None):
|
||||
super(TransformerEncoderPixelDecoder, self).__init__(
|
||||
in_channels,
|
||||
feat_channels,
|
||||
out_channels,
|
||||
norm_cfg,
|
||||
act_cfg,
|
||||
init_cfg=init_cfg)
|
||||
self.last_feat_conv = None
|
||||
|
||||
self.encoder = build_transformer_layer_sequence(encoder)
|
||||
self.encoder_embed_dims = self.encoder.embed_dims
|
||||
assert self.encoder_embed_dims == feat_channels, 'embed_dims({}) of ' \
|
||||
'tranformer encoder must equal to feat_channels({})'.format(
|
||||
feat_channels, self.encoder_embed_dims)
|
||||
self.positional_encoding = build_positional_encoding(
|
||||
positional_encoding)
|
||||
self.encoder_in_proj = Conv2d(
|
||||
in_channels[-1], feat_channels, kernel_size=1)
|
||||
self.encoder_out_proj = ConvModule(
|
||||
feat_channels,
|
||||
feat_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=self.use_bias,
|
||||
norm_cfg=norm_cfg,
|
||||
act_cfg=act_cfg)
|
||||
|
||||
def init_weights(self):
|
||||
"""Initialize weights."""
|
||||
for i in range(0, self.num_inputs - 2):
|
||||
kaiming_init(self.lateral_convs[i].conv, a=1)
|
||||
kaiming_init(self.output_convs[i].conv, a=1)
|
||||
|
||||
kaiming_init(self.mask_feature, a=1)
|
||||
kaiming_init(self.encoder_in_proj, a=1)
|
||||
kaiming_init(self.encoder_out_proj.conv, a=1)
|
||||
|
||||
def forward(self, feats, img_metas):
|
||||
"""
|
||||
Args:
|
||||
feats (list[Tensor]): Feature maps of each level. Each has
|
||||
shape of [bs, c, h, w].
|
||||
img_metas (list[dict]): List of image information. Pass in
|
||||
for creating more accurate padding mask.
|
||||
|
||||
Returns:
|
||||
tuple: a tuple containing the following:
|
||||
|
||||
- mask_feature (Tensor): shape [bs, c, h, w].
|
||||
- memory (Tensor): shape [bs, c, h, w].
|
||||
"""
|
||||
feat_last = feats[-1]
|
||||
bs, c, h, w = feat_last.shape
|
||||
input_img_h, input_img_w = img_metas[0]['pad_shape'][:-1]
|
||||
# input_img_h, input_img_w = img_metas[0]['batch_input_shape']
|
||||
padding_mask = feat_last.new_ones((bs, input_img_h, input_img_w),
|
||||
dtype=torch.float32)
|
||||
for i in range(bs):
|
||||
img_h, img_w, _ = img_metas[i]['img_shape']
|
||||
padding_mask[i, :img_h, :img_w] = 0
|
||||
padding_mask = F.interpolate(
|
||||
padding_mask.unsqueeze(1),
|
||||
size=feat_last.shape[-2:],
|
||||
mode='nearest').to(torch.bool).squeeze(1)
|
||||
|
||||
pos_embed = self.positional_encoding(padding_mask)
|
||||
feat_last = self.encoder_in_proj(feat_last)
|
||||
# [bs, c, h, w] -> [nq, bs, dim]
|
||||
feat_last = feat_last.flatten(2).permute(2, 0, 1)
|
||||
pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
|
||||
padding_mask = padding_mask.flatten(1) # [bs, h, w] -> [bs, h*w]
|
||||
memory = self.encoder(
|
||||
query=feat_last,
|
||||
key=None,
|
||||
value=None,
|
||||
query_pos=pos_embed,
|
||||
query_key_padding_mask=padding_mask)
|
||||
# [nq, bs, em] -> [bs, c, h, w]
|
||||
memory = memory.permute(1, 2, 0).view(bs, self.encoder_embed_dims, h,
|
||||
w)
|
||||
y = self.encoder_out_proj(memory)
|
||||
for i in range(self.num_inputs - 2, -1, -1):
|
||||
x = feats[i]
|
||||
cur_fpn = self.lateral_convs[i](x)
|
||||
y = cur_fpn + \
|
||||
F.interpolate(y, size=cur_fpn.shape[-2:], mode='nearest')
|
||||
y = self.output_convs[i](y)
|
||||
|
||||
mask_feature = self.mask_feature(y)
|
||||
return mask_feature, memory
|
||||
5
segmentation/mmseg_custom/models/segmentors/__init__.py
Normal file
5
segmentation/mmseg_custom/models/segmentors/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .encoder_decoder_mask2former import EncoderDecoderMask2Former
|
||||
from .encoder_decoder_mask2former_aug import EncoderDecoderMask2FormerAug
|
||||
|
||||
__all__ = ['EncoderDecoderMask2Former', 'EncoderDecoderMask2FormerAug']
|
||||
@@ -0,0 +1,285 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from mmseg.core import add_prefix
|
||||
from mmseg.models import builder
|
||||
from mmseg.models.builder import SEGMENTORS
|
||||
from mmseg.models.segmentors.base import BaseSegmentor
|
||||
from mmseg.ops import resize
|
||||
|
||||
|
||||
@SEGMENTORS.register_module()
|
||||
class EncoderDecoderMask2Former(BaseSegmentor):
|
||||
"""Encoder Decoder segmentors.
|
||||
|
||||
EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.
|
||||
Note that auxiliary_head is only used for deep supervision during training,
|
||||
which could be dumped during inference.
|
||||
"""
|
||||
def __init__(self,
|
||||
backbone,
|
||||
decode_head,
|
||||
neck=None,
|
||||
auxiliary_head=None,
|
||||
train_cfg=None,
|
||||
test_cfg=None,
|
||||
pretrained=None,
|
||||
init_cfg=None):
|
||||
super(EncoderDecoderMask2Former, self).__init__(init_cfg)
|
||||
if pretrained is not None:
|
||||
assert backbone.get('pretrained') is None, \
|
||||
'both backbone and segmentor set pretrained weight'
|
||||
backbone.pretrained = pretrained
|
||||
self.backbone = builder.build_backbone(backbone)
|
||||
if neck is not None:
|
||||
self.neck = builder.build_neck(neck)
|
||||
decode_head.update(train_cfg=train_cfg)
|
||||
decode_head.update(test_cfg=test_cfg)
|
||||
self._init_decode_head(decode_head)
|
||||
self._init_auxiliary_head(auxiliary_head)
|
||||
|
||||
self.train_cfg = train_cfg
|
||||
self.test_cfg = test_cfg
|
||||
|
||||
assert self.with_decode_head
|
||||
|
||||
def _init_decode_head(self, decode_head):
|
||||
"""Initialize ``decode_head``"""
|
||||
self.decode_head = builder.build_head(decode_head)
|
||||
self.align_corners = self.decode_head.align_corners
|
||||
self.num_classes = self.decode_head.num_classes
|
||||
|
||||
def _init_auxiliary_head(self, auxiliary_head):
|
||||
"""Initialize ``auxiliary_head``"""
|
||||
if auxiliary_head is not None:
|
||||
if isinstance(auxiliary_head, list):
|
||||
self.auxiliary_head = nn.ModuleList()
|
||||
for head_cfg in auxiliary_head:
|
||||
self.auxiliary_head.append(builder.build_head(head_cfg))
|
||||
else:
|
||||
self.auxiliary_head = builder.build_head(auxiliary_head)
|
||||
|
||||
def extract_feat(self, img):
|
||||
"""Extract features from images."""
|
||||
x = self.backbone(img)
|
||||
if self.with_neck:
|
||||
x = self.neck(x)
|
||||
return x
|
||||
|
||||
def encode_decode(self, img, img_metas):
|
||||
"""Encode images with backbone and decode into a semantic segmentation
|
||||
map of the same size as input."""
|
||||
x = self.extract_feat(img)
|
||||
out = self._decode_head_forward_test(x, img_metas)
|
||||
out = resize(
|
||||
input=out,
|
||||
size=img.shape[2:],
|
||||
mode='bilinear',
|
||||
align_corners=self.align_corners)
|
||||
return out
|
||||
|
||||
def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg,
|
||||
**kwargs):
|
||||
"""Run forward function and calculate loss for decode head in
|
||||
training."""
|
||||
losses = dict()
|
||||
loss_decode = self.decode_head.forward_train(x, img_metas,
|
||||
gt_semantic_seg, **kwargs)
|
||||
|
||||
losses.update(add_prefix(loss_decode, 'decode'))
|
||||
return losses
|
||||
|
||||
def _decode_head_forward_test(self, x, img_metas):
|
||||
"""Run forward function and calculate loss for decode head in
|
||||
inference."""
|
||||
seg_logits = self.decode_head.forward_test(x, img_metas, self.test_cfg)
|
||||
return seg_logits
|
||||
|
||||
def _auxiliary_head_forward_train(self, x, img_metas, gt_semantic_seg):
|
||||
"""Run forward function and calculate loss for auxiliary head in
|
||||
training."""
|
||||
losses = dict()
|
||||
if isinstance(self.auxiliary_head, nn.ModuleList):
|
||||
for idx, aux_head in enumerate(self.auxiliary_head):
|
||||
loss_aux = aux_head.forward_train(x, img_metas,
|
||||
gt_semantic_seg,
|
||||
self.train_cfg)
|
||||
losses.update(add_prefix(loss_aux, f'aux_{idx}'))
|
||||
else:
|
||||
loss_aux = self.auxiliary_head.forward_train(
|
||||
x, img_metas, gt_semantic_seg, self.train_cfg)
|
||||
losses.update(add_prefix(loss_aux, 'aux'))
|
||||
|
||||
return losses
|
||||
|
||||
def forward_dummy(self, img):
|
||||
"""Dummy forward function."""
|
||||
seg_logit = self.encode_decode(img, None)
|
||||
|
||||
return seg_logit
|
||||
|
||||
def forward_train(self, img, img_metas, gt_semantic_seg, **kwargs):
|
||||
"""Forward function for training.
|
||||
|
||||
Args:
|
||||
img (Tensor): Input images.
|
||||
img_metas (list[dict]): List of image info dict where each dict
|
||||
has: 'img_shape', 'scale_factor', 'flip', and may also contain
|
||||
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
|
||||
For details on the values of these keys see
|
||||
`mmseg/datasets/pipelines/formatting.py:Collect`.
|
||||
gt_semantic_seg (Tensor): Semantic segmentation masks
|
||||
used if the architecture supports semantic segmentation task.
|
||||
|
||||
Returns:
|
||||
dict[str, Tensor]: a dictionary of loss components
|
||||
"""
|
||||
|
||||
x = self.extract_feat(img)
|
||||
|
||||
losses = dict()
|
||||
|
||||
loss_decode = self._decode_head_forward_train(x, img_metas,
|
||||
gt_semantic_seg,
|
||||
**kwargs)
|
||||
losses.update(loss_decode)
|
||||
|
||||
if self.with_auxiliary_head:
|
||||
loss_aux = self._auxiliary_head_forward_train(
|
||||
x, img_metas, gt_semantic_seg)
|
||||
losses.update(loss_aux)
|
||||
|
||||
return losses
|
||||
|
||||
# TODO refactor
|
||||
def slide_inference(self, img, img_meta, rescale):
|
||||
"""Inference by sliding-window with overlap.
|
||||
|
||||
If h_crop > h_img or w_crop > w_img, the small patch will be used to
|
||||
decode without padding.
|
||||
"""
|
||||
|
||||
h_stride, w_stride = self.test_cfg.stride
|
||||
h_crop, w_crop = self.test_cfg.crop_size
|
||||
batch_size, _, h_img, w_img = img.size()
|
||||
num_classes = self.num_classes
|
||||
h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
|
||||
w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
|
||||
preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
|
||||
count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
|
||||
for h_idx in range(h_grids):
|
||||
for w_idx in range(w_grids):
|
||||
y1 = h_idx * h_stride
|
||||
x1 = w_idx * w_stride
|
||||
y2 = min(y1 + h_crop, h_img)
|
||||
x2 = min(x1 + w_crop, w_img)
|
||||
y1 = max(y2 - h_crop, 0)
|
||||
x1 = max(x2 - w_crop, 0)
|
||||
crop_img = img[:, :, y1:y2, x1:x2]
|
||||
crop_seg_logit = self.encode_decode(crop_img, img_meta)
|
||||
preds += F.pad(crop_seg_logit,
|
||||
(int(x1), int(preds.shape[3] - x2), int(y1),
|
||||
int(preds.shape[2] - y2)))
|
||||
|
||||
count_mat[:, :, y1:y2, x1:x2] += 1
|
||||
assert (count_mat == 0).sum() == 0
|
||||
if torch.onnx.is_in_onnx_export():
|
||||
# cast count_mat to constant while exporting to ONNX
|
||||
count_mat = torch.from_numpy(
|
||||
count_mat.cpu().detach().numpy()).to(device=img.device)
|
||||
preds = preds / count_mat
|
||||
if rescale:
|
||||
preds = resize(
|
||||
preds,
|
||||
size=img_meta[0]['ori_shape'][:2],
|
||||
mode='bilinear',
|
||||
align_corners=self.align_corners,
|
||||
warning=False)
|
||||
return preds
|
||||
|
||||
def whole_inference(self, img, img_meta, rescale):
|
||||
"""Inference with full image."""
|
||||
|
||||
seg_logit = self.encode_decode(img, img_meta)
|
||||
if rescale:
|
||||
# support dynamic shape for onnx
|
||||
if torch.onnx.is_in_onnx_export():
|
||||
size = img.shape[2:]
|
||||
else:
|
||||
size = img_meta[0]['ori_shape'][:2]
|
||||
seg_logit = resize(
|
||||
seg_logit,
|
||||
size=size,
|
||||
mode='bilinear',
|
||||
align_corners=self.align_corners,
|
||||
warning=False)
|
||||
|
||||
return seg_logit
|
||||
|
||||
def inference(self, img, img_meta, rescale):
|
||||
"""Inference with slide/whole style.
|
||||
|
||||
Args:
|
||||
img (Tensor): The input image of shape (N, 3, H, W).
|
||||
img_meta (dict): Image info dict where each dict has: 'img_shape',
|
||||
'scale_factor', 'flip', and may also contain
|
||||
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
|
||||
For details on the values of these keys see
|
||||
`mmseg/datasets/pipelines/formatting.py:Collect`.
|
||||
rescale (bool): Whether rescale back to original shape.
|
||||
|
||||
Returns:
|
||||
Tensor: The output segmentation map.
|
||||
"""
|
||||
|
||||
assert self.test_cfg.mode in ['slide', 'whole']
|
||||
ori_shape = img_meta[0]['ori_shape']
|
||||
assert all(_['ori_shape'] == ori_shape for _ in img_meta)
|
||||
if self.test_cfg.mode == 'slide':
|
||||
seg_logit = self.slide_inference(img, img_meta, rescale)
|
||||
else:
|
||||
seg_logit = self.whole_inference(img, img_meta, rescale)
|
||||
output = F.softmax(seg_logit, dim=1)
|
||||
flip = img_meta[0]['flip']
|
||||
if flip:
|
||||
flip_direction = img_meta[0]['flip_direction']
|
||||
assert flip_direction in ['horizontal', 'vertical']
|
||||
if flip_direction == 'horizontal':
|
||||
output = output.flip(dims=(3,))
|
||||
elif flip_direction == 'vertical':
|
||||
output = output.flip(dims=(2,))
|
||||
|
||||
return output
|
||||
|
||||
def simple_test(self, img, img_meta, rescale=True):
|
||||
"""Simple test with single image."""
|
||||
seg_logit = self.inference(img, img_meta, rescale)
|
||||
seg_pred = seg_logit.argmax(dim=1)
|
||||
if torch.onnx.is_in_onnx_export():
|
||||
# our inference backend only support 4D output
|
||||
seg_pred = seg_pred.unsqueeze(0)
|
||||
return seg_pred
|
||||
seg_pred = seg_pred.cpu().numpy()
|
||||
# unravel batch dim
|
||||
seg_pred = list(seg_pred)
|
||||
return seg_pred
|
||||
|
||||
def aug_test(self, imgs, img_metas, rescale=True):
|
||||
"""Test with augmentations.
|
||||
|
||||
Only rescale=True is supported.
|
||||
"""
|
||||
# aug_test rescale all imgs back to ori_shape for now
|
||||
assert rescale
|
||||
# to save memory, we get augmented seg logit inplace
|
||||
seg_logit = self.inference(imgs[0], img_metas[0], rescale)
|
||||
for i in range(1, len(imgs)):
|
||||
cur_seg_logit = self.inference(imgs[i], img_metas[i], rescale)
|
||||
seg_logit += cur_seg_logit
|
||||
seg_logit /= len(imgs)
|
||||
seg_pred = seg_logit.argmax(dim=1)
|
||||
seg_pred = seg_pred.cpu().numpy()
|
||||
# unravel batch dim
|
||||
seg_pred = list(seg_pred)
|
||||
return seg_pred
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user