dataset_type = 'ADE20KDataset'
data_root = 'datasets/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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=16,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='images/training',
               ann_dir='annotations/training',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation',
        pipeline=test_pipeline))

project/_base_/datasets/ade20k.py
JittorCV/jittorseg
dataset_type = 'ADE20KDataset'
data_root = 'datasets/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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=16,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='images/training',
               ann_dir='annotations/training',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='images/validation',
        ann_dir='annotations/validation',
        pipeline=test_pipeline))

project/_base_/datasets/ade20k_640x640.py
JittorCV/jittorseg
dataset_type = 'CityscapesDataset'
data_root = 'datasets/cityscapes/'
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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=8,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='leftImg8bit/train',
               ann_dir='gtFine/train',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline))

project/_base_/datasets/cityscapes_1024x1024.py
JittorCV/jittorseg
dataset_type = 'CityscapesDataset'
data_root = 'datasets/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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=8,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='leftImg8bit/train',
               ann_dir='gtFine/train',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline))
project/_base_/datasets/cityscapes_512x1024.py
JittorCV/jittorseg
dataset_type = 'CityscapesDataset'
data_root = 'datasets/cityscapes/'
img_norm_cfg = dict(mean=[123.675, 116.28, 103.53],
                    std=[58.395, 57.12, 57.375],
                    to_rgb=True)
crop_size = (769, 769)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(type='Resize', img_scale=(2049, 1025), 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, 1025),
        # 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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=8,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='leftImg8bit/train',
               ann_dir='gtFine/train',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline))

project/_base_/datasets/cityscapes_769x769.py
JittorCV/jittorseg
dataset_type = 'iSAIDDataset'
data_root = 'datasets/iSAID_Patches'

img_norm_cfg = dict(mean=[123.675, 116.28, 103.53],
                    std=[58.395, 57.12, 57.375],
                    to_rgb=True)

crop_size = (896, 896)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(type='Resize', img_scale=(896, 896), 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=(896, 896),
        # 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']),
        ])
]

dataset = dict(
    train=dict(type=dataset_type,
               batch_size=16,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='train/images',
               ann_dir='train/Semantic_masks',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='val/images',
        ann_dir='val/Semantic_masks',
        pipeline=test_pipeline))

project/_base_/datasets/isaid_869x869.py
JittorCV/jittorseg
dataset_type = 'LoveDADataset'
data_root = 'dataset/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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=16,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='img_dir/train',
               ann_dir='ann_dir/train',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        data_root=data_root,
        img_dir='img_dir/val',
        ann_dir='ann_dir/val',
        pipeline=test_pipeline))

project/_base_/datasets/loveda.py
JittorCV/jittorseg
dataset_type = 'PascalVOCDataset'
data_root = 'datasets/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', 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']),
        ])
]
dataset = dict(
    train=dict(type=dataset_type,
               batch_size=16,
               num_workers=8,
               shuffle=True,
               drop_last=False,
               data_root=data_root,
               img_dir='JPEGImages',
               ann_dir='SegmentationClass',
               split='ImageSets/Segmentation/train.txt',
               pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        # Fixed to one
        batch_size=1,
        num_workers=1,
        shuffle=False,
        drop_last=False,
        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)
    )

project/_base_/datasets/voc.py
JittorCV/jittorseg
_base_ = './pascal_voc12.py'
# dataset settings
img_norm_cfg = dict(mean=[123.675, 116.28, 103.53],
                    std=[58.395, 57.12, 57.375],
                    to_rgb=True)
crop_size = (321, 321)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=True),
    dict(type='Resize', img_scale=(1280, 321), 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=(1280, 321),
        # 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']),
        ])
]
project/_base_/datasets/voc_321x321.py
JittorCV/jittorseg
_base_ = './pascal_voc12.py'
dataset = dict(
    train=dict(
        ann_dir=['SegmentationClass', 'SegmentationClassAug'],
        split=[
            'ImageSets/Segmentation/train.txt',
            'ImageSets/Segmentation/aug.txt'
        ]))
project/_base_/datasets/voc_aug.py
JittorCV/jittorseg
logger = dict(type="RunLogger")
log_interval = 50
project/_base_/default_runtime.py
JittorCV/jittorseg
max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000
# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
# scheduler
scheduler = dict(type='PolyLR',
                 max_steps=max_iter,
                 power=0.9,
                 min_lr=1e-4)

project/_base_/schedules/schedule_160k.py
JittorCV/jittorseg
max_iter = 40000
eval_interval = 2000
checkpoint_interval = 2000
# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
# scheduler
scheduler = dict(type='PolyLR',
                 max_steps=max_iter,
                 power=0.9,
                 min_lr=1e-4)

project/_base_/schedules/schedule_40k.py
JittorCV/jittorseg
max_iter = 80000
eval_interval = 4000
checkpoint_interval = 4000
# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
# scheduler
scheduler = dict(type='PolyLR',
                 max_steps=max_iter,
                 power=0.9,
                 min_lr=1e-4)

project/_base_/schedules/schedule_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]
# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='ANNHead',
                     in_channels=[1024, 2048],
                     in_index=[2, 3],
                     channels=512,
                     project_channels=256,
                     query_scales=(1, ),
                     key_pool_scales=(1, 3, 6, 8),
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/ann/ann_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k_640x640.py', '../_base_/default_runtime.py'
]

norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://beit_base_patch16_224_pt22k_ft22k.pkl',
    backbone=dict(type='BEiT',
                  img_size=(640, 640),
                  patch_size=16,
                  in_channels=3,
                  embed_dims=768,
                  num_layers=12,
                  num_heads=12,
                  mlp_ratio=4,
                  out_indices=(3, 5, 7, 11),
                  qv_bias=True,
                  attn_drop_rate=0.0,
                  drop_path_rate=0.1,
                  norm_cfg=dict(type='LN', eps=1e-6),
                  act_cfg=dict(type='GELU'),
                  norm_eval=False,
                  init_values=0.1),
    neck=dict(type='Feature2Pyramid', embed_dim=768, rescales=[4, 2, 1, 0.5]),
    decode_head=dict(type='UPerHead',
                     in_channels=[768, 768, 768, 768],
                     in_index=[0, 1, 2, 3],
                     pool_scales=(1, 2, 3, 6),
                     channels=768,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=768,
                        in_index=2,
                        channels=256,
                        num_convs=1,
                        concat_input=False,
                        dropout_ratio=0.1,
                        num_classes=150,
                        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='slide', crop_size=(640, 640), stride=(426, 426)))

parameter_groups_generator = dict(type="LRDecayParameterGroupsGenerator",
                                  paramwise_cfg=dict(num_layers=12,
                                                     decay_rate=0.9))

optimizer = dict(
    type='CustomAdamW',
    lr=3e-5,
    betas=(0.9, 0.999),
    weight_decay=0.05,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/beit/upernet_beit-base_8x2_640x640_160k_ade20k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='CCHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     recurrence=2,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/ccnet/ccnet_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py',
    '../_base_/default_runtime.py',
]

# model settings
model = dict(
    type='EncoderDecoder',
    pretrained=
    'jittorhub://convnext-base_3rdparty_32xb128-noema_in1k_20220301-2a0ee547.pkl',
    backbone=dict(type='ConvNeXt',
                  arch='base',
                  out_indices=[0, 1, 2, 3],
                  drop_path_rate=0.4,
                  layer_scale_init_value=1.0,
                  gap_before_final_norm=False),
    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=150,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0)),
    auxiliary_head=dict(type='FCNHead',
                        in_channels=512,
                        in_index=2,
                        channels=256,
                        num_convs=1,
                        concat_input=False,
                        dropout_ratio=0.1,
                        num_classes=150,
                        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='slide', crop_size=(512, 512), stride=(341, 341)))

parameter_groups_generator = dict(type="LRDecayParameterGroupsGenerator",
                                  paramwise_cfg={
                                      'decay_rate': 0.9,
                                      'decay_type': 'stage_wise',
                                      'num_layers': 12
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.0001,
    betas=(0.9, 0.999),
    weight_decay=0.05,
)

max_iter = 160000
scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/convnext/upernet_convnext_base_512x512_ade20k_160k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='DAHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     pam_channels=64,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/danet/danet_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='ASPPHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     dilations=(1, 12, 24, 36),
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/deeplabv3/deeplabv3_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='DepthwiseSeparableASPPHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     dilations=(1, 12, 24, 36),
                     c1_in_channels=256,
                     c1_channels=48,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/deeplabv3plus/deeplabv3plus_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/cityscapes_769x769.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_40k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='EAHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     dropout_ratio=0.1,
                     num_classes=19,
                     align_corners=True,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0,
                                      class_weight=[
                                          0.8373, 0.918, 0.866, 1.0345, 1.0166,
                                          0.9969, 0.9754, 1.0489, 0.8786,
                                          1.0023, 0.9539, 0.9843, 1.1116,
                                          0.9037, 1.0865, 1.0955, 1.0865,
                                          1.1529, 1.0507
                                      ]),
                     sampler=dict(type='OHEMPixelSampler',
                                  thresh=0.7,
                                  min_kept=100000)),
    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=True,
                        loss_decode=dict(type='CrossEntropyLoss',
                                         use_sigmoid=False,
                                         loss_weight=0.4)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='slide', crop_size=(769, 769), stride=(513, 513)))

project/eanet/eanet_r50-d8_769x769_cityscapes_40k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/cityscapes_512x1024.py',
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='EMAHead',
                     in_channels=2048,
                     in_index=3,
                     channels=256,
                     ema_channels=512,
                     num_bases=64,
                     num_stages=3,
                     momentum=0.1,
                     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'))

project/emanet/emanet_r50-d8_512x1024_cityscapes_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/cityscapes_512x1024.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='FCNHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     num_convs=2,
                     concat_input=True,
                     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'))

project/fcn/fcn_r50-d8_512x1024_cityscapes_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='FCNHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     num_convs=2,
                     concat_input=True,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/fcn/fcn_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/isaid_869x869.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]


# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='FCNHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     num_convs=2,
                     concat_input=True,
                     dropout_ratio=0.1,
                     num_classes=16,
                     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=16,
                        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'))

project/fcn/fcn_r50-d8_896x896_isaid_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_160k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='GCHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     ratio=1 / 4.,
                     pooling_type='att',
                     fusion_types=('channel_add', ),
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/gcnet/gcnet_r50-d8_512x512_ade20k_160k.py
JittorCV/jittorseg
_base_ = ['../_base_/datasets/ade20k.py', '../_base_/default_runtime.py']

norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mae_pretrain_vit_base.pkl',
    backbone=dict(type='MAE',
                  img_size=(512, 512),
                  patch_size=16,
                  in_channels=3,
                  embed_dims=768,
                  num_layers=12,
                  num_heads=12,
                  mlp_ratio=4,
                  out_indices=(3, 5, 7, 11),
                  attn_drop_rate=0.0,
                  drop_path_rate=0.1,
                  norm_cfg=dict(type='LN', eps=1e-6),
                  act_cfg=dict(type='GELU'),
                  norm_eval=False,
                  init_values=1.0),
    neck=dict(type='Feature2Pyramid', embed_dim=768, rescales=[4, 2, 1, 0.5]),
    decode_head=dict(type='UPerHead',
                     in_channels=[768, 768, 768, 768],
                     in_index=[0, 1, 2, 3],
                     pool_scales=(1, 2, 3, 6),
                     channels=768,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=768,
                        in_index=2,
                        channels=256,
                        num_convs=1,
                        concat_input=False,
                        dropout_ratio=0.1,
                        num_classes=150,
                        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='slide', crop_size=(512, 512), stride=(341, 341)))

parameter_groups_generator = dict(type="LRDecayParameterGroupsGenerator",
                                  paramwise_cfg=dict(num_layers=12,
                                                     decay_rate=0.65))

optimizer = dict(
    type='CustomAdamW',
    lr=1e-4,
    betas=(0.9, 0.999),
    weight_decay=0.05,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/mae/upernet_mae-base_fp16_8x2_512x512_160k_ade20k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    # pretrained='jittorhub://mobilenet_v2.pkl',
    backbone=dict(type='MobileNetV2',
                  widen_factor=1.,
                  strides=(1, 2, 2, 1, 1, 1, 1),
                  dilations=(1, 1, 1, 2, 2, 4, 4),
                  out_indices=(1, 2, 4, 6)),
    decode_head=dict(type='FCNHead',
                     in_channels=320,
                     in_index=3,
                     channels=512,
                     num_convs=2,
                     concat_input=True,
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=96,
                        in_index=2,
                        channels=256,
                        num_convs=1,
                        concat_input=False,
                        dropout_ratio=0.1,
                        num_classes=150,
                        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'))

project/mobilenet_v2/fcn_m-v2-d8_512x512_ade20k_160k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='NLHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     dropout_ratio=0.1,
                     reduction=2,
                     use_scale=True,
                     mode='embedded_gaussian',
                     num_classes=150,
                     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=150,
                        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'))

project/nonlocal_net/nonlocal_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_160k.py'
]

norm_cfg = dict(type='BN')
model = dict(
    type='CascadeEncoderDecoder',
    num_stages=2,
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    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,
                  contract_dilation=True),
    neck=dict(type='FPN',
              in_channels=[256, 512, 1024, 2048],
              out_channels=256,
              num_outs=4),
    decode_head=[
        dict(type='FPNHead',
             in_channels=[256, 256, 256, 256],
             in_index=[0, 1, 2, 3],
             feature_strides=[4, 8, 16, 32],
             channels=128,
             dropout_ratio=-1,
             num_classes=150,
             norm_cfg=norm_cfg,
             align_corners=False,
             loss_decode=dict(type='CrossEntropyLoss',
                              use_sigmoid=False,
                              loss_weight=1.0)),
        dict(type='PointHead',
             in_channels=[256],
             in_index=[0],
             channels=256,
             num_fcs=3,
             coarse_pred_each_layer=True,
             dropout_ratio=-1,
             num_classes=150,
             align_corners=False,
             loss_decode=dict(type='CrossEntropyLoss',
                              use_sigmoid=False,
                              loss_weight=1.0))
    ],
    # model training and testing settings
    train_cfg=dict(num_points=2048,
                   oversample_ratio=3,
                   importance_sample_ratio=0.75),
    test_cfg=dict(mode='whole',
                  subdivision_steps=2,
                  subdivision_num_points=8196,
                  scale_factor=2))

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=200,
                 warmup_ratio=1e-6,
                 max_steps=160000,
                 power=1.0,
                 min_lr=0)

project/point_rend/pointrend_r50_512x512_ade20k_160k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/cityscapes_512x1024.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='PSPHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     pool_scales=(1, 2, 3, 6),
                     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'))

project/pspnet/pspnet_r50-d8_512x1024_cityscapes_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]
# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='PSPHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     pool_scales=(1, 2, 3, 6),
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/pspnet/pspnet_r50-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/loveda.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    backbone=dict(type='ResNetV1c',
                  depth=50,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='PSPHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     pool_scales=(1, 2, 3, 6),
                     dropout_ratio=0.1,
                     num_classes=7,
                     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=7,
                        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'))

project/pspnet/pspnet_r50-d8_512x512_loveda_80k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py',
    '../_base_/schedules/schedule_80k.py'
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnest101.pkl',
    backbone=dict(type='ResNeSt',
                  depth=101,
                  stem_channels=128,
                  radix=2,
                  reduction_factor=4,
                  avg_down_stride=True,
                  num_stages=4,
                  out_indices=(0, 1, 2, 3),
                  dilations=(1, 1, 2, 4),
                  strides=(1, 2, 1, 1),
                  norm_cfg=norm_cfg,
                  norm_eval=False,
                  contract_dilation=True),
    decode_head=dict(type='PSPHead',
                     in_channels=2048,
                     in_index=3,
                     channels=512,
                     pool_scales=(1, 2, 3, 6),
                     dropout_ratio=0.1,
                     num_classes=150,
                     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=150,
                        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'))

project/resnest/pspnet_s101-d8_512x512_ade20k_80k.py
JittorCV/jittorseg
_base_ = ['../../_base_/datasets/ade20k.py', '../../_base_/default_runtime.py']

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mit_b0.pkl',
    backbone=dict(type='mit_b0'),
    decode_head=dict(type='SegFormerHead',
                     in_channels=[32, 64, 160, 256],
                     in_index=[0, 1, 2, 3],
                     feature_strides=[4, 8, 16, 32],
                     channels=128,
                     dropout_ratio=0.1,
                     num_classes=150,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     decoder_params=dict(embed_dim=256),
                     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'))

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']),
        ])
]

dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segformer/b0/segformer_b0_512x512_ade_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/cityscapes_1024x1024.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_b.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.2,
                  depths=[3, 3, 12, 3]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     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),
                     ham_channels=512),
    # model training and testing settings
    train_cfg=dict(),
    # test_cfg=dict(mode='whole'))
    test_cfg=dict(mode='slide', crop_size=(1024, 1024), stride=(768, 768)))

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, 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='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/base/segnext_base_1024x1024_cityscapes_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/ade20k.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_b.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.2,
                  depths=[3, 3, 12, 3]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=512,
                     dropout_ratio=0.1,
                     num_classes=150,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=512),
    # model training and testing settings
    train_cfg=dict(),
    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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/base/segnext_base_512x512_ade_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/isaid_869x869.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_b.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.2,
                  depths=[3, 3, 12, 3]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=512,
                     dropout_ratio=0.1,
                     num_classes=16,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=512),
    # model training and testing settings
    train_cfg=dict(),
    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=(896, 896),
        # 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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/base/segnext_base_896x896_isaid_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/cityscapes_1024x1024.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_l.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.3,
                  depths=[3, 5, 27, 3]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=1024,
                     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),
                     ham_channels=1024),
    # model training and testing settings
    train_cfg=dict(),
    # test_cfg=dict(mode='whole'))
    test_cfg=dict(mode='slide', crop_size=(1024, 1024), stride=(768, 768)))

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, 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='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/large/segnext_large_1024x1024_cityscapes_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/ade20k.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_l.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.3,
                  depths=[3, 5, 27, 3]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=1024,
                     dropout_ratio=0.1,
                     num_classes=150,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=1024),
    # model training and testing settings
    train_cfg=dict(),
    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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/large/segnext_large_512x512_ade_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/isaid_869x869.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_l.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.3,
                  depths=[3, 5, 27, 3]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=1024,
                     dropout_ratio=0.1,
                     num_classes=16,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=1024),
    # model training and testing settings
    train_cfg=dict(),
    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=(896, 896),
        # 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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/large/segnext_large_896x896_isaid_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/cityscapes_1024x1024.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_s.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.1,
                  depths=[2, 2, 4, 2]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[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),
                     ham_channels=256,
                     ham_kwargs=dict(MD_R=16)),
    # model training and testing settings
    train_cfg=dict(),
    # test_cfg=dict(mode='whole'))
    test_cfg=dict(mode='slide', crop_size=(1024, 1024), stride=(768, 768)))

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, 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='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/small/segnext_small_1024x1024_cityscapes_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/ade20k.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_s.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.1,
                  depths=[2, 2, 4, 2]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=256,
                     dropout_ratio=0.1,
                     num_classes=150,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=256,
                     ham_kwargs=dict(MD_R=16)),
    # model training and testing settings
    train_cfg=dict(),
    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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/small/segnext_small_512x512_ade_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/isaid_869x869.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_s.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[64, 128, 320, 512],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.1,
                  depths=[2, 2, 4, 2]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[128, 320, 512],
                     in_index=[1, 2, 3],
                     channels=256,
                     dropout_ratio=0.1,
                     num_classes=16,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=256,
                     ham_kwargs=dict(MD_R=16)),
    # model training and testing settings
    train_cfg=dict(),
    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=(896, 896),
        # 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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/small/segnext_small_896x896_isaid_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/cityscapes_1024x1024.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_t.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[32, 64, 160, 256],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.1,
                  depths=[3, 3, 5, 2]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[64, 160, 256],
                     in_index=[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),
                     ham_channels=256,
                     ham_kwargs=dict(MD_R=16)),
    # model training and testing settings
    train_cfg=dict(),
    # test_cfg=dict(mode='whole'))
    test_cfg=dict(mode='slide', crop_size=(1024, 1024), stride=(768, 768)))

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, 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='ResizeToMultiple', size_divisor=32),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/tiny/segnext_tiny_1024x1024_cityscapes_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/ade20k.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_t.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[32, 64, 160, 256],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.1,
                  depths=[3, 3, 5, 2]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[64, 160, 256],
                     in_index=[1, 2, 3],
                     channels=256,
                     dropout_ratio=0.1,
                     num_classes=150,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=256,
                     ham_kwargs=dict(MD_R=16)),
    # model training and testing settings
    train_cfg=dict(),
    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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/tiny/segnext_tiny_512x512_ade_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/isaid_869x869.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='GN', num_groups=32)
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://mscan_t.pkl',
    backbone=dict(type='MSCAN',
                  embed_dims=[32, 64, 160, 256],
                  mlp_ratios=[8, 8, 4, 4],
                  drop_rate=0.0,
                  drop_path_rate=0.1,
                  depths=[3, 3, 5, 2]),
    decode_head=dict(type='LightHamHead',
                     in_channels=[64, 160, 256],
                     in_index=[1, 2, 3],
                     channels=256,
                     dropout_ratio=0.1,
                     num_classes=16,
                     norm_cfg=norm_cfg,
                     align_corners=False,
                     loss_decode=dict(type='CrossEntropyLoss',
                                      use_sigmoid=False,
                                      loss_weight=1.0),
                     ham_channels=256,
                     ham_kwargs=dict(MD_R=16)),
    # model training and testing settings
    train_cfg=dict(),
    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=(896, 896),
        # 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']),
        ])
]
dataset = dict(
    val=dict(pipeline=test_pipeline))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_block': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.),
                                      'head': dict(lr_mult=10.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/segnext/tiny/segnext_tiny_896x896_isaid_160k.py
JittorCV/jittorseg
_base_ = [
    '../../_base_/datasets/ade20k.py',
    '../../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='BN')
backbone_norm_cfg = dict(type='LN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://swin_tiny_patch4_window7_224.pkl',
    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=150,
                     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=150,
                        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'))

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'absolute_pos_embed':
                                      dict(decay_mult=0.),
                                      'relative_position_bias_table':
                                      dict(decay_mult=0.),
                                      'norm':
                                      dict(decay_mult=0.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.00006,
    betas=(0.9, 0.999),
    weight_decay=0.01,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/swin/tiny/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py',
    '../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://resnet50_v1c-2cccc1ad.pkl',
    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,
                  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=150,
                     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=150,
                        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'))

optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR', max_steps=max_iter, power=0.9, min_lr=1e-4)

project/upernet/upernet_r50_512x512_ade20k_160k.py
JittorCV/jittorseg
_base_ = [
    '../_base_/datasets/ade20k.py',
    '../_base_/default_runtime.py',
]

# model settings
norm_cfg = dict(type='BN')
model = dict(
    type='EncoderDecoder',
    pretrained='jittorhub://vit_base_p16_224-80ecf9dd.pkl',
    backbone=dict(
        type='VisionTransformer',
        img_size=(512, 512),
        patch_size=16,
        in_channels=3,
        embed_dims=768,
        num_layers=12,
        num_heads=12,
        mlp_ratio=4,
        out_indices=(2, 5, 8, 11),
        qkv_bias=True,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0,
        with_cls_token=True,
        norm_cfg=dict(type='LN', eps=1e-6),
        act_cfg=dict(type='GELU'),
        norm_eval=False,
        final_norm=True,
        interpolate_mode='bicubic'),
    neck=dict(
        type='MultiLevelNeck',
        in_channels=[768, 768, 768, 768],
        out_channels=768,
        scales=[4, 2, 1, 0.5]),
    decode_head=dict(
        type='UPerHead',
        in_channels=[768, 768, 768, 768],
        in_index=[0, 1, 2, 3],
        pool_scales=(1, 2, 3, 6),
        channels=512,
        dropout_ratio=0.1,
        num_classes=150,
        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=768,
        in_index=3,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=150,
        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'))  # yapf: disable

parameter_groups_generator = dict(type="CustomPrameterGroupsGenerator",
                                  custom_keys={
                                      'pos_embed': dict(decay_mult=0.),
                                      'cls_token': dict(decay_mult=0.),
                                      'norm': dict(decay_mult=0.)
                                  })

optimizer = dict(
    type='CustomAdamW',
    lr=0.0001,
    betas=(0.9, 0.999),
    weight_decay=0.05,
)

max_iter = 160000
eval_interval = 8000
checkpoint_interval = 8000

scheduler = dict(type='PolyLR',
                 warmup='linear',
                 warmup_iters=1500,
                 warmup_ratio=1e-6,
                 max_steps=max_iter,
                 power=1.0,
                 min_lr=0)

project/vit/upernet_vit-b16_ln_mln_512x512_ade20k_160k.py
JittorCV/jittorseg
import argparse
import os.path as osp
from jittorseg.utils.general import mkdir_or_exist, scandir
from cityscapesscripts.preparation.json2labelImg import json2labelImg
from tqdm import tqdm


def convert_json_to_label(json_file):
    label_file = json_file.replace('_polygons.json', '_labelTrainIds.png')
    json2labelImg(json_file, label_file, 'trainIds')


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert Cityscapes annotations to TrainIds')
    parser.add_argument('cityscapes_path', help='cityscapes data path')
    parser.add_argument('--gt-dir', default='gtFine', type=str)
    parser.add_argument('-o', '--out-dir', help='output path')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    cityscapes_path = args.cityscapes_path
    out_dir = args.out_dir if args.out_dir else cityscapes_path
    mkdir_or_exist(out_dir)

    gt_dir = osp.join(cityscapes_path, args.gt_dir)

    for poly in tqdm(scandir(gt_dir, '_polygons.json', recursive=True)):
        poly_file = osp.join(gt_dir, poly)
        convert_json_to_label(poly_file)

    split_names = ['train', 'val', 'test']

    for split in split_names:
        filenames = []
        for poly in scandir(osp.join(gt_dir, split),
                            '_polygons.json',
                            recursive=True):
            filenames.append(poly.replace('_gtFine_polygons.json', ''))
        with open(osp.join(out_dir, f'{split}.txt'), 'w') as f:
            f.writelines(f + '\n' for f in filenames)


if __name__ == '__main__':
    main()

python/datasets_convert_helper/cityscapes.py
JittorCV/jittorseg
# From https://github.com/lxtGH/PFSegNets/blob/master/tools/split_iSAID.py
import cv2
import os
import numpy as np
from natsort import natsorted
from glob import glob
from shutil import copyfile
import argparse
from multiprocessing.pool import ThreadPool
from PIL import Image

# BGR
mask_mapping = {
    (0, 0, 0): 0,
    (63, 0, 0): 1,
    (63, 63, 0): 2,
    (0, 63, 0): 3,
    (127, 63, 0): 4,
    (191, 63, 0): 5,
    (255, 63, 0): 6,
    (63, 127, 0): 7,
    (127, 127, 0): 8,
    (127, 0, 0): 9,
    (191, 0, 0): 10,
    (255, 0, 0): 11,
    (127, 191, 0): 12,
    (191, 127, 0): 13,
    (255, 127, 0): 14,
    (155, 100, 0): 15
}

pool = ThreadPool(16)


def parse_args():
    parser = argparse.ArgumentParser(description='Crop iSAID Dataset')
    parser.add_argument('--src',
                        default='./datasets/iSAID',
                        type=str,
                        help='path for the original dataset')
    parser.add_argument('--target',
                        default='./datasets/iSAID_Patches',
                        type=str,
                        help='path for saving the new dataset')
    parser.add_argument('--patch_width',
                        default=800,
                        type=int,
                        help='Width of the cropped image patch')
    parser.add_argument('--patch_height',
                        default=800,
                        type=int,
                        help='Height of the cropped image patch')
    parser.add_argument('--overlap_area',
                        default=200,
                        type=int,
                        help='Overlap area')
    args = parser.parse_args()
    return args


def slide_crop(file_, src_path, target_path, patch_H, patch_W, overlap, extra):
    if file_ == 'P1527' or file_ == 'P1530':
        return 0
    filename = file_ + extra + '.png'
    full_filename = src_path + '/' + filename
    img = cv2.imread(full_filename)
    img_H, img_W, _ = img.shape
    if extra == '_instance_color_RGB':
        mask_gray = np.zeros(img.shape[:2])
        for k, v in mask_mapping.items():
            mask_gray[(img == k).all(axis=2)] = v
        img = mask_gray

    if img_H > patch_H and img_W > patch_W:
        for x in range(0, img_W, patch_W - overlap):
            for y in range(0, img_H, patch_H - overlap):
                x_str = x
                x_end = x + patch_W
                if x_end > img_W:
                    diff_x = x_end - img_W
                    x_str -= diff_x
                    x_end = img_W
                y_str = y
                y_end = y + patch_H
                if y_end > img_H:
                    diff_y = y_end - img_H
                    y_str -= diff_y
                    y_end = img_H
                if extra == '_instance_color_RGB':
                    patch = img[y_str:y_end, x_str:x_end]
                else:
                    patch = img[y_str:y_end, x_str:x_end, :]
                image = file_ + '_' + str(y_str) + '_' + str(
                    y_end) + '_' + str(x_str) + '_' + str(
                        x_end) + extra + '.png'
                save_path_image = target_path + '/' + image
                if extra == '_instance_color_RGB':
                    mask_gray = Image.fromarray(patch.astype('uint8')).convert('L')
                    mask_gray.save(save_path_image)
                else:
                    cv2.imwrite(save_path_image, patch)
    else:
        if extra == '_instance_color_RGB':
            mask_gray = Image.fromarray(img.astype('uint8')).convert('L')
            mask_gray.save(target_path + '/' + filename)
        else:
            copyfile(full_filename, target_path + '/' + filename)


def main():
    args = parse_args()
    src = args.src
    target = args.target
    sub_set = ['train', 'val', 'test']
    image_sub_folder = ['images', 'Semantic_masks']
    patch_H, patch_W = args.patch_width, args.patch_height  # image patch width and height
    overlap = args.overlap_area  # overlap area

    for set in sub_set:
        for folder in image_sub_folder:
            if set == 'test' and folder == 'Semantic_masks':
                continue
            if folder == 'Semantic_masks':
                extra = '_instance_color_RGB'
            else:
                extra = ''

            src_path = src + "/" + set + "/" + folder
            target_path = target + "/" + set + "/" + folder
            os.makedirs(target_path, exist_ok=True)

            files = glob(src_path + "/*.png")
            files_tmp = []
            for i in files:
                file = os.path.split(i)[-1]
                if '_' in file:
                    files_tmp.append(file.split('_')[0])
                else:
                    files_tmp.append(file.split('.')[0])
            files = natsorted(files_tmp)
            for file_ in files:
                pool.apply_async(slide_crop, (
                    file_,
                    src_path,
                    target_path,
                    patch_H,
                    patch_W,
                    overlap,
                    extra,
                ))
    pool.close()
    pool.join()


if __name__ == '__main__':
    main()

python/datasets_convert_helper/isaid.py
JittorCV/jittorseg
import argparse
import os
import os.path as osp
import shutil
import tempfile
import zipfile
from jittorseg.utils.general import mkdir_or_exist


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert LoveDA dataset to mmsegmentation format')
    parser.add_argument('dataset_path', help='LoveDA folder path')
    parser.add_argument('--tmp_dir', help='path of the temporary directory')
    parser.add_argument('-o', '--out_dir', help='output path')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    dataset_path = args.dataset_path
    if args.out_dir is None:
        out_dir = osp.join('data', 'loveDA')
    else:
        out_dir = args.out_dir

    print('Making directories...')
    mkdir_or_exist(out_dir)
    mkdir_or_exist(osp.join(out_dir, 'img_dir'))
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'train'))
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'val'))
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'test'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'train'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'val'))

    assert 'Train.zip' in os.listdir(dataset_path), \
        'Train.zip is not in {}'.format(dataset_path)
    assert 'Val.zip' in os.listdir(dataset_path), \
        'Val.zip is not in {}'.format(dataset_path)
    assert 'Test.zip' in os.listdir(dataset_path), \
        'Test.zip is not in {}'.format(dataset_path)

    with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
        for dataset in ['Train', 'Val', 'Test']:
            zip_file = zipfile.ZipFile(
                os.path.join(dataset_path, dataset + '.zip'))
            zip_file.extractall(tmp_dir)
            data_type = dataset.lower()
            for location in ['Rural', 'Urban']:
                for image_type in ['images_png', 'masks_png']:
                    if image_type == 'images_png':
                        dst = osp.join(out_dir, 'img_dir', data_type)
                    else:
                        dst = osp.join(out_dir, 'ann_dir', data_type)
                    if dataset == 'Test' and image_type == 'masks_png':
                        continue
                    else:
                        src_dir = osp.join(tmp_dir, dataset, location,
                                           image_type)
                        src_lst = os.listdir(src_dir)
                        for file in src_lst:
                            shutil.move(osp.join(src_dir, file), dst)
        print('Removing the temporary files...')

    print('Done!')


if __name__ == '__main__':
    main()

python/datasets_convert_helper/loveda.py
JittorCV/jittorseg
import argparse
import glob
import math
import os
import os.path as osp
import tempfile
import zipfile
import numpy as np
from jittorseg.utils.general import mkdir_or_exist
import cv2


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert potsdam dataset to mmsegmentation format')
    parser.add_argument('dataset_path', help='potsdam folder path')
    parser.add_argument('--tmp_dir', help='path of the temporary directory')
    parser.add_argument('-o', '--out_dir', help='output path')
    parser.add_argument('--clip_size',
                        type=int,
                        help='clipped size of image after preparation',
                        default=512)
    parser.add_argument('--stride_size',
                        type=int,
                        help='stride of clipping original images',
                        default=256)
    args = parser.parse_args()
    return args


def clip_big_image(image_path, clip_save_dir, args, to_label=False):
    # Original image of Potsdam dataset is very large, thus pre-processing
    # of them is adopted. Given fixed clip size and stride size to generate
    # clipped image, the intersection　of width and height is determined.
    # For example, given one 5120 x 5120 original image, the clip size is
    # 512 and stride size is 256, thus it would generate 20x20 = 400 images
    # whose size are all 512x512.
    image = cv2.imread(image_path)

    h, w, c = image.shape
    clip_size = args.clip_size
    stride_size = args.stride_size

    num_rows = math.ceil((h - clip_size) / stride_size) if math.ceil(
        (h - clip_size) /
        stride_size) * stride_size + clip_size >= h else math.ceil(
            (h - clip_size) / stride_size) + 1
    num_cols = math.ceil((w - clip_size) / stride_size) if math.ceil(
        (w - clip_size) /
        stride_size) * stride_size + clip_size >= w else math.ceil(
            (w - clip_size) / stride_size) + 1

    x, y = np.meshgrid(np.arange(num_cols + 1), np.arange(num_rows + 1))
    xmin = x * clip_size
    ymin = y * clip_size

    xmin = xmin.ravel()
    ymin = ymin.ravel()
    xmin_offset = np.where(xmin + clip_size > w, w - xmin - clip_size,
                           np.zeros_like(xmin))
    ymin_offset = np.where(ymin + clip_size > h, h - ymin - clip_size,
                           np.zeros_like(ymin))
    boxes = np.stack([
        xmin + xmin_offset, ymin + ymin_offset,
        np.minimum(xmin + clip_size, w),
        np.minimum(ymin + clip_size, h)
    ],
                     axis=1)

    if to_label:
        color_map = np.array([[0, 0, 0], [255, 255, 255], [255, 0, 0],
                              [255, 255, 0], [0, 255, 0], [0, 255, 255],
                              [0, 0, 255]])
        flatten_v = np.matmul(image.reshape(-1, c),
                              np.array([2, 3, 4]).reshape(3, 1))
        out = np.zeros_like(flatten_v)
        for idx, class_color in enumerate(color_map):
            value_idx = np.matmul(class_color,
                                  np.array([2, 3, 4]).reshape(3, 1))
            out[flatten_v == value_idx] = idx
        image = out.reshape(h, w)

    for box in boxes:
        start_x, start_y, end_x, end_y = box
        clipped_image = image[start_y:end_y,
                              start_x:end_x] if to_label else image[
                                  start_y:end_y, start_x:end_x, :]
        idx_i, idx_j = osp.basename(image_path).split('_')[2:4]
        cv2.imwrite(
            clipped_image.astype(np.uint8),
            osp.join(
                clip_save_dir,
                f'{idx_i}_{idx_j}_{start_x}_{start_y}_{end_x}_{end_y}.png'))


def main():
    args = parse_args()
    splits = {
        'train': [
            '2_10', '2_11', '2_12', '3_10', '3_11', '3_12', '4_10', '4_11',
            '4_12', '5_10', '5_11', '5_12', '6_10', '6_11', '6_12', '6_7',
            '6_8', '6_9', '7_10', '7_11', '7_12', '7_7', '7_8', '7_9'
        ],
        'val': [
            '5_15', '6_15', '6_13', '3_13', '4_14', '6_14', '5_14', '2_13',
            '4_15', '2_14', '5_13', '4_13', '3_14', '7_13'
        ]
    }

    dataset_path = args.dataset_path
    if args.out_dir is None:
        out_dir = osp.join('data', 'potsdam')
    else:
        out_dir = args.out_dir

    print('Making directories...')
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'train'))
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'val'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'train'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'val'))

    zipp_list = glob.glob(os.path.join(dataset_path, '*.zip'))
    print('Find the data', zipp_list)

    for zipp in zipp_list:
        with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
            zip_file = zipfile.ZipFile(zipp)
            zip_file.extractall(tmp_dir)
            src_path_list = glob.glob(os.path.join(tmp_dir, '*.tif'))
            if not len(src_path_list):
                sub_tmp_dir = os.path.join(tmp_dir, os.listdir(tmp_dir)[0])
                src_path_list = glob.glob(os.path.join(sub_tmp_dir, '*.tif'))

            for i, src_path in enumerate(src_path_list):
                idx_i, idx_j = osp.basename(src_path).split('_')[2:4]
                data_type = 'train' if f'{idx_i}_{idx_j}' in splits[
                    'train'] else 'val'
                if 'label' in src_path:
                    dst_dir = osp.join(out_dir, 'ann_dir', data_type)
                    clip_big_image(src_path, dst_dir, args, to_label=True)
                else:
                    dst_dir = osp.join(out_dir, 'img_dir', data_type)
                    clip_big_image(src_path, dst_dir, args, to_label=False)

    print('Removing the temporary files...')

    print('Done!')


if __name__ == '__main__':
    main()

python/datasets_convert_helper/potsdam.py
JittorCV/jittorseg
import argparse
import glob
import math
import os
import os.path as osp
import tempfile
import zipfile

import cv2
from jittorseg.utils.general import mkdir_or_exist
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert vaihingen dataset to mmsegmentation format')
    parser.add_argument('dataset_path', help='vaihingen folder path')
    parser.add_argument('--tmp_dir', help='path of the temporary directory')
    parser.add_argument('-o', '--out_dir', help='output path')
    parser.add_argument(
        '--clip_size',
        type=int,
        help='clipped size of image after preparation',
        default=512)
    parser.add_argument(
        '--stride_size',
        type=int,
        help='stride of clipping original images',
        default=256)
    args = parser.parse_args()
    return args


def clip_big_image(image_path, clip_save_dir, to_label=False):
    # Original image of Vaihingen dataset is very large, thus pre-processing
    # of them is adopted. Given fixed clip size and stride size to generate
    # clipped image, the intersection　of width and height is determined.
    # For example, given one 5120 x 5120 original image, the clip size is
    # 512 and stride size is 256, thus it would generate 20x20 = 400 images
    # whose size are all 512x512.
    image = cv2.imread(image_path)

    h, w, c = image.shape
    cs = args.clip_size
    ss = args.stride_size

    num_rows = math.ceil((h - cs) / ss) if math.ceil(
        (h - cs) / ss) * ss + cs >= h else math.ceil((h - cs) / ss) + 1
    num_cols = math.ceil((w - cs) / ss) if math.ceil(
        (w - cs) / ss) * ss + cs >= w else math.ceil((w - cs) / ss) + 1

    x, y = np.meshgrid(np.arange(num_cols + 1), np.arange(num_rows + 1))
    xmin = x * cs
    ymin = y * cs

    xmin = xmin.ravel()
    ymin = ymin.ravel()
    xmin_offset = np.where(xmin + cs > w, w - xmin - cs, np.zeros_like(xmin))
    ymin_offset = np.where(ymin + cs > h, h - ymin - cs, np.zeros_like(ymin))
    boxes = np.stack([
        xmin + xmin_offset, ymin + ymin_offset,
        np.minimum(xmin + cs, w),
        np.minimum(ymin + cs, h)
    ],
                     axis=1)

    if to_label:
        color_map = np.array([[0, 0, 0], [255, 255, 255], [255, 0, 0],
                              [255, 255, 0], [0, 255, 0], [0, 255, 255],
                              [0, 0, 255]])
        flatten_v = np.matmul(
            image.reshape(-1, c),
            np.array([2, 3, 4]).reshape(3, 1))
        out = np.zeros_like(flatten_v)
        for idx, class_color in enumerate(color_map):
            value_idx = np.matmul(class_color,
                                  np.array([2, 3, 4]).reshape(3, 1))
            out[flatten_v == value_idx] = idx
        image = out.reshape(h, w)

    for box in boxes:
        start_x, start_y, end_x, end_y = box
        clipped_image = image[start_y:end_y,
                              start_x:end_x] if to_label else image[
                                  start_y:end_y, start_x:end_x, :]
        area_idx = osp.basename(image_path).split('_')[3].strip('.tif')
        cv2.imwrite(
            clipped_image.astype(np.uint8),
            osp.join(clip_save_dir,
                     f'{area_idx}_{start_x}_{start_y}_{end_x}_{end_y}.png'))


def main():
    splits = {
        'train': [
            'area1', 'area11', 'area13', 'area15', 'area17', 'area21',
            'area23', 'area26', 'area28', 'area3', 'area30', 'area32',
            'area34', 'area37', 'area5', 'area7'
        ],
        'val': [
            'area6', 'area24', 'area35', 'area16', 'area14', 'area22',
            'area10', 'area4', 'area2', 'area20', 'area8', 'area31', 'area33',
            'area27', 'area38', 'area12', 'area29'
        ],
    }

    dataset_path = args.dataset_path
    if args.out_dir is None:
        out_dir = osp.join('data', 'vaihingen')
    else:
        out_dir = args.out_dir

    print('Making directories...')
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'train'))
    mkdir_or_exist(osp.join(out_dir, 'img_dir', 'val'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'train'))
    mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'val'))

    zipp_list = glob.glob(os.path.join(dataset_path, '*.zip'))
    print('Find the data', zipp_list)

    with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
        for zipp in zipp_list:
            zip_file = zipfile.ZipFile(zipp)
            zip_file.extractall(tmp_dir)
            src_path_list = glob.glob(os.path.join(tmp_dir, '*.tif'))
            if 'ISPRS_semantic_labeling_Vaihingen' in zipp:
                src_path_list = glob.glob(
                    os.path.join(os.path.join(tmp_dir, 'top'), '*.tif'))
            if 'ISPRS_semantic_labeling_Vaihingen_ground_truth_eroded_COMPLETE' in zipp:  # noqa
                src_path_list = glob.glob(os.path.join(tmp_dir, '*.tif'))
                # delete unused area9 ground truth
                for area_ann in src_path_list:
                    if 'area9' in area_ann:
                        src_path_list.remove(area_ann)
            for i, src_path in enumerate(src_path_list):
                area_idx = osp.basename(src_path).split('_')[3].strip('.tif')
                data_type = 'train' if area_idx in splits['train'] else 'val'
                if 'noBoundary' in src_path:
                    dst_dir = osp.join(out_dir, 'ann_dir', data_type)
                    clip_big_image(src_path, dst_dir, to_label=True)
                else:
                    dst_dir = osp.join(out_dir, 'img_dir', data_type)
                    clip_big_image(src_path, dst_dir, to_label=False)

        print('Removing the temporary files...')

    print('Done!')


if __name__ == '__main__':
    args = parse_args()
    main()

python/datasets_convert_helper/vaihingen.py
JittorCV/jittorseg
import argparse
import os.path as osp
from jittorseg.utils.general import mkdir_or_exist, scandir
import numpy as np
from PIL import Image
from scipy.io import loadmat

AUG_LEN = 10582


def convert_mat(mat_file, in_dir, out_dir):
    data = loadmat(osp.join(in_dir, mat_file))
    mask = data['GTcls'][0]['Segmentation'][0].astype(np.uint8)
    seg_filename = osp.join(out_dir, mat_file.replace('.mat', '.png'))
    Image.fromarray(mask).save(seg_filename, 'PNG')


def generate_aug_list(merged_list, excluded_list):
    return list(set(merged_list) - set(excluded_list))


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert PASCAL VOC annotations')
    parser.add_argument('devkit_path', help='pascal voc devkit path')
    parser.add_argument('aug_path', help='pascal voc aug path')
    parser.add_argument('-o', '--out_dir', help='output path')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    devkit_path = args.devkit_path
    aug_path = args.aug_path
    if args.out_dir is None:
        out_dir = osp.join(devkit_path, 'VOC2012', 'SegmentationClassAug')
    else:
        out_dir = args.out_dir
    mkdir_or_exist(out_dir)
    in_dir = osp.join(aug_path, 'dataset', 'cls')

    for i in list(scandir(in_dir, suffix='.mat')):
        convert_mat(i, in_dir, out_dir)

    full_aug_list = []
    with open(osp.join(aug_path, 'dataset', 'train.txt')) as f:
        full_aug_list += [line.strip() for line in f]
    with open(osp.join(aug_path, 'dataset', 'val.txt')) as f:
        full_aug_list += [line.strip() for line in f]

    with open(
            osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation',
                     'train.txt')) as f:
        ori_train_list = [line.strip() for line in f]
    with open(
            osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation',
                     'val.txt')) as f:
        val_list = [line.strip() for line in f]

    aug_train_list = generate_aug_list(ori_train_list + full_aug_list,
                                       val_list)
    assert len(aug_train_list) == AUG_LEN, 'len(aug_train_list) != {}'.format(
        AUG_LEN)

    with open(
            osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation',
                     'trainaug.txt'), 'w') as f:
        f.writelines(line + '\n' for line in aug_train_list)

    aug_list = generate_aug_list(full_aug_list, ori_train_list + val_list)
    assert len(aug_list) == AUG_LEN - len(
        ori_train_list), 'len(aug_list) != {}'.format(AUG_LEN -
                                                      len(ori_train_list))
    with open(
            osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation', 'aug.txt'),
            'w') as f:
        f.writelines(line + '\n' for line in aug_list)

    print('Done!')


if __name__ == '__main__':
    main()

python/datasets_convert_helper/voc_aug.py
JittorCV/jittorseg
from . import models
from . import runner
from . import config
from . import datasets
from . import utils
from . import optims
from . import sampler

# version must use ' instead of "
__version__ = '0.1.0.0'

python/jittorseg/__init__.py
JittorCV/jittorseg
# reference maskrcnn_benchmark default
from collections import OrderedDict
from jittorseg.utils.general import check_file
import os
import yaml
import copy
from importlib import import_module
import sys
import inspect

__all__ = ["get_cfg", "init_cfg", "save_cfg", "print_cfg", "update_cfg"]
BASE_KEY = "_base_"
COVER_KEY = "_cover_"
IGNORE_KEY = "_ignore_"


class Config(OrderedDict):
    def __init__(self, *args):
        super().__init__()
        if len(args) == 1:
            self.load_from_file(args[0])
        else:
            assert (len(args) == 0)

    def __getattr__(self, name):
        if name in self:
            return self[name]
        return None

    def __setattr__(self, name, value):
        self[name] = value

    # Load config file to a Dict without considering BASE_KEY
    @staticmethod
    def _load_dict_from_file_no_base(filename):
        if check_file(filename, ext=[".yaml"]):
            with open(filename, "r") as f:
                cfg = yaml.safe_load(f.read())
        elif check_file(filename, ext=[".py"]):
            f_dir = os.path.dirname(filename)
            f_name = os.path.basename(filename)
            module_name = f_name[:-3]

            # temp_module_name = osp.splitext(temp_config_name)[0]
            sys.path.insert(0, f_dir)
            # Config._validate_py_syntax(filename)
            mod = import_module(module_name)
            sys.path.pop(0)
            cfg = {
                name: value
                for name, value in mod.__dict__.items()
                if not name.startswith('__')
            }
            # delete imported module
            del sys.modules[module_name]
        else:
            assert (False), "unsupported config type."
        return cfg

    # Load config file to a Dict with considering BASE_KEY
    @staticmethod
    def _load_dict_from_file(filename):
        cfg = Config._load_dict_from_file_no_base(filename)
        cfg_dir = os.path.dirname(filename)
        if (BASE_KEY in cfg):
            if (isinstance(cfg[BASE_KEY], list)):
                base_filenames = cfg[BASE_KEY]
            else:
                assert (isinstance(cfg[BASE_KEY], str))
                base_filenames = [cfg[BASE_KEY]]
            cfg_base = {}
            for bfn in base_filenames:
                Config.merge_dict_b2a(
                    cfg_base,
                    Config._load_dict_from_file(os.path.join(cfg_dir, bfn)))
            cfg.pop(BASE_KEY)
            Config.merge_dict_b2a(cfg_base, cfg)
            cfg = cfg_base
        return cfg

    @staticmethod
    def merge_dict_b2a(a, b):
        def clear_cover_key(a):
            if (not isinstance(a, dict)):
                return a
            out = copy.deepcopy(a)
            if (COVER_KEY in out):
                out.pop(COVER_KEY)
            for k, v in out.items():
                out[k] = clear_cover_key(v)
            return out

        assert (isinstance(a, dict))
        assert (isinstance(b, dict))
        if (COVER_KEY in b):
            a.clear()
            temp = clear_cover_key(copy.deepcopy(b))
            a.update(temp)
            return
        for k, v in b.items():
            if ((not (k in a))
                    or (isinstance(v, dict) and v.get(COVER_KEY, False))
                    or (not isinstance(v, dict))
                    or (not isinstance(a[k], dict))):
                a[k] = clear_cover_key(copy.deepcopy(v))
            else:
                Config.merge_dict_b2a(a[k], v)

    def load_from_file(self, filename):
        cfg = Config._load_dict_from_file(filename)
        self.clear()
        self.update(self.dfs(cfg))
        if (self.name is None):
            self.name = os.path.splitext(os.path.basename(filename))[0]
        if self.work_dir is None:
            self.work_dir = f"work_dirs/{self.name}"

    def dfs(self, cfg_other):
        if isinstance(cfg_other, dict):
            now_cfg = Config()
            for k, d in cfg_other.items():
                if (inspect.ismodule(d)):
                    continue
                now_cfg[k] = self.dfs(d)
        elif isinstance(cfg_other, list):
            now_cfg = [
                self.dfs(d) for d in cfg_other if (not inspect.ismodule(d))
            ]
        else:
            now_cfg = copy.deepcopy(cfg_other)
        return now_cfg

    def dump(self):
        """convert Config to dict"""
        now = dict()
        for k, d in self.items():
            if isinstance(d, Config):
                d = d.dump()
            if isinstance(d, list):
                d = [dd.dump() if isinstance(dd, Config) else dd for dd in d]
            now[k] = d
        return now


_cfg = Config()


def init_cfg(filename):
    print("Loading config from: ", filename)
    _cfg.load_from_file(filename)


def get_cfg():
    return _cfg


def update_cfg(**kwargs):
    _cfg.update(kwargs)


def save_cfg(save_file):
    with open(save_file, "w") as f:
        f.write(yaml.dump(_cfg.dump()))


def print_cfg():
    data = yaml.dump(_cfg.dump())
    # TODO: data keys are not sorted
    print(data)

python/jittorseg/config/config.py
JittorCV/jittorseg
from .custom import CustomDataset
from .isaid import iSAIDDataset
from .ade import ADE20KDataset
from .voc import PascalVOCDataset
from .cityscapes import CityscapesDataset
from .loveda import LoveDADataset
from .isprs import ISPRSDataset
from .potsdam import PotsdamDataset

__all__ = [
    'CustomDataset', 'iSAIDDataset', 'ADE20KDataset', 'PascalVOCDataset',
    'CityscapesDataset', 'LoveDADataset', 'ISPRSDataset', 'PotsdamDataset'
]
python/jittorseg/datasets/__init__.py
JittorCV/jittorseg
from jittorseg.utils.registry import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class ADE20KDataset(CustomDataset):
    CLASSES = (
        'wall', 'building', 'sky', 'floor', 'tree', 'ceiling', 'road', 'bed ',
        'windowpane', 'grass', 'cabinet', 'sidewalk', 'person', 'earth',
        'door', 'table', 'mountain', 'plant', 'curtain', 'chair', 'car',
        'water', 'painting', 'sofa', 'shelf', 'house', 'sea', 'mirror', 'rug',
        'field', 'armchair', 'seat', 'fence', 'desk', 'rock', 'wardrobe',
        'lamp', 'bathtub', 'railing', 'cushion', 'base', 'box', 'column',
        'signboard', 'chest of drawers', 'counter', 'sand', 'sink',
        'skyscraper', 'fireplace', 'refrigerator', 'grandstand', 'path',
        'stairs', 'runway', 'case', 'pool table', 'pillow', 'screen door',
        'stairway', 'river', 'bridge', 'bookcase', 'blind', 'coffee table',
        'toilet', 'flower', 'book', 'hill', 'bench', 'countertop', 'stove',
        'palm', 'kitchen island', 'computer', 'swivel chair', 'boat', 'bar',
        'arcade machine', 'hovel', 'bus', 'towel', 'light', 'truck', 'tower',
        'chandelier', 'awning', 'streetlight', 'booth', 'television receiver',
        'airplane', 'dirt track', 'apparel', 'pole', 'land', 'bannister',
        'escalator', 'ottoman', 'bottle', 'buffet', 'poster', 'stage', 'van',
        'ship', 'fountain', 'conveyer belt', 'canopy', 'washer', 'plaything',
        'swimming pool', 'stool', 'barrel', 'basket', 'waterfall', 'tent',
        'bag', 'minibike', 'cradle', 'oven', 'ball', 'food', 'step', 'tank',
        'trade name', 'microwave', 'pot', 'animal', 'bicycle', 'lake',
        'dishwasher', 'screen', 'blanket', 'sculpture', 'hood', 'sconce',
        'vase', 'traffic light', 'tray', 'ashcan', 'fan', 'pier', 'crt screen',
        'plate', 'monitor', 'bulletin board', 'shower', 'radiator', 'glass',
        'clock', 'flag')

    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],
               [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
               [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
               [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
               [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
               [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
               [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
               [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
               [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
               [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
               [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
               [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
               [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
               [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
               [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
               [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
               [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
               [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
               [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
               [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
               [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
               [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
               [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
               [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
               [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
               [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
               [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
               [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
               [102, 255, 0], [92, 0, 255]]

    def __init__(self, **kwargs):
        super(ADE20KDataset, self).__init__(img_suffix='.jpg',
                                            seg_map_suffix='.png',
                                            reduce_zero_label=True,
                                            **kwargs)
python/jittorseg/datasets/ade.py
JittorCV/jittorseg
import os.path as osp
import tempfile
import numpy as np
from PIL import Image
from tqdm import tqdm
from jittorseg.utils.registry import DATASETS
from jittorseg.utils.general import scandir, mkdir_or_exist
from .custom import CustomDataset


@DATASETS.register_module()
class CityscapesDataset(CustomDataset):
    CLASSES = ('road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
               'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
               'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle',
               'bicycle')

    PALETTE = [[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156],
               [190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0],
               [107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60],
               [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100],
               [0, 80, 100], [0, 0, 230], [119, 11, 32]]

    def __init__(self, **kwargs):
        super(CityscapesDataset,
              self).__init__(img_suffix='_leftImg8bit.png',
                             seg_map_suffix='_gtFine_labelTrainIds.png',
                             **kwargs)

    @staticmethod
    def _convert_to_label_id(result):
        """Convert trainId to id for cityscapes."""
        if isinstance(result, str):
            result = np.load(result)
        import cityscapesscripts.helpers.labels as CSLabels
        result_copy = result.copy()
        for trainId, label in CSLabels.trainId2label.items():
            result_copy[result == trainId] = label.id

        return result_copy

    def results2img(self, results, imgfile_prefix, to_label_id):
        mkdir_or_exist(imgfile_prefix)
        result_files = []
        for idx in tqdm(range(len(self))):
            result = results[idx]
            if to_label_id:
                result = self._convert_to_label_id(result)
            filename = self.img_infos[idx]['filename']
            basename = osp.splitext(osp.basename(filename))[0]

            png_filename = osp.join(imgfile_prefix, f'{basename}.png')

            output = Image.fromarray(result.astype(np.uint8)).convert('P')
            import cityscapesscripts.helpers.labels as CSLabels
            palette = np.zeros((len(CSLabels.id2label), 3), dtype=np.uint8)
            for label_id, label in CSLabels.id2label.items():
                palette[label_id] = label.color

            output.putpalette(palette)
            output.save(png_filename)
            result_files.append(png_filename)

        return result_files

    def format_results(self, results, imgfile_prefix=None, to_label_id=True):
        assert isinstance(results, list), 'results must be a list'
        assert len(results) == len(self), (
            'The length of results is not equal to the dataset len: '
            f'{len(results)} != {len(self)}')

        if imgfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            imgfile_prefix = tmp_dir.name
        else:
            tmp_dir = None
        result_files = self.results2img(results, imgfile_prefix, to_label_id)

        return result_files, tmp_dir

    def evaluate(self,
                 results,
                 metric='mIoU',
                 logger=None,
                 imgfile_prefix=None):
        eval_results = dict()
        metrics = metric.copy() if isinstance(metric, list) else [metric]
        if 'cityscapes' in metrics:
            eval_results.update(
                self._evaluate_cityscapes(results, logger, imgfile_prefix))
            metrics.remove('cityscapes')
        if len(metrics) > 0:
            eval_results.update(
                super(CityscapesDataset,
                      self).evaluate(results, metrics, logger))

        return eval_results

    def _evaluate_cityscapes(self, results, logger, imgfile_prefix):
        try:
            import cityscapesscripts.evaluation.evalPixelLevelSemanticLabeling as CSEval  # noqa
        except ImportError:
            raise ImportError('Please run "pip install cityscapesscripts" to '
                              'install cityscapesscripts first.')
        msg = 'Evaluating in Cityscapes style'
        logger.log({'msg:': msg})

        result_files, tmp_dir = self.format_results(results, imgfile_prefix)

        if tmp_dir is None:
            result_dir = imgfile_prefix
        else:
            result_dir = tmp_dir.name

        eval_results = dict()
        logger.log({'Evaluating results under ': result_dir})

        CSEval.args.evalInstLevelScore = True
        CSEval.args.predictionPath = osp.abspath(result_dir)
        CSEval.args.evalPixelAccuracy = True
        CSEval.args.JSONOutput = False

        seg_map_list = []
        pred_list = []

        # when evaluating with official cityscapesscripts,
        # **_gtFine_labelIds.png is used
        for seg_map in scandir(self.ann_dir,
                               'gtFine_labelIds.png',
                               recursive=True):
            seg_map_list.append(osp.join(self.ann_dir, seg_map))
            pred_list.append(CSEval.getPrediction(CSEval.args, seg_map))

        eval_results.update(
            CSEval.evaluateImgLists(pred_list, seg_map_list, CSEval.args))

        if tmp_dir is not None:
            tmp_dir.cleanup()

        return eval_results
python/jittorseg/datasets/cityscapes.py
JittorCV/jittorseg
import jittor as jt
from jittor.dataset import Dataset
import numpy as np
from jittorseg.utils.registry import DATASETS
import warnings
import os.path as osp
from collections import OrderedDict
from prettytable import PrettyTable
from jittorseg.utils.general import is_list_of, scandir, list_from_file
from jittorseg.utils.metrics import intersect_and_union, pre_eval_to_metrics, eval_metrics
from .pipelines import Compose, LoadAnnotations


@DATASETS.register_module()
class CustomDataset(Dataset):

    CLASSES = None

    PALETTE = None

    def __init__(self,
                 pipeline,
                 img_dir,
                 img_suffix='.jpg',
                 ann_dir=None,
                 seg_map_suffix='.png',
                 batch_size=1,
                 num_workers=0,
                 shuffle=False,
                 drop_last=False,
                 split=None,
                 data_root=None,
                 test_mode=False,
                 ignore_index=255,
                 reduce_zero_label=False,
                 classes=None,
                 palette=None,
                 gt_seg_map_loader_cfg=None,
                 buffer_size=536870912,
                 stop_grad=True,
                 keep_numpy_array=False,
                 endless=False):
        super(CustomDataset, self).__init__(batch_size=batch_size,
                                            num_workers=num_workers,
                                            shuffle=shuffle,
                                            drop_last=drop_last,
                                            buffer_size=buffer_size,
                                            stop_grad=stop_grad,
                                            keep_numpy_array=keep_numpy_array,
                                            endless=endless)

        self.pipeline = Compose(pipeline)
        self.img_dir = img_dir
        self.img_suffix = img_suffix
        self.ann_dir = ann_dir
        self.seg_map_suffix = seg_map_suffix
        self.split = split
        self.data_root = data_root
        self.test_mode = test_mode
        self.ignore_index = ignore_index
        self.reduce_zero_label = reduce_zero_label
        self.label_map = None
        self.CLASSES, self.PALETTE = self.get_classes_and_palette(
            classes, palette)
        self.gt_seg_map_loader = LoadAnnotations(
        ) if gt_seg_map_loader_cfg is None else LoadAnnotations(
            **gt_seg_map_loader_cfg)

        if test_mode:
            assert self.CLASSES is not None, \
                '`cls.CLASSES` or `classes` should be specified when testing'

        # join paths if data_root is specified
        if self.data_root is not None:
            if not osp.isabs(self.img_dir):
                self.img_dir = osp.join(self.data_root, self.img_dir)
            if not (self.ann_dir is None or osp.isabs(self.ann_dir)):
                self.ann_dir = osp.join(self.data_root, self.ann_dir)
            if not (self.split is None or osp.isabs(self.split)):
                self.split = osp.join(self.data_root, self.split)

        # load annotations
        self.img_infos = self.load_annotations(self.img_dir, self.img_suffix,
                                               self.ann_dir,
                                               self.seg_map_suffix, self.split)

    def __len__(self):
        return len(self.img_infos)

    def load_annotations(self, img_dir, img_suffix, ann_dir, seg_map_suffix,
                         split):

        img_infos = []
        if split is not None:
            with open(split) as f:
                for line in f:
                    img_name = line.strip()
                    img_info = dict(filename=img_name + img_suffix)
                    if ann_dir is not None:
                        seg_map = img_name + seg_map_suffix
                        img_info['ann'] = dict(seg_map=seg_map)
                    img_infos.append(img_info)
        else:
            for img in scandir(img_dir, img_suffix, recursive=True):
                img_info = dict(filename=img)
                if ann_dir is not None:
                    seg_map = img.replace(img_suffix, seg_map_suffix)
                    img_info['ann'] = dict(seg_map=seg_map)
                img_infos.append(img_info)
            img_infos = sorted(img_infos, key=lambda x: x['filename'])
        if jt.rank == 0:
            print(f'Loaded {len(img_infos)} images')
        return img_infos

    def get_ann_info(self, idx):
        return self.img_infos[idx]['ann']

    def pre_pipeline(self, results):
        results['seg_fields'] = []
        results['img_prefix'] = self.img_dir
        results['seg_prefix'] = self.ann_dir
        if self.custom_classes:
            results['label_map'] = self.label_map

    def __getitem__(self, idx):
        if self.test_mode:
            return self.prepare_test_img(idx)
        else:
            return self.prepare_train_img(idx)

    def prepare_train_img(self, idx):
        img_info = self.img_infos[idx]
        ann_info = self.get_ann_info(idx)
        results = dict(img_info=img_info, ann_info=ann_info)
        self.pre_pipeline(results)
        return self.pipeline(results)

    def prepare_test_img(self, idx):
        img_info = self.img_infos[idx]
        results = dict(img_info=img_info)
        self.pre_pipeline(results)
        return self.pipeline(results)

    def format_results(self, results, imgfile_prefix, indices=None, **kwargs):
        raise NotImplementedError

    def get_gt_seg_map_by_idx(self, index):
        """Get one ground truth segmentation map for evaluation."""
        ann_info = self.get_ann_info(index)
        results = dict(ann_info=ann_info)
        self.pre_pipeline(results)
        self.gt_seg_map_loader(results)
        return results['gt_semantic_seg']

    def get_gt_seg_maps(self, efficient_test=None):
        """Get ground truth segmentation maps for evaluation."""
        if efficient_test is not None:
            warnings.warn(
                'DeprecationWarning: ``efficient_test`` has been deprecated '
                'since MMSeg v0.16, the ``get_gt_seg_maps()`` is CPU memory '
                'friendly by default. ')

        for idx in range(len(self)):
            ann_info = self.get_ann_info(idx)
            results = dict(ann_info=ann_info)
            self.pre_pipeline(results)
            self.gt_seg_map_loader(results)
            yield results['gt_semantic_seg']

    def pre_eval(self, preds, indices):

        # In order to compat with batch inference
        if not isinstance(indices, list):
            indices = [indices]
        if not isinstance(preds, list):
            preds = [preds]

        pre_eval_results = []

        for pred, index in zip(preds, indices):
            seg_map = self.get_gt_seg_map_by_idx(index)
            pre_eval_results.append(
                intersect_and_union(
                    pred,
                    seg_map,
                    len(self.CLASSES),
                    self.ignore_index,
                    # as the labels has been converted when dataset initialized
                    # in `get_palette_for_custom_classes ` this `label_map`
                    # should be `dict()`, see
                    # https://github.com/open-mmlab/mmsegmentation/issues/1415
                    # for more ditails
                    label_map=dict(),
                    reduce_zero_label=self.reduce_zero_label))

        return pre_eval_results

    def get_classes_and_palette(self, classes=None, palette=None):
        if classes is None:
            self.custom_classes = False
            return self.CLASSES, self.PALETTE

        self.custom_classes = True
        if isinstance(classes, str):
            # take it as a file path
            class_names = list_from_file(classes)
        elif isinstance(classes, (tuple, list)):
            class_names = classes
        else:
            raise ValueError(f'Unsupported type {type(classes)} of classes.')

        if self.CLASSES:
            if not set(class_names).issubset(self.CLASSES):
                raise ValueError('classes is not a subset of CLASSES.')

            # dictionary, its keys are the old label ids and its values
            # are the new label ids.
            # used for changing pixel labels in load_annotations.
            self.label_map = {}
            for i, c in enumerate(self.CLASSES):
                if c not in class_names:
                    self.label_map[i] = -1
                else:
                    self.label_map[i] = class_names.index(c)

        palette = self.get_palette_for_custom_classes(class_names, palette)

        return class_names, palette

    def get_palette_for_custom_classes(self, class_names, palette=None):
        if self.label_map is not None:
            # return subset of palette
            palette = []
            for old_id, new_id in sorted(self.label_map.items(),
                                         key=lambda x: x[1]):
                if new_id != -1:
                    palette.append(self.PALETTE[old_id])
            palette = type(self.PALETTE)(palette)

        elif palette is None:
            if self.PALETTE is None:
                # Get random state before set seed, and restore
                # random state later.
                # It will prevent loss of randomness, as the palette
                # may be different in each iteration if not specified.
                # See: https://github.com/open-mmlab/mmdetection/issues/5844
                state = np.random.get_state()
                np.random.seed(42)
                # random palette
                palette = np.random.randint(0, 255, size=(len(class_names), 3))
                np.random.set_state(state)
            else:
                palette = self.PALETTE

        return palette

    def evaluate(self,
                 results,
                 metric='mIoU',
                 logger=None,
                 gt_seg_maps=None,
                 **kwargs):
        if isinstance(metric, str):
            metric = [metric]
        allowed_metrics = ['mIoU', 'mDice', 'mFscore']
        if not set(metric).issubset(set(allowed_metrics)):
            raise KeyError('metric {} is not supported'.format(metric))
        eval_results = {}

        if is_list_of(results, np.ndarray) or is_list_of(results, str):
            if gt_seg_maps is None:
                gt_seg_maps = self.get_gt_seg_maps()
            num_classes = len(self.CLASSES)
            ret_metrics = eval_metrics(
                results,
                gt_seg_maps,
                num_classes,
                self.ignore_index,
                metric,
                label_map=dict(),
                reduce_zero_label=self.reduce_zero_label)
        else:
            ret_metrics = pre_eval_to_metrics(results, metric)
        if self.CLASSES is None:
            class_names = tuple(range(num_classes))
        else:
            class_names = self.CLASSES

        ret_metrics_summary = OrderedDict({
            ret_metric: np.round(np.nanmean(ret_metric_value) * 100, 2)
            for ret_metric, ret_metric_value in ret_metrics.items()
        })

        ret_metrics.pop('aAcc', None)
        ret_metrics_class = OrderedDict({
            ret_metric: np.round(ret_metric_value * 100, 2)
            for ret_metric, ret_metric_value in ret_metrics.items()
        })
        ret_metrics_class.update({'Class': class_names})
        ret_metrics_class.move_to_end('Class', last=False)

        # for logger
        class_table_data = PrettyTable()
        for key, val in ret_metrics_class.items():
            class_table_data.add_column(key, val)

        summary_table_data = PrettyTable()
        for key, val in ret_metrics_summary.items():
            if key == 'aAcc':
                summary_table_data.add_column(key, [val])
            else:
                summary_table_data.add_column('m' + key, [val])
        logger.log({'per class results:\n': class_table_data.get_string()})
        logger.log({'Summary:\n': summary_table_data.get_string()})
        # each metric dict
        for key, value in ret_metrics_summary.items():
            if key == 'aAcc':
                eval_results[key] = value / 100.0
            else:
                eval_results['m' + key] = value / 100.0

        ret_metrics_class.pop('Class', None)
        for key, value in ret_metrics_class.items():
            eval_results.update({
                key + '.' + str(name): value[idx] / 100.0
                for idx, name in enumerate(class_names)
            })

        return eval_results

    def collate_batch(self, batch):
        img_metas = []
        img = []
        gt_semantic_seg = []
        test = False
        if 'gt_semantic_seg' not in batch[0].keys():
            test = True
        if not test:
            for i in batch:
                img.append(i['img'])
                img_metas.append(i['img_metas'])
                gt_semantic_seg.append(i['gt_semantic_seg'])

            return {
                'img': jt.stack(img, dim=0),
                'img_metas': img_metas,
                'gt_semantic_seg': jt.stack(gt_semantic_seg, dim=0)
            }
        else:
            for i in range(len(batch[0]['img'])):
                img.append(batch[0]['img'][i].unsqueeze(0))
                img_metas.append(batch[0]['img_metas'][i])
            return {'img': img, 'img_metas': img_metas}
python/jittorseg/datasets/custom.py
JittorCV/jittorseg
from jittorseg.utils.registry import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class iSAIDDataset(CustomDataset):
    CLASSES = ('background', 'ship', 'store_tank', 'baseball_diamond',
               'tennis_court', 'basketball_court', 'Ground_Track_Field',
               'Bridge', 'Large_Vehicle', 'Small_Vehicle', 'Helicopter',
               'Swimming_pool', 'Roundabout', 'Soccer_ball_field', 'plane',
               'Harbor')
    PALETTE = [[0, 0, 0], [0, 0, 63], [0, 63, 63], [0, 63, 0], [0, 63, 127],
               [0, 63, 191], [0, 63, 255], [0, 127, 63], [0, 127, 127],
               [0, 0, 127], [0, 0, 191], [0, 0, 255], [0, 191, 127],
               [0, 127, 191], [0, 127, 255], [0, 100, 155]]

    def __init__(self, **kwargs):
        super(iSAIDDataset,
              self).__init__(img_suffix='.png',
                             seg_map_suffix='_instance_color_RGB.png',
                             **kwargs)
python/jittorseg/datasets/isaid.py
JittorCV/jittorseg
from jittorseg.utils.registry import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class ISPRSDataset(CustomDataset):
    """ISPRS dataset.
    In segmentation map annotation for LoveDA, 0 is the ignore index.
    ``reduce_zero_label`` should be set to True. The ``img_suffix`` and
    ``seg_map_suffix`` are both fixed to '.png'.
    """
    CLASSES = ('impervious_surface', 'building', 'low_vegetation', 'tree',
               'car', 'clutter')

    PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0],
               [255, 255, 0], [255, 0, 0]]

    def __init__(self, **kwargs):
        super(ISPRSDataset, self).__init__(
            img_suffix='.png',
            seg_map_suffix='.png',
            reduce_zero_label=True,
            **kwargs)
python/jittorseg/datasets/isprs.py
JittorCV/jittorseg
import os.path as osp
from jittorseg.utils.registry import DATASETS
from .custom import CustomDataset
import numpy as np
from PIL import Image
from jittorseg.utils.general import mkdir_or_exist


@DATASETS.register_module()
class LoveDADataset(CustomDataset):
    """LoveDA dataset.
    In segmentation map annotation for LoveDA, 0 is the ignore index.
    ``reduce_zero_label`` should be set to True. The ``img_suffix`` and
    ``seg_map_suffix`` are both fixed to '.png'.
    """
    CLASSES = ('background', 'building', 'road', 'water', 'barren', 'forest',
               'agricultural')

    PALETTE = [[255, 255, 255], [255, 0, 0], [255, 255, 0], [0, 0, 255],
               [159, 129, 183], [0, 255, 0], [255, 195, 128]]

    def __init__(self, **kwargs):
        super(LoveDADataset, self).__init__(
            img_suffix='.png',
            seg_map_suffix='.png',
            reduce_zero_label=True,
            **kwargs)

    def results2img(self, results, imgfile_prefix, indices=None):
        """Write the segmentation results to images.
        Args:
            results (list[ndarray]): Testing results of the
                dataset.
            imgfile_prefix (str): The filename prefix of the png files.
                If the prefix is "somepath/xxx",
                the png files will be named "somepath/xxx.png".
            indices (list[int], optional): Indices of input results, if not
                set, all the indices of the dataset will be used.
                Default: None.
        Returns:
            list[str: str]: result txt files which contains corresponding
            semantic segmentation images.
        """

        mkdir_or_exist(imgfile_prefix)
        result_files = []
        for result, idx in zip(results, indices):

            filename = self.img_infos[idx]['filename']
            basename = osp.splitext(osp.basename(filename))[0]

            png_filename = osp.join(imgfile_prefix, f'{basename}.png')

            # The  index range of official requirement is from 0 to 6.
            output = Image.fromarray(result.astype(np.uint8))
            output.save(png_filename)
            result_files.append(png_filename)

        return result_files

    def format_results(self, results, imgfile_prefix, indices=None):
        """Format the results into dir (standard format for LoveDA evaluation).
        Args:
            results (list): Testing results of the dataset.
            imgfile_prefix (str): The prefix of images files. It
                includes the file path and the prefix of filename, e.g.,
                "a/b/prefix".
            indices (list[int], optional): Indices of input results,
                if not set, all the indices of the dataset will be used.
                Default: None.
        Returns:
            tuple: (result_files, tmp_dir), result_files is a list containing
                the image paths, tmp_dir is the temporal directory created
                for saving json/png files when img_prefix is not specified.
        """
        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.'

        result_files = self.results2img(results, imgfile_prefix, indices)

        return result_files
python/jittorseg/datasets/loveda.py
JittorCV/jittorseg
from .compose import Compose
from .formating import (Collect, to_tensor)
from .loading import LoadAnnotations, LoadImageFromFile
from .test_time_aug import MultiScaleFlipAug
from .transforms import (ResizeToMultiple, CLAHE, Normalize, Pad,
                         PhotoMetricDistortion, RandomCrop, RandomFlip,
                         RandomRotate, Rerange, Resize, RGB2Gray, SegRescale)
from .utils import imread, imresize, imrescale, imflip, impad_to_multiple, impad, imnormalize, imrotate, clahe, bgr2hsv, hsv2bgr

__all__ = [
    'Compose', 'to_tensor', 'ToTensor', 'ImageToTensor', 'Transpose',
    'Collect', 'LoadAnnotations', 'LoadImageFromFile', 'MultiScaleFlipAug',
    'ResizeToMultiple', 'Resize', 'RandomFlip', 'Pad', 'RandomCrop',
    'Normalize', 'SegRescale', 'PhotoMetricDistortion', 'RandomRotate',
    'AdjustGamma', 'CLAHE', 'Rerange', 'RGB2Gray', 'imread', 'imresize',
    'imrescale', 'imflip', 'impad_to_multiple', 'impad', 'imnormalize',
    'imrotate', 'clahe', 'bgr2hsv', 'hsv2bgr'
]
python/jittorseg/datasets/pipelines/__init__.py
JittorCV/jittorseg
import collections

from jittorseg.utils.registry import build_from_cfg, TRANSFORMS


@TRANSFORMS.register_module()
class Compose(object):
    def __init__(self, transforms):
        assert isinstance(transforms, collections.abc.Sequence)
        self.transforms = []
        for transform in transforms:
            if isinstance(transform, dict):
                transform = build_from_cfg(transform, TRANSFORMS)
                self.transforms.append(transform)
            elif callable(transform):
                self.transforms.append(transform)
            else:
                raise TypeError('transform must be callable or a dict')

    def __call__(self, data):
        for t in self.transforms:
            data = t(data)
            if data is None:
                return None
        return data

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += f'    {t}'
        format_string += '\n)'
        return format_string
python/jittorseg/datasets/pipelines/compose.py
JittorCV/jittorseg
import numpy as np

from jittorseg.utils.registry import TRANSFORMS
import jittor as jt


def to_tensor(data):
    return jt.Var(data)


@TRANSFORMS.register_module()
class DefaultFormatBundle(object):
    def __call__(self, results):

        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'] = to_tensor(img)
        if 'gt_semantic_seg' in results:
            # convert to long
            results['gt_semantic_seg'] = to_tensor(
                results['gt_semantic_seg'][None, ...].astype(np.int64))
        return results

    def __repr__(self):
        return self.__class__.__name__


@TRANSFORMS.register_module()
class ImageToTensor(object):
    def __init__(self, keys):
        self.keys = keys

    def __call__(self, results):
        for key in self.keys:
            img = results[key]
            if len(img.shape) < 3:
                img = np.expand_dims(img, -1)
            results[key] = to_tensor(img.transpose(2, 0, 1))
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(keys={self.keys})'


@TRANSFORMS.register_module()
class Collect(object):
    def __init__(self,
                 keys,
                 meta_keys=('filename', 'ori_filename', 'ori_shape',
                            'img_shape', 'pad_shape', 'scale_factor', 'flip',
                            'flip_direction', 'img_norm_cfg')):
        self.keys = keys
        self.meta_keys = meta_keys

    def __call__(self, results):
        data = {}
        img_meta = {}
        for key in self.meta_keys:
            img_meta[key] = results[key]
        data['img_metas'] = img_meta
        for key in self.keys:
            data[key] = results[key]
        return data

    def __repr__(self):
        return self.__class__.__name__ + f'(keys={self.keys}, meta_keys={self.meta_keys})'
python/jittorseg/datasets/pipelines/formating.py
JittorCV/jittorseg
import os.path as osp
import numpy as np
from jittorseg.utils.registry import TRANSFORMS
from .utils import imread


@TRANSFORMS.register_module()
class LoadImageFromFile(object):
    def __init__(self, to_float32=False, color_type='color', backend='cv2'):
        self.to_float32 = to_float32
        self.color_type = color_type
        self.backend = backend

    def __call__(self, results):
        if results.get('img_prefix') is not None:
            filename = osp.join(results['img_prefix'],
                                results['img_info']['filename'])
        else:
            filename = results['img_info']['filename']
        img = imread(filename, flag=self.color_type, backend=self.backend)

        if self.to_float32:
            img = img.astype(np.float32)

        results['filename'] = filename
        results['ori_filename'] = results['img_info']['filename']
        results['img'] = img
        results['img_shape'] = img.shape
        results['ori_shape'] = img.shape
        # Set initial values for default meta_keys
        results['pad_shape'] = img.shape
        results['scale_factor'] = 1.0
        num_channels = 1 if len(img.shape) < 3 else img.shape[2]
        results['img_norm_cfg'] = dict(mean=np.zeros(num_channels,
                                                     dtype=np.float32),
                                       std=np.ones(num_channels,
                                                   dtype=np.float32),
                                       to_rgb=False)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(to_float32={self.to_float32},'
        repr_str += f"color_type='{self.color_type}',"
        repr_str += f"backend='{self.backend}')"
        return repr_str


@TRANSFORMS.register_module()
class LoadAnnotations(object):
    def __init__(self, reduce_zero_label=False, backend='pillow'):
        self.reduce_zero_label = reduce_zero_label
        self.backend = backend

    def __call__(self, results):
        if results.get('seg_prefix', None) is not None:
            filename = osp.join(results['seg_prefix'],
                                results['ann_info']['seg_map'])
        else:
            filename = results['ann_info']['seg_map']
        gt_semantic_seg = imread(
            filename, flag='unchanged',
            backend=self.backend).squeeze().astype(np.uint8)
        # modify if custom classes
        if results.get('label_map', None) is not None:
            gt_semantic_seg_copy = gt_semantic_seg.copy()
            for old_id, new_id in results['label_map'].items():
                gt_semantic_seg[gt_semantic_seg_copy == old_id] = new_id
        # reduce zero_label
        if self.reduce_zero_label:
            # avoid using underflow conversion
            gt_semantic_seg[gt_semantic_seg == 0] = 255
            gt_semantic_seg = gt_semantic_seg - 1
            gt_semantic_seg[gt_semantic_seg == 254] = 255
        results['gt_semantic_seg'] = gt_semantic_seg
        results['seg_fields'].append('gt_semantic_seg')
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(reduce_zero_label={self.reduce_zero_label},'
        repr_str += f"backend='{self.backend}')"
        return repr_str
python/jittorseg/datasets/pipelines/loading.py
JittorCV/jittorseg
import warnings

from jittorseg.utils.registry import TRANSFORMS
from jittorseg.utils.general import is_list_of
from .compose import Compose


@TRANSFORMS.register_module()
class MultiScaleFlipAug(object):
    def __init__(self,
                 transforms,
                 img_scale,
                 img_ratios=None,
                 flip=False,
                 flip_direction='horizontal'):
        self.transforms = Compose(transforms)
        if img_ratios is not None:
            img_ratios = img_ratios if isinstance(img_ratios,
                                                  list) else [img_ratios]
            assert is_list_of(img_ratios, float)
        if img_scale is None:
            # mode 1: given img_scale=None and a range of image ratio
            self.img_scale = None
            assert is_list_of(img_ratios, float)
        elif isinstance(img_scale, tuple) and is_list_of(img_ratios, float):
            assert len(img_scale) == 2
            # mode 2: given a scale and a range of image ratio
            self.img_scale = [(int(img_scale[0] * ratio),
                               int(img_scale[1] * ratio))
                              for ratio in img_ratios]
        else:
            # mode 3: given multiple scales
            self.img_scale = img_scale if isinstance(img_scale,
                                                     list) else [img_scale]
        assert is_list_of(self.img_scale, tuple) or self.img_scale is None
        self.flip = flip
        self.img_ratios = img_ratios
        self.flip_direction = flip_direction if isinstance(
            flip_direction, list) else [flip_direction]
        assert is_list_of(self.flip_direction, str)
        if not self.flip and self.flip_direction != ['horizontal']:
            warnings.warn(
                'flip_direction has no effect when flip is set to False')
        if (self.flip
                and not any([t['type'] == 'RandomFlip' for t in transforms])):
            warnings.warn(
                'flip has no effect when RandomFlip is not in transforms')

    def __call__(self, results):
        aug_data = []
        if self.img_scale is None and is_list_of(self.img_ratios, float):
            h, w = results['img'].shape[:2]
            img_scale = [(int(w * ratio), int(h * ratio))
                         for ratio in self.img_ratios]
        else:
            img_scale = self.img_scale
        flip_aug = [False, True] if self.flip else [False]
        for scale in img_scale:
            for flip in flip_aug:
                for direction in self.flip_direction:
                    _results = results.copy()
                    _results['scale'] = scale
                    _results['flip'] = flip
                    _results['flip_direction'] = direction
                    data = self.transforms(_results)
                    aug_data.append(data)
        # list of dict to dict of list
        aug_data_dict = {key: [] for key in aug_data[0]}
        for data in aug_data:
            for key, val in data.items():
                aug_data_dict[key].append(val)
        return aug_data_dict

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(transforms={self.transforms}, '
        repr_str += f'img_scale={self.img_scale}, flip={self.flip})'
        repr_str += f'flip_direction={self.flip_direction}'
        return repr_str
python/jittorseg/datasets/pipelines/test_time_aug.py
JittorCV/jittorseg
import numpy as np
from numpy import random
from jittorseg.utils.general import is_list_of, is_tuple_of
from jittorseg.utils.registry import TRANSFORMS
from .utils import imresize, imresize_to_multiple, imrescale, imflip, impad_to_multiple, impad, imnormalize, imrotate, clahe, bgr2hsv, hsv2bgr


@TRANSFORMS.register_module()
class ResizeToMultiple(object):
    def __init__(self, size_divisor=32, interpolation=None):
        self.size_divisor = size_divisor
        self.interpolation = interpolation

    def __call__(self, results):
        # Align image to multiple of size divisor.
        img = results['img']
        img = imresize_to_multiple(img,
                                   self.size_divisor,
                                   scale_factor=1,
                                   interpolation=self.interpolation
                                   if self.interpolation else 'bilinear')

        results['img'] = img
        results['img_shape'] = img.shape
        results['pad_shape'] = img.shape

        # Align segmentation map to multiple of size divisor.
        for key in results.get('seg_fields', []):
            gt_seg = results[key]
            gt_seg = imresize_to_multiple(gt_seg,
                                          self.size_divisor,
                                          scale_factor=1,
                                          interpolation='nearest')
            results[key] = gt_seg

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += (f'(size_divisor={self.size_divisor}, '
                     f'interpolation={self.interpolation})')
        return repr_str


@TRANSFORMS.register_module()
class Resize(object):
    def __init__(self,
                 img_scale=None,
                 multiscale_mode='range',
                 ratio_range=None,
                 keep_ratio=True,
                 min_size=None):
        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 is_list_of(self.img_scale, tuple)

        if ratio_range is not None:
            # mode 1: given img_scale=None and a range of image ratio
            # mode 2: given a scale and a range of image ratio
            assert self.img_scale is None or len(self.img_scale) == 1
        else:
            # mode 3 and 4: 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.min_size = min_size

    @staticmethod
    def random_select(img_scales):
        assert 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):
        assert 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):
        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):
        if self.ratio_range is not None:
            if self.img_scale is None:
                h, w = results['img'].shape[:2]
                scale, scale_idx = self.random_sample_ratio((w, h),
                                                            self.ratio_range)
            else:
                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.min_size is not None:
                # TODO: Now 'min_size' is an 'int' which means the minimum
                # shape of images is (min_size, min_size, 3). 'min_size'
                # with tuple type will be supported, i.e. the width and
                # height are not equal.
                if min(results['scale']) < self.min_size:
                    new_short = self.min_size
                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 = imrescale(results['img'],
                                          results['scale'],
                                          return_scale=True)
            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 = 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):
        for key in results.get('seg_fields', []):
            if self.keep_ratio:
                gt_seg = imrescale(results[key],
                                   results['scale'],
                                   interpolation='nearest')
            else:
                gt_seg = imresize(results[key],
                                  results['scale'],
                                  interpolation='nearest')
            results[key] = gt_seg

    def __call__(self, results):
        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


@TRANSFORMS.register_module()
class RandomFlip(object):
    def __init__(self, prob=None, direction='horizontal'):
        self.prob = prob
        self.direction = direction
        if prob is not None:
            assert prob >= 0 and prob <= 1
        assert direction in ['horizontal', 'vertical']

    def __call__(self, results):
        if 'flip' not in results:
            flip = True if np.random.rand() < self.prob else False
            results['flip'] = flip
        if 'flip_direction' not in results:
            results['flip_direction'] = self.direction
        if results['flip']:
            # flip image
            results['img'] = imflip(results['img'],
                                    direction=results['flip_direction'])

            # flip segs
            for key in results.get('seg_fields', []):
                # use copy() to make numpy stride positive
                results[key] = imflip(
                    results[key], direction=results['flip_direction']).copy()
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(prob={self.prob})'


@TRANSFORMS.register_module()
class Pad(object):
    def __init__(self,
                 size=None,
                 size_divisor=None,
                 pad_val=0,
                 seg_pad_val=255):
        self.size = size
        self.size_divisor = size_divisor
        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 or size_divisor is not None
        assert size is None or size_divisor is None

    def _pad_img(self, results):
        if self.size is not None:
            padded_img = impad(results['img'],
                               shape=self.size,
                               pad_val=self.pad_val)
        elif self.size_divisor is not None:
            padded_img = impad_to_multiple(results['img'],
                                           self.size_divisor,
                                           pad_val=self.pad_val)
        results['img'] = padded_img
        results['pad_shape'] = padded_img.shape
        results['pad_fixed_size'] = self.size
        results['pad_size_divisor'] = self.size_divisor

    def _pad_seg(self, results):
        for key in results.get('seg_fields', []):
            results[key] = impad(results[key],
                                 shape=results['pad_shape'][:2],
                                 pad_val=self.seg_pad_val)

    def __call__(self, results):
        self._pad_img(results)
        self._pad_seg(results)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(size={self.size}, size_divisor={self.size_divisor}, ' \
                    f'pad_val={self.pad_val})'
        return repr_str


@TRANSFORMS.register_module()
class Normalize(object):
    def __init__(self, mean, std, to_rgb=True):
        self.mean = np.array(mean, dtype=np.float32)
        self.std = np.array(std, dtype=np.float32)
        self.to_rgb = to_rgb

    def __call__(self, results):
        results['img'] = imnormalize(results['img'], self.mean, self.std,
                                     self.to_rgb)
        results['img_norm_cfg'] = dict(mean=self.mean,
                                       std=self.std,
                                       to_rgb=self.to_rgb)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(mean={self.mean}, std={self.std}, to_rgb=' \
                    f'{self.to_rgb})'
        return repr_str


@TRANSFORMS.register_module()
class Rerange(object):
    def __init__(self, min_value=0, max_value=255):
        assert isinstance(min_value, float) or isinstance(min_value, int)
        assert isinstance(max_value, float) or isinstance(max_value, int)
        assert min_value < max_value
        self.min_value = min_value
        self.max_value = max_value

    def __call__(self, results):
        img = results['img']
        img_min_value = np.min(img)
        img_max_value = np.max(img)

        assert img_min_value < img_max_value
        # rerange to [0, 1]
        img = (img - img_min_value) / (img_max_value - img_min_value)
        # rerange to [min_value, max_value]
        img = img * (self.max_value - self.min_value) + self.min_value
        results['img'] = img

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(min_value={self.min_value}, max_value={self.max_value})'
        return repr_str


@TRANSFORMS.register_module()
class CLAHE(object):
    def __init__(self, clip_limit=40.0, tile_grid_size=(8, 8)):
        assert isinstance(clip_limit, (float, int))
        self.clip_limit = clip_limit
        assert is_tuple_of(tile_grid_size, int)
        assert len(tile_grid_size) == 2
        self.tile_grid_size = tile_grid_size

    def __call__(self, results):
        for i in range(results['img'].shape[2]):
            results['img'][:, :, i] = clahe(
                np.array(results['img'][:, :, i], dtype=np.uint8),
                self.clip_limit, self.tile_grid_size)

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(clip_limit={self.clip_limit}, '\
                    f'tile_grid_size={self.tile_grid_size})'
        return repr_str


@TRANSFORMS.register_module()
class RandomCrop(object):
    def __init__(self, crop_size, cat_max_ratio=1., ignore_index=255):
        assert crop_size[0] > 0 and crop_size[1] > 0
        self.crop_size = crop_size
        self.cat_max_ratio = cat_max_ratio
        self.ignore_index = ignore_index

    def get_crop_bbox(self, img):
        margin_h = max(img.shape[0] - self.crop_size[0], 0)
        margin_w = max(img.shape[1] - self.crop_size[1], 0)
        offset_h = np.random.randint(0, margin_h + 1)
        offset_w = np.random.randint(0, margin_w + 1)
        crop_y1, crop_y2 = offset_h, offset_h + self.crop_size[0]
        crop_x1, crop_x2 = offset_w, offset_w + self.crop_size[1]

        return crop_y1, crop_y2, crop_x1, crop_x2

    def crop(self, img, crop_bbox):
        crop_y1, crop_y2, crop_x1, crop_x2 = crop_bbox
        img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]
        return img

    def __call__(self, results):
        img = results['img']
        crop_bbox = self.get_crop_bbox(img)
        if self.cat_max_ratio < 1.:
            # Repeat 10 times
            for _ in range(10):
                seg_temp = self.crop(results['gt_semantic_seg'], crop_bbox)
                labels, cnt = np.unique(seg_temp, return_counts=True)
                cnt = cnt[labels != self.ignore_index]
                if len(cnt) > 1 and np.max(cnt) / np.sum(
                        cnt) < self.cat_max_ratio:
                    break
                crop_bbox = self.get_crop_bbox(img)

        # crop the image
        img = self.crop(img, crop_bbox)
        img_shape = img.shape
        results['img'] = img
        results['img_shape'] = img_shape

        # crop semantic seg
        for key in results.get('seg_fields', []):
            results[key] = self.crop(results[key], crop_bbox)

        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(crop_size={self.crop_size})'


@TRANSFORMS.register_module()
class RandomRotate(object):
    def __init__(self,
                 prob,
                 degree,
                 pad_val=0,
                 seg_pad_val=255,
                 center=None,
                 auto_bound=False):
        self.prob = prob
        assert prob >= 0 and prob <= 1
        if isinstance(degree, (float, int)):
            assert degree > 0, f'degree {degree} should be positive'
            self.degree = (-degree, degree)
        else:
            self.degree = degree
        assert len(self.degree) == 2, f'degree {self.degree} should be a ' \
                                      f'tuple of (min, max)'
        self.pal_val = pad_val
        self.seg_pad_val = seg_pad_val
        self.center = center
        self.auto_bound = auto_bound

    def __call__(self, results):
        rotate = True if np.random.rand() < self.prob else False
        degree = np.random.uniform(min(*self.degree), max(*self.degree))
        if rotate:
            # rotate image
            results['img'] = imrotate(results['img'],
                                      angle=degree,
                                      border_value=self.pal_val,
                                      center=self.center,
                                      auto_bound=self.auto_bound)

            # rotate segs
            for key in results.get('seg_fields', []):
                results[key] = imrotate(results[key],
                                        angle=degree,
                                        border_value=self.seg_pad_val,
                                        center=self.center,
                                        auto_bound=self.auto_bound,
                                        interpolation='nearest')
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(prob={self.prob}, ' \
                    f'degree={self.degree}, ' \
                    f'pad_val={self.pal_val}, ' \
                    f'seg_pad_val={self.seg_pad_val}, ' \
                    f'center={self.center}, ' \
                    f'auto_bound={self.auto_bound})'
        return repr_str


@TRANSFORMS.register_module()
class RandomRotate90(object):
    def __init__(self,
                 prob,
                 pad_val=0,
                 seg_pad_val=255,
                 center=None,
                 auto_bound=False):
        self.prob = prob
        assert prob >= 0 and prob <= 1
        self.pal_val = pad_val
        self.seg_pad_val = seg_pad_val
        self.center = center
        self.auto_bound = auto_bound

    def __call__(self, results):
        rotate = True if np.random.rand() < self.prob else False
        degree = np.random.randint(1, 4) * 90
        if rotate:
            # rotate image
            results['img'] = imrotate(results['img'],
                                      angle=degree,
                                      border_value=self.pal_val,
                                      center=self.center,
                                      auto_bound=self.auto_bound)

            # rotate segs
            for key in results.get('seg_fields', []):
                results[key] = imrotate(results[key],
                                        angle=degree,
                                        border_value=self.seg_pad_val,
                                        center=self.center,
                                        auto_bound=self.auto_bound,
                                        interpolation='nearest')
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(prob={self.prob}, ' \
                    f'degree={self.degree}, ' \
                    f'pad_val={self.pal_val}, ' \
                    f'seg_pad_val={self.seg_pad_val}, ' \
                    f'center={self.center}, ' \
                    f'auto_bound={self.auto_bound})'
        return repr_str


@TRANSFORMS.register_module()
class RGB2Gray(object):
    def __init__(self, out_channels=None, weights=(0.299, 0.587, 0.114)):
        assert out_channels is None or out_channels > 0
        self.out_channels = out_channels
        assert isinstance(weights, tuple)
        for item in weights:
            assert isinstance(item, (float, int))
        self.weights = weights

    def __call__(self, results):
        img = results['img']
        assert len(img.shape) == 3
        assert img.shape[2] == len(self.weights)
        weights = np.array(self.weights).reshape((1, 1, -1))
        img = (img * weights).sum(2, keepdims=True)
        if self.out_channels is None:
            img = img.repeat(weights.shape[2], axis=2)
        else:
            img = img.repeat(self.out_channels, axis=2)

        results['img'] = img
        results['img_shape'] = img.shape

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(out_channels={self.out_channels}, ' \
                    f'weights={self.weights})'
        return repr_str


@TRANSFORMS.register_module()
class SegRescale(object):
    def __init__(self, scale_factor=1):
        self.scale_factor = scale_factor

    def __call__(self, results):
        for key in results.get('seg_fields', []):
            if self.scale_factor != 1:
                results[key] = imrescale(results[key],
                                         self.scale_factor,
                                         interpolation='nearest')
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(scale_factor={self.scale_factor})'


@TRANSFORMS.register_module()
class PhotoMetricDistortion(object):
    def __init__(self,
                 brightness_delta=32,
                 contrast_range=(0.5, 1.5),
                 saturation_range=(0.5, 1.5),
                 hue_delta=18):
        self.brightness_delta = brightness_delta
        self.contrast_lower, self.contrast_upper = contrast_range
        self.saturation_lower, self.saturation_upper = saturation_range
        self.hue_delta = hue_delta

    def convert(self, img, alpha=1, beta=0):
        img = img.astype(np.float32) * alpha + beta
        img = np.clip(img, 0, 255)
        return img.astype(np.uint8)

    def brightness(self, img):
        if random.randint(2):
            return self.convert(img,
                                beta=random.uniform(-self.brightness_delta,
                                                    self.brightness_delta))
        return img

    def contrast(self, img):
        if random.randint(2):
            return self.convert(img,
                                alpha=random.uniform(self.contrast_lower,
                                                     self.contrast_upper))
        return img

    def saturation(self, img):
        if random.randint(2):
            img = bgr2hsv(img)
            img[:, :,
                1] = self.convert(img[:, :, 1],
                                  alpha=random.uniform(self.saturation_lower,
                                                       self.saturation_upper))
            img = hsv2bgr(img)
        return img

    def hue(self, img):
        if random.randint(2):
            img = bgr2hsv(img)
            img[:, :,
                0] = (img[:, :, 0].astype(int) +
                      random.randint(-self.hue_delta, self.hue_delta)) % 180
            img = hsv2bgr(img)
        return img

    def __call__(self, results):
        img = results['img']
        # random brightness
        img = self.brightness(img)

        # mode == 0 --> do random contrast first
        # mode == 1 --> do random contrast last
        mode = random.randint(2)
        if mode == 1:
            img = self.contrast(img)

        # random saturation
        img = self.saturation(img)

        # random hue
        img = self.hue(img)

        # random contrast
        if mode == 0:
            img = self.contrast(img)

        results['img'] = img
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += (f'(brightness_delta={self.brightness_delta}, '
                     f'contrast_range=({self.contrast_lower}, '
                     f'{self.contrast_upper}), '
                     f'saturation_range=({self.saturation_lower}, '
                     f'{self.saturation_upper}), '
                     f'hue_delta={self.hue_delta})')
        return repr_str
python/jittorseg/datasets/pipelines/transforms.py
JittorCV/jittorseg
import cv2
import numbers
from PIL import Image, ImageOps
import numpy as np
from jittorseg.utils.general import is_tuple_of
from jittorseg.utils.helpers import to_2tuple
from cv2 import (IMREAD_COLOR, IMREAD_GRAYSCALE, IMREAD_IGNORE_ORIENTATION,
                 IMREAD_UNCHANGED)
try:
    import tifffile
except ImportError:
    tifffile = None

supported_backends = ['cv2', 'pillow', 'tifffile']
pillow_interp_codes = {
    'nearest': Image.NEAREST,
    'bilinear': Image.BILINEAR,
    'bicubic': Image.BICUBIC,
    'box': Image.BOX,
    'lanczos': Image.LANCZOS,
    'hamming': Image.HAMMING
}

cv2_interp_codes = {
    'nearest': cv2.INTER_NEAREST,
    'bilinear': cv2.INTER_LINEAR,
    'bicubic': cv2.INTER_CUBIC,
    'area': cv2.INTER_AREA,
    'lanczos': cv2.INTER_LANCZOS4
}

imread_flags = {
    'color': IMREAD_COLOR,
    'grayscale': IMREAD_GRAYSCALE,
    'unchanged': IMREAD_UNCHANGED,
    'color_ignore_orientation': IMREAD_IGNORE_ORIENTATION | IMREAD_COLOR,
    'grayscale_ignore_orientation':
    IMREAD_IGNORE_ORIENTATION | IMREAD_GRAYSCALE
}
imread_backend = 'cv2'


def imread(img_path, flag='color', channel_order='bgr', backend=None):
    if backend is None:
        backend = imread_backend
    if backend not in supported_backends:
        raise ValueError(f'backend: {backend} is not supported. Supported '
                         "backends are 'cv2', 'pillow', 'tifffile'")

    if backend == 'pillow':
        img = Image.open(img_path)
        img = _pillow2array(img, flag, channel_order)
        return img
    elif backend == 'tifffile':
        img = tifffile.imread(img_path)
        return img
    else:
        flag = imread_flags[flag]
        img = cv2.imread(img_path, flag)
        if flag == IMREAD_COLOR and channel_order == 'rgb':
            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
        return img


def _pillow2array(img, flag='color', channel_order='bgr'):
    channel_order = channel_order.lower()
    if channel_order not in ['rgb', 'bgr']:
        raise ValueError('channel order must be either "rgb" or "bgr"')

    if flag == 'unchanged':
        array = np.array(img)
        if array.ndim >= 3 and array.shape[2] >= 3:  # color image
            array[:, :, :3] = array[:, :, (2, 1, 0)]  # RGB to BGR
    else:
        # Handle exif orientation tag
        if flag in ['color', 'grayscale']:
            img = ImageOps.exif_transpose(img)
        # If the image mode is not 'RGB', convert it to 'RGB' first.
        if img.mode != 'RGB':
            if img.mode != 'LA':
                # Most formats except 'LA' can be directly converted to RGB
                img = img.convert('RGB')
            else:
                # When the mode is 'LA', the default conversion will fill in
                #  the canvas with black, which sometimes shadows black objects
                #  in the foreground.
                #
                # Therefore, a random color (124, 117, 104) is used for canvas
                img_rgba = img.convert('RGBA')
                img = Image.new('RGB', img_rgba.size, (124, 117, 104))
                img.paste(img_rgba, mask=img_rgba.split()[3])  # 3 is alpha
        if flag in ['color', 'color_ignore_orientation']:
            array = np.array(img)
            if channel_order != 'rgb':
                array = array[:, :, ::-1]  # RGB to BGR
        elif flag in ['grayscale', 'grayscale_ignore_orientation']:
            img = img.convert('L')
            array = np.array(img)
        else:
            raise ValueError(
                'flag must be "color", "grayscale", "unchanged", '
                f'"color_ignore_orientation" or "grayscale_ignore_orientation"'
                f' but got {flag}')
    return array


def imresize_to_multiple(img,
                         divisor,
                         size=None,
                         scale_factor=None,
                         keep_ratio=False,
                         return_scale=False,
                         interpolation='bilinear',
                         out=None,
                         backend=None):
    h, w = img.shape[:2]
    if size is not None and scale_factor is not None:
        raise ValueError('only one of size or scale_factor should be defined')
    elif size is None and scale_factor is None:
        raise ValueError('one of size or scale_factor should be defined')
    elif size is not None:
        size = to_2tuple(size)
        if keep_ratio:
            size = rescale_size((w, h), size, return_scale=False)
    else:
        size = _scale_size((w, h), scale_factor)

    divisor = to_2tuple(divisor)
    size = tuple(int(np.ceil(s / d)) * d for s, d in zip(size, divisor))
    resized_img, w_scale, h_scale = imresize(img,
                                             size,
                                             return_scale=True,
                                             interpolation=interpolation,
                                             out=out,
                                             backend=backend)
    if return_scale:
        return resized_img, w_scale, h_scale
    else:
        return resized_img


def imresize(img,
             size,
             return_scale=False,
             interpolation='bilinear',
             out=None,
             backend=None):
    h, w = img.shape[:2]
    if backend is None:
        backend = imread_backend
    if backend not in ['cv2', 'pillow']:
        raise ValueError(f'backend: {backend} is not supported for resize.'
                         f"Supported backends are 'cv2', 'pillow'")

    if backend == 'pillow':
        assert img.dtype == np.uint8, 'Pillow backend only support uint8 type'
        pil_image = Image.fromarray(img)
        pil_image = pil_image.resize(size, pillow_interp_codes[interpolation])
        resized_img = np.array(pil_image)
    else:
        resized_img = cv2.resize(img,
                                 size,
                                 dst=out,
                                 interpolation=cv2_interp_codes[interpolation])
    if not return_scale:
        return resized_img
    else:
        w_scale = size[0] / w
        h_scale = size[1] / h
        return resized_img, w_scale, h_scale


def _scale_size(size, scale):
    if isinstance(scale, (float, int)):
        scale = (scale, scale)
    w, h = size
    return int(w * float(scale[0]) + 0.5), int(h * float(scale[1]) + 0.5)


def rescale_size(old_size, scale, return_scale=False):
    w, h = old_size
    if isinstance(scale, (float, int)):
        if scale <= 0:
            raise ValueError(f'Invalid scale {scale}, must be positive.')
        scale_factor = scale
    elif isinstance(scale, tuple):
        max_long_edge = max(scale)
        max_short_edge = min(scale)
        scale_factor = min(max_long_edge / max(h, w),
                           max_short_edge / min(h, w))
    else:
        raise TypeError(
            f'Scale must be a number or tuple of int, but got {type(scale)}')

    new_size = _scale_size((w, h), scale_factor)

    if return_scale:
        return new_size, scale_factor
    else:
        return new_size


def imrescale(img, scale, return_scale=False, interpolation='bilinear'):
    h, w = img.shape[:2]
    new_size, scale_factor = rescale_size((w, h), scale, return_scale=True)
    rescaled_img = imresize(img, new_size, interpolation=interpolation)
    if return_scale:
        return rescaled_img, scale_factor
    else:
        return rescaled_img


def imflip(img, direction='horizontal'):
    assert direction in ['horizontal', 'vertical', 'diagonal']
    if direction == 'horizontal':
        return np.flip(img, axis=1)
    elif direction == 'vertical':
        return np.flip(img, axis=0)
    else:
        return np.flip(img, axis=(0, 1))


def impad(img,
          *,
          shape=None,
          padding=None,
          pad_val=0,
          padding_mode='constant'):
    assert (shape is not None) ^ (padding is not None)
    if shape is not None:
        width = max(shape[1] - img.shape[1], 0)
        height = max(shape[0] - img.shape[0], 0)
        padding = (0, 0, width, height)

    # check pad_val
    if isinstance(pad_val, tuple):
        assert len(pad_val) == img.shape[-1]
    elif not isinstance(pad_val, numbers.Number):
        raise TypeError('pad_val must be a int or a tuple. '
                        f'But received {type(pad_val)}')

    # check padding
    if isinstance(padding, tuple) and len(padding) in [2, 4]:
        if len(padding) == 2:
            padding = (padding[0], padding[1], padding[0], padding[1])
    elif isinstance(padding, numbers.Number):
        padding = (padding, padding, padding, padding)
    else:
        raise ValueError('Padding must be a int or a 2, or 4 element tuple.'
                         f'But received {padding}')

    # check padding mode
    assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']

    border_type = {
        'constant': cv2.BORDER_CONSTANT,
        'edge': cv2.BORDER_REPLICATE,
        'reflect': cv2.BORDER_REFLECT_101,
        'symmetric': cv2.BORDER_REFLECT
    }
    img = cv2.copyMakeBorder(img,
                             padding[1],
                             padding[3],
                             padding[0],
                             padding[2],
                             border_type[padding_mode],
                             value=pad_val)

    return img


def impad_to_multiple(img, divisor, pad_val=0):
    pad_h = int(np.ceil(img.shape[0] / divisor)) * divisor
    pad_w = int(np.ceil(img.shape[1] / divisor)) * divisor
    return impad(img, shape=(pad_h, pad_w), pad_val=pad_val)


def imnormalize(img, mean, std, to_rgb=True):
    img = img.copy().astype(np.float32)
    return imnormalize_(img, mean, std, to_rgb)


def imnormalize_(img, mean, std, to_rgb=True):
    # cv2 inplace normalization does not accept uint8
    assert img.dtype != np.uint8
    mean = np.float64(mean.reshape(1, -1))
    stdinv = 1 / np.float64(std.reshape(1, -1))
    if to_rgb:
        cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
    cv2.subtract(img, mean, img)  # inplace
    cv2.multiply(img, stdinv, img)  # inplace
    return img


def clahe(img, clip_limit=40.0, tile_grid_size=(8, 8)):
    assert isinstance(img, np.ndarray)
    assert img.ndim == 2
    assert isinstance(clip_limit, (float, int))
    assert is_tuple_of(tile_grid_size, int)
    assert len(tile_grid_size) == 2

    clahe = cv2.createCLAHE(clip_limit, tile_grid_size)
    return clahe.apply(np.array(img, dtype=np.uint8))


def imrotate(img,
             angle,
             center=None,
             scale=1.0,
             border_value=0,
             interpolation='bilinear',
             auto_bound=False):
    if center is not None and auto_bound:
        raise ValueError('`auto_bound` conflicts with `center`')
    h, w = img.shape[:2]
    if center is None:
        center = ((w - 1) * 0.5, (h - 1) * 0.5)
    assert isinstance(center, tuple)

    matrix = cv2.getRotationMatrix2D(center, -angle, scale)
    if auto_bound:
        cos = np.abs(matrix[0, 0])
        sin = np.abs(matrix[0, 1])
        new_w = h * sin + w * cos
        new_h = h * cos + w * sin
        matrix[0, 2] += (new_w - w) * 0.5
        matrix[1, 2] += (new_h - h) * 0.5
        w = int(np.round(new_w))
        h = int(np.round(new_h))
    rotated = cv2.warpAffine(img,
                             matrix, (w, h),
                             flags=cv2_interp_codes[interpolation],
                             borderValue=border_value)
    return rotated


def convert_color_factory(src, dst):

    code = getattr(cv2, f'COLOR_{src.upper()}2{dst.upper()}')

    def convert_color(img):
        out_img = cv2.cvtColor(img, code)
        return out_img

    convert_color.__doc__ = f"""Convert a {src.upper()} image to {dst.upper()}
        image.

    Args:
        img (ndarray or str): The input image.

    Returns:
        ndarray: The converted {dst.upper()} image.
    """

    return convert_color


bgr2rgb = convert_color_factory('bgr', 'rgb')

rgb2bgr = convert_color_factory('rgb', 'bgr')

bgr2hsv = convert_color_factory('bgr', 'hsv')

hsv2bgr = convert_color_factory('hsv', 'bgr')

bgr2hls = convert_color_factory('bgr', 'hls')

hls2bgr = convert_color_factory('hls', 'bgr')
python/jittorseg/datasets/pipelines/utils.py
JittorCV/jittorseg
from jittorseg.utils.registry import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class PotsdamDataset(CustomDataset):
    """ISPRS Potsdam dataset.
    In segmentation map annotation for Potsdam dataset, 0 is the ignore index.
    ``reduce_zero_label`` should be set to True. The ``img_suffix`` and
    ``seg_map_suffix`` are both fixed to '.png'.
    """
    CLASSES = ('impervious_surface', 'building', 'low_vegetation', 'tree',
               'car', 'clutter')

    PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0],
               [255, 255, 0], [255, 0, 0]]

    def __init__(self, **kwargs):
        super(PotsdamDataset, self).__init__(
            img_suffix='.png',
            seg_map_suffix='.png',
            reduce_zero_label=True,
            **kwargs)

python/jittorseg/datasets/potsdam.py
JittorCV/jittorseg
import os.path as osp

from jittorseg.utils.registry import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class PascalVOCDataset(CustomDataset):
    CLASSES = ('background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
               'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
               'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
               'train', 'tvmonitor')

    PALETTE = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
               [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0],
               [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128],
               [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0],
               [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]]

    def __init__(self, split, **kwargs):
        super(PascalVOCDataset, self).__init__(img_suffix='.jpg',
                                               seg_map_suffix='.png',
                                               split=split,
                                               **kwargs)
        assert osp.exists(self.img_dir) and self.split is not None
python/jittorseg/datasets/voc.py
JittorCV/jittorseg
from .priority import Priority
from .hook import Hook
from .checkpoint_hook import CheckpointHook
from .eval_hook import EvalHook
from .log_printer_hook import LogPrinterHook
from .loss_parser_hook import LossParserHook
from .lr_scheduler_hook import LRSchedulerHook
from .optimizer_hook import OptimizerHook

python/jittorseg/hooks/__init__.py
JittorCV/jittorseg
from jittorseg.utils.registry import HOOKS
from .hook import Hook
from jittorseg.utils.general import check_interval


@HOOKS.register_module()
class CheckpointHook(Hook):
    def __init__(self, runner, priority, checkpoint_interval):
        super(CheckpointHook, self).__init__(runner=runner, priority=priority)
        self.checkpoint_interval = checkpoint_interval

    def after_iter(self):
        if check_interval(self.runner.iter + 1, self.checkpoint_interval):
            self.runner.save()

python/jittorseg/hooks/checkpoint_hook.py
JittorCV/jittorseg
import jittor as jt
from jittorseg.utils.registry import HOOKS
from .hook import Hook
from jittorseg.utils.general import check_interval
from tqdm import tqdm
from jittorseg.utils.general import np2tmp


@HOOKS.register_module()
class EvalHook(Hook):
    def __init__(self, runner, priority, eval_interval, efficient_val):
        super(EvalHook, self).__init__(runner=runner, priority=priority)
        self.eval_interval = eval_interval
        self.efficient_val = efficient_val

    def after_iter(self):
        if check_interval(self.runner.iter + 1, self.eval_interval):
            self.val()
            self.runner.model.train()

    @jt.no_grad()
    @jt.single_process_scope()
    def val(self):
        if self.runner.val_dataset is None:
            self.runner.logger.print_log("Please set Val dataset")
        else:
            self.runner.logger.print_log("Validating....")
            if self.runner.model.is_training():
                self.runner.model.eval()
            results = []
            for _, (data) in tqdm(enumerate(self.runner.val_dataset)):
                images = data['img']
                img_metas = data['img_metas']
                result = self.runner.model(images, img_metas, return_loss=False)
                jt.sync_all(True)
                if isinstance(result, list):
                    if self.efficient_val:
                        result = [np2tmp(_) for _ in result]
                    results.extend(result)
                else:
                    if self.efficient_val:
                        result = np2tmp(result)
                    results.append(result)

            self.runner.val_dataset.evaluate(results,
                                             metric='mIoU',
                                             logger=self.runner.logger)

python/jittorseg/hooks/eval_hook.py
JittorCV/jittorseg
from .priority import get_priority, Priority


class Hook():
    def __init__(self, runner, priority):
        self.runner = runner
        self._priority = get_priority(priority)

    stages = (
        'before_run', 'after_run',
        'before_epoch', 'after_epoch',
        'before_iter', 'after_iter'
    )

    @property
    def name(self):
        return self.__class__.__name__

    @property
    def priority(self):
        return self._priority

    def before_run(self):
        pass

    def after_run(self):
        pass

    def before_epoch(self):
        pass

    def after_epoch(self):
        pass

    def before_iter(self):
        pass

    def after_iter(self):
        pass

python/jittorseg/hooks/hook.py
JittorCV/jittorseg
import jittor as jt
from jittorseg.utils.registry import HOOKS
from .hook import Hook
from jittorseg.utils.general import check_interval, sync
import time
import datetime

@HOOKS.register_module()
class LogPrinterHook(Hook):
    def __init__(self, runner, priority, log_interval):
        super(LogPrinterHook, self).__init__(runner=runner, priority=priority)
        self.log_interval = log_interval

    def before_run(self):
        self.runner.logger.print_log("Begin running")

    def before_iter(self):
        self.runner.outputs["start_time"] = time.time()

    def after_iter(self):
        if check_interval(self.runner.iter + 1, self.log_interval):
            start_time = self.runner.outputs["start_time"]
            ptime = time.time() - start_time
            batch_idx = self.runner.outputs["batch_idx"]
            all_loss = self.runner.outputs["all_loss"]
            losses = self.runner.outputs["losses"]
            batch_size = self.runner.outputs["batch_size"]
            fps = batch_size * (batch_idx + 1) / ptime
            eta_time = (self.runner.total_iter - self.runner.iter - 1) * ptime / (batch_idx +
                                                                1)
            eta_str = str(datetime.timedelta(seconds=int(eta_time)))
            data = dict(name=self.runner.cfg.name,
                        lr=self.runner.optimizer.cur_lr(),
                        iter=self.runner.iter + 1,
                        epoch=self.runner.epoch,
                        batch_idx=batch_idx,
                        batch_size=batch_size,
                        total_loss=all_loss,
                        fps=fps,
                        eta=eta_str)
            data.update(losses)
            data = sync(data)
            # is_main use jt.rank==0, so it's scope must have no jt.Vars
            if jt.rank == 0:
                self.runner.logger.log(data)

    def after_run(self):
        self.runner.logger.print_log("After running")
python/jittorseg/hooks/log_printer_hook.py
JittorCV/jittorseg
from jittorseg.utils.registry import HOOKS
from .hook import Hook
from jittorseg.utils.general import parse_losses


@HOOKS.register_module()
class LossParserHook(Hook):
    def __init__(self, runner, priority):
        super(LossParserHook, self).__init__(runner=runner, priority=priority)

    def after_iter(self):
        all_loss, loss = parse_losses(self.runner.outputs["losses"])
        self.runner.outputs["all_loss"] = all_loss
        self.runner.outputs["loss"] = loss

python/jittorseg/hooks/loss_parser_hook.py
JittorCV/jittorseg
from jittorseg.utils.registry import HOOKS
from .hook import Hook


@HOOKS.register_module()
class LRSchedulerHook(Hook):
    def __init__(self, runner, priority, by_epoch):
        super(LRSchedulerHook, self).__init__(runner=runner, priority=priority)
        self.by_epoch = by_epoch

    def after_iter(self):
        self.runner.scheduler.step(self.runner.iter, self.runner.epoch, by_epoch=self.by_epoch)

python/jittorseg/hooks/lr_scheduler_hook.py
JittorCV/jittorseg
from jittorseg.utils.registry import HOOKS
from .hook import Hook


@HOOKS.register_module()
class OptimizerHook(Hook):
    def __init__(self, runner, priority):
        super(OptimizerHook, self).__init__(runner=runner, priority=priority)

    def after_iter(self):
        all_loss = self.runner.outputs["all_loss"]
        self.runner.optimizer.step(all_loss)

python/jittorseg/hooks/optimizer_hook.py
JittorCV/jittorseg
from enum import Enum


class Priority(Enum):
    """Hook priority levels.

    +--------------+------------+
    | Level        | Value      |
    +==============+============+
    | HIGHEST      | 0          |
    +--------------+------------+
    | VERY_HIGH    | 10         |
    +--------------+------------+
    | HIGH         | 30         |
    +--------------+------------+
    | ABOVE_NORMAL | 40         |
    +--------------+------------+
    | NORMAL       | 50         |
    +--------------+------------+
    | BELOW_NORMAL | 60         |
    +--------------+------------+
    | LOW          | 70         |
    +--------------+------------+
    | VERY_LOW     | 90         |
    +--------------+------------+
    | LOWEST       | 100        |
    +--------------+------------+
    """

    HIGHEST = 0
    VERY_HIGH = 10
    HIGH = 30
    ABOVE_NORMAL = 40
    NORMAL = 50
    BELOW_NORMAL = 60
    LOW = 70
    VERY_LOW = 90
    LOWEST = 100


def get_priority(priority):
    """Get priority value.

    Args:
        priority (int or str or :obj:`Priority`): Priority.

    Returns:
        int: The priority value.
    """
    if isinstance(priority, int):
        if priority < 0 or priority > 100:
            raise ValueError('priority must be between 0 and 100')
        return priority
    elif isinstance(priority, Priority):
        return priority.value
    elif isinstance(priority, str):
        return Priority[priority.upper()].value
    else:
        raise TypeError('priority must be an integer or Priority enum value')

python/jittorseg/hooks/priority.py
JittorCV/jittorseg
from .backbones import *
from .basic_operators import *
from .decode_heads import *
from .losses import *
from .necks import *
from .layers import *
from .segmentors import *

python/jittorseg/models/__init__.py
JittorCV/jittorseg
from .mix_transformer import *
from .resnet import *
from .mscan import MSCAN
from .swin import SwinTransformer
from .resnext import ResNeXt
from .resnest import ResNeSt
from .convnext import ConvNeXt
from .vit import VisionTransformer
from .beit import BEiT
from .mae import MAE
from .mobilenet_v2 import MobileNetV2

python/jittorseg/models/backbones/__init__.py
JittorCV/jittorseg
import numpy as np
import jittor as jt
from jittor import nn
from ..basic_operators import build_norm_layer, build_dropout
from jittorseg.utils.weight_init import (constant_init, kaiming_init, trunc_normal_)
from jittorseg.utils.helpers import to_2tuple
from jittor.nn import BatchNorm as _BatchNorm

from jittorseg.utils.registry import BACKBONES
from jittorseg.models.basic_operators.embed import PatchEmbed
from .vit import TransformerEncoderLayer as VisionTransformerEncoderLayer

try:
    from scipy import interpolate
except ImportError:
    interpolate = None


class BEiTAttention(nn.Module):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 window_size,
                 bias='qv_bias',
                 qk_scale=None,
                 attn_drop_rate=0.,
                 proj_drop_rate=0.,
                 **kwargs):
        super().__init__()
        self.embed_dims = embed_dims
        self.num_heads = num_heads
        head_embed_dims = embed_dims // num_heads
        self.bias = bias
        self.scale = qk_scale or head_embed_dims**-0.5

        qkv_bias = bias
        if bias == 'qv_bias':
            self._init_qv_bias()
            qkv_bias = False

        self.window_size = window_size
        self._init_rel_pos_embedding()

        self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop_rate)
        self.proj = nn.Linear(embed_dims, embed_dims)
        self.proj_drop = nn.Dropout(proj_drop_rate)

    def _init_qv_bias(self):
        self.q_bias = jt.zeros((self.embed_dims))
        self.v_bias = jt.zeros((self.embed_dims))

    def _init_rel_pos_embedding(self):
        Wh, Ww = self.window_size
        # cls to token & token 2 cls & cls to cls
        self.num_relative_distance = (2 * Wh - 1) * (2 * Ww - 1) + 3
        # relative_position_bias_table shape is (2*Wh-1 * 2*Ww-1 + 3, nH)
        self.relative_position_bias_table = jt.zeros(
            (self.num_relative_distance, self.num_heads))

        # get pair-wise relative position index for
        # each token inside the window
        coords_h = jt.arange(Wh)
        coords_w = jt.arange(Ww)
        # coords shape is (2, Wh, Ww)
        coords = jt.stack(jt.meshgrid([coords_h, coords_w]))
        # coords_flatten shape is (2, Wh*Ww)
        coords_flatten = jt.flatten(coords, 1)
        relative_coords = (coords_flatten[:, :, None] -
                           coords_flatten[:, None, :])
        # relative_coords shape is (Wh*Ww, Wh*Ww, 2)
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        # shift to start from 0
        relative_coords[:, :, 0] += Wh - 1
        relative_coords[:, :, 1] += Ww - 1
        relative_coords[:, :, 0] *= 2 * Ww - 1
        relative_position_index = jt.zeros(
            ((Wh * Ww + 1, ) * 2)).astype(relative_coords.dtype)
        # relative_position_index shape is (Wh*Ww, Wh*Ww)
        relative_position_index[1:, 1:] = relative_coords.sum(-1)
        relative_position_index[0, 0:] = self.num_relative_distance - 3
        relative_position_index[0:, 0] = self.num_relative_distance - 2
        relative_position_index[0, 0] = self.num_relative_distance - 1

        self.relative_position_index = relative_position_index

    def execute(self, x):
        """
        Args:
            x (Var): input features with shape of (num_windows*B, N, C).
        """
        B, N, C = x.shape
        if self.bias == 'qv_bias':
            k_bias = jt.zeros_like(self.v_bias)
            qkv_bias = jt.concat((self.q_bias, k_bias, self.v_bias))
            # linear
            qkv = x.matmul(self.qkv.weight.t()) + qkv_bias
        else:
            qkv = self.qkv(x)

        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))
        if self.relative_position_bias_table is not None:
            Wh = self.window_size[0]
            Ww = self.window_size[1]
            relative_position_bias = self.relative_position_bias_table[
                self.relative_position_index.view(-1)].view(
                    Wh * Ww + 1, Wh * Ww + 1, -1)
            relative_position_bias = relative_position_bias.permute(
                2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
            attn = attn + relative_position_bias.unsqueeze(0).float32()
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class BEiTTransformerEncoderLayer(VisionTransformerEncoderLayer):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 feedexecute_channels,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 num_fcs=2,
                 bias='qv_bias',
                 act_cfg=dict(type='GELU'),
                 norm_cfg=dict(type='LN'),
                 window_size=None,
                 attn_cfg=dict(),
                 ffn_cfg=dict(add_identity=False),
                 init_values=None):
        attn_cfg.update(dict(window_size=window_size, qk_scale=None))

        super(BEiTTransformerEncoderLayer,
              self).__init__(embed_dims=embed_dims,
                             num_heads=num_heads,
                             feedexecute_channels=feedexecute_channels,
                             attn_drop_rate=attn_drop_rate,
                             drop_path_rate=0.,
                             drop_rate=0.,
                             num_fcs=num_fcs,
                             qkv_bias=bias,
                             act_cfg=act_cfg,
                             norm_cfg=norm_cfg,
                             attn_cfg=attn_cfg,
                             ffn_cfg=ffn_cfg)

        # NOTE: drop path for stochastic depth, we shall see if
        # this is better than dropout here
        dropout_layer = dict(type='DropPath', p=drop_path_rate)
        self.drop_path = build_dropout(
            dropout_layer) if dropout_layer else nn.Identity()
        self.gamma_1 = init_values * jt.ones((embed_dims))
        self.gamma_2 = init_values * jt.ones((embed_dims))

    def build_attn(self, attn_cfg):
        self.attn = BEiTAttention(**attn_cfg)

    def execute(self, x):
        x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x)))
        x = x + self.drop_path(self.gamma_2 * self.ffn(self.norm2(x)))
        return x


@BACKBONES.register_module()
class BEiT(nn.Module):

    def __init__(self,
                 img_size=224,
                 patch_size=16,
                 in_channels=3,
                 embed_dims=768,
                 num_layers=12,
                 num_heads=12,
                 mlp_ratio=4,
                 out_indices=-1,
                 qv_bias=True,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 norm_cfg=dict(type='LN'),
                 act_cfg=dict(type='GELU'),
                 patch_norm=False,
                 final_norm=False,
                 num_fcs=2,
                 norm_eval=False,
                 pretrained=None,
                 init_values=0.1):
        super(BEiT, self).__init__()
        if isinstance(img_size, int):
            img_size = to_2tuple(img_size)
        elif isinstance(img_size, tuple):
            if len(img_size) == 1:
                img_size = to_2tuple(img_size[0])
            assert len(img_size) == 2, \
                f'The size of image should have length 1 or 2, ' \
                f'but got {len(img_size)}'

        self.in_channels = in_channels
        self.img_size = img_size
        self.patch_size = patch_size
        self.norm_eval = norm_eval
        self.pretrained = pretrained
        self.num_layers = num_layers
        self.embed_dims = embed_dims
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio
        self.attn_drop_rate = attn_drop_rate
        self.drop_path_rate = drop_path_rate
        self.num_fcs = num_fcs
        self.qv_bias = qv_bias
        self.act_cfg = act_cfg
        self.norm_cfg = norm_cfg
        self.patch_norm = patch_norm
        self.init_values = init_values
        self.window_size = (img_size[0] // patch_size,
                            img_size[1] // patch_size)
        self.patch_shape = self.window_size
        self.cls_token = jt.zeros((1, 1, embed_dims))

        self._build_patch_embedding()
        self._build_layers()

        if isinstance(out_indices, int):
            if out_indices == -1:
                out_indices = num_layers - 1
            self.out_indices = [out_indices]
        elif isinstance(out_indices, list) or isinstance(out_indices, tuple):
            self.out_indices = out_indices
        else:
            raise TypeError('out_indices must be type of int, list or tuple')

        self.final_norm = final_norm
        if final_norm:
            self.norm1_name, norm1 = build_norm_layer(norm_cfg,
                                                      embed_dims,
                                                      postfix=1)
            setattr(self, self.norm1_name, norm1)

    def _build_patch_embedding(self):
        """Build patch embedding layer."""
        self.patch_embed = PatchEmbed(
            in_channels=self.in_channels,
            embed_dims=self.embed_dims,
            conv_type='Conv2d',
            kernel_size=self.patch_size,
            stride=self.patch_size,
            padding=0,
            norm_cfg=self.norm_cfg if self.patch_norm else None)

    def _build_layers(self):
        """Build transformer encoding layers."""

        dpr = [
            x.item()
            for x in jt.linspace(0, self.drop_path_rate, self.num_layers)
        ]
        self.layers = nn.ModuleList()
        for i in range(self.num_layers):
            self.layers.append(
                BEiTTransformerEncoderLayer(
                    embed_dims=self.embed_dims,
                    num_heads=self.num_heads,
                    feedexecute_channels=self.mlp_ratio * self.embed_dims,
                    attn_drop_rate=self.attn_drop_rate,
                    drop_path_rate=dpr[i],
                    num_fcs=self.num_fcs,
                    bias='qv_bias' if self.qv_bias else False,
                    act_cfg=self.act_cfg,
                    norm_cfg=self.norm_cfg,
                    window_size=self.window_size,
                    init_values=self.init_values))

    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    def _geometric_sequence_interpolation(self, src_size, dst_size, sequence,
                                          num):
        """Get new sequence via geometric sequence interpolation.

        Args:
            src_size (int): Pos_embedding size in pre-trained model.
            dst_size (int): Pos_embedding size in the current model.
            sequence (Var): The relative position bias of the pretrain
                model after removing the extra tokens.
            num (int): Number of attention heads.
        Returns:
            new_sequence (Var): Geometric sequence interpolate the
                pre-trained relative position bias to the size of
                the current model.
        """

        def geometric_progression(a, r, n):
            return a * (1.0 - r**n) / (1.0 - r)

        # Here is a binary function.
        left, right = 1.01, 1.5
        while right - left > 1e-6:
            q = (left + right) / 2.0
            gp = geometric_progression(1, q, src_size // 2)
            if gp > dst_size // 2:
                right = q
            else:
                left = q
        # The position of each interpolated point is determined
        # by the ratio obtained by dichotomy.
        dis = []
        cur = 1
        for i in range(src_size // 2):
            dis.append(cur)
            cur += q**(i + 1)
        r_ids = [-_ for _ in reversed(dis)]
        x = r_ids + [0] + dis
        y = r_ids + [0] + dis
        t = dst_size // 2.0
        dx = np.arange(-t, t + 0.1, 1.0)
        dy = np.arange(-t, t + 0.1, 1.0)
        # Interpolation functions are being executed and called.
        new_sequence = []
        for i in range(num):
            if isinstance(sequence, jt.Var):
                z = sequence[:, i].view(src_size, src_size).float().numpy()
                f = interpolate.interp2d(x, y, z, kind='cubic')
                new_sequence.append(jt.Var(f(dx, dy)).view(-1, 1))
            else:
                z = jt.Var(sequence[:, i]).view(src_size,
                                                src_size).float().numpy()
                f = interpolate.interp2d(x, y, z, kind='cubic')
                new_sequence.append(jt.Var(f(dx, dy)).view(-1, 1))

        new_sequence = jt.concat(new_sequence, dim=-1)
        return new_sequence

    def resize_rel_pos_embed(self, checkpoint):
        """Resize relative pos_embed weights.

        This function is modified from
        https://github.com/microsoft/unilm/blob/master/beit/semantic_segmentation/mmcv_custom/checkpoint.py.  # noqa: E501
        Copyright (c) Microsoft Corporation
        Licensed under the MIT License
        Args:
            checkpoint (dict): Key and value of the pretrain model.
        Returns:
            state_dict (dict): Interpolate the relative pos_embed weights
                in the pre-train model to the current model size.
        """
        if 'state_dict' in checkpoint:
            state_dict = checkpoint['state_dict']
        else:
            state_dict = checkpoint

        all_keys = list(state_dict.keys())
        for key in all_keys:
            if 'relative_position_index' in key:
                state_dict.pop(key)
            # In order to keep the center of pos_bias as consistent as
            # possible after interpolation, and vice versa in the edge
            # area, the geometric sequence interpolation method is adopted.
            if 'relative_position_bias_table' in key:
                rel_pos_bias = state_dict[key]
                if isinstance(rel_pos_bias, jt.Var):
                    src_num_pos, num_attn_heads = rel_pos_bias.size()
                    dst_num_pos, _ = self.state_dict()[key].size()
                else:
                    src_num_pos, num_attn_heads = rel_pos_bias.shape
                    dst_num_pos, _ = self.state_dict()[key].shape
                dst_patch_shape = self.patch_shape
                if dst_patch_shape[0] != dst_patch_shape[1]:
                    raise NotImplementedError()
                # Count the number of extra tokens.
                num_extra_tokens = dst_num_pos - (
                    dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
                src_size = int((src_num_pos - num_extra_tokens)**0.5)
                dst_size = int((dst_num_pos - num_extra_tokens)**0.5)
                if src_size != dst_size:
                    extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
                    rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
                    new_rel_pos_bias = self._geometric_sequence_interpolation(
                        src_size, dst_size, rel_pos_bias, num_attn_heads)
                    new_rel_pos_bias = jt.concat(
                        (new_rel_pos_bias, extra_tokens), dim=0)
                    state_dict[key] = new_rel_pos_bias

        return state_dict

    def init_weights(self, pretrained=None):

        def _init_weights(m):
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

        self.apply(_init_weights)
        if pretrained is None:
            trunc_normal_(self.cls_token, std=.02)
            for n, m in self.named_modules():
                if isinstance(m, nn.Linear):
                    trunc_normal_(m.weight, std=.02)
                    if m.bias is not None:
                        if 'ffn' in n:
                            nn.init.normal_(m.bias, mean=0., std=1e-6)
                        else:
                            nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Conv2d):
                    kaiming_init(m, mode='fan_in', bias=0.)
                elif isinstance(m, (_BatchNorm, nn.GroupNorm, nn.LayerNorm)):
                    constant_init(m, val=1.0, bias=0.)
        else:
            checkpoint = jt.load(pretrained)
            state_dict = self.resize_rel_pos_embed(checkpoint)
            self.load_parameters(state_dict)

    def execute(self, inputs):
        B = inputs.shape[0]
        x, hw_shape = self.patch_embed(inputs)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = jt.concat((cls_tokens, x), dim=1)

        outs = []
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i == len(self.layers) - 1:
                if self.final_norm:
                    x = self.norm1(x)
            if i in self.out_indices:
                # Remove class token and reshape token for decoder head
                out = x[:, 1:]
                B, _, C = out.shape
                out = out.reshape(B, hw_shape[0], hw_shape[1],
                                  C).permute(0, 3, 1, 2).contiguous()
                outs.append(out)

        return tuple(outs)

    def train(self, mode=True):
        super(BEiT, self).train(mode)
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, nn.LayerNorm):
                    m.eval()

python/jittorseg/models/backbones/beit.py
JittorCV/jittorseg
from functools import partial
from itertools import chain
from typing import Sequence
from collections import OrderedDict
import math

import jittor as jt
from jittor import nn
from jittorseg.utils.registry import BACKBONES
from jittorseg.utils.weight_init import trunc_normal_init, normal_init, constant_init


class LayerNorm2d(nn.LayerNorm):
    """LayerNorm on channels for 2d images.
    Args:
        num_channels (int): The number of channels of the input Var.
        eps (float): a value added to the denominator for numerical stability.
            Defaults to 1e-5.
        elementwise_affine (bool): a boolean value that when set to ``True``,
            this module has learnable per-element affine parameters initialized
            to ones (for weights) and zeros (for biases). Defaults to True.
    """

    def __init__(self, num_channels: int, **kwargs) -> None:
        super().__init__(num_channels, **kwargs)
        self.num_channels = self.normalized_shape[0]

    def execute(self, x):
        assert x.ndim == 4, 'LayerNorm2d only supports inputs with shape ' \
            f'(N, C, H, W), but got Var with shape {x.shape}'
        return nn.layer_norm(x.permute(0, 2, 3,
                                       1), self.normalized_shape, self.weight,
                             self.bias, self.eps).permute(0, 3, 1, 2)


class ConvNeXtBlock(nn.Module):
    """ConvNeXt Block.
    Args:
        in_channels (int): The number of input channels.
        norm_cfg (dict): The config dict for norm layers.
            Defaults to ``dict(type='LN2d', eps=1e-6)``.
        act_cfg (dict): The config dict for activation between pointwise
            convolution. Defaults to ``dict(type='GELU')``.
        mlp_ratio (float): The expansion ratio in both pointwise convolution.
            Defaults to 4.
        linear_pw_conv (bool): Whether to use linear layer to do pointwise
            convolution. More details can be found in the note.
            Defaults to True.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.
        layer_scale_init_value (float): Init value for Layer Scale.
            Defaults to 1e-6.
    Note:
        There are two equivalent implementations:
        1. DwConv -> LayerNorm -> 1x1 Conv -> GELU -> 1x1 Conv;
           all outputs are in (N, C, H, W).
        2. DwConv -> LayerNorm -> Permute to (N, H, W, C) -> Linear -> GELU
           -> Linear; Permute back
        As default, we use the second to align with the official repository.
        And it may be slightly faster.
    """

    def __init__(self,
                 in_channels,
                 mlp_ratio=4.,
                 linear_pw_conv=True,
                 drop_path_rate=0.,
                 layer_scale_init_value=1e-6):
        super().__init__()
        self.depthwise_conv = nn.Conv2d(in_channels,
                                        in_channels,
                                        kernel_size=7,
                                        padding=3,
                                        groups=in_channels)

        self.linear_pw_conv = linear_pw_conv
        self.norm = LayerNorm2d(in_channels)

        mid_channels = int(mlp_ratio * in_channels)
        if self.linear_pw_conv:
            # Use linear layer to do pointwise conv.
            pw_conv = nn.Linear
        else:
            pw_conv = partial(nn.Conv2d, kernel_size=1)

        self.pointwise_conv1 = pw_conv(in_channels, mid_channels)
        self.act = nn.GELU()
        self.pointwise_conv2 = pw_conv(mid_channels, in_channels)

        self.gamma = jt.ones(
            (in_channels)
        ) * layer_scale_init_value if layer_scale_init_value > 0 else None

        self.drop_path = nn.DropPath(
            drop_path_rate) if drop_path_rate > 0. else nn.Identity()

    def execute(self, x):
        shortcut = x
        x = self.depthwise_conv(x)
        x = self.norm(x)

        if self.linear_pw_conv:
            x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)

        x = self.pointwise_conv1(x)
        x = self.act(x)
        x = self.pointwise_conv2(x)

        if self.linear_pw_conv:
            x = x.permute(0, 3, 1, 2)  # permute back

        if self.gamma is not None:
            x = x.multiply(self.gamma.view(1, -1, 1, 1))

        x = shortcut + self.drop_path(x)
        return x


@BACKBONES.register_module()
class ConvNeXt(nn.Module):
    """ConvNeXt.
    A Jittor implementation of : `A ConvNet for the 2020s
    <https://arxiv.org/pdf/2201.03545.pdf>`_
    Modified from the `official repo
    <https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py>`_
    and `timm
    <https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/convnext.py>`_.
    Args:
        arch (str | dict): The model's architecture. If string, it should be
            one of architecture in ``ConvNeXt.arch_settings``. And if dict, it
            should include the following two keys:
            - depths (list[int]): Number of blocks at each stage.
            - channels (list[int]): The number of channels at each stage.
            Defaults to 'tiny'.
        in_channels (int): Number of input image channels. Defaults to 3.
        stem_patch_size (int): The size of one patch in the stem layer.
            Defaults to 4.
        norm_cfg (dict): The config dict for norm layers.
            Defaults to ``dict(type='LN2d', eps=1e-6)``.
        act_cfg (dict): The config dict for activation between pointwise
            convolution. Defaults to ``dict(type='GELU')``.
        linear_pw_conv (bool): Whether to use linear layer to do pointwise
            convolution. Defaults to True.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.
        layer_scale_init_value (float): Init value for Layer Scale.
            Defaults to 1e-6.
        out_indices (Sequence | int): Output from which stages.
            Defaults to -1, means the last stage.
        frozen_stages (int): Stages to be frozen (all param fixed).
            Defaults to 0, which means not freezing any parameters.
        gap_before_final_norm (bool): Whether to globally average the feature
            map before the final norm layer. In the official repo, it's only
            used in classification task. Defaults to True.
        init_cfg (dict, optional): Initialization config dict
    """  # noqa: E501
    arch_settings = {
        'tiny': {
            'depths': [3, 3, 9, 3],
            'channels': [96, 192, 384, 768]
        },
        'small': {
            'depths': [3, 3, 27, 3],
            'channels': [96, 192, 384, 768]
        },
        'base': {
            'depths': [3, 3, 27, 3],
            'channels': [128, 256, 512, 1024]
        },
        'large': {
            'depths': [3, 3, 27, 3],
            'channels': [192, 384, 768, 1536]
        },
        'xlarge': {
            'depths': [3, 3, 27, 3],
            'channels': [256, 512, 1024, 2048]
        },
    }

    def __init__(self,
                 arch='tiny',
                 in_channels=3,
                 stem_patch_size=4,
                 linear_pw_conv=True,
                 drop_path_rate=0.,
                 layer_scale_init_value=1e-6,
                 out_indices=-1,
                 frozen_stages=0,
                 gap_before_final_norm=True):
        super().__init__()

        if isinstance(arch, str):
            assert arch in self.arch_settings, \
                f'Unavailable arch, please choose from ' \
                f'({set(self.arch_settings)}) or pass a dict.'
            arch = self.arch_settings[arch]
        elif isinstance(arch, dict):
            assert 'depths' in arch and 'channels' in arch, \
                f'The arch dict must have "depths" and "channels", ' \
                f'but got {list(arch.keys())}.'

        self.depths = arch['depths']
        self.channels = arch['channels']
        assert (isinstance(self.depths, Sequence)
                and isinstance(self.channels, Sequence)
                and len(self.depths) == len(self.channels)), \
            f'The "depths" ({self.depths}) and "channels" ({self.channels}) ' \
            'should be both sequence with the same length.'

        self.num_stages = len(self.depths)

        if isinstance(out_indices, int):
            out_indices = [out_indices]
        assert isinstance(out_indices, Sequence), \
            f'"out_indices" must by a sequence or int, ' \
            f'get {type(out_indices)} instead.'
        for i, index in enumerate(out_indices):
            if index < 0:
                out_indices[i] = 4 + index
                assert out_indices[i] >= 0, f'Invalid out_indices {index}'
        self.out_indices = out_indices

        self.frozen_stages = frozen_stages
        self.gap_before_final_norm = gap_before_final_norm

        # stochastic depth decay rule
        dpr = [
            x.item() for x in jt.linspace(0, drop_path_rate, sum(self.depths))
        ]
        block_idx = 0

        # 4 downsample layers between stages, including the stem layer.
        self.downsample_layers = nn.ModuleList()
        stem = nn.Sequential(
            nn.Conv2d(in_channels,
                      self.channels[0],
                      kernel_size=stem_patch_size,
                      stride=stem_patch_size),
            LayerNorm2d(self.channels[0]),
        )
        self.downsample_layers.append(stem)

        # 4 feature resolution stages, each consisting of multiple residual
        # blocks
        self.stages = nn.ModuleList()

        for i in range(self.num_stages):
            depth = self.depths[i]
            channels = self.channels[i]

            if i >= 1:
                downsample_layer = nn.Sequential(
                    LayerNorm2d(self.channels[i - 1]),
                    nn.Conv2d(self.channels[i - 1],
                              channels,
                              kernel_size=2,
                              stride=2),
                )
                self.downsample_layers.append(downsample_layer)

            stage = nn.Sequential(*[
                ConvNeXtBlock(in_channels=channels,
                              drop_path_rate=dpr[block_idx + j],
                              linear_pw_conv=linear_pw_conv,
                              layer_scale_init_value=layer_scale_init_value)
                for j in range(depth)
            ])
            block_idx += depth

            self.stages.append(stage)

            if i in self.out_indices:
                norm_layer = LayerNorm2d(channels)
                layer_name = f'norm{i}'
                setattr(self, layer_name, norm_layer)

        self._freeze_stages()

    def execute(self, x):
        outs = []
        for i, stage in enumerate(self.stages):
            x = self.downsample_layers[i](x)
            x = stage(x)
            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                if self.gap_before_final_norm:
                    gap = x.mean([-2, -1], keepdim=True)
                    outs.append(norm_layer(gap).flatten(1))
                else:
                    outs.append(norm_layer(x))

        return tuple(outs)

    def init_weights(self, pretrained=None):
        if pretrained is None:
            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, val=1.0, bias=0.)
                elif isinstance(m, nn.Conv2d):
                    fan_out = m.kernel_size[0] * m.kernel_size[
                        1] * m.out_channels
                    fan_out //= m.groups
                    normal_init(m,
                                mean=0,
                                std=math.sqrt(2.0 / fan_out),
                                bias=0)
        elif isinstance(pretrained, str):
            ckpt = jt.load(pretrained)
            if 'state_dict' in ckpt:
                _state_dict = ckpt['state_dict']
            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()}
            self.load_parameters(state_dict)

    def _freeze_stages(self):
        for i in range(self.frozen_stages):
            downsample_layer = self.downsample_layers[i]
            stage = self.stages[i]
            downsample_layer.eval()
            stage.eval()
            for param in chain(downsample_layer.parameters(),
                               stage.parameters()):
                param.requires_grad = False

    def train(self, mode=True):
        super(ConvNeXt, self).train(mode)
        self._freeze_stages()

python/jittorseg/models/backbones/convnext.py
JittorCV/jittorseg
import math

import jittor as jt
from jittor import nn
from jittorseg.utils.weight_init import (constant_init, kaiming_init, trunc_normal_)
from jittor.nn import BatchNorm as _BatchNorm

from jittorseg.utils.registry import BACKBONES
from .beit import BEiT, BEiTAttention, BEiTTransformerEncoderLayer


class MAEAttention(BEiTAttention):
    """Multi-head self-attention with relative position bias used in MAE.

    This module is different from ``BEiTAttention`` by initializing the
    relative bias table with zeros.
    """

    def init_weights(self):
        """Initialize relative position bias with zeros."""

        # As MAE initializes relative position bias as zeros and this class
        # inherited from BEiT which initializes relative position bias
        # with `trunc_normal`, `init_weights` here does
        # nothing and just passes directly

        pass


class MAETransformerEncoderLayer(BEiTTransformerEncoderLayer):
    """Implements one encoder layer in Vision Transformer.

    This module is different from ``BEiTTransformerEncoderLayer`` by replacing
    ``BEiTAttention`` with ``MAEAttention``.
    """

    def build_attn(self, attn_cfg):
        self.attn = MAEAttention(**attn_cfg)


@BACKBONES.register_module()
class MAE(BEiT):
    """VisionTransformer with support for patch.

    Args:
        img_size (int | tuple): Input image size. Default: 224.
        patch_size (int): The patch size. Default: 16.
        in_channels (int): Number of input channels. Default: 3.
        embed_dims (int): embedding dimension. Default: 768.
        num_layers (int): depth of transformer. Default: 12.
        num_heads (int): number of attention heads. Default: 12.
        mlp_ratio (int): ratio of mlp hidden dim to embedding dim.
            Default: 4.
        out_indices (list | tuple | int): Output from which stages.
            Default: -1.
        attn_drop_rate (float): The drop out rate for attention layer.
            Default 0.0
        drop_path_rate (float): stochastic depth rate. Default 0.0.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='LN')
        act_cfg (dict): The activation config for FFNs.
            Default: dict(type='GELU').
        patch_norm (bool): Whether to add a norm in PatchEmbed Block.
            Default: False.
        final_norm (bool): Whether to add a additional layer to normalize
            final feature map. Default: False.
        num_fcs (int): The number of fully-connected layers for FFNs.
            Default: 2.
        norm_eval (bool): Whether to set norm layers to eval mode, namely,
            freeze running stats (mean and var). Note: Effect on Batch Norm
            and its variants only. Default: False.
        pretrained (str, optional): model pretrained path. Default: None.
        init_values (float): Initialize the values of Attention and FFN
            with learnable scaling. Defaults to 0.1.
        init_cfg (dict or list[dict], optional): Initialization config dict.
            Default: None.
    """

    def __init__(self,
                 img_size=224,
                 patch_size=16,
                 in_channels=3,
                 embed_dims=768,
                 num_layers=12,
                 num_heads=12,
                 mlp_ratio=4,
                 out_indices=-1,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 norm_cfg=dict(type='LN'),
                 act_cfg=dict(type='GELU'),
                 patch_norm=False,
                 final_norm=False,
                 num_fcs=2,
                 norm_eval=False,
                 init_values=0.1):
        super(MAE, self).__init__(img_size=img_size,
                                  patch_size=patch_size,
                                  in_channels=in_channels,
                                  embed_dims=embed_dims,
                                  num_layers=num_layers,
                                  num_heads=num_heads,
                                  mlp_ratio=mlp_ratio,
                                  out_indices=out_indices,
                                  qv_bias=False,
                                  attn_drop_rate=attn_drop_rate,
                                  drop_path_rate=drop_path_rate,
                                  norm_cfg=norm_cfg,
                                  act_cfg=act_cfg,
                                  patch_norm=patch_norm,
                                  final_norm=final_norm,
                                  num_fcs=num_fcs,
                                  norm_eval=norm_eval,
                                  init_values=init_values)

        self.cls_token = jt.zeros((1, 1, embed_dims))

        self.num_patches = self.patch_shape[0] * self.patch_shape[1]
        self.pos_embed = jt.zeros((1, self.num_patches + 1, embed_dims))

    def _build_layers(self):
        dpr = [
            x.item()
            for x in jt.linspace(0, self.drop_path_rate, self.num_layers)
        ]
        self.layers = nn.ModuleList()
        for i in range(self.num_layers):
            self.layers.append(
                MAETransformerEncoderLayer(
                    embed_dims=self.embed_dims,
                    num_heads=self.num_heads,
                    feedexecute_channels=self.mlp_ratio * self.embed_dims,
                    attn_drop_rate=self.attn_drop_rate,
                    drop_path_rate=dpr[i],
                    num_fcs=self.num_fcs,
                    bias=True,
                    act_cfg=self.act_cfg,
                    norm_cfg=self.norm_cfg,
                    window_size=self.patch_shape,
                    init_values=self.init_values))

    def fix_init_weight(self):
        """Rescale the initialization according to layer id.

        This function is copied from  https://github.com/microsoft/unilm/blob/master/beit/modeling_pretrain.py. # noqa: E501
        Copyright (c) Microsoft Corporation
        Licensed under the MIT License
        """

        def rescale(param, layer_id):
            param /= math.sqrt(2.0 * layer_id)

        for layer_id, layer in enumerate(self.layers):
            rescale(layer.attn.proj.weight.data, layer_id + 1)
            rescale(layer.ffn.layers[1].weight.data, layer_id + 1)

    def init_weights(self, pretrained=None):

        def _init_weights(m):
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

        self.apply(_init_weights)
        self.fix_init_weight()

        if pretrained is None:
            trunc_normal_(self.cls_token, std=.02)
            for n, m in self.named_modules():
                if isinstance(m, nn.Linear):
                    trunc_normal_(m.weight, std=.02)
                    if m.bias is not None:
                        if 'ffn' in n:
                            nn.init.normal_(m.bias, mean=0., std=1e-6)
                        else:
                            nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Conv2d):
                    kaiming_init(m, mode='fan_in', bias=0.)
                elif isinstance(m, (_BatchNorm, nn.GroupNorm, nn.LayerNorm)):
                    constant_init(m, val=1.0, bias=0.)
        else:
            checkpoint = jt.load(pretrained)
            state_dict = self.resize_rel_pos_embed(checkpoint)
            state_dict = self.resize_abs_pos_embed(state_dict)
            self.load_parameters(state_dict)

    def resize_abs_pos_embed(self, state_dict):
        if 'pos_embed' in state_dict:
            pos_embed_checkpoint = state_dict['pos_embed']
            embedding_size = pos_embed_checkpoint.shape[-1]
            num_extra_tokens = self.pos_embed.shape[-2] - self.num_patches
            # height (== width) for the checkpoint position embedding
            orig_size = int(
                (pos_embed_checkpoint.shape[-2] - num_extra_tokens)**0.5)
            # height (== width) for the new position embedding
            new_size = int(self.num_patches**0.5)
            # class_token and dist_token are kept unchanged
            if orig_size != new_size:
                extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
                # only the position tokens are interpolated
                pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
                if not isinstance(pos_tokens, jt.Var):
                    pos_tokens = jt.Var(pos_tokens)
                pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size,
                                                embedding_size).permute(
                                                    0, 3, 1, 2)
                pos_tokens = nn.interpolate(
                    pos_tokens,
                    size=(new_size, new_size),
                    mode='bicubic',
                    align_corners=False)
                pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
                new_pos_embed = jt.concat((extra_tokens, pos_tokens), dim=1)
                state_dict['pos_embed'] = new_pos_embed
        return state_dict

    def execute(self, inputs):
        B = inputs.shape[0]

        x, hw_shape = self.patch_embed(inputs)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = jt.concat((cls_tokens, x), dim=1)
        x = x + self.pos_embed

        outs = []
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i == len(self.layers) - 1:
                if self.final_norm:
                    x = self.norm1(x)
            if i in self.out_indices:
                out = x[:, 1:]
                B, _, C = out.shape
                out = out.reshape(B, hw_shape[0], hw_shape[1],
                                  C).permute(0, 3, 1, 2).contiguous()
                outs.append(out)

        return tuple(outs)

python/jittorseg/models/backbones/mae.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from functools import partial
from jittorseg.utils.helpers import to_2tuple
from jittorseg.utils.registry import BACKBONES
import math
from jittorseg.utils.weight_init import trunc_normal_init, normal_init, constant_init


class Mlp(nn.Module):
    def __init__(self,
                 in_features,
                 hidden_features=None,
                 out_features=None,
                 act_layer=nn.GELU,
                 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)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def execute(self, x, H, W):
        x = self.fc1(x)
        x = self.dwconv(x, H, W)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self,
                 dim,
                 num_heads=8,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop=0.,
                 proj_drop=0.,
                 sr_ratio=1):
        super().__init__()
        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."

        self.dim = dim
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.sr_ratio = sr_ratio
        if sr_ratio > 1:
            self.sr = nn.Conv2d(dim,
                                dim,
                                kernel_size=sr_ratio,
                                stride=sr_ratio)
            self.norm = nn.LayerNorm(dim)

    def execute(self, x, H, W):
        B, N, C = x.shape
        q = self.q(x).reshape(B, N, self.num_heads,
                              C // self.num_heads).permute(0, 2, 1, 3)

        if self.sr_ratio > 1:
            x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
            x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
            x_ = self.norm(x_)
            kv = self.kv(x_).reshape(B, -1, 2, self.num_heads,
                                     C // self.num_heads).permute(
                                         2, 0, 3, 1, 4)
        else:
            kv = self.kv(x).reshape(B, -1, 2, self.num_heads,
                                    C // self.num_heads).permute(
                                        2, 0, 3, 1, 4)
        k, v = kv[0], kv[1]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = nn.softmax(attn, dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)

        return x


class Block(nn.Module):
    def __init__(self,
                 dim,
                 num_heads,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop=0.,
                 attn_drop=0.,
                 drop_path=0.,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm,
                 sr_ratio=1):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim,
                              num_heads=num_heads,
                              qkv_bias=qkv_bias,
                              qk_scale=qk_scale,
                              attn_drop=attn_drop,
                              proj_drop=drop,
                              sr_ratio=sr_ratio)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = nn.DropPath(
            drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim,
                       hidden_features=mlp_hidden_dim,
                       act_layer=act_layer,
                       drop=drop)

    def execute(self, x, H, W):
        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))

        return x


class OverlapPatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self,
                 img_size=224,
                 patch_size=7,
                 stride=4,
                 in_chans=3,
                 embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)

        self.img_size = img_size
        self.patch_size = patch_size
        self.H, self.W = img_size[0] // patch_size[0], img_size[
            1] // patch_size[1]
        self.num_patches = self.H * self.W
        self.proj = nn.Conv2d(in_chans,
                              embed_dim,
                              kernel_size=patch_size,
                              stride=stride,
                              padding=(patch_size[0] // 2, patch_size[1] // 2))
        self.norm = nn.LayerNorm(embed_dim)

    def execute(self, x):
        x = self.proj(x)
        _, _, H, W = x.shape
        x = x.flatten(2).transpose(1, 2)
        x = self.norm(x)

        return x, H, W


class MixVisionTransformer(nn.Module):
    def __init__(self,
                 img_size=224,
                 patch_size=16,
                 in_chans=3,
                 num_classes=1000,
                 embed_dims=[64, 128, 256, 512],
                 num_heads=[1, 2, 4, 8],
                 mlp_ratios=[4, 4, 4, 4],
                 qkv_bias=False,
                 qk_scale=None,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 norm_layer=nn.LayerNorm,
                 depths=[3, 4, 6, 3],
                 sr_ratios=[8, 4, 2, 1]):
        super().__init__()
        self.num_classes = num_classes
        self.depths = depths

        # patch_embed
        self.patch_embed1 = OverlapPatchEmbed(img_size=img_size,
                                              patch_size=7,
                                              stride=4,
                                              in_chans=in_chans,
                                              embed_dim=embed_dims[0])
        self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4,
                                              patch_size=3,
                                              stride=2,
                                              in_chans=embed_dims[0],
                                              embed_dim=embed_dims[1])
        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8,
                                              patch_size=3,
                                              stride=2,
                                              in_chans=embed_dims[1],
                                              embed_dim=embed_dims[2])
        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16,
                                              patch_size=3,
                                              stride=2,
                                              in_chans=embed_dims[2],
                                              embed_dim=embed_dims[3])

        # transformer encoder
        dpr = [x.item() for x in jt.linspace(0, drop_path_rate, sum(depths))
               ]  # stochastic depth decay rule
        cur = 0
        self.block1 = nn.ModuleList([
            Block(dim=embed_dims[0],
                  num_heads=num_heads[0],
                  mlp_ratio=mlp_ratios[0],
                  qkv_bias=qkv_bias,
                  qk_scale=qk_scale,
                  drop=drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dpr[cur + i],
                  norm_layer=norm_layer,
                  sr_ratio=sr_ratios[0]) for i in range(depths[0])
        ])
        self.norm1 = norm_layer(embed_dims[0])

        cur += depths[0]
        self.block2 = nn.ModuleList([
            Block(dim=embed_dims[1],
                  num_heads=num_heads[1],
                  mlp_ratio=mlp_ratios[1],
                  qkv_bias=qkv_bias,
                  qk_scale=qk_scale,
                  drop=drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dpr[cur + i],
                  norm_layer=norm_layer,
                  sr_ratio=sr_ratios[1]) for i in range(depths[1])
        ])
        self.norm2 = norm_layer(embed_dims[1])

        cur += depths[1]
        self.block3 = nn.ModuleList([
            Block(dim=embed_dims[2],
                  num_heads=num_heads[2],
                  mlp_ratio=mlp_ratios[2],
                  qkv_bias=qkv_bias,
                  qk_scale=qk_scale,
                  drop=drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dpr[cur + i],
                  norm_layer=norm_layer,
                  sr_ratio=sr_ratios[2]) for i in range(depths[2])
        ])
        self.norm3 = norm_layer(embed_dims[2])

        cur += depths[2]
        self.block4 = nn.ModuleList([
            Block(dim=embed_dims[3],
                  num_heads=num_heads[3],
                  mlp_ratio=mlp_ratios[3],
                  qkv_bias=qkv_bias,
                  qk_scale=qk_scale,
                  drop=drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dpr[cur + i],
                  norm_layer=norm_layer,
                  sr_ratio=sr_ratios[3]) for i in range(depths[3])
        ])
        self.norm4 = norm_layer(embed_dims[3])

    def init_weights(self, pretrained=None):
        if pretrained is None:
            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, val=1.0, bias=0.)
                elif isinstance(m, nn.Conv2d):
                    fan_out = m.kernel_size[0] * m.kernel_size[
                        1] * m.out_channels
                    fan_out //= m.groups
                    normal_init(m,
                                mean=0,
                                std=math.sqrt(2.0 / fan_out),
                                bias=0)
        elif isinstance(pretrained, str):
            self.load_parameters(jt.load(pretrained))

    def execute_features(self, x):
        B = x.shape[0]
        outs = []

        # stage 1
        x, H, W = self.patch_embed1(x)
        for i, blk in enumerate(self.block1):
            x = blk(x, H, W)
        x = self.norm1(x)
        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2)
        outs.append(x)

        # stage 2
        x, H, W = self.patch_embed2(x)
        for i, blk in enumerate(self.block2):
            x = blk(x, H, W)
        x = self.norm2(x)
        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2)
        outs.append(x)

        # stage 3
        x, H, W = self.patch_embed3(x)
        for i, blk in enumerate(self.block3):
            x = blk(x, H, W)
        x = self.norm3(x)
        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2)
        outs.append(x)

        # stage 4
        x, H, W = self.patch_embed4(x)
        for i, blk in enumerate(self.block4):
            x = blk(x, H, W)
        x = self.norm4(x)
        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2)
        outs.append(x)

        return outs

    def execute(self, x):
        x = self.execute_features(x)

        return x


class DWConv(nn.Module):
    def __init__(self, dim=768):
        super(DWConv, self).__init__()
        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

    def execute(self, x, H, W):
        B, N, C = x.shape
        x = x.transpose(1, 2).view(B, C, H, W)
        x = self.dwconv(x)
        x = x.flatten(2).transpose(1, 2)

        return x


@BACKBONES.register_module()
class mit_b0(MixVisionTransformer):
    def __init__(self, **kwargs):
        super(mit_b0, self).__init__(patch_size=4,
                                     embed_dims=[32, 64, 160, 256],
                                     num_heads=[1, 2, 5, 8],
                                     mlp_ratios=[4, 4, 4, 4],
                                     qkv_bias=True,
                                     norm_layer=partial(nn.LayerNorm,
                                                        eps=1e-6),
                                     depths=[2, 2, 2, 2],
                                     sr_ratios=[8, 4, 2, 1],
                                     drop_rate=0.0,
                                     drop_path_rate=0.1)


@BACKBONES.register_module()
class mit_b1(MixVisionTransformer):
    def __init__(self, **kwargs):
        super(mit_b1, self).__init__(patch_size=4,
                                     embed_dims=[64, 128, 320, 512],
                                     num_heads=[1, 2, 5, 8],
                                     mlp_ratios=[4, 4, 4, 4],
                                     qkv_bias=True,
                                     norm_layer=partial(nn.LayerNorm,
                                                        eps=1e-6),
                                     depths=[2, 2, 2, 2],
                                     sr_ratios=[8, 4, 2, 1],
                                     drop_rate=0.0,
                                     drop_path_rate=0.1)


@BACKBONES.register_module()
class mit_b2(MixVisionTransformer):
    def __init__(self, **kwargs):
        super(mit_b2, self).__init__(patch_size=4,
                                     embed_dims=[64, 128, 320, 512],
                                     num_heads=[1, 2, 5, 8],
                                     mlp_ratios=[4, 4, 4, 4],
                                     qkv_bias=True,
                                     norm_layer=partial(nn.LayerNorm,
                                                        eps=1e-6),
                                     depths=[3, 4, 6, 3],
                                     sr_ratios=[8, 4, 2, 1],
                                     drop_rate=0.0,
                                     drop_path_rate=0.1)


@BACKBONES.register_module()
class mit_b3(MixVisionTransformer):
    def __init__(self, **kwargs):
        super(mit_b3, self).__init__(patch_size=4,
                                     embed_dims=[64, 128, 320, 512],
                                     num_heads=[1, 2, 5, 8],
                                     mlp_ratios=[4, 4, 4, 4],
                                     qkv_bias=True,
                                     norm_layer=partial(nn.LayerNorm,
                                                        eps=1e-6),
                                     depths=[3, 4, 18, 3],
                                     sr_ratios=[8, 4, 2, 1],
                                     drop_rate=0.0,
                                     drop_path_rate=0.1)


@BACKBONES.register_module()
class mit_b4(MixVisionTransformer):
    def __init__(self, **kwargs):
        super(mit_b4, self).__init__(patch_size=4,
                                     embed_dims=[64, 128, 320, 512],
                                     num_heads=[1, 2, 5, 8],
                                     mlp_ratios=[4, 4, 4, 4],
                                     qkv_bias=True,
                                     norm_layer=partial(nn.LayerNorm,
                                                        eps=1e-6),
                                     depths=[3, 8, 27, 3],
                                     sr_ratios=[8, 4, 2, 1],
                                     drop_rate=0.0,
                                     drop_path_rate=0.1)


@BACKBONES.register_module()
class mit_b5(MixVisionTransformer):
    def __init__(self, **kwargs):
        super(mit_b5, self).__init__(patch_size=4,
                                     embed_dims=[64, 128, 320, 512],
                                     num_heads=[1, 2, 5, 8],
                                     mlp_ratios=[4, 4, 4, 4],
                                     qkv_bias=True,
                                     norm_layer=partial(nn.LayerNorm,
                                                        eps=1e-6),
                                     depths=[3, 6, 40, 3],
                                     sr_ratios=[8, 4, 2, 1],
                                     drop_rate=0.0,
                                     drop_path_rate=0.1)

python/jittorseg/models/backbones/mix_transformer.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule
from jittor.nn import BatchNorm as _BatchNorm

from jittorseg.utils.registry import BACKBONES
from jittorseg.models.basic_operators.inverted_residual import InvertedResidual
from jittorseg.utils.helpers import make_divisible


@BACKBONES.register_module()
class MobileNetV2(nn.Module):
    # Parameters to build layers. 3 parameters are needed to construct a
    # layer, from left to right: expand_ratio, channel, num_blocks.
    arch_settings = [[1, 16, 1], [6, 24, 2], [6, 32, 3], [6, 64, 4],
                     [6, 96, 3], [6, 160, 3], [6, 320, 1]]

    def __init__(self,
                 widen_factor=1.,
                 strides=(1, 2, 2, 2, 1, 2, 1),
                 dilations=(1, 1, 1, 1, 1, 1, 1),
                 out_indices=(1, 2, 4, 6),
                 frozen_stages=-1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU6'),
                 norm_eval=False):
        super(MobileNetV2, self).__init__()

        self.widen_factor = widen_factor
        self.strides = strides
        self.dilations = dilations
        assert len(strides) == len(dilations) == len(self.arch_settings)
        self.out_indices = out_indices
        for index in out_indices:
            if index not in range(0, 7):
                raise ValueError('the item in out_indices must in '
                                 f'range(0, 7). But received {index}')

        if frozen_stages not in range(-1, 7):
            raise ValueError('frozen_stages must be in range(-1, 7). '
                             f'But received {frozen_stages}')
        self.out_indices = out_indices
        self.frozen_stages = frozen_stages
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.norm_eval = norm_eval

        self.in_channels = make_divisible(32 * widen_factor, 8)

        self.conv1 = ConvModule(in_channels=3,
                                out_channels=self.in_channels,
                                kernel_size=3,
                                stride=2,
                                padding=1,
                                conv_cfg=self.conv_cfg,
                                norm_cfg=self.norm_cfg,
                                act_cfg=self.act_cfg)

        self.layers = []

        for i, layer_cfg in enumerate(self.arch_settings):
            expand_ratio, channel, num_blocks = layer_cfg
            stride = self.strides[i]
            dilation = self.dilations[i]
            out_channels = make_divisible(channel * widen_factor, 8)
            inverted_res_layer = self.make_layer(out_channels=out_channels,
                                                 num_blocks=num_blocks,
                                                 stride=stride,
                                                 dilation=dilation,
                                                 expand_ratio=expand_ratio)
            layer_name = f'layer{i + 1}'
            setattr(self, layer_name, inverted_res_layer)
            self.layers.append(layer_name)

    def init_weights(self, pretrained=None):
        if isinstance(pretrained, str):
            self.load_parameters(jt.load(pretrained))

    def make_layer(self, out_channels, num_blocks, stride, dilation,
                   expand_ratio):
        layers = []
        for i in range(num_blocks):
            layers.append(
                InvertedResidual(self.in_channels,
                                 out_channels,
                                 stride if i == 0 else 1,
                                 expand_ratio=expand_ratio,
                                 dilation=dilation if i == 0 else 1,
                                 conv_cfg=self.conv_cfg,
                                 norm_cfg=self.norm_cfg,
                                 act_cfg=self.act_cfg))
            self.in_channels = out_channels

        return nn.Sequential(*layers)

    def execute(self, x):
        x = self.conv1(x)

        outs = []
        for i, layer_name in enumerate(self.layers):
            layer = getattr(self, layer_name)
            x = layer(x)
            if i in self.out_indices:
                outs.append(x)

        if len(outs) == 1:
            return outs[0]
        else:
            return tuple(outs)

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            for param in self.conv1.parameters():
                param.requires_grad = False
        for i in range(1, self.frozen_stages + 1):
            layer = getattr(self, f'layer{i}')
            layer.eval()
            for param in layer.parameters():
                param.requires_grad = False

    def train(self, mode=True):
        super(MobileNetV2, self).train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()

python/jittorseg/models/backbones/mobilenet_v2.py
JittorCV/jittorseg
import math
import jittor as jt
from jittor import nn
from jittorseg.utils.helpers import to_2tuple
from jittorseg.utils.registry import BACKBONES
from jittorseg.utils.weight_init import trunc_normal_init, normal_init, constant_init
from ..basic_operators import build_norm_layer


class Mlp(nn.Module):

    def __init__(self,
                 in_features,
                 hidden_features=None,
                 out_features=None,
                 act_layer=nn.GELU,
                 drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
        self.drop = nn.Dropout(drop)

    def execute(self, x):
        x = self.fc1(x)

        x = self.dwconv(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)

        return x


class StemConv(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 norm_cfg=dict(type='BN')):
        super(StemConv, self).__init__()

        self.proj = nn.Sequential(
            nn.Conv2d(in_channels,
                      out_channels // 2,
                      kernel_size=(3, 3),
                      stride=(2, 2),
                      padding=(1, 1)),
            build_norm_layer(norm_cfg, out_channels // 2)[1],
            nn.GELU(),
            nn.Conv2d(out_channels // 2,
                      out_channels,
                      kernel_size=(3, 3),
                      stride=(2, 2),
                      padding=(1, 1)),
            build_norm_layer(norm_cfg, out_channels)[1],
        )

    def execute(self, x):
        x = self.proj(x)
        _, _, H, W = x.size()
        x = x.flatten(2).transpose(1, 2)
        return x, H, W


class AttentionModule(nn.Module):

    def __init__(self, dim):
        super().__init__()
        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
        self.conv0_1 = nn.Conv2d(dim, dim, (1, 7), padding=(0, 3), groups=dim)
        self.conv0_2 = nn.Conv2d(dim, dim, (7, 1), padding=(3, 0), groups=dim)

        self.conv1_1 = nn.Conv2d(dim, dim, (1, 11), padding=(0, 5), groups=dim)
        self.conv1_2 = nn.Conv2d(dim, dim, (11, 1), padding=(5, 0), groups=dim)

        self.conv2_1 = nn.Conv2d(dim,
                                 dim, (1, 21),
                                 padding=(0, 10),
                                 groups=dim)
        self.conv2_2 = nn.Conv2d(dim,
                                 dim, (21, 1),
                                 padding=(10, 0),
                                 groups=dim)
        self.conv3 = nn.Conv2d(dim, dim, 1)

    def execute(self, x):
        u = x.clone()
        attn = self.conv0(x)

        attn_0 = self.conv0_1(attn)
        attn_0 = self.conv0_2(attn_0)

        attn_1 = self.conv1_1(attn)
        attn_1 = self.conv1_2(attn_1)

        attn_2 = self.conv2_1(attn)
        attn_2 = self.conv2_2(attn_2)
        attn = attn + attn_0 + attn_1 + attn_2

        attn = self.conv3(attn)

        return attn * u


class SpatialAttention(nn.Module):

    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        self.proj_1 = nn.Conv2d(d_model, d_model, 1)
        self.activation = nn.GELU()
        self.spatial_gating_unit = AttentionModule(d_model)
        self.proj_2 = nn.Conv2d(d_model, d_model, 1)

    def execute(self, x):
        shorcut = x.clone()
        x = self.proj_1(x)
        x = self.activation(x)
        x = self.spatial_gating_unit(x)
        x = self.proj_2(x)
        x = x + shorcut
        return x


class Block(nn.Module):

    def __init__(self,
                 dim,
                 mlp_ratio=4.,
                 drop=0.,
                 drop_path=0.,
                 act_layer=nn.GELU,
                 norm_cfg=dict(type='BN')):
        super().__init__()
        self.norm1 = build_norm_layer(norm_cfg, dim)[1]
        self.attn = SpatialAttention(dim)
        self.drop_path = nn.DropPath(
            drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = build_norm_layer(norm_cfg, dim)[1]
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim,
                       hidden_features=mlp_hidden_dim,
                       act_layer=act_layer,
                       drop=drop)
        layer_scale_init_value = 1e-2
        self.layer_scale_1 = layer_scale_init_value * jt.ones((dim))
        self.layer_scale_2 = layer_scale_init_value * jt.ones((dim))

    def execute(self, x, H, W):
        B, N, C = x.shape
        x = x.permute(0, 2, 1).view(B, C, H, W)
        x = x + self.drop_path(
            self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) *
            self.attn(self.norm1(x)))
        x = x + self.drop_path(
            self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) *
            self.mlp(self.norm2(x)))
        x = x.view(B, C, N).permute(0, 2, 1)
        return x


class OverlapPatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self,
                 patch_size=7,
                 stride=4,
                 in_chans=3,
                 embed_dim=768,
                 norm_cfg=dict(type='BN')):
        super().__init__()
        patch_size = to_2tuple(patch_size)

        self.proj = nn.Conv2d(in_chans,
                              embed_dim,
                              kernel_size=patch_size,
                              stride=stride,
                              padding=(patch_size[0] // 2, patch_size[1] // 2))
        self.norm = build_norm_layer(norm_cfg, embed_dim)[1]

    def execute(self, x):
        x = self.proj(x)
        _, _, H, W = x.shape
        x = self.norm(x)

        x = x.flatten(2).transpose(1, 2)

        return x, H, W


@BACKBONES.register_module()
class MSCAN(nn.Module):

    def __init__(self,
                 in_chans=3,
                 embed_dims=[64, 128, 256, 512],
                 mlp_ratios=[4, 4, 4, 4],
                 drop_rate=0.,
                 drop_path_rate=0.,
                 depths=[3, 4, 6, 3],
                 num_stages=4,
                 norm_cfg=dict(type='BN')):
        super(MSCAN, self).__init__()

        self.depths = depths
        self.num_stages = num_stages

        dpr = [x.item() for x in jt.linspace(0, drop_path_rate, sum(depths))
               ]  # stochastic depth decay rule
        cur = 0

        for i in range(num_stages):
            if i == 0:
                patch_embed = StemConv(3, embed_dims[0], norm_cfg=norm_cfg)
            else:
                patch_embed = OverlapPatchEmbed(
                    patch_size=7 if i == 0 else 3,
                    stride=4 if i == 0 else 2,
                    in_chans=in_chans if i == 0 else embed_dims[i - 1],
                    embed_dim=embed_dims[i],
                    norm_cfg=norm_cfg)

            block = nn.ModuleList([
                Block(dim=embed_dims[i],
                      mlp_ratio=mlp_ratios[i],
                      drop=drop_rate,
                      drop_path=dpr[cur + j],
                      norm_cfg=norm_cfg) for j in range(depths[i])
            ])
            norm = nn.LayerNorm(embed_dims[i])
            cur += depths[i]

            setattr(self, f"patch_embed{i + 1}", patch_embed)
            setattr(self, f"block{i + 1}", block)
            setattr(self, f"norm{i + 1}", norm)

    def init_weights(self, pretrained=None):
        if pretrained is None:
            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, val=1.0, bias=0.)
                elif isinstance(m, nn.Conv2d):
                    fan_out = m.kernel_size[0] * m.kernel_size[
                        1] * m.out_channels
                    fan_out //= m.groups
                    normal_init(m,
                                mean=0,
                                std=math.sqrt(2.0 / fan_out),
                                bias=0)
        elif isinstance(pretrained, str):
            self.load_parameters(jt.load(pretrained))

    def execute(self, x):
        B = x.shape[0]
        outs = []

        for i in range(self.num_stages):
            patch_embed = getattr(self, f"patch_embed{i + 1}")
            block = getattr(self, f"block{i + 1}")
            norm = getattr(self, f"norm{i + 1}")
            x, H, W = patch_embed(x)
            for blk in block:
                x = blk(x, H, W)
            x = norm(x)
            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2)
            outs.append(x)

        return outs


class DWConv(nn.Module):

    def __init__(self, dim=768):
        super(DWConv, self).__init__()
        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

    def execute(self, x):
        x = self.dwconv(x)
        return x

python/jittorseg/models/backbones/mscan.py
JittorCV/jittorseg
import math
import jittor as jt
from jittor import nn
from ..basic_operators import build_conv_layer, build_norm_layer

from jittorseg.utils.registry import BACKBONES
from .resnet import ResLayer
from .resnet import Bottleneck as _Bottleneck
from .resnet import ResNetV1d


class RSoftmax(nn.Module):

    def __init__(self, radix, groups):
        super().__init__()
        self.radix = radix
        self.groups = groups

    def execute(self, x):
        batch = x.size(0)
        if self.radix > 1:
            x = x.view(batch, self.groups, self.radix, -1).transpose(1, 2)
            x = nn.softmax(x, dim=1)
            x = x.reshape(batch, -1)
        else:
            x = jt.sigmoid(x)
        return x


class SplitAttentionConv2d(nn.Module):

    def __init__(self,
                 in_channels,
                 channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 radix=2,
                 reduction_factor=4,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 dcn=None):
        super(SplitAttentionConv2d, self).__init__()
        inter_channels = max(in_channels * radix // reduction_factor, 32)
        self.radix = radix
        self.groups = groups
        self.channels = channels
        self.with_dcn = dcn is not None
        self.dcn = dcn
        self.conv = build_conv_layer(conv_cfg,
                                     in_channels,
                                     channels * radix,
                                     kernel_size,
                                     stride=stride,
                                     padding=padding,
                                     dilation=dilation,
                                     groups=groups * radix,
                                     bias=False)
        self.norm0_name, norm0 = build_norm_layer(norm_cfg,
                                                  channels * radix,
                                                  postfix=0)
        setattr(self, self.norm0_name, norm0)
        self.relu = nn.ReLU()
        self.fc1 = build_conv_layer(None,
                                    channels,
                                    inter_channels,
                                    1,
                                    groups=self.groups)
        self.norm1_name, norm1 = build_norm_layer(norm_cfg,
                                                  inter_channels,
                                                  postfix=1)
        setattr(self, self.norm1_name, norm1)
        self.fc2 = build_conv_layer(None,
                                    inter_channels,
                                    channels * radix,
                                    1,
                                    groups=self.groups)
        self.rsoftmax = RSoftmax(radix, groups)
        self.adaptive_avg_pool2d = nn.AdaptiveAvgPool2d(1)

    @property
    def norm0(self):
        """nn.Module: the normalization layer named "norm0" """
        return getattr(self, self.norm0_name)

    @property
    def norm1(self):
        """nn.Module: the normalization layer named "norm1" """
        return getattr(self, self.norm1_name)

    def execute(self, x):
        x = self.conv(x)
        x = self.norm0(x)
        x = self.relu(x)

        batch, rchannel = x.shape[:2]
        batch = x.size(0)
        if self.radix > 1:
            splits = x.view(batch, self.radix, -1, *x.shape[2:])
            gap = splits.sum(dim=1)
        else:
            gap = x
        gap = self.adaptive_avg_pool2d(gap)
        gap = self.fc1(gap)

        gap = self.norm1(gap)
        gap = self.relu(gap)

        atten = self.fc2(gap)
        atten = self.rsoftmax(atten).view(batch, -1, 1, 1)

        if self.radix > 1:
            attens = atten.view(batch, self.radix, -1, *atten.shape[2:])
            out = jt.sum(attens * splits, dim=1)
        else:
            out = atten * x
        return out


class Bottleneck(_Bottleneck):
    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 groups=1,
                 base_width=4,
                 base_channels=64,
                 radix=2,
                 reduction_factor=4,
                 avg_down_stride=True,
                 **kwargs):
        """Bottleneck block for ResNeSt."""
        super(Bottleneck, self).__init__(inplanes, planes, **kwargs)

        if groups == 1:
            width = self.planes
        else:
            width = math.floor(self.planes *
                               (base_width / base_channels)) * groups

        self.avg_down_stride = avg_down_stride and self.conv2_stride > 1

        self.norm1_name, norm1 = build_norm_layer(self.norm_cfg,
                                                  width,
                                                  postfix=1)
        self.norm3_name, norm3 = build_norm_layer(self.norm_cfg,
                                                  self.planes * self.expansion,
                                                  postfix=3)

        self.conv1 = build_conv_layer(self.conv_cfg,
                                      self.inplanes,
                                      width,
                                      kernel_size=1,
                                      stride=self.conv1_stride,
                                      bias=False)
        setattr(self, self.norm1_name, norm1)
        self.with_modulated_dcn = False
        self.conv2 = SplitAttentionConv2d(
            width,
            width,
            kernel_size=3,
            stride=1 if self.avg_down_stride else self.conv2_stride,
            padding=self.dilation,
            dilation=self.dilation,
            groups=groups,
            radix=radix,
            reduction_factor=reduction_factor,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            dcn=self.dcn)
        delattr(self, self.norm2_name)

        if self.avg_down_stride:
            self.avd_layer = nn.AvgPool2d(3, self.conv2_stride, padding=1)

        self.conv3 = build_conv_layer(self.conv_cfg,
                                      width,
                                      self.planes * self.expansion,
                                      kernel_size=1,
                                      bias=False)
        setattr(self, self.norm3_name, norm3)

    def execute(self, x):

        def _inner_execute(x):
            identity = x

            out = self.conv1(x)
            out = self.norm1(out)
            out = self.relu(out)

            out = self.conv2(out)

            if self.avg_down_stride:
                out = self.avd_layer(out)

            out = self.conv3(out)
            out = self.norm3(out)

            if self.downsample is not None:
                identity = self.downsample(x)

            out += identity

            return out

        out = _inner_execute(x)
        out = self.relu(out)

        return out


@BACKBONES.register_module()
class ResNeSt(ResNetV1d):
    """ResNeSt backbone.

    This backbone is the implementation of `ResNeSt:
    Split-Attention Networks <https://arxiv.org/abs/2004.08955>`_.

    Args:
        groups (int): Number of groups of Bottleneck. Default: 1
        base_width (int): Base width of Bottleneck. Default: 4
        radix (int): Radix of SpltAtConv2d. Default: 2
        reduction_factor (int): Reduction factor of inter_channels in
            SplitAttentionConv2d. Default: 4.
        avg_down_stride (bool): Whether to use average pool for stride in
            Bottleneck. Default: True.
        kwargs (dict): Keyword arguments for ResNet.
    """

    arch_settings = {
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3)),
        200: (Bottleneck, (3, 24, 36, 3))
    }

    def __init__(self,
                 groups=1,
                 base_width=4,
                 radix=2,
                 reduction_factor=4,
                 avg_down_stride=True,
                 **kwargs):
        self.groups = groups
        self.base_width = base_width
        self.radix = radix
        self.reduction_factor = reduction_factor
        self.avg_down_stride = avg_down_stride
        super(ResNeSt, self).__init__(**kwargs)

    def make_res_layer(self, **kwargs):
        """Pack all blocks in a stage into a ``ResLayer``."""
        return ResLayer(groups=self.groups,
                        base_width=self.base_width,
                        base_channels=self.base_channels,
                        radix=self.radix,
                        reduction_factor=self.reduction_factor,
                        avg_down_stride=self.avg_down_stride,
                        **kwargs)

python/jittorseg/models/backbones/resnest.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from jittorseg.utils.registry import BACKBONES
from ..basic_operators import build_conv_layer, build_norm_layer


class ResLayer(nn.Sequential):
    def __init__(self,
                 block,
                 inplanes,
                 planes,
                 num_blocks,
                 stride=1,
                 dilation=1,
                 avg_down=False,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 multi_grid=None,
                 contract_dilation=False,
                 **kwargs):
        self.block = block

        downsample = None
        if stride != 1 or inplanes != planes * block.expansion:
            downsample = []
            conv_stride = stride
            if avg_down:
                conv_stride = 1
                downsample.append(
                    nn.AvgPool2d(kernel_size=stride,
                                 stride=stride,
                                 ceil_mode=True,
                                 count_include_pad=False))
            downsample.extend([
                build_conv_layer(conv_cfg,
                                 inplanes,
                                 planes * block.expansion,
                                 kernel_size=1,
                                 stride=conv_stride,
                                 bias=False),
                build_norm_layer(norm_cfg, planes * block.expansion)[1]
            ])
            downsample = nn.Sequential(*downsample)

        layers = []
        if multi_grid is None:
            if dilation > 1 and contract_dilation:
                first_dilation = dilation // 2
            else:
                first_dilation = dilation
        else:
            first_dilation = multi_grid[0]
        layers.append(
            block(inplanes=inplanes,
                  planes=planes,
                  stride=stride,
                  dilation=first_dilation,
                  downsample=downsample,
                  **kwargs))
        inplanes = planes * block.expansion
        for i in range(1, num_blocks):
            layers.append(
                block(
                    inplanes=inplanes,
                    planes=planes,
                    stride=1,
                    dilation=dilation if multi_grid is None else multi_grid[i],
                    **kwargs))
        super(ResLayer, self).__init__(*layers)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 dilation=1,
                 downsample=None,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 dcn=None):
        super(BasicBlock, self).__init__()
        assert dcn is None, 'Not implemented yet.'

        self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1)
        self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2)

        self.conv1 = build_conv_layer(conv_cfg,
                                      inplanes,
                                      planes,
                                      3,
                                      stride=stride,
                                      padding=dilation,
                                      dilation=dilation,
                                      bias=False)
        setattr(self, self.norm1_name, norm1)
        self.conv2 = build_conv_layer(conv_cfg,
                                      planes,
                                      planes,
                                      3,
                                      padding=1,
                                      bias=False)
        setattr(self, self.norm2_name, norm2)

        self.relu = nn.ReLU()
        self.downsample = downsample
        self.stride = stride
        self.dilation = dilation

    @property
    def norm1(self):
        """nn.Module: normalization layer after the first convolution layer"""
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        """nn.Module: normalization layer after the second convolution layer"""
        return getattr(self, self.norm2_name)

    def execute(self, x):

        def _inner_execute(x):
            identity = x

            out = self.conv1(x)
            out = self.norm1(out)
            out = self.relu(out)

            out = self.conv2(out)
            out = self.norm2(out)

            if self.downsample is not None:
                identity = self.downsample(x)

            out += identity

            return out

        out = _inner_execute(x)

        out = self.relu(out)

        return out


class Bottleneck(nn.Module):

    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 dilation=1,
                 downsample=None,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 dcn=None):
        super(Bottleneck, self).__init__()
        assert dcn is None or isinstance(dcn, dict)

        self.inplanes = inplanes
        self.planes = planes
        self.stride = stride
        self.dilation = dilation
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.dcn = dcn
        self.with_dcn = dcn is not None

        self.conv1_stride = 1
        self.conv2_stride = stride

        self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1)
        self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2)
        self.norm3_name, norm3 = build_norm_layer(norm_cfg,
                                                  planes * self.expansion,
                                                  postfix=3)

        self.conv1 = nn.Conv2d(inplanes,
                               planes,
                               kernel_size=1,
                               stride=self.conv1_stride,
                               bias=False)
        self.conv1 = build_conv_layer(conv_cfg,
                                      inplanes,
                                      planes,
                                      kernel_size=1,
                                      stride=self.conv1_stride,
                                      bias=False)
        setattr(self, self.norm1_name, norm1)
        fallback_on_stride = False
        if self.with_dcn:
            fallback_on_stride = dcn.pop('fallback_on_stride', False)
        if not self.with_dcn or fallback_on_stride:
            self.conv2 = build_conv_layer(conv_cfg,
                                          planes,
                                          planes,
                                          kernel_size=3,
                                          stride=self.conv2_stride,
                                          padding=dilation,
                                          dilation=dilation,
                                          bias=False)
        else:
            assert self.conv_cfg is None, 'conv_cfg must be None for DCN'
            self.conv2 = build_conv_layer(dcn,
                                          planes,
                                          planes,
                                          kernel_size=3,
                                          stride=self.conv2_stride,
                                          padding=dilation,
                                          dilation=dilation,
                                          bias=False)

        setattr(self, self.norm2_name, norm2)
        self.conv3 = build_conv_layer(conv_cfg,
                                      planes,
                                      planes * self.expansion,
                                      kernel_size=1,
                                      bias=False)
        setattr(self, self.norm3_name, norm3)

        self.relu = nn.ReLU()
        self.downsample = downsample

    @property
    def norm1(self):
        """nn.Module: normalization layer after the first convolution layer"""
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        """nn.Module: normalization layer after the second convolution layer"""
        return getattr(self, self.norm2_name)

    @property
    def norm3(self):
        """nn.Module: normalization layer after the third convolution layer"""
        return getattr(self, self.norm3_name)

    def execute(self, x):

        def _inner_execute(x):
            identity = x

            out = self.conv1(x)
            out = self.norm1(out)
            out = self.relu(out)

            out = self.conv2(out)
            out = self.norm2(out)
            out = self.relu(out)

            out = self.conv3(out)
            out = self.norm3(out)

            if self.downsample is not None:
                identity = self.downsample(x)

            out += identity

            return out

        out = _inner_execute(x)

        out = self.relu(out)

        return out


@BACKBONES.register_module()
class ResNet(nn.Module):
    arch_settings = {
        18: (BasicBlock, (2, 2, 2, 2)),
        34: (BasicBlock, (3, 4, 6, 3)),
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3))
    }

    def __init__(self,
                 depth,
                 in_channels=3,
                 stem_channels=64,
                 base_channels=64,
                 num_stages=4,
                 strides=(1, 2, 2, 2),
                 dilations=(1, 1, 1, 1),
                 out_indices=(0, 1, 2, 3),
                 deep_stem=False,
                 avg_down=False,
                 frozen_stages=-1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN', requires_grad=True),
                 norm_eval=False,
                 dcn=None,
                 stage_with_dcn=(False, False, False, False),
                 multi_grid=None,
                 contract_dilation=False,
                 zero_init_residual=True):
        super(ResNet, self).__init__()
        if depth not in self.arch_settings:
            raise KeyError(f'invalid depth {depth} for resnet')
        self.depth = depth
        self.stem_channels = stem_channels
        self.base_channels = base_channels
        self.num_stages = num_stages
        assert num_stages >= 1 and num_stages <= 4
        self.strides = strides
        self.dilations = dilations
        assert len(strides) == len(dilations) == num_stages
        self.out_indices = out_indices
        assert max(out_indices) < num_stages
        self.deep_stem = deep_stem
        self.avg_down = avg_down
        self.frozen_stages = frozen_stages
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.norm_eval = norm_eval
        self.dcn = dcn
        self.stage_with_dcn = stage_with_dcn
        if dcn is not None:
            assert len(stage_with_dcn) == num_stages
        self.multi_grid = multi_grid
        self.contract_dilation = contract_dilation
        self.zero_init_residual = zero_init_residual
        self.block, stage_blocks = self.arch_settings[depth]
        self.stage_blocks = stage_blocks[:num_stages]
        self.inplanes = stem_channels

        self._make_stem_layer(in_channels, stem_channels)

        self.res_layers = []
        for i, num_blocks in enumerate(self.stage_blocks):
            stride = strides[i]
            dilation = dilations[i]
            dcn = self.dcn if self.stage_with_dcn[i] else None
            # multi grid is applied to last layer only
            stage_multi_grid = multi_grid if i == len(
                self.stage_blocks) - 1 else None
            planes = base_channels * 2**i
            res_layer = self.make_res_layer(
                block=self.block,
                inplanes=self.inplanes,
                planes=planes,
                num_blocks=num_blocks,
                stride=stride,
                dilation=dilation,
                avg_down=self.avg_down,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                dcn=dcn,
                multi_grid=stage_multi_grid,
                contract_dilation=contract_dilation)
            self.inplanes = planes * self.block.expansion
            layer_name = f'layer{i+1}'
            setattr(self, layer_name, res_layer)
            # self.add_module(layer_name, res_layer)
            self.res_layers.append(layer_name)

        self._freeze_stages()

        self.feat_dim = self.block.expansion * base_channels * 2**(
            len(self.stage_blocks) - 1)

    def make_res_layer(self, **kwargs):
        """Pack all blocks in a stage into a ``ResLayer``."""
        return ResLayer(**kwargs)

    @property
    def norm1(self):
        """nn.Module: the normalization layer named "norm1" """
        return getattr(self, self.norm1_name)

    def _make_stem_layer(self, in_channels, stem_channels):
        """Make stem layer for ResNet."""
        if self.deep_stem:
            self.stem = nn.Sequential(
                build_conv_layer(self.conv_cfg,
                                 in_channels,
                                 stem_channels // 2,
                                 kernel_size=3,
                                 stride=2,
                                 padding=1,
                                 bias=False),
                build_norm_layer(self.norm_cfg, stem_channels // 2)[1],
                nn.ReLU(),
                build_conv_layer(self.conv_cfg,
                                 stem_channels // 2,
                                 stem_channels // 2,
                                 kernel_size=3,
                                 stride=1,
                                 padding=1,
                                 bias=False),
                build_norm_layer(self.norm_cfg, stem_channels // 2)[1],
                nn.ReLU(),
                build_conv_layer(self.conv_cfg,
                                 stem_channels // 2,
                                 stem_channels,
                                 kernel_size=3,
                                 stride=1,
                                 padding=1,
                                 bias=False),
                build_norm_layer(self.norm_cfg, stem_channels)[1], nn.ReLU())
        else:
            self.conv1 = build_conv_layer(self.conv_cfg,
                                          in_channels,
                                          stem_channels,
                                          kernel_size=7,
                                          stride=2,
                                          padding=3,
                                          bias=False)
            self.norm1_name, norm1 = build_norm_layer(self.norm_cfg,
                                                      stem_channels,
                                                      postfix=1)
            setattr(self, self.norm1_name, norm1)
            self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            if self.deep_stem:
                self.stem.eval()
                for param in self.stem.parameters():
                    param.requires_grad = False
            else:
                self.norm1.eval()
                for m in [self.conv1, self.norm1]:
                    for param in m.parameters():
                        param.requires_grad = False

        for i in range(1, self.frozen_stages + 1):
            m = getattr(self, f'layer{i}')
            m.eval()
            for param in m.parameters():
                param.requires_grad = False

    def init_weights(self, pretrained=None):
        if isinstance(pretrained, str):
            self.load_parameters(jt.load(pretrained))

    def execute(self, x):
        if self.deep_stem:
            x = self.stem(x)
        else:
            x = self.conv1(x)
            x = self.norm1(x)
            x = self.relu(x)
        x = self.maxpool(x)
        outs = []
        for i, layer_name in enumerate(self.res_layers):
            res_layer = getattr(self, layer_name)
            x = res_layer(x)
            if i in self.out_indices:
                outs.append(x)
        return tuple(outs)

    def train(self, mode=True):
        super(ResNet, self).train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                # trick: eval have effect on BatchNorm only
                if isinstance(m, nn.BatchNorm2d):
                    m.eval()


@BACKBONES.register_module()
class ResNetV1c(ResNet):

    def __init__(self, **kwargs):
        super(ResNetV1c, self).__init__(deep_stem=True,
                                        avg_down=False,
                                        **kwargs)


@BACKBONES.register_module()
class ResNetV1d(ResNet):

    def __init__(self, **kwargs):
        super(ResNetV1d, self).__init__(deep_stem=True,
                                        avg_down=True,
                                        **kwargs)

python/jittorseg/models/backbones/resnet.py
JittorCV/jittorseg
import math
from jittorseg.utils.registry import BACKBONES
from .resnet import ResLayer
from ..basic_operators import build_conv_layer, build_norm_layer
from .resnet import Bottleneck as _Bottleneck
from .resnet import ResNet


class Bottleneck(_Bottleneck):

    def __init__(self,
                 inplanes,
                 planes,
                 groups=1,
                 base_width=4,
                 base_channels=64,
                 **kwargs):
        super(Bottleneck, self).__init__(inplanes, planes, **kwargs)

        if groups == 1:
            width = self.planes
        else:
            width = math.floor(self.planes *
                               (base_width / base_channels)) * groups

        self.norm1_name, norm1 = build_norm_layer(self.norm_cfg,
                                                  width,
                                                  postfix=1)
        self.norm2_name, norm2 = build_norm_layer(self.norm_cfg,
                                                  width,
                                                  postfix=2)
        self.norm3_name, norm3 = build_norm_layer(self.norm_cfg,
                                                  self.planes * self.expansion,
                                                  postfix=3)

        self.conv1 = build_conv_layer(self.conv_cfg,
                                      self.inplanes,
                                      width,
                                      kernel_size=1,
                                      stride=self.conv1_stride,
                                      bias=False)
        setattr(self, self.norm1_name, norm1)
        fallback_on_stride = False
        self.with_modulated_dcn = False
        if self.with_dcn:
            fallback_on_stride = self.dcn.pop('fallback_on_stride', False)
        if not self.with_dcn or fallback_on_stride:
            self.conv2 = build_conv_layer(self.conv_cfg,
                                          width,
                                          width,
                                          kernel_size=3,
                                          stride=self.conv2_stride,
                                          padding=self.dilation,
                                          dilation=self.dilation,
                                          groups=groups,
                                          bias=False)
        else:
            assert self.conv_cfg is None, 'conv_cfg must be None for DCN'
            self.conv2 = build_conv_layer(self.dcn,
                                          width,
                                          width,
                                          kernel_size=3,
                                          stride=self.conv2_stride,
                                          padding=self.dilation,
                                          dilation=self.dilation,
                                          groups=groups,
                                          bias=False)

        setattr(self, self.norm2_name, norm2)
        self.conv3 = build_conv_layer(self.conv_cfg,
                                      width,
                                      self.planes * self.expansion,
                                      kernel_size=1,
                                      bias=False)
        setattr(self, self.norm3_name, norm3)


@BACKBONES.register_module()
class ResNeXt(ResNet):
    arch_settings = {
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3))
    }

    def __init__(self, groups=1, base_width=4, **kwargs):
        self.groups = groups
        self.base_width = base_width
        super(ResNeXt, self).__init__(**kwargs)

    def make_res_layer(self, **kwargs):
        """Pack all blocks in a stage into a ``ResLayer``"""
        return ResLayer(groups=self.groups,
                        base_width=self.base_width,
                        base_channels=self.base_channels,
                        **kwargs)

python/jittorseg/models/backbones/resnext.py
JittorCV/jittorseg
from collections import OrderedDict
from copy import deepcopy

import jittor as jt
from jittor import nn
from jittorseg.utils.helpers import to_2tuple
from jittorseg.utils.registry import BACKBONES
from ..basic_operators import build_norm_layer, build_dropout
from jittorseg.utils.weight_init import trunc_normal_init, constant_init, trunc_normal_
from jittorseg.models.basic_operators.embed import PatchEmbed, PatchMerging, FFN


class WindowMSA(nn.Module):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 window_size,
                 qkv_bias=True,
                 qk_scale=None,
                 attn_drop_rate=0.,
                 proj_drop_rate=0.):

        super().__init__()
        self.embed_dims = embed_dims
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_embed_dims = embed_dims // num_heads
        self.scale = qk_scale or head_embed_dims**-0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = jt.zeros(
            ((2 * window_size[0] - 1) * (2 * window_size[1] - 1),
             num_heads))  # 2*Wh-1 * 2*Ww-1, nH

        # About 2x faster than original impl
        Wh, Ww = self.window_size
        rel_index_coords = self.double_step_seq(2 * Ww - 1, Wh, 1, Ww)
        rel_position_index = rel_index_coords + rel_index_coords.transpose(
            0, 1)
        rel_position_index = rel_position_index.flip(1)
        # self.register_buffer('relative_position_index', rel_position_index)
        self.relative_position_index = rel_position_index
        self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop_rate)
        self.proj = nn.Linear(embed_dims, embed_dims)
        self.proj_drop = nn.Dropout(proj_drop_rate)

        self.softmax = nn.Softmax(dim=-1)

    def init_weights(self):
        trunc_normal_init(self.relative_position_bias_table, std=.02)

    def execute(self, x, mask=None):
        """
        Args:

            x (Var): input features with shape of (num_windows*B, N, C)
            mask (Var | None, Optional): mask with shape of (num_windows,
                Wh*Ww, Wh*Ww), value should be between (-inf, 0].
        """
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads,
                                  C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))

        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index.view(-1)].view(
                self.window_size[0] * self.window_size[1],
                self.window_size[0] * self.window_size[1],
                -1)  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1)  # nH, Wh*Ww, Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B // nW, nW, self.num_heads, N,
                             N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
        attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    @staticmethod
    def double_step_seq(step1, len1, step2, len2):
        seq1 = jt.arange(0, step1 * len1, step1)
        seq2 = jt.arange(0, step2 * len2, step2)
        return (seq1[:, None] + seq2[None, :]).reshape(1, -1)


class ShiftWindowMSA(nn.Module):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 window_size,
                 shift_size=0,
                 qkv_bias=True,
                 qk_scale=None,
                 attn_drop_rate=0,
                 proj_drop_rate=0,
                 dropout_layer=dict(type='DropPath', p=0.)):
        super().__init__()

        self.window_size = window_size
        self.shift_size = shift_size
        assert 0 <= self.shift_size < self.window_size

        self.w_msa = WindowMSA(embed_dims=embed_dims,
                               num_heads=num_heads,
                               window_size=to_2tuple(window_size),
                               qkv_bias=qkv_bias,
                               qk_scale=qk_scale,
                               attn_drop_rate=attn_drop_rate,
                               proj_drop_rate=proj_drop_rate)

        self.drop = build_dropout(dropout_layer)

    def execute(self, query, hw_shape):
        B, L, C = query.shape
        H, W = hw_shape
        assert L == H * W, 'input feature has wrong size'
        query = query.view(B, H, W, C)

        # pad feature maps to multiples of window size
        pad_r = (self.window_size - W % self.window_size) % self.window_size
        pad_b = (self.window_size - H % self.window_size) % self.window_size
        query = nn.pad(query, (0, 0, 0, pad_r, 0, pad_b))
        H_pad, W_pad = query.shape[1], query.shape[2]

        # cyclic shift
        if self.shift_size > 0:
            shifted_query = jt.roll(query,
                                    shifts=(-self.shift_size,
                                            -self.shift_size),
                                    dims=(1, 2))

            # calculate attention mask for SW-MSA
            img_mask = jt.zeros((1, H_pad, W_pad, 1))
            h_slices = (slice(0, -self.window_size),
                        slice(-self.window_size,
                              -self.shift_size), slice(-self.shift_size, None))
            w_slices = (slice(0, -self.window_size),
                        slice(-self.window_size,
                              -self.shift_size), slice(-self.shift_size, None))
            cnt = 0
            for h in h_slices:
                for w in w_slices:
                    img_mask[:, h, w, :] = cnt
                    cnt += 1

            # nW, window_size, window_size, 1
            mask_windows = self.window_partition(img_mask)
            mask_windows = mask_windows.view(
                -1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0,
                                              float(-100.0)).masked_fill(
                                                  attn_mask == 0, float(0.0))
        else:
            shifted_query = query
            attn_mask = None

        # nW*B, window_size, window_size, C
        query_windows = self.window_partition(shifted_query)
        # nW*B, window_size*window_size, C
        query_windows = query_windows.view(-1, self.window_size**2, C)

        # W-MSA/SW-MSA (nW*B, window_size*window_size, C)
        attn_windows = self.w_msa(query_windows, mask=attn_mask)

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size,
                                         self.window_size, C)

        # B H' W' C
        shifted_x = self.window_reverse(attn_windows, H_pad, W_pad)
        # reverse cyclic shift
        if self.shift_size > 0:
            x = jt.roll(shifted_x,
                        shifts=(self.shift_size, self.shift_size),
                        dims=(1, 2))
        else:
            x = shifted_x

        if pad_r > 0 or pad_b:
            x = x[:, :H, :W, :]

        x = x.view(B, H * W, C)

        x = self.drop(x)
        return x

    def window_reverse(self, windows, H, W):
        window_size = self.window_size
        B = int(windows.shape[0] / (H * W / window_size / window_size))
        x = windows.view(B, H // window_size, W // window_size, window_size,
                         window_size, -1)
        x = x.permute(0, 1, 3, 2, 4, 5).view(B, H, W, -1)
        return x

    def window_partition(self, x):
        B, H, W, C = x.shape
        window_size = self.window_size
        x = x.view(B, H // window_size, window_size, W // window_size,
                   window_size, C)
        windows = x.permute(0, 1, 3, 2, 4, 5)
        windows = windows.view(-1, window_size, window_size, C)
        return windows


class SwinBlock(nn.Module):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 feedexecute_channels,
                 window_size=7,
                 shift=False,
                 qkv_bias=True,
                 qk_scale=None,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 act_cfg=dict(type='GELU'),
                 norm_cfg=dict(type='LN')):

        super(SwinBlock, self).__init__()

        self.norm1 = build_norm_layer(norm_cfg, embed_dims)[1]
        self.attn = ShiftWindowMSA(embed_dims=embed_dims,
                                   num_heads=num_heads,
                                   window_size=window_size,
                                   shift_size=window_size // 2 if shift else 0,
                                   qkv_bias=qkv_bias,
                                   qk_scale=qk_scale,
                                   attn_drop_rate=attn_drop_rate,
                                   proj_drop_rate=drop_rate,
                                   dropout_layer=dict(type='DropPath',
                                                      p=drop_path_rate))

        self.norm2 = build_norm_layer(norm_cfg, embed_dims)[1]
        self.ffn = FFN(embed_dims=embed_dims,
                       feedexecute_channels=feedexecute_channels,
                       num_fcs=2,
                       ffn_drop=drop_rate,
                       dropout_layer=dict(type='DropPath', p=drop_path_rate),
                       act_cfg=act_cfg,
                       add_identity=True)

    def execute(self, x, hw_shape):

        def _inner_execute(x):
            identity = x
            x = self.norm1(x)
            x = self.attn(x, hw_shape)

            x = x + identity

            identity = x
            x = self.norm2(x)
            x = self.ffn(x, identity=identity)

            return x

        x = _inner_execute(x)

        return x


class SwinBlockSequence(nn.Module):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 feedexecute_channels,
                 depth,
                 window_size=7,
                 qkv_bias=True,
                 qk_scale=None,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 downsample=None,
                 act_cfg=dict(type='GELU'),
                 norm_cfg=dict(type='LN')):
        super().__init__()

        if isinstance(drop_path_rate, list):
            drop_path_rates = drop_path_rate
            assert len(drop_path_rates) == depth
        else:
            drop_path_rates = [deepcopy(drop_path_rate) for _ in range(depth)]

        self.blocks = nn.ModuleList()
        for i in range(depth):
            block = SwinBlock(embed_dims=embed_dims,
                              num_heads=num_heads,
                              feedexecute_channels=feedexecute_channels,
                              window_size=window_size,
                              shift=False if i % 2 == 0 else True,
                              qkv_bias=qkv_bias,
                              qk_scale=qk_scale,
                              drop_rate=drop_rate,
                              attn_drop_rate=attn_drop_rate,
                              drop_path_rate=drop_path_rates[i],
                              act_cfg=act_cfg,
                              norm_cfg=norm_cfg)

            self.blocks.append(block)

        self.downsample = downsample

    def execute(self, x, hw_shape):
        for block in self.blocks:
            x = block(x, hw_shape)

        if self.downsample:
            x_down, down_hw_shape = self.downsample(x, hw_shape)
            return x_down, down_hw_shape, x, hw_shape
        else:
            return x, hw_shape, x, hw_shape


@BACKBONES.register_module()
class SwinTransformer(nn.Module):

    def __init__(self,
                 pretrain_img_size=224,
                 in_channels=3,
                 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.1,
                 use_abs_pos_embed=False,
                 act_cfg=dict(type='GELU'),
                 norm_cfg=dict(type='LN'),
                 frozen_stages=-1):
        self.frozen_stages = frozen_stages

        if isinstance(pretrain_img_size, int):
            pretrain_img_size = to_2tuple(pretrain_img_size)
        elif isinstance(pretrain_img_size, tuple):
            if len(pretrain_img_size) == 1:
                pretrain_img_size = to_2tuple(pretrain_img_size[0])
            assert len(pretrain_img_size) == 2, \
                f'The size of image should have length 1 or 2, ' \
                f'but got {len(pretrain_img_size)}'

        super(SwinTransformer, self).__init__()

        num_layers = len(depths)
        self.out_indices = out_indices
        self.use_abs_pos_embed = use_abs_pos_embed

        assert strides[0] == patch_size, 'Use non-overlapping patch embed.'

        self.patch_embed = PatchEmbed(
            in_channels=in_channels,
            embed_dims=embed_dims,
            conv_type='Conv2d',
            kernel_size=patch_size,
            stride=strides[0],
            padding='corner',
            norm_cfg=norm_cfg if patch_norm else None,
        )

        if self.use_abs_pos_embed:
            patch_row = pretrain_img_size[0] // patch_size
            patch_col = pretrain_img_size[1] // patch_size
            num_patches = patch_row * patch_col
            self.absolute_pos_embed = nn.Parameter(
                jt.zeros((1, num_patches, embed_dims)))

        self.drop_after_pos = nn.Dropout(p=drop_rate)

        # set stochastic depth decay rule
        total_depth = sum(depths)
        dpr = [x.item() for x in jt.linspace(0, drop_path_rate, total_depth)]

        self.stages = nn.ModuleList()
        in_channels = embed_dims
        for i in range(num_layers):
            if i < num_layers - 1:
                downsample = PatchMerging(
                    in_channels=in_channels,
                    out_channels=2 * in_channels,
                    stride=strides[i + 1],
                    norm_cfg=norm_cfg if patch_norm else None,
                )
            else:
                downsample = None

            stage = SwinBlockSequence(
                embed_dims=in_channels,
                num_heads=num_heads[i],
                feedexecute_channels=int(mlp_ratio * in_channels),
                depth=depths[i],
                window_size=window_size,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop_rate=drop_rate,
                attn_drop_rate=attn_drop_rate,
                drop_path_rate=dpr[sum(depths[:i]):sum(depths[:i + 1])],
                downsample=downsample,
                act_cfg=act_cfg,
                norm_cfg=norm_cfg,
            )

            self.stages.append(stage)
            if downsample:
                in_channels = downsample.out_channels

        self.num_features = [int(embed_dims * 2**i) for i in range(num_layers)]
        # Add a norm layer for each output
        for i in out_indices:
            layer = build_norm_layer(norm_cfg, self.num_features[i])[1]
            layer_name = f'norm{i}'
            setattr(self, layer_name, layer)

    def train(self, mode=True):
        """Convert the model into training mode while keep layers freezed."""
        super(SwinTransformer, self).train(mode)
        self._freeze_stages()

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.patch_embed.eval()
            for param in self.patch_embed.parameters():
                param.requires_grad = False
            if self.use_abs_pos_embed:
                self.absolute_pos_embed.requires_grad = False
            self.drop_after_pos.eval()

        for i in range(1, self.frozen_stages + 1):

            if (i - 1) in self.out_indices:
                norm_layer = getattr(self, f'norm{i-1}')
                norm_layer.eval()
                for param in norm_layer.parameters():
                    param.requires_grad = False

            m = self.stages[i - 1]
            m.eval()
            for param in m.parameters():
                param.requires_grad = False

    def init_weights(self, pretrained=None):
        if pretrained is None:
            if self.use_abs_pos_embed:
                trunc_normal_(self.absolute_pos_embed, std=0.02)
            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, val=1.0, bias=0.)
        else:
            use_shape = False
            if 'pkl' in pretrained:
                use_shape = True
            ckpt = jt.load(pretrained)
            if 'state_dict' in ckpt:
                _state_dict = ckpt['state_dict']
            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()}

            # reshape absolute position embedding
            if state_dict.get('absolute_pos_embed') is not None:
                absolute_pos_embed = state_dict['absolute_pos_embed']
                if use_shape:
                    N1, L, C1 = absolute_pos_embed.shape
                    N2, C2, H, W = self.absolute_pos_embed.shape
                else:
                    N1, L, C1 = absolute_pos_embed.size()
                    N2, C2, H, W = self.absolute_pos_embed.size()
                if N1 != N2 or C1 != C2 or L != H * W:
                    print('Error in loading absolute_pos_embed, pass')
                else:
                    state_dict['absolute_pos_embed'] = absolute_pos_embed.view(
                        N2, H, W, C2).permute(0, 3, 1, 2)

            # interpolate position bias table if needed
            relative_position_bias_table_keys = [
                k for k in state_dict.keys()
                if 'relative_position_bias_table' in k
            ]
            for table_key in relative_position_bias_table_keys:
                table_pretrained = state_dict[table_key]
                table_current = self.state_dict()[table_key]
                if use_shape:
                    L1, nH1 = table_pretrained.shape
                    L2, nH2 = table_current.shape
                else:
                    L1, nH1 = table_pretrained.size()
                    L2, nH2 = table_current.size()
                if nH1 != nH2:
                    print(f'Error in loading {table_key}, pass')
                elif L1 != L2:
                    S1 = int(L1**0.5)
                    S2 = int(L2**0.5)
                    table_pretrained_resized = nn.interpolate(
                        table_pretrained.permute(1, 0).reshape(1, nH1, S1, S1),
                        size=(S2, S2),
                        mode='bicubic')
                    state_dict[table_key] = table_pretrained_resized.view(
                        nH2, L2).permute(1, 0)

            # load state_dict
            self.load_parameters(state_dict)

    def execute(self, x):
        x, hw_shape = self.patch_embed(x)

        if self.use_abs_pos_embed:
            x = x + self.absolute_pos_embed
        x = self.drop_after_pos(x)

        outs = []
        for i, stage in enumerate(self.stages):
            x, hw_shape, out, out_hw_shape = stage(x, hw_shape)
            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                out = norm_layer(out)
                out = out.view(-1, *out_hw_shape,
                               self.num_features[i]).permute(0, 3, 1, 2)
                outs.append(out)

        return outs

python/jittorseg/models/backbones/swin.py
JittorCV/jittorseg
import math

import jittor as jt
from jittor import nn
from jittorseg.models.basic_operators.embed import PatchEmbed, FFN
from jittorseg.utils.weight_init import kaiming_init, constant_init, trunc_normal_
from jittorseg.utils.helpers import to_2tuple
from ..basic_operators import build_norm_layer
from jittor.nn import BatchNorm as _BatchNorm

from ..layers import resize, MultiHeadAttention
from jittorseg.utils.registry import BACKBONES


class TransformerEncoderLayer(nn.Module):

    def __init__(self,
                 embed_dims,
                 num_heads,
                 feedexecute_channels,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 num_fcs=2,
                 qkv_bias=True,
                 act_cfg=dict(type='GELU'),
                 norm_cfg=dict(type='LN'),
                 batch_first=True,
                 attn_cfg=dict(),
                 ffn_cfg=dict()):
        super(TransformerEncoderLayer, self).__init__()

        self.norm1_name, norm1 = build_norm_layer(norm_cfg,
                                                  embed_dims,
                                                  postfix=1)
        setattr(self, self.norm1_name, norm1)

        attn_cfg.update(
            dict(embed_dims=embed_dims,
                 num_heads=num_heads,
                 attn_drop=attn_drop_rate,
                 proj_drop=drop_rate,
                 batch_first=batch_first,
                 bias=qkv_bias))

        self.build_attn(attn_cfg)

        self.norm2_name, norm2 = build_norm_layer(norm_cfg,
                                                  embed_dims,
                                                  postfix=2)
        setattr(self, self.norm2_name, norm2)

        ffn_cfg.update(
            dict(embed_dims=embed_dims,
                 feedexecute_channels=feedexecute_channels,
                 num_fcs=num_fcs,
                 ffn_drop=drop_rate,
                 dropout_layer=dict(type='DropPath', p=drop_path_rate)
                 if drop_path_rate > 0 else None,
                 act_cfg=act_cfg))
        self.build_ffn(ffn_cfg)

    def build_attn(self, attn_cfg):
        self.attn = MultiHeadAttention(**attn_cfg)

    def build_ffn(self, ffn_cfg):
        self.ffn = FFN(**ffn_cfg)

    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        return getattr(self, self.norm2_name)

    def execute(self, x):

        def _inner_execute(x):
            x = self.attn(self.norm1(x), identity=x)
            x = self.ffn(self.norm2(x), identity=x)
            return x

        x = _inner_execute(x)
        return x


@BACKBONES.register_module()
class VisionTransformer(nn.Module):

    def __init__(self,
                 img_size=224,
                 patch_size=16,
                 in_channels=3,
                 embed_dims=768,
                 num_layers=12,
                 num_heads=12,
                 mlp_ratio=4,
                 out_indices=-1,
                 qkv_bias=True,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.,
                 with_cls_token=True,
                 output_cls_token=False,
                 norm_cfg=dict(type='LN'),
                 act_cfg=dict(type='GELU'),
                 patch_norm=False,
                 final_norm=False,
                 interpolate_mode='bicubic',
                 num_fcs=2,
                 norm_eval=False):
        super(VisionTransformer, self).__init__()

        if isinstance(img_size, int):
            img_size = to_2tuple(img_size)
        elif isinstance(img_size, tuple):
            if len(img_size) == 1:
                img_size = to_2tuple(img_size[0])
            assert len(img_size) == 2, \
                f'The size of image should have length 1 or 2, ' \
                f'but got {len(img_size)}'

        if output_cls_token:
            assert with_cls_token is True, f'with_cls_token must be True if' \
                f'set output_cls_token to True, but got {with_cls_token}'

        self.img_size = img_size
        self.patch_size = patch_size
        self.interpolate_mode = interpolate_mode
        self.norm_eval = norm_eval

        self.patch_embed = PatchEmbed(
            in_channels=in_channels,
            embed_dims=embed_dims,
            conv_type='Conv2d',
            kernel_size=patch_size,
            stride=patch_size,
            padding='corner',
            norm_cfg=norm_cfg if patch_norm else None)

        num_patches = (img_size[0] // patch_size) * \
            (img_size[1] // patch_size)

        self.with_cls_token = with_cls_token
        self.output_cls_token = output_cls_token
        self.cls_token = jt.zeros((1, 1, embed_dims))
        self.pos_embed = jt.zeros((1, num_patches + 1, embed_dims))
        self.drop_after_pos = nn.Dropout(p=drop_rate)

        if isinstance(out_indices, int):
            if out_indices == -1:
                out_indices = num_layers - 1
            self.out_indices = [out_indices]
        elif isinstance(out_indices, list) or isinstance(out_indices, tuple):
            self.out_indices = out_indices
        else:
            raise TypeError('out_indices must be type of int, list or tuple')

        dpr = [x.item() for x in jt.linspace(0, drop_path_rate, num_layers)
               ]  # stochastic depth decay rule

        self.layers = nn.ModuleList()
        for i in range(num_layers):
            self.layers.append(
                TransformerEncoderLayer(embed_dims=embed_dims,
                                        num_heads=num_heads,
                                        feedexecute_channels=mlp_ratio *
                                        embed_dims,
                                        attn_drop_rate=attn_drop_rate,
                                        drop_rate=drop_rate,
                                        drop_path_rate=dpr[i],
                                        num_fcs=num_fcs,
                                        qkv_bias=qkv_bias,
                                        act_cfg=act_cfg,
                                        norm_cfg=norm_cfg,
                                        batch_first=True))

        self.final_norm = final_norm
        if final_norm:
            self.norm1_name, norm1 = build_norm_layer(norm_cfg,
                                                      embed_dims,
                                                      postfix=1)
            setattr(self, self.norm1_name, norm1)

    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    def init_weights(self, pretrained=None):
        if pretrained is None:
            trunc_normal_(self.pos_embed, std=.02)
            trunc_normal_(self.cls_token, std=.02)
            for n, m in self.named_modules():
                if isinstance(m, nn.Linear):
                    trunc_normal_(m.weight, std=.02)
                    if m.bias is not None:
                        if 'ffn' in n:
                            nn.init.normal_(m.bias, mean=0., std=1e-6)
                        else:
                            nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Conv2d):
                    kaiming_init(m, mode='fan_in', bias=0.)
                elif isinstance(m, (_BatchNorm, nn.GroupNorm, nn.LayerNorm)):
                    constant_init(m, val=1.0, bias=0.)
        elif isinstance(pretrained, str):
            checkpoint = jt.load(pretrained)
            if 'state_dict' in checkpoint:
                state_dict = checkpoint['state_dict']
            else:
                state_dict = checkpoint

            if 'pos_embed' in state_dict.keys():
                if self.pos_embed.shape != state_dict['pos_embed'].shape:
                    h, w = self.img_size
                    pos_size = int(
                        math.sqrt(state_dict['pos_embed'].shape[1] - 1))
                    state_dict['pos_embed'] = self.resize_pos_embed(
                        state_dict['pos_embed'],
                        (h // self.patch_size, w // self.patch_size),
                        (pos_size, pos_size), self.interpolate_mode)
            self.load_parameters(state_dict)

    def _pos_embeding(self, patched_img, hw_shape, pos_embed):
        assert patched_img.ndim == 3 and pos_embed.ndim == 3, \
            'the shapes of patched_img and pos_embed must be [B, L, C]'
        x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
        if x_len != pos_len:
            if pos_len == (self.img_size[0] // self.patch_size) * (
                    self.img_size[1] // self.patch_size) + 1:
                pos_h = self.img_size[0] // self.patch_size
                pos_w = self.img_size[1] // self.patch_size
            else:
                raise ValueError(
                    'Unexpected shape of pos_embed, got {}.'.format(
                        pos_embed.shape))
            pos_embed = self.resize_pos_embed(pos_embed, hw_shape,
                                              (pos_h, pos_w),
                                              self.interpolate_mode)
        return self.drop_after_pos(patched_img + pos_embed)

    @staticmethod
    def resize_pos_embed(pos_embed, input_shpae, pos_shape, mode):
        assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]'
        if not isinstance(pos_embed, jt.Var):
            pos_embed = jt.Var(pos_embed)
        pos_h, pos_w = pos_shape
        cls_token_weight = pos_embed[:, 0]
        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
        pos_embed_weight = pos_embed_weight.reshape(
            1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
        pos_embed_weight = resize(pos_embed_weight,
                                  size=input_shpae,
                                  align_corners=False,
                                  mode=mode)
        cls_token_weight = cls_token_weight.unsqueeze(1)
        pos_embed_weight = jt.flatten(pos_embed_weight, 2).transpose(1, 2)
        pos_embed = jt.concat((cls_token_weight, pos_embed_weight), dim=1)
        return pos_embed

    def execute(self, inputs):
        B = inputs.shape[0]
        x, hw_shape = self.patch_embed(inputs)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = jt.concat((cls_tokens, x), dim=1)
        x = self._pos_embeding(x, hw_shape, self.pos_embed)

        if not self.with_cls_token:
            # Remove class token for transformer encoder input
            x = x[:, 1:]

        outs = []
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i == len(self.layers) - 1:
                if self.final_norm:
                    x = self.norm1(x)
            if i in self.out_indices:
                if self.with_cls_token:
                    # Remove class token and reshape token for decoder head
                    out = x[:, 1:]
                else:
                    out = x
                B, _, C = out.shape
                out = out.reshape(B, hw_shape[0], hw_shape[1],
                                  C).permute(0, 3, 1, 2)
                if self.output_cls_token:
                    out = [out, x[:, 0]]
                outs.append(out)

        return tuple(outs)

    def train(self, mode=True):
        super(VisionTransformer, self).train(mode)
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, nn.LayerNorm):
                    m.eval()

python/jittorseg/models/backbones/vit.py
JittorCV/jittorseg
from .conv import build_conv_layer
from .activation import build_activation_layer
from .norm import build_norm_layer
from .drop import build_dropout
from .padding import build_padding_layer
from .conv_module import ConvModule
from .depthwise_separable_conv_module import DepthwiseSeparableConvModule

python/jittorseg/models/basic_operators/__init__.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn

from jittorseg.utils.registry import ACTIVATION_LAYERS, build_from_cfg

for module in [
        nn.ReLU, nn.LeakyReLU, nn.PReLU, nn.ReLU6, nn.ELU, nn.Sigmoid, nn.Tanh
]:
    if module.__name__ == 'relu':
        ACTIVATION_LAYERS.register_module(name='ReLU', module=module)
    elif module.__name__ == 'relu6':
        ACTIVATION_LAYERS.register_module(name='ReLU6', module=module)
    else:
        ACTIVATION_LAYERS.register_module(module=module)


@ACTIVATION_LAYERS.register_module()
class HSigmoid(nn.Module):

    def __init__(self, bias=3.0, divisor=6.0, min_value=0.0, max_value=1.0):
        super().__init__()
        self.bias = bias
        self.divisor = divisor
        assert self.divisor != 0
        self.min_value = min_value
        self.max_value = max_value

    def execute(self, x):
        x = (x + self.bias) / self.divisor

        return x.clamp_(self.min_value, self.max_value)


@ACTIVATION_LAYERS.register_module()
class HSwish(nn.Module):

    def __init__(self):
        super().__init__()
        self.act = nn.ReLU6()

    def execute(self, x):
        return x * self.act(x + 3) / 6


@ACTIVATION_LAYERS.register_module(name='Clip')
@ACTIVATION_LAYERS.register_module()
class Clamp(nn.Module):

    def __init__(self, min=-1., max=1.):
        super().__init__()
        self.min = min
        self.max = max

    def execute(self, x):
        return jt.clamp(x, min_v=self.min, max_v=self.max)


class GELU(nn.Module):

    def execute(self, input):
        return nn.gelu(input)


ACTIVATION_LAYERS.register_module(module=GELU)


def build_activation_layer(cfg):
    return build_from_cfg(cfg, ACTIVATION_LAYERS)

python/jittorseg/models/basic_operators/activation.py
JittorCV/jittorseg
import math
from jittor import nn
from jittorseg.utils.registry import CONV_LAYERS
from typing import Tuple, Union

CONV_LAYERS.register_module('Conv1d', module=nn.Conv1d)
CONV_LAYERS.register_module('Conv2d', module=nn.Conv2d)
CONV_LAYERS.register_module('Conv3d', module=nn.Conv3d)
CONV_LAYERS.register_module('Conv', module=nn.Conv2d)


class Conv2dAdaptivePadding(nn.Conv2d):

    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: Union[int, Tuple[int, int]],
                 stride: Union[int, Tuple[int, int]] = 1,
                 padding: Union[int, Tuple[int, int]] = 0,
                 dilation: Union[int, Tuple[int, int]] = 1,
                 groups: int = 1,
                 bias: bool = True):
        super().__init__(in_channels, out_channels, kernel_size, stride, 0,
                         dilation, groups, bias)

    def execute(self, x):
        img_h, img_w = x.size()[-2:]
        kernel_h, kernel_w = self.weight.size()[-2:]
        stride_h, stride_w = self.stride
        output_h = math.ceil(img_h / stride_h)
        output_w = math.ceil(img_w / stride_w)
        pad_h = (max((output_h - 1) * self.stride[0] +
                     (kernel_h - 1) * self.dilation[0] + 1 - img_h, 0))
        pad_w = (max((output_w - 1) * self.stride[1] +
                     (kernel_w - 1) * self.dilation[1] + 1 - img_w, 0))
        if pad_h > 0 or pad_w > 0:
            x = nn.pad(x, [
                pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2
            ])
        return nn.conv2d(x, self.weight, self.bias, self.stride, self.padding,
                         self.dilation, self.groups)


CONV_LAYERS.register_module('Conv2dAdaptivePadding',
                            module=Conv2dAdaptivePadding)


def build_conv_layer(cfg, *args, **kwargs):
    """Build convolution layer.

    Args:
        cfg (None or dict): The conv layer config, which should contain:
            - type (str): Layer type.
            - layer args: Args needed to instantiate an conv layer.
        args (argument list): Arguments passed to the `__init__`
            method of the corresponding conv layer.
        kwargs (keyword arguments): Keyword arguments passed to the `__init__`
            method of the corresponding conv layer.

    Returns:
        nn.Module: Created conv layer.
    """
    if cfg is None:
        cfg_ = dict(type='Conv2d')
    else:
        if not isinstance(cfg, dict):
            raise TypeError('cfg must be a dict')
        if 'type' not in cfg:
            raise KeyError('the cfg dict must contain the key "type"')
        cfg_ = cfg.copy()

    layer_type = cfg_.pop('type')
    conv_layer = CONV_LAYERS.get(layer_type)

    layer = conv_layer(*args, **kwargs, **cfg_)

    return layer

python/jittorseg/models/basic_operators/conv.py
JittorCV/jittorseg
import warnings

from jittor import nn

from jittor.nn import BatchNorm as _BatchNorm
from jittor.nn import InstanceNorm as _InstanceNorm
from jittorseg.utils.weight_init import constant_init, kaiming_init
from .activation import build_activation_layer
from .conv import build_conv_layer
from .norm import build_norm_layer
from .padding import build_padding_layer


class ConvModule(nn.Module):
    _abbr_ = 'conv_block'

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 bias='auto',
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=dict(type='ReLU'),
                 padding_mode='zeros',
                 order=('conv', 'norm', 'act')):
        super().__init__()
        assert conv_cfg is None or isinstance(conv_cfg, dict)
        assert norm_cfg is None or isinstance(norm_cfg, dict)
        assert act_cfg is None or isinstance(act_cfg, dict)
        official_padding_mode = ['zeros', 'circular']
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.with_explicit_padding = padding_mode not in official_padding_mode
        self.order = order
        assert isinstance(self.order, tuple) and len(self.order) == 3
        assert set(order) == {'conv', 'norm', 'act'}

        self.with_norm = norm_cfg is not None
        self.with_activation = act_cfg is not None
        # if the conv layer is before a norm layer, bias is unnecessary.
        if bias == 'auto':
            bias = not self.with_norm
        self.with_bias = bias

        if self.with_explicit_padding:
            pad_cfg = dict(type=padding_mode)
            self.padding_layer = build_padding_layer(pad_cfg, padding)

        # reset padding to 0 for conv module
        conv_padding = 0 if self.with_explicit_padding else padding
        # build convolution layer
        self.conv = build_conv_layer(conv_cfg,
                                     in_channels,
                                     out_channels,
                                     kernel_size,
                                     stride=stride,
                                     padding=conv_padding,
                                     dilation=dilation,
                                     groups=groups,
                                     bias=bias)
        # export the attributes of self.conv to a higher level for convenience
        self.in_channels = self.conv.in_channels
        self.out_channels = self.conv.out_channels
        self.kernel_size = self.conv.kernel_size
        self.stride = self.conv.stride
        self.padding = padding
        self.dilation = self.conv.dilation
        self.groups = self.conv.groups

        # build normalization layers
        if self.with_norm:
            # norm layer is after conv layer
            if order.index('norm') > order.index('conv'):
                norm_channels = out_channels
            else:
                norm_channels = in_channels
            self.norm_name, norm = build_norm_layer(norm_cfg, norm_channels)
            setattr(self, self.norm_name, norm)
            if self.with_bias:
                if isinstance(norm, (_BatchNorm, _InstanceNorm)):
                    warnings.warn(
                        'Unnecessary conv bias before batch/instance norm')
        else:
            self.norm_name = None

        # build activation layer
        if self.with_activation:
            self.activate = build_activation_layer(act_cfg)

        # Use msra init by default
        self.init_weights()

    @property
    def norm(self):
        if self.norm_name:
            return getattr(self, self.norm_name)
        else:
            return None

    def init_weights(self):
        if not hasattr(self.conv, 'init_weights'):
            if self.with_activation and self.act_cfg['type'] == 'LeakyReLU':
                nonlinearity = 'leaky_relu'
                a = self.act_cfg.get('negative_slope', 0.01)
            else:
                nonlinearity = 'relu'
                a = 0
            kaiming_init(self.conv, a=a, nonlinearity=nonlinearity)
        if self.with_norm:
            constant_init(self.norm, 1, bias=0)

    def execute(self, x, activate=True, norm=True):
        for layer in self.order:
            if layer == 'conv':
                if self.with_explicit_padding:
                    x = self.padding_layer(x)
                x = self.conv(x)
            elif layer == 'norm' and norm and self.with_norm:
                x = self.norm(x)
            elif layer == 'act' and activate and self.with_activation:
                x = self.activate(x)
        return x

python/jittorseg/models/basic_operators/conv_module.py
JittorCV/jittorseg
from jittor import nn

from .conv_module import ConvModule


class DepthwiseSeparableConvModule(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 norm_cfg=None,
                 act_cfg=dict(type='ReLU'),
                 dw_norm_cfg='default',
                 dw_act_cfg='default',
                 pw_norm_cfg='default',
                 pw_act_cfg='default',
                 **kwargs):
        super().__init__()
        assert 'groups' not in kwargs, 'groups should not be specified'

        # if norm/activation config of depthwise/pointwise ConvModule is not
        # specified, use default config.
        dw_norm_cfg = dw_norm_cfg if dw_norm_cfg != 'default' else norm_cfg
        dw_act_cfg = dw_act_cfg if dw_act_cfg != 'default' else act_cfg
        pw_norm_cfg = pw_norm_cfg if pw_norm_cfg != 'default' else norm_cfg
        pw_act_cfg = pw_act_cfg if pw_act_cfg != 'default' else act_cfg

        # depthwise convolution
        self.depthwise_conv = ConvModule(
            in_channels,
            in_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=in_channels,
            norm_cfg=dw_norm_cfg,
            act_cfg=dw_act_cfg,
            **kwargs)

        self.pointwise_conv = ConvModule(
            in_channels,
            out_channels,
            1,
            norm_cfg=pw_norm_cfg,
            act_cfg=pw_act_cfg,
            **kwargs)

    def execute(self, x):
        x = self.depthwise_conv(x)
        x = self.pointwise_conv(x)
        return x

python/jittorseg/models/basic_operators/depthwise_separable_conv_module.py
JittorCV/jittorseg
from jittor.nn import Dropout, DropPath
from jittorseg.utils.registry import DROPOUT_LAYERS, build_from_cfg

DROPOUT_LAYERS.register_module(name='Dropout', module=Dropout)
DROPOUT_LAYERS.register_module(name='DropPath', module=DropPath)


def build_dropout(cfg, **default_args):
    """Builder for drop out layers."""
    return build_from_cfg(cfg, DROPOUT_LAYERS, **default_args)

python/jittorseg/models/basic_operators/drop.py
JittorCV/jittorseg
from typing import Sequence
import math
from jittor import nn
from jittorseg.utils.helpers import to_2tuple
from jittorseg.models.basic_operators import build_conv_layer, build_norm_layer, build_activation_layer, build_dropout


class AdaptivePadding(nn.Module):

    def __init__(self, kernel_size=1, stride=1, dilation=1, padding='corner'):

        super(AdaptivePadding, self).__init__()

        assert padding in ('same', 'corner')

        kernel_size = to_2tuple(kernel_size)
        stride = to_2tuple(stride)
        dilation = to_2tuple(dilation)

        self.padding = padding
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation

    def get_pad_shape(self, input_shape):
        input_h, input_w = input_shape
        kernel_h, kernel_w = self.kernel_size
        stride_h, stride_w = self.stride
        output_h = math.ceil(input_h / stride_h)
        output_w = math.ceil(input_w / stride_w)
        pad_h = max((output_h - 1) * stride_h +
                    (kernel_h - 1) * self.dilation[0] + 1 - input_h, 0)
        pad_w = max((output_w - 1) * stride_w +
                    (kernel_w - 1) * self.dilation[1] + 1 - input_w, 0)
        return pad_h, pad_w

    def execute(self, x):
        pad_h, pad_w = self.get_pad_shape(x.size()[-2:])
        if pad_h > 0 or pad_w > 0:
            if self.padding == 'corner':
                x = nn.pad(x, [0, pad_w, 0, pad_h])
            elif self.padding == 'same':
                x = nn.pad(x, [
                    pad_w // 2, pad_w - pad_w // 2, pad_h // 2,
                    pad_h - pad_h // 2
                ])
        return x


class PatchEmbed(nn.Module):

    def __init__(self,
                 in_channels=3,
                 embed_dims=768,
                 conv_type='Conv2d',
                 kernel_size=16,
                 stride=None,
                 padding='corner',
                 dilation=1,
                 bias=True,
                 norm_cfg=None,
                 input_size=None):
        super(PatchEmbed, self).__init__()

        self.embed_dims = embed_dims
        if stride is None:
            stride = kernel_size

        kernel_size = to_2tuple(kernel_size)
        stride = to_2tuple(stride)
        dilation = to_2tuple(dilation)

        if isinstance(padding, str):
            self.adap_padding = AdaptivePadding(kernel_size=kernel_size,
                                                stride=stride,
                                                dilation=dilation,
                                                padding=padding)
            # disable the padding of conv
            padding = 0
        else:
            self.adap_padding = None
        padding = to_2tuple(padding)

        self.projection = build_conv_layer(dict(type=conv_type),
                                           in_channels=in_channels,
                                           out_channels=embed_dims,
                                           kernel_size=kernel_size,
                                           stride=stride,
                                           padding=padding,
                                           dilation=dilation,
                                           bias=bias)

        if norm_cfg is not None:
            self.norm = build_norm_layer(norm_cfg, embed_dims)[1]
        else:
            self.norm = None

        if input_size:
            input_size = to_2tuple(input_size)
            # `init_out_size` would be used outside to
            # calculate the num_patches
            # when `use_abs_pos_embed` outside
            self.init_input_size = input_size
            if self.adap_padding:
                pad_h, pad_w = self.adap_padding.get_pad_shape(input_size)
                input_h, input_w = input_size
                input_h = input_h + pad_h
                input_w = input_w + pad_w
                input_size = (input_h, input_w)

            h_out = (input_size[0] + 2 * padding[0] - dilation[0] *
                     (kernel_size[0] - 1) - 1) // stride[0] + 1
            w_out = (input_size[1] + 2 * padding[1] - dilation[1] *
                     (kernel_size[1] - 1) - 1) // stride[1] + 1
            self.init_out_size = (h_out, w_out)
        else:
            self.init_input_size = None
            self.init_out_size = None

    def execute(self, x):
        """
        Args:
            x (Var): Has shape (B, C, H, W). In most case, C is 3.

        Returns:
            tuple: Contains merged results and its spatial shape.

                - x (Var): Has shape (B, out_h * out_w, embed_dims)
                - out_size (tuple[int]): Spatial shape of x, arrange as
                    (out_h, out_w).
        """

        if self.adap_padding:
            x = self.adap_padding(x)

        x = self.projection(x)
        out_size = (x.shape[2], x.shape[3])
        x = x.flatten(2).transpose(1, 2)
        if self.norm is not None:
            x = self.norm(x)
        return x, out_size


class PatchMerging(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=2,
                 stride=None,
                 padding='corner',
                 dilation=1,
                 bias=False,
                 norm_cfg=dict(type='LN')):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        if stride:
            stride = stride
        else:
            stride = kernel_size

        self.kernel_size = to_2tuple(kernel_size)
        self.stride = to_2tuple(stride)
        self.dilation = to_2tuple(dilation)

        if isinstance(padding, str):
            self.adap_padding = AdaptivePadding(kernel_size=self.kernel_size,
                                                stride=self.stride,
                                                dilation=self.dilation,
                                                padding=padding)
            # disable the padding of unfold
            padding = 0
        else:
            self.adap_padding = None

        self.padding = to_2tuple(padding)

        sample_dim = self.kernel_size[0] * self.kernel_size[1] * in_channels

        if norm_cfg is not None:
            self.norm = build_norm_layer(norm_cfg, sample_dim)[1]
        else:
            self.norm = None

        self.reduction = nn.Linear(sample_dim, out_channels, bias=bias)

    def execute(self, x, input_size):
        """
        Args:
            x (Var): Has shape (B, H*W, C_in).
            input_size (tuple[int]): The spatial shape of x, arrange as (H, W).
                Default: None.

        Returns:
            tuple: Contains merged results and its spatial shape.

                - x (Var): Has shape (B, Merged_H * Merged_W, C_out)
                - out_size (tuple[int]): Spatial shape of x, arrange as
                    (Merged_H, Merged_W).
        """
        B, L, C = x.shape
        assert isinstance(input_size, Sequence), f'Expect ' \
                                                 f'input_size is ' \
                                                 f'`Sequence` ' \
                                                 f'but get {input_size}'

        H, W = input_size
        assert L == H * W, 'input feature has wrong size'

        x = x.view(B, H, W, C).permute([0, 3, 1, 2])  # B, C, H, W
        # Use nn.Unfold to merge patch. About 25% faster than original method,
        # but need to modify pretrained model for compatibility

        if self.adap_padding:
            x = self.adap_padding(x)
            H, W = x.shape[-2:]

        x = nn.unfold(x,
                      kernel_size=self.kernel_size,
                      dilation=self.dilation,
                      padding=self.padding,
                      stride=self.stride)
        # if kernel_size=2 and stride=2, x should has shape (B, 4*C, H/2*W/2)

        out_h = (H + 2 * self.padding[0] - self.dilation[0] *
                 (self.kernel_size[0] - 1) - 1) // self.stride[0] + 1
        out_w = (W + 2 * self.padding[1] - self.dilation[1] *
                 (self.kernel_size[1] - 1) - 1) // self.stride[1] + 1

        output_size = (out_h, out_w)
        x = x.transpose(1, 2)  # B, H/2*W/2, 4*C
        x = self.norm(x) if self.norm else x
        x = self.reduction(x)
        return x, output_size


class FFN(nn.Module):

    def __init__(self,
                 embed_dims=256,
                 feedexecute_channels=1024,
                 num_fcs=2,
                 act_cfg=dict(type='ReLU', inplace=True),
                 ffn_drop=0.,
                 dropout_layer=None,
                 add_identity=True,
                 **kwargs):
        super().__init__()
        assert num_fcs >= 2, 'num_fcs should be no less ' \
                             f'than 2. got {num_fcs}.'
        self.embed_dims = embed_dims
        self.feedexecute_channels = feedexecute_channels
        self.num_fcs = num_fcs
        self.act_cfg = act_cfg
        self.activate = build_activation_layer(act_cfg)

        layers = []
        in_channels = embed_dims
        for _ in range(num_fcs - 1):
            layers.append(
                nn.Sequential(nn.Linear(in_channels, feedexecute_channels),
                              self.activate, nn.Dropout(ffn_drop)))
            in_channels = feedexecute_channels
        layers.append(nn.Linear(feedexecute_channels, embed_dims))
        layers.append(nn.Dropout(ffn_drop))
        self.layers = nn.Sequential(*layers)
        self.dropout_layer = build_dropout(
            dropout_layer) if dropout_layer else nn.Identity()
        self.add_identity = add_identity

    def execute(self, x, identity=None):
        out = self.layers(x)
        if not self.add_identity:
            return self.dropout_layer(out)
        if identity is None:
            identity = x
        return identity + self.dropout_layer(out)

python/jittorseg/models/basic_operators/embed.py
JittorCV/jittorseg
from jittorseg.models.basic_operators import ConvModule
from jittor import nn

from jittorseg.utils.helpers import make_divisible
from jittorseg.utils.general import is_tuple_of


class SELayer(nn.Module):
    def __init__(self,
                 channels,
                 ratio=16,
                 conv_cfg=None,
                 act_cfg=(dict(type='ReLU'),
                          dict(type='HSigmoid', bias=3.0, divisor=6.0))):
        super(SELayer, self).__init__()
        if isinstance(act_cfg, dict):
            act_cfg = (act_cfg, act_cfg)
        assert len(act_cfg) == 2
        assert is_tuple_of(act_cfg, dict)
        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
        self.conv1 = ConvModule(in_channels=channels,
                                out_channels=make_divisible(
                                    channels // ratio, 8),
                                kernel_size=1,
                                stride=1,
                                conv_cfg=conv_cfg,
                                act_cfg=act_cfg[0])
        self.conv2 = ConvModule(in_channels=make_divisible(
            channels // ratio, 8),
            out_channels=channels,
            kernel_size=1,
            stride=1,
            conv_cfg=conv_cfg,
            act_cfg=act_cfg[1])

    def execute(self, x):
        out = self.global_avgpool(x)
        out = self.conv1(out)
        out = self.conv2(out)
        return x * out


class InvertedResidual(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 stride,
                 expand_ratio,
                 dilation=1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU6'),
                 **kwargs):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2], f'stride must in [1, 2]. ' \
                                 f'But received {stride}.'
        self.use_res_connect = self.stride == 1 and in_channels == out_channels
        hidden_dim = int(round(in_channels * expand_ratio))

        layers = []
        if expand_ratio != 1:
            layers.append(
                ConvModule(
                    in_channels=in_channels,
                    out_channels=hidden_dim,
                    kernel_size=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg,
                    **kwargs))
        layers.extend([
            ConvModule(
                in_channels=hidden_dim,
                out_channels=hidden_dim,
                kernel_size=3,
                stride=stride,
                padding=dilation,
                dilation=dilation,
                groups=hidden_dim,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg,
                **kwargs),
            ConvModule(
                in_channels=hidden_dim,
                out_channels=out_channels,
                kernel_size=1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=None,
                **kwargs)
        ])
        self.conv = nn.Sequential(*layers)

    def execute(self, x):

        def _inner_execute(x):
            if self.use_res_connect:
                return x + self.conv(x)
            else:
                return self.conv(x)

        out = _inner_execute(x)

        return out


class InvertedResidualV3(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 mid_channels,
                 kernel_size=3,
                 stride=1,
                 se_cfg=None,
                 with_expand_conv=True,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU')):
        super(InvertedResidualV3, self).__init__()
        self.with_res_shortcut = (stride == 1 and in_channels == out_channels)
        assert stride in [1, 2]
        self.with_se = se_cfg is not None
        self.with_expand_conv = with_expand_conv

        if self.with_se:
            assert isinstance(se_cfg, dict)
        if not self.with_expand_conv:
            assert mid_channels == in_channels

        if self.with_expand_conv:
            self.expand_conv = ConvModule(
                in_channels=in_channels,
                out_channels=mid_channels,
                kernel_size=1,
                stride=1,
                padding=0,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg)
        self.depthwise_conv = ConvModule(
            in_channels=mid_channels,
            out_channels=mid_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=kernel_size // 2,
            groups=mid_channels,
            conv_cfg=dict(
                type='Conv2dAdaptivePadding') if stride == 2 else conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

        if self.with_se:
            self.se = SELayer(**se_cfg)

        self.linear_conv = ConvModule(
            in_channels=mid_channels,
            out_channels=out_channels,
            kernel_size=1,
            stride=1,
            padding=0,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=None)

    def execute(self, x):

        def _inner_execute(x):
            out = x

            if self.with_expand_conv:
                out = self.expand_conv(out)

            out = self.depthwise_conv(out)

            if self.with_se:
                out = self.se(out)

            out = self.linear_conv(out)

            if self.with_res_shortcut:
                return x + out
            else:
                return out

        out = _inner_execute(x)

        return out

python/jittorseg/models/basic_operators/inverted_residual.py
JittorCV/jittorseg
import inspect

from jittor import nn

from jittorseg.utils.general import is_tuple_of
from jittorseg.utils.registry import NORM_LAYERS
from jittor.nn import BatchNorm as _BatchNorm
from jittor.nn import InstanceNorm as _InstanceNorm

NORM_LAYERS.register_module('BN', module=nn.BatchNorm)
NORM_LAYERS.register_module('BN1d', module=nn.BatchNorm)
NORM_LAYERS.register_module('BN2d', module=nn.BatchNorm)
NORM_LAYERS.register_module('BN3d', module=nn.BatchNorm)
NORM_LAYERS.register_module('GN', module=nn.GroupNorm)
NORM_LAYERS.register_module('LN', module=nn.LayerNorm)
NORM_LAYERS.register_module('IN', module=nn.InstanceNorm)
NORM_LAYERS.register_module('IN1d', module=nn.InstanceNorm)
NORM_LAYERS.register_module('IN2d', module=nn.InstanceNorm)
NORM_LAYERS.register_module('IN3d', module=nn.InstanceNorm)


def infer_abbr(class_type):
    """Infer abbreviation from the class name.

    When we build a norm layer with `build_norm_layer()`, we want to preserve
    the norm type in variable names, e.g, self.bn1, self.gn. This method will
    infer the abbreviation to map class types to abbreviations.

    Rule 1: If the class has the property "_abbr_", return the property.
    Rule 2: If the parent class is _BatchNorm, GroupNorm, LayerNorm or
    InstanceNorm, the abbreviation of this layer will be "bn", "gn", "ln" and
    "in" respectively.
    Rule 3: If the class name contains "batch", "group", "layer" or "instance",
    the abbreviation of this layer will be "bn", "gn", "ln" and "in"
    respectively.
    Rule 4: Otherwise, the abbreviation falls back to "norm".

    Args:
        class_type (type): The norm layer type.

    Returns:
        str: The inferred abbreviation.
    """
    if not inspect.isclass(class_type):
        raise TypeError(
            f'class_type must be a type, but got {type(class_type)}')
    if hasattr(class_type, '_abbr_'):
        return class_type._abbr_
    if issubclass(class_type, _InstanceNorm):  # IN is a subclass of BN
        return 'in'
    elif issubclass(class_type, _BatchNorm):
        return 'bn'
    elif issubclass(class_type, nn.GroupNorm):
        return 'gn'
    elif issubclass(class_type, nn.LayerNorm):
        return 'ln'
    else:
        class_name = class_type.__name__.lower()
        if 'batch' in class_name:
            return 'bn'
        elif 'group' in class_name:
            return 'gn'
        elif 'layer' in class_name:
            return 'ln'
        elif 'instance' in class_name:
            return 'in'
        else:
            return 'norm_layer'


def build_norm_layer(cfg, num_features, postfix=''):
    """Build normalization layer.

    Args:
        cfg (dict): The norm layer config, which should contain:

            - type (str): Layer type.
            - layer args: Args needed to instantiate a norm layer.
            - requires_grad (bool, optional): Whether stop gradient updates.
        num_features (int): Number of input channels.
        postfix (int | str): The postfix to be appended into norm abbreviation
            to create named layer.

    Returns:
        tuple[str, nn.Module]: The first element is the layer name consisting
        of abbreviation and postfix, e.g., bn1, gn. The second element is the
        created norm layer.
    """
    if not isinstance(cfg, dict):
        raise TypeError('cfg must be a dict')
    if 'type' not in cfg:
        raise KeyError('the cfg dict must contain the key "type"')
    cfg_ = cfg.copy()

    layer_type = cfg_.pop('type')

    norm_layer = NORM_LAYERS.get(layer_type)
    abbr = infer_abbr(norm_layer)

    assert isinstance(postfix, (int, str))
    name = abbr + str(postfix)

    cfg_.setdefault('eps', 1e-5)
    if layer_type != 'GN':
        layer = norm_layer(num_features, **cfg_)
    else:
        assert 'num_groups' in cfg_
        layer = norm_layer(num_channels=num_features, **cfg_)

    return name, layer


def is_norm(layer, exclude=None):
    """Check if a layer is a normalization layer.

    Args:
        layer (nn.Module): The layer to be checked.
        exclude (type | tuple[type]): Types to be excluded.

    Returns:
        bool: Whether the layer is a norm layer.
    """
    if exclude is not None:
        if not isinstance(exclude, tuple):
            exclude = (exclude, )
        if not is_tuple_of(exclude, type):
            raise TypeError(
                f'"exclude" must be either None or type or a tuple of types, '
                f'but got {type(exclude)}: {exclude}')

    if exclude and isinstance(layer, exclude):
        return False

    all_norm_bases = (_BatchNorm, _InstanceNorm, nn.GroupNorm, nn.LayerNorm)
    return isinstance(layer, all_norm_bases)

python/jittorseg/models/basic_operators/norm.py
JittorCV/jittorseg
from jittor import nn
from jittorseg.utils.registry import PADDING_LAYERS

PADDING_LAYERS.register_module('zero', module=nn.ZeroPad2d)
PADDING_LAYERS.register_module('reflect', module=nn.ReflectionPad2d)
PADDING_LAYERS.register_module('replicate', module=nn.ReplicationPad2d)


def build_padding_layer(cfg, *args, **kwargs):
    if not isinstance(cfg, dict):
        raise TypeError('cfg must be a dict')
    if 'type' not in cfg:
        raise KeyError('the cfg dict must contain the key "type"')

    cfg_ = cfg.copy()
    padding_type = cfg_.pop('type')
    padding_layer = PADDING_LAYERS.get(padding_type)

    layer = padding_layer(*args, **kwargs, **cfg_)

    return layer

python/jittorseg/models/basic_operators/padding.py
JittorCV/jittorseg
from .segformer_head import SegFormerHead
from .fcn_head import FCNHead
from .psp_head import PSPHead
from .ham_head import LightHamHead
from .uper_head import UPerHead
from .ea_head import EAHead
from .cc_head import CCHead
from .da_head import DAHead
from .aspp_head import ASPPHead
from .sep_aspp_head import DepthwiseSeparableASPPHead
from .point_head import PointHead
from .fpn_head import FPNHead
from .nl_head import NonLocal2d
# from .lraspp_head import LRASPPHead
from .gc_head import GCHead
from .ema_head import EMAHead
from .ann_head import ANNHead

python/jittorseg/models/decode_heads/__init__.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from jittorseg.utils.registry import HEADS
from ..layers import SelfAttentionBlock as _SelfAttentionBlock
from .decode_head import BaseDecodeHead


class PPMConcat(nn.ModuleList):

    def __init__(self, pool_scales=(1, 3, 6, 8)):
        super(PPMConcat, self).__init__(
            [nn.AdaptiveAvgPool2d(pool_scale) for pool_scale in pool_scales])

    def execute(self, feats):
        """execute function."""
        ppm_outs = []
        for ppm in self:
            ppm_out = ppm(feats)
            ppm_outs.append(ppm_out.view(*feats.shape[:2], -1))
        concat_outs = jt.concat(ppm_outs, dim=2)
        return concat_outs


class SelfAttentionBlock(_SelfAttentionBlock):

    def __init__(self, low_in_channels, high_in_channels, channels,
                 out_channels, share_key_query, query_scale, key_pool_scales,
                 conv_cfg, norm_cfg, act_cfg):
        key_psp = PPMConcat(key_pool_scales)
        if query_scale > 1:
            query_downsample = nn.MaxPool2d(kernel_size=query_scale)
        else:
            query_downsample = None
        super(SelfAttentionBlock,
              self).__init__(key_in_channels=low_in_channels,
                             query_in_channels=high_in_channels,
                             channels=channels,
                             out_channels=out_channels,
                             share_key_query=share_key_query,
                             query_downsample=query_downsample,
                             key_downsample=key_psp,
                             key_query_num_convs=1,
                             key_query_norm=True,
                             value_out_num_convs=1,
                             value_out_norm=False,
                             matmul_norm=True,
                             with_out=True,
                             conv_cfg=conv_cfg,
                             norm_cfg=norm_cfg,
                             act_cfg=act_cfg)


class AFNB(nn.Module):

    def __init__(self, low_in_channels, high_in_channels, channels,
                 out_channels, query_scales, key_pool_scales, conv_cfg,
                 norm_cfg, act_cfg):
        super(AFNB, self).__init__()
        self.stages = nn.ModuleList()
        for query_scale in query_scales:
            self.stages.append(
                SelfAttentionBlock(low_in_channels=low_in_channels,
                                   high_in_channels=high_in_channels,
                                   channels=channels,
                                   out_channels=out_channels,
                                   share_key_query=False,
                                   query_scale=query_scale,
                                   key_pool_scales=key_pool_scales,
                                   conv_cfg=conv_cfg,
                                   norm_cfg=norm_cfg,
                                   act_cfg=act_cfg))
        self.bottleneck = ConvModule(out_channels + high_in_channels,
                                     out_channels,
                                     1,
                                     conv_cfg=conv_cfg,
                                     norm_cfg=norm_cfg,
                                     act_cfg=None)

    def execute(self, low_feats, high_feats):
        """execute function."""
        priors = [stage(high_feats, low_feats) for stage in self.stages]
        context = jt.stack(priors, dim=0).sum(dim=0)
        output = self.bottleneck(jt.concat([context, high_feats], 1))
        return output


class APNB(nn.Module):

    def __init__(self, in_channels, channels, out_channels, query_scales,
                 key_pool_scales, conv_cfg, norm_cfg, act_cfg):
        super(APNB, self).__init__()
        self.stages = nn.ModuleList()
        for query_scale in query_scales:
            self.stages.append(
                SelfAttentionBlock(low_in_channels=in_channels,
                                   high_in_channels=in_channels,
                                   channels=channels,
                                   out_channels=out_channels,
                                   share_key_query=True,
                                   query_scale=query_scale,
                                   key_pool_scales=key_pool_scales,
                                   conv_cfg=conv_cfg,
                                   norm_cfg=norm_cfg,
                                   act_cfg=act_cfg))
        self.bottleneck = ConvModule(2 * in_channels,
                                     out_channels,
                                     1,
                                     conv_cfg=conv_cfg,
                                     norm_cfg=norm_cfg,
                                     act_cfg=act_cfg)

    def execute(self, feats):
        """execute function."""
        priors = [stage(feats, feats) for stage in self.stages]
        context = jt.stack(priors, dim=0).sum(dim=0)
        output = self.bottleneck(jt.concat([context, feats], 1))
        return output


@HEADS.register_module()
class ANNHead(BaseDecodeHead):

    def __init__(self,
                 project_channels,
                 query_scales=(1, ),
                 key_pool_scales=(1, 3, 6, 8),
                 **kwargs):
        super(ANNHead, self).__init__(input_transform='multiple_select',
                                      **kwargs)
        assert len(self.in_channels) == 2
        low_in_channels, high_in_channels = self.in_channels
        self.project_channels = project_channels
        self.fusion = AFNB(low_in_channels=low_in_channels,
                           high_in_channels=high_in_channels,
                           out_channels=high_in_channels,
                           channels=project_channels,
                           query_scales=query_scales,
                           key_pool_scales=key_pool_scales,
                           conv_cfg=self.conv_cfg,
                           norm_cfg=self.norm_cfg,
                           act_cfg=self.act_cfg)
        self.bottleneck = ConvModule(high_in_channels,
                                     self.channels,
                                     3,
                                     padding=1,
                                     conv_cfg=self.conv_cfg,
                                     norm_cfg=self.norm_cfg,
                                     act_cfg=self.act_cfg)
        self.context = APNB(in_channels=self.channels,
                            out_channels=self.channels,
                            channels=project_channels,
                            query_scales=query_scales,
                            key_pool_scales=key_pool_scales,
                            conv_cfg=self.conv_cfg,
                            norm_cfg=self.norm_cfg,
                            act_cfg=self.act_cfg)

    def execute(self, inputs):
        """execute function."""
        low_feats, high_feats = self._transform_inputs(inputs)
        output = self.fusion(low_feats, high_feats)
        output = self.dropout(output)
        output = self.bottleneck(output)
        output = self.context(output)
        output = self.cls_seg(output)

        return output

python/jittorseg/models/decode_heads/ann_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn

from ..basic_operators import ConvModule
from ..layers import resize

from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead


class ASPPModule(nn.ModuleList):
    """Atrous Spatial Pyramid Pooling (ASPP) Module.
    Args:
        dilations (tuple[int]): Dilation rate of each layer.
        in_channels (int): Input channels.
        channels (int): Channels after modules, before conv_seg.
    """

    def __init__(self, dilations, in_channels, channels, conv_cfg, norm_cfg,
                 act_cfg):
        super(ASPPModule, self).__init__()
        self.dilations = dilations
        self.in_channels = in_channels
        self.channels = channels
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        for dilation in dilations:
            self.append(
                ConvModule(self.in_channels,
                           self.channels,
                           1 if dilation == 1 else 3,
                           dilation=dilation,
                           padding=0 if dilation == 1 else dilation,
                           conv_cfg=self.conv_cfg,
                           norm_cfg=self.norm_cfg,
                           act_cfg=self.act_cfg))

    def execute(self, x):
        aspp_outs = []
        for aspp_module in self:
            aspp_outs.append(aspp_module(x))

        return aspp_outs


@HEADS.register_module()
class ASPPHead(BaseDecodeHead):
    """Rethinking Atrous Convolution for Semantic Image Segmentation.
    This head is the implementation of `DeepLabV3
    <https://arxiv.org/abs/1706.05587>`_.
    Args:
        dilations (tuple[int]): Dilation rates for ASPP module.
            Default: (1, 6, 12, 18).
    """

    def __init__(self, dilations=(1, 6, 12, 18), **kwargs):
        super(ASPPHead, self).__init__(**kwargs)
        assert isinstance(dilations, (list, tuple))
        self.dilations = dilations
        self.image_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            ConvModule(self.in_channels,
                       self.channels,
                       1,
                       conv_cfg=self.conv_cfg,
                       norm_cfg=self.norm_cfg,
                       act_cfg=self.act_cfg))
        self.aspp_modules = ASPPModule(dilations,
                                       self.in_channels,
                                       self.channels,
                                       conv_cfg=self.conv_cfg,
                                       norm_cfg=self.norm_cfg,
                                       act_cfg=self.act_cfg)
        self.bottleneck = ConvModule((len(dilations) + 1) * self.channels,
                                     self.channels,
                                     3,
                                     padding=1,
                                     conv_cfg=self.conv_cfg,
                                     norm_cfg=self.norm_cfg,
                                     act_cfg=self.act_cfg)

    def _execute_feature(self, inputs):
        x = self._transform_inputs(inputs)
        aspp_outs = [
            resize(self.image_pool(x),
                   size=x.size()[2:],
                   mode='bilinear',
                   align_corners=self.align_corners)
        ]
        aspp_outs.extend(self.aspp_modules(x))
        aspp_outs = jt.concat(aspp_outs, dim=1)
        feats = self.bottleneck(aspp_outs)
        return feats

    def execute(self, inputs):
        """Forward function."""
        output = self._execute_feature(inputs)
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/aspp_head.py
JittorCV/jittorseg
from abc import ABCMeta, abstractmethod

from .decode_head import BaseDecodeHead


class BaseCascadeDecodeHead(BaseDecodeHead, metaclass=ABCMeta):
    def __init__(self, *args, **kwargs):
        super(BaseCascadeDecodeHead, self).__init__(*args, **kwargs)

    @abstractmethod
    def execute(self, inputs, prev_output):
        pass

    def execute_train(self, inputs, prev_output, img_metas, gt_semantic_seg,
                      train_cfg):
        seg_logits = self.execute(inputs, prev_output)
        losses = self.losses(seg_logits, gt_semantic_seg)

        return losses

    def execute_test(self, inputs, prev_output, img_metas, test_cfg):
        return self.execute(inputs, prev_output)

python/jittorseg/models/decode_heads/cascade_decode_head.py
JittorCV/jittorseg
import jittor as jt
from .fcn_head import FCNHead
from jittorseg.utils.registry import HEADS

try:
    from ..layers import CrissCrossAttention
except ModuleNotFoundError:
    CrissCrossAttention = None


@HEADS.register_module()
class CCHead(FCNHead):
    def __init__(self, recurrence=2, **kwargs):
        if CrissCrossAttention is None:
            raise RuntimeError('Please install mmcv-full for '
                               'CrissCrossAttention layers')
        super(CCHead, self).__init__(num_convs=2, **kwargs)
        self.recurrence = recurrence
        self.cca = CrissCrossAttention(self.channels)

    def execute(self, inputs):
        """Forward function."""
        x = self._transform_inputs(inputs)
        output = self.convs[0](x)
        for _ in range(self.recurrence):
            output = self.cca(output)
        output = self.convs[1](output)
        if self.concat_input:
            output = self.conv_cat(jt.concat([x, output], dim=1))
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/cc_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from jittorseg.utils.general import add_prefix
from jittorseg.utils.registry import HEADS
from ..layers import Scale, SelfAttentionBlock
from .decode_head import BaseDecodeHead


class PAM(SelfAttentionBlock):
    def __init__(self, in_channels, channels):
        super(PAM, self).__init__(
            key_in_channels=in_channels,
            query_in_channels=in_channels,
            channels=channels,
            out_channels=in_channels,
            share_key_query=False,
            query_downsample=None,
            key_downsample=None,
            key_query_num_convs=1,
            key_query_norm=False,
            value_out_num_convs=1,
            value_out_norm=False,
            matmul_norm=False,
            with_out=False,
            conv_cfg=None,
            norm_cfg=None,
            act_cfg=None)

        self.gamma = Scale(0)

    def execute(self, x):
        out = super(PAM, self).execute(x, x)

        out = self.gamma(out) + x
        return out


class CAM(nn.Module):
    """Channel Attention Module (CAM)"""

    def __init__(self):
        super(CAM, self).__init__()
        self.gamma = Scale(0)

    def execute(self, x):
        batch_size, channels, height, width = x.size()
        proj_query = x.view(batch_size, channels, -1)
        proj_key = x.view(batch_size, channels, -1).permute(0, 2, 1)
        energy = nn.bmm(proj_query, proj_key)
        energy_new = jt.max(
            energy, -1, keepdims=True)[0].expand_as(energy) - energy
        attention = nn.softmax(energy_new, dim=-1)
        proj_value = x.view(batch_size, channels, -1)

        out = nn.bmm(attention, proj_value)
        out = out.view(batch_size, channels, height, width)

        out = self.gamma(out) + x
        return out


@HEADS.register_module()
class DAHead(BaseDecodeHead):
    def __init__(self, pam_channels, **kwargs):
        super(DAHead, self).__init__(**kwargs)
        self.pam_channels = pam_channels
        self.pam_in_conv = ConvModule(
            self.in_channels,
            self.channels,
            3,
            padding=1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.pam = PAM(self.channels, pam_channels)
        self.pam_out_conv = ConvModule(
            self.channels,
            self.channels,
            3,
            padding=1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.pam_conv_seg = nn.Conv2d(
            self.channels, self.num_classes, kernel_size=1)

        self.cam_in_conv = ConvModule(
            self.in_channels,
            self.channels,
            3,
            padding=1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.cam = CAM()
        self.cam_out_conv = ConvModule(
            self.channels,
            self.channels,
            3,
            padding=1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.cam_conv_seg = nn.Conv2d(
            self.channels, self.num_classes, kernel_size=1)

    def pam_cls_seg(self, feat):
        """PAM feature classification."""
        if self.dropout is not None:
            feat = self.dropout(feat)
        output = self.pam_conv_seg(feat)
        return output

    def cam_cls_seg(self, feat):
        """CAM feature classification."""
        if self.dropout is not None:
            feat = self.dropout(feat)
        output = self.cam_conv_seg(feat)
        return output

    def execute(self, inputs):
        """Forward function."""
        x = self._transform_inputs(inputs)
        pam_feat = self.pam_in_conv(x)
        pam_feat = self.pam(pam_feat)
        pam_feat = self.pam_out_conv(pam_feat)
        pam_out = self.pam_cls_seg(pam_feat)

        cam_feat = self.cam_in_conv(x)
        cam_feat = self.cam(cam_feat)
        cam_feat = self.cam_out_conv(cam_feat)
        cam_out = self.cam_cls_seg(cam_feat)

        feat_sum = pam_feat + cam_feat
        pam_cam_out = self.cls_seg(feat_sum)

        return pam_cam_out, pam_out, cam_out

    def execute_test(self, inputs, img_metas, test_cfg):
        return self.execute(inputs)[0]

    def losses(self, seg_logit, seg_label):
        """Compute ``pam_cam``, ``pam``, ``cam`` loss."""
        pam_cam_seg_logit, pam_seg_logit, cam_seg_logit = seg_logit
        loss = dict()
        loss.update(
            add_prefix(
                super(DAHead, self).losses(pam_cam_seg_logit, seg_label),
                'pam_cam'))
        loss.update(
            add_prefix(
                super(DAHead, self).losses(pam_seg_logit, seg_label), 'pam'))
        loss.update(
            add_prefix(
                super(DAHead, self).losses(cam_seg_logit, seg_label), 'cam'))
        return loss

python/jittorseg/models/decode_heads/da_head.py
JittorCV/jittorseg
import warnings
import jittor as jt
from jittor import nn
from ..layers import resize
from jittorseg.utils.registry import build_from_cfg, LOSSES, PIXEL_SAMPLERS
from abc import ABCMeta, abstractmethod

from ..losses import accuracy


class BaseDecodeHead(nn.Module, metaclass=ABCMeta):

    def __init__(self,
                 in_channels,
                 channels,
                 *,
                 num_classes,
                 out_channels=None,
                 threshold=None,
                 dropout_ratio=0.1,
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=dict(type='ReLU'),
                 in_index=-1,
                 input_transform=None,
                 decoder_params=None,
                 loss_decode=dict(type='CrossEntropyLoss',
                                  use_sigmoid=False,
                                  loss_weight=1.0),
                 ignore_index=255,
                 sampler=None,
                 align_corners=False):
        super(BaseDecodeHead, self).__init__()
        self._init_inputs(in_channels, in_index, input_transform)
        self.channels = channels
        self.dropout_ratio = dropout_ratio
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.in_index = in_index

        self.ignore_index = ignore_index
        self.align_corners = align_corners

        if out_channels is None:
            if num_classes == 2:
                warnings.warn('For binary segmentation, we suggest using'
                              '`out_channels = 1` to define the output'
                              'channels of segmentor, and use `threshold`'
                              'to convert seg_logist into a prediction'
                              'applying a threshold')
            out_channels = num_classes

        if out_channels != num_classes and out_channels != 1:
            raise ValueError(
                'out_channels should be equal to num_classes,'
                'except binary segmentation set out_channels == 1 and'
                f'num_classes == 2, but got out_channels={out_channels}'
                f'and num_classes={num_classes}')

        if out_channels == 1 and threshold is None:
            threshold = 0.3
            warnings.warn('threshold is not defined for binary, and defaults'
                          'to 0.3')
        self.num_classes = num_classes
        self.out_channels = out_channels
        self.threshold = threshold

        if isinstance(loss_decode, dict):
            self.loss_decode = build_from_cfg(loss_decode, LOSSES)
        elif isinstance(loss_decode, (list, tuple)):
            self.loss_decode = nn.ModuleList()
            for loss in loss_decode:
                self.loss_decode.append(build_from_cfg(loss, LOSSES))
        else:
            raise TypeError(f'loss_decode must be a dict or sequence of dict,\
                but got {type(loss_decode)}')

        if sampler is not None:
            self.sampler = build_from_cfg(sampler, PIXEL_SAMPLERS, context=self)
        else:
            self.sampler = None

        self.conv_seg = nn.Conv2d(channels, self.out_channels, kernel_size=1)
        if dropout_ratio > 0:
            self.dropout = nn.Dropout(dropout_ratio)
        else:
            self.dropout = None

    def extra_repr(self):
        s = f'input_transform={self.input_transform}, ' \
            f'ignore_index={self.ignore_index}, ' \
            f'align_corners={self.align_corners}'
        return s

    def _init_inputs(self, in_channels, in_index, input_transform):
        if input_transform is not None:
            assert input_transform in ['resize_concat', 'multiple_select']
        self.input_transform = input_transform
        self.in_index = in_index
        if input_transform is not None:
            assert isinstance(in_channels, (list, tuple))
            assert isinstance(in_index, (list, tuple))
            assert len(in_channels) == len(in_index)
            if input_transform == 'resize_concat':
                self.in_channels = sum(in_channels)
            else:
                self.in_channels = in_channels
        else:
            assert isinstance(in_channels, int)
            assert isinstance(in_index, int)
            self.in_channels = in_channels

    def init_weights(self):
        pass

    def _transform_inputs(self, inputs):
        if self.input_transform == 'resize_concat':
            inputs = [inputs[i] for i in self.in_index]
            upsampled_inputs = [
                resize(input=x,
                       size=inputs[0].shape[2:],
                       mode='bilinear',
                       align_corners=self.align_corners) for x in inputs
            ]
            inputs = jt.concat(upsampled_inputs, dim=1)
        elif self.input_transform == 'multiple_select':
            inputs = [inputs[i] for i in self.in_index]
        else:
            inputs = inputs[self.in_index]

        return inputs

    @abstractmethod
    def execute(self, inputs):
        pass

    def execute_train(self, inputs, img_metas, gt_semantic_seg, train_cfg):
        seg_logits = self.execute(inputs)
        losses = self.losses(seg_logits, gt_semantic_seg)
        return losses

    def execute_test(self, inputs, img_metas, test_cfg):
        return self.execute(inputs)

    def cls_seg(self, feat):
        """Classify each pixel."""
        if self.dropout is not None:
            feat = self.dropout(feat)
        output = self.conv_seg(feat)
        return output

    def losses(self, seg_logit, seg_label):
        loss = dict()
        seg_logit = resize(input=seg_logit,
                           size=seg_label.shape[2:],
                           mode='bilinear',
                           align_corners=self.align_corners)
        if self.sampler is not None:
            seg_weight = self.sampler.sample(seg_logit, seg_label)
        else:
            seg_weight = None
        seg_label = seg_label.squeeze(1)

        if not isinstance(self.loss_decode, nn.ModuleList):
            losses_decode = [self.loss_decode]
        else:
            losses_decode = self.loss_decode
        for loss_decode in losses_decode:
            if loss_decode.loss_name not in loss:
                loss[loss_decode.loss_name] = loss_decode(
                    seg_logit,
                    seg_label,
                    weight=seg_weight,
                    ignore_index=self.ignore_index)
            else:
                loss[loss_decode.loss_name] += loss_decode(
                    seg_logit,
                    seg_label,
                    weight=seg_weight,
                    ignore_index=self.ignore_index)
        # TODO ignore_index
        loss['acc_seg'] = accuracy(seg_logit, seg_label)
        return loss

python/jittorseg/models/decode_heads/decode_head.py
JittorCV/jittorseg
from .decode_head import BaseDecodeHead
from jittorseg.utils.registry import HEADS
from ..layers import External_attention


@HEADS.register_module()
class EAHead(BaseDecodeHead):

    def __init__(self, **kwargs):
        super(EAHead, self).__init__(**kwargs)
        self.ea = External_attention(self.in_channels, self.channels)

    def execute(self, inputs):
        x = self._transform_inputs(inputs)
        x = self.ea(x)
        output = self.cls_seg(x)
        return output

python/jittorseg/models/decode_heads/ea_head.py
JittorCV/jittorseg
import math

import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead


class EMAModule(nn.Module):

    def __init__(self, channels, num_bases, num_stages, momentum):
        super(EMAModule, self).__init__()
        assert num_stages >= 1, 'num_stages must be at least 1!'
        self.num_bases = num_bases
        self.num_stages = num_stages
        self.momentum = momentum

        bases = jt.normal(0,
                          math.sqrt(2. / self.num_bases),
                          size=(1, channels, self.num_bases))
        # [1, channels, num_bases]
        bases = jt.normalize(bases, dim=1, p=2)
        self.bases = bases

    def execute(self, feats):
        batch_size, channels, height, width = feats.size()
        # [batch_size, channels, height*width]
        feats = feats.view(batch_size, channels, height * width)
        # [batch_size, channels, num_bases]
        bases = self.bases.repeat(batch_size, 1, 1)

        with jt.no_grad():
            for i in range(self.num_stages):
                # [batch_size, height*width, num_bases]
                attention = jt.linalg.einsum('bcn,bck->bnk', feats, bases)
                attention = nn.softmax(attention, dim=2)
                # l1 norm
                attention_normed = jt.normalize(attention, dim=1, p=1)
                # [batch_size, channels, num_bases]
                bases = jt.linalg.einsum('bcn,bnk->bck', feats,
                                         attention_normed)
                # l2 norm
                bases = jt.normalize(bases, dim=1, p=2)

        feats_recon = jt.linalg.einsum('bck,bnk->bcn', bases, attention)
        feats_recon = feats_recon.view(batch_size, channels, height, width)

        if self.train():
            bases = bases.mean(dim=0, keepdim=True)
            # l2 norm
            bases = jt.normalize(bases, dim=1, p=2)
            self.bases = (1 -
                          self.momentum) * self.bases + self.momentum * bases

        return feats_recon


@HEADS.register_module()
class EMAHead(BaseDecodeHead):

    def __init__(self,
                 ema_channels,
                 num_bases,
                 num_stages,
                 concat_input=True,
                 momentum=0.1,
                 **kwargs):
        super(EMAHead, self).__init__(**kwargs)
        self.ema_channels = ema_channels
        self.num_bases = num_bases
        self.num_stages = num_stages
        self.concat_input = concat_input
        self.momentum = momentum
        self.ema_module = EMAModule(self.ema_channels, self.num_bases,
                                    self.num_stages, self.momentum)

        self.ema_in_conv = ConvModule(self.in_channels,
                                      self.ema_channels,
                                      3,
                                      padding=1,
                                      conv_cfg=self.conv_cfg,
                                      norm_cfg=self.norm_cfg,
                                      act_cfg=self.act_cfg)
        # project (0, inf) -> (-inf, inf)
        self.ema_mid_conv = ConvModule(self.ema_channels,
                                       self.ema_channels,
                                       1,
                                       conv_cfg=self.conv_cfg,
                                       norm_cfg=None,
                                       act_cfg=None)
        for param in self.ema_mid_conv.parameters():
            param.requires_grad = False

        self.ema_out_conv = ConvModule(self.ema_channels,
                                       self.ema_channels,
                                       1,
                                       conv_cfg=self.conv_cfg,
                                       norm_cfg=self.norm_cfg,
                                       act_cfg=None)
        self.bottleneck = ConvModule(self.ema_channels,
                                     self.channels,
                                     3,
                                     padding=1,
                                     conv_cfg=self.conv_cfg,
                                     norm_cfg=self.norm_cfg,
                                     act_cfg=self.act_cfg)
        if self.concat_input:
            self.conv_cat = ConvModule(self.in_channels + self.channels,
                                       self.channels,
                                       kernel_size=3,
                                       padding=1,
                                       conv_cfg=self.conv_cfg,
                                       norm_cfg=self.norm_cfg,
                                       act_cfg=self.act_cfg)

    def execute(self, inputs):
        x = self._transform_inputs(inputs)
        feats = self.ema_in_conv(x)
        identity = feats
        feats = self.ema_mid_conv(feats)
        recon = self.ema_module(feats)
        recon = nn.relu(recon)
        recon = self.ema_out_conv(recon)
        output = nn.relu(identity + recon)
        output = self.bottleneck(output)
        if self.concat_input:
            output = self.conv_cat(jt.concat([x, output], dim=1))
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/ema_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead


@HEADS.register_module()
class FCNHead(BaseDecodeHead):
    def __init__(self,
                 num_convs=2,
                 kernel_size=3,
                 concat_input=True,
                 dilation=1,
                 **kwargs):
        assert num_convs >= 0 and dilation > 0 and isinstance(dilation, int)
        self.num_convs = num_convs
        self.concat_input = concat_input
        self.kernel_size = kernel_size
        super(FCNHead, self).__init__(**kwargs)
        if num_convs == 0:
            assert self.in_channels == self.channels

        conv_padding = (kernel_size // 2) * dilation
        convs = []
        convs.append(
            ConvModule(
                self.in_channels,
                self.channels,
                kernel_size=kernel_size,
                padding=conv_padding,
                dilation=dilation,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg))
        for i in range(num_convs - 1):
            convs.append(
                ConvModule(
                    self.channels,
                    self.channels,
                    kernel_size=kernel_size,
                    padding=conv_padding,
                    dilation=dilation,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    act_cfg=self.act_cfg))
        if num_convs == 0:
            self.convs = nn.Identity()
        else:
            self.convs = nn.Sequential(*convs)
        if self.concat_input:
            self.conv_cat = ConvModule(
                self.in_channels + self.channels,
                self.channels,
                kernel_size=kernel_size,
                padding=kernel_size // 2,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg)

    def _execute_feature(self, inputs):
        x = self._transform_inputs(inputs)
        feats = self.convs(x)
        if self.concat_input:
            feats = self.conv_cat(jt.concat([x, feats], dim=1))
        return feats

    def execute(self, inputs):
        output = self._execute_feature(inputs)
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/fcn_head.py
JittorCV/jittorseg
import numpy as np
from jittor import nn
from ..basic_operators import ConvModule

from ..layers import Upsample, resize
from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead


@HEADS.register_module()
class FPNHead(BaseDecodeHead):

    def __init__(self, feature_strides, **kwargs):
        super(FPNHead, self).__init__(input_transform='multiple_select',
                                      **kwargs)
        assert len(feature_strides) == len(self.in_channels)
        assert min(feature_strides) == feature_strides[0]
        self.feature_strides = feature_strides

        self.scale_heads = nn.ModuleList()
        for i in range(len(feature_strides)):
            head_length = max(
                1,
                int(np.log2(feature_strides[i]) - np.log2(feature_strides[0])))
            scale_head = []
            for k in range(head_length):
                scale_head.append(
                    ConvModule(
                        self.in_channels[i] if k == 0 else self.channels,
                        self.channels,
                        3,
                        padding=1,
                        conv_cfg=self.conv_cfg,
                        norm_cfg=self.norm_cfg,
                        act_cfg=self.act_cfg))
                if feature_strides[i] != feature_strides[0]:
                    scale_head.append(
                        Upsample(scale_factor=2,
                                 mode='bilinear',
                                 align_corners=self.align_corners))
            self.scale_heads.append(nn.Sequential(*scale_head))

    def execute(self, inputs):

        x = self._transform_inputs(inputs)

        output = self.scale_heads[0](x[0])
        for i in range(1, len(self.feature_strides)):
            # non inplace
            output = output + resize(self.scale_heads[i](x[i]),
                                     size=output.shape[2:],
                                     mode='bilinear',
                                     align_corners=self.align_corners)

        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/fpn_head.py
JittorCV/jittorseg
import jittor as jt
from ..layers import ContextBlock

from jittorseg.utils.registry import HEADS
from .fcn_head import FCNHead


@HEADS.register_module()
class GCHead(FCNHead):

    def __init__(self,
                 ratio=1 / 4.,
                 pooling_type='att',
                 fusion_types=('channel_add', ),
                 **kwargs):
        super(GCHead, self).__init__(num_convs=2, **kwargs)
        self.ratio = ratio
        self.pooling_type = pooling_type
        self.fusion_types = fusion_types
        self.gc_block = ContextBlock(in_channels=self.channels,
                                     ratio=self.ratio,
                                     pooling_type=self.pooling_type,
                                     fusion_types=self.fusion_types)

    def execute(self, inputs):
        x = self._transform_inputs(inputs)
        output = self.convs[0](x)
        output = self.gc_block(output)
        output = self.convs[1](output)
        if self.concat_input:
            output = self.conv_cat(jt.concat([x, output], dim=1))
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/gc_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..layers import resize
from jittorseg.utils.registry import HEADS
from ..basic_operators import ConvModule
from .decode_head import BaseDecodeHead


class _MatrixDecomposition2DBase(nn.Module):

    def __init__(self, args=dict()):
        super().__init__()

        self.spatial = args.setdefault('SPATIAL', True)

        self.S = args.setdefault('MD_S', 1)
        self.D = args.setdefault('MD_D', 512)
        self.R = args.setdefault('MD_R', 64)

        self.train_steps = args.setdefault('TRAIN_STEPS', 6)
        self.eval_steps = args.setdefault('EVAL_STEPS', 7)

        self.inv_t = args.setdefault('INV_T', 100)
        self.eta = args.setdefault('ETA', 0.9)

        self.rand_init = args.setdefault('RAND_INIT', True)

        print('spatial', self.spatial)
        print('S', self.S)
        print('D', self.D)
        print('R', self.R)
        print('train_steps', self.train_steps)
        print('eval_steps', self.eval_steps)
        print('inv_t', self.inv_t)
        print('eta', self.eta)
        print('rand_init', self.rand_init)

    def _build_bases(self, B, S, D, R):
        raise NotImplementedError

    def local_step(self, x, bases, coef):
        raise NotImplementedError

    @jt.no_grad()
    def local_inference(self, x, bases):
        # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
        coef = nn.bmm(x.transpose(1, 2), bases)
        coef = nn.softmax(self.inv_t * coef, dim=-1)

        steps = self.train_steps if self.is_training() else self.eval_steps
        for _ in range(steps):
            bases, coef = self.local_step(x, bases, coef)

        return bases, coef

    def compute_coef(self, x, bases, coef):
        raise NotImplementedError

    def execute(self, x, return_bases=False):
        B, C, H, W = x.shape

        # (B, C, H, W) -> (B * S, D, N)
        if self.spatial:
            D = C // self.S
            N = H * W
            x = x.view(B * self.S, D, N)
        else:
            D = H * W
            N = C // self.S
            x = x.view(B * self.S, N, D).transpose(1, 2)

        if not self.rand_init and not hasattr(self, 'bases'):
            bases = self._build_bases(1, self.S, D, self.R)
            self.register_buffer('bases', bases)

        # (S, D, R) -> (B * S, D, R)
        if self.rand_init:
            bases = self._build_bases(B, self.S, D, self.R)
        else:
            bases = self.bases.repeat(B, 1, 1)

        bases, coef = self.local_inference(x, bases)

        # (B * S, N, R)
        coef = self.compute_coef(x, bases, coef)

        # (B * S, D, R) @ (B * S, N, R)^T -> (B * S, D, N)
        x = nn.bmm(bases, coef.transpose(1, 2))

        # (B * S, D, N) -> (B, C, H, W)
        if self.spatial:
            x = x.view(B, C, H, W)
        else:
            x = x.transpose(1, 2).view(B, C, H, W)

        # (B * H, D, R) -> (B, H, N, D)
        bases = bases.view(B, self.S, D, self.R)

        return x


class NMF2D(_MatrixDecomposition2DBase):

    def __init__(self, args=dict()):
        super().__init__(args)

        self.inv_t = 1

    def _build_bases(self, B, S, D, R):
        bases = jt.rand((B * S, D, R))

        bases = jt.normalize(bases, dim=1)

        return bases

    def local_step(self, x, bases, coef):
        # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
        numerator = nn.bmm(x.transpose(1, 2), bases)
        # (B * S, N, R) @ [(B * S, D, R)^T @ (B * S, D, R)] -> (B * S, N, R)
        denominator = nn.bmm(coef, nn.bmm(bases.transpose(1, 2), bases))
        # Multiplicative Update
        coef = coef * numerator / (denominator + 1e-6)

        # (B * S, D, N) @ (B * S, N, R) -> (B * S, D, R)
        numerator = nn.bmm(x, coef)
        # (B * S, D, R) @ [(B * S, N, R)^T @ (B * S, N, R)] -> (B * S, D, R)
        denominator = nn.bmm(bases, nn.bmm(coef.transpose(1, 2), coef))
        # Multiplicative Update
        bases = bases * numerator / (denominator + 1e-6)

        return bases, coef

    def compute_coef(self, x, bases, coef):
        # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
        numerator = nn.bmm(x.transpose(1, 2), bases)
        # (B * S, N, R) @ (B * S, D, R)^T @ (B * S, D, R) -> (B * S, N, R)
        denominator = nn.bmm(coef, nn.bmm(bases.transpose(1, 2), bases))
        # multiplication update
        coef = coef * numerator / (denominator + 1e-6)

        return coef


class Hamburger(nn.Module):

    def __init__(self,
                 ham_channels=512,
                 ham_kwargs=dict(),
                 norm_cfg=None,
                 **kwargs):
        super().__init__()

        self.ham_in = ConvModule(ham_channels,
                                 ham_channels,
                                 1,
                                 norm_cfg=None,
                                 act_cfg=None)

        self.ham = NMF2D(ham_kwargs)

        self.ham_out = ConvModule(ham_channels,
                                  ham_channels,
                                  1,
                                  norm_cfg=norm_cfg,
                                  act_cfg=None)

    def execute(self, x):
        enjoy = self.ham_in(x)
        enjoy = nn.relu(enjoy)
        enjoy = self.ham(enjoy)
        enjoy = self.ham_out(enjoy)
        ham = nn.relu(x + enjoy)

        return ham


@HEADS.register_module()
class LightHamHead(BaseDecodeHead):

    def __init__(self, ham_channels=512, ham_kwargs=dict(), **kwargs):
        super(LightHamHead, self).__init__(input_transform='multiple_select',
                                           **kwargs)
        self.ham_channels = ham_channels

        self.squeeze = ConvModule(sum(self.in_channels),
                                  self.ham_channels,
                                  1,
                                  conv_cfg=self.conv_cfg,
                                  norm_cfg=self.norm_cfg,
                                  act_cfg=self.act_cfg)

        self.hamburger = Hamburger(ham_channels, ham_kwargs, **kwargs)

        self.align = ConvModule(self.ham_channels,
                                self.channels,
                                1,
                                conv_cfg=self.conv_cfg,
                                norm_cfg=self.norm_cfg,
                                act_cfg=self.act_cfg)

    def execute(self, inputs):
        inputs = self._transform_inputs(inputs)

        inputs = [
            resize(level,
                   size=inputs[0].shape[2:],
                   mode='bilinear',
                   align_corners=self.align_corners) for level in inputs
        ]

        inputs = jt.concat(inputs, dim=1)
        x = self.squeeze(inputs)

        x = self.hamburger(x)

        output = self.align(x)
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/ham_head.py
JittorCV/jittorseg
import jittor as jt
from ..layers import NonLocal2d

from jittorseg.utils.registry import HEADS
from .fcn_head import FCNHead


@HEADS.register_module()
class NLHead(FCNHead):
    def __init__(self,
                 reduction=2,
                 use_scale=True,
                 mode='embedded_gaussian',
                 **kwargs):
        super(NLHead, self).__init__(num_convs=2, **kwargs)
        self.reduction = reduction
        self.use_scale = use_scale
        self.mode = mode
        self.nl_block = NonLocal2d(in_channels=self.channels,
                                   reduction=self.reduction,
                                   use_scale=self.use_scale,
                                   conv_cfg=self.conv_cfg,
                                   norm_cfg=self.norm_cfg,
                                   mode=self.mode)

    def execute(self, inputs):
        """execute function."""
        x = self._transform_inputs(inputs)
        output = self.convs[0](x)
        output = self.nl_block(output)
        output = self.convs[1](output)
        if self.concat_input:
            output = self.conv_cat(jt.concat([x, output], dim=1))
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/nl_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from jittorseg.utils.registry import HEADS
from ..layers import resize
from ..losses import accuracy
from .cascade_decode_head import BaseCascadeDecodeHead


def calculate_uncertainty(seg_logits):
    top2_scores = jt.topk(seg_logits, k=2, dim=1)[0]
    return (top2_scores[:, 1] - top2_scores[:, 0]).unsqueeze(1)


def denormalize(grid):
    return grid * 2.0 - 1.0


def point_sample(input, points, align_corners=False, **kwargs):
    add_dim = False
    if points.ndim == 3:
        add_dim = True
        points = points.unsqueeze(2)
    output = nn.grid_sample(input,
                            denormalize(points),
                            align_corners=align_corners,
                            **kwargs)
    if add_dim:
        output = output.squeeze(3)
    return output


@HEADS.register_module()
class PointHead(BaseCascadeDecodeHead):

    def __init__(self,
                 num_fcs=3,
                 coarse_pred_each_layer=True,
                 conv_cfg=dict(type='Conv1d'),
                 norm_cfg=None,
                 act_cfg=dict(type='ReLU'),
                 **kwargs):
        super(PointHead, self).__init__(input_transform='multiple_select',
                                        conv_cfg=conv_cfg,
                                        norm_cfg=norm_cfg,
                                        act_cfg=act_cfg,
                                        **kwargs)

        self.num_fcs = num_fcs
        self.coarse_pred_each_layer = coarse_pred_each_layer

        fc_in_channels = sum(self.in_channels) + self.num_classes
        fc_channels = self.channels
        self.fcs = nn.ModuleList()
        for k in range(num_fcs):
            fc = ConvModule(fc_in_channels,
                            fc_channels,
                            kernel_size=1,
                            stride=1,
                            padding=0,
                            conv_cfg=conv_cfg,
                            norm_cfg=norm_cfg,
                            act_cfg=act_cfg)
            self.fcs.append(fc)
            fc_in_channels = fc_channels
            fc_in_channels += self.num_classes if self.coarse_pred_each_layer \
                else 0
        self.fc_seg = nn.Conv1d(fc_in_channels,
                                self.num_classes,
                                kernel_size=1,
                                stride=1,
                                padding=0)
        if self.dropout_ratio > 0:
            self.dropout = nn.Dropout(self.dropout_ratio)
        delattr(self, 'conv_seg')

    def cls_seg(self, feat):
        """Classify each pixel with fc."""
        if self.dropout is not None:
            feat = self.dropout(feat)
        output = self.fc_seg(feat)
        return output

    def execute(self, fine_grained_point_feats, coarse_point_feats):
        x = jt.concat([fine_grained_point_feats, coarse_point_feats], dim=1)
        for fc in self.fcs:
            x = fc(x)
            if self.coarse_pred_each_layer:
                x = jt.concat((x, coarse_point_feats), dim=1)
        return self.cls_seg(x)

    def _get_fine_grained_point_feats(self, x, points):
        fine_grained_feats_list = [
            point_sample(_, points, align_corners=self.align_corners)
            for _ in x
        ]
        if len(fine_grained_feats_list) > 1:
            fine_grained_feats = jt.concat(fine_grained_feats_list, dim=1)
        else:
            fine_grained_feats = fine_grained_feats_list[0]

        return fine_grained_feats

    def _get_coarse_point_feats(self, prev_output, points):
        coarse_feats = point_sample(prev_output,
                                    points,
                                    align_corners=self.align_corners)

        return coarse_feats

    def execute_train(self, inputs, prev_output, img_metas, gt_semantic_seg,
                      train_cfg):
        """execute function for training.
        Args:
            inputs (list[Tensor]): List of multi-level img features.
            prev_output (Tensor): The output of previous decode head.
            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.
            train_cfg (dict): The training config.

        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        x = self._transform_inputs(inputs)
        with jt.no_grad():
            points = self.get_points_train(prev_output,
                                           calculate_uncertainty,
                                           cfg=train_cfg)
        fine_grained_point_feats = self._get_fine_grained_point_feats(
            x, points)
        coarse_point_feats = self._get_coarse_point_feats(prev_output, points)
        point_logits = self.execute(fine_grained_point_feats,
                                    coarse_point_feats)
        point_label = point_sample(gt_semantic_seg.float(),
                                   points,
                                   mode='nearest',
                                   align_corners=self.align_corners)
        point_logits = point_logits.unsqueeze(-1)
        point_label = point_label.squeeze(1).unsqueeze(-1).long()

        losses = self.losses(point_logits, point_label)

        return losses

    def execute_test(self, inputs, prev_output, img_metas, test_cfg):
        x = self._transform_inputs(inputs)
        refined_seg_logits = prev_output.clone()
        for _ in range(test_cfg.subdivision_steps):
            refined_seg_logits = resize(refined_seg_logits,
                                        scale_factor=test_cfg.scale_factor,
                                        mode='bilinear',
                                        align_corners=self.align_corners)
            batch_size, channels, height, width = refined_seg_logits.shape
            point_indices, points = self.get_points_test(refined_seg_logits,
                                                         calculate_uncertainty,
                                                         cfg=test_cfg)
            fine_grained_point_feats = self._get_fine_grained_point_feats(
                x, points)
            coarse_point_feats = self._get_coarse_point_feats(
                prev_output, points)
            point_logits = self.execute(fine_grained_point_feats,
                                        coarse_point_feats)

            point_indices = point_indices.unsqueeze(1).expand(-1, channels, -1)
            refined_seg_logits = refined_seg_logits.reshape(
                batch_size, channels, height * width)
            refined_seg_logits = refined_seg_logits.scatter_(
                2, point_indices, point_logits)
            refined_seg_logits = refined_seg_logits.view(
                batch_size, channels, height, width)

        return refined_seg_logits

    def losses(self, point_logits, point_label):
        """Compute segmentation loss."""
        loss = dict()
        if not isinstance(self.loss_decode, nn.ModuleList):
            losses_decode = [self.loss_decode]
        else:
            losses_decode = self.loss_decode
        for loss_module in losses_decode:
            loss['point.' + loss_module.loss_name] = loss_module(
                point_logits, point_label, ignore_index=self.ignore_index)

        loss['acc_point'] = accuracy(point_logits,
                                     point_label,
                                     ignore_index=self.ignore_index)
        return loss

    def get_points_train(self, seg_logits, uncertainty_func, cfg):
        num_points = cfg.num_points
        oversample_ratio = cfg.oversample_ratio
        importance_sample_ratio = cfg.importance_sample_ratio
        assert oversample_ratio >= 1
        assert 0 <= importance_sample_ratio <= 1
        batch_size = seg_logits.shape[0]
        num_sampled = int(num_points * oversample_ratio)
        point_coords = jt.rand((batch_size, num_sampled, 2))
        point_logits = point_sample(seg_logits, point_coords)
        point_uncertainties = uncertainty_func(point_logits)
        num_uncertain_points = int(importance_sample_ratio * num_points)
        num_random_points = num_points - num_uncertain_points
        idx = jt.topk(point_uncertainties[:, 0, :],
                      k=num_uncertain_points,
                      dim=1)[1]
        shift = num_sampled * jt.arange(batch_size, dtype=jt.int64)
        idx += shift[:, None]
        point_coords = point_coords.view(-1, 2)[idx.view(-1), :].view(
            batch_size, num_uncertain_points, 2)
        if num_random_points > 0:
            rand_point_coords = jt.rand((batch_size, num_random_points, 2))
            point_coords = jt.concat((point_coords, rand_point_coords), dim=1)
        return point_coords

    def get_points_test(self, seg_logits, uncertainty_func, cfg):
        num_points = cfg.subdivision_num_points
        uncertainty_map = uncertainty_func(seg_logits)
        batch_size, _, height, width = uncertainty_map.shape
        h_step = 1.0 / height
        w_step = 1.0 / width

        uncertainty_map = uncertainty_map.view(batch_size, height * width)
        num_points = min(height * width, num_points)
        point_indices = uncertainty_map.topk(num_points, dim=1)[1]
        point_coords = jt.zeros((batch_size, num_points, 2)).float()
        point_coords[:, :, 0] = w_step / 2.0 + (point_indices %
                                                width).float() * w_step
        point_coords[:, :, 1] = h_step / 2.0 + (point_indices //
                                                width).float() * h_step
        return point_indices, point_coords

python/jittorseg/models/decode_heads/point_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from ..layers import resize
from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead


class PPM(nn.ModuleList):

    def __init__(self, pool_scales, in_channels, channels, conv_cfg, norm_cfg,
                 act_cfg, align_corners, **kwargs):
        super(PPM, self).__init__()
        self.pool_scales = pool_scales
        self.align_corners = align_corners
        self.in_channels = in_channels
        self.channels = channels
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        for pool_scale in pool_scales:
            self.append(
                nn.Sequential(
                    nn.AdaptiveAvgPool2d(pool_scale),
                    ConvModule(self.in_channels,
                               self.channels,
                               1,
                               conv_cfg=self.conv_cfg,
                               norm_cfg=self.norm_cfg,
                               act_cfg=self.act_cfg,
                               **kwargs)))

    def execute(self, x):
        ppm_outs = []
        for ppm in self:
            ppm_out = ppm(x)
            upsampled_ppm_out = resize(ppm_out,
                                       size=x.size()[2:],
                                       mode='bilinear',
                                       align_corners=self.align_corners)
            ppm_outs.append(upsampled_ppm_out)
        return ppm_outs


@HEADS.register_module()
class PSPHead(BaseDecodeHead):

    def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs):
        super(PSPHead, self).__init__(**kwargs)
        assert isinstance(pool_scales, (list, tuple))
        self.pool_scales = pool_scales
        self.psp_modules = PPM(self.pool_scales,
                               self.in_channels,
                               self.channels,
                               conv_cfg=self.conv_cfg,
                               norm_cfg=self.norm_cfg,
                               act_cfg=self.act_cfg,
                               align_corners=self.align_corners)
        self.bottleneck = ConvModule(self.in_channels +
                                     len(pool_scales) * self.channels,
                                     self.channels,
                                     3,
                                     padding=1,
                                     conv_cfg=self.conv_cfg,
                                     norm_cfg=self.norm_cfg,
                                     act_cfg=self.act_cfg)

    def execute(self, inputs):
        x = self._transform_inputs(inputs)
        psp_outs = [x]
        psp_outs.extend(self.psp_modules(x))
        psp_outs = jt.concat(psp_outs, dim=1)
        output = self.bottleneck(psp_outs)
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/psp_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from ..layers import resize
from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead


class MLP(nn.Module):
    """
    Linear Embedding
    """
    def __init__(self, input_dim=2048, embed_dim=768):
        super().__init__()
        self.proj = nn.Linear(input_dim, embed_dim)

    def execute(self, x):
        x = x.flatten(2).transpose(1, 2)
        x = self.proj(x)
        return x


@HEADS.register_module()
class SegFormerHead(BaseDecodeHead):
    """
    SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers
    """
    def __init__(self, feature_strides, **kwargs):
        super(SegFormerHead, self).__init__(input_transform='multiple_select',
                                            **kwargs)
        assert len(feature_strides) == len(self.in_channels)
        assert min(feature_strides) == feature_strides[0]
        self.feature_strides = feature_strides

        c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels

        decoder_params = kwargs['decoder_params']
        embedding_dim = decoder_params['embed_dim']

        self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim)
        self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim)
        self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim)
        self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim)

        self.linear_fuse = ConvModule(in_channels=embedding_dim * 4,
                                      out_channels=embedding_dim,
                                      kernel_size=1,
                                      norm_cfg=self.norm_cfg)

        self.linear_pred = nn.Conv2d(embedding_dim,
                                     self.num_classes,
                                     kernel_size=1)

    def execute(self, inputs):
        x = self._transform_inputs(inputs)  # len=4, 1/4,1/8,1/16,1/32
        c1, c2, c3, c4 = x

        n, _, h, w = c4.shape

        _c4 = self.linear_c4(c4).permute(0, 2,
                                         1).reshape(n, -1, c4.shape[2],
                                                    c4.shape[3])
        _c4 = resize(_c4,
                     size=c1.size()[2:],
                     mode='bilinear',
                     align_corners=False)

        _c3 = self.linear_c3(c3).permute(0, 2,
                                         1).reshape(n, -1, c3.shape[2],
                                                    c3.shape[3])
        _c3 = resize(_c3,
                     size=c1.size()[2:],
                     mode='bilinear',
                     align_corners=False)

        _c2 = self.linear_c2(c2).permute(0, 2,
                                         1).reshape(n, -1, c2.shape[2],
                                                    c2.shape[3])
        _c2 = resize(_c2,
                     size=c1.size()[2:],
                     mode='bilinear',
                     align_corners=False)

        _c1 = self.linear_c1(c1).permute(0, 2,
                                         1).reshape(n, -1, c1.shape[2],
                                                    c1.shape[3])

        _c = self.linear_fuse(jt.concat([_c4, _c3, _c2, _c1], dim=1))

        x = self.dropout(_c)
        x = self.linear_pred(x)

        return x

python/jittorseg/models/decode_heads/segformer_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..layers import resize
from ..basic_operators import ConvModule, DepthwiseSeparableConvModule
from jittorseg.utils.registry import HEADS
from .aspp_head import ASPPHead


class DepthwiseSeparableASPPModule(nn.ModuleList):
    """Atrous Spatial Pyramid Pooling (ASPP) Module with depthwise separable
    conv."""

    def __init__(self, dilations, in_channels, channels, norm_cfg, act_cfg):
        super(DepthwiseSeparableASPPModule, self).__init__()
        self.dilations = dilations
        self.in_channels = in_channels
        self.channels = channels
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        for i, dilation in enumerate(self.dilations):
            if dilation == 1:
                self.append(
                    ConvModule(self.in_channels,
                               self.channels,
                               1,
                               dilation=dilation,
                               padding=0,
                               norm_cfg=self.norm_cfg,
                               act_cfg=self.act_cfg))
            if dilation > 1:
                self.append(
                    DepthwiseSeparableConvModule(self.in_channels,
                                                 self.channels,
                                                 3,
                                                 dilation=dilation,
                                                 padding=dilation,
                                                 norm_cfg=self.norm_cfg,
                                                 act_cfg=self.act_cfg))

    def execute(self, x):
        aspp_outs = []
        for aspp_module in self:
            aspp_outs.append(aspp_module(x))
        return aspp_outs


@HEADS.register_module()
class DepthwiseSeparableASPPHead(ASPPHead):
    """Encoder-Decoder with Atrous Separable Convolution for Semantic Image
    Segmentation.
    This head is the implementation of `DeepLabV3+
    <https://arxiv.org/abs/1802.02611>`_.
    Args:
        c1_in_channels (int): The input channels of c1 decoder. If is 0,
            the no decoder will be used.
        c1_channels (int): The intermediate channels of c1 decoder.
    """

    def __init__(self, c1_in_channels, c1_channels, **kwargs):
        super(DepthwiseSeparableASPPHead, self).__init__(**kwargs)
        assert c1_in_channels >= 0
        self.aspp_modules = DepthwiseSeparableASPPModule(
            dilations=self.dilations,
            in_channels=self.in_channels,
            channels=self.channels,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        if c1_in_channels > 0:
            self.c1_bottleneck = ConvModule(c1_in_channels,
                                            c1_channels,
                                            1,
                                            conv_cfg=self.conv_cfg,
                                            norm_cfg=self.norm_cfg,
                                            act_cfg=self.act_cfg)
        else:
            self.c1_bottleneck = None
        self.sep_bottleneck = nn.Sequential(
            DepthwiseSeparableConvModule(self.channels + c1_channels,
                                         self.channels,
                                         3,
                                         padding=1,
                                         norm_cfg=self.norm_cfg,
                                         act_cfg=self.act_cfg),
            DepthwiseSeparableConvModule(self.channels,
                                         self.channels,
                                         3,
                                         padding=1,
                                         norm_cfg=self.norm_cfg,
                                         act_cfg=self.act_cfg))

    def execute(self, inputs):
        """execute function."""
        x = self._transform_inputs(inputs)
        aspp_outs = [
            resize(self.image_pool(x),
                   size=x.size()[2:],
                   mode='bilinear',
                   align_corners=self.align_corners)
        ]
        aspp_outs.extend(self.aspp_modules(x))
        aspp_outs = jt.concat(aspp_outs, dim=1)
        output = self.bottleneck(aspp_outs)
        if self.c1_bottleneck is not None:
            c1_output = self.c1_bottleneck(inputs[0])
            output = resize(input=output,
                            size=c1_output.shape[2:],
                            mode='bilinear',
                            align_corners=self.align_corners)
            output = jt.concat([output, c1_output], dim=1)
        output = self.sep_bottleneck(output)
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/sep_aspp_head.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule

from ..layers import resize
from jittorseg.utils.registry import HEADS
from .decode_head import BaseDecodeHead
from .psp_head import PPM


@HEADS.register_module()
class UPerHead(BaseDecodeHead):

    def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs):
        super(UPerHead, self).__init__(input_transform='multiple_select',
                                       **kwargs)
        # PSP Module
        self.psp_modules = PPM(pool_scales,
                               self.in_channels[-1],
                               self.channels,
                               conv_cfg=self.conv_cfg,
                               norm_cfg=self.norm_cfg,
                               act_cfg=self.act_cfg,
                               align_corners=self.align_corners)
        self.bottleneck = ConvModule(self.in_channels[-1] +
                                     len(pool_scales) * self.channels,
                                     self.channels,
                                     3,
                                     padding=1,
                                     conv_cfg=self.conv_cfg,
                                     norm_cfg=self.norm_cfg,
                                     act_cfg=self.act_cfg)
        # FPN Module
        self.lateral_convs = nn.ModuleList()
        self.fpn_convs = nn.ModuleList()
        for in_channels in self.in_channels[:-1]:  # skip the top layer
            l_conv = ConvModule(in_channels,
                                self.channels,
                                1,
                                conv_cfg=self.conv_cfg,
                                norm_cfg=self.norm_cfg,
                                act_cfg=self.act_cfg)
            fpn_conv = ConvModule(self.channels,
                                  self.channels,
                                  3,
                                  padding=1,
                                  conv_cfg=self.conv_cfg,
                                  norm_cfg=self.norm_cfg,
                                  act_cfg=self.act_cfg)
            self.lateral_convs.append(l_conv)
            self.fpn_convs.append(fpn_conv)

        self.fpn_bottleneck = ConvModule(len(self.in_channels) * self.channels,
                                         self.channels,
                                         3,
                                         padding=1,
                                         conv_cfg=self.conv_cfg,
                                         norm_cfg=self.norm_cfg,
                                         act_cfg=self.act_cfg)

    def psp_execute(self, inputs):
        """execute function of PSP module."""
        x = inputs[-1]
        psp_outs = [x]
        psp_outs.extend(self.psp_modules(x))
        psp_outs = jt.concat(psp_outs, dim=1)
        output = self.bottleneck(psp_outs)

        return output

    def _execute_feature(self, inputs):
        inputs = self._transform_inputs(inputs)

        # build laterals
        laterals = [
            lateral_conv(inputs[i])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        laterals.append(self.psp_execute(inputs))

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] = laterals[i - 1] + resize(
                laterals[i],
                size=prev_shape,
                mode='bilinear',
                align_corners=self.align_corners)

        # build outputs
        fpn_outs = [
            self.fpn_convs[i](laterals[i])
            for i in range(used_backbone_levels - 1)
        ]
        # append psp feature
        fpn_outs.append(laterals[-1])

        for i in range(used_backbone_levels - 1, 0, -1):
            fpn_outs[i] = resize(fpn_outs[i],
                                 size=fpn_outs[0].shape[2:],
                                 mode='bilinear',
                                 align_corners=self.align_corners)
        fpn_outs = jt.concat(fpn_outs, dim=1)
        feats = self.fpn_bottleneck(fpn_outs)
        return feats

    def execute(self, inputs):
        """execute function."""
        output = self._execute_feature(inputs)
        output = self.cls_seg(output)
        return output

python/jittorseg/models/decode_heads/uper_head.py
JittorCV/jittorseg
from .wrappers import Upsample, resize
from .external_attention import External_attention
from .cc_attention import CrissCrossAttention
from .scale import Scale
from .self_attention_block import SelfAttentionBlock
from .multi_head_attention import MultiHeadAttention
from .context_block import ContextBlock
from .non_local import NonLocal1d, NonLocal2d, NonLocal3d

__all__ = [
    'Upsample', 'resize', 'External_attention', 'CrissCrossAttention', 'Scale',
    'SelfAttentionBlock', 'MultiHeadAttention', 'ContextBlock', 'NonLocal1d',
    'NonLocal2d', 'NonLocal3d'
]

python/jittorseg/models/layers/__init__.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn

from .scale import Scale


def NEG_INF_DIAG(n):
    return jt.diag(jt.Var(float('-inf')).repeat(n), 0)


class CrissCrossAttention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.query_conv = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.key_conv = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.value_conv = nn.Conv2d(in_channels, in_channels, 1)
        self.gamma = Scale(0.)
        self.in_channels = in_channels

    def execute(self, x):
        B, C, H, W = x.size()
        query = self.query_conv(x)
        key = self.key_conv(x)
        value = self.value_conv(x)
        energy_H = jt.linalg.einsum('bchw,bciw->bwhi', query,
                                    key) + NEG_INF_DIAG(H)
        energy_H = energy_H.transpose(1, 2)
        energy_W = jt.linalg.einsum('bchw,bchj->bhwj', query, key)
        attn = nn.softmax(jt.concat([energy_H, energy_W], dim=-1),
                          dim=-1)  # [B,H,W,(H+W)]
        out = jt.linalg.einsum('bciw,bhwi->bchw', value, attn[..., :H])
        out += jt.linalg.einsum('bchj,bhwj->bchw', value, attn[..., H:])

        out = self.gamma(out) + x

        return out

    def __repr__(self):
        s = self.__class__.__name__
        s += f'(in_channels={self.in_channels})'
        return s

python/jittorseg/models/layers/cc_attention.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn

from jittorseg.utils.weight_init import constant_init, kaiming_init


def last_zero_init(m):
    if isinstance(m, nn.Sequential):
        constant_init(m[-1], val=0)
    else:
        constant_init(m, val=0)


class ContextBlock(nn.Module):
    _abbr_ = 'context_block'

    def __init__(self,
                 in_channels,
                 ratio,
                 pooling_type='att',
                 fusion_types=('channel_add', )):
        super().__init__()
        assert pooling_type in ['avg', 'att']
        assert isinstance(fusion_types, (list, tuple))
        valid_fusion_types = ['channel_add', 'channel_mul']
        assert all([f in valid_fusion_types for f in fusion_types])
        assert len(fusion_types) > 0, 'at least one fusion should be used'
        self.in_channels = in_channels
        self.ratio = ratio
        self.planes = int(in_channels * ratio)
        self.pooling_type = pooling_type
        self.fusion_types = fusion_types
        if pooling_type == 'att':
            self.conv_mask = nn.Conv2d(in_channels, 1, kernel_size=1)
        else:
            self.avg_pool = nn.AdaptiveAvgPool2d(1)
        if 'channel_add' in fusion_types:
            self.channel_add_conv = nn.Sequential(
                nn.Conv2d(self.in_channels, self.planes, kernel_size=1),
                nn.LayerNorm([self.planes, 1, 1]),
                nn.ReLU(),  # yapf: disable
                nn.Conv2d(self.planes, self.in_channels, kernel_size=1))
        else:
            self.channel_add_conv = None
        if 'channel_mul' in fusion_types:
            self.channel_mul_conv = nn.Sequential(
                nn.Conv2d(self.in_channels, self.planes, kernel_size=1),
                nn.LayerNorm([self.planes, 1, 1]),
                nn.ReLU(),  # yapf: disable
                nn.Conv2d(self.planes, self.in_channels, kernel_size=1))
        else:
            self.channel_mul_conv = None
        self.reset_parameters()

    def reset_parameters(self):
        if self.pooling_type == 'att':
            kaiming_init(self.conv_mask, mode='fan_in')
            self.conv_mask.inited = True

        if self.channel_add_conv is not None:
            last_zero_init(self.channel_add_conv)
        if self.channel_mul_conv is not None:
            last_zero_init(self.channel_mul_conv)

    def spatial_pool(self, x):
        batch, channel, height, width = x.size()
        if self.pooling_type == 'att':
            input_x = x
            # [N, C, H * W]
            input_x = input_x.view(batch, channel, height * width)
            # [N, 1, C, H * W]
            input_x = input_x.unsqueeze(1)
            # [N, 1, H, W]
            context_mask = self.conv_mask(x)
            # [N, 1, H * W]
            context_mask = context_mask.view(batch, 1, height * width)
            # [N, 1, H * W]
            context_mask = nn.softmax(context_mask, dim=2)
            # [N, 1, H * W, 1]
            context_mask = context_mask.unsqueeze(-1)
            # [N, 1, C, 1]
            context = jt.matmul(input_x, context_mask)
            # [N, C, 1, 1]
            context = context.view(batch, channel, 1, 1)
        else:
            # [N, C, 1, 1]
            context = self.avg_pool(x)

        return context

    def execute(self, x):
        # [N, C, 1, 1]
        context = self.spatial_pool(x)

        out = x
        if self.channel_mul_conv is not None:
            # [N, C, 1, 1]
            channel_mul_term = jt.sigmoid(self.channel_mul_conv(context))
            out = out * channel_mul_term
        if self.channel_add_conv is not None:
            # [N, C, 1, 1]
            channel_add_term = self.channel_add_conv(context)
            out = out + channel_add_term

        return out

python/jittorseg/models/layers/context_block.py
JittorCV/jittorseg
from jittor import Module, nn
from ..basic_operators import ConvModule


class External_attention(Module):
    '''
    Arguments:
        c (int): The input and output channel number.
    '''
    def __init__(self, in_channels, channels, k=256):
        super(External_attention, self).__init__()

        self.in_channels = in_channels
        self.channels = channels
        self.k = k

        self.conv1 = ConvModule(self.in_channels, self.channels, 1)

        self.linear_0 = ConvModule(self.channels, self.k, 1)

        self.linear_1 = ConvModule(self.k, self.channels, 1)

        self.conv2 = ConvModule(self.channels, self.channels, 1)

    def execute(self, x):
        x = self.conv1(x)
        idn = x
        b, c, h, w = x.size()
        x = self.linear_0(x)  # b, k, h, w
        x = x.view(b, self.k, h * w)  # b * k * n

        x = nn.softmax(x, dim=-1)  # b, k, n
        x = x / (1e-9 + x.sum(dim=1, keepdims=True))  # b, k, n

        x = x.view(b, self.k, h, w)
        x = self.linear_1(x)  # b, c, h, w

        x = x + idn
        x = self.conv2(x)
        return x

python/jittorseg/models/layers/external_attention.py
JittorCV/jittorseg
import warnings
import jittor as jt
from jittor import nn
from ..basic_operators import build_dropout


class MultiHeadAttention(nn.Module):
    def __init__(self,
                 embed_dims,
                 num_heads,
                 attn_drop=0.,
                 proj_drop=0.,
                 dropout_layer=dict(type='Dropout', p=0.),
                 init_cfg=None,
                 batch_first=False,
                 **kwargs):
        super().__init__(init_cfg)
        self.embed_dims = embed_dims
        self.num_heads = num_heads
        self.batch_first = batch_first

        # TODO: jt.attention.MultiheadAttention have inconsistent with torch.nn.MultiheadAttention
        self.attn = jt.attention.MultiheadAttention(embed_dims,
                                                    num_heads,
                                                    dropout=attn_drop,
                                                    **kwargs)

        self.proj_drop = nn.Dropout(proj_drop)
        self.dropout_layer = build_dropout(
            dropout_layer) if dropout_layer else nn.Identity()

    def execute(self,
                query,
                key=None,
                value=None,
                identity=None,
                query_pos=None,
                key_pos=None,
                attn_mask=None,
                key_padding_mask=None,
                **kwargs):
        """Forward function for `MultiheadAttention`.

        **kwargs allow passing a more general data flow when combining
        with other operations in `transformerlayer`.

        Args:
            query (Var): The input query with shape [num_queries, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_queries embed_dims].
            key (Var): The key Var with shape [num_keys, bs,
                embed_dims] if self.batch_first is False, else
                [bs, num_keys, embed_dims] .
                If None, the ``query`` will be used. Defaults to None.
            value (Var): The value Var with same shape as `key`.
                Same in `nn.MultiheadAttention.forward`. Defaults to None.
                If None, the `key` will be used.
            identity (Var): This Var, with the same shape as x,
                will be used for the identity link.
                If None, `x` will be used. Defaults to None.
            query_pos (Var): The positional encoding for query, with
                the same shape as `x`. If not None, it will
                be added to `x` before forward function. Defaults to None.
            key_pos (Var): The positional encoding for `key`, with the
                same shape as `key`. Defaults to None. If not None, it will
                be added to `key` before forward function. If None, and
                `query_pos` has the same shape as `key`, then `query_pos`
                will be used for `key_pos`. Defaults to None.
            attn_mask (Var): ByteVar mask with shape [num_queries,
                num_keys]. Same in `nn.MultiheadAttention.forward`.
                Defaults to None.
            key_padding_mask (Var): ByteVar with shape [bs, num_keys].
                Defaults to None.

        Returns:
            Var: forwarded results with shape
            [num_queries, bs, embed_dims]
            if self.batch_first is False, else
            [bs, num_queries embed_dims].
        """

        if key is None:
            key = query
        if value is None:
            value = key
        if identity is None:
            identity = query
        if key_pos is None:
            if query_pos is not None:
                # use query_pos if key_pos is not available
                if query_pos.shape == key.shape:
                    key_pos = query_pos
                else:
                    warnings.warn(f'position encoding of key is'
                                  f'missing in {self.__class__.__name__}.')
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos

        if self.batch_first:
            query = query.transpose(0, 1)
            key = key.transpose(0, 1)
            value = value.transpose(0, 1)

        out = self.attn(query=query,
                        key=key,
                        value=value,
                        attn_mask=attn_mask,
                        key_padding_mask=key_padding_mask)[0]

        if self.batch_first:
            out = out.transpose(0, 1)

        return identity + self.dropout_layer(self.proj_drop(out))

python/jittorseg/models/layers/multi_head_attention.py
JittorCV/jittorseg
from abc import ABCMeta

import jittor as jt
from jittor import nn

from jittorseg.utils.weight_init import constant_init, normal_init
from ..basic_operators import ConvModule


class _NonLocalNd(nn.Module, metaclass=ABCMeta):

    def __init__(self,
                 in_channels,
                 reduction=2,
                 use_scale=True,
                 conv_cfg=None,
                 norm_cfg=None,
                 mode='embedded_gaussian',
                 **kwargs):
        super().__init__()
        self.in_channels = in_channels
        self.reduction = reduction
        self.use_scale = use_scale
        self.inter_channels = max(in_channels // reduction, 1)
        self.mode = mode

        if mode not in [
                'gaussian', 'embedded_gaussian', 'dot_product', 'concatenation'
        ]:
            raise ValueError("Mode should be in 'gaussian', 'concatenation', "
                             f"'embedded_gaussian' or 'dot_product', but got "
                             f'{mode} instead.')

        # g, theta, phi are defaulted as `nn.ConvNd`.
        # Here we use ConvModule for potential usage.
        self.g = ConvModule(self.in_channels,
                            self.inter_channels,
                            kernel_size=1,
                            conv_cfg=conv_cfg,
                            act_cfg=None)
        self.conv_out = ConvModule(self.inter_channels,
                                   self.in_channels,
                                   kernel_size=1,
                                   conv_cfg=conv_cfg,
                                   norm_cfg=norm_cfg,
                                   act_cfg=None)

        if self.mode != 'gaussian':
            self.theta = ConvModule(self.in_channels,
                                    self.inter_channels,
                                    kernel_size=1,
                                    conv_cfg=conv_cfg,
                                    act_cfg=None)
            self.phi = ConvModule(self.in_channels,
                                  self.inter_channels,
                                  kernel_size=1,
                                  conv_cfg=conv_cfg,
                                  act_cfg=None)

        if self.mode == 'concatenation':
            self.concat_project = ConvModule(self.inter_channels * 2,
                                             1,
                                             kernel_size=1,
                                             stride=1,
                                             padding=0,
                                             bias=False,
                                             act_cfg=dict(type='ReLU'))

        self.init_weights(**kwargs)

    def init_weights(self, std=0.01, zeros_init=True):
        if self.mode != 'gaussian':
            for m in [self.g, self.theta, self.phi]:
                normal_init(m.conv, std=std)
        else:
            normal_init(self.g.conv, std=std)
        if zeros_init:
            if self.conv_out.norm_cfg is None:
                constant_init(self.conv_out.conv, 0)
            else:
                constant_init(self.conv_out.norm, 0)
        else:
            if self.conv_out.norm_cfg is None:
                normal_init(self.conv_out.conv, std=std)
            else:
                normal_init(self.conv_out.norm, std=std)

    def gaussian(self, theta_x, phi_x):
        # NonLocal1d pairwise_weight: [N, H, H]
        # NonLocal2d pairwise_weight: [N, HxW, HxW]
        # NonLocal3d pairwise_weight: [N, TxHxW, TxHxW]
        pairwise_weight = jt.matmul(theta_x, phi_x)
        pairwise_weight = pairwise_weight.softmax(dim=-1)
        return pairwise_weight

    def embedded_gaussian(self, theta_x, phi_x):
        # NonLocal1d pairwise_weight: [N, H, H]
        # NonLocal2d pairwise_weight: [N, HxW, HxW]
        # NonLocal3d pairwise_weight: [N, TxHxW, TxHxW]
        pairwise_weight = jt.matmul(theta_x, phi_x)
        if self.use_scale:
            # theta_x.shape[-1] is `self.inter_channels`
            pairwise_weight /= theta_x.shape[-1]**0.5
        pairwise_weight = pairwise_weight.softmax(dim=-1)
        return pairwise_weight

    def dot_product(self, theta_x, phi_x):
        # NonLocal1d pairwise_weight: [N, H, H]
        # NonLocal2d pairwise_weight: [N, HxW, HxW]
        # NonLocal3d pairwise_weight: [N, TxHxW, TxHxW]
        pairwise_weight = jt.matmul(theta_x, phi_x)
        pairwise_weight /= pairwise_weight.shape[-1]
        return pairwise_weight

    def concatenation(self, theta_x, phi_x):
        # NonLocal1d pairwise_weight: [N, H, H]
        # NonLocal2d pairwise_weight: [N, HxW, HxW]
        # NonLocal3d pairwise_weight: [N, TxHxW, TxHxW]
        h = theta_x.size(2)
        w = phi_x.size(3)
        theta_x = theta_x.repeat(1, 1, 1, w)
        phi_x = phi_x.repeat(1, 1, h, 1)

        concat_feature = jt.concat([theta_x, phi_x], dim=1)
        pairwise_weight = self.concat_project(concat_feature)
        n, _, h, w = pairwise_weight.size()
        pairwise_weight = pairwise_weight.view(n, h, w)
        pairwise_weight /= pairwise_weight.shape[-1]

        return pairwise_weight

    def execute(self, x):
        # Assume `reduction = 1`, then `inter_channels = C`
        # or `inter_channels = C` when `mode="gaussian"`

        # NonLocal1d x: [N, C, H]
        # NonLocal2d x: [N, C, H, W]
        # NonLocal3d x: [N, C, T, H, W]
        n = x.size(0)

        # NonLocal1d g_x: [N, H, C]
        # NonLocal2d g_x: [N, HxW, C]
        # NonLocal3d g_x: [N, TxHxW, C]
        g_x = self.g(x).view(n, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        # NonLocal1d theta_x: [N, H, C], phi_x: [N, C, H]
        # NonLocal2d theta_x: [N, HxW, C], phi_x: [N, C, HxW]
        # NonLocal3d theta_x: [N, TxHxW, C], phi_x: [N, C, TxHxW]
        if self.mode == 'gaussian':
            theta_x = x.view(n, self.in_channels, -1)
            theta_x = theta_x.permute(0, 2, 1)
            if self.sub_sample:
                phi_x = self.phi(x).view(n, self.in_channels, -1)
            else:
                phi_x = x.view(n, self.in_channels, -1)
        elif self.mode == 'concatenation':
            theta_x = self.theta(x).view(n, self.inter_channels, -1, 1)
            phi_x = self.phi(x).view(n, self.inter_channels, 1, -1)
        else:
            theta_x = self.theta(x).view(n, self.inter_channels, -1)
            theta_x = theta_x.permute(0, 2, 1)
            phi_x = self.phi(x).view(n, self.inter_channels, -1)

        pairwise_func = getattr(self, self.mode)
        # NonLocal1d pairwise_weight: [N, H, H]
        # NonLocal2d pairwise_weight: [N, HxW, HxW]
        # NonLocal3d pairwise_weight: [N, TxHxW, TxHxW]
        pairwise_weight = pairwise_func(theta_x, phi_x)

        # NonLocal1d y: [N, H, C]
        # NonLocal2d y: [N, HxW, C]
        # NonLocal3d y: [N, TxHxW, C]
        y = jt.matmul(pairwise_weight, g_x)
        # NonLocal1d y: [N, C, H]
        # NonLocal2d y: [N, C, H, W]
        # NonLocal3d y: [N, C, T, H, W]
        y = y.permute(0, 2, 1).reshape(n, self.inter_channels, *x.size()[2:])

        output = x + self.conv_out(y)

        return output


class NonLocal1d(_NonLocalNd):

    def __init__(self,
                 in_channels,
                 sub_sample=False,
                 conv_cfg=dict(type='Conv1d'),
                 **kwargs):
        super().__init__(in_channels, conv_cfg=conv_cfg, **kwargs)

        self.sub_sample = sub_sample

        if sub_sample:
            max_pool_layer = nn.MaxPool1d(kernel_size=2)
            self.g = nn.Sequential(self.g, max_pool_layer)
            if self.mode != 'gaussian':
                self.phi = nn.Sequential(self.phi, max_pool_layer)
            else:
                self.phi = max_pool_layer


class NonLocal2d(_NonLocalNd):
    _abbr_ = 'nonlocal_block'

    def __init__(self,
                 in_channels,
                 sub_sample=False,
                 conv_cfg=dict(type='Conv2d'),
                 **kwargs):
        super().__init__(in_channels, conv_cfg=conv_cfg, **kwargs)

        self.sub_sample = sub_sample

        if sub_sample:
            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
            self.g = nn.Sequential(self.g, max_pool_layer)
            if self.mode != 'gaussian':
                self.phi = nn.Sequential(self.phi, max_pool_layer)
            else:
                self.phi = max_pool_layer


class NonLocal3d(_NonLocalNd):

    def __init__(self,
                 in_channels,
                 sub_sample=False,
                 conv_cfg=dict(type='Conv3d'),
                 **kwargs):
        super().__init__(in_channels, conv_cfg=conv_cfg, **kwargs)
        self.sub_sample = sub_sample

        if sub_sample:
            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
            self.g = nn.Sequential(self.g, max_pool_layer)
            if self.mode != 'gaussian':
                self.phi = nn.Sequential(self.phi, max_pool_layer)
            else:
                self.phi = max_pool_layer

python/jittorseg/models/layers/non_local.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn


class Scale(nn.Module):
    """A learnable scale parameter.

    This layer scales the input by a learnable factor. It multiplies a
    learnable scale parameter of shape (1,) with input of any shape.

    Args:
        scale (float): Initial value of scale factor. Default: 1.0
    """

    def __init__(self, scale=1.0):
        super().__init__()
        self.scale = jt.Var(scale)

    def execute(self, x):
        return x * self.scale

python/jittorseg/models/layers/scale.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from ..basic_operators import ConvModule
from jittorseg.utils.weight_init import constant_init


class SelfAttentionBlock(nn.Module):
    """General self-attention block/non-local block.

    Please refer to https://arxiv.org/abs/1706.03762 for details about key,
    query and value.

    Args:
        key_in_channels (int): Input channels of key feature.
        query_in_channels (int): Input channels of query feature.
        channels (int): Output channels of key/query transform.
        out_channels (int): Output channels.
        share_key_query (bool): Whether share projection weight between key
            and query projection.
        query_downsample (nn.Module): Query downsample module.
        key_downsample (nn.Module): Key downsample module.
        key_query_num_convs (int): Number of convs for key/query projection.
        value_num_convs (int): Number of convs for value projection.
        matmul_norm (bool): Whether normalize attention map with sqrt of
            channels
        with_out (bool): Whether use out projection.
        conv_cfg (dict|None): Config of conv layers.
        norm_cfg (dict|None): Config of norm layers.
        act_cfg (dict|None): Config of activation layers.
    """

    def __init__(self, key_in_channels, query_in_channels, channels,
                 out_channels, share_key_query, query_downsample,
                 key_downsample, key_query_num_convs, value_out_num_convs,
                 key_query_norm, value_out_norm, matmul_norm, with_out,
                 conv_cfg, norm_cfg, act_cfg):
        super(SelfAttentionBlock, self).__init__()
        if share_key_query:
            assert key_in_channels == query_in_channels
        self.key_in_channels = key_in_channels
        self.query_in_channels = query_in_channels
        self.out_channels = out_channels
        self.channels = channels
        self.share_key_query = share_key_query
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.key_project = self.build_project(key_in_channels,
                                              channels,
                                              num_convs=key_query_num_convs,
                                              use_conv_module=key_query_norm,
                                              conv_cfg=conv_cfg,
                                              norm_cfg=norm_cfg,
                                              act_cfg=act_cfg)
        if share_key_query:
            self.query_project = self.key_project
        else:
            self.query_project = self.build_project(
                query_in_channels,
                channels,
                num_convs=key_query_num_convs,
                use_conv_module=key_query_norm,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg)
        self.value_project = self.build_project(
            key_in_channels,
            channels if with_out else out_channels,
            num_convs=value_out_num_convs,
            use_conv_module=value_out_norm,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)
        if with_out:
            self.out_project = self.build_project(
                channels,
                out_channels,
                num_convs=value_out_num_convs,
                use_conv_module=value_out_norm,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg)
        else:
            self.out_project = None

        self.query_downsample = query_downsample
        self.key_downsample = key_downsample
        self.matmul_norm = matmul_norm

        self.init_weights()

    def init_weights(self):
        """Initialize weight of later layer."""
        if self.out_project is not None:
            if not isinstance(self.out_project, ConvModule):
                constant_init(self.out_project, 0)

    def build_project(self, in_channels, channels, num_convs, use_conv_module,
                      conv_cfg, norm_cfg, act_cfg):
        """Build projection layer for key/query/value/out."""
        if use_conv_module:
            convs = [
                ConvModule(in_channels,
                           channels,
                           1,
                           conv_cfg=conv_cfg,
                           norm_cfg=norm_cfg,
                           act_cfg=act_cfg)
            ]
            for _ in range(num_convs - 1):
                convs.append(
                    ConvModule(channels,
                               channels,
                               1,
                               conv_cfg=conv_cfg,
                               norm_cfg=norm_cfg,
                               act_cfg=act_cfg))
        else:
            convs = [nn.Conv2d(in_channels, channels, 1)]
            for _ in range(num_convs - 1):
                convs.append(nn.Conv2d(channels, channels, 1))
        if len(convs) > 1:
            convs = nn.Sequential(*convs)
        else:
            convs = convs[0]
        return convs

    def execute(self, query_feats, key_feats):
        """Forward function."""
        batch_size = query_feats.size(0)
        query = self.query_project(query_feats)
        if self.query_downsample is not None:
            query = self.query_downsample(query)
        query = query.reshape(*query.shape[:2], -1)
        query = query.permute(0, 2, 1)

        key = self.key_project(key_feats)
        value = self.value_project(key_feats)
        if self.key_downsample is not None:
            key = self.key_downsample(key)
            value = self.key_downsample(value)
        key = key.reshape(*key.shape[:2], -1)
        value = value.reshape(*value.shape[:2], -1)
        value = value.permute(0, 2, 1)

        sim_map = jt.matmul(query, key)
        if self.matmul_norm:
            sim_map = (self.channels**-.5) * sim_map
        sim_map = nn.softmax(sim_map, dim=-1)

        context = jt.matmul(sim_map, value)
        context = context.permute(0, 2, 1)
        context = context.reshape(batch_size, -1, *query_feats.shape[2:])
        if self.out_project is not None:
            context = self.out_project(context)
        return context

python/jittorseg/models/layers/self_attention_block.py
JittorCV/jittorseg
import warnings
from jittor import nn
from jittor import Function


# TODO Save memory
class Resize(Function):

    def execute(self, input, size, scale_factor, mode, align_corners):
        self.input_size = input.shape[2:]
        self.scale_factor = scale_factor
        self.mode = mode
        self.align_corners = align_corners
        return nn.interpolate(input, size, scale_factor, mode, align_corners)

    def grad(self, grad_output):
        return nn.interpolate(grad_output, self.input_size, self.scale_factor,
                              self.mode, self.align_corners)


interpolate = Resize.apply


def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    if size is not None:
        size = tuple(int(x) for x in size)
    return interpolate(input, size, scale_factor, mode, align_corners)


class Upsample(nn.Module):

    def __init__(self,
                 size=None,
                 scale_factor=None,
                 mode='nearest',
                 align_corners=None):
        super(Upsample, self).__init__()
        self.size = size
        if isinstance(scale_factor, tuple):
            self.scale_factor = tuple(float(factor) for factor in scale_factor)
        else:
            self.scale_factor = float(scale_factor) if scale_factor else None
        self.mode = mode
        self.align_corners = align_corners

    def execute(self, x):
        if not self.size:
            size = [int(t * self.scale_factor) for t in x.shape[-2:]]
        else:
            size = self.size
        return resize(x, size, None, self.mode, self.align_corners)

python/jittorseg/models/layers/wrappers.py
JittorCV/jittorseg
from .accuracy import Accuracy, accuracy
from .cross_entropy_loss import (CrossEntropyLoss, cross_entropy)
from .utils import reduce_loss, weight_reduce_loss, weighted_loss

__all__ = [
    'accuracy', 'Accuracy', 'cross_entropy', 'CrossEntropyLoss', 'reduce_loss',
    'weight_reduce_loss', 'weighted_loss'
]

python/jittorseg/models/losses/__init__.py
JittorCV/jittorseg
from jittor import nn
import jittor as jt

eps = jt.Var(1.1920928955078125e-07)


def accuracy(pred, target, topk=1, thresh=None, ignore_index=None):
    assert isinstance(topk, (int, tuple))
    if isinstance(topk, int):
        topk = (topk, )
        return_single = True
    else:
        return_single = False

    maxk = max(topk)
    if pred.size(0) == 0:
        accu = [jt.Var(0.) for i in range(len(topk))]
        return accu[0] if return_single else accu
    assert pred.ndim == target.ndim + 1
    assert pred.size(0) == target.size(0)
    assert maxk <= pred.size(1), \
        f'maxk {maxk} exceeds pred dimension {pred.size(1)}'
    pred_value, pred_label = pred.topk(maxk, dim=1)
    # transpose to shape (maxk, N, ...)
    pred_label = pred_label.transpose(0, 1)
    correct = pred_label == (target.unsqueeze(0).expand_as(pred_label))
    if thresh is not None:
        # Only prediction values larger than thresh are counted as correct
        correct = correct & (pred_value > thresh).t()
    if ignore_index is not None:
        correct = correct[:, target != ignore_index]
    res = []
    for k in topk:
        # Avoid causing ZeroDivisionError when all pixels
        # of an image are ignored
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdims=True) + eps
        if ignore_index is not None:
            total_num = target[target != ignore_index].numel() + eps
        else:
            total_num = target.numel() + eps
        res.append(correct_k.multiply(100.0 / total_num))
    return res[0] if return_single else res


class Accuracy(nn.Module):

    def __init__(self, topk=(1, ), thresh=None, ignore_index=None):
        super().__init__()
        self.topk = topk
        self.thresh = thresh
        self.ignore_index = ignore_index

    def execute(self, pred, target):
        return accuracy(pred, target, self.topk, self.thresh,
                        self.ignore_index)

python/jittorseg/models/losses/accuracy.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn
from jittorseg.utils.registry import LOSSES
from .utils import weight_reduce_loss
from jittor.nn import cross_entropy_loss
import warnings


def cross_entropy(pred,
                  label,
                  weight=None,
                  class_weight=None,
                  reduction='mean',
                  avg_factor=None,
                  ignore_index=-100,
                  avg_non_ignore=False):
    loss = cross_entropy_loss(pred,
                              label,
                              weight=class_weight,
                              reduction='none',
                              ignore_index=ignore_index)

    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


@LOSSES.register_module()
class CrossEntropyLoss(nn.Module):
    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 = 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 Jittor official '
                'cross_entropy, set ``avg_non_ignore=True``.')

        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 execute(self,
                cls_score,
                label,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                ignore_index=-100,
                **kwargs):
        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 = jt.Var(self.class_weight)
        else:
            class_weight = None
        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):
        return self._loss_name

python/jittorseg/models/losses/cross_entropy_loss.py
JittorCV/jittorseg
import functools
import warnings


def get_enum(reduction: str) -> int:
    if reduction == 'none':
        ret = 0
    elif reduction == 'mean':
        ret = 1
    elif reduction == 'elementwise_mean':
        warnings.warn(
            "reduction='elementwise_mean' is deprecated, please use reduction='mean' instead."
        )
        ret = 1
    elif reduction == 'sum':
        ret = 2
    else:
        ret = -1  # TODO: remove once JIT exceptions support control flow
        raise ValueError(
            "{} is not a valid value for reduction".format(reduction))
    return ret


def reduce_loss(loss, reduction):
    reduction_enum = get_enum(reduction)
    # none: 0, elementwise_mean:1, sum: 2
    if reduction_enum == 0:
        return loss
    elif reduction_enum == 1:
        return loss.mean()
    elif reduction_enum == 2:
        return loss.sum()


def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None):
    # if weight is specified, apply element-wise weight
    if weight is not None:
        assert weight.ndim == loss.ndim
        if weight.ndim > 1:
            assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
        loss = loss * weight

    # if avg_factor is not specified, just reduce the loss
    if avg_factor is None:
        loss = reduce_loss(loss, reduction)
    else:
        # if reduction is mean, then average the loss by avg_factor
        if reduction == 'mean':
            loss = loss.sum() / avg_factor
        # if reduction is 'none', then do nothing, otherwise raise an error
        elif reduction != 'none':
            raise ValueError('avg_factor can not be used with reduction="sum"')
    return loss


def weighted_loss(loss_func):
    @functools.wraps(loss_func)
    def wrapper(pred,
                target,
                weight=None,
                reduction='mean',
                avg_factor=None,
                **kwargs):
        # get element-wise loss
        loss = loss_func(pred, target, **kwargs)
        loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
        return loss

    return wrapper

python/jittorseg/models/losses/utils.py
JittorCV/jittorseg
from .multilevel_neck import MultiLevelNeck
from .featurepyramid import Feature2Pyramid
from .fpn import FPN
python/jittorseg/models/necks/__init__.py
JittorCV/jittorseg
from jittor import nn
from ..basic_operators import build_norm_layer

from jittorseg.utils.registry import NECKS


@NECKS.register_module()
class Feature2Pyramid(nn.Module):
    def __init__(self,
                 embed_dim,
                 rescales=[4, 2, 1, 0.5],
                 norm_cfg=dict(type='BN')):
        super(Feature2Pyramid, self).__init__()
        self.rescales = rescales
        self.upsample_4x = None
        for k in self.rescales:
            if k == 4:
                self.upsample_4x = nn.Sequential(
                    nn.ConvTranspose2d(
                        embed_dim, embed_dim, kernel_size=2, stride=2),
                    build_norm_layer(norm_cfg, embed_dim)[1],
                    nn.GELU(),
                    nn.ConvTranspose2d(
                        embed_dim, embed_dim, kernel_size=2, stride=2),
                )
            elif k == 2:
                self.upsample_2x = nn.Sequential(
                    nn.ConvTranspose2d(
                        embed_dim, embed_dim, kernel_size=2, stride=2))
            elif k == 1:
                self.identity = nn.Identity()
            elif k == 0.5:
                self.downsample_2x = nn.MaxPool2d(kernel_size=2, stride=2)
            elif k == 0.25:
                self.downsample_4x = nn.MaxPool2d(kernel_size=4, stride=4)
            else:
                raise KeyError(f'invalid {k} for feature2pyramid')

    def execute(self, inputs):
        assert len(inputs) == len(self.rescales)
        outputs = []
        if self.upsample_4x is not None:
            ops = [
                self.upsample_4x, self.upsample_2x, self.identity,
                self.downsample_2x
            ]
        else:
            ops = [
                self.upsample_2x, self.identity, self.downsample_2x,
                self.downsample_4x
            ]
        for i in range(len(inputs)):
            outputs.append(ops[i](inputs[i]))
        return tuple(outputs)

python/jittorseg/models/necks/featurepyramid.py
JittorCV/jittorseg
from jittor import nn
from ..basic_operators import ConvModule

from ..layers import resize
from jittorseg.utils.registry import NECKS


@NECKS.register_module()
class FPN(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 num_outs,
                 start_level=0,
                 end_level=-1,
                 add_extra_convs=False,
                 extra_convs_on_inputs=False,
                 relu_before_extra_convs=False,
                 no_norm_on_lateral=False,
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=None,
                 upsample_cfg=dict(mode='nearest')):
        super(FPN, self).__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.relu_before_extra_convs = relu_before_extra_convs
        self.no_norm_on_lateral = no_norm_on_lateral
        self.fp16_enabled = False
        self.upsample_cfg = upsample_cfg.copy()

        if end_level == -1:
            self.backbone_end_level = self.num_ins
            assert num_outs >= self.num_ins - start_level
        else:
            # if end_level < inputs, no extra level is allowed
            self.backbone_end_level = end_level
            assert end_level <= len(in_channels)
            assert num_outs == end_level - start_level
        self.start_level = start_level
        self.end_level = end_level
        self.add_extra_convs = add_extra_convs
        assert isinstance(add_extra_convs, (str, bool))
        if isinstance(add_extra_convs, str):
            # Extra_convs_source choices: 'on_input', 'on_lateral', 'on_output'
            assert add_extra_convs in ('on_input', 'on_lateral', 'on_output')
        elif add_extra_convs:  # True
            if extra_convs_on_inputs:
                # For compatibility with previous release
                # TODO: deprecate `extra_convs_on_inputs`
                self.add_extra_convs = 'on_input'
            else:
                self.add_extra_convs = 'on_output'

        self.lateral_convs = nn.ModuleList()
        self.fpn_convs = nn.ModuleList()

        for i in range(self.start_level, self.backbone_end_level):
            l_conv = ConvModule(
                in_channels[i],
                out_channels,
                1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg if not self.no_norm_on_lateral else None,
                act_cfg=act_cfg)
            fpn_conv = ConvModule(out_channels,
                                  out_channels,
                                  3,
                                  padding=1,
                                  conv_cfg=conv_cfg,
                                  norm_cfg=norm_cfg,
                                  act_cfg=act_cfg)

            self.lateral_convs.append(l_conv)
            self.fpn_convs.append(fpn_conv)

        # add extra conv layers (e.g., RetinaNet)
        extra_levels = num_outs - self.backbone_end_level + self.start_level
        if self.add_extra_convs and extra_levels >= 1:
            for i in range(extra_levels):
                if i == 0 and self.add_extra_convs == 'on_input':
                    in_channels = self.in_channels[self.backbone_end_level - 1]
                else:
                    in_channels = out_channels
                extra_fpn_conv = ConvModule(in_channels,
                                            out_channels,
                                            3,
                                            stride=2,
                                            padding=1,
                                            conv_cfg=conv_cfg,
                                            norm_cfg=norm_cfg,
                                            act_cfg=act_cfg)
                self.fpn_convs.append(extra_fpn_conv)

    def execute(self, inputs):
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            # In some cases, fixing `scale factor` (e.g. 2) is preferred, but
            #  it cannot co-exist with `size` in `F.interpolate`.
            if 'scale_factor' in self.upsample_cfg:
                laterals[i - 1] = laterals[i - 1] + resize(
                    laterals[i], **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] = laterals[i - 1] + resize(
                    laterals[i], size=prev_shape, **self.upsample_cfg)

        # build outputs
        # part 1: from original levels
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]
        # part 2: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(nn.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    extra_source = inputs[self.backbone_end_level - 1]
                elif self.add_extra_convs == 'on_lateral':
                    extra_source = laterals[-1]
                elif self.add_extra_convs == 'on_output':
                    extra_source = outs[-1]
                else:
                    raise NotImplementedError
                outs.append(self.fpn_convs[used_backbone_levels](extra_source))
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](nn.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)

python/jittorseg/models/necks/fpn.py
JittorCV/jittorseg
from jittor import nn
from jittorseg.utils.weight_init import xavier_init
from ..layers import resize
from jittorseg.utils.registry import NECKS
import collections


@NECKS.register_module()
class MultiLevelNeck(nn.Module):
    """MultiLevelNeck.

    A neck structure connect vit backbone and decoder_heads.

    Args:
        in_channels (List[int]): Number of input channels per scale.
        out_channels (int): Number of output channels (used at each scale).
        scales (List[float]): Scale factors for each input feature map.
            Default: [0.5, 1, 2, 4]
        norm_cfg (dict): Config dict for normalization layer. Default: None.
        act_cfg (dict): Config dict for activation layer in ConvModule.
            Default: None.
    """

    def __init__(self, in_channels, out_channels, scales=[0.5, 1, 2, 4]):
        super(MultiLevelNeck, self).__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.scales = scales
        self.num_outs = len(scales)
        self.lateral_convs = nn.ModuleList()
        self.convs = nn.ModuleList()
        for in_channel in in_channels:

            self.lateral_convs.append(
                nn.Sequential(
                    collections.OrderedDict([
                        ('conv', nn.Conv(in_channel, out_channels, 1)),
                    ])))

        for _ in range(self.num_outs):
            self.convs.append(
                nn.Sequential(
                    collections.OrderedDict([('conv',
                                              nn.Conv(out_channels,
                                                      out_channels,
                                                      3,
                                                      padding=1))])))

    # default init_weights for conv(msra) and norm in ConvModule
    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution='uniform')

    def execute(self, inputs):
        assert len(inputs) == len(self.in_channels)
        inputs = [
            lateral_conv(inputs[i])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]
        # for len(inputs) not equal to self.num_outs
        if len(inputs) == 1:
            inputs = [inputs[0] for _ in range(self.num_outs)]
        outs = []
        for i in range(self.num_outs):
            x_resize = resize(inputs[i],
                              scale_factor=self.scales[i],
                              mode='bilinear')
            outs.append(self.convs[i](x_resize))
        return tuple(outs)

python/jittorseg/models/necks/multilevel_neck.py
JittorCV/jittorseg
from .encoder_decoder import EncoderDecoder
from .cascade_encoder_decoder import CascadeEncoderDecoder

__all__ = ['EncoderDecoder', 'CascadeEncoderDecoder']

python/jittorseg/models/segmentors/__init__.py
JittorCV/jittorseg
import logging
import warnings
from abc import ABCMeta, abstractmethod
from collections import OrderedDict

import cv2
import numpy as np
import jittor as jt
from jittor import nn


class BaseSegmentor(nn.Module):

    __metaclass__ = ABCMeta

    def __init__(self):
        super(BaseSegmentor, self).__init__()

    @property
    def with_neck(self):
        return hasattr(self, 'neck') and self.neck is not None

    @property
    def with_auxiliary_head(self):
        return hasattr(self,
                       'auxiliary_head') and self.auxiliary_head is not None

    @property
    def with_decode_head(self):
        return hasattr(self, 'decode_head') and self.decode_head is not None

    @abstractmethod
    def extract_feat(self, imgs):
        pass

    @abstractmethod
    def encode_decode(self, img, img_metas):
        pass

    @abstractmethod
    def execute_train(self, imgs, img_metas, **kwargs):
        pass

    @abstractmethod
    def simple_test(self, img, img_meta, **kwargs):
        pass

    @abstractmethod
    def aug_test(self, imgs, img_metas, **kwargs):
        pass

    def init_weights(self, pretrained=None):
        if pretrained is not None:
            logger = logging.getLogger()
            logger.info(f'load model from: {pretrained}')

    def execute_test(self, imgs, img_metas, **kwargs):
        for var, name in [(imgs, 'imgs'), (img_metas, 'img_metas')]:
            if not isinstance(var, list):
                raise TypeError(f'{name} must be a list, but got '
                                f'{type(var)}')

        num_augs = len(imgs)
        if num_augs != len(img_metas):
            raise ValueError(f'num of augmentations ({len(imgs)}) != '
                             f'num of image meta ({len(img_metas)})')
        # all images in the same aug batch all of the same ori_shape and pad
        if num_augs == 1:
            return self.simple_test(imgs[0], img_metas[0], **kwargs)
        else:
            return self.aug_test(imgs, img_metas, **kwargs)

    def execute(self, img, img_metas, return_loss=True, **kwargs):
        if return_loss:
            return self.execute_train(img, img_metas, **kwargs)
        else:
            return self.execute_test(img, img_metas, **kwargs)

    @staticmethod
    def _parse_losses(losses):
        log_vars = OrderedDict()
        for loss_name, loss_value in losses.items():
            if isinstance(loss_value, jt.Var):
                log_vars[loss_name] = loss_value.mean()
            elif isinstance(loss_value, list):
                log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
            else:
                raise TypeError(f'{loss_name} is not a var or list of vars')

        loss = sum(_value for _key, _value in log_vars.items()
                   if 'loss' in _key)

        log_vars['loss'] = loss
        for loss_name, loss_value in log_vars.items():
            log_vars[loss_name] = loss_value.item()

        return loss, log_vars

    def show_result(self,
                    img,
                    result,
                    palette=None,
                    show=False,
                    out_file=None,
                    opacity=0.5):
        img = cv2.imread(img)
        img = img.copy()
        seg = result[0]
        if palette is None:
            if self.PALETTE is None:
                # Get random state before set seed,
                # and restore random state later.
                # It will prevent loss of randomness, as the palette
                # may be different in each iteration if not specified.
                # See: https://github.com/open-mmlab/mmdetection/issues/5844
                state = np.random.get_state()
                np.random.seed(42)
                # random palette
                palette = np.random.randint(
                    0, 255, size=(len(self.CLASSES), 3))
                np.random.set_state(state)
            else:
                palette = self.PALETTE
        palette = np.array(palette)
        assert palette.shape[0] == len(self.CLASSES)
        assert palette.shape[1] == 3
        assert len(palette.shape) == 2
        assert 0 < opacity <= 1.0
        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
        for label, color in enumerate(palette):
            color_seg[seg == label, :] = color
        # convert to BGR
        color_seg = color_seg[..., ::-1]

        img = img * (1 - opacity) + color_seg * opacity
        img = img.astype(np.uint8)
        # if out_file specified, do not show image in window
        if out_file is not None:
            show = False

        if out_file is not None:
            cv2.imwrite(out_file, img)

        if not (show or out_file):
            warnings.warn('show==False and out_file is not specified, only '
                          'result image will be returned')
            return img

python/jittorseg/models/segmentors/base.py
JittorCV/jittorseg
from jittor import nn

from jittorseg.utils.general import add_prefix
from ..layers import resize
from jittorseg.utils.registry import MODELS, build_from_cfg, HEADS
from .encoder_decoder import EncoderDecoder


@MODELS.register_module()
class CascadeEncoderDecoder(EncoderDecoder):

    def __init__(self,
                 num_stages,
                 backbone,
                 decode_head,
                 neck=None,
                 auxiliary_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        self.num_stages = num_stages
        super(CascadeEncoderDecoder,
              self).__init__(backbone=backbone,
                             decode_head=decode_head,
                             neck=neck,
                             auxiliary_head=auxiliary_head,
                             train_cfg=train_cfg,
                             test_cfg=test_cfg,
                             pretrained=pretrained)

    def _init_decode_head(self, decode_head):
        """Initialize ``decode_head``"""
        assert isinstance(decode_head, list)
        assert len(decode_head) == self.num_stages
        self.decode_head = nn.ModuleList()
        for i in range(self.num_stages):
            self.decode_head.append(build_from_cfg(decode_head[i], HEADS))
        self.align_corners = self.decode_head[-1].align_corners
        self.num_classes = self.decode_head[-1].num_classes
        self.out_channels = self.decode_head[-1].out_channels

    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[0].execute_test(x, img_metas, self.test_cfg)
        for i in range(1, self.num_stages):
            out = self.decode_head[i].execute_test(x, out, img_metas,
                                                   self.test_cfg)
        out = resize(input=out,
                     size=img.shape[2:],
                     mode='bilinear',
                     align_corners=self.align_corners)
        return out

    def _decode_head_execute_train(self, x, img_metas, gt_semantic_seg):
        """Run execute function and calculate loss for decode head in
        training."""
        losses = dict()

        loss_decode = self.decode_head[0].execute_train(
            x, img_metas, gt_semantic_seg, self.train_cfg)

        losses.update(add_prefix(loss_decode, 'decode_0'))

        for i in range(1, self.num_stages):
            # execute test again, maybe unnecessary for most methods.
            if i == 1:
                prev_outputs = self.decode_head[0].execute_test(
                    x, img_metas, self.test_cfg)
            else:
                prev_outputs = self.decode_head[i - 1].execute_test(
                    x, prev_outputs, img_metas, self.test_cfg)
            loss_decode = self.decode_head[i].execute_train(
                x, prev_outputs, img_metas, gt_semantic_seg, self.train_cfg)
            losses.update(add_prefix(loss_decode, f'decode_{i}'))

        return losses

python/jittorseg/models/segmentors/cascade_encoder_decoder.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn

from jittorseg.utils.general import add_prefix
from ..layers import resize
from jittorseg.utils.registry import MODELS, build_from_cfg, BACKBONES, NECKS, HEADS
from .base import BaseSegmentor


@MODELS.register_module()
class EncoderDecoder(BaseSegmentor):

    def __init__(self,
                 backbone,
                 decode_head,
                 neck=None,
                 auxiliary_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        super(EncoderDecoder, self).__init__()
        self.backbone = build_from_cfg(backbone, BACKBONES)
        if neck is not None:
            self.neck = build_from_cfg(neck, NECKS)
        self._init_decode_head(decode_head)
        self._init_auxiliary_head(auxiliary_head)

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

        self.init_weights(pretrained=pretrained)

        assert self.with_decode_head

    def _init_decode_head(self, decode_head):
        self.decode_head = build_from_cfg(decode_head, HEADS)
        self.align_corners = self.decode_head.align_corners
        self.num_classes = self.decode_head.num_classes
        self.out_channels = self.decode_head.out_channels

    def _init_auxiliary_head(self, 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(build_from_cfg(head_cfg, HEADS))
            else:
                self.auxiliary_head = build_from_cfg(auxiliary_head, HEADS)

    def init_weights(self, pretrained=None):
        """Initialize the weights in backbone and heads.

        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Defaults to None.
        """

        super(EncoderDecoder, self).init_weights(pretrained)
        self.backbone.init_weights(pretrained=pretrained)
        if isinstance(self.decode_head, nn.ModuleList):
            for decode_head in self.decode_head:
                decode_head.init_weights()
        else:
            self.decode_head.init_weights()
        if self.with_auxiliary_head:
            if isinstance(self.auxiliary_head, nn.ModuleList):
                for aux_head in self.auxiliary_head:
                    aux_head.init_weights()
            else:
                self.auxiliary_head.init_weights()

    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_execute_test(x, img_metas)
        out = resize(input=out,
                     size=img.shape[2:],
                     mode='bilinear',
                     align_corners=self.align_corners)
        return out

    def _decode_head_execute_train(self, x, img_metas, gt_semantic_seg):
        losses = dict()
        loss_decode = self.decode_head.execute_train(x, img_metas,
                                                     gt_semantic_seg,
                                                     self.train_cfg)

        losses.update(add_prefix(loss_decode, 'decode'))
        return losses

    def _decode_head_execute_test(self, x, img_metas):
        seg_logits = self.decode_head.execute_test(x, img_metas, self.test_cfg)
        return seg_logits

    def _auxiliary_head_execute_train(self, x, img_metas, gt_semantic_seg):
        losses = dict()
        if isinstance(self.auxiliary_head, nn.ModuleList):
            for idx, aux_head in enumerate(self.auxiliary_head):
                loss_aux = aux_head.execute_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.execute_train(
                x, img_metas, gt_semantic_seg, self.train_cfg)
            losses.update(add_prefix(loss_aux, 'aux'))

        return losses

    def execute_dummy(self, img):
        seg_logit = self.encode_decode(img, None)

        return seg_logit

    def execute_train(self, img, img_metas, gt_semantic_seg):
        x = self.extract_feat(img)

        losses = dict()

        loss_decode = self._decode_head_execute_train(x, img_metas,
                                                      gt_semantic_seg)
        losses.update(loss_decode)
        if self.with_auxiliary_head:
            loss_aux = self._auxiliary_head_execute_train(
                x, img_metas, gt_semantic_seg)
            losses.update(loss_aux)
        return losses

    # TODO refactor
    def slide_inference(self, img, img_meta, rescale):
        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()
        out_channels = self.out_channels
        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 = jt.zeros((batch_size, out_channels, h_img, w_img))
        count_mat = jt.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 += nn.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
        preds = preds / count_mat
        if rescale:
            resize_shape = img_meta['img_shape'][:2]
            preds = preds[:, :, :resize_shape[0], :resize_shape[1]]
            preds = resize(preds,
                           size=img_meta['ori_shape'][:2],
                           mode='bilinear',
                           align_corners=self.align_corners,
                           warning=False)
        return preds

    def whole_inference(self, img, img_meta, rescale):
        seg_logit = self.encode_decode(img, img_meta)
        if rescale:
            resize_shape = img_meta['img_shape'][:2]
            seg_logit = seg_logit[:, :, :resize_shape[0], :resize_shape[1]]
            size = img_meta['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):
        assert self.test_cfg.mode in ['slide', 'whole']
        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 = nn.softmax(seg_logit, dim=1)
        flip = img_meta['flip']
        if flip:
            flip_direction = img_meta['flip_direction']
            assert flip_direction in ['horizontal', 'vertical']
            if flip_direction == 'horizontal':
                output = jt.flip(output, dim=(3, ))
            elif flip_direction == 'vertical':
                output = jt.flip(output, dim=(2, ))

        return output

    def simple_test(self, img, img_meta, rescale=True):
        seg_logit = self.inference(img, img_meta, rescale)
        if self.out_channels == 1:
            seg_pred = (seg_logit >
                        self.decode_head.threshold).to(seg_logit).squeeze(1)
        else:
            seg_pred = seg_logit.argmax(dim=1)[0]
        # unravel batch dim
        seg_pred = list(seg_pred.numpy())
        return seg_pred

    def aug_test(self, imgs, img_metas, rescale=True):
        # 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)
        if self.out_channels == 1:
            seg_pred = (seg_logit >
                        self.decode_head.threshold).to(seg_logit).squeeze(1)
        else:
            seg_pred = seg_logit.argmax(dim=1)[0]
        # unravel batch dim
        seg_pred = list(seg_pred.numpy())
        return seg_pred

python/jittorseg/models/segmentors/encoder_decoder.py
JittorCV/jittorseg
from .lr_scheduler import *
from .optimizer import *
from .prameter_groups_generator import *
from .lr_decay_parameter_groups_generator import *
python/jittorseg/optims/__init__.py
JittorCV/jittorseg
from jittorseg.utils.registry import MODELS
import json


def get_layer_id_for_convnext(var_name, max_layer_id):
    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 = max_layer_id
        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 = max_layer_id
        return layer_id
    else:
        return max_layer_id + 1


def get_stage_id_for_convnext(var_name, max_stage_id):
    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 max_stage_id - 1


def get_layer_id_for_vit(var_name, max_layer_id):
    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.layers'):
        layer_id = int(var_name.split('.')[2])
        return layer_id + 1
    else:
        return max_layer_id - 1


@MODELS.register_module()
def LRDecayParameterGroupsGenerator(named_params,
                                    model,
                                    paramwise_cfg={},
                                    logger=None):
    num_layers = paramwise_cfg.get('num_layers') + 2
    decay_rate = paramwise_cfg.get('decay_rate')
    decay_type = paramwise_cfg.get('decay_type', 'layer_wise')

    parameter_groups = {}
    normal_group_list = []
    custom_group_list = []
    for p in named_params:
        name, param = p
        if not param.requires_grad:
            normal_group_list.append({'params': [param]})
            continue
        if len(param.shape) == 1 or name.endswith('.bias') or name in (
                'pos_embed', 'cls_token'):
            group_name = 'no_decay'
            decay_mult = 0.
        else:
            group_name = 'decay'
            decay_mult = 1
        if 'layer_wise' in decay_type:
            if 'ConvNeXt' in model.backbone.__class__.__name__:
                layer_id = get_layer_id_for_convnext(
                    name, paramwise_cfg.get('num_layers'))
            elif 'BEiT' in model.backbone.__class__.__name__ or \
                    'MAE' in model.backbone.__class__.__name__:
                layer_id = get_layer_id_for_vit(name, num_layers)
            else:
                raise NotImplementedError()
        elif decay_type == 'stage_wise':
            if 'ConvNeXt' in model.backbone.__class__.__name__:
                layer_id = get_stage_id_for_convnext(name, num_layers)
            else:
                raise NotImplementedError()
        scale = decay_rate**(num_layers - layer_id - 1)
        group_name = f'layer_{layer_id}_{group_name}'
        if group_name not in parameter_groups.keys():
            parameter_groups[group_name] = {
                'decay_mult': decay_mult,
                'lr_mult': scale,
                'param_names': []
            }
        parameter_groups[group_name]['param_names'].append(name)
        custom_group_list.append({
            'decay_mult': decay_mult,
            'params': [param],
            'lr_mult': scale,
        })
    logger.log({'parameter_groups': json.dumps(parameter_groups, indent=2)})
    return normal_group_list + custom_group_list

python/jittorseg/optims/lr_decay_parameter_groups_generator.py
JittorCV/jittorseg
# reference: https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/hooks/lr_updater.py
from jittorseg.utils.registry import SCHEDULERS
import math


@SCHEDULERS.register_module()
class WarmUpLR(object):
    def __init__(self,
                 optimizer,
                 warmup_ratio=1.0 / 3,
                 warmup_iters=500,
                 warmup=None):
        self.optimizer = optimizer
        self.warmup_ratio = warmup_ratio
        self.warmup_iters = warmup_iters
        self.warmup = warmup
        self.base_lr = optimizer.lr
        self.base_lr_pg = [
            pg.get("lr", optimizer.lr) for pg in optimizer.param_groups
        ]
        self.step(0, 0)

    def get_warmup_lr(self, lr, cur_iters):
        if self.warmup == 'constant':
            k = self.warmup_ratio
        elif self.warmup == 'linear':
            k = 1 - (1 - cur_iters / self.warmup_iters) * (1 -
                                                           self.warmup_ratio)
        elif self.warmup == 'exp':
            k = self.warmup_ratio**(1 - cur_iters / self.warmup_iters)
        return k * self.get_lr(lr, cur_iters)

    def get_lr(self, lr, steps):
        return lr

    def _update_lr(self, steps, get_lr_func):
        self.optimizer.lr = get_lr_func(self.base_lr, steps)
        for i, param_group in enumerate(self.optimizer.param_groups):
            param_group["lr"] = get_lr_func(self.base_lr_pg[i], steps)

    def step(self, iters, epochs, by_epoch=True):
        if self.warmup is not None:
            if iters >= self.warmup_iters:
                if by_epoch:
                    self._update_lr(epochs, self.get_lr)
                else:
                    self._update_lr(iters, self.get_lr)
            else:
                self._update_lr(iters, self.get_warmup_lr)
        else:
            if by_epoch:
                self._update_lr(epochs, self.get_lr)
            else:
                self._update_lr(iters, self.get_lr)

    def parameters(self):
        return {
            key: value
            for key, value in self.__dict__.items() if key != 'optimizer'
        }

    def load_parameters(self, data):
        if isinstance(data, dict):
            for k, d in data.items():
                if k in self.__dict__:
                    self.__dict__[k] = d


@SCHEDULERS.register_module()
class WarmUpLRGroup(object):
    def __init__(self,
                 optimizer,
                 warmup_ratio=1.0 / 3,
                 warmup_ratio_pg=None,
                 warmup_init_lr_pg=None,
                 warmup_iters=500,
                 warmup=None,
                 warmup_pg=None,
                 warmup_initial_momentum=None):
        self.optimizer = optimizer
        self.base_lr = optimizer.lr
        self.base_lr_pg = [
            pg.get("lr", optimizer.lr) for pg in optimizer.param_groups
        ]
        self.warmup = warmup
        self.warmup_ratio = warmup_ratio

        self.warmup_pg = warmup_pg if warmup_pg is not None else [
            warmup for _ in self.base_lr_pg
        ]

        if warmup_init_lr_pg or warmup_ratio_pg:
            assert (warmup_init_lr_pg is None) ^ (warmup_ratio_pg is
                                                  None), 'only one can be set'
            if warmup_init_lr_pg:
                self.warmup_ratio_pg = [
                    warmup_init_lr_pg[i] / self.base_lr_pg[i]
                    for i in range(len(warmup_init_lr_pg))
                ]
            else:
                self.warmup_ratio_pg = warmup_ratio_pg
            assert len(self.warmup_ratio_pg) == len(
                self.base_lr_pg), 'these two must be the same'
        else:
            self.warmup_ratio_pg = [self.warmup_ratio for _ in self.base_lr_pg]

        if warmup_initial_momentum is not None:
            self.warmup_initial_momentum = warmup_initial_momentum
            self.base_mom_pg = [
                pg.get('momentum', -1.) for pg in optimizer.param_groups
            ]
        else:
            self.warmup_initial_momentum = None
        self.warmup_iters = warmup_iters
        self.step(0, 0)

    def get_warmup_lr(self, lr, cur_iters, warmup=None, ratio=None, epochs=0):
        if warmup == 'constant':
            k = ratio
        elif warmup == 'linear':
            # in the linear case, the ratio is adjusted based on current learning rate.
            ratio *= lr / self.get_lr(lr, epochs)
            k = 1 - (1 - cur_iters / self.warmup_iters) * (1 - ratio)
        elif warmup == 'exp':
            k = ratio**(1 - cur_iters / self.warmup_iters)
        return k * lr

    def get_warmup_mom(self, momentum, cur_iters):
        # only supports linear
        return self.warmup_initial_momentum + (
            momentum - self.warmup_initial_momentum) * (cur_iters /
                                                        self.warmup_iters)

    def get_lr(self, lr, steps, warmup=None, ratio=None):
        return lr

    def _update_lr(self, steps, get_lr_func):
        self.optimizer.lr = get_lr_func(self.base_lr, steps)
        for i, param_group in enumerate(self.optimizer.param_groups):
            param_group["lr"] = get_lr_func(self.base_lr_pg[i], steps)

    def _update_mom(self, steps, get_mom_func):
        for i, param_group in enumerate(self.optimizer.param_groups):
            if 'momentum' in param_group:
                param_group['momentum'] = get_mom_func(self.base_mom_pg[i],
                                                       steps)

    def _update_warmup_lr(self, steps, get_lr_func, epochs=0):
        self.optimizer.lr = get_lr_func(self.base_lr,
                                        steps,
                                        warmup=self.warmup,
                                        ratio=self.warmup_ratio)
        for i, param_group in enumerate(self.optimizer.param_groups):
            param_group["lr"] = get_lr_func(self.base_lr_pg[i],
                                            steps,
                                            warmup=self.warmup_pg[i],
                                            ratio=self.warmup_ratio_pg[i],
                                            epochs=epochs)

    def step(self, iters, epochs, by_epoch=True):
        if self.warmup is not None:
            if iters >= self.warmup_iters:
                if by_epoch:
                    self._update_lr(epochs, self.get_lr)
                else:
                    self._update_lr(iters - self.warmup_iters, self.get_lr)
            else:
                self._update_warmup_lr(iters,
                                       self.get_warmup_lr,
                                       epochs=epochs)
                if self.warmup_initial_momentum:
                    self._update_mom(iters, self.get_warmup_mom)
        else:
            if by_epoch:
                self._update_lr(epochs, self.get_lr)
            else:
                self._update_lr(iters, self.get_lr)

    def parameters(self):
        return {
            key: value
            for key, value in self.__dict__.items() if key != 'optimizer'
        }

    def load_parameters(self, data):
        if isinstance(data, dict):
            for k, d in data.items():
                if k in self.__dict__:
                    self.__dict__[k] = d


@SCHEDULERS.register_module()
class StepLR(WarmUpLR):
    def __init__(self, milestones, gamma=0.1, min_lr=None, **kwargs):
        if isinstance(milestones, list):
            assert all([s > 0 for s in milestones])
        elif isinstance(milestones, int):
            assert milestones > 0
        else:
            raise TypeError('"step" must be a list or integer')
        self.milestones = milestones
        self.gamma = gamma
        self.min_lr = min_lr
        super(StepLR, self).__init__(**kwargs)

    def get_lr(self, base_lr, steps):
        # calculate exponential term
        if isinstance(self.milestones, int):
            exp = steps // self.milestones
        else:
            exp = len(self.milestones)
            for i, s in enumerate(self.milestones):
                if steps < s:
                    exp = i
                    break

        lr = base_lr * (self.gamma**exp)
        if self.min_lr is not None:
            # clip to a minimum value
            lr = max(lr, self.min_lr)
        return lr


@SCHEDULERS.register_module()
class CosineAnnealingLR(WarmUpLR):
    def __init__(self, max_steps, min_lr=None, min_lr_ratio=None, **kwargs):
        assert (min_lr is None) ^ (min_lr_ratio is None)
        self.min_lr = min_lr
        self.min_lr_ratio = min_lr_ratio
        self.max_steps = max_steps
        super(CosineAnnealingLR, self).__init__(**kwargs)

    def get_lr(self, base_lr, steps):
        if self.min_lr_ratio is not None:
            target_lr = base_lr * self.min_lr_ratio
        else:
            target_lr = self.min_lr
        cos_out = math.cos(math.pi * (steps / self.max_steps)) + 1
        lr = target_lr + 0.5 * (base_lr - target_lr) * cos_out
        return lr


@SCHEDULERS.register_module()
class CosineAnnealingLRGroup(WarmUpLRGroup):
    def __init__(self, max_steps, min_lr=None, min_lr_ratio=None, **kwargs):
        assert (min_lr is None) ^ (min_lr_ratio is None)
        self.min_lr = min_lr
        self.min_lr_ratio = min_lr_ratio
        self.max_steps = max_steps
        super(CosineAnnealingLRGroup, self).__init__(**kwargs)

    def get_lr(self, base_lr, steps):
        if self.min_lr_ratio is not None:
            target_lr = base_lr * self.min_lr_ratio
        else:
            target_lr = self.min_lr
        cos_out = math.cos(math.pi * (steps / self.max_steps)) + 1
        lr = target_lr + 0.5 * (base_lr - target_lr) * cos_out
        return lr


@SCHEDULERS.register_module()
class ExpLR(WarmUpLR):
    def __init__(self, gamma, **kwargs):
        self.gamma = gamma
        super(ExpLR, self).__init__(**kwargs)

    def get_lr(self, base_lr, steps):
        return base_lr * self.gamma**steps


@SCHEDULERS.register_module()
class PolyLR(WarmUpLR):
    def __init__(self, max_steps, power=1., min_lr=0., **kwargs):
        self.power = power
        self.min_lr = min_lr
        self.max_steps = max_steps
        super(PolyLR, self).__init__(**kwargs)

    def get_lr(self, base_lr, steps):
        coeff = (1 - steps / self.max_steps)**self.power
        return (base_lr - self.min_lr) * coeff + self.min_lr


@SCHEDULERS.register_module()
class InvLR(WarmUpLR):
    def __init__(self, gamma, power=1., **kwargs):
        self.gamma = gamma
        self.power = power
        super(InvLR, self).__init__(**kwargs)

    def get_lr(self, base_lr, steps):
        return base_lr * (1 + self.gamma * steps)**(-self.power)

python/jittorseg/optims/lr_scheduler.py
JittorCV/jittorseg
from jittorseg.utils.registry import OPTIMS
import jittor as jt
from jittor import optim


class Optimizer(object):
    def parameters(self):
        data = {}
        for k, d in self.__dict__.items():
            if k == "param_groups":
                continue
            data[k] = d
        return data

    def load_parameters(self, data):
        if isinstance(data, dict):
            for k, d in data.items():
                if k in self.__dict__:
                    self.__dict__[k] = d

    def cur_lr(self):
        return self.param_groups[0].get("lr", self.lr)


@OPTIMS.register_module()
class SGD(optim.SGD, Optimizer):
    def __init__(self,
                 params,
                 lr,
                 momentum=0,
                 weight_decay=0,
                 dampening=0,
                 nesterov=False,
                 grad_clip=None):
        super(SGD, self).__init__(params, lr, momentum, weight_decay,
                                  dampening, nesterov)
        self.grad_clip = grad_clip

    def pre_step(self, loss, retain_graph=False):
        super(SGD, self).pre_step(loss)
        if self.grad_clip is not None:
            self.clip_grad_norm(**self.grad_clip)


@OPTIMS.register_module()
class GradMutilpySGD(optim.SGD, Optimizer):
    def __init__(self, grad_clip=None, **kwargs):
        super(GradMutilpySGD, self).__init__(**kwargs)
        self.grad_clip = grad_clip

    def step(self, loss):
        if loss is not None:
            self.pre_step(loss)
        if self.grad_clip is not None:
            self.clip_grad_norm(**self.grad_clip)
        for pg in self.param_groups:
            # get arguments from each param_groups
            lr = pg.get("lr", self.lr)
            momentum = pg.get("momentum", self.momentum)
            weight_decay = pg.get("weight_decay", self.weight_decay)
            dampening = pg.get("dampening", self.dampening)
            nesterov = pg.get("nesterov", self.nesterov)

            m = pg.get("grad_mutilpy", 1)
            # optimize main body
            for p, g, v in zip(pg["params"], pg["grads"], pg["values"]):
                if p.is_stop_grad():
                    continue
                dp = p * weight_decay + g * m
                v.update(momentum * v + dp * (1 - dampening))
                if nesterov:
                    p.update(p - (dp + momentum * v) * lr)
                else:
                    p.update(p - v * lr)
        self.zero_grad()


@OPTIMS.register_module()
class AdamW(optim.AdamW, Optimizer):
    def __init__(self,
                 params,
                 lr,
                 eps=1e-8,
                 betas=(0.9, 0.999),
                 weight_decay=0):
        super(AdamW, self).__init__(params, lr, eps, betas, weight_decay)

    def pre_step(self, loss, retain_graph=False):
        super(AdamW, self).pre_step(loss)


@OPTIMS.register_module()
class CustomAdamW(AdamW, Optimizer):
    def __init__(self,
                 params,
                 lr,
                 eps=1e-8,
                 betas=(0.9, 0.999),
                 weight_decay=0):
        super(CustomAdamW, self).__init__(params, lr, eps, betas, weight_decay)

    def step(self, loss):
        if loss is not None:
            self.pre_step(loss)

        n = float(self.n_step)
        for pg in self.param_groups:
            # get arguments from each param_groups
            lr = pg.get("lr", self.lr) * pg.get("lr_mult", 1)
            eps = pg.get("eps", self.eps)
            weight_decay = pg.get("weight_decay", self.weight_decay) * pg.get(
                "decay_mult", 1)
            b0, b1 = pg.get("betas", self.betas)
            for p, g, v, m in zip(pg["params"], pg["grads"], pg["values"],
                                  pg["m"]):
                if p.is_stop_grad():
                    continue
                p.update(p * (1 - lr * weight_decay))
                bias_correction1 = 1 - b0**n
                bias_correction2 = 1 - b1**n
                m.update(b0 * m + (1 - b0) * g)  # exp_avg
                v.update(b1 * v + (1 - b1) * g * g)  # exp_avg_sq
                denom = jt.sqrt(v) / jt.sqrt(bias_correction2) + eps
                step_size = lr / bias_correction1
                p.update(p - step_size * m / denom)
        self.post_step()

python/jittorseg/optims/optimizer.py
JittorCV/jittorseg
from jittorseg.utils.registry import MODELS


@MODELS.register_module()
def CustomPrameterGroupsGenerator(named_params, model, custom_keys={}, logger=None):
    def get_custom_parameter_groups(name):
        for ck in custom_keys.keys():
            if ck in name:
                return custom_keys[ck]
        return None

    normal_group_list = []
    custom_group_list = []

    for p in named_params:
        name, param = p
        custom_group = get_custom_parameter_groups(name)
        if custom_group is not None:
            tmp = {}
            tmp['params'] = [param]
            for i in custom_group.keys():
                tmp[i] = custom_group.get(i)
            custom_group_list.append(tmp)
            continue
        normal_group_list.append({'params': [param]})
    return normal_group_list + custom_group_list

python/jittorseg/optims/prameter_groups_generator.py
JittorCV/jittorseg
from .train_runner import TrainRunner
from .eval_runner import EvalRunner
from .test_runner import TestRunner

python/jittorseg/runner/__init__.py
JittorCV/jittorseg
import jittor as jt
import jittorseg
from jittorseg.config import get_cfg, save_cfg, print_cfg
from jittorseg.utils.registry import build_from_cfg, MODELS, SCHEDULERS, HOOKS, OPTIMS
from jittorseg.utils.general import build_file, current_time, check_file, search_ckpt, is_method_overridden
from jittorseg.hooks import Hook

from abc import ABCMeta, abstractmethod


class BaseRunner(metaclass=ABCMeta):
    def __init__(self):
        cfg = get_cfg()
        self.cfg = cfg
        # self.flip_test = [] if cfg.flip_test is None else cfg.flip_test
        self.work_dir = cfg.work_dir

        self.max_epoch = cfg.max_epoch
        self.max_iter = cfg.max_iter
        assert (self.max_iter is None) ^ (
                self.max_epoch is None), "You must set max_iter or max_epoch"
        self.iter = 0
        self.epoch = 0

        self.resume_path = cfg.resume_path
        self.efficient_val = cfg.efficient_val

        self.logger = build_from_cfg(self.cfg.logger,
                                     HOOKS,
                                     work_dir=self.work_dir)

        self.model = build_from_cfg(cfg.model, MODELS)
        if jt.rank == 0:
            print_cfg()
            self.logger.log({'model': self.model})

        if (cfg.parameter_groups_generator):
            params = build_from_cfg(cfg.parameter_groups_generator,
                                    MODELS,
                                    named_params=self.model.named_parameters(),
                                    model=self.model,
                                    logger=self.logger)
        else:
            params = self.model.parameters()
        self.optimizer = build_from_cfg(cfg.optimizer, OPTIMS, params=params)
        self.scheduler = build_from_cfg(cfg.scheduler,
                                        SCHEDULERS,
                                        optimizer=self.optimizer)
        self.hooks = []
        self.outputs = {}

        save_file = build_file(self.work_dir, prefix="config.yaml")
        save_cfg(save_file)

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=True)

        # TODO: Resuming is auto-executed now. If the `pretrained_weights` is set, the weights from pretrained model
        #  will be overwrited by the resumed one.

        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()

    @abstractmethod
    def run(self, *args, **kwargs):
        pass

    # def test_time(self):
    #     warmup = 10
    #     rerun = 100
    #     self.model.train()
    #     for batch_idx, (data) in enumerate(self.train_dataset):
    #         break
    #     images = data['img']
    #     img_metas = data['img_metas']
    #     gt = data['gt_semantic_seg']
    #     print("warmup...")
    #     for i in tqdm(range(warmup)):
    #         losses = self.model(img=images,
    #                             img_metas=img_metas,
    #                             gt_semantic_seg=gt)
    #         all_loss, losses = parse_losses(losses)
    #         self.optimizer.step(all_loss)
    #         self.scheduler.step(self.iter, self.epoch, by_epoch=True)
    #     jt.sync_all(True)
    #     print("testing...")
    #     start_time = time.time()
    #     for i in tqdm(range(rerun)):
    #         losses = self.model(img=images,
    #                             img_metas=img_metas,
    #                             gt_semantic_seg=gt)
    #         all_loss, losses = parse_losses(losses)
    #         self.optimizer.step(all_loss)
    #         self.scheduler.step(self.iter, self.epoch, by_epoch=True)
    #     jt.sync_all(True)
    #     batch_size = len(gt) * jt.world_size
    #     ptime = time.time() - start_time
    #     fps = batch_size * rerun / ptime
    #     print("FPS:", fps)

    @jt.single_process_scope()
    def save(self):
        save_data = {
            "meta": {
                "jittorseg_version": jittorseg.__version__,
                "epoch": self.epoch,
                "iter": self.iter + 1,
                "max_iter": self.max_iter,
                "max_epoch": self.max_epoch,
                "save_time": current_time(),
                "config": self.cfg.dump()
            },
            "model": self.model.state_dict(),
            "scheduler": self.scheduler.parameters(),
            "optimizer": self.optimizer.parameters()
        }

        save_file = build_file(self.work_dir,
                               prefix=f"checkpoints/ckpt_{self.iter}.pkl")
        jt.save(save_data, save_file)
        print("saved")

    def load(self, load_path, model_only=False):
        resume_data = jt.load(load_path)

        if not model_only:
            meta = resume_data.get("meta", dict())
            self.epoch = meta.get("epoch", self.epoch)
            self.iter = meta.get("iter", self.iter)
            self.max_iter = meta.get("max_iter", self.max_iter)
            self.max_epoch = meta.get("max_epoch", self.max_epoch)
            self.scheduler.load_parameters(resume_data.get(
                "scheduler", dict()))
            self.optimizer.load_parameters(resume_data.get(
                "optimizer", dict()))
        if "model" in resume_data:
            state_dict = resume_data["model"]
        elif "state_dict" in resume_data:
            state_dict = resume_data["state_dict"]
        else:
            state_dict = resume_data
        self.model.load_parameters(state_dict)
        self.logger.print_log(
            f"Missing key: {self.model.state_dict().keys() - state_dict.keys()}"
        )
        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path)

    def register_hook(self, hook_cfg):
        hook = build_from_cfg(hook_cfg, HOOKS, runner=self)
        idx = 0
        ok = False
        for h in self.hooks:
            if hook.priority < h.priority:
                ok = True
                break
            idx += 1
        if ok:
            self.hooks.insert(idx, hook)
        else:
            self.hooks.append(hook)

    def call_hook(self, hook_stage):
        for h in self.hooks:
            getattr(h, hook_stage)()

    def hook_info(self):
        info = {}
        for stage in Hook.stages:
            info[stage] = []
            for h in self.hooks:
                if is_method_overridden(stage, Hook, h):
                    info[stage].append((h.name, h.priority))

        info_str = "\n"
        for k, v in info.items():
            info_str += f"Stage {k}:\n"
            info_str += f"{'Name':^20}|{'Prio':^10}\n"
            info_str += '-' * 30 + '\n'
            for pair in v:
                info_str += f"{pair[0]:^20}|{pair[1]:^10}\n"
            info_str += '-' * 30 + '\n'
        return info_str









python/jittorseg/runner/base_runner.py
JittorCV/jittorseg
from jittorseg.utils.registry import build_from_cfg, DATASETS
from .base_runner import BaseRunner
from jittorseg.hooks import Priority


class EvalRunner(BaseRunner):

    def __init__(self):
        super().__init__()
        self.val_dataset = build_from_cfg(self.cfg.dataset.val, DATASETS)
        self.register_hook(dict(type="EvalHook", priority=Priority.NORMAL, eval_interval=1, efficient_val=self.cfg.efficient_val))

    def run(self):
        self.call_hook("after_iter")



python/jittorseg/runner/eval_runner.py
JittorCV/jittorseg
import jittor as jt
from .base_runner import BaseRunner
from tqdm import tqdm
from jittorseg.utils.registry import build_from_cfg, DATASETS
from jittorseg.utils.visualize import visualize_result


class TestRunner(BaseRunner):

    def __init__(self, save_dir):
        super().__init__()
        self.test_dataset = build_from_cfg(self.cfg.dataset.test, DATASETS)
        self.save_dir = save_dir

    @jt.no_grad()
    @jt.single_process_scope()
    def test(self):
        if self.test_dataset is None:
            self.logger.print_log("Please set Test dataset")
        else:
            self.logger.print_log("Testing...")
            self.model.eval()
            for _, (data) in tqdm(enumerate(self.test_dataset)):
                images = data['img']
                img_metas = data['img_metas']
                results = self.model(images, img_metas, return_loss=False)
                visualize_result(results[0],
                                 palette=self.test_dataset.PALETTE,
                                 save_dir=self.save_dir,
                                 file_name=img_metas[0]['ori_filename'])

    def run(self):
        self.test()


python/jittorseg/runner/test_runner.py
JittorCV/jittorseg
import jittor as jt
from jittorseg.utils.registry import build_from_cfg, DATASETS
from .base_runner import BaseRunner
from jittorseg.hooks import Priority


class TrainRunner(BaseRunner):

    def __init__(self):
        super().__init__()
        self.train_dataset = build_from_cfg(self.cfg.dataset.train,
                                            DATASETS,
                                            drop_last=jt.in_mpi)
        self.val_dataset = build_from_cfg(self.cfg.dataset.val, DATASETS)

        if self.max_epoch:
            if (self.train_dataset):
                self.total_iter = self.max_epoch * len(self.train_dataset)
            else:
                self.total_iter = 0
        else:
            self.total_iter = self.max_iter

        self.register_train_hooks()
        self.register_custom_hooks()
        # if jt.rank == 0:
        self.logger.print_log(self.hook_info())

    def run(self):
        self.call_hook("before_run")
        while not self.finish:
            self.train()
        self.call_hook("after_run")

    @property
    def finish(self):
        if self.max_epoch:
            return self.epoch >= self.max_epoch
        else:
            return self.iter >= self.max_iter

    def train(self):
        self.model.train()
        self.call_hook("before_epoch")
        for batch_idx, (data) in enumerate(self.train_dataset):
            images = data['img']
            img_metas = data['img_metas']
            gt = data['gt_semantic_seg']
            self.outputs["batch_size"] = len(gt) * jt.world_size
            self.outputs["batch_idx"] = batch_idx
            self.outputs["data"] = data
            self.call_hook("before_iter")
            losses = self.model(img=images,
                                img_metas=img_metas,
                                gt_semantic_seg=gt)
            self.outputs["losses"] = losses
            self.call_hook("after_iter")
            self.iter += 1

            if self.finish:
                break
        self.call_hook("after_epoch")
        self.epoch += 1

    def register_train_hooks(self):
        self.register_hook(dict(type="LossParserHook", priority=Priority.VERY_HIGH))
        self.register_hook(dict(type="OptimizerHook", priority=Priority.HIGH))
        self.register_hook(dict(type="LRSchedulerHook", priority=Priority.ABOVE_NORMAL, by_epoch=False))
        self.register_hook(dict(type="LogPrinterHook", priority=Priority.NORMAL, log_interval=self.cfg.log_interval))
        self.register_hook(dict(type="CheckpointHook", priority=Priority.BELOW_NORMAL,
                                checkpoint_interval=self.cfg.checkpoint_interval))
        self.register_hook(dict(type="EvalHook", priority=Priority.LOW, eval_interval=self.cfg.eval_interval,
                                efficient_val=self.cfg.efficient_val))

    def register_custom_hooks(self):
        if self.cfg.custom_hooks is not None:
            assert isinstance(self.cfg.custom_hooks, list), "`custom_hooks` should be a list"
            for h in self.cfg.custom_hooks:
                self.register_hook(h)

python/jittorseg/runner/train_runner.py
JittorCV/jittorseg
import jittor as jt
from jittor import nn

from jittorseg.utils.registry import PIXEL_SAMPLERS
from .base_pixel_sampler import BasePixelSampler


@PIXEL_SAMPLERS.register_module()
class OHEMPixelSampler(BasePixelSampler):
    def __init__(self, context, thresh=None, min_kept=100000):
        super(OHEMPixelSampler, self).__init__()
        self.context = context
        assert min_kept > 1
        self.thresh = thresh
        self.min_kept = min_kept

    def sample(self, seg_logit, seg_label):
        with jt.no_grad():
            assert seg_logit.shape[2:] == seg_label.shape[2:]
            assert seg_label.shape[1] == 1
            seg_label = seg_label.squeeze(1).long()
            batch_kept = self.min_kept * seg_label.size(0)
            valid_mask = seg_label != self.context.ignore_index
            seg_weight = jt.zeros(seg_label.size()).astype(seg_logit.dtype)
            valid_seg_weight = seg_weight[valid_mask]
            if self.thresh is not None:
                seg_prob = nn.softmax(seg_logit, dim=1)

                tmp_seg_label = seg_label.clone().unsqueeze(1)
                tmp_seg_label[tmp_seg_label == self.context.ignore_index] = 0
                seg_prob = jt.gather(seg_prob, 1, tmp_seg_label).squeeze(1)
                sort_indices, sort_prob = seg_prob[valid_mask].argsort()

                if sort_prob.numel() > 0:
                    min_threshold = sort_prob[min(batch_kept,
                                                  sort_prob.numel() - 1)]
                else:
                    min_threshold = 0.0
                threshold = max(min_threshold, self.thresh)
                valid_seg_weight[seg_prob[valid_mask] < threshold] = 1.
            else:
                if not isinstance(self.context.loss_decode, nn.ModuleList):
                    losses_decode = [self.context.loss_decode]
                else:
                    losses_decode = self.context.loss_decode
                losses = 0.0
                for loss_module in losses_decode:
                    losses += loss_module(
                        seg_logit,
                        seg_label,
                        weight=None,
                        ignore_index=self.context.ignore_index,
                        reduction_override='none')

                sort_indices, _ = losses[valid_mask].argsort(descending=True)
                valid_seg_weight[sort_indices[:batch_kept]] = 1.

            seg_weight[valid_mask] = valid_seg_weight

            return seg_weight

python/jittorseg/sampler/ohem_pixel_sampler.py
JittorCV/jittorseg
from .logger import *
from .registry import Registry, build_from_cfg
python/jittorseg/utils/__init__.py
JittorCV/jittorseg
import jittor as jt
import time
import warnings
import numpy as np
import random
import os
import glob
from functools import partial
from six.moves import map, zip
import numpy
import tempfile
from pathlib import Path
import os.path as osp
from collections import abc


def is_seq_of(seq, expected_type, seq_type=None):
    if seq_type is None:
        exp_seq_type = abc.Sequence
    else:
        assert isinstance(seq_type, type)
        exp_seq_type = seq_type
    if not isinstance(seq, exp_seq_type):
        return False
    for item in seq:
        if not isinstance(item, expected_type):
            return False
    return True


def is_tuple_of(seq, expected_type):
    return is_seq_of(seq, expected_type, seq_type=tuple)


def is_list_of(seq, expected_type):
    return is_seq_of(seq, expected_type, seq_type=list)


def list_from_file(filename, prefix='', offset=0, max_num=0, encoding='utf-8'):
    cnt = 0
    item_list = []
    with open(filename, 'r', encoding=encoding) as f:
        for _ in range(offset):
            f.readline()
        for line in f:
            if 0 < max_num <= cnt:
                break
            item_list.append(prefix + line.rstrip('\n\r'))
            cnt += 1
    return item_list


def mkdir_or_exist(dir_name, mode=0o777):
    if dir_name == '':
        return
    dir_name = osp.expanduser(dir_name)
    os.makedirs(dir_name, mode=mode, exist_ok=True)


def scandir(dir_path, suffix=None, recursive=False):
    if isinstance(dir_path, (str, Path)):
        dir_path = str(dir_path)
    else:
        raise TypeError('"dir_path" must be a string or Path object')

    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
        raise TypeError('"suffix" must be a string or tuple of strings')

    root = dir_path

    def _scandir(dir_path, suffix, recursive):
        for entry in os.scandir(dir_path):
            if not entry.name.startswith('.') and entry.is_file():
                rel_path = osp.relpath(entry.path, root)
                if suffix is None or rel_path.endswith(suffix):
                    yield rel_path
            elif recursive and os.path.isdir(entry.path):
                # scan recursively if entry.path is a directory
                yield from _scandir(entry.path,
                                    suffix=suffix,
                                    recursive=recursive)

    return _scandir(dir_path, suffix=suffix, recursive=recursive)


def to_jt_var(data):
    """
        convert data to jt_array
    """

    def _to_jt_var(data):
        if isinstance(data, (list, tuple)):
            data = [_to_jt_var(d) for d in data]
        elif isinstance(data, dict):
            data = {k: _to_jt_var(d) for k, d in data.items()}
        elif isinstance(data, numpy.ndarray):
            data = jt.array(data)
        elif not isinstance(data, (int, float, str, np.ndarray)):
            raise ValueError(f"{type(data)} is not supported")
        return data

    return _to_jt_var(data)


def sync(data, reduce_mode="mean", to_numpy=True):
    def _sync(data):
        if isinstance(data, (list, tuple)):
            data = [_sync(d) for d in data]
        elif isinstance(data, dict):
            data = {k: _sync(d) for k, d in data.items()}
        elif isinstance(data, jt.Var):
            if jt.in_mpi:
                data = data.mpi_all_reduce(reduce_mode)
            if to_numpy:
                data = data.numpy()
        elif not isinstance(data, (int, float, str, np.ndarray)):
            raise ValueError(f"{type(data)} is not supported")
        return data

    return _sync(data)


def multi_apply(func, *args, **kwargs):
    pfunc = partial(func, **kwargs) if kwargs else func
    map_results = map(pfunc, *args)
    return tuple(map(list, zip(*map_results)))


def unmap(data, count, inds, fill=0):
    if data.ndim == 1:
        ret = jt.full((count,), fill, dtype=data.dtype)
        ret[inds] = data
    else:
        new_size = (count,) + data.size()[1:]
        ret = jt.full(new_size, fill, dtype=data.dtype)
        ret[inds, :] = data
    return ret


def parse_losses(losses):
    _losses = dict()
    for loss_name, loss_value in losses.items():
        if isinstance(loss_value, jt.Var):
            _losses[loss_name] = loss_value.mean()
        elif isinstance(loss_value, list):
            _losses[loss_name] = sum(_loss.mean() for _loss in loss_value)
        else:
            raise TypeError(
                '{} is not a Var or list of Vars'.format(loss_name))

    total_loss = sum(_value for _key, _value in _losses.items()
                     if 'loss' in _key)
    return total_loss, _losses


def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    jt.seed(seed)


def current_time():
    return time.asctime(time.localtime(time.time()))


def check_file(file, ext=None):
    if file is None:
        return False
    if not os.path.exists(file):
        warnings.warn(f"{file} is not exists")
        return False
    if not os.path.isfile(file):
        warnings.warn(f"{file} must be a file")
        return False
    if ext:
        if not os.path.splitext(file)[1] in ext:
            # warnings.warn(f"the type of {file} must be in {ext}")
            return False
    return True


def build_file(work_dir, prefix):
    work_dir = os.path.abspath(work_dir)
    prefixes = prefix.split("/")
    file_name = prefixes[-1]
    prefix = "/".join(prefixes[:-1])
    if len(prefix) > 0:
        work_dir = os.path.join(work_dir, prefix)
    os.makedirs(work_dir, exist_ok=True)
    file = os.path.join(work_dir, file_name)
    return file


def check_interval(step, step_interval):
    if step is None or step_interval is None:
        return False
    if step and step % step_interval == 0:
        return True
    return False


def check_dir(work_dir):
    os.makedirs(work_dir, exist_ok=True)


def list_files(file_dir):
    if os.path.isfile(file_dir):
        return [file_dir]

    filenames = []
    for f in os.listdir(file_dir):
        ff = os.path.join(file_dir, f)
        if os.path.isfile(ff):
            filenames.append(ff)
        elif os.path.isdir(ff):
            filenames.extend(list_files(ff))

    return filenames


def is_img(f):
    ext = os.path.splitext(f)[1]
    return ext.lower() in [".jpg", ".bmp", ".jpeg", ".png", "tiff"]


def list_images(img_dir):
    img_files = []
    for img_d in img_dir.split(","):
        if len(img_d) == 0:
            continue
        if not os.path.exists(img_d):
            raise f"{img_d} not exists"
        img_d = os.path.abspath(img_d)
        img_files.extend([f for f in list_files(img_d) if is_img(f)])
    return img_files


def search_ckpt(work_dir):
    files = glob.glob(os.path.join(work_dir, "checkpoints/ckpt_*.pkl"))
    if len(files) == 0:
        return None
    files = sorted(files, key=lambda x: int(x.split("_")[-1].split(".pkl")[0]))
    return files[-1]


def add_prefix(inputs, prefix):
    outputs = dict()
    for name, value in inputs.items():
        outputs[f'{prefix}.{name}'] = value

    return outputs


def np2tmp(array, temp_file_name=None):
    if temp_file_name is None:
        temp_file_name = tempfile.NamedTemporaryFile(suffix='.npy',
                                                     delete=False).name
    np.save(temp_file_name, array)
    return temp_file_name


def is_method_overridden(method, base_class, derived_class):
    """Check if a method of base class is overridden in derived class.

    Args:
        method (str): the method name to check.
        base_class (type): the class of the base class.
        derived_class (type | Any): the class or instance of the derived class.
    """
    assert isinstance(base_class, type), \
        "base_class doesn't accept instance, Please pass class instead."

    if not isinstance(derived_class, type):
        derived_class = derived_class.__class__

    base_method = getattr(base_class, method)
    derived_method = getattr(derived_class, method)
    return derived_method != base_method

python/jittorseg/utils/general.py
JittorCV/jittorseg
import cv2
import os
import jittor as jt
from jittorseg.runner import TestRunner
from jittorseg.config import init_cfg, update_cfg, get_cfg
from jittorseg.datasets.pipelines import Compose


class InferenceSegmentor:
    def __init__(self, config_file, checkpoint_file, save_dir):
        init_cfg(config_file)
        if len(checkpoint_file) > 0:
            update_cfg(resume_path=checkpoint_file)

        self.runner = TestRunner(save_dir=save_dir)
        self.runner.model.eval()
        self.transforms = Compose(get_cfg().test_pipeline[1:])
        self.palette = self.runner.test_dataset.PALETTE
        self.runner.model.CLASSES = self.runner.test_dataset.CLASSES
        self.runner.model.PALETTE = self.runner.test_dataset.PALETTE
        self.save_dir = save_dir

    def load_img(self, results):
        if isinstance(results['img'], str):
            results['filename'] = results['img']
            results['ori_filename'] = results['img']
        else:
            results['filename'] = None
            results['ori_filename'] = None
        img = cv2.imread(results['img'])
        results['img'] = img
        results['img_shape'] = img.shape
        results['ori_shape'] = img.shape
        return results

    @jt.no_grad()
    @jt.single_process_scope()
    def infer(self, img):
        data = dict(img=img)
        data = self.transforms(self.load_img(data))
        data['img'][0] = data['img'][0].unsqueeze(0)
        results = self.runner.model(**data, return_loss=False, rescale=True)
        results = self.runner.model.show_result(img, results, out_file=os.path.join(self.save_dir, img[img.rfind('/') + 1:]))
        return results

python/jittorseg/utils/inference.py
JittorCV/jittorseg
from jittorseg.utils.general import build_file, current_time
from .registry import HOOKS, build_from_cfg
import time
import os
from tensorboardX import SummaryWriter
from jittorseg.config import get_cfg


@HOOKS.register_module()
class TextLogger:
    def __init__(self, work_dir):
        save_file = build_file(
            work_dir,
            prefix="textlog/log_" +
            time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime()) + ".txt")
        self.log_file = open(save_file, "a")

    def log(self, data):
        msg = ",".join([f"{k}:{d}" for k, d in data.items()])
        msg = current_time() + ' ' + msg + "\n"
        self.log_file.write(msg)
        self.log_file.flush()


@HOOKS.register_module()
class TensorboardLogger:
    def __init__(self, work_dir):
        self.cfg = get_cfg()
        tensorboard_dir = os.path.join(work_dir, "tensorboard")
        self.writer = SummaryWriter(tensorboard_dir, flush_secs=10)

    def log(self, data):
        if "iter" in data.keys():
            step = data["iter"]
            for k, d in data.items():
                if k in ["iter", "epoch", "batch_idx", "times", "batch_size"]:
                    continue
                if isinstance(d, str):
                    continue
                self.writer.add_scalar(k, d, global_step=step)


@HOOKS.register_module()
class RunLogger:
    def __init__(self, work_dir, loggers=["TextLogger", "TensorboardLogger"]):
        self.loggers = [
            build_from_cfg(log, HOOKS, work_dir=work_dir) for log in loggers
        ]

    def log(self, data, **kwargs):
        data.update(kwargs)
        data = {
            k: d.item() if hasattr(d, "item") else d
            for k, d in data.items()
        }
        for logger in self.loggers:
            logger.log(data)
        self.print_log(data)

    def get_time(self, s):
        s = int(s)
        days = s // 60 // 60 // 24
        hours = s // 60 // 60 % 24
        minutes = s // 60 % 60
        seconds = s % 60
        return f' [{days}D:{hours}H:{minutes}M:{seconds}S] '

    def print_log(self, msg):
        if isinstance(msg, dict):
            msgs = []
            for k, d in msg.items():
                if (k == "remain_time"):
                    msgs.append(f" {k}:{self.get_time(d)}")
                else:
                    msgs.append(f" {k}:{d:.7f}"
                                if isinstance(d, float) else f" {k}:{d}")
            msg = ",".join(msgs)
        print(current_time(), msg)

python/jittorseg/utils/logger.py
JittorCV/jittorseg
import numpy as np
from collections import OrderedDict
from PIL import Image


def f_score(precision, recall, beta=1):
    score = (1 + beta**2) * (precision * recall) / (
        (beta**2 * precision) + recall)
    return score


def intersect_and_union(pred_label,
                        label,
                        num_classes,
                        ignore_index,
                        label_map=dict(),
                        reduce_zero_label=False):
    if isinstance(pred_label, str):
        pred_label = np.load(pred_label)
    if isinstance(label, str):
        label = Image.open(label)
        label = np.array(label)
        if label.ndim >= 3 and label.shape[2] >= 3:  # color image
            label[:, :, :3] = label[:, :, (2, 1, 0)]  # RGB to BGR

    if label_map is not None:
        label_copy = label.copy()
        for old_id, new_id in label_map.items():
            label[label_copy == old_id] = new_id
    if reduce_zero_label:
        label[label == 0] = 255
        label = label - 1
        label[label == 254] = 255

    mask = (label != ignore_index)
    pred_label = pred_label[mask]
    label = label[mask]

    intersect = pred_label[pred_label == label]
    area_intersect, _ = np.histogram(intersect,
                                     bins=np.arange(num_classes + 1))
    area_pred_label, _ = np.histogram(pred_label,
                                      bins=np.arange(num_classes + 1))
    area_label, _ = np.histogram(label, bins=np.arange(num_classes + 1))
    area_union = area_pred_label + area_label - area_intersect
    return area_intersect, area_union, area_pred_label, area_label


def total_intersect_and_union(results,
                              gt_seg_maps,
                              num_classes,
                              ignore_index,
                              label_map=dict(),
                              reduce_zero_label=False):
    total_area_intersect = np.zeros((num_classes, ), dtype=np.float64)
    total_area_union = np.zeros((num_classes, ), dtype=np.float64)
    total_area_pred_label = np.zeros((num_classes, ), dtype=np.float64)
    total_area_label = np.zeros((num_classes, ), dtype=np.float64)
    for result, gt_seg_map in zip(results, gt_seg_maps):
        area_intersect, area_union, area_pred_label, area_label = \
            intersect_and_union(
                result, gt_seg_map, num_classes, ignore_index,
                label_map, reduce_zero_label)
        total_area_intersect += area_intersect
        total_area_union += area_union
        total_area_pred_label += area_pred_label
        total_area_label += area_label
    return total_area_intersect, total_area_union, total_area_pred_label, \
        total_area_label


def mean_iou(results,
             gt_seg_maps,
             num_classes,
             ignore_index,
             nan_to_num=None,
             label_map=dict(),
             reduce_zero_label=False):
    iou_result = eval_metrics(results=results,
                              gt_seg_maps=gt_seg_maps,
                              num_classes=num_classes,
                              ignore_index=ignore_index,
                              metrics=['mIoU'],
                              nan_to_num=nan_to_num,
                              label_map=label_map,
                              reduce_zero_label=reduce_zero_label)
    return iou_result


def mean_dice(results,
              gt_seg_maps,
              num_classes,
              ignore_index,
              nan_to_num=None,
              label_map=dict(),
              reduce_zero_label=False):
    dice_result = eval_metrics(results=results,
                               gt_seg_maps=gt_seg_maps,
                               num_classes=num_classes,
                               ignore_index=ignore_index,
                               metrics=['mDice'],
                               nan_to_num=nan_to_num,
                               label_map=label_map,
                               reduce_zero_label=reduce_zero_label)
    return dice_result


def mean_fscore(results,
                gt_seg_maps,
                num_classes,
                ignore_index,
                nan_to_num=None,
                label_map=dict(),
                reduce_zero_label=False,
                beta=1):
    fscore_result = eval_metrics(results=results,
                                 gt_seg_maps=gt_seg_maps,
                                 num_classes=num_classes,
                                 ignore_index=ignore_index,
                                 metrics=['mFscore'],
                                 nan_to_num=nan_to_num,
                                 label_map=label_map,
                                 reduce_zero_label=reduce_zero_label,
                                 beta=beta)
    return fscore_result


def eval_metrics(results,
                 gt_seg_maps,
                 num_classes,
                 ignore_index,
                 metrics=['mIoU'],
                 nan_to_num=None,
                 label_map=dict(),
                 reduce_zero_label=False,
                 beta=1):
    total_area_intersect, total_area_union, total_area_pred_label, \
        total_area_label = total_intersect_and_union(
            results, gt_seg_maps, num_classes, ignore_index, label_map,
            reduce_zero_label)
    ret_metrics = total_area_to_metrics(total_area_intersect, total_area_union,
                                        total_area_pred_label,
                                        total_area_label, metrics, nan_to_num,
                                        beta)

    return ret_metrics


def pre_eval_to_metrics(pre_eval_results,
                        metrics=['mIoU'],
                        nan_to_num=None,
                        beta=1):
    # convert list of tuples to tuple of lists, e.g.
    # [(A_1, B_1, C_1, D_1), ...,  (A_n, B_n, C_n, D_n)] to
    # ([A_1, ..., A_n], ..., [D_1, ..., D_n])
    pre_eval_results = tuple(zip(*pre_eval_results))
    assert len(pre_eval_results) == 4

    total_area_intersect = sum(pre_eval_results[0])
    total_area_union = sum(pre_eval_results[1])
    total_area_pred_label = sum(pre_eval_results[2])
    total_area_label = sum(pre_eval_results[3])

    ret_metrics = total_area_to_metrics(total_area_intersect, total_area_union,
                                        total_area_pred_label,
                                        total_area_label, metrics, nan_to_num,
                                        beta)

    return ret_metrics


def total_area_to_metrics(total_area_intersect,
                          total_area_union,
                          total_area_pred_label,
                          total_area_label,
                          metrics=['mIoU'],
                          nan_to_num=None,
                          beta=1):
    if isinstance(metrics, str):
        metrics = [metrics]
    allowed_metrics = ['mIoU', 'mDice', 'mFscore']
    if not set(metrics).issubset(set(allowed_metrics)):
        raise KeyError('metrics {} is not supported'.format(metrics))

    all_acc = total_area_intersect.sum() / total_area_label.sum()
    ret_metrics = OrderedDict({'aAcc': all_acc})
    for metric in metrics:
        if metric == 'mIoU':
            iou = total_area_intersect / total_area_union
            acc = total_area_intersect / total_area_label
            ret_metrics['IoU'] = iou
            ret_metrics['Acc'] = acc
        elif metric == 'mDice':
            dice = 2 * total_area_intersect / (total_area_pred_label +
                                               total_area_label)
            acc = total_area_intersect / total_area_label
            ret_metrics['Dice'] = dice
            ret_metrics['Acc'] = acc
        elif metric == 'mFscore':
            precision = total_area_intersect / total_area_pred_label
            recall = total_area_intersect / total_area_label
            f_value = np.array(
                [f_score(x[0], x[1], beta) for x in zip(precision, recall)])
            ret_metrics['Fscore'] = f_value
            ret_metrics['Precision'] = precision
            ret_metrics['Recall'] = recall

    if nan_to_num is not None:
        ret_metrics = OrderedDict({
            metric: np.nan_to_num(metric_value, nan=nan_to_num)
            for metric, metric_value in ret_metrics.items()
        })
    return ret_metrics

python/jittorseg/utils/metrics.py
JittorCV/jittorseg
class Registry:
    def __init__(self):
        self._modules = {}

    def register_module(self, name=None, module=None):
        def _register_module(module):
            key = name
            if key is None:
                key = module.__name__
            assert key not in self._modules, f"{key} is already registered."
            self._modules[key] = module
            return module

        if module is not None:
            return _register_module(module)

        return _register_module

    def get(self, name):
        assert name in self._modules, f"{name} is not registered."
        return self._modules[name]


def build_from_cfg(cfg, registry, **kwargs):
    if isinstance(cfg, str):
        return registry.get(cfg)(**kwargs)
    elif isinstance(cfg, dict):
        args = cfg.copy()
        args.update(kwargs)
        obj_type = args.pop('type')
        obj_cls = registry.get(obj_type)
        try:
            module = obj_cls(**args)
        except TypeError as e:
            if "<class" not in str(e):
                e = f"{obj_cls}.{e}"
            raise TypeError(e)

        return module
    elif isinstance(cfg, list):
        from jittor import nn
        return nn.Sequential(
            [build_from_cfg(c, registry, **kwargs) for c in cfg])
    elif cfg is None:
        return None
    else:
        raise TypeError(f"type {type(cfg)} not support")


DATASETS = Registry()
TRANSFORMS = Registry()
MODELS = Registry()
BACKBONES = Registry()
HEADS = Registry()
LOSSES = Registry()
OPTIMS = Registry()
NECKS = Registry()
SCHEDULERS = Registry()
HOOKS = Registry()
PIXEL_SAMPLERS = Registry()
# basic_operators
CONV_LAYERS = Registry()
ACTIVATION_LAYERS = Registry()
NORM_LAYERS = Registry()
DROPOUT_LAYERS = Registry()
PADDING_LAYERS = Registry()

python/jittorseg/utils/registry.py
JittorCV/jittorseg
import cv2
import os
import numpy as np
import warnings


def visualize_result(seg, palette=None, save_dir=None, file_name=None):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)
    out_file = os.path.join(save_dir, file_name[file_name.rfind("/") + 1:])
    if palette is None:
        palette = np.random.randint(0, 255, size=(255, 3))
        warnings.warn('palette is not specified, random palette is used')

    palette = np.array(palette)
    assert palette.shape[1] == 3
    assert len(palette.shape) == 2
    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    for label, color in enumerate(palette):
        color_seg[seg == label, :] = color
    # convert to BGR
    color_seg = color_seg[..., ::-1]
    color_seg = color_seg.astype(np.uint8)
    cv2.imwrite(out_file, color_seg)

python/jittorseg/utils/visualize.py
JittorCV/jittorseg
import jittor as jt
from jittor import init, nn
import numpy as np
import math
from scipy import special
import warnings


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
            "The distribution of values may be incorrect.",
            stacklevel=2)

    with jt.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        init.uniform_(tensor, 2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor = jt.array(special.erfinv(tensor.numpy()))
        # Transform to proper mean, std
        tensor.multiply(std * math.sqrt(2.))
        tensor.add(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp(min_v=a, max_v=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


def normal_init(module, mean=0, std=1, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        init.gauss_(module.weight, mean, std)
    if hasattr(module, 'bias') and isinstance(
            module.bias, jt.Var) and module.bias is not None:
        init.constant_(module.bias, bias)


def constant_init(module, val, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        init.constant_(module.weight, val)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def xavier_init(module, gain=1, bias=0, distribution='normal'):
    assert distribution in ['uniform', 'normal']
    if hasattr(module, 'weight') and module.weight is not None:
        if distribution == 'uniform':
            init.xavier_uniform_(module.weight, gain=gain)
        else:
            init.xavier_normal_(module.weight, gain=gain)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def trunc_normal_init(module: nn.Module,
                      mean: float = 0,
                      std: float = 1,
                      a: float = -2,
                      b: float = 2,
                      bias: float = 0) -> None:
    if hasattr(module, 'weight') and module.weight is not None:
        trunc_normal_(module.weight, mean, std, a, b)  # type: ignore
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)  # type: ignore


def uniform_init(module, a=0, b=1, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        init.uniform_(module.weight, a, b)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def kaiming_init(module,
                 a=0,
                 mode='fan_out',
                 nonlinearity='relu',
                 bias=0,
                 distribution='normal'):
    assert distribution in ['uniform', 'normal']
    if hasattr(module, 'weight') and module.weight is not None:
        if distribution == 'uniform':
            init.kaiming_uniform_(module.weight,
                                  a=a,
                                  mode=mode,
                                  nonlinearity=nonlinearity)
        else:
            init.kaiming_normal_(module.weight,
                                 a=a,
                                 mode=mode,
                                 nonlinearity=nonlinearity)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def caffe2_xavier_init(module, bias=0):
    kaiming_init(module,
                 a=1,
                 mode='fan_in',
                 nonlinearity='leaky_relu',
                 bias=bias,
                 distribution='uniform')


def bias_init_with_prob(prior_prob):
    """initialize conv/fc bias value according to a given probability value."""
    bias_init = float(-np.log((1 - prior_prob) / prior_prob))
    return bias_init

python/jittorseg/utils/weight_init.py
JittorCV/jittorseg
from setuptools import find_packages, setup
import os

path = os.path.join(os.path.dirname(__file__), "python")
with open(os.path.join(path, "jittorseg/__init__.py"), "r", encoding='utf8') as fh:
    for line in fh:
        if line.startswith('__version__'):
            version = line.split("'")[1]
            break
    else:
        raise RuntimeError("Unable to find version string.")


def parse_requirements(fname='requirements.txt', with_version=True):
    """Parse the package dependencies listed in a requirements file but strips
    specific versioning information.

    Args:
        fname (str): path to requirements file
        with_version (bool, default=False): if True include version specs

    Returns:
        List[str]: list of requirements items

    CommandLine:
        python -c "import setup; print(setup.parse_requirements())"
    """
    import sys
    from os.path import exists
    import re
    require_fpath = fname

    def parse_line(line):
        """Parse information from a line in a requirements text file."""
        if line.startswith('-r '):
            # Allow specifying requirements in other files
            target = line.split(' ')[1]
            for info in parse_require_file(target):
                yield info
        else:
            info = {'line': line}
            if line.startswith('-e '):
                info['package'] = line.split('#egg=')[1]
            else:
                # Remove versioning from the package
                pat = '(' + '|'.join(['>=', '==', '>']) + ')'
                parts = re.split(pat, line, maxsplit=1)
                parts = [p.strip() for p in parts]

                info['package'] = parts[0]
                if len(parts) > 1:
                    op, rest = parts[1:]
                    if ';' in rest:
                        # Handle platform specific dependencies
                        # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies
                        version, platform_deps = map(str.strip,
                                                     rest.split(';'))
                        info['platform_deps'] = platform_deps
                    else:
                        version = rest  # NOQA
                    info['version'] = (op, version)
            yield info

    def parse_require_file(fpath):
        with open(fpath, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    for info in parse_line(line):
                        yield info

    def gen_packages_items():
        if exists(require_fpath):
            for info in parse_require_file(require_fpath):
                parts = [info['package']]
                if with_version and 'version' in info:
                    parts.extend(info['version'])
                if not sys.version.startswith('3.4'):
                    # apparently package_deps are broken in 3.4
                    platform_deps = info.get('platform_deps')
                    if platform_deps is not None:
                        parts.append(';' + platform_deps)
                item = ''.join(parts)
                yield item

    packages = list(gen_packages_items())
    return packages


if __name__ == '__main__':
    setup(name='jittorseg',
          version=version,
          description='Semantic Segmentation Toolbox By Jittor',
          author="Jittor Group",
          author_email="jittor@qq.com",
          keywords='computer vision, semantic segmentation',
          url="http://jittor.com",
          python_requires='>=3.7',
          packages=find_packages(path),
          package_dir={'': "python"},
          license='Apache License 2.0',
          install_requires=parse_requirements('requirements.txt'),
          zip_safe=False)

setup.py
JittorCV/jittorseg
from jittorseg.utils.inference import InferenceSegmentor


def main():
    config_file = 'project/fcn/fcn_r50-d8_512x1024_cityscapes_80k.py'
    ckp_file = 'work_dirs/fcn_r50-d8_512x1024_cityscapes_80k/checkpoints/ckpt_80000.pkl'
    save_dir = './'
    image = 'cityscapes/leftImg8bit/val/munster/munster_000069_000019_leftImg8bit.png'

    inference_segmentor = InferenceSegmentor(config_file, ckp_file, save_dir)
    inference_segmentor.infer(image)


if __name__ == "__main__":
    main()

tools/demo.py
JittorCV/jittorseg
from utils.config_process import parser
import jittor as jt
from jittorseg.runner import TrainRunner, EvalRunner, TestRunner
from jittorseg.config import init_cfg
from jittorseg.config.config import update_cfg

jt.cudnn.set_max_workspace_ratio(0.0)


def main():
    args = parser()
    if not args.no_cuda:
        jt.flags.use_cuda = 1

    assert args.task in [
        "train", "val", "test"
    ], f"{args.task} not support, please choose [train,val,test]"

    if args.config_file:
        init_cfg(args.config_file)

    if args.resume:
        update_cfg(resume_path=args.resume)
    if args.efficient_val:
        update_cfg(efficient_val=args.efficient_val)

    runner = None

    if args.task == "train":
        runner = TrainRunner()
    elif args.task == "val":
        runner = EvalRunner()
    elif args.task == "test":
        runner = TestRunner(save_dir=args.save_dir)

    runner.run()


if __name__ == "__main__":
    main()

tools/run_net.py
JittorCV/jittorseg
import argparse


def parser():
    parser = argparse.ArgumentParser(
        description="OpenJittorVision Semantic segmentation Training")
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "--task",
        default="train",
        help="train,val test",
        type=str,
    )
    parser.add_argument(
        "--resume",
        default=None,
        help="resume path",
        type=str,
    )
    parser.add_argument(
        "--save-dir",
        default="./results",
        type=str,
    )
    parser.add_argument("--no_cuda", action='store_true')
    parser.add_argument("--efficient_val", action='store_true')
    return parser.parse_args()

tools/utils/config_process.py
JittorCV/jittorseg