# dataset settings
dataset_type = 'Dataset3dsam'
data_root = '/data/sdb/user/Deeplesion/'
anno_root = '/data/sdb/user/mmdet/data/'
train_pipeline = [
    dict(type='LoadTioImage'),
    dict(type='RescaleIntensity'),
    dict(type='ComputeAugParam')
]
test_pipeline = [
    dict(type='LoadTioImage'),
    dict(type='RescaleIntensity')
]
data = dict(
    samples_per_gpu=2,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,
        data_dir=data_root + 'nii/',
        index_file=anno_root + 'deeplesion_filename.csv',
        pipeline=train_pipeline,
    ),
    val=dict(
        type=dataset_type,
        data_dir=data_root + 'Images_nifti/',
        index_file=anno_root + 'dlt.csv',
        pipeline=test_pipeline,
    ),
    test=dict(
        type=dataset_type,
        data_dir=data_root + 'Images_nifti/',
        index_file=anno_root + '/filename_test.csv',
        pipeline=test_pipeline,
    ),
)
evaluation = dict(interval=1)

configs/_base_/datasets/deeplesion_sam.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
_base_ = './sam_deeplesion.py'
view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="Crop"),
    # dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    # dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', 'meshgrid', 'valid'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info"
         ), )
]
view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="Crop"),
    # dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    # dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', 'meshgrid', 'valid'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info"

         ), )
]
train_pipeline = [
    dict(type='LoadTioImage'),
    # dict(type='RescaleIntensity'),
    # dict(type='Crop100'),
    dict(type='CropBackground'),
    dict(type='ComputeAugParam_sample'),
    dict(
        type="MultiBranch", view1=view1_pipline, view2=view2_pipline
    ),
]
test_pipeline = [
    dict(type='LoadTestTioImage', landmark=False),
    dict(type='Resample'),
    dict(type='RescaleIntensity'),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img'])
]
model = dict(
    train_cfg=dict(
        pre_select_pos_number=2000,
        after_select_pos_number=100,
        pre_select_neg_number=2000,
        after_select_neg_number=500,
        positive_distance=2.,
        ignore_distance=20.,
        coarse_positive_distance=25.,
        coarse_ignore_distance=5.,
        coarse_z_thres=6.,
        coarse_pre_select_neg_number=250,
        coarse_after_select_neg_number=200,
        coarse_global_select_number=1000,
        temperature=0.5
    ),
    test_cfg=dict(
        save_path='/data/sdd/user/results/result-dlt/',
        # save_path='/data/sdd/user/results/landmark_n/',
        output_embedding=True
    ))
lymphnode_data_root = '/data/sdd/user/processed_data/lymphnode/nii/'
anno_root = '/data/sdd/user/processed_data/ind_files/'
dlt_eval_data_root = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
landmark_eval_data_root = '/data/sdd/user/rawdata/chestCT/CT_nifti_ln_instance/'
landmark_ce_eval_data_root = '/data/sdd/user/rawdata/chestCT/RTCT2CE_nifti_nnunet/'
deeplesion_data_root = '/data/sdd/user/processed_data/Deeplesion/nii/'
word_data_root = '/data/sde/user/WORD-V0.1.0/imagesTr/'
data = dict(
    samples_per_gpu=6,
    workers_per_gpu=12,
    train=dict(
        data_dir=lymphnode_data_root,
        index_file=anno_root + 'lymphnode_filename.csv',
        pipeline=train_pipeline,
    ),
    val=dict(
        data_dir=landmark_eval_data_root,
        index_file=anno_root + 'landmark.csv',
        pipeline=test_pipeline,
    ),
    test=dict(
        # data_dir=word_data_root,
        # index_file=anno_root + 'word_filename.csv',
        data_dir=dlt_eval_data_root,
        # index_file=anno_root + 'deeplesion_dlt_jimmy.csv',
        index_file=anno_root + 'dlt.csv',
        pipeline=test_pipeline,
    ), )

fp16 = dict(loss_scale="dynamic")
# optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
# lr_config = dict(
#     policy='step',
#     warmup='linear',
#     warmup_iters=500,
#     warmup_ratio=0.001,
#     step=[18000,19000])

configs/sam/sam_NIHLN.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
_base_ = [
    '../_base_/datasets/deeplesion_sam.py',
    '../_base_/default_runtime.py'
]

model = dict(
    type='Sam',
    backbone=dict(
        type='ResNet3d',
        pretrained2d=True,
        pretrained='torchvision://resnet18',
        depth=18,
        in_channels=1,
        spatial_strides=(2, 2, 2, 2),
        temporal_strides=(1, 1, 1, 2),
        conv1_kernel=(3, 7, 7),
        conv1_stride_t=1,
        conv1_stride_s=1,
        pool1_stride_t=1,
        with_pool1=False,
        with_pool2=True,
        conv_cfg=dict(type='Conv3d'),
        inflate=((0, 0), (0, 0), (1, 1), (1, 1)),
        norm_eval=False,
        zero_init_residual=False),
    neck=dict(
        type='FPN3d',
        end_level=3,
        in_channels=[64, 128, 256],
        out_channels=128,
        num_outs=3,
        conv_cfg=dict(type='Conv3d')),
    read_out_head=dict(
        type='FPN3d',
        end_level=1,
        in_channels=[512],
        out_channels=128,
        num_outs=1,
        conv_cfg=dict(type='Conv3d')),
    # model training and testing settings
    train_cfg=dict(
        pre_select_pos_number=2000,
        after_select_pos_number=100,
        pre_select_neg_number=2000,
        after_select_neg_number=500,
        positive_distance=5.,
        ignore_distance=3.,
        coarse_positive_distance=25.,
        coarse_ignore_distance=5.,
        coarse_z_thres=6.,
        coarse_pre_select_neg_number=250,
        coarse_after_select_neg_number=200,
        coarse_global_select_number=300,
        temperature=0.5
    ),
    test_cfg=dict(
        save_path='/data/sdc/user/result-dlt/',
        output_embedding=False
    ))
view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="Crop"),
    # dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', 'meshgrid', 'valid'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info"
         ), )
]
view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="Crop"),
    # dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', 'meshgrid', 'valid'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info"

         ), )
]
train_pipeline = [
    dict(type='LoadTioImage'),
    # dict(type='RescaleIntensity'),
    dict(type='CropBackground'),
    dict(type='ComputeAugParam'),
    dict(
        type="MultiBranch", view1=view1_pipline, view2=view2_pipline
    ),
]
test_pipeline = [
    dict(type='LoadTestTioImage'),
    dict(type='Resample'),
    dict(type='RescaleIntensity'),
    dict(type="GenerateMetaInfo"),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img'])
]
data = dict(
    samples_per_gpu=3,
    workers_per_gpu=12,
    train=dict(
        pipeline=train_pipeline,
    ),
    val=dict(
        pipeline=test_pipeline,
    )
)
find_unused_parameters = False

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[])
runner = dict(type="IterBasedRunner", max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=20)

configs/sam/sam_deeplesion.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
_base_ = [
    '../_base_/datasets/deeplesion_sam.py',
    '../_base_/default_runtime.py'
]

model = dict(
    type='Sam_uniform_cross_volume_dual_head_mean_vector',
    backbone=dict(
        type='ResNet3d',
        pretrained2d=True,
        pretrained='torchvision://resnet18',
        depth=18,
        in_channels=1,
        spatial_strides=(2, 2, 2, 2),
        temporal_strides=(1, 1, 1, 2),
        conv1_kernel=(3, 7, 7),
        conv1_stride_t=1,
        conv1_stride_s=1,
        pool1_stride_t=1,
        with_pool1=False,
        with_pool2=True,
        conv_cfg=dict(type='Conv3d'),
        inflate=((0, 0), (0, 0), (1, 1), (1, 1)),
        # norm_cfg = dict(type='GN',num_groups=32, requires_grad=True),
        norm_eval=False,
        zero_init_residual=False),
    neck=dict(
        type='FPN3d',
        start_level=0,
        end_level=3,
        in_channels=[64, 128, 256],
        out_channels=128,
        num_outs=3,
        conv_cfg=dict(type='Conv3d')),
    # sem_neck=dict(
    #     type='FPN3d',
    #     start_level=0,
    #     end_level=4,
    #     in_channels=[64, 128, 256,512],
    #     out_channels=128,
    #     num_outs=4,
    #     conv_cfg=dict(type='Conv3d')),
    read_out_head=dict(
        type='FPN3d',
        end_level=1,
        in_channels=[512],
        out_channels=128,
        num_outs=1,
        conv_cfg=dict(type='Conv3d')),
    superloss=dict(type='SupConMeanLoss',
                   temperature=0.1,
                   base_temperature=1.0,
                   ),
    # model training and testing settings
    train_cfg=dict(
        intra_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            coarse_global_select_number=1000,
            temperature=0.5),
        inter_cfg=dict(
            sampling=True,
            ignore_by_distance=False,
            sampling_pos_number=2000,
            sampling_neg_number=3000,
            positive_distance=16.,
            ignore_distance=8,
            no_overlap_pos_number=100,
            no_overlap_neg_number=200,
            temperature=0.5)
    ),
    test_cfg=dict(
        # save_path='/data/sdd/user/results/result-word/',
        # save_path='/data/sdd/user/results/result-landmark_n/',
        save_path='/data/sdd/user/results/result-dlt/',
        output_embedding=True
    ))

intra_view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    # dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    # dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'
         ), )
]
intra_view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    # dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    # dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'

         ), )
]
intra_train_pipeline = [
    dict(type='LoadTioImage'),
    # dict(type='RescaleIntensity'),
    # dict(type='Crop100'),
    dict(type='CropBackground'),
    dict(type='ComputeAugParam_sample'),
    dict(
        type="MultiBranch", view1=intra_view1_pipline, view2=intra_view2_pipline
    ),
]

inter_view1_pipline = [
    dict(type="ExtraAttrs", tag="im1"),

    dict(type="SeperateInfoByImage"),
    dict(type='LoadTioImageWithMask', with_mesh=False),
    dict(type="ExtraAttrs", style="inter"),

    # dict(type='Resample', norm_spacing=(4., 4., 4.)),
    # dict(type="Crop", switch='fix'),
    dict(type="Resample", norm_spacing=(2., 2., 2.), intra_volume=False, crop_artifacts=False),
    # dict(type="RandomElasticDeformation"),
    dict(type='Cropz', x_slice=180, y_slice=180, z_slice=120),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    # dict(type='RandomCrop3d', thres=0.3),
    # dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo", with_mask=True),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "volume_size",
             "paired",
             "style",
             "mask"
         ), )
]
inter_view2_pipline = [
    dict(type="ExtraAttrs", tag="im2"),

    dict(type="SeperateInfoByImage"),
    dict(type='LoadTioImageWithMask', with_mesh=False),
    dict(type="ExtraAttrs", style="inter"),
    # dict(type="RandomAffine3d"),

    dict(type="Resample", norm_spacing=(2., 2., 2.), intra_volume=False, crop_artifacts=False),
    dict(type='Cropz', x_slice=180, y_slice=180, z_slice=120),
    # dict(type="Crop", switch='fix'),
    # dict(type='Resample', norm_spacing=(4., 4., 4.)),
    # dict(type="RandomElasticDeformation"),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    # dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo", with_mask=True),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img', 'mask'],
         meta_keys=(
             "filename",
             "tag",
             "volume_size",
             "paired",
             "style",
             "mask"
         ), )
]
inter_train_pipeline = [
    dict(type='GeneratePairedImagesMaskInfo'),
    dict(
        type="MultiBranch", view1=inter_view1_pipline, view2=inter_view2_pipline
    ),
]
test_pipeline = [
    dict(type='LoadTestTioImage', landmark=False),
    dict(type='Resample', norm_spacing=(2., 2., 2.)),
    dict(type='RescaleIntensity'),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img'])
]
chest_inter_data_root = '/data/sdd/user/processed_data/chestct/nii-ce-resample/'
chest_inter_mask_root = '/data/sdd/user/processed_data/chestct/nii-ce-resample-mask/'
lymphnode_intra_data_root = '/data/sdd/user/processed_data/lymphnode/nii/'

anno_root = '/data/sdd/user/processed_data/ind_files/'
dlt_eval_data_root = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
landmark_n_eval_data_root = '/data/sdd/user/rawdata/chestCT/CT_nifti_ln_instance/'
deeplesion_data_root = '/data/sdd/user/processed_data/Deeplesion/nii/'
landmark_ce_eval_data_root = '/data/sdd/user/rawdata/chestCT/RTCT2CE_nifti_nnunet/'

flare_inter_data_root = '/data/sdd/user/processed_data/flare/nii-resample'
flare_inter_mask_root = '/data/sdd/user/processed_data/flare/nii-resample-mask/'

totolseg_inter_data_root = '/data/sdd/user/processed_data/totolsegmentaion/nii-resample'
totolseg_inter_mask_root = '/data/sdd/user/processed_data/totolsegmentaion/nii-resample-mask'

word_data_root = '/data/sde/user/WORD-V0.1.0/imagesTr/'
lymphnode_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=500,
    data_dir=lymphnode_intra_data_root,
    index_file=anno_root + 'lymphnode_filename.csv',
    pipeline=intra_train_pipeline,
)
chest_inter_set = dict(
    type='Dataset3dsamCrossVolume',
    multisets=True,
    set_length=500,
    data_dir=chest_inter_data_root,
    mask_dir=chest_inter_mask_root,
    index_file=anno_root + 'chest_ce_filename.csv',
    pipeline=inter_train_pipeline,
)

flare_inter_set = dict(
    type='Dataset3dsamCrossVolume',
    multisets=True,
    set_length=2000,
    data_dir=flare_inter_data_root,
    mask_dir=flare_inter_mask_root,
    index_file=anno_root + 'flare_filename.csv',
    pipeline=inter_train_pipeline,
)

totalseg_inter_set = dict(
    type='Dataset3dsamCrossVolume',
    multisets=True,
    set_length=1200,
    data_dir=totolseg_inter_data_root,
    mask_dir=totolseg_inter_mask_root,
    index_file=anno_root + 'totolseg_filename.csv',
    pipeline=inter_train_pipeline,
)

data = dict(
    samples_per_gpu=5,
    workers_per_gpu=12,
    train=dict(
        type='ConcatDataset',
        datasets=[lymphnode_intra_set, totalseg_inter_set]
    ),
    val=dict(
        data_dir=landmark_ce_eval_data_root,
        index_file=anno_root + 'landmark_ce.csv',
        pipeline=test_pipeline,
    ),

    test=dict(
        # data_dir=word_data_root,
        # index_file=anno_root + 'word_filename.csv',
        data_dir=dlt_eval_data_root,
        index_file=anno_root + 'dlt.csv',
        pipeline=test_pipeline,
    ), )
#     data_dir=landmark_n_eval_data_root,
#     index_file=anno_root + 'landmark.csv',
#     pipeline=test_pipeline,
# ), )
find_unused_parameters = True

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[])
runner = dict(type="IterBasedRunner", max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=20)
fp16 = dict(loss_scale="dynamic")

configs/samv2/samv2_NIHLN.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
_base_ = [
    '../_base_/datasets/deeplesion_sam.py',
    '../_base_/default_runtime.py'
]

model = dict(
    type='Sam_uniform_cross_volume_dual_head_mri_ct',
    backbone=dict(
        type='ResNet3d',
        pretrained2d=True,
        pretrained='torchvision://resnet18',
        depth=18,
        in_channels=1,
        spatial_strides=(2, 2, 2, 2),
        temporal_strides=(1, 1, 1, 2),
        conv1_kernel=(3, 7, 7),
        conv1_stride_t=1,
        conv1_stride_s=1,
        pool1_stride_t=1,
        with_pool1=False,
        with_pool2=True,
        conv_cfg=dict(type='Conv3d'),
        inflate=((0, 0), (0, 0), (1, 1), (1, 1)),
        # norm_cfg = dict(type='GN',num_groups=32, requires_grad=True),
        norm_eval=False,
        zero_init_residual=False),
    neck=dict(
        type='FPN3d',
        start_level=0,
        end_level=3,
        in_channels=[64, 128, 256],
        out_channels=128,
        num_outs=3,
        conv_cfg=dict(type='Conv3d')),
    read_out_head=dict(
        type='FPN3d',
        end_level=1,
        in_channels=[512],
        out_channels=128,
        num_outs=1,
        conv_cfg=dict(type='Conv3d')),
    # model training and testing settings
    train_cfg=dict(
        meta_cfg=dict(
            prob=0.5,
        ),
        intra_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            coarse_global_select_number=1000,
            fine_temperature=0.5,
            coarse_temperature=0.5),
        inter_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            fine_temperature=0.5,
            coarse_temperature=0.5)
    ),
    test_cfg=dict(
        save_path='/data/sdd/user/results/result-ctmri-landmark/',
        # save_path='/data/sdd/user/results/result-landmark_n/',
        # save_path='/data/sdd/user/results/result-dlt/',
        output_embedding=True
    ))

intra_view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    # dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'
         ), )
]
intra_view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    # dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'

         ), )
]
intra_train_pipeline = [
    dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type='ComputeAugParam_masksample', patch_size=(96, 96, 48)),
    dict(
        type="MultiBranch", view1=intra_view1_pipline, view2=intra_view2_pipline
    ),
]

inter_view1_pipline = [
    dict(type="ExtraAttrs", tag="CT"),
    dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type="ExtraAttrs", style="inter"),
    # dict(type='Resample', norm_spacing=(4., 4., 4.)),
    # dict(type="Crop", switch='fix'),
    dict(type="DynamicSpacing"),
    dict(type="Resample", norm_spacing=(2., 2., 2.), intra_volume=False, crop_artifacts=False, dynamic_spacing=True),
    dict(type="RandomAffine3d"),
    # dict(type="RandomElasticDeformation"),
    dict(type='RescaleIntensity'),
    dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    # dict(type='RandomCrop3d', thres=0.3),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo", with_mask=True),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "volume_size",
             "style",
             "mask",
             "meshgrid",
         ), )
]
inter_view2_pipline = [
    dict(type="ExtraAttrs", tag="MRI"),
    dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type="ExtraAttrs", style="inter"),
    # dict(type="RandomAffine3d"),
    dict(type="DynamicSpacing"),
    dict(type="Resample", norm_spacing=(2., 2., 2.), intra_volume=False, crop_artifacts=False, dynamic_spacing=True),
    dict(type="RandomAffine3d"),
    # dict(type="Crop", switch='fix'),
    # dict(type='Resample', norm_spacing=(4., 4., 4.)),
    # dict(type="RandomElasticDeformation"),
    dict(type='DynamicRescaleIntensity'),
    dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo", with_mask=True),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img', 'mask'],
         meta_keys=(
             "filename",
             "tag",
             "volume_size",
             "style",
             "mask",
             "meshgrid",
         ), )
]
inter_train_pipeline = [
    # dict(type="stoppoint"),
    dict(
        type="MultiBranch", view1=inter_view1_pipline, view2=inter_view2_pipline
    ),
]

test_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_ct.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    # dict(type='DynamicRescaleIntensity'),
    dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=False),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
test_ct_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_ct.nii.gz'),
    dict(type='Resample', norm_spacing=(2, 2, 2), crop_artifacts=False),
    # dict(type='DynamicRescaleIntensity'),
    dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=False),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
test_mri_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_t1.nii.gz'),
    dict(type='Resample', norm_spacing=(2, 2, 2), crop_artifacts=False),
    dict(type='DynamicRescaleIntensity'),
    # dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=True),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
chest_inter_data_root = '/data/sdd/user/processed_data/chestct/nii-ce-resample/'
chest_inter_mask_root = '/data/sdd/user/processed_data/chestct/nii-ce-resample-mask/'
lymphnode_intra_data_root = '/data/sdd/user/processed_data/lymphnode/nii/'

anno_root = '/data/sdd/user/processed_data/ind_files/'
dlt_eval_data_root = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
landmark_n_eval_data_root = '/data/sdd/user/rawdata/chestCT/CT_nifti_ln_instance/'
deeplesion_data_root = '/data/sdd/user/processed_data/Deeplesion/nii/'
landmark_ce_eval_data_root = '/data/sdd/user/rawdata/chestCT/RTCT2CE_nifti_nnunet/'

flare_inter_data_root = '/data/sdd/user/processed_data/flare/nii-resample'
flare_inter_mask_root = '/data/sdd/user/processed_data/flare/nii-resample-mask/'

totolseg_inter_data_root = '/data/sdd/user/processed_data/totolsegmentaion/nii-resample'
totolseg_inter_mask_root = '/data/sdd/user/processed_data/totolsegmentaion/nii-resample-mask'

headneck_data_root = '/data/sdd/user/processed_data/head_neck_CT_MRI/image/'
headneck_mask_root = '/data/sdd/user/processed_data/head_neck_CT_MRI/mask/'

headneck_landmrk_root = '/data/sdd/user/rawdata/CT-MRI-landmark/'

word_data_root = '/data/sde/user/WORD-V0.1.0/imagesTr/'

ct_mri_landmark_ct_set = dict(
    type='Dataset3dsam',
    data_dir=headneck_landmrk_root,
    index_file=anno_root + 'ct_mri_landmark_test.csv',
    pipeline=test_ct_pipeline,
)
ct_mri_landmark_mri_set = dict(
    type='Dataset3dsam',
    data_dir=headneck_landmrk_root,
    index_file=anno_root + 'ct_mri_landmark_test.csv',
    pipeline=test_mri_pipeline,
)

lymphnode_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=500,
    data_dir=lymphnode_intra_data_root,
    index_file=anno_root + 'lymphnode_filename.csv',
    pipeline=intra_train_pipeline,
)

headneck_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=headneck_data_root,
    index_file=anno_root + 'CT_MRI_filename.csv',
    pipeline=intra_train_pipeline,
    mask_dir=headneck_mask_root,
)

headneck_inter_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=headneck_data_root,
    index_file=anno_root + 'CT_MRI_filename.csv',
    pipeline=inter_train_pipeline,
    mask_dir=headneck_mask_root,
)

data = dict(
    samples_per_gpu=4,
    workers_per_gpu=12,
    train=dict(
        type='ConcatDataset',
        datasets=[headneck_intra_set, headneck_inter_set]
    ),
    val=dict(
        data_dir=landmark_ce_eval_data_root,
        index_file=anno_root + 'landmark_ce.csv',
        pipeline=test_pipeline,
    ),

    test=dict(
        type='ConcatDataset',
        datasets=[ct_mri_landmark_ct_set, ct_mri_landmark_mri_set]
    ), )
#     data_dir=landmark_n_eval_data_root,
#     index_file=anno_root + 'landmark.csv',
#     pipeline=test_pipeline,
# ), )
find_unused_parameters = True

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[])
runner = dict(type="IterBasedRunner", max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=20)
fp16 = dict(loss_scale="dynamic")

configs/samv2/samv2_ct_mri_infer.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
_base_ = [
    '../_base_/datasets/deeplesion_sam.py',
    '../_base_/default_runtime.py'
]

model = dict(
    type='Sam_uniform_cross_volume_dual_head_mri_ct',
    backbone=dict(
        type='ResNet3d',
        pretrained2d=True,
        pretrained='torchvision://resnet18',
        depth=18,
        in_channels=1,
        spatial_strides=(2, 2, 2, 2),
        temporal_strides=(1, 1, 1, 2),
        conv1_kernel=(3, 7, 7),
        conv1_stride_t=1,
        conv1_stride_s=1,
        pool1_stride_t=1,
        with_pool1=False,
        with_pool2=True,
        conv_cfg=dict(type='Conv3d'),
        inflate=((0, 0), (0, 0), (1, 1), (1, 1)),
        # norm_cfg = dict(type='GN',num_groups=32, requires_grad=True),
        norm_eval=False,
        zero_init_residual=False),
    neck=dict(
        type='FPN3d',
        start_level=0,
        end_level=3,
        in_channels=[64, 128, 256],
        out_channels=128,
        num_outs=3,
        conv_cfg=dict(type='Conv3d')),
    read_out_head=dict(
        type='FPN3d',
        end_level=1,
        in_channels=[512],
        out_channels=128,
        num_outs=1,
        conv_cfg=dict(type='Conv3d')),
    # model training and testing settings
    train_cfg=dict(
        meta_cfg=dict(
            prob=1.1,
        ),
        intra_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            coarse_global_select_number=1000,
            fine_temperature=0.5,
            coarse_temperature=0.5),
        inter_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            fine_temperature=0.5,
            coarse_temperature=0.5)
    ),
    test_cfg=dict(
        save_path='/data/sdd/user/results/result-ctmri-landmark/',
        output_embedding=True
    ))

intra_lymphnode_view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'
         ), )
]
intra_lymphnode_view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'

         ), )
]
intra_abd_view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity_modality', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'
         ), )
]
intra_abd_view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity_modality', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'

         ), )
]

intra_train_lymphnode_pipeline = [
    # dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type='ComputeAugParam_sample', patch_size=(96, 96, 32)),
    dict(
        type="MultiBranch", view1=intra_lymphnode_view1_pipline, view2=intra_lymphnode_view2_pipline
    ),
]

intra_train_abd_pipeline = [
    # dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='AddSuffixByProb', suffix=['_0000.nii.gz', '_0001.nii.gz'], p=[0.6, 0.4]),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type='ComputeAugParam_masksample', patch_size=(96, 96, 32)),
    dict(
        type="MultiBranch", view1=intra_abd_view1_pipline, view2=intra_abd_view2_pipline
    ),
]


inter_view1_pipline = [
    # dict(type="stoppoint"),
    dict(type="ExtraAttrs", tag="CT"),
    dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type="ExtraAttrs", style="inter"),
    dict(type="DynamicSpacing"),
    dict(type="Resample", norm_spacing=(2., 2., 2.), intra_volume=False, crop_artifacts=False, dynamic_spacing=True),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    dict(type='RescaleIntensity', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo", with_mask=True),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "volume_size",
             "style",
             "mask",
             "meshgrid",
             "resample_spacing"
         ), )
]
inter_view2_pipline = [
    dict(type="ExtraAttrs", tag="MRI"),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='AddSuffix', suffix='_0001.nii.gz'),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type="ExtraAttrs", style="inter"),
    dict(type="DynamicSpacing"),
    dict(type="Resample", norm_spacing=(2., 2., 2.), intra_volume=False, crop_artifacts=False, dynamic_spacing=True),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    dict(type='DynamicRescaleIntensity', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo", with_mask=True),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img', 'mask'],
         meta_keys=(
             "filename",
             "tag",
             "volume_size",
             "style",
             "mask",
             "meshgrid",
             "resample_spacing"
         ), )
]
inter_train_pipeline = [
    dict(
        type="MultiBranch", view1=inter_view1_pipline, view2=inter_view2_pipline
    ),
]


test_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_ct.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    # dict(type='DynamicRescaleIntensity'),
    dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=False),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
test_ct_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_ct.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    # dict(type='DynamicRescaleIntensity'),
    dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=False),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
test_mri_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_t1.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    dict(type='DynamicRescaleIntensity'),
    # dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=True),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]

lymphnode_intra_data_root = '/data/sdd/user/processed_data/lymphnode/nii-crop/'
lymphnode_intra_mask_root = '/data/sdd/user/processed_data/lymphnode/nii-crop-mask/'

anno_root = '/data/sdd/user/processed_data/ind_files/'
dlt_eval_data_root = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
landmark_n_eval_data_root = '/data/sdd/user/rawdata/chestCT/CT_nifti_ln_instance/'

abd_data_root = '/data/sdd/user/processed_data/abd_CT_MRI/image/'
abd_mask_root = '/data/sdd/user/processed_data/abd_CT_MRI/mask-slic/'

abd_inter_data_root = '/data/sde/user/abd_registered/processed_0/image-regis_0/'
abd_inter_mask_root = '/data/sde/user/abd_registered/processed_0/mask-regis_0/'

headneck_landmrk_root = '/data/sdd/user/rawdata/CT-MRI-landmark/'

lymphnode_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=lymphnode_intra_data_root,
    index_file=anno_root + 'lymphnode_filename.csv',
    pipeline=intra_train_lymphnode_pipeline,
    mask_dir=lymphnode_intra_mask_root,
)
abd_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=abd_data_root,
    index_file=anno_root + 'CT_MRI_abd_filename.csv',
    pipeline=intra_train_abd_pipeline,
    mask_dir=abd_mask_root,
)
abd_inter_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=abd_inter_data_root,
    index_file=anno_root + 'CT_MRI_abd_inter_0_filename.csv',
    pipeline=inter_train_pipeline,
    mask_dir=abd_inter_mask_root,
)
data = dict(
    samples_per_gpu=4,
    workers_per_gpu=12,
    train=dict(
        type='ConcatDataset',
        # datasets=[abd_inter_set]
        datasets=[abd_intra_set, lymphnode_intra_set,abd_inter_set]
    ),
    val=dict(
        data_dir=intra_train_abd_pipeline,
        index_file=anno_root + 'landmark_ce.csv',
        pipeline=test_pipeline,
    ),

    test=dict(
        type='ConcatDataset',
        datasets=[intra_train_abd_pipeline, intra_train_abd_pipeline]
    ), )
find_unused_parameters = True

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[])
runner = dict(type="IterBasedRunner", max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=20)
fp16 = dict(loss_scale="dynamic")

configs/samv2/samv2_ct_mri_inter.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
_base_ = [
    '../_base_/datasets/deeplesion_sam.py',
    '../_base_/default_runtime.py'
]

model = dict(
    type='Sam_uniform_cross_volume_dual_head_mri_ct',
    backbone=dict(
        type='ResNet3d',
        pretrained2d=True,
        pretrained='torchvision://resnet18',
        depth=18,
        in_channels=1,
        spatial_strides=(2, 2, 2, 2),
        temporal_strides=(1, 1, 1, 2),
        conv1_kernel=(3, 7, 7),
        conv1_stride_t=1,
        conv1_stride_s=1,
        pool1_stride_t=1,
        with_pool1=False,
        with_pool2=True,
        conv_cfg=dict(type='Conv3d'),
        inflate=((0, 0), (0, 0), (1, 1), (1, 1)),
        # norm_cfg = dict(type='GN',num_groups=32, requires_grad=True),
        norm_eval=False,
        zero_init_residual=False),
    neck=dict(
        type='FPN3d',
        start_level=0,
        end_level=3,
        in_channels=[64, 128, 256],
        out_channels=128,
        num_outs=3,
        conv_cfg=dict(type='Conv3d')),
    read_out_head=dict(
        type='FPN3d',
        end_level=1,
        in_channels=[512],
        out_channels=128,
        num_outs=1,
        conv_cfg=dict(type='Conv3d')),
    # model training and testing settings
    train_cfg=dict(
        meta_cfg=dict(
            prob=-0.1,
        ),
        intra_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            coarse_global_select_number=1000,
            fine_temperature=0.5,
            coarse_temperature=0.5),
        inter_cfg=dict(
            pre_select_pos_number=2000,
            after_select_pos_number=100,
            pre_select_neg_number=2000,
            after_select_neg_number=500,
            positive_distance=2.,
            ignore_distance=20.,
            coarse_positive_distance=25.,
            coarse_ignore_distance=5.,
            coarse_z_thres=6.,
            coarse_pre_select_neg_number=250,
            coarse_after_select_neg_number=200,
            fine_temperature=0.5,
            coarse_temperature=0.5)
    ),
    test_cfg=dict(
        save_path='/data/sdd/user/results/result-ctmri-landmark/',
        output_embedding=True
    ))

intra_lymphnode_view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'
         ), )
]
intra_lymphnode_view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'

         ), )
]
intra_abd_view1_pipline = [
    dict(type="ExtraAttrs", tag="view1"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    # dict(type="RandomElasticDeformation"),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity_modality', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'
         ), )
]
intra_abd_view2_pipline = [
    dict(type="ExtraAttrs", tag="view2"),
    dict(type="ExtraAttrs", style="intra"),
    dict(type="Crop"),
    dict(type="RandomAffine3d", degrees=(-5, 5, -5, 5, -5, 5)),
    dict(type="Resample"),
    dict(type="Crop", switch='fix'),
    dict(type='RescaleIntensity_modality', out_min_max=(0, 1), bias=0),
    dict(type='RandHistogramShift'),
    # dict(type="RandomNoise3d"),
    dict(type="RandomBlur3d"),
    dict(type='RescaleIntensity', out_min_max=(0, 255), in_min_max=(0, 1), bias=50),
    dict(type="GenerateMeshGrid"),
    dict(type="GenerateMetaInfo"),
    dict(type='DefaultFormatBundle3d', switch='v2'),
    dict(type='Collect3d',
         keys=['img'],
         meta_keys=(
             "filename",
             "tag",
             "crop_info",
             "style",
             'meshgrid',
             'valid'

         ), )
]

intra_train_lymphnode_pipeline = [
    # dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type='ComputeAugParam_masksample', patch_size=(96, 96, 32)),
    dict(
        type="MultiBranch", view1=intra_lymphnode_view1_pipline, view2=intra_lymphnode_view2_pipline
    ),
]

intra_train_abd_pipeline = [
    # dict(type='AddSuffix', suffix='_0000.nii.gz'),
    dict(type='AddSuffixByProb', suffix=['_0000.nii.gz', '_0001.nii.gz'], p=[0.6, 0.4]),
    dict(type='LoadTioImageWithMask', with_mesh=True),
    dict(type='ComputeAugParam_masksample', patch_size=(96, 96, 32)),
    dict(
        type="MultiBranch", view1=intra_abd_view1_pipline, view2=intra_abd_view2_pipline
    ),
]

test_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_ct.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    # dict(type='DynamicRescaleIntensity'),
    dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=False),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
test_ct_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_ct.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    # dict(type='DynamicRescaleIntensity'),
    dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=False),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]
test_mri_pipeline = [
    # dict(type='AddSuffix', suffix='_t1.nii.gz'),
    # dict(type='AddSuffixByProb', suffix=['_0001.nii.gz', '_0002.nii.gz'], p=[0.5, 0.5]),
    dict(type='LoadTestTioImage', landmark=True, suffix='_t1.nii.gz'),
    dict(type='Resample', norm_spacing=(2., 2., 2.), crop_artifacts=False),
    dict(type='DynamicRescaleIntensity'),
    # dict(type='RescaleIntensity'),
    dict(type="IsMRI", is_mri=True),
    dict(type="GenerateMetaInfo", is_train=False),
    # dict(type='DefaultFormatBundle3d'),
    dict(type='Collect3d',
         keys=['img', ],
         meta_keys=(
             "filename",
             "is_mri",
         ), )
]

lymphnode_intra_data_root = '/data/sdd/user/processed_data/lymphnode/nii-crop/'
lymphnode_intra_mask_root = '/data/sdd/user/processed_data/lymphnode/nii-crop-mask/'

anno_root = '/data/sdd/user/processed_data/ind_files/'
dlt_eval_data_root = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
landmark_n_eval_data_root = '/data/sdd/user/rawdata/chestCT/CT_nifti_ln_instance/'

abd_data_root = '/data/sdd/user/processed_data/abd_CT_MRI/image/'
abd_mask_root = '/data/sdd/user/processed_data/abd_CT_MRI/mask-slic/'

abd_inter_data_root = '/data/sdd/user/processed_data/abd_CT_MRI/image-regis_0/'
abd_inter_mask_root = '/data/sdd/user/processed_data/abd_CT_MRI/mask-regis_0/'

headneck_landmrk_root = '/data/sdd/user/rawdata/CT-MRI-landmark/'

lymphnode_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=lymphnode_intra_data_root,
    index_file=anno_root + 'lymphnode_filename.csv',
    pipeline=intra_train_lymphnode_pipeline,
    mask_dir=lymphnode_intra_mask_root,
)
abd_intra_set = dict(
    type='Dataset3dsam',
    multisets=True,
    set_length=100,
    data_dir=abd_data_root,
    index_file=anno_root + 'CT_MRI_abd_filename.csv',
    pipeline=intra_train_abd_pipeline,
    mask_dir=abd_mask_root,
)
data = dict(
    samples_per_gpu=4,
    workers_per_gpu=12,
    train=dict(
        type='ConcatDataset',
        # datasets=[abd_inter_set]
        datasets=[abd_intra_set, lymphnode_intra_set]
    ),
    val=dict(
        data_dir=intra_train_abd_pipeline,
        index_file=anno_root + 'landmark_ce.csv',
        pipeline=test_pipeline,
    ),

    test=dict(
        type='ConcatDataset',
        datasets=[intra_train_abd_pipeline, intra_train_abd_pipeline]
    ), )
find_unused_parameters = True

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[])
runner = dict(type="IterBasedRunner", max_iters=20000)
checkpoint_config = dict(by_epoch=False, interval=1000, max_keep_ckpts=20)
fp16 = dict(loss_scale="dynamic")

configs/samv2/samv2_ct_mri_intra.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .deeds_preprocess_no_mask_origin_image import padding_img
misc/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import nibabel as nib
import os
import torch
import numpy as np

target_image = '/data/sdd/user/testt'

source_image = '/data/sdd/user/tests'

save_path = '/data/sdd/user/DEEDs/deedsBCV'
s_name = 'ct.nii.gz'
s_mask_name = 'ct_mask.nii.gz'
t_name = 'mri.nii.gz'
t_mask_name = 'mri_mask.nii.gz'


def padding_img(source_image=source_image, target_image=target_image, save_path=save_path, s_name=s_name,
                t_name=t_name, source_mask=None, target_mask=None):
    s_img = tio.ScalarImage(source_image)
    t_img = tio.ScalarImage(target_image)
    if not source_mask is None:
        s_mask = tio.LabelMap(source_mask)
        t_mask = tio.LabelMap(target_mask)

    # s_mask_data = s_mask.data
    # s_img.data[s_mask_data>4] = -1024
    # s_img.data[s_mask_data<4] = -1024
    #
    # t_mask_data = t_mask.data
    # t_img.data[t_mask_data>4] = -1024
    # t_img.data[t_mask_data<4] = -1024

    ToCanonical = tio.ToCanonical()  # this will cause bias shift, please check it
    s_img = ToCanonical(s_img)
    s_img.data = torch.flip(s_img.data, (1, 2))
    s_img.affine = np.array([-s_img.affine[0, :], -s_img.affine[1, :], s_img.affine[2, :], s_img.affine[3, :]])
    t_img = ToCanonical(t_img)
    t_img.data = torch.flip(t_img.data, (1, 2))
    t_img.affine = np.array([-t_img.affine[0, :], -t_img.affine[1, :], t_img.affine[2, :], t_img.affine[3, :]])
    if not source_mask is None:
        s_mask = ToCanonical(s_mask)
        s_mask.data = torch.flip(s_mask.data, (1, 2))
        s_mask.affine = np.array([-s_mask.affine[0, :], -s_mask.affine[1, :], s_mask.affine[2, :], s_mask.affine[3, :]])
        t_mask = ToCanonical(t_mask)
        t_mask.data = torch.flip(t_mask.data, (1, 2))
        t_mask.affine = np.array([-t_mask.affine[0, :], -t_mask.affine[1, :], t_mask.affine[2, :], t_mask.affine[3, :]])

    resampletrans = tio.Resample(s_img.spacing)
    # resampletrans = tio.Resample((2.,2.,4.))
    # s_img = resampletrans(s_img)
    # s_img_resample.data[s_img_resample.data > 3000] = 3000
    t_img_resample = resampletrans(t_img)
    if not source_mask is None:
        t_mask_resample = resampletrans(t_mask)

    # organ_ind = torch.where(s_img_resample.data>0)
    # organ_ind = torch.stack(organ_ind[1:])

    # organ_ind.min(dim=1)

    s_img_shape = np.array(s_img.shape[1:])
    t_img_shape = np.array(t_img_resample.shape[1:])

    all_shape = np.stack([s_img_shape, t_img_shape])
    max_shape = all_shape.max(axis=0)

    max_shape_normed = np.ceil(np.array(s_img.spacing) / 2. * max_shape)
    max_shape_normed = max_shape_normed + (3 - max_shape_normed % 3)
    max_shape_padding = np.round(max_shape_normed / np.array(s_img.spacing) * 2.)
    max_shape = max_shape_padding.astype(np.int)

    s_img_paddings = max_shape - s_img_shape
    s_img_paddings_half = np.floor(s_img_paddings / 2).astype('int')

    t_img_paddings = max_shape - t_img_shape
    t_img_paddings_half = np.floor(t_img_paddings / 2).astype('int')

    s_padded_img = np.zeros(max_shape, dtype=s_img.data[0, :, :, :].numpy().dtype)
    t_padded_img = np.zeros(max_shape, dtype=t_img_resample.data[0, :, :, :].numpy().dtype)

    s_padded_img[:] = -1024
    t_padded_img[:] = t_img_resample.data.numpy().min()

    if not source_mask is None:
        s_padding_mask = np.zeros(max_shape, dtype=s_mask.data[0, :, :, :].numpy().dtype)
        t_padding_mask = np.zeros(max_shape, dtype=t_mask_resample.data[0, :, :, :].numpy().dtype)

    # s_padded_img[s_img_paddings_half[0]:s_img_paddings_half[0] + s_img_shape[0],
    # s_img_paddings_half[1]:s_img_paddings_half[1] + s_img_shape[1],
    # s_img_paddings_half[2]:s_img_paddings_half[2] + s_img_shape[2]] = s_img_resample.data[0, :, :, :].numpy()
    # s_padded_mask[s_mask_paddings_half[0]:s_mask_paddings_half[0] + s_mask_shape[0],
    # s_mask_paddings_half[1]:s_mask_paddings_half[1] + s_mask_shape[1],
    # s_mask_paddings_half[2]:s_mask_paddings_half[2] + s_mask_shape[2]] = s_mask_resample.data[0, :, :, :].numpy()
    #
    # t_padded_img[t_img_paddings_half[0]:t_img_paddings_half[0] + t_img_shape[0],
    # t_img_paddings_half[1]:t_img_paddings_half[1] + t_img_shape[1],
    # t_img_paddings_half[2]:t_img_paddings_half[2] + t_img_shape[2]] = t_img_resample.data[0, :, :, :].numpy()
    # t_padded_mask[t_mask_paddings_half[0]:t_mask_paddings_half[0] + t_mask_shape[0],
    # t_mask_paddings_half[1]:t_mask_paddings_half[1] + t_mask_shape[1],
    # t_mask_paddings_half[2]:t_mask_paddings_half[2] + t_mask_shape[2]] = t_mask_resample.data[0, :, :, :].numpy()

    s_padded_img[:s_img_shape[0],
    : s_img_shape[1],
    : s_img_shape[2]] = s_img.data[0, :, :, :].numpy()

    t_padded_img[: t_img_shape[0],
    : t_img_shape[1],
    :t_img_shape[2]] = t_img_resample.data[0, :, :, :].numpy()

    tmp1 = s_img.affine[0, 3].copy()
    tmp2 = s_img.affine[1, 3].copy()
    tmp3 = s_img.affine[2, 3].copy()
    s_img.affine[0, 3] = tmp2
    s_img.affine[1, 3] = tmp1
    nii_img = nib.Nifti1Image(s_padded_img, affine=s_img.affine)
    nib.save(nii_img, os.path.join(save_path, s_name))

    tmp1 = t_img_resample.affine[0, 3].copy()
    tmp2 = t_img_resample.affine[1, 3].copy()
    t_img_resample.affine[0, 3] = tmp2
    t_img_resample.affine[1, 3] = tmp1
    nii_img = nib.Nifti1Image(t_padded_img, affine=t_img_resample.affine)
    nib.save(nii_img, os.path.join(save_path, t_name))

    if not source_mask is None:
        s_padding_mask[:s_img_shape[0],
        : s_img_shape[1],
        : s_img_shape[2]] = s_mask.data[0, :, :, :].numpy()

        t_padding_mask[: t_img_shape[0],
        : t_img_shape[1],
        :t_img_shape[2]] = t_mask_resample.data[0, :, :, :].numpy()

        tmp1 = s_mask.affine[0, 3].copy()
        tmp2 = s_mask.affine[1, 3].copy()
        tmp3 = s_img.affine[2, 3].copy()
        s_mask.affine[0, 3] = tmp2
        s_mask.affine[1, 3] = tmp1
        nii_img = nib.Nifti1Image(s_padding_mask, affine=s_mask.affine)
        nib.save(nii_img, os.path.join(save_path, s_mask_name))

        tmp1 = t_mask_resample.affine[0, 3].copy()
        tmp2 = t_mask_resample.affine[1, 3].copy()
        t_mask_resample.affine[0, 3] = tmp2
        t_mask_resample.affine[1, 3] = tmp1
        nii_img = nib.Nifti1Image(t_padding_mask, affine=t_mask_resample.affine)
        nib.save(nii_img, os.path.join(save_path, t_mask_name))


if __name__ == '__main__':
    padding_img(source_image=source_image, target_image=target_image, save_path=save_path, s_name=s_name,
                t_name=t_name)

misc/deeds_preprocess_no_mask_origin_image.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import csv
import os
import nibabel as nib

data_root = '/data/sdd/user/rawdata/lymphnode/Mediastinum/images_nii/'

save_path = '/data/sdd/user/processed_data/lymphnode/nii/'
if not os.path.exists(save_path):
    os.mkdir(save_path)

all_files = os.listdir(data_root)
print('test')

with open("/data/sdd/user/processed_data/lymphnode_m_filename.csv", "w") as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["name"])
    i = 0
    for name in all_files:
        filepath = os.path.join(data_root, name)
        img = tio.ScalarImage(filepath)
        resampletrans = tio.Resample((1.6, 1.6, 1.6))
        img_resample = resampletrans(img)

        nii_img = nib.Nifti1Image(img_resample.data[0, :, :, :].numpy(), affine=img_resample.affine)
        nib.save(nii_img, os.path.join(save_path, name))
        writer.writerow([name])
        print(i)
        i = i + 1

misc/lymphnode_preprocess.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import csv
import os
import nibabel as nib
from sam.datasets.piplines.process3d import RemoveCTBed,CropBackground,Resample

data_root = '/data/sdd/user/rawdata/lymphnode/Abdomen/images_nii/'

save_path = '/data/sdd/user/processed_data/lymphnode/nii-crop/'
mask_save_path = '/data/sdd/user/processed_data/lymphnode/nii-crop-mask/'
if not os.path.exists(save_path):
    os.mkdir(save_path)

if not os.path.exists(mask_save_path):
    os.mkdir(mask_save_path)

all_files = os.listdir(data_root)
print('test')


for name in all_files:
    print(name)
    filepath = os.path.join(data_root, name)
    img = tio.ScalarImage(filepath)
    subject = tio.Subject(image = img)
    data = {}
    data['subject'] = subject
    removebedtrans = RemoveCTBed()
    data = removebedtrans(data)
    resampletrans = Resample(norm_spacing=(1.6, 1.6, 1.6))
    data = resampletrans(data)
    croptrans = CropBackground(by_mask=True)
    data = croptrans(data)

    nii_img = nib.Nifti1Image(data['subject'].image.data[0,:,:,:].numpy(), affine=data['subject'].image.affine)
    nib.save(nii_img, os.path.join(save_path, name))

    mask_img = nib.Nifti1Image(data['subject']['body_mask'].data[0,:,:,:].numpy(), affine=data['subject'].image.affine)
    nib.save(mask_img, os.path.join(mask_save_path, name))


misc/lymphnode_preprocess_crop.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import csv
import os
import os.path as osp
import nibabel as nib
from sam.datasets.piplines.process3d import RemoveCTBed, CropBackground, Resample
from multiprocessing import Pool, set_start_method


root_fd = '/data/user/codes/SAM_release/data/'
dataset_name = 'NIH_lymph_node'
data_root = osp.join(root_fd, 'raw_data', dataset_name)
save_path = osp.join(root_fd, 'processed_data', dataset_name, 'nii-crop')
mask_save_path = osp.join(root_fd, 'processed_data', dataset_name, 'nii-crop-mask')
ind_save_path = osp.join(root_fd, 'processed_data', dataset_name, 'ind_files')
os.makedirs(save_path, exist_ok=True)
os.makedirs(mask_save_path, exist_ok=True)
os.makedirs(ind_save_path, exist_ok=True)

norm_spacing = 1.6


def process(name):
    print(name)
    filepath = os.path.join(data_root, name)
    img = tio.ScalarImage(filepath)
    subject = tio.Subject(image=img)
    data = {}
    data['subject'] = subject
    removebedtrans = RemoveCTBed()
    data = removebedtrans(data)
    resampletrans = Resample(norm_spacing=(norm_spacing,)*3)
    data = resampletrans(data)
    # croptrans = CropBackground(by_mask=True)  # whether crop background
    # data = croptrans(data)
    nii_img = nib.Nifti1Image(data['subject'].image.data[0, :, :, :].numpy(), affine=data['subject'].image.affine)
    nib.save(nii_img, os.path.join(save_path, name))

    mask_img = nib.Nifti1Image(data['subject']['body_mask'].data[0, :, :, :].numpy(),
                               affine=data['subject'].image.affine)
    nib.save(mask_img, os.path.join(mask_save_path, name))
    return name


def run_pool():
    set_start_method('spawn')
    all_file_name = os.listdir(data_root)
    with Pool(processes=20) as pool:
        allresults = pool.map(process, all_file_name)
    return allresults


if __name__ == '__main__':
    results = run_pool()
    with open(osp.join(ind_save_path, "lymphnode_abd_filename.csv"), "w") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["name"])
        for nii_name in results:
            if not nii_name:
                continue
            else:
                writer.writerow([nii_name])
    print('done')

misc/lymphnode_preprocess_crop_multi_process.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import SimpleITK
import csv
import os
import numpy as np
import nibabel as nib

data_root = '/data/sdd/user/rawdata/lymphnode/Mediastinum/images_nii/'

save_path = '/data/sdd/user/processed_data/lymphnode/nii-1mm/'
if not os.path.exists(save_path):
    os.mkdir(save_path)

all_files = os.listdir(data_root)
print('test')

# with open("/data/sdb/user/mmdet/data/lymphnode_filename_med.csv", "w") as csvfile:
#     writer = csv.writer(csvfile)
#     writer.writerow(["name"])
#     i = 0
i = 0
for name in all_files:
    filepath = os.path.join(data_root, name)
    img = tio.ScalarImage(filepath)
    resampletrans = tio.Resample((1., 1., 1.))
    img_resample = resampletrans(img)

    nii_img = nib.Nifti1Image(img_resample.data[0, :, :, :].numpy(), affine=img_resample.affine)
    nib.save(nii_img, os.path.join(save_path, name))
    # writer.writerow([name])
    print(i)
    i = i + 1

misc/lymphnodenii.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import matplotlib.pyplot as plt
import numpy as np

from skimage import data
from skimage import color
from skimage import morphology
from skimage import segmentation
import torchio as tio
import nibabel as nib
import os
import torch


class RandHistogramShift():
    """
    Apply random nonlinear transform to the image's intensity histogram.
    Args:
        num_control_points: number of control points governing the nonlinear intensity mapping.
            a smaller number of control points allows for larger intensity shifts. if two values provided, number of
            control points selecting from range (min_value, max_value).
        prob: probability of histogram shift.
    """

    def __init__(self, num_control_points=(5, 15), prob: float = 1) -> None:

        if isinstance(num_control_points, int):
            if num_control_points <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (num_control_points, num_control_points)
        else:
            if len(num_control_points) != 2:
                raise ValueError("num_control points should be a number or a pair of numbers")
            if min(num_control_points) <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (min(num_control_points), max(num_control_points))
        self.prob = prob

    def interp(self, x, xp, fp):
        ns = torch if isinstance(x, torch.Tensor) else np
        if isinstance(x, np.ndarray):
            # approx 2x faster than code below for ndarray
            return np.interp(x, xp, fp)

        m = (fp[1:] - fp[:-1]) / (xp[1:] - xp[:-1])
        b = fp[:-1] - (m * xp[:-1])

        indices = ns.searchsorted(xp.reshape(-1), x.reshape(-1)) - 1
        indices = ns.clip(indices, 0, len(m) - 1)

        f = (m[indices] * x.reshape(-1) + b[indices]).reshape(x.shape)
        f[x < xp[0]] = fp[0]
        f[x > xp[-1]] = fp[-1]
        return f

    def randomize(self, flag) -> None:
        num_control_point = np.random.randint(self.num_control_points[0], self.num_control_points[1] + 1)
        if flag == 0:
            reference_control_points = np.linspace(0.17, 0.5, num_control_point)
            reference_control_points_start = np.linspace(0., 0.17, 5)
            reference_control_points_end = np.linspace(0.5, 1., 5)
        else:
            reference_control_points = np.linspace(0.5, 0.67, num_control_point)
            reference_control_points_start = np.linspace(0., 0.5, 5)
            reference_control_points_end = np.linspace(0.67, 1., 5)
        self.reference_control_points = np.concatenate(
            [reference_control_points_start[:-1], reference_control_points, reference_control_points_end[1:]])
        self.floating_control_points = np.copy(self.reference_control_points)
        for i in range(1, num_control_point - 1):
            self.floating_control_points[i] = np.random.uniform(
                self.floating_control_points[i - 1], self.floating_control_points[i + 1]
            )

    def __call__(self, data):
        if np.random.rand() < self.prob:
            img_t = data['subject'].image.data
            img_mask = data['subject'].mask.data
            all_parts = img_mask.max()
            img_t[img_t > 2048] = 2048
            img_t[img_t < -1024] = -1024
            img_t = (img_t + 1024) / 3072.
            # img_t[img_t > 300] = 300
            # img_t[img_t < -200] = -200
            # img_t = (img_t + 200) / 500.

            img_t_flip = 1 - img_t
            img_min, img_max = img_t.min(), img_t.max()
            img_flip_min, img_flip_max = img_t_flip.min(), img_t_flip.max()
            xps = []
            yps = []
            reference_control_points_scaled_all = []
            floating_control_points_scaled_all = []
            flip_flags = []
            for i in range(all_parts):
                if np.random.rand() < 0.3:
                    flag = 1
                    use_min = img_flip_min
                    use_max = img_flip_max
                else:
                    flag = 0
                    use_min = img_min
                    use_max = img_max
                flip_flags.append(flag)
                self.randomize(flag)
                xp = torch.tensor(self.reference_control_points, dtype=img_t.dtype)
                yp = torch.tensor(self.floating_control_points, dtype=img_t.dtype)
                xps.append(xp)
                yps.append(yp)

                # if np.random.rand() < 0.01:  # contrast strength
                #     reference_control_points_scaled = xp * (use_max - use_min) + use_min
                #     floating_control_points_scaled = yp * (use_max - use_min) + use_min
                # else:
                reference_control_points_scaled = xp
                floating_control_points_scaled = yp
                reference_control_points_scaled_all.append(reference_control_points_scaled)
                floating_control_points_scaled_all.append(floating_control_points_scaled)

            img_final = torch.zeros_like(img_t)
            img_final = torch.where(img_mask > 0, img_final, img_t)
            for i in range(all_parts):
                if flip_flags[i] == 0:
                    img_ted = self.interp(img_t, reference_control_points_scaled_all[i],
                                          floating_control_points_scaled_all[i])
                else:
                    img_ted = self.interp(img_t_flip, reference_control_points_scaled_all[i],
                                          floating_control_points_scaled_all[i])
                img_final = torch.where(img_mask == i + 1, img_ted, img_final)
            # import matplotlib.pyplot as plt
            # plt.imshow(img_t[0, :, :, 90], cmap='gray', vmin=0, vmax=1)
            # plt.show()
            data['subject'].image.data = img_final
        return data


if __name__ == '__main__':
    image_path = '/data/sdd/user/rawdata/CT-MRI-Abd/nii/'

    save_path = '/data/sdd/user/rawdata/CT-MRI-Abd/nii-mask-slic'
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    all_img = os.listdir(image_path)
    all_ct_img = [img for img in all_img if '_0000' in img]
    all_mri_img = [img for img in all_img if '_0001' in img]

    for img in all_ct_img:
        print(img)
        image = tio.ScalarImage(os.path.join(image_path, img))
        numpy_img = image.data.numpy()
        numpy_img = numpy_img[0, :, :, :]
        numpy_img = numpy_img.transpose(1, 0, 2)
        mask = numpy_img > -300
        mask_processed = np.copy(mask)
        for i in range(mask.shape[2]):
            mask_processed[:, :, i] = morphology.closing(mask[:, :, i], morphology.disk(3))
            mask_processed[:, :, i] = morphology.opening(mask_processed[:, :, i], morphology.disk(3))
            mask_processed[:, :, i] = morphology.remove_small_holes(
                morphology.remove_small_objects(mask_processed[:, :, i], min_size=2000), area_threshold=25000)
        m_slic = segmentation.slic(numpy_img, n_segments=15, mask=mask_processed, multichannel=False, start_label=1,
                                   compactness=0.01)
        m_slic = m_slic.transpose(1, 0, 2)
        m_slic = m_slic.astype(np.int32)
        nii_img = nib.Nifti1Image(m_slic, affine=image.affine)
        nib.save(nii_img, os.path.join(save_path, img))

misc/make_slic_mask.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import os
import numpy as np
import pickle
import nibabel as nib
import torch
from multiprocessing import Pool, set_start_method
import csv

data_path = '/data/sde/user/Totalsegmentator/Totalsegmentator_dataset'
save_path = '/data/sdd/user/processed_data/totolsegmentaion/nii-resample'
mask_save_path = '/data/sdd/user/processed_data/totolsegmentaion/nii-resample-mask'

#this contains all mask name, you need to generate it
seg_labels = pickle.load(open('/data/sdd/user/processed_data/ind_files/seg_labels.pkl', 'rb'))
# all_seg_files = os.listdir(os.path.join(data_path,case,'segmentations'))
# all_seg_name = [seg_name.split('.',1)[0] for seg_name in all_seg_files]
# f = open('/data/sdd/user/processed_data/ind_files/seg_labels.pkl','wb')
# pickle.dump(all_seg_name,f)


if not os.path.exists(save_path):
    os.mkdir(save_path)
if not os.path.exists(mask_save_path):
    os.mkdir(mask_save_path)


def process(case):
    print(case)
    img = tio.ScalarImage(os.path.join(data_path, case, 'ct.nii.gz'))
    data = img.data
    # assert img.orientation == ('R','A','S')
    ToCanonical = tio.ToCanonical()
    img = ToCanonical(img)
    img_data = img.data
    mask_all = torch.zeros_like(img_data)
    for label, seg_mask in enumerate(seg_labels):
        mask = tio.LabelMap(os.path.join(data_path, case, 'segmentations/', seg_mask + '.nii.gz'))
        mask = ToCanonical(mask)
        mask_data = mask.data
        mask_all[mask_data == 1] = label + 1
    img.data = torch.flip(img.data, (1, 2))
    img.affine = np.array([-img.affine[0, :], -img.affine[1, :], img.affine[2, :], img.affine[3, :]])
    mask_all = torch.flip(mask_all, (1, 2))

    nii_img = nib.Nifti1Image(img.data[0, :, :, :].numpy(), affine=img.affine)
    nib.save(nii_img, os.path.join(save_path, case + '.nii.gz'))
    nii_img = nib.Nifti1Image(mask_all[0, :, :, :].numpy(), affine=img.affine)
    nib.save(nii_img, os.path.join(mask_save_path, case + '.nii.gz'))
    return case + '.nii.gz'


def run_pool():
    set_start_method('spawn')
    all_cases = os.listdir(data_path)
    all_cases = [case for case in all_cases if '.csv' not in case]
    with Pool(processes=20) as pool:
        allresults = pool.map(process, all_cases)
    return allresults


if __name__ == '__main__':
    results = run_pool()
    with open("/data/sdd/user/processed_data/ind_files/totolseg_filename.csv", "w") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["name"])
        for nii_name in results:
            if not nii_name:
                continue
            else:
                writer.writerow([nii_name])

misc/merge_mask_label_multi_process.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
import torchio as tio
import os
import numpy as np

ct_landmark_path = '/data/sdd/user/DEEDs/deedsBCV/ct_mask.nii.gz'
mri_landmark_path = '/data/sdd/user/DEEDs/deedsBCV/nonlinear_mri_ct_deformed_seg.nii.gz'

ct_landmark = tio.LabelMap(ct_landmark_path)
t1_landmark = tio.LabelMap(mri_landmark_path)

ct_landmark_numpy = ct_landmark.data.numpy()[0, :, :, :]
ct_landmark_numpy = ct_landmark_numpy.transpose(1, 0, 2)
all_ct_landmarks = []
for j in range(1, ct_landmark_numpy.max() + 1):
    loc = np.where(ct_landmark_numpy == j)
    loc = np.stack(loc)
    loc = loc.mean(axis=1)
    all_ct_landmarks.append(loc)
all_ct_landmarks = np.stack(all_ct_landmarks)

all_t1_landmarks = []
t1_landmark_numpy = t1_landmark.data.numpy()[0, :, :, :]
t1_landmark_numpy = t1_landmark_numpy.transpose(1, 0, 2)
all_t1_landmarks = []
for j in range(1, ct_landmark_numpy.max() + 1):
    loc = np.where(t1_landmark_numpy == j)
    loc = np.stack(loc)
    loc = loc.mean(axis=1)
    all_t1_landmarks.append(loc)
all_t1_landmarks = np.stack(all_t1_landmarks)

ct_not_nan_index = np.unique(np.where(np.isnan(all_ct_landmarks) == False)[0])
all_ct_landmarks = all_ct_landmarks[ct_not_nan_index, :]
all_t1_landmarks = all_t1_landmarks[ct_not_nan_index, :]
dis_case = []
for j in range(all_ct_landmarks.shape[0]):
    gth = all_ct_landmarks[j,:]
    pre = all_t1_landmarks[j,:]
    if np.isnan(pre).any():
        continue
    dis = np.linalg.norm(gth - pre)
    # print(str(j)+','+str(dis))
    dis_case.append(dis)
dis_here = np.array(dis_case)
# print(dis_case)
print(dis_here.mean(), dis_here.std(), dis_here.max())

print('done')
misc/metric/compute_landmark_distance.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import os
import torchio as tio
import numpy as np
import torch
import nibabel as nib
from scipy import ndimage
from sam.datasets.piplines.process3d import Resample
from multiprocessing import Pool, set_start_method
import csv
from skimage import morphology
from skimage import segmentation

data_path = '/data/sde/user/abd_registered/abd_registered_1_1.8/'
img_save_path = '/data/sde/user/abd_registered/processed_1_1.8/image-regis_1/'
mask_save_path = '/data/sde/user/abd_registered/processed_1_1.8/mask-regis_1/'

if not os.path.exists(img_save_path):
    os.mkdir(img_save_path)
if not os.path.exists(mask_save_path):
    os.mkdir(mask_save_path)


def process(name):
    print(name)
    ct_img = tio.ScalarImage(os.path.join(data_path, name + '_0000.nii.gz'))
    t1_img = tio.ScalarImage(os.path.join(data_path, name + '_0001.nii.gz'))
    # t2_img = tio.ScalarImage(os.path.join(data_path, name + '_0002.nii.gz'))

    t1_data = t1_img.data.numpy()
    # t2_data = t2_img.data.numpy()
    ct_data = ct_img.data.numpy()
    t1_data = t1_data[0, :, :, :]
    t1_data = t1_data.transpose(1, 0, 2)
    # t2_data = t2_data[0, :, :, :]
    # t2_data = t2_data.transpose(1, 0, 2)
    ct_data = ct_data[0, :, :, :]
    ct_data = ct_data.transpose(1, 0, 2)
    t1_min_max = np.percentile(t1_data, (0.2, 99.8))
    # t2_min_max = np.percentile(t2_data, (2, 99.8))
    t1_mask_data = t1_data > (t1_min_max[0] + 30)
    # t2_mask_data = t2_data > (t2_min_max[0] + 50)
    ct_mask_data = (ct_data > -300)

    assert ct_img.orientation == ('L', 'P', 'S')

    ct_mask_processed = np.copy(ct_mask_data)
    t1_mask_processed = np.copy(t1_mask_data)
    # t2_mask_processed = np.copy(t2_mask_data)

    for i in range(ct_mask_data.shape[2]):
        ct_mask_processed[:, :, i] = morphology.closing(ct_mask_data[:, :, i], morphology.disk(3))
        ct_mask_processed[:, :, i] = morphology.opening(ct_mask_processed[:, :, i], morphology.disk(3))
        ct_mask_processed[:, :, i] = morphology.remove_small_holes(
            morphology.remove_small_objects(ct_mask_processed[:, :, i], min_size=2000), area_threshold=30000)
    # ct_slic = segmentation.slic(ct_data, n_segments=1, mask=ct_mask_processed, multichannel=False, start_label=1,
    #                            compactness=0.01)

    for i in range(t1_mask_data.shape[2]):
        t1_mask_processed[:, :, i] = morphology.closing(t1_mask_data[:, :, i], morphology.disk(3))
        t1_mask_processed[:, :, i] = morphology.opening(t1_mask_processed[:, :, i], morphology.disk(3))
        t1_mask_processed[:, :, i] = morphology.remove_small_holes(
            morphology.remove_small_objects(t1_mask_processed[:, :, i], min_size=2000), area_threshold=30000)
    # t1_slic = segmentation.slic(t1_data, n_segments=1, mask=t1_mask_processed, multichannel=False, start_label=1,
    #                            compactness=0.05)

    # for i in range(t2_mask_data.shape[2]):
    #     t2_mask_processed[:, :, i] = morphology.closing(t2_mask_data[:, :, i], morphology.disk(3))
    #     t2_mask_processed[:, :, i] = morphology.opening(t2_mask_processed[:, :, i], morphology.disk(3))
    #     t2_mask_processed[:, :, i] = morphology.remove_small_holes(
    #         morphology.remove_small_objects(t2_mask_processed[:, :, i], min_size=2000), area_threshold=25000)
    # t2_slic = segmentation.slic(t1_data, n_segments=15, mask=t2_mask_processed, multichannel=False, start_label=1,
    #                            compactness=0.05)

    t1_mask = tio.LabelMap(tensor=t1_mask_processed.transpose(1, 0, 2)[np.newaxis, :, :, :], affine=t1_img.affine)
    # t2_mask = tio.LabelMap(tensor=t2_slic.transpose(1, 0, 2)[np.newaxis, :, :, :], affine=t2_img.affine)
    ct_mask = tio.LabelMap(tensor=ct_mask_processed.transpose(1, 0, 2)[np.newaxis, :, :, :], affine=ct_img.affine)

    subject = tio.Subject(image=ct_img,
                          t1=t1_img,
                          # t2=t2_img,
                          ct_mask=ct_mask,
                          t1_mask=t1_mask, )
    # t2_mask=t2_mask)

    # ToCanonical = tio.ToCanonical()
    # subject = ToCanonical(subject)
    # print(subject.orientation)
    # if subject.orientation == ('R', 'A', 'S'):
    #     img_data = img_tio.data
    #     img_tio.data = torch.flip(img_data, (1, 2))
    #     img_tio.affine = np.array(
    #         [-img_tio.affine[0, :], -img_tio.affine[1, :], img_tio.affine[2, :], img_tio.affine[3, :]])
    data = {}
    resampletrans = Resample(norm_spacing=(1.6, 1.6, 1.6), crop_artifacts=False)
    data['subject'] = subject
    data = resampletrans(data)
    subject = data['subject']

    nii_img = nib.Nifti1Image(subject.image.data[0, :, :, :].numpy(), affine=subject.image.affine)
    nib.save(nii_img, os.path.join(img_save_path, name + '_0000.nii.gz'))
    nii_img = nib.Nifti1Image(subject.t1.data[0, :, :, :].numpy(), affine=subject.t1.affine)
    nib.save(nii_img, os.path.join(img_save_path, name + '_0001.nii.gz'))
    # nii_img = nib.Nifti1Image(subject.t2.data[0, :, :, :].numpy(), affine=subject.t2.affine)
    # nib.save(nii_img, os.path.join(img_save_path, name+'_0002.nii.gz'))

    nii_img = nib.Nifti1Image(subject.ct_mask.data[0, :, :, :].numpy().astype(np.int16), affine=subject.ct_mask.affine)
    nib.save(nii_img, os.path.join(mask_save_path, name + '_0000.nii.gz'))
    nii_img = nib.Nifti1Image(subject.t1_mask.data[0, :, :, :].numpy().astype(np.int16), affine=subject.t1_mask.affine)
    nib.save(nii_img, os.path.join(mask_save_path, name + '_0001.nii.gz'))
    zs = np.where(subject.t1_mask.data[0, :, :, :] == 1)[2]
    # print(zs.max()-zs.min())
    # nii_img = nib.Nifti1Image(subject.t2_mask.data[0, :, :, :].numpy().astype(np.int16), affine=subject.t2_mask.affine)
    # nib.save(nii_img, os.path.join(mask_save_path, name + '_0002.nii.gz'))
    return name


def run_pool():
    set_start_method('spawn')
    all_elements = []
    # deeplesion_nii_path = '/data/sdb/user/Deeplesion/Images_nifti'
    # all_niis = os.listdir(deeplesion_nii_path)
    all_file_name = os.listdir(data_path)
    prefix_name = [name.split('_', 1)[0] for name in all_file_name]
    prefix_name = np.unique(prefix_name)
    with Pool(processes=20) as pool:
        allresults = pool.map(process, prefix_name)
    return allresults


if __name__ == '__main__':
    results = run_pool()
    # with open("/data/sdd/user/processed_data/ind_files/CT_MRI_abd_inter_0_filename.csv", "w") as csvfile:
    #     writer = csv.writer(csvfile)
    #     writer.writerow(["name"])
    #     for nii_name in results:
    #         if not nii_name:
    #             continue
    #         else:
    #             writer.writerow([nii_name])

misc/mri_ct_multi_process.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import os
import torchio as tio
import numpy as np
import torch
import nibabel as nib
from scipy import ndimage
from skimage import morphology
from skimage import segmentation


data_path = '/data/sde/user/abd_registered/abd_registered_0/'
img_save_path = '/data/sde/user/abd_registered/processed_0/image-regis_0/'
mask_save_path = '/data/sde/user/abd_registered/processed_0/mask-regis_0/'

if not os.path.exists(img_save_path):
    os.mkdir(img_save_path)
if not os.path.exists(mask_save_path):
    os.mkdir(mask_save_path)

all_file_name = os.listdir(data_path)

prefix_name = [name.split('_', 1)[0] for name in all_file_name]
prefix_name = np.unique(prefix_name)

for name in prefix_name:
    if '970107' not in name:
        continue
    print(name)
    ct_img = tio.ScalarImage(os.path.join(data_path, name + '_0000.nii.gz'))
    t1_img = tio.ScalarImage(os.path.join(data_path, name + '_0001.nii.gz'))
    # t2_img = tio.ScalarImage(os.path.join(data_path, name + '_0002.nii.gz'))

    t1_data = t1_img.data.numpy()
    t1_data = t1_data[0, :, :, :]
    t1_data = t1_data.transpose(1, 0, 2)
    # t2_data = t2_img.data.numpy()
    ct_data = ct_img.data.numpy()
    ct_data = ct_data[0, :, :, :]
    ct_data = ct_data.transpose(1, 0, 2)
    assert ct_img.orientation == ('L', 'P', 'S')
    t1_min_max = np.percentile(t1_data, (2, 99.8))
    # t2_min_max = np.percentile(t2_data, (2, 99.8))
    t1_mask_data = t1_data > (t1_min_max[0] + 50)
    # t2_mask_data = t2_data > (t2_min_max[0] + 10)
    ct_mask_data = (ct_data > -300)
    ct_mask_processed = np.copy(ct_mask_data)
    t1_mask_processed = np.copy(t1_mask_data)

    for i in range(ct_mask_data.shape[2]):
        ct_mask_processed[:, :, i] = morphology.closing(ct_mask_data[:, :, i], morphology.disk(3))
        ct_mask_processed[:, :, i] = morphology.opening(ct_mask_processed[:, :, i], morphology.disk(3))
        ct_mask_processed[:, :, i] = morphology.remove_small_holes(
            morphology.remove_small_objects(ct_mask_processed[:, :, i], min_size=2000), area_threshold=25000)
    # ct_slic = segmentation.slic(ct_data, n_segments=15, mask=ct_mask_processed, multichannel=False, start_label=1,
    #                            compactness=0.01)

    for i in range(t1_mask_data.shape[2]):
        t1_mask_processed[:, :, i] = morphology.closing(t1_mask_data[:, :, i], morphology.disk(3))
        t1_mask_processed[:, :, i] = morphology.opening(t1_mask_processed[:, :, i], morphology.disk(3))
        t1_mask_processed[:, :, i] = morphology.remove_small_holes(
            morphology.remove_small_objects(t1_mask_processed[:, :, i], min_size=2000), area_threshold=25000)
    # t1_slic = segmentation.slic(t1_data, n_segments=15, mask=t1_mask_processed, multichannel=False, start_label=1,
    #                            compactness=0.05)

    t1_mask = tio.LabelMap(tensor=t1_mask_processed.transpose(1, 0, 2)[np.newaxis, :, :, :], affine=t1_img.affine)
    # t2_mask = tio.LabelMap(tensor=t2_mask_data, affine=t2_img.affine)
    ct_mask = tio.LabelMap(tensor=ct_mask_processed.transpose(1, 0, 2)[np.newaxis, :, :, :], affine=ct_img.affine)

    subject = tio.Subject(ct=ct_img,
                          t1=t1_img,
                          ct_mask=ct_mask,
                          t1_mask=t1_mask, )

    # ToCanonical = tio.ToCanonical()
    # subject = ToCanonical(subject)
    resampletrans = tio.Resample((1.6, 1.6, 1.6))
    subject = resampletrans(subject)

    nii_img = nib.Nifti1Image(subject.ct.data[0, :, :, :].numpy(), affine=subject.ct.affine)
    nib.save(nii_img, os.path.join(img_save_path, name + '_0000.nii.gz'))
    nii_img = nib.Nifti1Image(subject.t1.data[0, :, :, :].numpy(), affine=subject.t1.affine)
    nib.save(nii_img, os.path.join(img_save_path, name + '_0001.nii.gz'))
    # nii_img = nib.Nifti1Image(subject.t2.data[0, :, :, :].numpy(), affine=subject.t2.affine)
    # nib.save(nii_img, os.path.join(img_save_path, name+'_0002.nii.gz'))

    nii_img = nib.Nifti1Image(subject.ct_mask.data[0, :, :, :].numpy().astype(np.int16), affine=subject.ct_mask.affine)
    nib.save(nii_img, os.path.join(mask_save_path, name + '_0000.nii.gz'))
    nii_img = nib.Nifti1Image(subject.t1_mask.data[0, :, :, :].numpy().astype(np.int16), affine=subject.t1_mask.affine)
    nib.save(nii_img, os.path.join(mask_save_path, name + '_0001.nii.gz'))
    # nii_img = nib.Nifti1Image(subject.t2_mask.data[0, :, :, :].numpy(), affine=subject.t2_mask.affine)
    # nib.save(nii_img, os.path.join(mask_save_path, name + '_0002.nii.gz'))

    print('done')

print('done')

misc/mri_ct_process.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import matplotlib.pyplot as plt
import numpy as np

from skimage import data
from skimage import color
from skimage import morphology
from skimage import segmentation
import torchio as tio
import nibabel as nib
import os
import torch


class RandHistogramShift():
    """
    Apply random nonlinear transform to the image's intensity histogram.
    Args:
        num_control_points: number of control points governing the nonlinear intensity mapping.
            a smaller number of control points allows for larger intensity shifts. if two values provided, number of
            control points selecting from range (min_value, max_value).
        prob: probability of histogram shift.
    """

    def __init__(self, num_control_points=(5, 15), prob: float = 0.9) -> None:

        if isinstance(num_control_points, int):
            if num_control_points <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (num_control_points, num_control_points)
        else:
            if len(num_control_points) != 2:
                raise ValueError("num_control points should be a number or a pair of numbers")
            if min(num_control_points) <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (min(num_control_points), max(num_control_points))
        self.prob = prob

    def interp(self, x, xp, fp):
        ns = torch if isinstance(x, torch.Tensor) else np
        if isinstance(x, np.ndarray):
            # approx 2x faster than code below for ndarray
            return np.interp(x, xp, fp)

        m = (fp[1:] - fp[:-1]) / (xp[1:] - xp[:-1])
        b = fp[:-1] - (m * xp[:-1])

        indices = ns.searchsorted(xp.reshape(-1), x.reshape(-1)) - 1
        indices = ns.clip(indices, 0, len(m) - 1)

        f = (m[indices] * x.reshape(-1) + b[indices]).reshape(x.shape)
        f[x < xp[0]] = fp[0]
        f[x > xp[-1]] = fp[-1]
        return f

    def randomize(self) -> None:
        num_control_point = np.random.randint(self.num_control_points[0], self.num_control_points[1] + 1)
        self.reference_control_points = np.linspace(0, 1, num_control_point)
        self.floating_control_points = np.copy(self.reference_control_points)
        for i in range(1, num_control_point - 1):
            self.floating_control_points[i] = np.random.uniform(
                self.floating_control_points[i - 1], self.floating_control_points[i + 1]
            )

    def __call__(self, data):
        if np.random.rand() < self.prob:
            img_t = data['subject'].image.data
            img_mask = data['subject'].mask.data
            self.randomize()
            if self.reference_control_points is None or self.floating_control_points is None:
                raise RuntimeError("please call the `randomize()` function first.")
            if np.random.rand() < 0.5:  # flip intensity
                img_t2 = 1 - img_t
            else:
                img_t2 = img_t

            img_min, img_max = img_t2.min(), img_t2.max()
            xp = torch.tensor(self.reference_control_points, dtype=img_t.dtype)
            yp = torch.tensor(self.floating_control_points, dtype=img_t.dtype)
            if np.random.rand() < 0.5:  # contrast strength
                reference_control_points_scaled = xp * (img_max - img_min) + img_min
                floating_control_points_scaled = yp * (img_max - img_min) + img_min
            else:
                reference_control_points_scaled = xp
                floating_control_points_scaled = yp
            img_ted = self.interp(img_t2, reference_control_points_scaled, floating_control_points_scaled)
            img_t = torch.where(img_mask > 0, img_ted, img_t)
            # import matplotlib.pyplot as plt
            # plt.imshow(img_t[0, :, :, 90], cmap='gray', vmin=0, vmax=1)
            # plt.show()
            data['subject'].image.data = img_t
        return data


if __name__ == '__main__':
    image_path = '/data/sdd/user/rawdata/CT-MRI-Abd/nii/1058150_0000.nii.gz'
    save_path = '/data/sdd/user/rawdata/CT-MRI-Abd/'
    image = tio.ScalarImage(image_path)
    numpy_img = image.data.numpy()
    numpy_img = numpy_img[0, :, :, :]
    numpy_img = numpy_img.transpose(1, 0, 2)

    body_thres = (600, 50)
    # img_min_max = np.percentile(numpy_img, (2, 99.8)) # for MRI image
    img_min_max = (-1024, 3071)  # for CT image
    mask = numpy_img > (img_min_max[0] + body_thres[0])
    mask_processed = np.copy(mask)
    numpy_img[numpy_img < img_min_max[0]] = img_min_max[0]
    numpy_img[numpy_img > img_min_max[1]] = img_min_max[1]

    numpy_img = (numpy_img - img_min_max[0]) / (img_min_max[1] - img_min_max[0])
    for i in range(mask.shape[2]):
        mask_processed[:, :, i] = morphology.closing(mask[:, :, i], morphology.disk(3))
        mask_processed[:, :, i] = morphology.opening(mask_processed[:, :, i], morphology.disk(3))
        mask_processed[:, :, i] = morphology.remove_small_holes(
            morphology.remove_small_objects(mask_processed[:, :, i], min_size=2000), area_threshold=25000)
    # Superpixel
    # m_slic = segmentation.slic(numpy_img, n_segments=15, mask=mask_processed, multichannel=False, start_label=1,
    #                            compactness=0.05)
    # m_slic = m_slic.transpose(1, 0, 2)
    # m_slic = m_slic.astype(np.int32)
    # nii_img = nib.Nifti1Image(m_slic, affine=image.affine)
    # nib.save(nii_img, os.path.join(save_path, '820691_0000_slic' + '.nii.gz'))

    mask_processed = mask_processed.transpose(1, 0, 2)
    mask_processed = mask_processed.astype(np.int32)
    nii_img = nib.Nifti1Image(mask_processed, affine=image.affine)
    nib.save(nii_img, os.path.join(save_path, '1058150_0000_mask' + '.nii.gz'))

    numpy_img = numpy_img.transpose(1, 0, 2)
    subject = tio.Subject(image=tio.ScalarImage(tensor=numpy_img[np.newaxis, :, :, :], affine=image.affine),
                          mask=tio.LabelMap(tensor=mask_processed[np.newaxis, :, :, :], affine=image.affine))
    ranodmlize = RandHistogramShift()
    data = {}
    data['subject'] = subject
    data = ranodmlize(data)
    processed_img = data['subject'].image.data.numpy()[0, :, :, :]
    nii_img = nib.Nifti1Image(processed_img, affine=image.affine)
    nib.save(nii_img, os.path.join(save_path, '1058150_0000_aug' + '.nii.gz'))

    print('done')

misc/testslic.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import xml.etree.ElementTree as ET

import SimpleITK as sitk
import cv2
import numpy as np

"""
date:2020/5/6
"""


def read_file(filepath):
    """
    read the information in the .annot (itk-snap annotation file)
    :param filepath: the directionary of your targeted file
    :return: RECIST: x11,y11,x12,y12,x21,y21,x22,y22
    slice_th: the slice of RECIST
    """
    tree = ET.parse(filepath)
    root = tree.getroot()
    recist = [[]]
    slice_th = 0
    slice_recist = []
    z_point = []
    count = 0
    for child in root:
        if child.tag == 'folder':
            for grandkid in child:
                if grandkid.tag == 'folder':
                    color = list(map(float, grandkid[0].attrib['value'].split(' ')))
                    if count == 0:
                        count = count + 1
                        coloruse = color
                    if color[0] == coloruse[0] and color[1] == coloruse[1] and color[2] == coloruse[2]:
                        try:
                            point1 = list(map(float, grandkid[2].attrib['value'].split(' ')))
                            point2 = list(map(float, grandkid[3].attrib['value'].split(' ')))
                            if (point1[0] != point2[0] or point1[1] != point2[1] or point1[-1] != point2[-1]):
                                if slice_th != 0 and point1[-1] == slice_th:
                                    recist[-1] += point1[0:2] + point2[0:2]
                                elif slice_th != 0:
                                    recist.append(point1[0:2] + point2[0:2])
                                    slice_th = point1[-1]
                                    slice_recist.append(int(slice_th))
                                else:
                                    slice_th = point1[-1]
                                    recist[0] = point1[0:2] + point2[0:2]
                                    slice_recist.append(int(slice_th))
                            else:
                                z_point.append(point1[-1])
                        except:
                            print('There is no point')
    # return recist,slice_recist
    return recist, min(z_point), max(z_point), slice_recist


def read_file1(filepath):
    """
    read the information in the .annot (itk-snap annotation file)
    :param filepath: the directionary of your targeted file
    :return: RECIST: x11,y11,x12,y12,x21,y21,x22,y22
    slice_th: the slice of RECIST
    """
    tree = ET.parse(filepath)
    root = tree.getroot()
    recist = [[]]
    slice_th = []
    slice_recist = []
    z_point = []
    count = 0
    for child in root:
        if child.tag == 'folder':
            for grandkid in child:
                if grandkid.tag == 'folder':
                    color = list(map(float, grandkid[0].attrib['value'].split(' ')))
                    if count == 0:
                        count = count + 1
                        coloruse = color
                    if color[0] == coloruse[0] and color[1] == coloruse[1] and color[2] == coloruse[2]:
                        try:
                            point1 = list(map(float, grandkid[2].attrib['value'].split(' ')))
                            point2 = list(map(float, grandkid[3].attrib['value'].split(' ')))
                            if (point1[0] != point2[0] or point1[1] != point2[1] or point1[-1] != point2[-1]):
                                slice_th_set = set(slice_th)
                                if len(slice_th) != 0 and point1[-1] in slice_th_set:
                                    p = slice_th.index(point1[-1])
                                    recist[p] += point1[0:2] + point2[0:2]
                                elif len(slice_th) != 0:
                                    recist.append(point1[0:2] + point2[0:2])
                                    slice_th.append(point1[-1])
                                    slice_recist.append(int(point1[-1]))
                                else:
                                    slice_th.append(point1[-1])
                                    recist[0] = point1[0:2] + point2[0:2]
                                    slice_recist.append(int(point1[-1]))
                            else:
                                z_point.append(point1[-1])
                        except:
                            print('There is no point')
    return recist, slice_recist
    # return recist,min(z_point),max(z_point),slice_recist


def gen_mask_polygon_from_recist(recist):
    """
    Generate ellipse from RECIST
    :param recist:
    :return:
    """
    assert len(recist) == 8
    x11, y11, x12, y12, x21, y21, x22, y22 = recist
    axis1 = np.linalg.solve(np.array([[x11, y11], [x12, y12]]), np.array([1, 1]))
    axis2 = np.linalg.solve(np.array([[x21, y21], [x22, y22]]), np.array([1, 1]))
    center = np.linalg.solve(np.array([[axis1[0], axis1[1]], [axis2[0], axis2[1]]]), np.array([1, 1]))
    center_recist = recist - np.tile(center, (4,))
    center_recist = np.reshape(center_recist, (4, 2))
    pt_angles = np.arctan2(center_recist[:, 1], center_recist[:, 0])
    pt_lens = np.sqrt(np.sum(center_recist ** 2, axis=1))
    ord = [0, 2, 1, 3, 0]
    grid = .1
    rotated_pts = []
    for p in range(4):
        if (pt_angles[ord[p]] < pt_angles[ord[p + 1]] and pt_angles[ord[p + 1]] - pt_angles[ord[p]] < np.pi) \
                or (pt_angles[ord[p]] - pt_angles[ord[p + 1]] > np.pi):
            angles = np.arange(0, np.pi / 2, grid)
        else:
            angles = np.arange(0, -np.pi / 2, -grid)
        xs = np.cos(angles) * pt_lens[ord[p]]
        ys = np.sin(angles) * pt_lens[ord[p + 1]]
        r = pt_angles[ord[p]]
        rotated_pts1 = np.matmul(np.array([[np.cos(r), -np.sin(r)], [np.sin(r), np.cos(r)]]), np.vstack((xs, ys)))
        rotated_pts.append(rotated_pts1)
    rotated_pts = np.hstack(rotated_pts)
    decentered_pts = rotated_pts + center.reshape((2, 1))
    polygon = decentered_pts.transpose().ravel()
    return polygon.tolist()


def getcross(recist, z):
    """
    get the recist crossline as [(x1,y1,z1),...,(xn,yn,zn)]
    :param recist:
    :param z:
    :return:
    """
    x11, y11, x12, y12, x21, y21, x22, y22 = recist
    axis1 = np.linalg.solve(np.array([[x11, y11], [x12, y12]]), np.array([1, 1]))
    axis2 = np.linalg.solve(np.array([[x21, y21], [x22, y22]]), np.array([1, 1]))
    a1 = axis1[0]
    b1 = axis1[1]
    a2 = axis2[0]
    b2 = axis2[1]
    line = []

    z = int(z)
    if abs(a1 / b1) < 1:
        for x in range(int(min(x11, x12) + 1), int(max(x11, x12) - 1)):
            line.append((int(x), int(-a1 / b1 * x + 1 / b1), z))
            # for numpy array use '[]'
            # for simpleitk image change '[]' to '()'
    else:
        for y in range(int(min(y11, y12) + 1), int(max(y11, y12) - 1)):
            line.append((int(-b1 / a1 * y + 1 / a1), int(y), z))

    if abs(a2 / b2) < 1:
        for x in range(int(min(x21, x22) + 1), int(max(x21, x22) - 1)):
            line.append((int(x), int(-a2 / b2 * x + 1 / b2), z))
    else:
        for y in range(int(min(y21, y22) + 1), int(max(y21, y22) - 1)):
            line.append((int(-b2 / a2 * y + 1 / a2), int(y), z))
    return line


def get_polygon_and_line(recist, recist_slice):
    polygon_all = []
    line_all = []
    for i in range(len(recist)):
        polygon_tmp = gen_mask_polygon_from_recist(recist[i])
        line_tmp = getcross(recist[i], recist_slice[i])
        polygon_all.append(polygon_tmp)
        line_all.append(line_tmp)
    return polygon_all, line_all


def get_cross_from_annot(recist):
    """
    support multi-slices annotation and on each slice support multi-crosses
    :param recist: get from read_file1
    :return: cross_list [i][j] the jth cross on ith slice
    """
    cross_list = []
    for i in range(recist.__len__()):
        recist_sovle = recist[i]
        recist_sovle = np.asarray(recist_sovle)
        line_num = int(recist_sovle.__len__() / 4)
        recist_sovle = np.reshape(recist_sovle, (line_num, -1))

        total = 0
        for j in range(line_num):
            line1 = recist_sovle[j]
            for k in range(j + 1, line_num):
                line2 = recist_sovle[k]

                x1_min = min(line1[0], line1[2])
                x1_max = max(line1[0], line1[2])
                y1_min = min(line1[1], line1[3])
                y1_max = max(line1[1], line1[3])

                x2_min = min(line2[0], line2[2])
                x2_max = max(line2[0], line2[2])
                y2_min = min(line2[1], line2[3])
                y2_max = max(line2[1], line2[3])

                x_min = max(x1_min, x2_min)
                x_max = min(x1_max, x2_max)
                y_min = max(y1_min, y2_min)
                y_max = min(y1_max, y2_max)

                axis1 = np.linalg.solve(np.array([[line1[0], line1[1]], [line1[2], line1[3]]]), np.array([1, 1]))
                axis2 = np.linalg.solve(np.array([[line2[0], line2[1]], [line2[2], line2[3]]]), np.array([1, 1]))
                if max(abs(axis1[0] - axis2[0]), abs(axis1[1] - axis2[1])) < 0.001:  # singular value#
                    continue
                center = np.linalg.solve(np.array([[axis1[0], axis1[1]], [axis2[0], axis2[1]]]), np.array([1, 1]))

                if center[0] <= x_max + 0.01 and center[
                    0] >= x_min - 0.01:  # relaxing by 0.01 for vertical and horizontal lines
                    if center[1] <= y_max + 0.01 and center[1] >= y_min - 0.01:
                        # cross_list.append([line1[0],line1[1],line1[2],line1[3],line2[0],line2[1],line2[2],line2[3]])
                        if total == 0:
                            total = 1
                            cross_list.append(
                                [[line1[0], line1[1], line1[2], line1[3], line2[0], line2[1], line2[2], line2[3]]])
                        else:
                            cross_list[-1] += [
                                [line1[0], line1[1], line1[2], line1[3], line2[0], line2[1], line2[2], line2[3]]]
    return cross_list


def ImageResample(sitk_image, new_spacing=[1.0, 1.0, 3.0], is_label=False):
    '''
        sitk_image:
        new_spacing: x,y,z
        is_label: if True, using Interpolator `sitk.sitkNearestNeighbor`
        '''
    size = np.array(sitk_image.GetSize())
    spacing = np.array(sitk_image.GetSpacing())
    new_spacing = np.array(new_spacing)
    new_size = size * spacing / new_spacing
    new_spacing_refine = size * spacing / new_size
    new_spacing_refine = [float(s) for s in new_spacing_refine]
    new_size = [int(s) for s in new_size]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputDirection(sitk_image.GetDirection())
    resample.SetOutputOrigin(sitk_image.GetOrigin())
    resample.SetSize(new_size)
    resample.SetOutputSpacing(new_spacing_refine)

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        # resample.SetInterpolator(sitk.sitkNearestNeighbor)
        # resample.SetInterpolator(sitk.sitkBSpline)
        resample.SetInterpolator(sitk.sitkLinear)

    newimage = resample.Execute(sitk_image)
    return newimage


def bb_intersection_over_union(boxA, boxB):
    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2] + boxA[0], boxB[2] + boxB[0])
    yB = min(boxA[3] + boxA[1], boxB[3] + boxB[1])

    # compute the area of intersection rectangle
    interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
    if interArea == 0:
        iou1 = 0
        iou2 = 0
    # compute the area of both the prediction and ground-truth
    # rectangles
    # boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
    # boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))
    boxAArea = abs(boxA[2] * boxA[3])
    boxBArea = abs(boxB[2] * boxB[3])

    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    # iou = interArea / float(boxAArea + boxBArea - interArea)
    iou1 = float(interArea) / float(boxAArea)
    iou2 = float(interArea) / float(boxBArea)
    # return the intersection over union value
    return iou1, iou2


def draw(rects, color, ocv):
    for r in rects:
        p1 = (r[0], r[1])
        p2 = (r[0] + r[2], r[1] + r[3])
        cv2.rectangle(ocv, p1, p2, color, 2)


if __name__ == '__main__':
    anno_path = '/data/sdd/user/processed_data/anno/000133_02_01_038-050.annot'
    a = read_file(anno_path)
    print('done')

misc/utils.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .models.frameworks import sam
from .models.backbones import resnet3d  # ,unet3d,Unetrpp,loftr,convnext,convnext3d
from .models.necks import fpn3d
from .datasets import dataset3dsam, dataset3dsam_cross_volume
from .apis import train_detector

sam/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .train import train_detector
sam/apis/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings

import numpy as np
import torch
import torch.distributed as dist
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (DistSamplerSeedHook, EpochBasedRunner,
                         Fp16OptimizerHook, OptimizerHook, build_optimizer,
                         build_runner, get_dist_info)

from mmdet.core import DistEvalHook, EvalHook
from mmdet.datasets import build_dataset, replace_ImageToTensor
from sam.datasets import build_dataloader
from mmdet.utils import find_latest_checkpoint, get_root_logger


def init_random_seed(seed=None, device='cuda'):
    """Initialize random seed.

    If the seed is not set, the seed will be automatically randomized,
    and then broadcast to all processes to prevent some potential bugs.

    Args:
        seed (int, Optional): The seed. Default to None.
        device (str): The device where the seed will be put on.
            Default to 'cuda'.

    Returns:
        int: Seed to be used.
    """
    if seed is not None:
        return seed

    # Make sure all ranks share the same random seed to prevent
    # some potential bugs. Please refer to
    # https://github.com/open-mmlab/mmdetection/issues/6339
    rank, world_size = get_dist_info()
    seed = np.random.randint(2**31)
    if world_size == 1:
        return seed

    if rank == 0:
        random_num = torch.tensor(seed, dtype=torch.int32, device=device)
    else:
        random_num = torch.tensor(0, dtype=torch.int32, device=device)
    dist.broadcast(random_num, src=0)
    return random_num.item()


def set_random_seed(seed, deterministic=False):
    """Set random seed.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
            to True and `torch.backends.cudnn.benchmark` to False.
            Default: False.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def train_detector(model,
                   dataset,
                   cfg,
                   distributed=False,
                   validate=False,
                   timestamp=None,
                   meta=None):
    logger = get_root_logger(log_level=cfg.log_level)

    # prepare data loaders
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
    if 'imgs_per_gpu' in cfg.data:
        logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
                       'Please use "samples_per_gpu" instead')
        if 'samples_per_gpu' in cfg.data:
            logger.warning(
                f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
                f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
                f'={cfg.data.imgs_per_gpu} is used in this experiments')
        else:
            logger.warning(
                'Automatically set "samples_per_gpu"="imgs_per_gpu"='
                f'{cfg.data.imgs_per_gpu} in this experiments')
        cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu

    runner_type = 'EpochBasedRunner' if 'runner' not in cfg else cfg.runner[
        'type']
    data_loaders = [
        build_dataloader(
            ds,
            cfg.data.samples_per_gpu,
            cfg.data.workers_per_gpu,
            # `num_gpus` will be ignored if distributed
            num_gpus=len(cfg.gpu_ids),
            dist=distributed,
            seed=cfg.seed,
            runner_type=runner_type,
            persistent_workers=cfg.data.get('persistent_workers', False))
        for ds in dataset
    ]

    # put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        # Sets the `find_unused_parameters` parameter in
        # torch.nn.parallel.DistributedDataParallel
        model = MMDistributedDataParallel(
            model.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False,
            find_unused_parameters=find_unused_parameters)
    else:
        model = MMDataParallel(
            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)

    # build runner
    optimizer = build_optimizer(model, cfg.optimizer)

    if 'runner' not in cfg:
        cfg.runner = {
            'type': 'EpochBasedRunner',
            'max_epochs': cfg.total_epochs
        }
        warnings.warn(
            'config is now expected to have a `runner` section, '
            'please set `runner` in your config.', UserWarning)
    else:
        if 'total_epochs' in cfg:
            assert cfg.total_epochs == cfg.runner.max_epochs

    runner = build_runner(
        cfg.runner,
        default_args=dict(
            model=model,
            optimizer=optimizer,
            work_dir=cfg.work_dir,
            logger=logger,
            meta=meta))

    # an ugly workaround to make .log and .log.json filenames the same
    runner.timestamp = timestamp

    # fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
    elif distributed and 'type' not in cfg.optimizer_config:
        optimizer_config = OptimizerHook(**cfg.optimizer_config)
    else:
        optimizer_config = cfg.optimizer_config

    # register hooks
    runner.register_training_hooks(
        cfg.lr_config,
        optimizer_config,
        cfg.checkpoint_config,
        cfg.log_config,
        cfg.get('momentum_config', None),
        custom_hooks_config=cfg.get('custom_hooks', None))

    if distributed:
        if isinstance(runner, EpochBasedRunner):
            runner.register_hook(DistSamplerSeedHook())

    # register eval hooks
    if validate:
        # Support batch_size > 1 in validation
        val_samples_per_gpu = cfg.data.val.pop('samples_per_gpu', 1)
        if val_samples_per_gpu > 1:
            # Replace 'ImageToTensor' to 'DefaultFormatBundle'
            cfg.data.val.pipeline = replace_ImageToTensor(
                cfg.data.val.pipeline)
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        val_dataloader = build_dataloader(
            val_dataset,
            samples_per_gpu=val_samples_per_gpu,
            workers_per_gpu=cfg.data.workers_per_gpu,
            dist=distributed,
            shuffle=False)
        eval_cfg = cfg.get('evaluation', {})
        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        eval_hook = DistEvalHook if distributed else EvalHook
        # In this PR (https://github.com/open-mmlab/mmcv/pull/1193), the
        # priority of IterTimerHook has been modified from 'NORMAL' to 'LOW'.
        runner.register_hook(
            eval_hook(val_dataloader, **eval_cfg), priority='LOW')

    resume_from = None
    if cfg.resume_from is None and cfg.get('auto_resume'):
        resume_from = find_latest_checkpoint(cfg.work_dir)
    if resume_from is not None:
        cfg.resume_from = resume_from

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    runner.run(data_loaders, cfg.workflow)

sam/apis/train.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .piplines import *
from .builder import build_dataloader
sam/datasets/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# modified by Xiaoyu Bai
import copy
import platform
import random
import warnings
from functools import partial

import numpy as np
# from mmcv.parallel import collate
from .collect import collate  # use modified collate fn
from mmcv.runner import get_dist_info
from mmcv.utils import TORCH_VERSION, Registry, build_from_cfg, digit_version
from torch.utils.data import DataLoader

from mmdet.datasets.samplers import (DistributedGroupSampler, DistributedSampler,
                                     GroupSampler, InfiniteBatchSampler,
                                     InfiniteGroupBatchSampler)
from .samplers.infinite_balanced_sampler import InfiniteBalancedBatchSampler

if platform.system() != 'Windows':
    # https://github.com/pytorch/pytorch/issues/973
    import resource

    rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
    base_soft_limit = rlimit[0]
    hard_limit = rlimit[1]
    soft_limit = min(max(4096, base_soft_limit), hard_limit)
    resource.setrlimit(resource.RLIMIT_NOFILE, (soft_limit, hard_limit))

DATASETS = Registry('dataset')
PIPELINES = Registry('pipeline')


def _concat_dataset(cfg, default_args=None):
    from mmdet.datasets.dataset_wrappers import ConcatDataset
    ann_files = cfg['ann_file']
    img_prefixes = cfg.get('img_prefix', None)
    seg_prefixes = cfg.get('seg_prefix', None)
    proposal_files = cfg.get('proposal_file', None)
    separate_eval = cfg.get('separate_eval', True)

    datasets = []
    num_dset = len(ann_files)
    for i in range(num_dset):
        data_cfg = copy.deepcopy(cfg)
        # pop 'separate_eval' since it is not a valid key for common datasets.
        if 'separate_eval' in data_cfg:
            data_cfg.pop('separate_eval')
        data_cfg['ann_file'] = ann_files[i]
        if isinstance(img_prefixes, (list, tuple)):
            data_cfg['img_prefix'] = img_prefixes[i]
        if isinstance(seg_prefixes, (list, tuple)):
            data_cfg['seg_prefix'] = seg_prefixes[i]
        if isinstance(proposal_files, (list, tuple)):
            data_cfg['proposal_file'] = proposal_files[i]
        datasets.append(build_dataset(data_cfg, default_args))

    return ConcatDataset(datasets, separate_eval)


def build_dataset(cfg, default_args=None):
    from mmdet.datasets.dataset_wrappers import (ConcatDataset, RepeatDataset,
                                                 ClassBalancedDataset, MultiImageMixDataset)
    if isinstance(cfg, (list, tuple)):
        dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg])
    elif cfg['type'] == 'ConcatDataset':
        dataset = ConcatDataset(
            [build_dataset(c, default_args) for c in cfg['datasets']],
            cfg.get('separate_eval', True))
    elif cfg['type'] == 'RepeatDataset':
        dataset = RepeatDataset(
            build_dataset(cfg['dataset'], default_args), cfg['times'])
    elif cfg['type'] == 'ClassBalancedDataset':
        dataset = ClassBalancedDataset(
            build_dataset(cfg['dataset'], default_args), cfg['oversample_thr'])
    elif cfg['type'] == 'MultiImageMixDataset':
        cp_cfg = copy.deepcopy(cfg)
        cp_cfg['dataset'] = build_dataset(cp_cfg['dataset'])
        cp_cfg.pop('type')
        dataset = MultiImageMixDataset(**cp_cfg)
    elif isinstance(cfg.get('ann_file'), (list, tuple)):
        dataset = _concat_dataset(cfg, default_args)
    else:
        dataset = build_from_cfg(cfg, DATASETS, default_args)

    return dataset


def build_dataloader(dataset,
                     samples_per_gpu,
                     workers_per_gpu,
                     num_gpus=1,
                     dist=True,
                     shuffle=True,
                     seed=None,
                     runner_type='EpochBasedRunner',
                     persistent_workers=False,
                     **kwargs):
    """Build PyTorch DataLoader.

    In distributed training, each GPU/process has a dataloader.
    In non-distributed training, there is only one dataloader for all GPUs.

    Args:
        dataset (Dataset): A PyTorch dataset.
        samples_per_gpu (int): Number of training samples on each GPU, i.e.,
            batch size of each GPU.
        workers_per_gpu (int): How many subprocesses to use for data loading
            for each GPU.
        num_gpus (int): Number of GPUs. Only used in non-distributed training.
        dist (bool): Distributed training/test or not. Default: True.
        shuffle (bool): Whether to shuffle the data at every epoch.
            Default: True.
        seed (int, Optional): Seed to be used. Default: None.
        runner_type (str): Type of runner. Default: `EpochBasedRunner`
        persistent_workers (bool): If True, the data loader will not shutdown
            the worker processes after a dataset has been consumed once.
            This allows to maintain the workers `Dataset` instances alive.
            This argument is only valid when PyTorch>=1.7.0. Default: False.
        kwargs: any keyword argument to be used to initialize DataLoader

    Returns:
        DataLoader: A PyTorch dataloader.
    """
    rank, world_size = get_dist_info()

    if dist:
        # When model is :obj:`DistributedDataParallel`,
        # `batch_size` of :obj:`dataloader` is the
        # number of training samples on each GPU.
        batch_size = samples_per_gpu
        num_workers = workers_per_gpu
    else:
        # When model is obj:`DataParallel`
        # the batch size is samples on all the GPUS
        batch_size = num_gpus * samples_per_gpu
        num_workers = num_gpus * workers_per_gpu

    if runner_type == 'IterBasedRunner':
        # this is a batch samplers, which can yield
        # a mini-batch indices each time.
        # it can be used in both `DataParallel` and
        # `DistributedDataParallel`
        batch_sampler = InfiniteBalancedBatchSampler(  # for SAMV2
            dataset, batch_size, world_size, rank, seed=seed, ratio=[5, 1])
        # if shuffle:
        #     batch_sampler = InfiniteGroupBatchSampler(
        #         dataset, batch_size, world_size, rank, seed=seed)
        # else:
        #     batch_sampler = InfiniteBatchSampler(
        #         dataset,
        #         batch_size,
        #         world_size,
        #         rank,
        #         seed=seed,
        #         shuffle=False)
        batch_size = 1
        sampler = None
    else:
        if dist:
            # DistributedGroupSampler will definitely shuffle the data to
            # satisfy that images on each GPU are in the same group
            if shuffle:
                sampler = DistributedGroupSampler(
                    dataset, samples_per_gpu, world_size, rank, seed=seed)
            else:
                sampler = DistributedSampler(
                    dataset, world_size, rank, shuffle=False, seed=seed)
        else:
            sampler = GroupSampler(dataset,
                                   samples_per_gpu) if shuffle else None
        batch_sampler = None

    init_fn = partial(
        worker_init_fn, num_workers=num_workers, rank=rank,
        seed=seed) if seed is not None else None

    if (TORCH_VERSION != 'parrots'
            and digit_version(TORCH_VERSION) >= digit_version('1.7.0')):
        kwargs['persistent_workers'] = persistent_workers
    elif persistent_workers is True:
        warnings.warn('persistent_workers is invalid because your pytorch '
                      'version is lower than 1.7.0')

    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=num_workers,
        batch_sampler=batch_sampler,
        collate_fn=partial(collate, samples_per_gpu=samples_per_gpu, flatten=True),
        pin_memory=False,
        worker_init_fn=init_fn,
        **kwargs)

    return data_loader


def worker_init_fn(worker_id, num_workers, rank, seed):
    # The seed of each worker equals to
    # num_worker * rank + worker_id + user_seed
    worker_seed = num_workers * rank + worker_id + seed
    np.random.seed(worker_seed)
    random.seed(worker_seed)

sam/datasets/builder.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
from collections.abc import Mapping, Sequence

import torch
import torch.nn.functional as F
from torch.utils.data.dataloader import default_collate

from mmcv.parallel.data_container import DataContainer


def collate(batch, samples_per_gpu=1, flatten=False):
    """Puts each data field into a tensor/DataContainer with outer dimension
    batch size.

    Extend default_collate to add support for
    :type:`~mmcv.parallel.DataContainer`. There are 3 cases.

    1. cpu_only = True, e.g., meta data
    2. cpu_only = False, stack = True, e.g., images tensors
    3. cpu_only = False, stack = False, e.g., gt bboxes
    """
    if not isinstance(batch, Sequence):
        raise TypeError(f"{batch.dtype} is not supported.")

    if isinstance(batch[0], DataContainer):
        stacked = []
        if batch[0].cpu_only:
            for i in range(0, len(batch), samples_per_gpu):
                stacked.append(
                    [sample.data for sample in batch[i : i + samples_per_gpu]]
                )
            return DataContainer(
                stacked, batch[0].stack, batch[0].padding_value, cpu_only=True
            )
        elif batch[0].stack:
            for i in range(0, len(batch), samples_per_gpu):
                assert isinstance(batch[i].data, torch.Tensor)

                if batch[i].pad_dims is not None:
                    ndim = batch[i].dim()
                    assert ndim > batch[i].pad_dims
                    max_shape = [0 for _ in range(batch[i].pad_dims)]
                    for dim in range(1, batch[i].pad_dims + 1):
                        max_shape[dim - 1] = batch[i].size(-dim)
                    for sample in batch[i : i + samples_per_gpu]:
                        for dim in range(0, ndim - batch[i].pad_dims):
                            assert batch[i].size(dim) == sample.size(dim)
                        for dim in range(1, batch[i].pad_dims + 1):
                            max_shape[dim - 1] = max(
                                max_shape[dim - 1], sample.size(-dim)
                            )
                    padded_samples = []
                    for sample in batch[i : i + samples_per_gpu]:
                        pad = [0 for _ in range(batch[i].pad_dims * 2)]
                        for dim in range(1, batch[i].pad_dims + 1):
                            pad[2 * dim - 1] = max_shape[dim - 1] - sample.size(-dim)
                        padded_samples.append(
                            F.pad(sample.data, pad, value=sample.padding_value)
                        )
                    stacked.append(default_collate(padded_samples))
                elif batch[i].pad_dims is None:
                    stacked.append(
                        default_collate(
                            [sample.data for sample in batch[i : i + samples_per_gpu]]
                        )
                    )
                else:
                    raise ValueError("pad_dims should be either None or integers (1-3)")

        else:
            for i in range(0, len(batch), samples_per_gpu):
                stacked.append(
                    [sample.data for sample in batch[i : i + samples_per_gpu]]
                )
        return DataContainer(stacked, batch[0].stack, batch[0].padding_value)
    elif any([isinstance(b, Sequence) for b in batch]):
        if flatten:
            flattened = []
            for b in batch:
                if isinstance(b, Sequence):
                    flattened.extend(b)
                else:
                    flattened.extend([b])
            return collate(flattened, len(flattened))
        else:
            transposed = zip(*batch)
            return [collate(samples, samples_per_gpu) for samples in transposed]
    elif isinstance(batch[0], Mapping):
        return {
            key: collate([d[key] for d in batch], samples_per_gpu) for key in batch[0]
        }
    else:
        return default_collate(batch)
sam/datasets/collect.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai

import numpy as np
import logging
import csv
from mmdet.datasets.builder import DATASETS
from mmdet.datasets.pipelines import Compose
import pickle


@DATASETS.register_module()
class Dataset3dsam(object):
    CLASSES = None

    def __init__(
            self, data_dir, index_file, annot=None, pipeline=None, test_mode=False, multisets=False, set_length=1000,
            mask_dir=None):

        self.data_path = data_dir
        self.with_mask = False
        if not mask_dir is None:
            self.with_mask = True
            self.mask_path = mask_dir
        self.loaddatalist(index_file)
        self.classes = ['lesion']
        self.num_classes = len(self.classes)
        self.num_niis = len(self.filename)
        self.multisets = multisets
        self.sample = False
        if self.multisets:
            self.set_length = set_length
        if not annot == None:
            self.loadannot(annot)
        self.pipline = pipeline
        self.test_mode = test_mode
        self.pipline = Compose(pipeline)
        self._set_group_flag()
        self.logger = logging.getLogger(__name__)

    def __len__(self):
        if self.multisets:
            return self.set_length
        else:
            return self.num_niis

    def _set_group_flag(self):
        """Set flag according to image aspect ratio.
        Images with aspect ratio greater than 1 will be set as group 1,
        otherwise group 0.
        """
        self.flag = np.zeros(len(self), dtype=np.uint8)

    def pre_pipeline(self, results):
        return results

    def __getitem__(self, index):
        if self.multisets:
            loc = np.random.randint(0, self.num_niis)
        else:
            loc = index
        image_fn = self.filename[loc]

        if self.test_mode:
            data = self.prepare_test_img(self.data_path, image_fn)
            return data
        while True:
            if self.sample:
                anno = self.annot[image_fn.split('.', 1)[0]]
                data = self.prepare_train_img(self.data_path, image_fn, anno)
            if self.with_mask:
                data = self.prepare_train_img(self.data_path, image_fn, mask_path = self.mask_path)
            else:
                data = self.prepare_train_img(self.data_path, image_fn)
            return data

    def prepare_train_img(self, data_path, image_fn, anno=None,mask_path=None):
        if anno is not None:
            data = {'data_path': data_path,
                    'image_fn': image_fn,
                    'anno': anno}
        if mask_path is not None:
            data = {'data_path': data_path,
                    'mask_path': mask_path,
                    'image_fn': image_fn}
        else:
            data = {'data_path': data_path,
                    'image_fn': image_fn}
        data = self.pipline(data)
        return data

    def prepare_test_img(self, data_path, image_fn):
        data = {'data_path': data_path,
                'image_fn': image_fn}
        data = self.pipline(data)
        return data

    def loaddatalist(self, path):
        info = []
        with open(path) as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                info.append(row)
        info = info[1:]
        self.filename = np.array([row[0] for row in info])

    def loadannot(self, annot):
        f = open(annot, 'rb')
        self.annot = pickle.load(f)
        self.sample = True

sam/datasets/dataset3dsam.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai

import numpy as np
import logging
import csv
from mmdet.datasets.builder import DATASETS
from mmdet.datasets.pipelines import Compose
import pickle


@DATASETS.register_module()
class Dataset3dsamCrossVolume(object):
    CLASSES = None

    def __init__(
            self, data_dir, index_file, pipeline=None, test_mode=False, multisets=False, set_length=1000,
            mask_dir=None):

        self.data_path = data_dir
        self.with_mask = False
        if not mask_dir is None:
            self.with_mask = True
            self.mask_path = mask_dir
        self.loaddatalist(index_file)
        self.classes = ['lesion']
        self.num_classes = len(self.classes)
        self.num_niis = len(self.filename)
        self.multisets = multisets
        self.sample = False
        if self.multisets:
            self.set_length = set_length
        self.pipline = pipeline
        self.test_mode = test_mode
        self.pipline = Compose(pipeline)
        self._set_group_flag()
        self.logger = logging.getLogger(__name__)

    def __len__(self):
        if self.multisets:
            return self.set_length
        else:
            return self.num_niis

    def _set_group_flag(self):
        """Set flag according to image aspect ratio.
        Images with aspect ratio greater than 1 will be set as group 1,
        otherwise group 0.
        """
        self.flag = np.zeros(len(self), dtype=np.uint8)

    def pre_pipeline(self, results):
        return results

    def __getitem__(self, index):
        # if self.multisets:
        #     loc = np.random.randint(0, self.num_niis)
        # else:
        loc1 = np.random.randint(0, self.num_niis)
        loc2 = np.random.randint(0, self.num_niis)
        image_fn1 = self.filename[loc1]
        image_fn2 = self.filename[loc2]


        if self.test_mode:
            data = self.prepare_test_img(self.data_path, image_fn1)
            return data
        while True:
            if self.with_mask:
                data = self.prepare_train_img(self.data_path, image_fn1,image_fn2, mask_path = self.mask_path)
            else:
                data = self.prepare_train_img(self.data_path, image_fn1,image_fn2)
            return data

    def prepare_train_img(self, data_path, image_fn1,image_fn2, mask_path=None):
        if mask_path is not None:
            data = {'data_path': data_path,
                    'mask_path': mask_path,
                    'image_fn1': image_fn1,
                    'image_fn2': image_fn2}
        else:
            data = {'data_path': data_path,
                    'image_fn1': image_fn1,
                    'image_fn2': image_fn2}
        data = self.pipline(data)
        return data

    def prepare_test_img(self, data_path, image_fn):
        data = {'data_path': data_path,
                'image_fn': image_fn}
        data = self.pipline(data)
        return data

    def loaddatalist(self, path):
        info = []
        with open(path) as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                info.append(row)
        info = info[1:]
        self.filename = np.array([row[0] for row in info])


sam/datasets/dataset3dsam_cross_volume.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .formating3d import  *
from .process3d import *
from .process3d_mask import *
from .process3d_cross_mod import *
sam/datasets/piplines/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# modified by xiaoyu Bai
from collections.abc import Sequence

import mmcv
import numpy as np
import torch
from mmcv.parallel import DataContainer as DC
from mmdet.datasets.builder import PIPELINES


def to_tensor(data):
    """Convert objects of various python types to :obj:`torch.Tensor`.

    Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
    :class:`Sequence`, :class:`int` and :class:`float`.

    Args:
        data (torch.Tensor | numpy.ndarray | Sequence | int | float): Data to
            be converted.
    """

    if isinstance(data, torch.Tensor):
        return data
    elif isinstance(data, np.ndarray):
        return torch.from_numpy(data)
    elif isinstance(data, Sequence) and not mmcv.is_str(data):
        return torch.tensor(data)
    elif isinstance(data, int):
        return torch.LongTensor([data])
    elif isinstance(data, float):
        return torch.FloatTensor([data])
    else:
        raise TypeError(f'type {type(data)} cannot be converted to tensor.')


@PIPELINES.register_module()
class DefaultFormatBundle3d:
    """Default formatting bundle.

    It simplifies the pipeline of formatting common fields, including "img",
    "proposals", "gt_bboxes", "gt_labels", "gt_masks" and "gt_semantic_seg".
    These fields are formatted as follows.

    - img: (1)transpose, (2)to tensor, (3)to DataContainer (stack=True)
    - proposals: (1)to tensor, (2)to DataContainer
    - gt_bboxes: (1)to tensor, (2)to DataContainer
    - gt_bboxes_ignore: (1)to tensor, (2)to DataContainer
    - gt_labels: (1)to tensor, (2)to DataContainer
    - gt_masks: (1)to tensor, (2)to DataContainer (cpu_only=True)
    - gt_semantic_seg: (1)unsqueeze dim-0 (2)to tensor, \
                       (3)to DataContainer (stack=True)

    Args:
        img_to_float (bool): Whether to force the image to be converted to
            float type. Default: True.
        pad_val (dict): A dict for padding value in batch collating,
            the default value is `dict(img=0, masks=0, seg=255)`.
            Without this argument, the padding value of "gt_semantic_seg"
            will be set to 0 by default, which should be 255.
    """

    def __init__(self,
                 img_to_float=True,
                 pad_val=dict(img=0, masks=0, seg=255),switch='v1'):
        self.img_to_float = img_to_float
        self.pad_val = pad_val
        assert switch in ['v1','v2']
        self.switch = switch

    def __call__(self, results):
        """Call function to transform and format common fields in results.

        Args:
            results (dict): Result dict contains the data to convert.

        Returns:
            dict: The result dict contains the data that is formatted with \
                default bundle.
        """
        if self.switch == 'v1':
            if 'img' in results:
                img = results['img']
                results['img'] = DC(
                    to_tensor(img), padding_value=self.pad_val['img'], stack=True, pad_dims=3)
            if 'meshgrid' in results:
                meshgrid = results['meshgrid']
                results['meshgrid'] = DC(
                    to_tensor(meshgrid), padding_value=self.pad_val['img'], stack=True,pad_dims=3)
            if 'valid' in results:
                valid = results['valid']
                results['valid'] = DC(
                    to_tensor(valid), padding_value=self.pad_val['img'], stack=True, pad_dims=3)
            if 'mask' in results:
                mask = results['mask']
                results['mask'] = DC(
                    to_tensor(mask), padding_value=self.pad_val['masks'], stack=False, pad_dims=3)
        elif self.switch == 'v2':
            if 'img' in results:
                img = results['img']
                results['img'] = DC(
                    to_tensor(img), padding_value=self.pad_val['img'], stack=False, pad_dims=3)
            if 'meshgrid' in results:
                meshgrid = results['meshgrid']
                results['meshgrid'] = to_tensor(meshgrid)
            if 'valid' in results:
                valid = results['valid']
                results['valid'] = to_tensor(valid)
            if 'mask' in results:
                mask = results['mask']
                results['mask'] = to_tensor(mask)
        return results


@PIPELINES.register_module()
class IsMRI:
    def __init__(self,is_mri=False):
        self.is_mri = is_mri
    def __call__(self, results):
        results['is_mri'] = self.is_mri
        return results

@PIPELINES.register_module()
class Collect3d:
    """Collect data from the loader relevant to the specific task.

    This is usually the last stage of the data loader pipeline. Typically keys
    is set to some subset of "img", "proposals", "gt_bboxes",
    "gt_bboxes_ignore", "gt_labels", and/or "gt_masks".

    The "img_meta" item is always populated.  The contents of the "img_meta"
    dictionary depends on "meta_keys". By default this includes:

        - "img_shape": shape of the image input to the network as a tuple \
            (h, w, c).  Note that images may be zero padded on the \
            bottom/right if the batch tensor is larger than this shape.

        - "scale_factor": a float indicating the preprocessing scale

        - "flip": a boolean indicating if image flip transform was used

        - "filename": path to the image file

        - "ori_shape": original shape of the image as a tuple (h, w, c)

        - "pad_shape": image shape after padding

        - "img_norm_cfg": a dict of normalization information:

            - mean - per channel mean subtraction
            - std - per channel std divisor
            - to_rgb - bool indicating if bgr was converted to rgb

    Args:
        keys (Sequence[str]): Keys of results to be collected in ``data``.
        meta_keys (Sequence[str], optional): Meta keys to be converted to
            ``mmcv.DataContainer`` and collected in ``data[img_metas]``.
            Default: ``('filename', 'ori_filename', 'ori_shape', 'img_shape',
            'pad_shape', 'scale_factor', 'flip', 'flip_direction',
            'img_norm_cfg')``
    """

    def __init__(self,
                 keys,
                 meta_keys=('filename',)
                 ):
        self.keys = keys
        self.meta_keys = meta_keys

    def __call__(self, results):
        """Call function to collect keys in results. The keys in ``meta_keys``
        will be converted to :obj:mmcv.DataContainer.

        Args:
            results (dict): Result dict contains the data to collect.

        Returns:
            dict: The result dict contains the following keys

                - keys in``self.keys``
                - ``img_metas``
        """

        data = {}
        img_meta = {}
        for key in self.meta_keys:
            img_meta[key] = results[key]
        data['img_metas'] = DC([img_meta], cpu_only=True)
        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})'

sam/datasets/piplines/formating3d.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Xiaoyu Bai, Medical AI Lab, Alibaba DAMO Academy

from mmdet.datasets.builder import PIPELINES
import torchio as tio
import numpy as np
from mmdet.datasets.pipelines import Compose
import copy
import torch
import os
import nibabel as nib
from skimage import measure
from scipy import ndimage
from skimage import morphology
from skimage import segmentation


class Crop_mod(tio.Crop):
    def apply_transform(self, sample):
        low = self.bounds_parameters[::2]
        high = self.bounds_parameters[1::2]
        index_ini = low
        index_fin = high
        # index_fin = np.array(sample.spatial_shape) - high
        for image in self.get_images(sample):
            new_origin = nib.affines.apply_affine(image.affine, index_ini)
            new_affine = image.affine.copy()
            new_affine[:3, 3] = new_origin
            i0, j0, k0 = index_ini
            i1, j1, k1 = index_fin
            image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone())
            image.affine = new_affine
        return sample


def get_range(mask, margin=0):
    """Get up, down, left, right extreme coordinates of a binary mask"""
    idx = np.nonzero(mask)
    u = max(0, idx[0].min() - margin)
    d = min(mask.shape[0] - 1, idx[0].max() + margin)
    l = max(0, idx[1].min() - margin)
    r = min(mask.shape[1] - 1, idx[1].max() + margin)
    return [u, d, l, r]


def meshgrid3d(shape, spacing):
    y_ = np.linspace(0., (shape[1] - 1) * spacing[0], shape[1])
    x_ = np.linspace(0., (shape[2] - 1) * spacing[1], shape[2])
    z_ = np.linspace(0., (shape[3] - 1) * spacing[2], shape[3])

    x, y, z = np.meshgrid(x_, y_, z_, indexing='xy')
    return [y, x, z]


def makelabel(shape):
    all_ = np.ones((1, shape[1] * shape[2] * shape[3]))
    all = all_.reshape(-1, shape[1], shape[2], shape[3])
    all[0, int(shape[1] / 2), int(shape[2] / 2), int(shape[3] / 2)] = 0
    return all


@PIPELINES.register_module()
class RemoveCTBed():
    def __init__(self, thres=-900):
        self.thres = thres

    def __call__(self, data):
        subject = data['subject']
        img = subject.image.data
        # img = img[:,:,:,:-1] # tio resample may generate allzero slice at the end
        mask = img > self.thres
        mask_np = mask[0, :, :, :].data.numpy()
        mask_processed = np.copy(mask_np)
        for i in range(mask_np.shape[2]):
            mask_processed[:, :, i] = morphology.closing(mask_np[:, :, i], morphology.disk(3))
            mask_processed[:, :, i] = morphology.opening(mask_processed[:, :, i], morphology.disk(3))
            mask_processed[:, :, i] = morphology.remove_small_holes(
                morphology.remove_small_objects(mask_processed[:, :, i], min_size=2000), area_threshold=40000)

        # mask_np = ndimage.binary_opening(mask_np, structure=np.ones((3, 3, 3)))
        # # mask_np = ndimage.binary_erosion(mask_np, structure=np.ones((5, 5, 5)))
        # mask_np = ndimage.binary_fill_holes(mask_np)
        # l, n = measure.label(mask_np, return_num=True)
        # regions = measure.regionprops(l)
        # max_area = 0
        # max_l = 0
        # for region in regions:
        #     if region.area > max_area:
        #         max_area = region.area
        #         max_l = region.label
        # ind = 1 - (l == max_l)
        # mask_np[ind == 1] = 0

        mask_np = mask_processed[np.newaxis, :, :, :]

        data['body_mask'] = mask_np

        # subject_new = tio.Subject(
        #     image=subject.image,
        #     label=tio.LabelMap(tensor=mask_np, affine=subject.image.affine)
        # )
        subject['body_mask'] = tio.LabelMap(tensor=mask_np, affine=subject.image.affine)
        #
        # l,n = measure.label(mask_np,return_num =True)
        # regions = measure.regionprops(l)
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RandomElasticDeformation():
    def __call__(self, data):
        subject = data['subject']
        random_elastic_transform = tio.RandomElasticDeformation(locked_borders=1)
        subject = random_elastic_transform(subject)
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class MultiBranch(object):
    def __init__(self, **transform_group):
        self.transform_group = {k: Compose(v) for k, v in transform_group.items()}

    def __call__(self, data):
        multi_results = []
        for k, v in self.transform_group.items():
            res = v(copy.deepcopy(data))
            if res is None:
                return None
            # res["img_metas"]["tag"] = k
            multi_results.append(res)
        return multi_results


@PIPELINES.register_module()
class LoadTestTioImage:
    """Load image as torchio subject.
    """

    def __init__(self, re_orient=True, landmark=False, suffix=None):
        self.re_orient = re_orient
        self.landmark = landmark
        if suffix is not None:
            self.suffix = suffix
        else:
            self.suffix = None

    def __call__(self, data):
        # time1 = time.time()
        data_path = data['data_path']
        image_fn = data['image_fn']
        if self.landmark:  # the fold structure of ChestCT landmark dataset is /fold/casefold/casenii
            if self.suffix is not None:
                img_tio = tio.ScalarImage(os.path.join(data_path, image_fn.split('.', 1)[0], image_fn + self.suffix))
                data['image_fn'] = image_fn + self.suffix
            else:
                img_tio = tio.ScalarImage(os.path.join(data_path, image_fn.split('.', 1)[0], image_fn))
        else:
            img_tio = tio.ScalarImage(data_path + image_fn)
        ToCanonical = tio.ToCanonical()
        img_tio = ToCanonical(img_tio)
        img_tio.data = torch.flip(img_tio.data, (1, 2))
        img_tio.affine = np.array(
            [-img_tio.affine[0, :], -img_tio.affine[1, :], img_tio.affine[2, :], img_tio.affine[3, :]])

        if self.re_orient:
            img_data = img_tio.data
            img_tio.data = img_data.permute(0, 2, 1, 3)
            img_tio.affine = np.array(
                [img_tio.affine[1, :], img_tio.affine[0, :], img_tio.affine[2, :], img_tio.affine[3, :]])
        img_tio_shape = img_tio.data.shape  # tio use lazy loading, so we read image shape from data.shape to force it load image here
        img_tio_spacing = img_tio.spacing
        subject = tio.Subject(
            image=img_tio
        )
        data['subject'] = subject
        # time2 = time.time()
        # print('load image info time cost:',time2-time1)
        return data


# @PIPELINES.register_module()
# class LoadTioImageWithMask:
#     """Load image as torchio subject.
#     """
#
#     def __init__(self, re_orient=True):
#         self.re_orient = re_orient
#
#     def __call__(self, data):
#         # time1 = time.time()
#         data_path = data['data_path']
#         image_fn = data['image_fn']
#         if 'mask' in data.keys():
#             mask_path = data['mask_path']
#             mask_tio = tio.LabelMap(os.path.join(mask_path, image_fn))
#
#         img_tio = tio.ScalarImage(os.path.join(data_path, image_fn))
#
#         # for NIH lymph node dataset covert orientation to PLS+
#         if self.re_orient:
#             img_data = img_tio.data
#             img_tio.data = img_data.permute(0, 2, 1, 3)
#             img_tio.affine = np.array(
#                 [img_tio.affine[1, :], img_tio.affine[0, :], img_tio.affine[2, :], img_tio.affine[3, :]])
#             if 'mask' in data.keys():
#                 mask_data = mask_tio.data
#                 mask_tio.data = mask_data.permute(0, 2, 1, 3)
#                 mask_tio.affine = np.array(
#                     [mask_tio.affine[1, :], mask_tio.affine[0, :], mask_tio.affine[2, :], mask_tio.affine[3, :]])
#         img_tio_shape = img_tio.data.shape  # tio use lazy loading, so we read image shape from data.shape to force it load image here
#         img_tio_spacing = img_tio.spacing
#         meshs = meshgrid3d(img_tio_shape, img_tio_spacing)
#         labels = makelabel(img_tio_shape)
#         if 'mask' in data.keys():
#             subject = tio.Subject(
#                 image=img_tio,
#                 mask=mask_tio,
#                 label_tio_y=tio.ScalarImage(
#                     tensor=meshs[0].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
#                     affine=img_tio.affine),
#                 label_tio_x=tio.ScalarImage(
#                     tensor=meshs[1].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
#                     affine=img_tio.affine),
#                 label_tio_z=tio.ScalarImage(
#                     tensor=meshs[2].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
#                     affine=img_tio.affine),
#                 labels_map=tio.LabelMap(tensor=labels, affine=img_tio.affine)
#                 # labels_map=tio.ScalarImage(tensor=labels, affine=img_tio.affine)
#             )
#         else:
#             subject = tio.Subject(
#                 image=img_tio,
#                 label_tio_y=tio.ScalarImage(
#                     tensor=meshs[0].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
#                     affine=img_tio.affine),
#                 label_tio_x=tio.ScalarImage(
#                     tensor=meshs[1].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
#                     affine=img_tio.affine),
#                 label_tio_z=tio.ScalarImage(
#                     tensor=meshs[2].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
#                     affine=img_tio.affine),
#                 labels_map=tio.LabelMap(tensor=labels, affine=img_tio.affine)
#                 # labels_map=tio.ScalarImage(tensor=labels, affine=img_tio.affine)
#             )
#         data['subject'] = subject
#         return data


@PIPELINES.register_module()
class GeneratePairedImagesInfo:
    def __call__(self, data):
        data1 = {'data_path': data['data_path'], 'image_fn': data['image_fn1']}
        data2 = {'data_path': data['data_path'], 'image_fn': data['image_fn2']}
        data['paired'] = [data['image_fn1'], data['image_fn2']]
        data['data1'] = data1
        data['data2'] = data2
        return data


@PIPELINES.register_module()
class SeperateInfoByImage:
    def __call__(self, data):
        assert 'tag' in data.keys()
        paired = data['paired']
        tag = data['tag']
        anchor = data['anchor']
        mask_path = data['mask_path']
        if data['tag'] == 'im1':
            data = data['data1']
        elif data['tag'] == 'im2':
            data = data['data2']
        else:
            return NotImplementedError
        data['paired'] = paired
        data['tag'] = tag
        data['anchor'] = anchor
        data['mask_path'] = mask_path
        return data


@PIPELINES.register_module()
class LoadTioImage:
    """Load image as torchio subject.
    """

    def __init__(self, re_orient=True, with_mesh=True):
        self.re_orient = re_orient
        self.with_mesh = with_mesh

    def __call__(self, data):
        # time1 = time.time()
        data_path = data['data_path']
        image_fn = data['image_fn']
        img_tio = tio.ScalarImage(os.path.join(data_path, image_fn))
        # for NIH lymph node dataset covert orientation to PLS+
        if self.re_orient:
            img_data = img_tio.data
            img_tio.data = img_data.permute(0, 2, 1, 3)
            img_tio.affine = np.array(
                [img_tio.affine[1, :], img_tio.affine[0, :], img_tio.affine[2, :], img_tio.affine[3, :]])
        img_tio_shape = img_tio.data.shape  # tio use lazy loading, so we read image shape from data.shape to force it load image here
        img_tio_spacing = img_tio.spacing
        if self.with_mesh:
            meshs = meshgrid3d(img_tio_shape, img_tio_spacing)
            labels = makelabel(img_tio_shape)
            subject = tio.Subject(
                image=img_tio,
                label_tio_y=tio.ScalarImage(
                    tensor=meshs[0].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                    affine=img_tio.affine),
                label_tio_x=tio.ScalarImage(
                    tensor=meshs[1].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                    affine=img_tio.affine),
                label_tio_z=tio.ScalarImage(
                    tensor=meshs[2].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                    affine=img_tio.affine),
                labels_map=tio.LabelMap(tensor=labels, affine=img_tio.affine)
                # labels_map=tio.ScalarImage(tensor=labels, affine=img_tio.affine)
            )
        else:
            subject = tio.Subject(
                image=img_tio,
            )
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class ExtraAttrs(object):
    def __init__(self, **attrs):
        self.attrs = attrs

    def __call__(self, data):
        for k, v in self.attrs.items():
            assert k not in data
            data[k] = v
        return data


@PIPELINES.register_module()
class RandomBlur3d():
    def __init__(self, std=(0, 4)):
        self.std = std

    def __call__(self, data):
        subject = data['subject']
        image = subject.image
        randomblur_transform = tio.RandomBlur(std=self.std)
        image = randomblur_transform(image)
        subject.image.data = image.data
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RandomBiasField():
    def __init__(self, coefficients=0.5, order=3):
        self.coefficients = coefficients
        self.order = order

    def __call__(self, data):
        subject = data['subject']
        image = subject.image
        randombias_transform = tio.RandomBiasField(coefficients=self.coefficients, order=self.order)
        image = randombias_transform(image)
        subject.image.data = image.data
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RandHistogramShift_block():
    """
    Apply random nonlinear transform to the image's intensity histogram.
    Args:
        num_control_points: number of control points governing the nonlinear intensity mapping.
            a smaller number of control points allows for larger intensity shifts. if two values provided, number of
            control points selecting from range (min_value, max_value).
        prob: probability of histogram shift.
    """

    def __init__(self, num_control_points=(5, 15), prob: float = 0.6) -> None:

        if isinstance(num_control_points, int):
            if num_control_points <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (num_control_points, num_control_points)
        else:
            if len(num_control_points) != 2:
                raise ValueError("num_control points should be a number or a pair of numbers")
            if min(num_control_points) <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (min(num_control_points), max(num_control_points))
        self.prob = prob

    def interp(self, x, xp, fp):
        ns = torch if isinstance(x, torch.Tensor) else np
        if isinstance(x, np.ndarray):
            # approx 2x faster than code below for ndarray
            return np.interp(x, xp, fp)

        m = (fp[1:] - fp[:-1]) / (xp[1:] - xp[:-1])
        b = fp[:-1] - (m * xp[:-1])

        indices = ns.searchsorted(xp.reshape(-1), x.reshape(-1)) - 1
        indices = ns.clip(indices, 0, len(m) - 1)

        f = (m[indices] * x.reshape(-1) + b[indices]).reshape(x.shape)
        f[x < xp[0]] = fp[0]
        f[x > xp[-1]] = fp[-1]
        return f

    def randomize(self, flag, is_mri=False) -> None:
        num_control_point = np.random.randint(self.num_control_points[0], self.num_control_points[1] + 1)
        random_s = np.random.uniform(0.1, 0.5)
        random_e = np.random.uniform(0.6, 0.8)
        if flag == 0:
            if is_mri:
                reference_control_points = np.linspace(0., 1., num_control_point)
            else:
                # reference_control_points = np.linspace(0.17, 0.5, num_control_point)
                # reference_control_points_start = np.linspace(0., 0.17, 5)
                # reference_control_points_end = np.linspace(0.5, 1., 5)
                reference_control_points = np.linspace(random_s, random_e, num_control_point)
                reference_control_points_start = np.linspace(0., random_s, 5)
                reference_control_points_end = np.linspace(random_e, 1., 5)
        else:
            if is_mri:
                reference_control_points = np.linspace(0., 1., num_control_point)
            else:
                # reference_control_points = np.linspace(0.5, 0.67, num_control_point)
                # reference_control_points_start = np.linspace(0., 0.5, 5)
                # reference_control_points_end = np.linspace(0.67, 1., 5)
                reference_control_points = np.linspace(random_s, random_e, num_control_point)
                reference_control_points_start = np.linspace(0., random_s, 5)
                reference_control_points_end = np.linspace(random_e, 1., 5)
        if is_mri:
            self.reference_control_points = reference_control_points
        else:
            self.reference_control_points = np.concatenate(
                [reference_control_points_start[:-1], reference_control_points, reference_control_points_end[1:]])
        self.floating_control_points = np.copy(self.reference_control_points)
        for i in range(1, num_control_point - 1):
            self.floating_control_points[i] = np.random.uniform(
                self.floating_control_points[i - 1], self.floating_control_points[i + 1]
            )

    def __call__(self, data, is_mri=False):
        if np.random.rand() < self.prob:
            img_t = data['subject'].image.data
            img_mask = data['subject'].mask.data
            modality = data['mod']
            if modality == 'mri':
                is_mri = True
            all_parts = img_mask.max()
            if all_parts < 1:
                return data
            img_t_flip = 1 - img_t
            img_min, img_max = img_t.min(), img_t.max()
            img_flip_min, img_flip_max = img_t_flip.min(), img_t_flip.max()
            xps = []
            yps = []
            reference_control_points_scaled_all = []
            floating_control_points_scaled_all = []
            flip_flags = []
            for i in range(all_parts):
                if np.random.rand() < 0.3:
                    flag = 1
                    use_min = img_flip_min
                    use_max = img_flip_max
                else:
                    flag = 0
                    use_min = img_min
                    use_max = img_max
                flip_flags.append(flag)
                self.randomize(flag, is_mri=is_mri)
                xp = torch.tensor(self.reference_control_points, dtype=img_t.dtype)
                yp = torch.tensor(self.floating_control_points, dtype=img_t.dtype)
                xps.append(xp)
                yps.append(yp)

                if np.random.rand() < 0.2:  # contrast strength
                    reference_control_points_scaled = xp * (use_max - use_min) + use_min
                    floating_control_points_scaled = yp * (use_max - use_min) + use_min
                else:
                    reference_control_points_scaled = xp
                    floating_control_points_scaled = yp
                reference_control_points_scaled_all.append(reference_control_points_scaled)
                floating_control_points_scaled_all.append(floating_control_points_scaled)

            img_final = torch.zeros_like(img_t)
            img_final = torch.where(img_mask > 0, img_final, img_t)
            for i in range(all_parts):
                if flip_flags[i] == 0:
                    img_ted = self.interp(img_t, reference_control_points_scaled_all[i],
                                          floating_control_points_scaled_all[i])
                else:
                    img_ted = self.interp(img_t_flip, reference_control_points_scaled_all[i],
                                          floating_control_points_scaled_all[i])
                img_final = torch.where(img_mask == i + 1, img_ted, img_final)
            # import matplotlib.pyplot as plt
            # plt.imshow(img_t[0, :, :, 90], cmap='gray', vmin=0, vmax=1)
            # plt.show()
            data['subject'].image.data = img_final
        return data


@PIPELINES.register_module()
class RandHistogramShift():
    """
    Apply random nonlinear transform to the image's intensity histogram.
    Args:
        num_control_points: number of control points governing the nonlinear intensity mapping.
            a smaller number of control points allows for larger intensity shifts. if two values provided, number of
            control points selecting from range (min_value, max_value).
        prob: probability of histogram shift.
    """

    def __init__(self, num_control_points=(5, 15), prob: float = 0.9) -> None:

        if isinstance(num_control_points, int):
            if num_control_points <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (num_control_points, num_control_points)
        else:
            if len(num_control_points) != 2:
                raise ValueError("num_control points should be a number or a pair of numbers")
            if min(num_control_points) <= 2:
                raise ValueError("num_control_points should be greater than or equal to 3")
            self.num_control_points = (min(num_control_points), max(num_control_points))
        self.prob = prob

    def interp(self, x, xp, fp):
        ns = torch if isinstance(x, torch.Tensor) else np
        if isinstance(x, np.ndarray):
            # approx 2x faster than code below for ndarray
            return np.interp(x, xp, fp)

        m = (fp[1:] - fp[:-1]) / (xp[1:] - xp[:-1])
        b = fp[:-1] - (m * xp[:-1])

        indices = ns.searchsorted(xp.reshape(-1), x.reshape(-1)) - 1
        indices = ns.clip(indices, 0, len(m) - 1)

        f = (m[indices] * x.reshape(-1) + b[indices]).reshape(x.shape)
        f[x < xp[0]] = fp[0]
        f[x > xp[-1]] = fp[-1]
        return f

    def randomize(self) -> None:
        num_control_point = np.random.randint(self.num_control_points[0], self.num_control_points[1] + 1)
        self.reference_control_points = np.linspace(0, 1, num_control_point)
        self.floating_control_points = np.copy(self.reference_control_points)
        for i in range(1, num_control_point - 1):
            self.floating_control_points[i] = np.random.uniform(
                self.floating_control_points[i - 1], self.floating_control_points[i + 1]
            )

    def __call__(self, data):
        if np.random.rand() < self.prob:
            img_t = data['subject'].image.data
            img_mask = data['subject'].mask.data
            self.randomize()
            if self.reference_control_points is None or self.floating_control_points is None:
                raise RuntimeError("please call the `randomize()` function first.")
            if np.random.rand() < 0.3:  # flip intensity
                img_t2 = 1 - img_t
            else:
                img_t2 = img_t

            img_min, img_max = img_t2.min(), img_t2.max()
            xp = torch.tensor(self.reference_control_points, dtype=img_t.dtype)
            yp = torch.tensor(self.floating_control_points, dtype=img_t.dtype)
            if np.random.rand() < 0.3:  # contrast strength
                reference_control_points_scaled = xp * (img_max - img_min) + img_min
                floating_control_points_scaled = yp * (img_max - img_min) + img_min
            else:
                reference_control_points_scaled = xp
                floating_control_points_scaled = yp
            img_ted = self.interp(img_t2, reference_control_points_scaled, floating_control_points_scaled)
            img_t = torch.where(img_mask > 0, img_ted, img_t)
            # import matplotlib.pyplot as plt
            # plt.imshow(img_t[0, :, :, 90], cmap='gray', vmin=0, vmax=1)
            # plt.show()
            data['subject'].image.data = img_t
        return data


@PIPELINES.register_module()
class RandomGamma():
    def __init__(self, log_gamma=(-1, 1)):
        self.log_gamma = log_gamma

    def __call__(self, data):
        subject = data['subject']
        image = subject.image
        randomgamma_transform = tio.RandomGamma(log_gamma=self.log_gamma)
        image = randomgamma_transform(image)
        subject.image.data = image.data
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RandomNoise3d():
    def __init__(self, std=(0, 0.001)):
        self.std = std

    def __call__(self, data):
        subject = data['subject']
        image = subject.image
        randomnoise_transform = tio.RandomNoise(mean=0, std=self.std)
        image = randomnoise_transform(image)
        subject.image.data = image.data
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class Resample():
    def __init__(self, norm_spacing=(2., 2., 2.), intra_volume=True, crop_artifacts=True, dynamic_spacing=False):
        self.norm_spacing = norm_spacing
        self.intra_volume = intra_volume
        self.crop_artifacts = crop_artifacts
        self.dynamic_spacing = dynamic_spacing

    def __call__(self, data):
        subject = data['subject']
        keys = list(data.keys())
        if 'tag' in keys and self.intra_volume:
            tag = data['tag']
            spacing_matrix = data['spacing_matrix']
            if tag == 'view1':
                spacing = spacing_matrix[0]
            else:
                spacing = spacing_matrix[1]
        elif self.dynamic_spacing:
            spacing = data['resample_spacing']
        else:
            spacing = np.array(self.norm_spacing)
        resample_transform = tio.Resample(target=spacing)
        subject = resample_transform(subject)
        if self.crop_artifacts:
            img_data = subject['image'].data
            crop_flag = 0
            crop_info = np.array(subject.shape[1:]).astype(int)
            if (img_data[0, :, :, -1] == 0).all():
                crop_flag = 1
                # crop_info[-1] = crop_info[-1] - 1
                crop_info[2] = crop_info[2] - 1
            if (img_data[0, :, -1, :] == 0).all():
                crop_flag = 1
                # crop_info[-1] = crop_info[-1] - 1
                crop_info[1] = crop_info[1] - 1
            if (img_data[0, -1, :, :] == 0).all():
                crop_flag = 1
                # crop_info[-1] = crop_info[-1] - 1
                crop_info[0] = crop_info[0] - 1
            if crop_flag == 1:
                crop_transform = Crop_mod([0, crop_info[0], 0, crop_info[1], 0, crop_info[2]])  # y1y2x1x2z1z2
                subject = crop_transform(subject)
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class GenerateMeshGrid():
    def __call__(self, data):
        sub = data['subject']
        image = sub.image.data
        y_sub = sub.label_tio_y.data
        x_sub = sub.label_tio_x.data
        z_sub = sub.label_tio_z.data
        pos_mm = torch.cat((y_sub, x_sub, z_sub), axis=0)
        # pos1_mm = pos1_mm.transpose((1, 2, 3, 0))
        sub_valid = sub.labels_map.data
        data['img'] = image.permute(0, 3, 1, 2)
        data['meshgrid'] = pos_mm.permute(0, 3, 1, 2)
        data['valid'] = sub_valid.permute(0, 3, 1, 2)
        data['spacing'] = (sub.image.spacing[2],sub.image.spacing[0],sub.image.spacing[1])
        return data


@PIPELINES.register_module()
class GenerateMetaInfo():
    def __init__(self, is_train=True, with_mask=False):
        self.is_train = is_train
        self.with_mask = with_mask

    def __call__(self, data):
        data['filename'] = data['image_fn']
        keys = list(data.keys())
        if 'tag' in keys and 'crop_matrix_origin' in keys:
            tag = data['tag']
            if tag == 'view1':
                data['crop_info'] = data['crop_matrix_origin'][0:2, :]
            else:
                data['crop_info'] = data['crop_matrix_origin'][2:, :]
        if 'img' not in keys:
            sub = data['subject']
            image = sub.image.data
            if self.is_train:
                data['img'] = image.permute(0, 3, 1, 2)
            else:
                data['img'] = [image.permute(0, 3, 1, 2)]
        if self.with_mask:
            sub = data['subject']
            mask = sub.mask.data
            data['mask'] = mask.permute(0, 3, 1, 2)
        data['volume_size'] = data['img'][0].shape

        return data


@PIPELINES.register_module()
class ComputeCorrespond():
    def __init__(self, shape=(96, 96, 24)):
        self.shape = shape

    def __call__(self, data):
        subject = data['subject']
        resize_transform = tio.Resize(self.shape)
        subject = resize_transform(subject)
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class Resize3d():
    def __init__(self, shape=(96, 96, 24)):
        self.shape = shape

    def __call__(self, data):
        subject = data['subject']
        resize_transform = tio.Resize(self.shape)
        subject = resize_transform(subject)
        data['subject'] = subject

        return data


@PIPELINES.register_module()
class Cropz():
    def __init__(self, x_slice=100, y_slice=100, z_slice=100):
        self.x_slice = x_slice
        self.y_slice = y_slice
        self.z_slice = z_slice

    def __call__(self, data):
        subject = data['subject']
        shape = subject.image.shape
        flag = False
        crop_info = [0, shape[1],
                     0, shape[2],
                     0, shape[3]]
        if shape[3] > self.z_slice + 10:
            seedz = np.random.randint(0, shape[3] - self.z_slice - 1)
            crop_info[4] = seedz
            crop_info[5] = seedz + self.z_slice
            flag = True
        if shape[1] > self.y_slice + 10:
            seedy = np.random.randint(0, shape[1] - self.y_slice - 1)
            crop_info[0] = seedy
            crop_info[1] = seedy + self.y_slice
            flag = True
        if shape[2] > self.x_slice + 10:
            seedx = np.random.randint(0, shape[2] - self.x_slice - 1)
            crop_info[2] = seedx
            crop_info[3] = seedx + self.x_slice
            flag = True
        if flag:
            crop_info = np.array(crop_info).astype(int)
            crop_transform = Crop_mod(crop_info)  # y1y2x1x2z1z2
            subject = crop_transform(subject)
            data['subject'] = subject
        return data


@PIPELINES.register_module()
class CropBackground():
    def __init__(self, thres=-500, by_mask=False):
        self.thres = thres
        self.by_mask = by_mask

    def __call__(self, data):
        subject = data['subject']
        img = subject.image.data
        if not self.by_mask:
            # img = img[:,:,:,:-1] # tio resample may generate allzero slice at the end
            mask = img > self.thres
        else:
            mask = subject['body_mask'].data
        foreground_region = torch.where(mask == 1)
        crop_info = (foreground_region[1].min(), foreground_region[1].max(),
                     foreground_region[2].min(), foreground_region[2].max(),
                     foreground_region[3].min(),
                     foreground_region[3].max())

        crop_info = np.array(crop_info).astype(int)
        crop_transform = Crop_mod(crop_info)  # y1y2x1x2z1z2
        subject = crop_transform(subject)
        data['subject'] = subject
        data['crop_info'] = crop_info
        return data


@PIPELINES.register_module()
class RandomCrop3d():
    def __init__(self, thres=0.99, min_z=18, max_z=32):
        self.thres = thres
        self.min_z = min_z
        self.max_z = max_z

    def __call__(self, data):
        subject = data['subject']
        img_shape = subject.image.shape
        p = np.random.uniform(0, 1)
        z_use = np.random.randint(self.min_z, self.max_z)
        if p < self.thres and (img_shape[3] - z_use - 1) > 0:
            seed = np.random.randint(0, img_shape[3] - z_use - 1)
            crop_info = np.array([0, img_shape[1], 0, img_shape[2], seed, seed + z_use])
            crop_transform = Crop_mod(crop_info)  # y1y2x1x2z1z2
            subject = crop_transform(subject)
            data['subject'] = subject
        return data


@PIPELINES.register_module()
class Crop():
    def __init__(self, switch='origin', by_view=True, fix_size=(96, 96, 32)):
        self.switch = switch
        self.fix_size = fix_size
        self.by_view = by_view
        assert self.switch in ['origin', 'fix']

    def __call__(self, data):
        # time1 = time.time()
        subject = data['subject']
        if 'crop_matrix_origin' in data.keys():
            crop_matrix_origin = data['crop_matrix_origin']
        if 'tag' in data.keys():
            tag = data['tag']
            if tag == 'view1':
                view = 0
            else:
                view = 2
        if self.switch == 'origin':
            if self.by_view:
                crop_info = (crop_matrix_origin[view][0],
                             crop_matrix_origin[view + 1][0],
                             crop_matrix_origin[view][1],
                             crop_matrix_origin[view + 1][1],
                             crop_matrix_origin[view][2],
                             crop_matrix_origin[view + 1][2])
            else:
                crop_info = (crop_matrix_origin[0][0],
                             crop_matrix_origin[1][0],
                             crop_matrix_origin[0][1],
                             crop_matrix_origin[1][1],
                             crop_matrix_origin[0][2],
                             crop_matrix_origin[1][2])
        else:
            subject_shape = subject.shape[1:]
            assert subject_shape >= self.fix_size, print(data['image_fn'])

            crop_info = (int(subject_shape[0] / 2) - self.fix_size[0] / 2,
                         int(subject_shape[0] / 2) + self.fix_size[0] / 2,
                         int(subject_shape[1] / 2) - self.fix_size[1] / 2,
                         int(subject_shape[1] / 2) + self.fix_size[1] / 2,
                         int(subject_shape[2] / 2) - self.fix_size[2] / 2,
                         int(subject_shape[2] / 2) + self.fix_size[2] / 2)
        crop_info = np.array(crop_info).astype(int)
        crop_transform = Crop_mod(crop_info)  # y1y2x1x2z1z2
        subject = crop_transform(subject)
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RandomAffine3d():
    def __init__(self, scales=0, degrees=(-10, 10, -10, 10, -10, 10), translation=0):
        self.scales = scales
        self.degrees = degrees
        self.translation = translation

    def __call__(self, data):
        subject = data['subject']
        random_affine_transform = tio.RandomAffine(scales=self.scales, degrees=self.degrees,
                                                   translation=self.translation)
        subject = random_affine_transform(subject)
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class DynamicRescaleIntensity():
    def __init__(self, out_min_max=(0, 255), bias=50):
        self.out_min_max = out_min_max
        self.bias = bias

    def __call__(self, data):
        subject = data['subject']
        img = subject.image.data.type(torch.FloatTensor)
        img_numpy = img.numpy()
        # in_min_max = torch.quantile(img.type(torch.FloatTensor), torch.tensor([0.05,0.99])) # torch quantile like np.quantile not np.percentail
        in_min_max = np.percentile(img_numpy, (0.2, 99.8))
        if in_min_max[1] < 1:  # we found for some normalized MR, we should not re-normalize the intensity
            in_min_max[1] = 1
        img[img < in_min_max[0]] = in_min_max[0]
        img[img > in_min_max[1]] = in_min_max[1]
        subject.image.data = img
        rescale_transform = tio.RescaleIntensity(out_min_max=self.out_min_max,
                                                 in_min_max=(int(in_min_max[0]), int(in_min_max[1])))
        image = rescale_transform(subject.image)
        subject.image.data = image.data - self.bias
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RescaleIntensity_modality():
    def __init__(self, out_min_max=(0, 255), in_min_max=(-1024, 3071), bias=50):
        self.out_min_max = out_min_max
        self.in_min_max = in_min_max
        self.bias = bias

    def __call__(self, data):
        subject = data['subject']
        modality = data['mod']
        img = subject.image.data.type(torch.FloatTensor)
        img_numpy = img.numpy()
        if modality == 'ct':
            img[img < self.in_min_max[0]] = self.in_min_max[0]
            img[img > self.in_min_max[1]] = self.in_min_max[1]
            subject.image.data = img
            rescale_transform = tio.RescaleIntensity(out_min_max=self.out_min_max, in_min_max=self.in_min_max)

        elif modality == 'mri':
            in_min_max = np.percentile(img_numpy, (0.2, 99.8))
            img[img < in_min_max[0]] = in_min_max[0]
            img[img > in_min_max[1]] = in_min_max[1]
            subject.image.data = img
            rescale_transform = tio.RescaleIntensity(out_min_max=self.out_min_max,
                                                     in_min_max=(int(in_min_max[0]), int(in_min_max[1])))
        image = rescale_transform(subject.image)
        subject.image.data = image.data - self.bias
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class RescaleIntensity():
    def __init__(self, out_min_max=(0, 255), in_min_max=(-1024, 3071), bias=50):
        self.out_min_max = out_min_max
        self.in_min_max = in_min_max
        self.bias = bias

    def __call__(self, data):
        subject = data['subject']
        img = subject.image.data
        img[img < self.in_min_max[0]] = self.in_min_max[0]
        img[img > self.in_min_max[1]] = self.in_min_max[1]
        subject.image.data = img
        rescale_transform = tio.RescaleIntensity(out_min_max=self.out_min_max, in_min_max=self.in_min_max)
        image = rescale_transform(subject.image)
        subject.image.data = image.data - self.bias
        data['subject'] = subject
        return data


@PIPELINES.register_module()
class ComputeAugParam_sample():
    def __init__(self, scale=(0.8, 1.2), standard_spacing=(2., 2., 2.), patch_size=(96, 96, 32)):
        self.scale = scale
        self.standard_spacing = standard_spacing
        self.patch_size = patch_size
        # self.sample = sample

    def __call__(self, data):
        if 'anno' in data.keys():
            sample = True
        else:
            sample = False
        if 'mask' in data['subject'].keys():
            sample_by_mask = True
        else:
            sample_by_mask = False
        view1_scale = np.random.uniform(self.scale[0], self.scale[1])
        view2_scale = np.random.uniform(self.scale[0], self.scale[1])
        # view2_scale = view1_scale
        origin_spacing = data['subject'].image.spacing
        origin_size = data['subject'].image.shape
        # print(origin_size)
        standard_size = [np.floor(origin_spacing[0] / self.standard_spacing[0] * origin_size[1]).astype(int),
                         np.floor(origin_spacing[1] / self.standard_spacing[1] * origin_size[2]).astype(int),
                         np.floor(origin_spacing[2] / self.standard_spacing[2] * origin_size[3]).astype(int)]
        view1_patch_size = np.ceil(np.array(self.patch_size) * view1_scale).astype(int)  # y,x,z
        view2_patch_size = np.ceil(np.array(self.patch_size) * view2_scale).astype(int)
        affine_standard_to_origin = np.array([self.standard_spacing[0] / origin_spacing[0], 0, 0, 0,
                                              0, self.standard_spacing[1] / origin_spacing[1], 0, 0,
                                              0, 0, self.standard_spacing[2] / origin_spacing[2], 0,
                                              0, 0, 0, 1]).reshape(4, 4).transpose()
        view1_crop_size = np.ceil(view1_patch_size * 1.05).astype(int)
        view2_crop_size = np.ceil(view2_patch_size * 1.05).astype(int)

        margin_x = 1
        margin_y = 1
        margin_z = 1
        rangex_min = margin_x
        rangex_max = standard_size[1] - margin_x
        rangey_min = margin_y
        rangey_max = standard_size[0] - margin_y
        rangez_min = margin_z
        rangez_max = standard_size[2] - margin_z

        p_intersect = np.random.rand(1)
        if p_intersect <= 0.2:
            view1_left_up = (np.random.randint(rangey_min, rangey_max - view1_crop_size[0]),
                             np.random.randint(rangex_min, rangex_max - view1_crop_size[1]),
                             np.random.randint(rangez_min, rangez_max - view1_crop_size[2]))
            view1_left_up = np.asarray(view1_left_up)
            view2_left_up = (np.random.randint(rangey_min, rangey_max - view2_crop_size[0]),
                             np.random.randint(rangex_min, rangex_max - view2_crop_size[1]),
                             np.random.randint(rangez_min, rangez_max - view2_crop_size[2]))
            view2_left_up = np.asarray(view2_left_up)
        else:
            if sample == True:
                crop_info = data['crop_info']
                p_sample = np.random.rand(1)
                anno = data['anno']
                anno_cropped = {}
                for key in anno.keys():
                    anno_cropped[int((key - crop_info[4]) * origin_spacing[0] / self.standard_spacing[0])] = []
                    for bbx in anno[key]:
                        bbx_tmp = np.array(bbx)
                        bbx_tmp = bbx_tmp - [crop_info[2], crop_info[0]] * 2
                        bbx_tmp = bbx_tmp * ([origin_spacing[1] / self.standard_spacing[1],
                                              origin_spacing[0] / self.standard_spacing[0]] * 2)
                        anno_cropped[int((key - crop_info[4]) * origin_spacing[0] / self.standard_spacing[0])].append(
                            bbx_tmp.tolist())

                zs = list(anno_cropped.keys())
                zs = [z for z in zs if z > 1 and z < standard_size[2] - 1]
                seed_ind = np.random.randint(0, zs.__len__())
                boxes_len = anno_cropped[zs[seed_ind]].__len__()
                box_ind = np.random.randint(0, boxes_len)
            else:
                p_sample = 0.1
            if p_sample < 0.5:
                seed_margin_xy = 10
                seed_margin_z = 4
                seed_point = (np.random.randint(rangey_min + seed_margin_xy, rangey_max - seed_margin_xy),
                              np.random.randint(rangex_min + seed_margin_xy, rangex_max - seed_margin_xy),
                              np.random.randint(rangez_min + seed_margin_z, rangez_max - seed_margin_z))
                seed_point = np.array(seed_point)
            else:
                bbx = anno_cropped[zs[seed_ind]][box_ind]
                seed_point = ((bbx[1] + bbx[3]) / 2, (bbx[0] + bbx[2]) / 2, zs[seed_ind])
                seed_point = np.array(seed_point, dtype=int)

            assert max(seed_point[2] - view1_crop_size[2], rangez_min) < min(rangez_max - view1_crop_size[2],
                                                                             seed_point[2]), data['image_fn']

            view1_left_up = (np.random.randint(max(seed_point[0] - view1_crop_size[0], rangey_min),
                                               min(rangey_max - view1_crop_size[0], seed_point[0])),
                             np.random.randint(max(seed_point[1] - view1_crop_size[1], rangex_min),
                                               min(rangex_max - view1_crop_size[1], seed_point[1])),
                             np.random.randint(max(seed_point[2] - view1_crop_size[2], rangez_min),
                                               min(rangez_max - view1_crop_size[2], seed_point[2])))
            view1_left_up = np.asarray(view1_left_up)

            view2_left_up = (np.random.randint(max(seed_point[0] - view2_crop_size[0], rangey_min),
                                               min(rangey_max - view2_crop_size[0], seed_point[0])),
                             np.random.randint(max(seed_point[1] - view2_crop_size[1], rangex_min),
                                               min(rangex_max - view2_crop_size[1], seed_point[1])),
                             np.random.randint(max(seed_point[2] - view2_crop_size[2], rangez_min),
                                               min(rangez_max - view2_crop_size[2], seed_point[2])))
            view2_left_up = np.asarray(view2_left_up)

        view1_y1x1z1_crop = view1_left_up
        view1_y2x2z2_crop = view1_left_up + view1_crop_size
        view2_y1x1z1_crop = view2_left_up
        view2_y2x2z2_crop = view2_left_up + view2_crop_size

        crop_matrix = np.array([view1_y1x1z1_crop, view1_y2x2z2_crop, view2_y1x1z1_crop, view2_y2x2z2_crop])

        crop_matrix_origin = np.dot(crop_matrix, affine_standard_to_origin[:3, :3])
        crop_matrix_origin[::2, :] = np.floor(crop_matrix_origin[::2, :])
        crop_matrix_origin[1::2, :] = np.ceil(crop_matrix_origin[1::2, :])
        crop_matrix_origin = crop_matrix_origin.astype(int)
        view1_spacing = np.array(self.standard_spacing) * view1_scale
        view2_spacing = np.array(self.standard_spacing) * view2_scale
        spacing_matrix = np.array([view1_spacing, view2_spacing])

        data['crop_matrix_origin'] = crop_matrix_origin
        data['spacing_matrix'] = spacing_matrix
        return data


@PIPELINES.register_module()
class ComputeAugParam_masksample():
    def __init__(self, scale=(0.8, 1.2), standard_spacing=(2., 2., 2.), patch_size=(96, 96, 32)):
        self.scale = scale
        self.standard_spacing = standard_spacing
        self.patch_size = patch_size
        # self.sample = sample

    def __call__(self, data):
        if 'mask' in data['subject'].keys():
            sample_by_mask = True
        else:
            sample_by_mask = False
        view1_scale = np.random.uniform(self.scale[0], self.scale[1])
        view2_scale = np.random.uniform(self.scale[0], self.scale[1])
        # view2_scale = view1_scale
        origin_spacing = data['subject'].image.spacing
        origin_size = data['subject'].image.shape
        # print(origin_size)
        standard_size = [np.floor(origin_spacing[0] / self.standard_spacing[0] * origin_size[1]).astype(int),
                         np.floor(origin_spacing[1] / self.standard_spacing[1] * origin_size[2]).astype(int),
                         np.floor(origin_spacing[2] / self.standard_spacing[2] * origin_size[3]).astype(int)]
        view1_patch_size = np.ceil(np.array(self.patch_size) * view1_scale).astype(int)  # y,x,z
        view2_patch_size = np.ceil(np.array(self.patch_size) * view2_scale).astype(int)
        affine_standard_to_origin = np.array([self.standard_spacing[0] / origin_spacing[0], 0, 0, 0,
                                              0, self.standard_spacing[1] / origin_spacing[1], 0, 0,
                                              0, 0, self.standard_spacing[2] / origin_spacing[2], 0,
                                              0, 0, 0, 1]).reshape(4, 4).transpose()
        view1_crop_size = np.ceil(view1_patch_size * 1.05).astype(int)
        view2_crop_size = np.ceil(view2_patch_size * 1.05).astype(int)

        margin_x = 1
        margin_y = 1
        margin_z = 1
        rangex_min = margin_x
        rangex_max = standard_size[1] - margin_x
        rangey_min = margin_y
        rangey_max = standard_size[0] - margin_y
        rangez_min = margin_z
        rangez_max = standard_size[2] - margin_z


        p_intersect = np.random.rand(1)
        if p_intersect <= 0.1:
            view1_left_up = (np.random.randint(rangey_min, rangey_max - view1_crop_size[0]),
                             np.random.randint(rangex_min, rangex_max - view1_crop_size[1]),
                             np.random.randint(rangez_min, rangez_max - view1_crop_size[2]))
            view1_left_up = np.asarray(view1_left_up)
            view2_left_up = (np.random.randint(rangey_min, rangey_max - view2_crop_size[0]),
                             np.random.randint(rangex_min, rangex_max - view2_crop_size[1]),
                             np.random.randint(rangez_min, rangez_max - view2_crop_size[2]))
            view2_left_up = np.asarray(view2_left_up)
        else:
            if sample_by_mask == True:
                p_sample = np.random.rand(1)
                mask = data['subject']['mask'].data[0, :, :, :].clone().numpy()
                # print(mask.shape)
            else:
                p_sample = 0.1
            if p_sample <= 0.1:
                seed_margin_xy = 10
                seed_margin_z = 4
                seed_point = (np.random.randint(rangey_min + seed_margin_xy, rangey_max - seed_margin_xy),
                              np.random.randint(rangex_min + seed_margin_xy, rangex_max - seed_margin_xy),
                              np.random.randint(rangez_min + seed_margin_z, rangez_max - seed_margin_z))
                seed_point = np.array(seed_point)
            else:
                mask[:, :, :5] = 0
                mask[:, :, -5::1] = 0
                mask[:5, :, :] = 0
                mask[-5::1, :, :] = 0
                mask[:, :5, :] = 0
                mask[:, -5::1, :] = 0
                ys, xs, zs = np.where(mask > 0)
                body_len = ys.shape[0]
                ind = np.random.randint(0, body_len - 1)
                seed_point = np.array([ys[ind], xs[ind], zs[ind]], dtype=int)
                seed_point = np.dot(seed_point, np.linalg.inv(affine_standard_to_origin[:3, :3])).astype(int)

            assert max(seed_point[2] - view1_crop_size[2], rangez_min) < min(rangez_max - view1_crop_size[2],
                                                                             seed_point[2]), [rangez_max, seed_point[2],
                                                                                              view1_crop_size[2],data['image_fn']]
            assert max(seed_point[1] - view1_crop_size[1], rangex_min) < min(rangex_max - view1_crop_size[1],
                                                                             seed_point[1]), [rangex_max,data['image_fn']]
            assert max(seed_point[0] - view1_crop_size[0], rangey_min) < min(rangey_max - view1_crop_size[0],
                                                                             seed_point[0]), [rangey_max,view1_crop_size,data['image_fn']]

            assert max(seed_point[2] - view2_crop_size[2], rangez_min) < min(rangez_max - view2_crop_size[2],
                                                                             seed_point[2]), [rangez_max, seed_point[2],
                                                                                              [view2_crop_size[2]],data['image_fn']]
            assert max(seed_point[1] - view2_crop_size[1], rangex_min) < min(rangex_max - view2_crop_size[1],
                                                                             seed_point[1]), [rangex_max,data['image_fn']]
            assert max(seed_point[0] - view2_crop_size[0], rangey_min) < min(rangey_max - view2_crop_size[0],
                                                                             seed_point[0]), [rangey_max,view1_crop_size,data['image_fn']]

            view1_left_up = (np.random.randint(max(seed_point[0] - view1_crop_size[0], rangey_min),
                                               min(rangey_max - view1_crop_size[0], seed_point[0])),
                             np.random.randint(max(seed_point[1] - view1_crop_size[1], rangex_min),
                                               min(rangex_max - view1_crop_size[1], seed_point[1])),
                             np.random.randint(max(seed_point[2] - view1_crop_size[2], rangez_min),
                                               min(rangez_max - view1_crop_size[2], seed_point[2])))
            view1_left_up = np.asarray(view1_left_up)

            view2_left_up = (np.random.randint(max(seed_point[0] - view2_crop_size[0], rangey_min),
                                               min(rangey_max - view2_crop_size[0], seed_point[0])),
                             np.random.randint(max(seed_point[1] - view2_crop_size[1], rangex_min),
                                               min(rangex_max - view2_crop_size[1], seed_point[1])),
                             np.random.randint(max(seed_point[2] - view2_crop_size[2], rangez_min),
                                               min(rangez_max - view2_crop_size[2], seed_point[2])))
            view2_left_up = np.asarray(view2_left_up)

        view1_y1x1z1_crop = view1_left_up
        view1_y2x2z2_crop = view1_left_up + view1_crop_size
        view2_y1x1z1_crop = view2_left_up
        view2_y2x2z2_crop = view2_left_up + view2_crop_size

        crop_matrix = np.array([view1_y1x1z1_crop, view1_y2x2z2_crop, view2_y1x1z1_crop, view2_y2x2z2_crop])

        crop_matrix_origin = np.dot(crop_matrix, affine_standard_to_origin[:3, :3])
        crop_matrix_origin[::2, :] = np.floor(crop_matrix_origin[::2, :])
        crop_matrix_origin[1::2, :] = np.ceil(crop_matrix_origin[1::2, :])
        crop_matrix_origin = crop_matrix_origin.astype(int)
        view1_spacing = np.array(self.standard_spacing) * view1_scale
        view2_spacing = np.array(self.standard_spacing) * view2_scale
        spacing_matrix = np.array([view1_spacing, view2_spacing])

        data['crop_matrix_origin'] = crop_matrix_origin
        data['spacing_matrix'] = spacing_matrix
        return data


@PIPELINES.register_module()
class ComputeAugParam():
    def __init__(self, scale=(0.8, 1.2), standard_spacing=(2., 2., 2.), patch_size=(96, 96, 32)):
        self.scale = scale
        self.standard_spacing = standard_spacing
        self.patch_size = patch_size
        # self.sample = sample

    def __call__(self, data):
        if 'anno' in data.keys():
            sample = True
        else:
            sample = False
        view1_scale = np.random.uniform(self.scale[0], self.scale[1])
        view2_scale = np.random.uniform(self.scale[0], self.scale[1])
        # view2_scale = view1_scale
        origin_spacing = data['subject'].image.spacing
        origin_size = data['subject'].image.shape
        # print(origin_size)
        standard_size = [np.floor(origin_spacing[0] / self.standard_spacing[0] * origin_size[1]).astype(int),
                         np.floor(origin_spacing[1] / self.standard_spacing[1] * origin_size[2]).astype(int),
                         np.floor(origin_spacing[2] / self.standard_spacing[2] * origin_size[3]).astype(int)]
        view1_patch_size = np.ceil(np.array(self.patch_size) * view1_scale).astype(int)  # y,x,z
        view2_patch_size = np.ceil(np.array(self.patch_size) * view2_scale).astype(int)
        affine_standard_to_origin = np.array([self.standard_spacing[0] / origin_spacing[0], 0, 0, 0,
                                              0, self.standard_spacing[1] / origin_spacing[1], 0, 0,
                                              0, 0, self.standard_spacing[2] / origin_spacing[2], 0,
                                              0, 0, 0, 1]).reshape(4, 4).transpose()
        view1_crop_size = np.ceil(view1_patch_size * 1.05).astype(int)
        view2_crop_size = np.ceil(view2_patch_size * 1.05).astype(int)

        margin_x = 1
        margin_y = 1
        margin_z = 1
        rangex_min = margin_x
        rangex_max = standard_size[1] - margin_x
        rangey_min = margin_y
        rangey_max = standard_size[0] - margin_y
        rangez_min = margin_z
        rangez_max = standard_size[2] - margin_z

        y_interval = rangey_max - rangey_min - (view1_crop_size[0] + view2_crop_size[0])
        x_interval = rangex_max - rangex_min - (view1_crop_size[1] + view2_crop_size[1])
        z_interval = rangez_max - rangez_min - (view1_crop_size[2] + view2_crop_size[2])

        if y_interval >= 2:
            p = np.random.rand(1)
            if p <= 0.8:
                intersection_y = np.random.randint(20, min(view1_crop_size[0], view2_crop_size[0]) - 2)
            else:
                intersection_y = np.random.randint(-y_interval + 1, -y_interval + 10)
        else:
            intersection_y = np.random.randint(-y_interval + 1, min(view1_crop_size[0], view2_crop_size[0]) - 1)

        if x_interval >= 2:
            p = np.random.rand(1)
            if p <= 0.8:
                intersection_x = np.random.randint(20, min(view1_crop_size[1], view2_crop_size[1]) - 2)
            else:
                intersection_x = np.random.randint(-x_interval + 1, -x_interval + 10)
        else:
            intersection_x = np.random.randint(-x_interval + 5, min(view1_crop_size[1], view2_crop_size[1]) - 2)

        if z_interval > 2:
            p = np.random.rand(1)
            if p <= 0.8:
                intersection_z = np.random.randint(5, min(view1_crop_size[2], view2_crop_size[2]) - 2)
            else:
                intersection_z = np.random.randint(-z_interval + 1, -1)
        else:
            intersection_z = np.random.randint(-z_interval + 1, min(view1_crop_size[2], view2_crop_size[2]))

        overall_crop_size_y = (view1_crop_size[0] + view2_crop_size[0]) - intersection_y
        overall_crop_size_x = (view1_crop_size[1] + view2_crop_size[1]) - intersection_x
        overall_crop_size_z = (view1_crop_size[2] + view2_crop_size[2]) - intersection_z

        overall_x_min_left_up_corner = rangex_min
        overall_x_max_left_up_corner = rangex_max - overall_crop_size_x
        overall_y_min_left_up_corner = rangey_min
        overall_y_max_left_up_corner = rangey_max - overall_crop_size_y
        overall_z_min_left_up_corner = rangez_min
        overall_z_max_left_up_corner = rangez_max - overall_crop_size_z
        assert overall_x_min_left_up_corner <= overall_x_max_left_up_corner and \
               overall_y_min_left_up_corner <= overall_y_max_left_up_corner and \
               overall_z_min_left_up_corner <= overall_z_max_left_up_corner, \
            data['image_fn']

        overall_crop_left_up = (np.random.randint(overall_y_min_left_up_corner, overall_y_max_left_up_corner),
                                np.random.randint(overall_x_min_left_up_corner, overall_x_max_left_up_corner),
                                np.random.randint(overall_z_min_left_up_corner, overall_z_max_left_up_corner))
        overall_crop_left_up = np.asarray(overall_crop_left_up)

        view1_x_min = 0
        view1_x_max = overall_crop_size_x - view1_crop_size[1]
        view1_y_min = 0
        view1_y_max = overall_crop_size_y - view1_crop_size[0]
        view1_z_min = 0
        view1_z_max = overall_crop_size_z - view1_crop_size[2]

        assert view1_z_min < view1_z_max and view1_y_min < view1_y_max and view1_x_min < view1_x_max, data[
            'image_fn']

        view2_x_min = 0
        view2_x_max = overall_crop_size_x - view2_crop_size[1]
        view2_y_min = 0
        view2_y_max = overall_crop_size_y - view2_crop_size[0]
        view2_z_min = 0
        view2_z_max = overall_crop_size_z - view2_crop_size[2]

        assert view2_z_min < view2_z_max and view2_y_min < view2_y_max and view2_x_min < view2_x_max, data[
            'image_fn']

        view1_left_up = (np.random.randint(view1_y_min, view1_y_max),
                         np.random.randint(view1_x_min, view1_x_max),
                         np.random.randint(view1_z_min, view1_z_max))
        view1_left_up = np.asarray(view1_left_up)

        view2_left_up = (np.random.randint(view2_y_min, view2_y_max),
                         np.random.randint(view2_x_min, view2_x_max),
                         np.random.randint(view2_z_min, view2_z_max))
        view2_left_up = np.asarray(view2_left_up)
        view1_y1x1z1_crop = view1_left_up + overall_crop_left_up
        view1_y2x2z2_crop = view1_left_up + overall_crop_left_up + view1_crop_size
        view2_y1x1z1_crop = view2_left_up + overall_crop_left_up
        view2_y2x2z2_crop = view2_left_up + overall_crop_left_up + view2_crop_size

        crop_matrix = np.array([view1_y1x1z1_crop, view1_y2x2z2_crop, view2_y1x1z1_crop, view2_y2x2z2_crop])

        crop_matrix_origin = np.dot(crop_matrix, affine_standard_to_origin[:3, :3])
        crop_matrix_origin[::2, :] = np.floor(crop_matrix_origin[::2, :])
        crop_matrix_origin[1::2, :] = np.ceil(crop_matrix_origin[1::2, :])
        crop_matrix_origin = crop_matrix_origin.astype(int)
        view1_spacing = np.array(self.standard_spacing) * view1_scale
        view2_spacing = np.array(self.standard_spacing) * view2_scale
        spacing_matrix = np.array([view1_spacing, view2_spacing])

        data['crop_matrix_origin'] = crop_matrix_origin
        data['spacing_matrix'] = spacing_matrix
        return data


@PIPELINES.register_module()
class compute_crop_patch():
    def __init__(self, standard_spacing=(2., 2., 2.), patch_size=(96, 96, 96), aspect_ratio=[0.5, 1.5]):  # yxz
        self.standard_spacing = standard_spacing
        self.standard_patch_size = patch_size
        self.aspect_ratio = aspect_ratio

    def __call__(self, data):
        origin_spacing = data['subject'].image.spacing
        origin_size = data['subject'].image.shape
        standard_size = [np.floor(origin_spacing[0] / self.standard_spacing[0] * origin_size[1]).astype(int),
                         np.floor(origin_spacing[1] / self.standard_spacing[1] * origin_size[2]).astype(int),
                         np.floor(origin_spacing[2] / self.standard_spacing[2] * origin_size[3]).astype(int)]
        # print(standard_size)
        affine_standard_to_origin = np.array([self.standard_spacing[0] / origin_spacing[0], 0, 0, 0,
                                              0, self.standard_spacing[1] / origin_spacing[1], 0, 0,
                                              0, 0, self.standard_spacing[2] / origin_spacing[2], 0,
                                              0, 0, 0, 1]).reshape(4, 4).transpose()

        patchsize_y = self.standard_patch_size[0] * np.random.uniform(self.aspect_ratio[0], self.aspect_ratio[1])
        patchsize_x = self.standard_patch_size[1] * np.random.uniform(self.aspect_ratio[0], self.aspect_ratio[1])
        patchsize_z = self.standard_patch_size[0] * self.standard_patch_size[1] * self.standard_patch_size[
            2] / patchsize_y / patchsize_x
        patchsize_y = int(patchsize_y)
        patchsize_x = int(patchsize_x)
        patchsize_z = int(patchsize_z)
        patchsize_y = min(standard_size[0] - 12, patchsize_y)
        patchsize_x = min(standard_size[1] - 12, patchsize_x)
        patchsize_z = min(standard_size[2] - 12, patchsize_z)
        use_patch_size = (patchsize_y, patchsize_x, patchsize_z)

        left_up = np.array((np.random.randint(5, standard_size[0] - patchsize_y - 5),
                            np.random.randint(5, standard_size[1] - patchsize_x - 5),
                            np.random.randint(5, standard_size[2] - patchsize_z - 5)))
        right_bottom = left_up + np.array(use_patch_size)
        crop_matrix = np.array([left_up, right_bottom])
        crop_matrix_remap = np.dot(crop_matrix, affine_standard_to_origin[:3, :3])
        crop_matrix_remap[0, :] = np.floor(crop_matrix_remap[0, :]).astype('int')
        crop_matrix_remap[1, :] = np.ceil(crop_matrix_remap[1, :]).astype('int')
        data['crop_matrix_origin'] = crop_matrix_remap
        center = np.array((crop_matrix_remap[0, :] + crop_matrix_remap[1, :]) / 2).astype(int)
        data['subject']['labels_map'].data[0, center[0], center[1], center[2]] = 0
        return data


@PIPELINES.register_module()
class ComputeAugParamSAMV2():
    def __init__(self, scale=(0.8, 1.2), standard_spacing=(4., 4., 4.)):
        self.scale = scale
        self.standard_spacing = standard_spacing

    def __call__(self, data):
        view1_scale = np.random.uniform(self.scale[0], self.scale[1])
        view2_scale = np.random.uniform(self.scale[0], self.scale[1])
        # view1_scale = 0.8
        # view2_scale = 1.2
        view1_spacing = np.array(self.standard_spacing) * view1_scale
        view2_spacing = np.array(self.standard_spacing) * view2_scale
        spacing_matrix = np.array([view1_spacing, view2_spacing])
        data['spacing_matrix'] = spacing_matrix
        return data

sam/datasets/piplines/process3d.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Xiaoyu Bai, Medical AI Lab, Alibaba DAMO Academy

from mmdet.datasets.builder import PIPELINES
import torchio as tio
import numpy as np
from mmdet.datasets.pipelines import Compose
import copy
import torch
import os
import nibabel as nib
from skimage import measure
from scipy import ndimage
# import ipdb


@PIPELINES.register_module()
class DynamicSpacing:
    def __init__(self, norm_spacing=(2., 2., 2.), scale=(0.8, 1.2)):
        self.norm_spacing = norm_spacing
        self.scale = scale
    def __call__(self, data):
        scale = np.random.uniform(self.scale[0], self.scale[1])
        data['resample_spacing'] = np.array(self.norm_spacing) * scale
        return data


@PIPELINES.register_module()
class SeperateCTMRIImage:
    def __call__(self, data):
        mask_path = data['mask_path']
        if data['tag'] == 'im1':
            data = data['data1']
        elif data['tag'] == 'im2':
            data = data['data2']
        else:
            return NotImplementedError
        data['mask_path'] = mask_path
        return data

@PIPELINES.register_module()
class AddSuffixByProb:
    def __init__(self, suffix: [str], p: [float]):
        self.suffix = suffix
        self.p = np.array(p)
        self.cum_p = np.cumsum(p)

    def __call__(self, data):
        porb = np.random.random()
        ind = np.searchsorted(self.cum_p, porb)
        data['image_fn'] = data['image_fn'] + self.suffix[ind]
        if 'mod' not in data.keys():
            if '0000' in self.suffix[ind]:
                data['mod'] = 'ct'
            elif '0001' in  self.suffix[ind] or '0002' in  self.suffix[ind]:
                data['mod'] = 'mri'
        return data


@PIPELINES.register_module()
class stoppoint:
    def __call__(self, data):
        print('stoppoint')
        return data

sam/datasets/piplines/process3d_cross_mod.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai

from mmdet.datasets.builder import PIPELINES
from .process3d import meshgrid3d, makelabel
import torchio as tio
import numpy as np
from mmdet.datasets.pipelines import Compose
import copy
import torch
import os
import nibabel as nib
from skimage import measure
from scipy import ndimage


@PIPELINES.register_module()
class GeneratePairedImagesMaskInfo:
    def __call__(self, data):
        anchor = 1  # select anchor organ for crop
        data1 = {'data_path': data['data_path'], 'image_fn': data['image_fn1']}
        data2 = {'data_path': data['data_path'], 'image_fn': data['image_fn2']}
        data['paired'] = [data['image_fn1'], data['image_fn2']]
        data['data1'] = data1
        data['data2'] = data2
        data['anchor'] = anchor
        return data


@PIPELINES.register_module()
class GenerateCTMRIPairInfo:
    def __call__(self, data):
        data1 = {'data_path': data['data_path'], 'image_fn': data['image_fn']}
        data2 = {'data_path': data['data_path'], 'image_fn': data['image_fn']}
        data['data1'] = data1
        data['data2'] = data2
        return data



@PIPELINES.register_module()
class AddSuffix:
    def __init__(self, suffix: str):
        self.suffix = suffix

    def __call__(self, data):
        data['image_fn'] = data['image_fn'] + self.suffix
        return data


@PIPELINES.register_module()
class LoadTioImageWithMask:
    """Load image as torchio subject.
    """

    def __init__(self, re_orient=True, with_mesh=False):
        self.re_orient = re_orient
        self.with_mesh = with_mesh

    def __call__(self, data):
        # time1 = time.time()
        data_path = data['data_path']
        image_fn = data['image_fn']
        if 'mask_path' in data.keys():
            mask_path = data['mask_path']
            mask_tio = tio.LabelMap(os.path.join(mask_path, image_fn))

        img_tio = tio.ScalarImage(os.path.join(data_path, image_fn))

        # print(img_tio.orientation,image_fn)

        # for NIH lymph node dataset covert orientation to PLS+
        if self.re_orient:
            img_data = img_tio.data
            img_tio.data = img_data.permute(0, 2, 1, 3)
            img_tio.affine = np.array(
                [img_tio.affine[1, :], img_tio.affine[0, :], img_tio.affine[2, :], img_tio.affine[3, :]])
            if 'mask_path' in data.keys():
                mask_data = mask_tio.data
                mask_tio.data = mask_data.permute(0, 2, 1, 3)
                mask_tio.affine = np.array(
                    [mask_tio.affine[1, :], mask_tio.affine[0, :], mask_tio.affine[2, :], mask_tio.affine[3, :]])
        img_tio_shape = img_tio.data.shape  # tio use lazy loading, so we read image shape from data.shape to force it load image here
        img_tio_spacing = img_tio.spacing
        if self.with_mesh:
            meshs = meshgrid3d(img_tio_shape, img_tio_spacing) #3 channels y x z
            labels = makelabel(img_tio_shape)
        if 'mask_path' in data.keys():
            if self.with_mesh:
                subject = tio.Subject(
                    image=img_tio,
                    mask=mask_tio,
                    label_tio_y=tio.ScalarImage(
                        tensor=meshs[0].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                        affine=img_tio.affine),
                    label_tio_x=tio.ScalarImage(
                        tensor=meshs[1].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                        affine=img_tio.affine),
                    label_tio_z=tio.ScalarImage(
                        tensor=meshs[2].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                        affine=img_tio.affine),
                    labels_map=tio.LabelMap(tensor=labels, affine=img_tio.affine)
                    # labels_map=tio.ScalarImage(tensor=labels, affine=img_tio.affine)
                )
            else:
                subject = tio.Subject(
                    image=img_tio,
                    mask=mask_tio)

        else:
            if self.with_mesh:
                subject = tio.Subject(
                    image=img_tio,
                    label_tio_y=tio.ScalarImage(
                        tensor=meshs[0].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                        affine=img_tio.affine),
                    label_tio_x=tio.ScalarImage(
                        tensor=meshs[1].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                        affine=img_tio.affine),
                    label_tio_z=tio.ScalarImage(
                        tensor=meshs[2].reshape(1, img_tio_shape[1], img_tio_shape[2], img_tio_shape[3]),
                        affine=img_tio.affine),
                    labels_map=tio.LabelMap(tensor=labels, affine=img_tio.affine)
                    # labels_map=tio.ScalarImage(tensor=labels, affine=img_tio.affine)
                )
            else:
                subject = tio.Subject(
                    image=img_tio
                )

        data['subject'] = subject
        return data

sam/datasets/piplines/process3d_mask.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
import itertools

# import ipdb
import numpy as np
import torch
from mmcv.runner import get_dist_info
from torch.utils.data.sampler import Sampler
import numpy


class InfiniteBalancedBatchSampler(Sampler):
    """Similar to `BatchSampler` warping a `DistributedSampler. It is designed
    iteration-based runners like `IterBasedRunner` and yields a mini-batch
    indices each time.

    The implementation logic is referred to
    https://github.com/facebookresearch/detectron2/blob/main/detectron2/data/samplers/grouped_batch_sampler.py

    Args:
        dataset (object): The dataset.
        batch_size (int): When model is :obj:`DistributedDataParallel`,
            it is the number of training samples on each GPU,
            When model is :obj:`DataParallel`, it is
            `num_gpus * samples_per_gpu`.
            Default : 1.
        world_size (int, optional): Number of processes participating in
            distributed training. Default: None.
        rank (int, optional): Rank of current process. Default: None.
        seed (int): Random seed. Default: 0.
        shuffle (bool): Whether shuffle the dataset or not. Default: True.
    """  # noqa: W605

    def __init__(self,
                 dataset,
                 batch_size=1,
                 world_size=None,
                 rank=None,
                 seed=0,
                 shuffle=True,
                 ratio=[5, 1]):
        _rank, _world_size = get_dist_info()
        if world_size is None:
            world_size = _world_size
        if rank is None:
            rank = _rank
        self.rank = rank
        self.world_size = world_size
        self.dataset = dataset
        self.batch_size = batch_size
        self.ratio = ratio
        # assert self.batch_size % (self.ratio[0] + self.ratio[1]) == 0
        self.seed = seed if seed is not None else 0
        self.shuffle = shuffle
        # self.size = dataset.datasets[0].__len__()
        self.size = dataset.cummulative_sizes[-2]
        self.cross_range = dataset.cummulative_sizes
        self.indices = self._indices_of_rank()

    def _infinite_indices(self):
        """Infinitely yield a sequence of indices."""
        g = torch.Generator()
        g.manual_seed(self.seed)
        while True:
            if self.shuffle:
                yield from torch.randperm(self.size, generator=g).tolist()

            else:
                yield from torch.arange(self.size).tolist()

    def _indices_of_rank(self):
        """Slice the infinite indices by rank."""
        yield from itertools.islice(self._infinite_indices(), self.rank, None,
                                    self.world_size)

    def __iter__(self):
        # once batch size is reached, yield the indices
        batch_buffer = []
        for idx in self.indices:
            if len(batch_buffer) >= self.batch_size - 1:
            # if len(batch_buffer) >= self.batch_size:  # first intra training no inter data used
                idx = np.random.randint(self.cross_range[-2],self.cross_range[-1])
            batch_buffer.append(idx)
            if len(batch_buffer) == self.batch_size:
                yield batch_buffer
                batch_buffer = []

    def __len__(self):
        """Length of base dataset."""
        return self.size

    def set_epoch(self, epoch):
        """Not supported in `IterationBased` runner."""
        raise NotImplementedError

sam/datasets/samplers/infinite_balanced_sampler.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .frameworks import sam
sam/models/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import warnings

import torch.nn as nn
import torch.utils.checkpoint as cp
from mmcv.cnn import (ConvModule, NonLocal3d, build_activation_layer,
                      constant_init, kaiming_init)
from mmcv.runner import _load_checkpoint, load_checkpoint
from mmcv.utils import _BatchNorm
from torch.nn.modules.utils import _ntuple, _triple

from mmdet.utils import get_root_logger
from mmdet.models.builder import BACKBONES

# try:
#     from mmdet.models import BACKBONES as MMDET_BACKBONES
#     from mmdet.models.builder import SHARED_HEADS as MMDET_SHARED_HEADS
#     mmdet_imported = True
# except (ImportError, ModuleNotFoundError):
#     mmdet_imported = False


class BasicBlock3d(nn.Module):
    """BasicBlock 3d block for ResNet3D.

    Args:
        inplanes (int): Number of channels for the input in first conv3d layer.
        planes (int): Number of channels produced by some norm/conv3d layers.
        spatial_stride (int): Spatial stride in the conv3d layer. Default: 1.
        temporal_stride (int): Temporal stride in the conv3d layer. Default: 1.
        dilation (int): Spacing between kernel elements. Default: 1.
        downsample (nn.Module | None): Downsample layer. Default: None.
        style (str): ``pytorch`` or ``caffe``. If set to "pytorch", the
            stride-two layer is the 3x3 conv layer, otherwise the stride-two
            layer is the first 1x1 conv layer. Default: 'pytorch'.
        inflate (bool): Whether to inflate kernel. Default: True.
        non_local (bool): Determine whether to apply non-local module in this
            block. Default: False.
        non_local_cfg (dict): Config for non-local module. Default: ``dict()``.
        conv_cfg (dict): Config dict for convolution layer.
            Default: ``dict(type='Conv3d')``.
        norm_cfg (dict): Config for norm layers. required keys are ``type``,
            Default: ``dict(type='BN3d')``.
        act_cfg (dict): Config dict for activation layer.
            Default: ``dict(type='ReLU')``.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.
    """
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 spatial_stride=1,
                 temporal_stride=1,
                 dilation=1,
                 downsample=None,
                 style='pytorch',
                 inflate=True,
                 non_local=False,
                 non_local_cfg=dict(),
                 conv_cfg=dict(type='Conv3d'),
                 norm_cfg=dict(type='BN3d'),
                 act_cfg=dict(type='ReLU'),
                 with_cp=False,
                 **kwargs):
        super().__init__()
        assert style in ['pytorch', 'caffe']
        # make sure that only ``inflate_style`` is passed into kwargs
        assert set(kwargs).issubset(['inflate_style'])

        self.inplanes = inplanes
        self.planes = planes
        self.spatial_stride = spatial_stride
        self.temporal_stride = temporal_stride
        self.dilation = dilation
        self.style = style
        self.inflate = inflate
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.with_cp = with_cp
        self.non_local = non_local
        self.non_local_cfg = non_local_cfg

        self.conv1_stride_s = spatial_stride
        self.conv2_stride_s = 1
        self.conv1_stride_t = temporal_stride
        self.conv2_stride_t = 1

        if self.inflate:
            conv1_kernel_size = (3, 3, 3)
            conv1_padding = (1, dilation, dilation)
            conv2_kernel_size = (3, 3, 3)
            conv2_padding = (1, 1, 1)
        else:
            conv1_kernel_size = (1, 3, 3)
            conv1_padding = (0, dilation, dilation)
            conv2_kernel_size = (1, 3, 3)
            conv2_padding = (0, 1, 1)

        self.conv1 = ConvModule(
            inplanes,
            planes,
            conv1_kernel_size,
            stride=(self.conv1_stride_t, self.conv1_stride_s,
                    self.conv1_stride_s),
            padding=conv1_padding,
            dilation=(1, dilation, dilation),
            bias=False,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.conv2 = ConvModule(
            planes,
            planes * self.expansion,
            conv2_kernel_size,
            stride=(self.conv2_stride_t, self.conv2_stride_s,
                    self.conv2_stride_s),
            padding=conv2_padding,
            bias=False,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=None)

        self.downsample = downsample
        self.relu = build_activation_layer(self.act_cfg)

        if self.non_local:
            self.non_local_block = NonLocal3d(self.conv2.norm.num_features,
                                              **self.non_local_cfg)

    def forward(self, x):
        """Defines the computation performed at every call."""

        def _inner_forward(x):
            """Forward wrapper for utilizing checkpoint."""
            identity = x

            out = self.conv1(x)
            out = self.conv2(out)

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

            out = out + identity
            return out

        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(_inner_forward, x)
        else:
            out = _inner_forward(x)
        out = self.relu(out)

        if self.non_local:
            out = self.non_local_block(out)

        return out


class Bottleneck3d(nn.Module):
    """Bottleneck 3d block for ResNet3D.

    Args:
        inplanes (int): Number of channels for the input in first conv3d layer.
        planes (int): Number of channels produced by some norm/conv3d layers.
        spatial_stride (int): Spatial stride in the conv3d layer. Default: 1.
        temporal_stride (int): Temporal stride in the conv3d layer. Default: 1.
        dilation (int): Spacing between kernel elements. Default: 1.
        downsample (nn.Module | None): Downsample layer. Default: None.
        style (str): ``pytorch`` or ``caffe``. If set to "pytorch", the
            stride-two layer is the 3x3 conv layer, otherwise the stride-two
            layer is the first 1x1 conv layer. Default: 'pytorch'.
        inflate (bool): Whether to inflate kernel. Default: True.
        inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines the
            kernel sizes and padding strides for conv1 and conv2 in each block.
            Default: '3x1x1'.
        non_local (bool): Determine whether to apply non-local module in this
            block. Default: False.
        non_local_cfg (dict): Config for non-local module. Default: ``dict()``.
        conv_cfg (dict): Config dict for convolution layer.
            Default: ``dict(type='Conv3d')``.
        norm_cfg (dict): Config for norm layers. required keys are ``type``,
            Default: ``dict(type='BN3d')``.
        act_cfg (dict): Config dict for activation layer.
            Default: ``dict(type='ReLU')``.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.
    """
    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 spatial_stride=1,
                 temporal_stride=1,
                 dilation=1,
                 downsample=None,
                 style='pytorch',
                 inflate=True,
                 inflate_style='3x1x1',
                 non_local=False,
                 non_local_cfg=dict(),
                 conv_cfg=dict(type='Conv3d'),
                 norm_cfg=dict(type='BN3d'),
                 act_cfg=dict(type='ReLU'),
                 with_cp=False):
        super().__init__()
        assert style in ['pytorch', 'caffe']
        assert inflate_style in ['3x1x1', '3x3x3']

        self.inplanes = inplanes
        self.planes = planes
        self.spatial_stride = spatial_stride
        self.temporal_stride = temporal_stride
        self.dilation = dilation
        self.style = style
        self.inflate = inflate
        self.inflate_style = inflate_style
        self.norm_cfg = norm_cfg
        self.conv_cfg = conv_cfg
        self.act_cfg = act_cfg
        self.with_cp = with_cp
        self.non_local = non_local
        self.non_local_cfg = non_local_cfg

        if self.style == 'pytorch':
            self.conv1_stride_s = 1
            self.conv2_stride_s = spatial_stride
            self.conv1_stride_t = 1
            self.conv2_stride_t = temporal_stride
        else:
            self.conv1_stride_s = spatial_stride
            self.conv2_stride_s = 1
            self.conv1_stride_t = temporal_stride
            self.conv2_stride_t = 1

        if self.inflate:
            if inflate_style == '3x1x1':
                conv1_kernel_size = (3, 1, 1)
                conv1_padding = (1, 0, 0)
                conv2_kernel_size = (1, 3, 3)
                conv2_padding = (0, dilation, dilation)
            else:
                conv1_kernel_size = (1, 1, 1)
                conv1_padding = (0, 0, 0)
                conv2_kernel_size = (3, 3, 3)
                conv2_padding = (1, dilation, dilation)
        else:
            conv1_kernel_size = (1, 1, 1)
            conv1_padding = (0, 0, 0)
            conv2_kernel_size = (1, 3, 3)
            conv2_padding = (0, dilation, dilation)

        self.conv1 = ConvModule(
            inplanes,
            planes,
            conv1_kernel_size,
            stride=(self.conv1_stride_t, self.conv1_stride_s,
                    self.conv1_stride_s),
            padding=conv1_padding,
            bias=False,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.conv2 = ConvModule(
            planes,
            planes,
            conv2_kernel_size,
            stride=(self.conv2_stride_t, self.conv2_stride_s,
                    self.conv2_stride_s),
            padding=conv2_padding,
            dilation=(1, dilation, dilation),
            bias=False,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.conv3 = ConvModule(
            planes,
            planes * self.expansion,
            1,
            bias=False,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            # No activation in the third ConvModule for bottleneck
            act_cfg=None)

        self.downsample = downsample
        self.relu = build_activation_layer(self.act_cfg)

        if self.non_local:
            self.non_local_block = NonLocal3d(self.conv3.norm.num_features,
                                              **self.non_local_cfg)

    def forward(self, x):
        """Defines the computation performed at every call."""

        def _inner_forward(x):
            """Forward wrapper for utilizing checkpoint."""
            identity = x

            out = self.conv1(x)
            out = self.conv2(out)
            out = self.conv3(out)

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

            out = out + identity
            return out

        if self.with_cp and x.requires_grad:
            out = cp.checkpoint(_inner_forward, x)
        else:
            out = _inner_forward(x)
        out = self.relu(out)

        if self.non_local:
            out = self.non_local_block(out)

        return out


@BACKBONES.register_module()
class ResNet3d(nn.Module):
    """ResNet 3d backbone.

    Args:
        depth (int): Depth of resnet, from {18, 34, 50, 101, 152}.
        pretrained (str | None): Name of pretrained model.
        stage_blocks (tuple | None): Set number of stages for each res layer.
            Default: None.
        pretrained2d (bool): Whether to load pretrained 2D model.
            Default: True.
        in_channels (int): Channel num of input features. Default: 3.
        base_channels (int): Channel num of stem output features. Default: 64.
        out_indices (Sequence[int]): Indices of output feature. Default: (3, ).
        num_stages (int): Resnet stages. Default: 4.
        spatial_strides (Sequence[int]):
            Spatial strides of residual blocks of each stage.
            Default: ``(1, 2, 2, 2)``.
        temporal_strides (Sequence[int]):
            Temporal strides of residual blocks of each stage.
            Default: ``(1, 1, 1, 1)``.
        dilations (Sequence[int]): Dilation of each stage.
            Default: ``(1, 1, 1, 1)``.
        conv1_kernel (Sequence[int]): Kernel size of the first conv layer.
            Default: ``(3, 7, 7)``.
        conv1_stride_s (int): Spatial stride of the first conv layer.
            Default: 2.
        conv1_stride_t (int): Temporal stride of the first conv layer.
            Default: 1.
        pool1_stride_s (int): Spatial stride of the first pooling layer.
            Default: 2.
        pool1_stride_t (int): Temporal stride of the first pooling layer.
            Default: 1.
        with_pool2 (bool): Whether to use pool2. Default: True.
        style (str): `pytorch` or `caffe`. If set to "pytorch", the stride-two
            layer is the 3x3 conv layer, otherwise the stride-two layer is
            the first 1x1 conv layer. Default: 'pytorch'.
        frozen_stages (int): Stages to be frozen (all param fixed). -1 means
            not freezing any parameters. Default: -1.
        inflate (Sequence[int]): Inflate Dims of each block.
            Default: (1, 1, 1, 1).
        inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines the
            kernel sizes and padding strides for conv1 and conv2 in each block.
            Default: '3x1x1'.
        conv_cfg (dict): Config for conv layers. required keys are ``type``
            Default: ``dict(type='Conv3d')``.
        norm_cfg (dict): Config for norm layers. required keys are ``type`` and
            ``requires_grad``.
            Default: ``dict(type='BN3d', requires_grad=True)``.
        act_cfg (dict): Config dict for activation layer.
            Default: ``dict(type='ReLU', inplace=True)``.
        norm_eval (bool): Whether to set BN layers to eval mode, namely, freeze
            running stats (mean and var). Default: False.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.
        non_local (Sequence[int]): Determine whether to apply non-local module
            in the corresponding block of each stages. Default: (0, 0, 0, 0).
        non_local_cfg (dict): Config for non-local module. Default: ``dict()``.
        zero_init_residual (bool):
            Whether to use zero initialization for residual block,
            Default: True.
        kwargs (dict, optional): Key arguments for "make_res_layer".
    """

    arch_settings = {
        18: (BasicBlock3d, (2, 2, 2, 2)),
        34: (BasicBlock3d, (3, 4, 6, 3)),
        50: (Bottleneck3d, (3, 4, 6, 3)),
        101: (Bottleneck3d, (3, 4, 23, 3)),
        152: (Bottleneck3d, (3, 8, 36, 3))
    }

    def __init__(self,
                 depth,
                 pretrained,
                 stage_blocks=None,
                 pretrained2d=True,
                 in_channels=3,
                 num_stages=4,
                 base_channels=64,
                 out_indices=(0,1,2,3),
                 spatial_strides=(1, 2, 2, 2),
                 temporal_strides=(1, 1, 1, 1),
                 dilations=(1, 1, 1, 1),
                 conv1_kernel=(3, 7, 7),
                 conv1_stride_s=2,
                 conv1_stride_t=1,
                 pool1_stride_s=2,
                 pool1_stride_t=1,
                 with_pool1=True,
                 with_pool2=True,
                 style='pytorch',
                 frozen_stages=-1,
                 inflate=(1, 1, 1, 1),
                 inflate_style='3x1x1',
                 conv_cfg=dict(type='Conv3d'),
                 norm_cfg=dict(type='BN3d', requires_grad=True),
                 act_cfg=dict(type='ReLU', inplace=True),
                 norm_eval=False,
                 with_cp=False,
                 non_local=(0, 0, 0, 0),
                 non_local_cfg=dict(),
                 zero_init_residual=True,
                 **kwargs):
        super().__init__()
        if depth not in self.arch_settings:
            raise KeyError(f'invalid depth {depth} for resnet')
        self.depth = depth
        self.pretrained = pretrained
        self.pretrained2d = pretrained2d
        self.in_channels = in_channels
        self.base_channels = base_channels
        self.num_stages = num_stages
        assert 1 <= num_stages <= 4
        self.stage_blocks = stage_blocks
        self.out_indices = out_indices
        assert max(out_indices) < num_stages
        self.spatial_strides = spatial_strides
        self.temporal_strides = temporal_strides
        self.dilations = dilations
        assert len(spatial_strides) == len(temporal_strides) == len(
            dilations) == num_stages
        if self.stage_blocks is not None:
            assert len(self.stage_blocks) == num_stages

        self.conv1_kernel = conv1_kernel
        self.conv1_stride_s = conv1_stride_s
        self.conv1_stride_t = conv1_stride_t
        self.pool1_stride_s = pool1_stride_s
        self.pool1_stride_t = pool1_stride_t
        self.with_pool1 = with_pool1
        self.with_pool2 = with_pool2
        self.style = style
        self.frozen_stages = frozen_stages
        self.stage_inflations = _ntuple(num_stages)(inflate)
        self.non_local_stages = _ntuple(num_stages)(non_local)
        self.inflate_style = inflate_style
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.norm_eval = norm_eval
        self.with_cp = with_cp
        self.zero_init_residual = zero_init_residual

        self.block, stage_blocks = self.arch_settings[depth]

        if self.stage_blocks is None:
            self.stage_blocks = stage_blocks[:num_stages]

        self.inplanes = self.base_channels

        self.non_local_cfg = non_local_cfg

        self._make_stem_layer()

        self.res_layers = []
        for i, num_blocks in enumerate(self.stage_blocks):
            spatial_stride = spatial_strides[i]
            temporal_stride = temporal_strides[i]
            dilation = dilations[i]
            planes = self.base_channels * 2**i
            res_layer = self.make_res_layer(
                self.block,
                self.inplanes,
                planes,
                num_blocks,
                spatial_stride=spatial_stride,
                temporal_stride=temporal_stride,
                dilation=dilation,
                style=self.style,
                norm_cfg=self.norm_cfg,
                conv_cfg=self.conv_cfg,
                act_cfg=self.act_cfg,
                non_local=self.non_local_stages[i],
                non_local_cfg=self.non_local_cfg,
                inflate=self.stage_inflations[i],
                inflate_style=self.inflate_style,
                with_cp=with_cp,
                **kwargs)
            self.inplanes = planes * self.block.expansion
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, res_layer)
            self.res_layers.append(layer_name)

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

    @staticmethod
    def make_res_layer(block,
                       inplanes,
                       planes,
                       blocks,
                       spatial_stride=1,
                       temporal_stride=1,
                       dilation=1,
                       style='pytorch',
                       inflate=1,
                       inflate_style='3x1x1',
                       non_local=0,
                       non_local_cfg=dict(),
                       norm_cfg=None,
                       act_cfg=None,
                       conv_cfg=None,
                       with_cp=False,
                       **kwargs):
        """Build residual layer for ResNet3D.

        Args:
            block (nn.Module): Residual module to be built.
            inplanes (int): Number of channels for the input feature
                in each block.
            planes (int): Number of channels for the output feature
                in each block.
            blocks (int): Number of residual blocks.
            spatial_stride (int | Sequence[int]): Spatial strides in
                residual and conv layers. Default: 1.
            temporal_stride (int | Sequence[int]): Temporal strides in
                residual and conv layers. Default: 1.
            dilation (int): Spacing between kernel elements. Default: 1.
            style (str): ``pytorch`` or ``caffe``. If set to ``pytorch``,
                the stride-two layer is the 3x3 conv layer, otherwise
                the stride-two layer is the first 1x1 conv layer.
                Default: ``pytorch``.
            inflate (int | Sequence[int]): Determine whether to inflate
                for each block. Default: 1.
            inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines
                the kernel sizes and padding strides for conv1 and conv2
                in each block. Default: '3x1x1'.
            non_local (int | Sequence[int]): Determine whether to apply
                non-local module in the corresponding block of each stages.
                Default: 0.
            non_local_cfg (dict): Config for non-local module.
                Default: ``dict()``.
            conv_cfg (dict | None): Config for norm layers. Default: None.
            norm_cfg (dict | None): Config for norm layers. Default: None.
            act_cfg (dict | None): Config for activate layers. Default: None.
            with_cp (bool | None): Use checkpoint or not. Using checkpoint
                will save some memory while slowing down the training speed.
                Default: False.

        Returns:
            nn.Module: A residual layer for the given config.
        """
        inflate = inflate if not isinstance(inflate,
                                            int) else (inflate, ) * blocks
        non_local = non_local if not isinstance(
            non_local, int) else (non_local, ) * blocks
        assert len(inflate) == blocks and len(non_local) == blocks
        downsample = None
        if spatial_stride != 1 or inplanes != planes * block.expansion:
            downsample = ConvModule(
                inplanes,
                planes * block.expansion,
                kernel_size=1,
                stride=(temporal_stride, spatial_stride, spatial_stride),
                bias=False,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=None)

        layers = []
        layers.append(
            block(
                inplanes,
                planes,
                spatial_stride=spatial_stride,
                temporal_stride=temporal_stride,
                dilation=dilation,
                downsample=downsample,
                style=style,
                inflate=(inflate[0] == 1),
                inflate_style=inflate_style,
                non_local=(non_local[0] == 1),
                non_local_cfg=non_local_cfg,
                norm_cfg=norm_cfg,
                conv_cfg=conv_cfg,
                act_cfg=act_cfg,
                with_cp=with_cp,
                **kwargs))
        inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(
                block(
                    inplanes,
                    planes,
                    spatial_stride=1,
                    temporal_stride=1,
                    dilation=dilation,
                    style=style,
                    inflate=(inflate[i] == 1),
                    inflate_style=inflate_style,
                    non_local=(non_local[i] == 1),
                    non_local_cfg=non_local_cfg,
                    norm_cfg=norm_cfg,
                    conv_cfg=conv_cfg,
                    act_cfg=act_cfg,
                    with_cp=with_cp,
                    **kwargs))

        return nn.Sequential(*layers)

    @staticmethod
    def _inflate_conv_params(conv3d, state_dict_2d, module_name_2d,
                             inflated_param_names,flag = 0):
        """Inflate a conv module from 2d to 3d.

        Args:
            conv3d (nn.Module): The destination conv3d module.
            state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
            module_name_2d (str): The name of corresponding conv module in the
                2d model.
            inflated_param_names (list[str]): List of parameters that have been
                inflated.
        """
        weight_2d_name = module_name_2d + '.weight'

        conv2d_weight = state_dict_2d[weight_2d_name]
        kernel_t = conv3d.weight.data.shape[2]

        if flag == 1:
            new_weight = conv2d_weight.data.unsqueeze(1).expand_as(
                conv3d.weight)
        else:
            new_weight = conv2d_weight.data.unsqueeze(2).expand_as(
                conv3d.weight) / kernel_t
        conv3d.weight.data.copy_(new_weight)
        inflated_param_names.append(weight_2d_name)

        if getattr(conv3d, 'bias') is not None:
            bias_2d_name = module_name_2d + '.bias'
            conv3d.bias.data.copy_(state_dict_2d[bias_2d_name])
            inflated_param_names.append(bias_2d_name)

    @staticmethod
    def _inflate_bn_params(bn3d, state_dict_2d, module_name_2d,
                           inflated_param_names):
        """Inflate a norm module from 2d to 3d.

        Args:
            bn3d (nn.Module): The destination bn3d module.
            state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
            module_name_2d (str): The name of corresponding bn module in the
                2d model.
            inflated_param_names (list[str]): List of parameters that have been
                inflated.
        """
        for param_name, param in bn3d.named_parameters():
            param_2d_name = f'{module_name_2d}.{param_name}'
            param_2d = state_dict_2d[param_2d_name]
            if param.data.shape != param_2d.shape:
                warnings.warn(f'The parameter of {module_name_2d} is not'
                              'loaded due to incompatible shapes. ')
                return

            param.data.copy_(param_2d)
            inflated_param_names.append(param_2d_name)

        for param_name, param in bn3d.named_buffers():
            param_2d_name = f'{module_name_2d}.{param_name}'
            # some buffers like num_batches_tracked may not exist in old
            # checkpoints
            if param_2d_name in state_dict_2d:
                param_2d = state_dict_2d[param_2d_name]
                param.data.copy_(param_2d)
                inflated_param_names.append(param_2d_name)

    @staticmethod
    def _inflate_weights(self, logger):
        """Inflate the resnet2d parameters to resnet3d.

        The differences between resnet3d and resnet2d mainly lie in an extra
        axis of conv kernel. To utilize the pretrained parameters in 2d model,
        the weight of conv2d models should be inflated to fit in the shapes of
        the 3d counterpart.

        Args:
            logger (logging.Logger): The logger used to print
                debugging information.
        """

        state_dict_r2d = _load_checkpoint(self.pretrained)
        if 'state_dict' in state_dict_r2d:
            state_dict_r2d = state_dict_r2d['state_dict']

        inflated_param_names = []
        for name, module in self.named_modules():
            if isinstance(module, ConvModule):
                # we use a ConvModule to wrap conv+bn+relu layers, thus the
                # name mapping is needed
                if 'downsample' in name:
                    # layer{X}.{Y}.downsample.conv->layer{X}.{Y}.downsample.0
                    original_conv_name = name + '.0'
                    # layer{X}.{Y}.downsample.bn->layer{X}.{Y}.downsample.1
                    original_bn_name = name + '.1'
                else:
                    # layer{X}.{Y}.conv{n}.conv->layer{X}.{Y}.conv{n}
                    original_conv_name = name
                    # layer{X}.{Y}.conv{n}.bn->layer{X}.{Y}.bn{n}
                    original_bn_name = name.replace('conv', 'bn')
                if original_conv_name + '.weight' not in state_dict_r2d:
                    logger.warning(f'Module not exist in the state_dict_r2d'
                                   f': {original_conv_name}')
                else:
                    shape_2d = state_dict_r2d[original_conv_name +
                                              '.weight'].shape
                    shape_3d = module.conv.weight.data.shape
                    if original_conv_name =='conv1':
                        self._inflate_conv_params(module.conv, state_dict_r2d,
                                                  original_conv_name,
                                                  inflated_param_names,flag=1)
                    else:
                        if shape_2d != shape_3d[:2] + shape_3d[3:]:
                            logger.warning(f'Weight shape mismatch for '
                                           f': {original_conv_name} : '
                                           f'3d weight shape: {shape_3d}; '
                                           f'2d weight shape: {shape_2d}. ')
                        else:
                            self._inflate_conv_params(module.conv, state_dict_r2d,
                                                      original_conv_name,
                                                      inflated_param_names)

                if original_bn_name + '.weight' not in state_dict_r2d:
                    logger.warning(f'Module not exist in the state_dict_r2d'
                                   f': {original_bn_name}')
                else:
                    self._inflate_bn_params(module.bn, state_dict_r2d,
                                            original_bn_name,
                                            inflated_param_names)

        # check if any parameters in the 2d checkpoint are not loaded
        remaining_names = set(
            state_dict_r2d.keys()) - set(inflated_param_names)
        if remaining_names:
            logger.info(f'These parameters in the 2d checkpoint are not loaded'
                        f': {remaining_names}')

    def inflate_weights(self, logger):
        self._inflate_weights(self, logger)

    def _make_stem_layer(self):
        """Construct the stem layers consists of a conv+norm+act module and a
        pooling layer."""
        self.conv1 = ConvModule(
            self.in_channels,
            self.base_channels,
            kernel_size=self.conv1_kernel,
            stride=(self.conv1_stride_t, self.conv1_stride_s,
                    self.conv1_stride_s),
            padding=tuple([(k - 1) // 2 for k in _triple(self.conv1_kernel)]),
            bias=False,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.maxpool = nn.MaxPool3d(
            kernel_size=(1, 3, 3),
            stride=(self.pool1_stride_t, self.pool1_stride_s,
                    self.pool1_stride_s),
            padding=(0, 1, 1))

        self.pool2 = nn.MaxPool3d(kernel_size=(2, 1, 1), stride=(2, 1, 1))

    def _freeze_stages(self):
        """Prevent all the parameters from being optimized before
        ``self.frozen_stages``."""
        if self.frozen_stages >= 0:
            self.conv1.eval()
            for param in self.conv1.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

    @staticmethod
    def _init_weights(self, pretrained=None):
        """Initiate the parameters either from existing checkpoint or from
        scratch.

        Args:
            pretrained (str | None): The path of the pretrained weight. Will
                override the original `pretrained` if set. The arg is added to
                be compatible with mmdet. Default: None.
        """
        if pretrained:
            self.pretrained = pretrained
        if isinstance(self.pretrained, str):
            logger = get_root_logger()
            logger.info(f'load model from: {self.pretrained}')

            if self.pretrained2d:
                # Inflate 2D model into 3D model.
                self.inflate_weights(logger)

            else:
                # Directly load 3D model.
                load_checkpoint(
                    self, self.pretrained, strict=False, logger=logger)

        elif self.pretrained is None:
            for m in self.modules():
                if isinstance(m, nn.Conv3d):
                    kaiming_init(m)
                elif isinstance(m, _BatchNorm):
                    constant_init(m, 1)

            if self.zero_init_residual:
                for m in self.modules():
                    if isinstance(m, Bottleneck3d):
                        constant_init(m.conv3.bn, 0)
                    elif isinstance(m, BasicBlock3d):
                        constant_init(m.conv2.bn, 0)
        else:
            raise TypeError('pretrained must be a str or None')

    def init_weights(self, pretrained=None):
        self._init_weights(self, pretrained)

    def forward(self, x):
        """Defines the computation performed at every call.

        Args:
            x (torch.Tensor): The input data.

        Returns:
            torch.Tensor: The feature of the input
            samples extracted by the backbone.
        """
        x = self.conv1(x)
        if self.with_pool1:
            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 == 0 and self.with_pool2:
                x = self.pool2(x)
            if i in self.out_indices:
                outs.append(x)
        if len(outs) == 1:
            return outs[0]

        return tuple(outs)

    def train(self, mode=True):
        """Set the optimization status when training."""
        super().train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()


@BACKBONES.register_module()
class ResNet3dLayer(nn.Module):
    """ResNet 3d Layer.

    Args:
        depth (int): Depth of resnet, from {18, 34, 50, 101, 152}.
        pretrained (str | None): Name of pretrained model.
        pretrained2d (bool): Whether to load pretrained 2D model.
            Default: True.
        stage (int): The index of Resnet stage. Default: 3.
        base_channels (int): Channel num of stem output features. Default: 64.
        spatial_stride (int): The 1st res block's spatial stride. Default 2.
        temporal_stride (int): The 1st res block's temporal stride. Default 1.
        dilation (int): The dilation. Default: 1.
        style (str): `pytorch` or `caffe`. If set to "pytorch", the stride-two
            layer is the 3x3 conv layer, otherwise the stride-two layer is
            the first 1x1 conv layer. Default: 'pytorch'.
        all_frozen (bool): Frozen all modules in the layer. Default: False.
        inflate (int): Inflate Dims of each block. Default: 1.
        inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines the
            kernel sizes and padding strides for conv1 and conv2 in each block.
            Default: '3x1x1'.
        conv_cfg (dict): Config for conv layers. required keys are ``type``
            Default: ``dict(type='Conv3d')``.
        norm_cfg (dict): Config for norm layers. required keys are ``type`` and
            ``requires_grad``.
            Default: ``dict(type='BN3d', requires_grad=True)``.
        act_cfg (dict): Config dict for activation layer.
            Default: ``dict(type='ReLU', inplace=True)``.
        norm_eval (bool): Whether to set BN layers to eval mode, namely, freeze
            running stats (mean and var). Default: False.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Default: False.
        zero_init_residual (bool):
            Whether to use zero initialization for residual block,
            Default: True.
        kwargs (dict, optional): Key arguments for "make_res_layer".
    """

    def __init__(self,
                 depth,
                 pretrained,
                 pretrained2d=True,
                 stage=3,
                 base_channels=64,
                 spatial_stride=2,
                 temporal_stride=1,
                 dilation=1,
                 style='pytorch',
                 all_frozen=False,
                 inflate=1,
                 inflate_style='3x1x1',
                 conv_cfg=dict(type='Conv3d'),
                 norm_cfg=dict(type='BN3d', requires_grad=True),
                 act_cfg=dict(type='ReLU', inplace=True),
                 norm_eval=False,
                 with_cp=False,
                 zero_init_residual=True,
                 **kwargs):

        super().__init__()
        self.arch_settings = ResNet3d.arch_settings
        assert depth in self.arch_settings

        self.make_res_layer = ResNet3d.make_res_layer
        self._inflate_conv_params = ResNet3d._inflate_conv_params
        self._inflate_bn_params = ResNet3d._inflate_bn_params
        self._inflate_weights = ResNet3d._inflate_weights
        self._init_weights = ResNet3d._init_weights

        self.depth = depth
        self.pretrained = pretrained
        self.pretrained2d = pretrained2d
        self.stage = stage
        # stage index is 0 based
        assert 0 <= stage <= 3
        self.base_channels = base_channels

        self.spatial_stride = spatial_stride
        self.temporal_stride = temporal_stride
        self.dilation = dilation

        self.style = style
        self.all_frozen = all_frozen

        self.stage_inflation = inflate
        self.inflate_style = inflate_style
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.norm_eval = norm_eval
        self.with_cp = with_cp
        self.zero_init_residual = zero_init_residual

        block, stage_blocks = self.arch_settings[depth]
        stage_block = stage_blocks[stage]
        planes = 64 * 2**stage
        inplanes = 64 * 2**(stage - 1) * block.expansion

        res_layer = self.make_res_layer(
            block,
            inplanes,
            planes,
            stage_block,
            spatial_stride=spatial_stride,
            temporal_stride=temporal_stride,
            dilation=dilation,
            style=self.style,
            norm_cfg=self.norm_cfg,
            conv_cfg=self.conv_cfg,
            act_cfg=self.act_cfg,
            inflate=self.stage_inflation,
            inflate_style=self.inflate_style,
            with_cp=with_cp,
            **kwargs)

        self.layer_name = f'layer{stage + 1}'
        self.add_module(self.layer_name, res_layer)

    def inflate_weights(self, logger):
        self._inflate_weights(self, logger)

    def _freeze_stages(self):
        """Prevent all the parameters from being optimized before
        ``self.frozen_stages``."""
        if self.all_frozen:
            layer = getattr(self, self.layer_name)
            layer.eval()
            for param in layer.parameters():
                param.requires_grad = False

    def init_weights(self, pretrained=None):
        self._init_weights(self, pretrained)

    def forward(self, x):
        """Defines the computation performed at every call.

        Args:
            x (torch.Tensor): The input data.

        Returns:
            torch.Tensor: The feature of the input
            samples extracted by the backbone.
        """
        res_layer = getattr(self, self.layer_name)
        out = res_layer(x)
        return out

    def train(self, mode=True):
        """Set the optimization status when training."""
        super().train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()


# if mmdet_imported:
#     MMDET_SHARED_HEADS.register_module()(ResNet3dLayer)
#     MMDET_BACKBONES.register_module()(ResNet3d)

sam/models/backbones/resnet3d.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .sam import Sam
# from .sam_transformer import Sam_transformer
# from .sam_single_scale import Sam_single_scale
# from .sam_v2_coarse_deform import Sam_coarse
# from .sam_v2_coarse_cycle import Sam_coarse_cyc
# from .sam_v2_fine import Sam_fine
# from .sam_v2_uniform import Sam_uniform
# from .sam_v2_uniform_bn import Sam_uniform_bn
# from .sam_v2_uniform_cross_volume_mask import Sam_uniform_cross_volume
# from .sam_v2_uniform_cross_volume_mask_neg_loss import Sam_uniform_cross_volume_neg
# from .sam_v2_uniform_cross_volume_mask_dual_head import Sam_uniform_cross_volume_dual_head
from .sam_v2_uniform_cross_volume_mask_dual_head_mean_vector import Sam_uniform_cross_volume_dual_head_mean_vector
# from .sam_v2_uniform_cross_volume_mask_dual_head_mean_vector_col import \
#     Sam_uniform_cross_volume_dual_head_mean_vector_col
# from .sam_v2_uniform_cross_volume_mask_dual_head_col_embeddings import Sam_uniform_cross_volume_dual_head_col_embeddings
# from .sam_v2_uniform_cross_volume_mask_dual_head_embeddings import Sam_uniform_cross_volume_dual_head_embeddings
from .sam_v2_uniform_cross_volume_mask_dual_head_MRI_CT import Sam_uniform_cross_volume_dual_head_mri_ct
# from .sam_v2_uniform_cross_volume_mask_dual_head_MRI_CT_transfer import \
#     Sam_uniform_cross_volume_dual_head_mri_ct_transfer
# from .sam_v2_uniform_cross_volume_mask_dual_head_MRI_CT_dual_encoder import \
#     Sam_uniform_cross_volume_dual_head_mri_ct_daul_encoder


__all__ = ['Sam',
           'Sam_uniform_cross_volume_dual_head_mean_vector',
           'Sam_uniform_cross_volume_dual_head_mri_ct',
           ]

sam/models/frameworks/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai
# Rearrange the ignore computation into useindex, save computation cost
import os.path
import warnings

# import ipdb
import torch
from mmdet.models.builder import DETECTORS, build_backbone, build_head, build_neck
from mmdet.models.detectors.base import BaseDetector
import torch.nn.functional as F
from torch import linalg as LA
import time
import pickle


class nodes():
    def __init__(self,
                 source,
                 target,
                 view='view1'):
        self.source = source
        self.target = target
        self.view = view
        self.ignore_list = []
        self.source_ignore_list = []

    def add_ignore(self, target_list):
        if self.target in target_list:
            target_list.remove(self.target)
            self.ignore_list = target_list

    def add_source_ignore(self, source_list):
        if self.source in source_list:
            source_list.remove(self.source)
        self.source_ignore_list = source_list


@DETECTORS.register_module()
class Sam(BaseDetector):
    def __init__(self,
                 backbone,
                 neck=None,
                 read_out_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 init_cfg=None):
        super(Sam, self).__init__(init_cfg)
        self.backbone = build_backbone(backbone)
        self.backbone.init_weights()
        if neck is not None:
            self.neck = build_neck(neck)
        self.criterion = torch.nn.CrossEntropyLoss().cuda()

        self.read_out_head = build_neck(read_out_head)
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

    def extract_feat(self, img):
        """Directly extract features from the backbone+neck."""
        x = self.backbone(img)
        out1 = self.neck(x[:3])[0]
        out2 = self.read_out_head([x[-1]])[0]
        out1 = F.normalize(out1, dim=1)
        out2 = F.normalize(out2, dim=1)
        out1 = out1.type(torch.half)
        out2 = out2.type(torch.half)
        return [out1, out2]

    def forward_dummy(self, img):
        """Used for computing network flops.

        See `mmdetection/tools/analysis_tools/get_flops.py`
        """
        outs = ()
        # backbone
        x = self.extract_feat(img)
        outs = outs + x

        return outs

    def forward_train(self,
                      img,
                      img_metas,
                      meshgrid,
                      valid,
                      **kwargs):

        """
        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.

            img_metas (list[dict]): list of image info dict where each dict
                has: 'img_shape', 'img_name',etc.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """

        x = self.extract_feat(img)
        losses = dict()
        loss = self.loss(x, meshgrid, img_metas, valid, img)
        losses['loss'] = loss

        return losses

    def loss(self, feats, meshgrid, img_metas, valid, *args):

        N, C, D, H, W = meshgrid.shape
        output_half = F.interpolate(meshgrid, size=(int(D / 2), int(H / 2), int(W / 2)), mode='trilinear', align_corners=False)
        output_416 = F.interpolate(meshgrid, size=(int(D / 4), int(H / 16), int(W / 16)), mode='trilinear', align_corners=False)
        valid_half = F.interpolate(valid, size=(int(D / 2), int(H / 2), int(W / 2)))
        valid_416 = F.interpolate(valid, size=(int(D / 4), int(H / 16), int(W / 16)))

        output_half = output_half.type(torch.half)
        output_416 = output_416.type(torch.half)
        valid_half = valid_half.type(torch.half)
        valid_416 = valid_416.type(torch.half)

        # fine_time_start = time.time()

        for i in range(int(N / 2)):
            result = self.single_fine_loss(feats[0][2 * i:2 * (i + 1)], feats[1][2 * i:2 * (i + 1)],
                                           output_half[2 * i:2 * (i + 1)], output_416[2 * i:2 * (i + 1)],
                                           valid_half[2 * i:2 * (i + 1)], valid_416[2 * i:2 * (i + 1)],
                                           (args[0], feats))
            if i == 0:
                fine_loss = result['loss']
            else:
                fine_loss += result['loss']
        fine_loss = fine_loss / int(N / 2)
        # fine_time_end = time.time()
        # print(f'fine time:{fine_time_end - fine_time_start}')
        # # # ipdb.set_trace()
        result = self.coarse_loss(feats[1], output_416, valid_416)
        # coarse_time_end = time.time()
        # print(f'coarse time:{coarse_time_end - fine_time_end}')
        coarse_loss = result['loss']

        loss = fine_loss + coarse_loss
        return loss

    def single_fine_loss(self, fine_feat, coarse_feat, fine_grid, coarse_grid, fine_vaild, coarse_valid, *args):

        view_1_fine = fine_feat[0, :, :, :, :].view(128, -1)
        view_2_fine = fine_feat[1, :, :, :, :].view(128, -1)

        out = dict()

        fine_local_grid = self.meshgrid3d(fine_feat.shape[2:])  # z y x

        p1_y_min = fine_grid[0, 0, fine_vaild[0, 0] > 0].min()
        p1_y_max = fine_grid[0, 0, fine_vaild[0, 0] > 0].max()
        p1_x_min = fine_grid[0, 1, fine_vaild[0, 0] > 0].min()
        p1_x_max = fine_grid[0, 1, fine_vaild[0, 0] > 0].max()
        p1_z_min = fine_grid[0, 2, fine_vaild[0, 0] > 0].min()
        p1_z_max = fine_grid[0, 2, fine_vaild[0, 0] > 0].max()

        p2_y_min = fine_grid[1, 0, fine_vaild[1, 0] > 0].min()
        p2_y_max = fine_grid[1, 0, fine_vaild[1, 0] > 0].max()
        p2_x_min = fine_grid[1, 1, fine_vaild[1, 0] > 0].min()
        p2_x_max = fine_grid[1, 1, fine_vaild[1, 0] > 0].max()
        p2_z_min = fine_grid[1, 2, fine_vaild[1, 0] > 0].min()
        p2_z_max = fine_grid[1, 2, fine_vaild[1, 0] > 0].max()

        intersection_y = [max(p1_y_min, p2_y_min), min(p1_y_max, p2_y_max)]
        intersection_x = [max(p1_x_min, p2_x_min), min(p1_x_max, p2_x_max)]
        intersection_z = [max(p1_z_min, p2_z_min), min(p1_z_max, p2_z_max)]

        if intersection_y[0] > intersection_y[1] or intersection_x[0] > intersection_x[1] or intersection_z[0] > \
                intersection_z[1]:

            # no overlap
            index_use = torch.randperm(fine_feat.shape[2] * fine_feat.shape[3] * fine_feat.shape[4])
            q_index = index_use[:self.train_cfg.after_select_pos_number]
            fine_local_grid_flatten = fine_local_grid.view(-1, 3)

            q_loc = fine_local_grid_flatten[q_index, :].type(torch.LongTensor)
            pos1_mm_use = fine_grid[0, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]
            pos1_mm_all = fine_grid[0, :, :, :, :].view(3, -1)
            dist1_1 = LA.norm((pos1_mm_use.view(-1, pos1_mm_use.shape[1], 1) - \
                               pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)
            ind_ignores_1_1 = torch.where(dist1_1 < self.train_cfg.ignore_distance)
            ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
            use_mask_1_1 = torch.ones_like(dist1_1)
            use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0

            q_fine_feat = fine_feat[0, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]
            # index_use = torch.randperm(fine_feat.shape[2] * fine_feat.shape[3] * fine_feat.shape[4])
            # q_index = index_use[:self.train_cfg.after_select_pos_number]
            # q_loc = fine_local_grid_flatten[q_index, :].type(torch.LongTensor)
            # q_fine_feat[1, :, :] = fine_feat[1, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]

            # q_fine_feat = q_fine_feat.permute(0, 2, 1).reshape(-1, 128)
            q_fine_feat = q_fine_feat.permute(1, 0)

            neg_fine = torch.einsum("nc,ck->nk", q_fine_feat, view_1_fine)
            neg_fine = neg_fine * use_mask_1_1
            inner = torch.einsum("nc,nc->n", q_fine_feat, q_fine_feat).view(-1, 1)

            neg_fine_use_index = neg_fine.topk(self.train_cfg.after_select_neg_number)[1]
            neg_fine_use = torch.zeros((neg_fine.shape[0], self.train_cfg.after_select_neg_number),
                                       device=fine_feat.device)
            for i in range(neg_fine.shape[0]):
                neg_fine_use[i, :] = neg_fine[i, neg_fine_use_index[i, :]]
            logits = torch.cat([inner, neg_fine_use], dim=1)
            logits = logits / self.train_cfg.temperature
            labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
            loss = self.criterion(logits, labels)
        else:
            # have overlap
            intersection_volume_sub1 = (fine_grid[0, 0] >= intersection_y[0]) * (
                    fine_grid[0, 0] <= intersection_y[1]) \
                                       * (fine_grid[0, 1] >= intersection_x[0]) * (
                                               fine_grid[0, 1] <= intersection_x[1]) \
                                       * (fine_grid[0, 2] >= intersection_z[0]) * (
                                               fine_grid[0, 2] <= intersection_z[1]) * fine_vaild[0, 0]

            intersection_volume_sub2 = (fine_grid[1, 0] >= intersection_y[0]) * (
                    fine_grid[1, 0] <= intersection_y[1]) \
                                       * (fine_grid[1, 1] >= intersection_x[0]) * (
                                               fine_grid[1, 1] <= intersection_x[1]) \
                                       * (fine_grid[1, 2] >= intersection_z[0]) * (
                                               fine_grid[1, 2] <= intersection_z[1]) * fine_vaild[1, 0]
            pos1_index_overlap = fine_local_grid[intersection_volume_sub1 > 0, :]
            pos2_index_overlap = fine_local_grid[intersection_volume_sub2 > 0, :]
            pos_index_all = fine_local_grid.view(-1, 3)
            # yxz pos_mm_overlap      zyx pos_index
            pos1_mm_overlap = fine_grid[0, :, intersection_volume_sub1 > 0]
            pos2_mm_overlap = fine_grid[1, :, intersection_volume_sub2 > 0]
            view1_num_points = pos1_mm_overlap.shape[1]
            view2_num_points = pos2_mm_overlap.shape[1]

            if view1_num_points < self.train_cfg.pre_select_pos_number:
                points1_select = torch.randperm(view1_num_points)
            else:
                points1_select = torch.randperm(view1_num_points)[:self.train_cfg.pre_select_pos_number]
            if view2_num_points < self.train_cfg.pre_select_pos_number:
                points2_select = torch.randperm(view2_num_points)
            else:
                points2_select = torch.randperm(view2_num_points)[:self.train_cfg.pre_select_pos_number]

            pos1_use = pos1_mm_overlap[:, points1_select]
            pos2_use = pos2_mm_overlap[:, points2_select]

            pos1_mm_all = fine_grid[0, :, :, :, :].view(3, -1)
            pos2_mm_all = fine_grid[1, :, :, :, :].view(3, -1)

            # time1 = time.time()

            with torch.no_grad():
                # dist1 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                #                  pos2_mm_center.view(-1, 1, pos2_mm_center.shape[1])), dim=0)
                # dist2 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                #                  pos1_mm_center.view(-1, 1, pos1_mm_center.shape[1])), dim=0)
                dist1_1 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                                   pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)

                dist1_2 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                                   pos2_mm_all.view(-1, 1, pos2_mm_all.shape[1])), dim=0)
                dist2_1 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                                   pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)
                dist2_2 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                                   pos2_mm_all.view(-1, 1, pos2_mm_all.shape[1])), dim=0)

            # time2 = time.time()
            # print('dist time:',time2-time1)

            pos1 = torch.unique(torch.where(dist1_2 < self.train_cfg.positive_distance)[0])
            pos2 = torch.unique(torch.where(dist2_1 < self.train_cfg.positive_distance)[0])

            if pos1.shape[0] == 0:
                print(dist1_2.min())
                pos1 = torch.unique(torch.where(dist1_2 < dist1_2.min() + 0.5)[0])
                pos2 = torch.unique(torch.where(dist2_1 < dist2_1.min() + 0.5)[0])
            if pos1.shape[0] <= self.train_cfg.after_select_pos_number:
                pos1 = pos1
            else:
                pos1 = pos1[torch.randperm(pos1.shape[0])[:self.train_cfg.after_select_pos_number]]

            if pos2.shape[0] <= self.train_cfg.after_select_pos_number:
                pos2 = pos2
            else:
                pos2 = pos2[torch.randperm(pos2.shape[0])[:self.train_cfg.after_select_pos_number]]

            dist1_1 = dist1_1[pos1, :]
            dist1_2 = dist1_2[pos1, :]
            dist2_1 = dist2_1[pos2, :]
            dist2_2 = dist2_2[pos2, :]

            pos1_source = points1_select[pos1]
            pos2_source = points2_select[pos2]
            pos1_target = dist1_2.min(dim=1)[1]
            pos2_target = dist2_1.min(dim=1)[1]

            ind_ignores_1_1 = torch.where(dist1_1 < self.train_cfg.ignore_distance)
            ind_ignores_1_2 = torch.where(dist1_2 < self.train_cfg.ignore_distance)
            ind_ignores_2_1 = torch.where(dist2_1 < self.train_cfg.ignore_distance)
            ind_ignores_2_2 = torch.where(dist2_2 < self.train_cfg.ignore_distance)
            ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
            ind_ignores_1_2 = torch.stack(ind_ignores_1_2)
            ind_ignores_2_1 = torch.stack(ind_ignores_2_1)
            ind_ignores_2_2 = torch.stack(ind_ignores_2_2)

            use_mask_1_1 = torch.ones_like(dist1_1)
            use_mask_1_2 = torch.ones_like(dist1_2)
            use_mask_2_1 = torch.ones_like(dist2_1)
            use_mask_2_2 = torch.ones_like(dist2_2)
            use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0
            use_mask_1_2[ind_ignores_1_2[0, :], ind_ignores_1_2[1, :]] = 0
            use_mask_2_1[ind_ignores_2_1[0, :], ind_ignores_2_1[1, :]] = 0
            use_mask_2_2[ind_ignores_2_2[0, :], ind_ignores_2_2[1, :]] = 0

            neg_mask1 = torch.cat([use_mask_1_1, use_mask_1_2], dim=1)
            neg_mask2 = torch.cat([use_mask_2_1, use_mask_2_2], dim=1)

            coarse_feat_resample = F.interpolate(coarse_feat, fine_feat.shape[2:], mode='trilinear', align_corners=False)
            coarse_feat_resample = F.normalize(coarse_feat_resample, dim=1)
            view_1_coarse = coarse_feat_resample[0, :, :, :, :].view(128, -1)
            view_2_coarse = coarse_feat_resample[1, :, :, :, :].view(128, -1)

            q_view1_location = pos1_index_overlap[pos1_source].type(torch.LongTensor)
            k_view1_location = pos_index_all[pos1_target].type(torch.LongTensor)
            q_view2_location = pos2_index_overlap[pos2_source].type(torch.LongTensor)
            k_view2_location = pos_index_all[pos2_target].type(torch.LongTensor)

            q_view1_fine_feat = fine_feat[0, :, q_view1_location[:, 0], q_view1_location[:, 1],
                                q_view1_location[:, 2]].transpose(0, 1)
            k_view1_fine_feat = fine_feat[1, :, k_view1_location[:, 0], k_view1_location[:, 1],
                                k_view1_location[:, 2]].transpose(0, 1)
            q_view2_fine_feat = fine_feat[1, :, q_view2_location[:, 0], q_view2_location[:, 1],
                                q_view2_location[:, 2]].transpose(0, 1)
            k_view2_fine_feat = fine_feat[0, :, k_view2_location[:, 0], k_view2_location[:, 1],
                                k_view2_location[:, 2]].transpose(0, 1)

            q_view1_coarse_feat = coarse_feat_resample[0, :, q_view1_location[:, 0], q_view1_location[:, 1],
                                  q_view1_location[:, 2]].transpose(0, 1)
            q_view2_coarse_feat = coarse_feat_resample[1, :, q_view2_location[:, 0], q_view2_location[:, 1],
                                  q_view2_location[:, 2]].transpose(0, 1)

            inner_view1 = torch.einsum("nc,nc->n", q_view1_fine_feat, k_view1_fine_feat).view(-1, 1)
            inner_view2 = torch.einsum("nc,nc->n", q_view2_fine_feat, k_view2_fine_feat).view(-1, 1)

            neg_fine_view1 = torch.einsum("nc,ck->nk", q_view1_fine_feat, torch.cat((view_1_fine, view_2_fine), dim=1))
            neg_coarse_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                            torch.cat((view_1_coarse, view_2_coarse), dim=1))
            neg_fine_view2 = torch.einsum("nc,ck->nk", q_view2_fine_feat, torch.cat((view_1_fine, view_2_fine), dim=1))
            neg_coarse_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                            torch.cat((view_1_coarse, view_2_coarse), dim=1))

            # print('compute pos and neg time:',time4-time3)

            neg_all_view1 = neg_fine_view1 + neg_coarse_view1
            neg_all_view2 = neg_fine_view2 + neg_coarse_view2
            neg_all_view1 = neg_all_view1 * neg_mask1
            neg_all_view2 = neg_all_view2 * neg_mask2

            neg_candidate_view1_index = neg_all_view1.topk(self.train_cfg.pre_select_neg_number, dim=1)[1]
            neg_candidate_view2_index = neg_all_view2.topk(self.train_cfg.pre_select_neg_number, dim=1)[
                1]  # find diverse and hard negatives

            neg_use_view1 = torch.zeros((q_view1_fine_feat.shape[0], self.train_cfg.after_select_neg_number),
                                        device=neg_all_view1.device)
            neg_use_view2 = torch.zeros((q_view2_fine_feat.shape[0], self.train_cfg.after_select_neg_number),
                                        device=neg_all_view2.device)
            for i in range(q_view1_fine_feat.shape[0]):
                use_index = neg_candidate_view1_index[i, torch.randperm(neg_candidate_view1_index[i, :].shape[0])[
                                                         :self.train_cfg.after_select_neg_number]]
                neg_use_view1[i, :] = neg_fine_view1[i, use_index]

            for i in range(q_view2_fine_feat.shape[0]):
                use_index = neg_candidate_view2_index[i, torch.randperm(neg_candidate_view2_index[i, :].shape[0])[
                                                         :self.train_cfg.after_select_neg_number]]

                neg_use_view2[i, :] = neg_fine_view2[i, use_index]
            # neg_use_view1 = neg_use_view1.to(inner_view1.device)
            # neg_use_view2 = neg_use_view2.to(inner_view2.device)

            logits_view1 = torch.cat([inner_view1, neg_use_view1], dim=1)
            logits_view2 = torch.cat([inner_view2, neg_use_view2], dim=1)
            logits = torch.cat([logits_view1, logits_view2], dim=0)
            logits = logits / self.train_cfg.temperature

            labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
            loss = self.criterion(logits, labels)
        out['loss'] = loss
        return out

    def coarse_loss(self, coarse_feat, coarse_grid, coarse_valid, *args):  # coarse_grid[B,C,D,H,W] C:y,x,z
        coarse_feat_flatten = coarse_feat.view(coarse_feat.shape[0], 128, -1)
        coarse_feat_flatten = coarse_feat_flatten.permute(1, 0, 2)
        coarse_feat_flatten = coarse_feat_flatten.reshape(128, -1)

        out = dict()

        for i in range(int(coarse_feat.shape[0] / 2)):
            view1_feat_coarse = coarse_feat[2 * i, :, :, :, :].view(128, -1)
            view2_feat_coarse = coarse_feat[2 * i + 1, :, :, :, :].view(128, -1)
            view1_loc_coarse = coarse_grid[2 * i, :, :, :, :].view(3, -1)
            view2_loc_coarse = coarse_grid[2 * i + 1, :, :, :, :].view(3, -1)

            global_use_index_list = [j + 2 * i * view1_feat_coarse.shape[1] for j in
                                     range(0, 2 * view1_feat_coarse.shape[1])]
            global_index_list = [j for j in range(0, coarse_feat.shape[0] * view1_feat_coarse.shape[1])]
            global_search_index_list = list(set(global_index_list) - set(global_use_index_list))
            global_search_index = torch.tensor(global_search_index_list).to(view1_feat_coarse.device)

            with torch.no_grad():
                dist = LA.norm((view1_loc_coarse.view(-1, view1_loc_coarse.shape[1], 1) - \
                                view2_loc_coarse.view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
                dist_z_1 = LA.norm((view1_loc_coarse[2, :].view(-1, view1_loc_coarse.shape[1], 1) - \
                                    view2_loc_coarse[2, :].view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
                dist_z_2 = LA.norm((view2_loc_coarse[2, :].view(-1, view2_loc_coarse.shape[1], 1) - \
                                    view1_loc_coarse[2, :].view(-1, 1, view1_loc_coarse.shape[1])), dim=0)
                dist_1_1 = LA.norm((view1_loc_coarse.view(-1, view1_loc_coarse.shape[1], 1) - \
                                    view1_loc_coarse.view(-1, 1, view1_loc_coarse.shape[1])), dim=0)
                dist_2_2 = LA.norm((view2_loc_coarse.view(-1, view2_loc_coarse.shape[1], 1) - \
                                    view2_loc_coarse.view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
            pos_view1_coarse_can = torch.unique(torch.where(dist < self.train_cfg.coarse_positive_distance)[0])

            key_view1_coarse_can = dist[pos_view1_coarse_can].min(dim=1)[1]
            pos_view1_coarse = pos_view1_coarse_can[dist_z_1[pos_view1_coarse_can, key_view1_coarse_can] < 8]
            key_view1_coarse = key_view1_coarse_can[dist_z_1[pos_view1_coarse_can, key_view1_coarse_can] < 8]
            pos_view2_coarse_can = torch.unique(
                torch.where(dist.transpose(0, 1) < self.train_cfg.coarse_positive_distance)[0])
            key_view2_coarse_can = dist.transpose(0, 1)[pos_view2_coarse_can].min(dim=1)[1]
            pos_view2_coarse = pos_view2_coarse_can[
                dist_z_2[pos_view2_coarse_can, key_view2_coarse_can] < self.train_cfg.coarse_z_thres]
            key_view2_coarse = key_view2_coarse_can[
                dist_z_2[pos_view2_coarse_can, key_view2_coarse_can] < self.train_cfg.coarse_z_thres]

            if pos_view1_coarse.shape[0] > 0:
                q_view1_coarse_feat = view1_feat_coarse[:, pos_view1_coarse].transpose(0, 1)
                k_view1_coarse_feat = view2_feat_coarse[:, key_view1_coarse].transpose(0, 1)
                q_view2_coarse_feat = view2_feat_coarse[:, pos_view2_coarse].transpose(0, 1)
                k_view2_coarse_feat = view1_feat_coarse[:, key_view2_coarse].transpose(0, 1)

                # all_nodes_view1 = []
                # all_nodes_view2 = []
                dist_1_2 = dist.clone()
                dist_2_1 = dist.transpose(0, 1).clone()
                dist_1_2 = dist_1_2[pos_view1_coarse, :]
                dist_1_1 = dist_1_1[pos_view1_coarse, :]
                dist_2_1 = dist_2_1[pos_view2_coarse, :]
                dist_2_2 = dist_2_2[pos_view2_coarse, :]

                ind_ignores_1_1 = torch.where(dist_1_1 < self.train_cfg.coarse_ignore_distance)
                ind_ignores_1_2 = torch.where(dist_1_2 < self.train_cfg.coarse_ignore_distance)
                ind_ignores_2_1 = torch.where(dist_2_1 < self.train_cfg.coarse_ignore_distance)
                ind_ignores_2_2 = torch.where(dist_2_2 < self.train_cfg.coarse_ignore_distance)
                ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
                ind_ignores_1_2 = torch.stack(ind_ignores_1_2)
                ind_ignores_2_1 = torch.stack(ind_ignores_2_1)
                ind_ignores_2_2 = torch.stack(ind_ignores_2_2)

                use_mask_1_1 = torch.ones_like(dist_1_1)
                use_mask_1_2 = torch.ones_like(dist_1_2)
                use_mask_2_1 = torch.ones_like(dist_2_1)
                use_mask_2_2 = torch.ones_like(dist_2_2)

                use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0
                use_mask_1_2[ind_ignores_1_2[0, :], ind_ignores_1_2[1, :]] = 0
                use_mask_2_1[ind_ignores_2_1[0, :], ind_ignores_2_1[1, :]] = 0
                use_mask_2_2[ind_ignores_2_2[0, :], ind_ignores_2_2[1, :]] = 0

                neg_mask1 = torch.cat([use_mask_1_1, use_mask_1_2], dim=1)
                neg_mask2 = torch.cat([use_mask_2_1, use_mask_2_2], dim=1)

                inner_view1 = torch.einsum("nc,nc->n", q_view1_coarse_feat, k_view1_coarse_feat).view(-1, 1)
                inner_view2 = torch.einsum("nc,nc->n", q_view2_coarse_feat, k_view2_coarse_feat).view(-1, 1)

                neg_coarse_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                                torch.cat((view1_feat_coarse, view2_feat_coarse), dim=1))
                neg_coarse_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                                torch.cat((view1_feat_coarse, view2_feat_coarse), dim=1))

                neg_coarse_view1 = neg_coarse_view1 * neg_mask1
                neg_coarse_view2 = neg_coarse_view2 * neg_mask2

                neg_candidate_view1_coarse_index = \
                    neg_coarse_view1.topk(self.train_cfg.coarse_pre_select_neg_number, dim=1)[1]
                neg_candidate_view2_coarse_index = \
                    neg_coarse_view2.topk(self.train_cfg.coarse_pre_select_neg_number, dim=1)[1]

                neg_use_view1_coarse = torch.zeros(
                    (q_view1_coarse_feat.shape[0], self.train_cfg.coarse_after_select_neg_number))
                neg_use_view2_coarse = torch.zeros(
                    (q_view2_coarse_feat.shape[0], self.train_cfg.coarse_after_select_neg_number))

                for i in range(q_view1_coarse_feat.shape[0]):
                    use_index = neg_candidate_view1_coarse_index[
                                    i, torch.randperm(neg_candidate_view1_coarse_index.shape[1])][
                                :self.train_cfg.coarse_after_select_neg_number]
                    neg_use_view1_coarse[i, :use_index.shape[0]] = neg_coarse_view1[i, use_index]

                for i in range(q_view2_coarse_feat.shape[0]):
                    use_index = neg_candidate_view2_coarse_index[
                                    i, torch.randperm(neg_candidate_view2_coarse_index.shape[1])][
                                :self.train_cfg.coarse_after_select_neg_number]
                    neg_use_view2_coarse[i, :use_index.shape[0]] = neg_coarse_view2[i, use_index]

                neg_coarse_global_index_1 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[:self.train_cfg.coarse_global_select_number]]
                neg_coarse_global_1 = coarse_feat_flatten[:, neg_coarse_global_index_1]
                neg_coarse_global_index_2 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[:self.train_cfg.coarse_global_select_number]]
                neg_coarse_global_2 = coarse_feat_flatten[:, neg_coarse_global_index_2]

                neg_coarse_global_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                                       neg_coarse_global_1)
                neg_coarse_global_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                                       neg_coarse_global_2)

                neg_view1_coarse_concate = torch.cat(
                    [neg_use_view1_coarse.to(view1_feat_coarse.device), neg_coarse_global_view1], dim=1)
                neg_view2_coarse_concate = torch.cat(
                    [neg_use_view2_coarse.to(view1_feat_coarse.device), neg_coarse_global_view2], dim=1)

                logits_view1 = torch.cat([inner_view1, neg_view1_coarse_concate], dim=1)
                logits_view2 = torch.cat([inner_view2, neg_view2_coarse_concate], dim=1)
                logits = torch.cat([logits_view1, logits_view2], dim=0)
                logits = logits / self.train_cfg.temperature
                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
                loss = self.criterion(logits, labels)
                if out.keys().__len__() == 0:
                    out['loss'] = loss
                else:
                    out['loss'] += loss
            else:
                neg_coarse_global_index_1 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[:self.train_cfg.coarse_global_select_number]]
                neg_coarse_global_1 = coarse_feat_flatten[:, neg_coarse_global_index_1]
                neg_coarse_global_index_2 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[:self.train_cfg.coarse_global_select_number]]
                neg_coarse_global_2 = coarse_feat_flatten[:, neg_coarse_global_index_2]

                neg_coarse_global_view1 = torch.einsum("nc,ck->nk", view1_feat_coarse.transpose(0, 1),
                                                       neg_coarse_global_1)
                neg_coarse_global_view2 = torch.einsum("nc,ck->nk", view2_feat_coarse.transpose(0, 1),
                                                       neg_coarse_global_2)

                logits_view1 = torch.einsum("nc,ck->nk", view1_feat_coarse.transpose(0, 1),
                                            torch.cat([view1_feat_coarse, view2_feat_coarse], dim=1)).topk(
                    self.train_cfg.coarse_after_select_neg_number + 1,
                    dim=1)[
                    0]
                logits_view2 = torch.einsum("nc,ck->nk", view2_feat_coarse.transpose(0, 1),
                                            torch.cat([view1_feat_coarse, view2_feat_coarse], dim=1)).topk(
                    self.train_cfg.coarse_after_select_neg_number + 1,
                    dim=1)[
                    0]

                logits_view1 = torch.cat([logits_view1, neg_coarse_global_view1], dim=1)
                logits_view2 = torch.cat([logits_view2, neg_coarse_global_view2], dim=1)
                logits = torch.cat([logits_view1, logits_view2], dim=0)
                logits = logits / self.train_cfg.temperature
                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
                loss = self.criterion(logits, labels)
                if out.keys().__len__() == 0:
                    out['loss'] = loss
                else:
                    out['loss'] += loss

        out['loss'] = out['loss'] / (coarse_feat.shape[0] / 2)
        return out

    def meshgrid3d(self, shape):
        z_ = torch.linspace(0., shape[0] - 1, shape[0])
        y_ = torch.linspace(0., shape[1] - 1, shape[1])
        x_ = torch.linspace(0., shape[2] - 1, shape[2])
        z, y, x = torch.meshgrid(z_, y_, x_)
        return torch.stack((z, y, x), 3)

    async def async_simple_test(self,
                                img,
                                img_meta,
                                proposals=None,
                                rescale=False):
        """Async test without augmentation."""
        assert self.with_bbox, 'Bbox head must be implemented.'
        x = self.extract_feat(img)

        if proposals is None:
            proposal_list = await self.rpn_head.async_simple_test_rpn(
                x, img_meta)
        else:
            proposal_list = proposals

        return await self.roi_head.async_simple_test(
            x, proposal_list, img_meta, rescale=rescale)

    def simple_test(self, img, img_metas, proposals=None, rescale=False):
        """Test without augmentation."""
        x = self.extract_feat(img)
        outs = []
        out1 = x[0].data.cpu().numpy()
        out2 = x[1].data.cpu().numpy()
        outs = [out1, out2, img.data.cpu().numpy(), img_metas[0]['filename'].split('.', 1)[0]]
        output_embedding = self.test_cfg.get('output_embedding', True)
        if not output_embedding:
            if not os.path.exists(self.test_cfg.save_path):
                os.mkdir(self.test_cfg.save_path)
            outfilename = self.test_cfg.save_path + \
                          img_metas[0]['filename'].split('.', 1)[0] + '.pkl'
            f = open(outfilename, 'wb')
            pickle.dump(outs, f)
            return [
                x[0][0, 0, 0, 0, 0].data.cpu()]  # we have saved the data into harddisk, this is just for fit the code
        else:
            return outs

    def aug_test(self, imgs, img_metas, **kwargs):
        return self.simple_test(imgs, img_metas, **kwargs)

sam/models/frameworks/sam.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai
# rearrange the ignore computation into useindex, save computation cost
import os.path
import warnings

import torch
from mmdet.models.builder import DETECTORS, build_backbone, build_head, build_neck
from mmdet.models.detectors.base import BaseDetector
import torch.nn.functional as F
from torch import linalg as LA
import time
import pickle
import numpy as np

import math
import matplotlib.pyplot as plt
from ..losses import SupConLoss
import random


@DETECTORS.register_module()
class Sam_uniform_cross_volume_dual_head_mri_ct(BaseDetector):
    def __init__(self,
                 backbone,
                 neck=None,
                 read_out_head=None,
                 sem_neck=None,
                 train_cfg=None,
                 test_cfg=None,
                 init_cfg=None):
        super(Sam_uniform_cross_volume_dual_head_mri_ct, self).__init__(init_cfg)
        self.backbone = build_backbone(backbone)
        self.backbone.init_weights()
        if neck is not None:
            self.neck = build_neck(neck)
        self.criterion = torch.nn.CrossEntropyLoss().cuda()
        self.supcriterion = SupConLoss()

        self.read_out_head = build_neck(read_out_head)
        if sem_neck is not None:
            self.semantic_head = build_neck(sem_neck)
        else:
            self.semantic_head = build_neck(neck)
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

    def extract_feat(self, img, normalize=True):
        """Directly extract features from the backbone+neck."""
        x = self.backbone(img)
        out1 = self.neck(x[:self.neck.end_level])[0]
        out2 = self.read_out_head([x[-1]])[0]
        out3 = self.semantic_head(x[:self.semantic_head.end_level])[0]
        if normalize:
            out1 = F.normalize(out1, dim=1)
            out2 = F.normalize(out2, dim=1)
            out3 = F.normalize(out3, dim=1)
        out1 = out1.type(torch.half)
        out2 = out2.type(torch.half)
        out3 = out3.type(torch.half)
        return [out1, out2, out3]

    def forward_dummy(self, img):
        """Used for computing network flops.

        See `mmdetection/tools/analysis_tools/get_flops.py`
        """
        outs = ()
        # backbone
        x = self.extract_feat(img)
        outs = outs + x

        return outs

    def forward_train(self,
                      img,
                      img_metas,
                      **kwargs):

        """
        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.

            img_metas (list[dict]): list of image info dict where each dict
                has: 'img_shape', 'img_name',etc.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        all_intra_info = []
        all_intra_img = []
        all_inter_info = []
        all_inter_img = []

        p = np.random.rand()

        for i in range(img_metas.__len__()):
            im_info = img_metas[i][0]
            if im_info['style'] == 'intra':
                all_intra_info.append(im_info)
                all_intra_img.append(img[i])
            if im_info['style'] == 'inter':
                all_inter_info.append(im_info)
                all_inter_img.append(img[i])

        intra_imgs = torch.stack(all_intra_img)
        mesh_grids = []
        valids = []
        for i in range(all_intra_info.__len__()):
            mesh = all_intra_info[i]['meshgrid']
            mesh_grids.append(mesh)
            valid = all_intra_info[i]['valid']
            valids.append(valid)
            all_intra_info[i].pop('meshgrid')
            all_intra_info[i].pop('valid')
        mesh_grids = torch.stack(mesh_grids).to(intra_imgs.device)
        valids = torch.stack(valids).to(intra_imgs.device)

        losses = dict()
        if p < self.train_cfg.meta_cfg.prob:
            inter_dict = {}
            for image, info in zip(all_inter_img, all_inter_info):
                pair_name = info['filename'].split('_', 1)[0]
                if pair_name not in inter_dict.keys():
                    inter_dict[pair_name] = {}
                if info['tag'] == 'CT':
                    inter_dict[pair_name]['CT'] = {}
                    inter_dict[pair_name]['CT']['filename'] = info['filename']
                    inter_dict[pair_name]['CT']['mask'] = info['mask']
                    inter_dict[pair_name]['CT']['meshgrid'] = info['meshgrid']
                    inter_dict[pair_name]['CT']['image'] = image
                    inter_dict[pair_name]['CT']['spacing'] =  info['resample_spacing']

                if info['tag'] == 'MRI':
                    inter_dict[pair_name]['MRI'] = {}
                    inter_dict[pair_name]['MRI']['filename'] = info['filename']
                    inter_dict[pair_name]['MRI']['mask'] = info['mask']
                    inter_dict[pair_name]['MRI']['meshgrid'] = info['meshgrid']
                    inter_dict[pair_name]['MRI']['image'] = image
                    inter_dict[pair_name]['MRI']['spacing'] = info['resample_spacing']

            for pair in inter_dict:
                mri_mask = inter_dict[pair]['MRI']['mask']
                mri_mesh = inter_dict[pair]['MRI']['meshgrid']
                ct_mask = inter_dict[pair]['CT']['mask']
                ct_mesh = inter_dict[pair]['CT']['meshgrid']
                mri_inds = torch.where(mri_mask == 1)
                mri_min = torch.tensor([ele.min() for ele in mri_inds])
                mri_max = torch.tensor([ele.max() for ele in mri_inds])
                mri_loc_z_y_x = torch.stack([mri_min, mri_max])[:, 1:].type(torch.float32).permute(1, 0)
                ct_inds = torch.where(ct_mask == 1)
                ct_min = torch.tensor([ele.min() for ele in ct_inds])
                ct_max = torch.tensor([ele.max() for ele in ct_inds])
                ct_loc_z_y_x = torch.stack([ct_min, ct_max])[:, 1:].type(torch.float32).permute(1, 0)

                # if mri_loc_z_y_x[0,1] - mri_loc_z_y_x[0,0] > 90:
                #     half =(mri_loc_z_y_x[0,1] - mri_loc_z_y_x[0,0])/2
                #     p_c = np.random.rand()
                #     if p_c < 0.5:
                #         mri_loc_z_y_x[0, 1] = mri_loc_z_y_x[0,0] + half
                #     else:
                #         mri_loc_z_y_x[0, 0] += half
                # spacing is y x z
                ct_spacing = inter_dict[pair]['CT']['spacing']
                mri_spacing =inter_dict[pair]['MRI']['spacing']
                # print(ct_spacing,mri_spacing)
                ratio = mri_spacing / ct_spacing
                ratio = torch.tensor([ratio[2], ratio[0], ratio[1]],dtype=mri_loc_z_y_x.dtype)
                trans_M = torch.zeros((3, 3), dtype=ratio.dtype).scatter_(0, torch.arange(0, 3).view(-1, 3),
                                                                          ratio.view(-1, 3))
                # print(trans_M)
                mri_shape = torch.tensor(mri_mask.shape[1:], dtype=torch.float32)
                ct_shape = torch.tensor(ct_mask.shape[1:], dtype=torch.float32)
                mri_ct_map_loc = torch.matmul(trans_M, mri_loc_z_y_x)
                mri_ct_map_loc = torch.where(mri_ct_map_loc < 0, torch.tensor(0, dtype=mri_ct_map_loc.dtype),
                                             mri_ct_map_loc)
                mri_ct_map_loc[:, 1] = torch.where(mri_ct_map_loc[:, 1] >= ct_shape - 1, ct_shape - 1,
                                                   mri_ct_map_loc[:, 1])
                mri_loc_z_y_x_remap = torch.matmul(torch.linalg.inv(trans_M), mri_ct_map_loc)
                margin = torch.tensor([104, 128, 128])
                possible_range = mri_loc_z_y_x_remap[:, 1] - mri_loc_z_y_x_remap[:, 0] - margin

                if possible_range[0] >= 1:
                    z_s = np.random.randint(0, int(possible_range[0]))
                    z_end = z_s + margin[0]
                else:
                    z_s = 0
                    z_end = int(mri_loc_z_y_x_remap[0, 1] - mri_loc_z_y_x_remap[0, 0])

                if possible_range[1] >= 1:
                    y_s = np.random.randint(0, int(possible_range[1]))
                    y_end = y_s + margin[1]
                else:
                    y_s = 0
                    y_end = int(mri_loc_z_y_x_remap[1, 1] - mri_loc_z_y_x_remap[1, 0])
                if possible_range[2] >= 1:
                    x_s = np.random.randint(0, int(possible_range[2]))
                    x_end = x_s + margin[2]
                else:
                    x_s = 0
                    x_end = int(mri_loc_z_y_x_remap[2, 1] - mri_loc_z_y_x_remap[2, 0])

                mri_crop_region = torch.tensor([z_s,z_end,y_s,y_end,x_s,x_end]).reshape(3,2)
                mri_crop_region = mri_crop_region +  mri_loc_z_y_x_remap[:,0:1]

                ct_crop_region = torch.matmul(trans_M, mri_crop_region)
                jitter_shift = 10
                jitter_ct_loc = torch.randint(-jitter_shift, jitter_shift, (3, 2))
                # print('before', mri_loc_z_y_x, mri_ct_map_loc, inter_dict[pair]['CT']['filename'])
                mri_ct_map_loc = ct_crop_region + jitter_ct_loc
                jitter_mri_loc = torch.randint(-jitter_shift, jitter_shift, (3, 2))
                mri_loc_z_y_x = mri_crop_region + jitter_mri_loc
                mri_ct_map_loc = torch.where(mri_ct_map_loc<0,torch.tensor(0,dtype=mri_ct_map_loc.dtype),mri_ct_map_loc)
                mri_loc_z_y_x = torch.where(mri_loc_z_y_x<0,torch.tensor(0,dtype=mri_loc_z_y_x.dtype),mri_loc_z_y_x)
                mri_loc_z_y_x[:,1] = torch.where(mri_loc_z_y_x[:, 1] >= mri_shape - 1, mri_shape - 1, mri_loc_z_y_x[:, 1])
                mri_ct_map_loc[:,1] = torch.where(mri_ct_map_loc[:, 1] >= ct_shape - 1, ct_shape - 1, mri_ct_map_loc[:, 1])
                mri_ct_map_loc = mri_ct_map_loc.type(torch.int32)
                mri_loc_z_y_x = mri_loc_z_y_x.type(torch.int32)
                inter_dict[pair]['MRI']['crop_info'] = mri_loc_z_y_x
                inter_dict[pair]['CT']['crop_info'] = mri_ct_map_loc
            all_inter_feats = []
            all_inter_meshgrid = []
            all_inter_mask = []
            for pair in inter_dict:
                CT_image = inter_dict[pair]['CT']['image']
                CT_crop_loc = inter_dict[pair]['CT']['crop_info']
                CT_mask = inter_dict[pair]['CT']['mask']
                CT_mesh = inter_dict[pair]['CT']['meshgrid']
                MRI_image = inter_dict[pair]['MRI']['image']
                MRI_crop_loc = inter_dict[pair]['MRI']['crop_info']
                MRI_mask = inter_dict[pair]['MRI']['mask']
                MRI_mesh = inter_dict[pair]['MRI']['meshgrid']
                feat = self.extract_feat(
                    CT_image[:, CT_crop_loc[0, 0]:CT_crop_loc[0, 1], CT_crop_loc[1, 0]:CT_crop_loc[1, 1],
                    CT_crop_loc[2, 0]:CT_crop_loc[2, 1]].unsqueeze(0))
                all_inter_feats.append(feat)
                all_inter_meshgrid.append(
                    CT_mesh[:, CT_crop_loc[0, 0]:CT_crop_loc[0, 1], CT_crop_loc[1, 0]:CT_crop_loc[1, 1],
                    CT_crop_loc[2, 0]:CT_crop_loc[2, 1]])
                all_inter_mask.append(
                    CT_mask[:, CT_crop_loc[0, 0]:CT_crop_loc[0, 1], CT_crop_loc[1, 0]:CT_crop_loc[1, 1],
                    CT_crop_loc[2, 0]:CT_crop_loc[2, 1]])
                feat = self.extract_feat(
                    MRI_image[:, MRI_crop_loc[0, 0]:MRI_crop_loc[0, 1], MRI_crop_loc[1, 0]:MRI_crop_loc[1, 1],
                    MRI_crop_loc[2, 0]:MRI_crop_loc[2, 1]].unsqueeze(0))
                all_inter_feats.append(feat)
                all_inter_meshgrid.append(
                    MRI_mesh[:, MRI_crop_loc[0, 0]:MRI_crop_loc[0, 1], MRI_crop_loc[1, 0]:MRI_crop_loc[1, 1],
                    MRI_crop_loc[2, 0]:MRI_crop_loc[2, 1]])
                all_inter_mask.append(
                    MRI_mask[:, MRI_crop_loc[0, 0]:MRI_crop_loc[0, 1], MRI_crop_loc[1, 0]:MRI_crop_loc[1, 1],
                    MRI_crop_loc[2, 0]:MRI_crop_loc[2, 1]])

            loss = self.loss_inter(all_inter_feats, all_inter_meshgrid, all_inter_mask)
        else:
            feats = self.extract_feat(intra_imgs)
            loss = self.loss_intra(feats, mesh_grids, all_intra_info, valids)
        losses['loss'] = loss
        return losses

    def loss_intra(self, feats, meshgrid, img_metas, valid, *args):

        N, C, D, H, W = meshgrid.shape
        output_half = F.interpolate(meshgrid, size=(int(D / 2), int(H / 2), int(W / 2)), mode='trilinear')
        output_416 = F.interpolate(meshgrid, size=(int(D / 4), int(H / 16), int(W / 16)), mode='trilinear')
        valid_half = F.interpolate(valid, size=(int(D / 2), int(H / 2), int(W / 2)))
        valid_416 = F.interpolate(valid, size=(int(D / 4), int(H / 16), int(W / 16)))

        output_half = output_half.type(torch.half)
        output_416 = output_416.type(torch.half)
        valid_half = valid_half.type(torch.half)
        valid_416 = valid_416.type(torch.half)

        # fine_time_start = time.time()

        for i in range(int(N / 2)):
            result = self.single_intra_fine_loss(feats[0][2 * i:2 * (i + 1)], feats[1][2 * i:2 * (i + 1)],
                                                 output_half[2 * i:2 * (i + 1)], output_416[2 * i:2 * (i + 1)],
                                                 valid_half[2 * i:2 * (i + 1)], valid_416[2 * i:2 * (i + 1)])
            if i == 0:
                fine_loss = result['loss']
            else:
                fine_loss += result['loss']
        fine_loss = fine_loss / int(N / 2)
        result = self.single_intra_coarse_loss(feats[1], output_416, valid_416)
        coarse_loss = result['loss']

        loss = fine_loss + coarse_loss
        return loss

    def loss_inter(self, feats, meshs, masks):
        N = feats.__len__()

        for i in range(int(N / 2)):
            result_fine = self.single_inter_fine_loss([feats[2 * i], feats[2 * i + 1]],
                                                 [meshs[2 * i], meshs[2 * i + 1]], [masks[2 * i], masks[2 * i + 1]])
            if i == 0:
                fine_loss = result_fine['loss']
            else:
                fine_loss += result_fine['loss']
        fine_loss = fine_loss / int(N / 2)
        for i in range(int(N / 2)):
            result_coarse = self.single_inter_coarse_loss([feats[2 * i], feats[2 * i + 1]],
                                                   [meshs[2 * i], meshs[2 * i + 1]], [masks[2 * i], masks[2 * i + 1]])
            if i == 0:
                coarse_loss = result_coarse['loss']
            else:
                coarse_loss += result_coarse['loss']
        coarse_loss = coarse_loss / int(N / 2)
        loss = fine_loss + coarse_loss
        return loss

    def single_inter_fine_loss(self, feat, mesh, mask):
        out = dict()
        ct_fine_feat = feat[0][0]
        ct_coarse_feat = feat[0][1]
        ct_mesh = mesh[0].to(ct_fine_feat.device)
        ct_mask = mask[0].to(ct_fine_feat.device)
        ct_local_grid = self.meshgrid3d(ct_fine_feat.shape[2:]).to(ct_fine_feat.device)
        mri_fine_feat = feat[1][0]
        mri_coarse_feat = feat[1][1]
        mri_mesh = mesh[1].to(mri_fine_feat.device)
        mri_mask = mask[1].to(mri_fine_feat.device)
        mri_local_grid = self.meshgrid3d(mri_fine_feat.shape[2:])

        ct_fine_mesh = F.interpolate(ct_mesh.unsqueeze(0), size=ct_fine_feat.shape[2:], mode='trilinear').type(
            torch.half)
        ct_coarse_mesh = F.interpolate(ct_mesh.unsqueeze(0), size=ct_coarse_feat.shape[2:], mode='trilinear').type(
            torch.half)
        mri_fine_mesh = F.interpolate(mri_mesh.unsqueeze(0), size=mri_fine_feat.shape[2:], mode='trilinear').type(
            torch.half)
        mri_coarse_mesh = F.interpolate(mri_mesh.unsqueeze(0), size=mri_coarse_feat.shape[2:], mode='trilinear').type(
            torch.half)
        mri_fine_mask = F.interpolate(mri_mask.type(mri_mesh.dtype).unsqueeze(0), size=mri_fine_feat.shape[2:]).type(
            torch.half)

        mri_fine_foregrond_points = mri_fine_mesh[:, :, mri_fine_mask[0, 0, :, :, :] == 1].squeeze(0)
        mri_fine_foregrond_points_loc = mri_local_grid[mri_fine_mask[0, 0, :, :, :] == 1, :]
        ct_fine_points_loc = ct_local_grid.view(-1, 3)
        mri_fine_points_loc = mri_local_grid.view(-1,3)

        points_select = torch.randperm(mri_fine_foregrond_points.shape[1])[
                        :self.train_cfg.inter_cfg.pre_select_pos_number]

        mri_fine_select_points = mri_fine_foregrond_points[:, points_select]
        mri_fine_all_points = mri_fine_mesh.view(3, -1)
        ct_fine_all_points = ct_fine_mesh.view(3, -1)

        with torch.no_grad():
            mri_ct_dist = LA.norm((mri_fine_select_points.view(-1, mri_fine_select_points.shape[1], 1) - \
                                   ct_fine_all_points.view(-1, 1, ct_fine_all_points.shape[1])), dim=0)
            mri_mri_dist = LA.norm((mri_fine_select_points.view(-1, mri_fine_select_points.shape[1], 1) - \
                                    mri_fine_all_points.view(-1, 1, mri_fine_all_points.shape[1])), dim=0)

        pos_match = torch.unique(torch.where(mri_ct_dist < self.train_cfg.inter_cfg.positive_distance)[0])

        if pos_match.shape[0] == 0:
            pos_match = torch.unique(torch.where(mri_ct_dist < mri_ct_dist.min() + 0.5)[0])
        if pos_match.shape[0] <= self.train_cfg.inter_cfg.after_select_pos_number:
            pos_match = pos_match
        else:
            pos_match = pos_match[torch.randperm(pos_match.shape[0])[:self.train_cfg.intra_cfg.after_select_pos_number]]

        mri_ct_dist = mri_ct_dist[pos_match, :]
        mri_mri_dist = mri_mri_dist[pos_match, :]

        mri_source = points_select[pos_match]
        ct_target = mri_ct_dist.min(dim=1)[1]

        mri_mri_ignore = torch.where(mri_mri_dist < self.train_cfg.inter_cfg.ignore_distance)
        mri_ct_ignore = torch.where(mri_ct_dist < self.train_cfg.inter_cfg.ignore_distance)
        mri_mri_ignore = torch.stack(mri_mri_ignore)
        mri_ct_ignore = torch.stack(mri_ct_ignore)

        mri_mri_mask = torch.ones_like(mri_mri_dist)
        mri_ct_mask = torch.ones_like(mri_ct_dist)

        mri_mri_mask[mri_mri_ignore[0, :], mri_mri_ignore[1, :]] = 0
        mri_ct_mask[mri_ct_ignore[0, :], mri_ct_ignore[1, :]] = 0

        neg_mask_1 = torch.cat([mri_mri_mask, mri_ct_mask], dim=1)

        mri_coarse_feat_resample = F.interpolate(mri_coarse_feat, mri_fine_feat.shape[2:], mode='trilinear')
        ct_coarse_feat_resample = F.interpolate(ct_coarse_feat, ct_fine_feat.shape[2:], mode='trilinear')
        mri_coarse_feat_resample = F.normalize(mri_coarse_feat_resample, dim=1)
        ct_coarse_feat_resample = F.normalize(ct_coarse_feat_resample, dim=1)

        mri_location = mri_fine_foregrond_points_loc[mri_source].type(torch.LongTensor)
        ct_location = ct_fine_points_loc[ct_target].type(torch.LongTensor)


        ct_fine_selection_points = ct_fine_mesh[0,:,ct_location[:,0],ct_location[:,1],ct_location[:,2]]

        with torch.no_grad():
            ct_mri_dist = LA.norm((ct_fine_selection_points.view(-1, ct_fine_selection_points.shape[1], 1) - \
                                   mri_fine_all_points.view(-1, 1, mri_fine_all_points.shape[1])), dim=0)
            ct_ct_dist = LA.norm((ct_fine_selection_points.view(-1, ct_fine_selection_points.shape[1], 1) - \
                                    ct_fine_all_points.view(-1, 1, ct_fine_all_points.shape[1])), dim=0)

        ct_ct_ignore = torch.where(ct_ct_dist < self.train_cfg.inter_cfg.ignore_distance)
        ct_mri_ignore = torch.where(ct_mri_dist < self.train_cfg.inter_cfg.ignore_distance)
        ct_ct_ignore = torch.stack(ct_ct_ignore)
        ct_mri_ignore = torch.stack(ct_mri_ignore)

        ct_ct_mask = torch.ones_like(ct_ct_dist)
        ct_mri_mask = torch.ones_like(ct_mri_dist)

        ct_ct_mask[ct_ct_ignore[0, :], ct_ct_ignore[1, :]] = 0
        ct_mri_mask[ct_mri_ignore[0, :], ct_mri_ignore[1, :]] = 0

        neg_mask_2 = torch.cat([ct_mri_mask,ct_ct_mask], dim=1)

        q_mri_fine_feat = mri_fine_feat[0, :, mri_location[:, 0], mri_location[:, 1],
                          mri_location[:, 2]].transpose(0, 1)
        k_ct_fine_feat = ct_fine_feat[0, :, ct_location[:, 0], ct_location[:, 1],
                         ct_location[:, 2]].transpose(0, 1)

        q_mri_coarse_feat = mri_coarse_feat_resample[0, :, mri_location[:, 0], mri_location[:, 1],
                            mri_location[:, 2]].transpose(0, 1)
        k_ct_coarse_feat = ct_coarse_feat_resample[0, :, ct_location[:, 0], ct_location[:, 1],
                         ct_location[:, 2]].transpose(0, 1)

        inner_view = torch.einsum("nc,nc->n", q_mri_fine_feat, k_ct_fine_feat).view(-1, 1)

        neg_fine_view_1 = torch.einsum("nc,ck->nk", q_mri_fine_feat, torch.cat(
            (mri_fine_feat[0, :, :, :, :].view(128, -1), ct_fine_feat[0, :, :, :, :].view(128, -1)), dim=1))
        neg_coarse_view_1 = torch.einsum("nc,ck->nk", q_mri_coarse_feat,
                                       torch.cat((mri_coarse_feat_resample[0, :, :, :, :].view(128, -1),
                                                  ct_coarse_feat_resample[0, :, :, :, :].view(128, -1)), dim=1))

        # print('compute pos and neg time:',time4-time3)

        neg_all_view_1 = neg_fine_view_1 + neg_coarse_view_1
        neg_all_view_1 = neg_all_view_1 * neg_mask_1

        neg_candidate_view_index_1 = neg_all_view_1.topk(self.train_cfg.inter_cfg.pre_select_neg_number, dim=1)[1]

        neg_use_view_1 = torch.zeros((q_mri_fine_feat.shape[0], self.train_cfg.inter_cfg.after_select_neg_number),
                                   device=neg_all_view_1.device)

        for i in range(q_mri_fine_feat.shape[0]):
            use_index = neg_candidate_view_index_1[i, torch.randperm(neg_candidate_view_index_1[i, :].shape[0])[
                                                    :self.train_cfg.inter_cfg.after_select_neg_number]]
            neg_use_view_1[i, :] = neg_all_view_1[i, use_index]

        logits_view_1 = torch.cat([inner_view, neg_use_view_1], dim=1)

        neg_fine_view_2 = torch.einsum("nc,ck->nk", k_ct_fine_feat, torch.cat(
            (mri_fine_feat[0, :, :, :, :].view(128, -1), ct_fine_feat[0, :, :, :, :].view(128, -1)), dim=1))
        neg_coarse_view_2 = torch.einsum("nc,ck->nk", k_ct_coarse_feat,
                                         torch.cat((mri_coarse_feat_resample[0, :, :, :, :].view(128, -1),
                                                    ct_coarse_feat_resample[0, :, :, :, :].view(128, -1)), dim=1))

        # print('compute pos and neg time:',time4-time3)

        neg_all_view_2 = neg_fine_view_2 + neg_coarse_view_2
        neg_all_view_2 = neg_all_view_2 * neg_mask_2

        neg_candidate_view_index_2 = neg_all_view_2.topk(self.train_cfg.inter_cfg.pre_select_neg_number, dim=1)[1]

        neg_use_view_2 = torch.zeros((k_ct_fine_feat.shape[0], self.train_cfg.inter_cfg.after_select_neg_number),
                                     device=neg_all_view_2.device)

        for i in range(q_mri_fine_feat.shape[0]):
            use_index = neg_candidate_view_index_2[i, torch.randperm(neg_candidate_view_index_2[i, :].shape[0])[
                                                      :self.train_cfg.inter_cfg.after_select_neg_number]]
            neg_use_view_2[i, :] = neg_fine_view_2[i, use_index]

        logits_view_2 = torch.cat([inner_view, neg_use_view_2], dim=1)




        logits = torch.cat([logits_view_1, logits_view_2], dim=0)
        logits = logits / self.train_cfg.inter_cfg.fine_temperature

        labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
        loss = self.criterion(logits, labels)
        out['loss'] = loss
        return out

    def single_inter_coarse_loss(self, feat, mesh, mask):
        out = dict()
        ct_fine_feat = feat[0][0]
        ct_coarse_feat = feat[0][1]
        ct_mesh = mesh[0].to(ct_fine_feat.device)
        ct_mask = mask[0].to(ct_fine_feat.device)
        ct_local_grid = self.meshgrid3d(ct_fine_feat.shape[2:]).to(ct_fine_feat.device)
        mri_fine_feat = feat[1][0]
        mri_coarse_feat = feat[1][1]
        mri_mesh = mesh[1].to(mri_fine_feat.device)
        mri_mask = mask[1].to(mri_fine_feat.device)
        mri_local_grid = self.meshgrid3d(mri_fine_feat.shape[2:])

        ct_fine_mesh = F.interpolate(ct_mesh.unsqueeze(0), size=ct_fine_feat.shape[2:], mode='trilinear').type(
            torch.half)
        ct_coarse_mesh = F.interpolate(ct_mesh.unsqueeze(0), size=ct_coarse_feat.shape[2:], mode='trilinear').type(
            torch.half)
        mri_fine_mesh = F.interpolate(mri_mesh.unsqueeze(0), size=mri_fine_feat.shape[2:], mode='trilinear').type(
            torch.half)
        mri_coarse_mesh = F.interpolate(mri_mesh.unsqueeze(0), size=mri_coarse_feat.shape[2:], mode='trilinear').type(
            torch.half)
        mri_fine_mask = F.interpolate(mri_mask.type(mri_mesh.dtype).unsqueeze(0), size=mri_fine_feat.shape[2:]).type(
            torch.half)

        mri_coarse_mesh_flatten =mri_coarse_mesh.view(3,-1)
        ct_coarse_mesh_flatten = ct_coarse_mesh.view(3,-1)

        with torch.no_grad():
            dist_mri_ct = LA.norm((mri_coarse_mesh_flatten.view(-1, mri_coarse_mesh_flatten.shape[1], 1) - \
                            ct_coarse_mesh_flatten.view(-1, 1, ct_coarse_mesh_flatten.shape[1])), dim=0)
            dist_z = LA.norm((mri_coarse_mesh_flatten[2, :].view(-1, mri_coarse_mesh_flatten.shape[1], 1) - \
                                ct_coarse_mesh_flatten[2, :].view(-1, 1, ct_coarse_mesh_flatten.shape[1])), dim=0)

            dist_mri_mri = LA.norm((mri_coarse_mesh_flatten.view(-1, mri_coarse_mesh_flatten.shape[1], 1) - \
                                mri_coarse_mesh_flatten.view(-1, 1, mri_coarse_mesh_flatten.shape[1])), dim=0)
            dist_ct_ct = LA.norm((ct_coarse_mesh_flatten.view(-1, ct_coarse_mesh_flatten.shape[1], 1) - \
                                ct_coarse_mesh_flatten.view(-1, 1, ct_coarse_mesh_flatten.shape[1])), dim=0)
            dist_ct_mri = dist_mri_ct.transpose(1,0)
        pos_coarse_can = torch.unique(
            torch.where(dist_mri_ct < self.train_cfg.inter_cfg.coarse_positive_distance)[0])

        key_coarse_can = dist_mri_ct[pos_coarse_can].min(dim=1)[1]
        pos_coarse = pos_coarse_can[dist_z[pos_coarse_can, key_coarse_can] < 8]
        key_coarse = key_coarse_can[dist_z[pos_coarse_can, key_coarse_can] < 8]

        mri_coarse_feat_flatten = mri_coarse_feat.view(mri_coarse_feat.shape[1],-1)
        ct_coarse_feat_flatten = ct_coarse_feat.view(ct_coarse_feat.shape[1],-1)

        q_mri_coarse_feat = mri_coarse_feat_flatten[:, pos_coarse].transpose(0, 1)
        k_ct_coarse_feat = ct_coarse_feat_flatten[:, key_coarse].transpose(0, 1)


        dist_mri_ct = dist_mri_ct[pos_coarse, :]
        dist_mri_mri = dist_mri_mri[pos_coarse, :]
        dist_ct_ct = dist_ct_ct[key_coarse,:]
        dist_ct_mri = dist_ct_mri[key_coarse,:]

        ind_ignores_mri_mri = torch.where(dist_mri_mri < self.train_cfg.inter_cfg.coarse_ignore_distance)
        ind_ignores_mri_ct = torch.where(dist_mri_ct < self.train_cfg.inter_cfg.coarse_ignore_distance)
        ind_ignores_ct_ct = torch.where(dist_ct_ct < self.train_cfg.inter_cfg.coarse_ignore_distance)
        ind_ignores_ct_mri = torch.where(dist_ct_mri < self.train_cfg.inter_cfg.coarse_ignore_distance)

        ind_ignores_mri_mri = torch.stack(ind_ignores_mri_mri)
        ind_ignores_mri_ct = torch.stack(ind_ignores_mri_ct)
        ind_ignores_ct_ct = torch.stack(ind_ignores_ct_ct)
        ind_ignores_ct_mri = torch.stack(ind_ignores_ct_mri)

        use_mask_mri_mri = torch.ones_like(dist_mri_mri)
        use_mask_mri_ct = torch.ones_like(dist_mri_ct)
        use_mask_ct_ct = torch.ones_like(dist_ct_ct)
        use_mask_ct_mri = torch.ones_like(dist_ct_mri)

        use_mask_mri_mri[ind_ignores_mri_mri[0, :], ind_ignores_mri_mri[1, :]] = 0
        use_mask_mri_ct[ind_ignores_mri_ct[0, :], ind_ignores_mri_ct[1, :]] = 0
        use_mask_ct_ct[ind_ignores_ct_ct[0, :], ind_ignores_ct_ct[1, :]] = 0
        use_mask_ct_mri[ind_ignores_ct_mri[0,:],ind_ignores_ct_mri[1,:]] = 0
        # # neg_mask = torch.cat([use_mask_mri_mri, use_mask_mri_ct], dim=1)
        # neg_mask =  use_mask_mri_ct
        neg_mask_1 = torch.cat([use_mask_mri_mri, use_mask_mri_ct], dim=1)
        neg_mask_2 = torch.cat([use_mask_ct_mri, use_mask_ct_ct], dim=1)

        inner_view = torch.einsum("nc,nc->n", q_mri_coarse_feat, k_ct_coarse_feat).view(-1, 1)


        neg_coarse_view_1 = torch.einsum("nc,ck->nk", q_mri_coarse_feat,
                                        torch.cat((mri_coarse_feat_flatten, ct_coarse_feat_flatten), dim=1))

        neg_coarse_view_2 = torch.einsum("nc,ck->nk", k_ct_coarse_feat,
                                        torch.cat((mri_coarse_feat_flatten, ct_coarse_feat_flatten), dim=1))
        # neg_coarse_view = torch.einsum("nc,ck->nk", q_mri_coarse_feat,ct_coarse_feat_flatten)

        neg_coarse_view_1 = neg_coarse_view_1 * neg_mask_1
        neg_coarse_view_2 = neg_coarse_view_2 * neg_mask_2


        neg_candidate_coarse_index_1 = \
            neg_coarse_view_1.topk(self.train_cfg.inter_cfg.coarse_pre_select_neg_number, dim=1)[1]

        neg_use_coarse_1 = torch.zeros(
            (q_mri_coarse_feat.shape[0], self.train_cfg.inter_cfg.coarse_after_select_neg_number),device=neg_coarse_view_1.device)

        for i in range(q_mri_coarse_feat.shape[0]):
            use_index = neg_candidate_coarse_index_1[i,torch.randperm(neg_candidate_coarse_index_1.shape[1])][
                        :self.train_cfg.inter_cfg.coarse_after_select_neg_number]
            neg_use_coarse_1[i, :use_index.shape[0]] = neg_coarse_view_1[i, use_index]

        neg_candidate_coarse_index_2 = \
            neg_coarse_view_2.topk(self.train_cfg.inter_cfg.coarse_pre_select_neg_number, dim=1)[1]

        neg_use_coarse_2 = torch.zeros(
            (k_ct_coarse_feat.shape[0], self.train_cfg.inter_cfg.coarse_after_select_neg_number),
            device=neg_coarse_view_2.device)

        for i in range(q_mri_coarse_feat.shape[0]):
            use_index = neg_candidate_coarse_index_2[i, torch.randperm(neg_candidate_coarse_index_2.shape[1])][
                        :self.train_cfg.inter_cfg.coarse_after_select_neg_number]
            neg_use_coarse_2[i, :use_index.shape[0]] = neg_coarse_view_1[i, use_index]

        logits_view_1 = torch.cat([inner_view, neg_use_coarse_1], dim=1)
        logits_view_2 = torch.cat([inner_view, neg_use_coarse_2], dim=1)
        logits_view = torch.cat([logits_view_1,logits_view_2],dim=0)
        logits = logits_view / self.train_cfg.inter_cfg.coarse_temperature
        labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
        loss = self.criterion(logits, labels)
        out['loss'] = loss
        return out


    def single_intra_fine_loss(self, fine_feat, coarse_feat, fine_grid, coarse_grid, fine_vaild, coarse_valid, *args):

        view_1_fine = fine_feat[0, :, :, :, :].view(128, -1)
        view_2_fine = fine_feat[1, :, :, :, :].view(128, -1)

        out = dict()

        fine_local_grid = self.meshgrid3d(fine_feat.shape[2:])  # z y x

        p1_y_min = fine_grid[0, 0, fine_vaild[0, 0] > 0].min()
        p1_y_max = fine_grid[0, 0, fine_vaild[0, 0] > 0].max()
        p1_x_min = fine_grid[0, 1, fine_vaild[0, 0] > 0].min()
        p1_x_max = fine_grid[0, 1, fine_vaild[0, 0] > 0].max()
        p1_z_min = fine_grid[0, 2, fine_vaild[0, 0] > 0].min()
        p1_z_max = fine_grid[0, 2, fine_vaild[0, 0] > 0].max()

        p2_y_min = fine_grid[1, 0, fine_vaild[1, 0] > 0].min()
        p2_y_max = fine_grid[1, 0, fine_vaild[1, 0] > 0].max()
        p2_x_min = fine_grid[1, 1, fine_vaild[1, 0] > 0].min()
        p2_x_max = fine_grid[1, 1, fine_vaild[1, 0] > 0].max()
        p2_z_min = fine_grid[1, 2, fine_vaild[1, 0] > 0].min()
        p2_z_max = fine_grid[1, 2, fine_vaild[1, 0] > 0].max()

        intersection_y = [max(p1_y_min, p2_y_min), min(p1_y_max, p2_y_max)]
        intersection_x = [max(p1_x_min, p2_x_min), min(p1_x_max, p2_x_max)]
        intersection_z = [max(p1_z_min, p2_z_min), min(p1_z_max, p2_z_max)]

        intersection_volume_sub1 = (fine_grid[0, 0] >= intersection_y[0]) * (
                fine_grid[0, 0] <= intersection_y[1]) \
                                   * (fine_grid[0, 1] >= intersection_x[0]) * (
                                           fine_grid[0, 1] <= intersection_x[1]) \
                                   * (fine_grid[0, 2] >= intersection_z[0]) * (
                                           fine_grid[0, 2] <= intersection_z[1]) * fine_vaild[0, 0]

        intersection_volume_sub2 = (fine_grid[1, 0] >= intersection_y[0]) * (
                fine_grid[1, 0] <= intersection_y[1]) \
                                   * (fine_grid[1, 1] >= intersection_x[0]) * (
                                           fine_grid[1, 1] <= intersection_x[1]) \
                                   * (fine_grid[1, 2] >= intersection_z[0]) * (
                                           fine_grid[1, 2] <= intersection_z[1]) * fine_vaild[1, 0]
        pos1_index_overlap = fine_local_grid[intersection_volume_sub1 > 0, :]
        pos2_index_overlap = fine_local_grid[intersection_volume_sub2 > 0, :]
        pos_index_all = fine_local_grid.view(-1, 3)
        # yxz pos_mm_overlap      zyx pos_index
        pos1_mm_overlap = fine_grid[0, :, intersection_volume_sub1 > 0]
        pos2_mm_overlap = fine_grid[1, :, intersection_volume_sub2 > 0]
        view1_num_points = pos1_mm_overlap.shape[1]
        view2_num_points = pos2_mm_overlap.shape[1]

        if intersection_y[0] > intersection_y[1] or intersection_x[0] > intersection_x[1] or intersection_z[0] > \
                intersection_z[1] or view1_num_points == 0 or view2_num_points==0:

            # no overlap
            index_use = torch.randperm(fine_feat.shape[2] * fine_feat.shape[3] * fine_feat.shape[4])
            q_index = index_use[:self.train_cfg.intra_cfg.after_select_pos_number]
            fine_local_grid_flatten = fine_local_grid.view(-1, 3)

            q_loc = fine_local_grid_flatten[q_index, :].type(torch.LongTensor)
            pos1_mm_use = fine_grid[0, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]
            pos1_mm_all = fine_grid[0, :, :, :, :].view(3, -1)
            dist1_1 = LA.norm((pos1_mm_use.view(-1, pos1_mm_use.shape[1], 1) - \
                               pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)
            ind_ignores_1_1 = torch.where(dist1_1 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
            use_mask_1_1 = torch.ones_like(dist1_1)
            use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0

            q_fine_feat = fine_feat[0, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]
            q_fine_feat = q_fine_feat.permute(1, 0)

            neg_fine = torch.einsum("nc,ck->nk", q_fine_feat, view_1_fine)
            neg_fine = neg_fine * use_mask_1_1
            inner = torch.einsum("nc,nc->n", q_fine_feat, q_fine_feat).view(-1, 1)

            neg_fine_use_index = neg_fine.topk(self.train_cfg.intra_cfg.after_select_neg_number)[1]
            neg_fine_use = torch.zeros((neg_fine.shape[0], self.train_cfg.intra_cfg.after_select_neg_number),
                                       device=fine_feat.device)
            for i in range(neg_fine.shape[0]):
                neg_fine_use[i, :] = neg_fine[i, neg_fine_use_index[i, :]]
            logits = torch.cat([inner, neg_fine_use], dim=1)
            logits = logits / self.train_cfg.intra_cfg.fine_temperature
            labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
            loss = self.criterion(logits, labels)
        else:
            # have overlap
            # intersection_volume_sub1 = (fine_grid[0, 0] >= intersection_y[0]) * (
            #         fine_grid[0, 0] <= intersection_y[1]) \
            #                            * (fine_grid[0, 1] >= intersection_x[0]) * (
            #                                    fine_grid[0, 1] <= intersection_x[1]) \
            #                            * (fine_grid[0, 2] >= intersection_z[0]) * (
            #                                    fine_grid[0, 2] <= intersection_z[1]) * fine_vaild[0, 0]
            #
            # intersection_volume_sub2 = (fine_grid[1, 0] >= intersection_y[0]) * (
            #         fine_grid[1, 0] <= intersection_y[1]) \
            #                            * (fine_grid[1, 1] >= intersection_x[0]) * (
            #                                    fine_grid[1, 1] <= intersection_x[1]) \
            #                            * (fine_grid[1, 2] >= intersection_z[0]) * (
            #                                    fine_grid[1, 2] <= intersection_z[1]) * fine_vaild[1, 0]
            # pos1_index_overlap = fine_local_grid[intersection_volume_sub1 > 0, :]
            # pos2_index_overlap = fine_local_grid[intersection_volume_sub2 > 0, :]
            # pos_index_all = fine_local_grid.view(-1, 3)
            # # yxz pos_mm_overlap      zyx pos_index
            # pos1_mm_overlap = fine_grid[0, :, intersection_volume_sub1 > 0]
            # pos2_mm_overlap = fine_grid[1, :, intersection_volume_sub2 > 0]
            # view1_num_points = pos1_mm_overlap.shape[1]
            # view2_num_points = pos2_mm_overlap.shape[1]

            if view1_num_points < self.train_cfg.intra_cfg.pre_select_pos_number:
                points1_select = torch.randperm(view1_num_points)
            else:
                points1_select = torch.randperm(view1_num_points)[:self.train_cfg.intra_cfg.pre_select_pos_number]
            if view2_num_points < self.train_cfg.intra_cfg.pre_select_pos_number:
                points2_select = torch.randperm(view2_num_points)
            else:
                points2_select = torch.randperm(view2_num_points)[:self.train_cfg.intra_cfg.pre_select_pos_number]

            pos1_use = pos1_mm_overlap[:, points1_select]
            pos2_use = pos2_mm_overlap[:, points2_select]

            pos1_mm_all = fine_grid[0, :, :, :, :].view(3, -1)
            pos2_mm_all = fine_grid[1, :, :, :, :].view(3, -1)

            # time1 = time.time()

            with torch.no_grad():

                dist1_1 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                                   pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)

                dist1_2 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                                   pos2_mm_all.view(-1, 1, pos2_mm_all.shape[1])), dim=0)
                dist2_1 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                                   pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)
                dist2_2 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                                   pos2_mm_all.view(-1, 1, pos2_mm_all.shape[1])), dim=0)

            # time2 = time.time()
            # print('dist time:',time2-time1)

            pos1 = torch.unique(torch.where(dist1_2 < self.train_cfg.intra_cfg.positive_distance)[0])
            pos2 = torch.unique(torch.where(dist2_1 < self.train_cfg.intra_cfg.positive_distance)[0])

            if pos1.shape[0] == 0:
                print(dist1_2.min())
                pos1 = torch.unique(torch.where(dist1_2 < dist1_2.min() + 0.5)[0])
                pos2 = torch.unique(torch.where(dist2_1 < dist2_1.min() + 0.5)[0])
            if pos1.shape[0] <= self.train_cfg.intra_cfg.after_select_pos_number:
                pos1 = pos1
            else:
                pos1 = pos1[torch.randperm(pos1.shape[0])[:self.train_cfg.intra_cfg.after_select_pos_number]]

            if pos2.shape[0] <= self.train_cfg.intra_cfg.after_select_pos_number:
                pos2 = pos2
            else:
                pos2 = pos2[torch.randperm(pos2.shape[0])[:self.train_cfg.intra_cfg.after_select_pos_number]]

            dist1_1 = dist1_1[pos1, :]
            dist1_2 = dist1_2[pos1, :]
            dist2_1 = dist2_1[pos2, :]
            dist2_2 = dist2_2[pos2, :]

            pos1_source = points1_select[pos1]
            pos2_source = points2_select[pos2]
            pos1_target = dist1_2.min(dim=1)[1]
            pos2_target = dist2_1.min(dim=1)[1]

            ind_ignores_1_1 = torch.where(dist1_1 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_1_2 = torch.where(dist1_2 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_2_1 = torch.where(dist2_1 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_2_2 = torch.where(dist2_2 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
            ind_ignores_1_2 = torch.stack(ind_ignores_1_2)
            ind_ignores_2_1 = torch.stack(ind_ignores_2_1)
            ind_ignores_2_2 = torch.stack(ind_ignores_2_2)

            use_mask_1_1 = torch.ones_like(dist1_1)
            use_mask_1_2 = torch.ones_like(dist1_2)
            use_mask_2_1 = torch.ones_like(dist2_1)
            use_mask_2_2 = torch.ones_like(dist2_2)
            use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0
            use_mask_1_2[ind_ignores_1_2[0, :], ind_ignores_1_2[1, :]] = 0
            use_mask_2_1[ind_ignores_2_1[0, :], ind_ignores_2_1[1, :]] = 0
            use_mask_2_2[ind_ignores_2_2[0, :], ind_ignores_2_2[1, :]] = 0

            neg_mask1 = torch.cat([use_mask_1_1, use_mask_1_2], dim=1)
            neg_mask2 = torch.cat([use_mask_2_1, use_mask_2_2], dim=1)

            coarse_feat_resample = F.interpolate(coarse_feat, fine_feat.shape[2:], mode='trilinear')
            coarse_feat_resample = F.normalize(coarse_feat_resample, dim=1)
            view_1_coarse = coarse_feat_resample[0, :, :, :, :].view(128, -1)
            view_2_coarse = coarse_feat_resample[1, :, :, :, :].view(128, -1)

            q_view1_location = pos1_index_overlap[pos1_source].type(torch.LongTensor)
            k_view1_location = pos_index_all[pos1_target].type(torch.LongTensor)
            q_view2_location = pos2_index_overlap[pos2_source].type(torch.LongTensor)
            k_view2_location = pos_index_all[pos2_target].type(torch.LongTensor)

            q_view1_fine_feat = fine_feat[0, :, q_view1_location[:, 0], q_view1_location[:, 1],
                                q_view1_location[:, 2]].transpose(0, 1)
            k_view1_fine_feat = fine_feat[1, :, k_view1_location[:, 0], k_view1_location[:, 1],
                                k_view1_location[:, 2]].transpose(0, 1)
            q_view2_fine_feat = fine_feat[1, :, q_view2_location[:, 0], q_view2_location[:, 1],
                                q_view2_location[:, 2]].transpose(0, 1)
            k_view2_fine_feat = fine_feat[0, :, k_view2_location[:, 0], k_view2_location[:, 1],
                                k_view2_location[:, 2]].transpose(0, 1)

            q_view1_coarse_feat = coarse_feat_resample[0, :, q_view1_location[:, 0], q_view1_location[:, 1],
                                  q_view1_location[:, 2]].transpose(0, 1)
            q_view2_coarse_feat = coarse_feat_resample[1, :, q_view2_location[:, 0], q_view2_location[:, 1],
                                  q_view2_location[:, 2]].transpose(0, 1)

            inner_view1 = torch.einsum("nc,nc->n", q_view1_fine_feat, k_view1_fine_feat).view(-1, 1)
            inner_view2 = torch.einsum("nc,nc->n", q_view2_fine_feat, k_view2_fine_feat).view(-1, 1)

            neg_fine_view1 = torch.einsum("nc,ck->nk", q_view1_fine_feat, torch.cat((view_1_fine, view_2_fine), dim=1))
            neg_coarse_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                            torch.cat((view_1_coarse, view_2_coarse), dim=1))
            neg_fine_view2 = torch.einsum("nc,ck->nk", q_view2_fine_feat, torch.cat((view_1_fine, view_2_fine), dim=1))
            neg_coarse_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                            torch.cat((view_1_coarse, view_2_coarse), dim=1))

            # print('compute pos and neg time:',time4-time3)

            neg_all_view1 = neg_fine_view1 + neg_coarse_view1
            neg_all_view2 = neg_fine_view2 + neg_coarse_view2
            neg_all_view1 = neg_all_view1 * neg_mask1
            neg_all_view2 = neg_all_view2 * neg_mask2

            neg_candidate_view1_index = neg_all_view1.topk(self.train_cfg.intra_cfg.pre_select_neg_number, dim=1)[1]
            neg_candidate_view2_index = neg_all_view2.topk(self.train_cfg.intra_cfg.pre_select_neg_number, dim=1)[
                1]  # find diverse and hard negatives

            neg_use_view1 = torch.zeros((q_view1_fine_feat.shape[0], self.train_cfg.intra_cfg.after_select_neg_number),
                                        device=neg_all_view1.device)
            neg_use_view2 = torch.zeros((q_view2_fine_feat.shape[0], self.train_cfg.intra_cfg.after_select_neg_number),
                                        device=neg_all_view2.device)
            for i in range(q_view1_fine_feat.shape[0]):
                use_index = neg_candidate_view1_index[i, torch.randperm(neg_candidate_view1_index[i, :].shape[0])[
                                                         :self.train_cfg.intra_cfg.after_select_neg_number]]
                neg_use_view1[i, :] = neg_fine_view1[i, use_index]

            for i in range(q_view2_fine_feat.shape[0]):
                use_index = neg_candidate_view2_index[i, torch.randperm(neg_candidate_view2_index[i, :].shape[0])[
                                                         :self.train_cfg.intra_cfg.after_select_neg_number]]

                neg_use_view2[i, :] = neg_fine_view2[i, use_index]

            logits_view1 = torch.cat([inner_view1, neg_use_view1], dim=1)
            logits_view2 = torch.cat([inner_view2, neg_use_view2], dim=1)
            logits = torch.cat([logits_view1, logits_view2], dim=0)
            logits = logits / self.train_cfg.intra_cfg.fine_temperature

            labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
            loss = self.criterion(logits, labels)
        out['loss'] = loss
        return out

    def single_intra_coarse_loss(self, coarse_feat, coarse_grid, coarse_valid, *args):  # coarse_grid[B,C,D,H,W] C:y,x,z
        coarse_feat_flatten = coarse_feat.view(coarse_feat.shape[0], 128, -1)
        coarse_feat_flatten = coarse_feat_flatten.permute(1, 0, 2)
        coarse_feat_flatten = coarse_feat_flatten.reshape(128, -1)

        out = dict()

        for i in range(int(coarse_feat.shape[0] / 2)):
            view1_feat_coarse = coarse_feat[2 * i, :, :, :, :].view(128, -1)
            view2_feat_coarse = coarse_feat[2 * i + 1, :, :, :, :].view(128, -1)
            view1_loc_coarse = coarse_grid[2 * i, :, :, :, :].view(3, -1)
            view2_loc_coarse = coarse_grid[2 * i + 1, :, :, :, :].view(3, -1)

            global_use_index_list = [j + 2 * i * view1_feat_coarse.shape[1] for j in
                                     range(0, 2 * view1_feat_coarse.shape[1])]
            global_index_list = [j for j in range(0, coarse_feat.shape[0] * view1_feat_coarse.shape[1])]
            global_search_index_list = list(set(global_index_list) - set(global_use_index_list))
            global_search_index = torch.tensor(global_search_index_list).to(view1_feat_coarse.device)

            with torch.no_grad():
                dist = LA.norm((view1_loc_coarse.view(-1, view1_loc_coarse.shape[1], 1) - \
                                view2_loc_coarse.view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
                dist_z_1 = LA.norm((view1_loc_coarse[2, :].view(-1, view1_loc_coarse.shape[1], 1) - \
                                    view2_loc_coarse[2, :].view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
                dist_z_2 = LA.norm((view2_loc_coarse[2, :].view(-1, view2_loc_coarse.shape[1], 1) - \
                                    view1_loc_coarse[2, :].view(-1, 1, view1_loc_coarse.shape[1])), dim=0)
                dist_1_1 = LA.norm((view1_loc_coarse.view(-1, view1_loc_coarse.shape[1], 1) - \
                                    view1_loc_coarse.view(-1, 1, view1_loc_coarse.shape[1])), dim=0)
                dist_2_2 = LA.norm((view2_loc_coarse.view(-1, view2_loc_coarse.shape[1], 1) - \
                                    view2_loc_coarse.view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
            pos_view1_coarse_can = torch.unique(
                torch.where(dist < self.train_cfg.intra_cfg.coarse_positive_distance)[0])

            key_view1_coarse_can = dist[pos_view1_coarse_can].min(dim=1)[1]
            pos_view1_coarse = pos_view1_coarse_can[dist_z_1[pos_view1_coarse_can, key_view1_coarse_can] < 8]
            key_view1_coarse = key_view1_coarse_can[dist_z_1[pos_view1_coarse_can, key_view1_coarse_can] < 8]
            pos_view2_coarse_can = torch.unique(
                torch.where(dist.transpose(0, 1) < self.train_cfg.intra_cfg.coarse_positive_distance)[0])
            key_view2_coarse_can = dist.transpose(0, 1)[pos_view2_coarse_can].min(dim=1)[1]
            pos_view2_coarse = pos_view2_coarse_can[
                dist_z_2[pos_view2_coarse_can, key_view2_coarse_can] < self.train_cfg.intra_cfg.coarse_z_thres]
            key_view2_coarse = key_view2_coarse_can[
                dist_z_2[pos_view2_coarse_can, key_view2_coarse_can] < self.train_cfg.intra_cfg.coarse_z_thres]

            if pos_view1_coarse.shape[0] > 0:
                q_view1_coarse_feat = view1_feat_coarse[:, pos_view1_coarse].transpose(0, 1)
                k_view1_coarse_feat = view2_feat_coarse[:, key_view1_coarse].transpose(0, 1)
                q_view2_coarse_feat = view2_feat_coarse[:, pos_view2_coarse].transpose(0, 1)
                k_view2_coarse_feat = view1_feat_coarse[:, key_view2_coarse].transpose(0, 1)

                # all_nodes_view1 = []
                # all_nodes_view2 = []
                dist_1_2 = dist.clone()
                dist_2_1 = dist.transpose(0, 1).clone()
                dist_1_2 = dist_1_2[pos_view1_coarse, :]
                dist_1_1 = dist_1_1[pos_view1_coarse, :]
                dist_2_1 = dist_2_1[pos_view2_coarse, :]
                dist_2_2 = dist_2_2[pos_view2_coarse, :]

                ind_ignores_1_1 = torch.where(dist_1_1 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_1_2 = torch.where(dist_1_2 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_2_1 = torch.where(dist_2_1 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_2_2 = torch.where(dist_2_2 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
                ind_ignores_1_2 = torch.stack(ind_ignores_1_2)
                ind_ignores_2_1 = torch.stack(ind_ignores_2_1)
                ind_ignores_2_2 = torch.stack(ind_ignores_2_2)

                use_mask_1_1 = torch.ones_like(dist_1_1)
                use_mask_1_2 = torch.ones_like(dist_1_2)
                use_mask_2_1 = torch.ones_like(dist_2_1)
                use_mask_2_2 = torch.ones_like(dist_2_2)

                use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0
                use_mask_1_2[ind_ignores_1_2[0, :], ind_ignores_1_2[1, :]] = 0
                use_mask_2_1[ind_ignores_2_1[0, :], ind_ignores_2_1[1, :]] = 0
                use_mask_2_2[ind_ignores_2_2[0, :], ind_ignores_2_2[1, :]] = 0

                neg_mask1 = torch.cat([use_mask_1_1, use_mask_1_2], dim=1)
                neg_mask2 = torch.cat([use_mask_2_1, use_mask_2_2], dim=1)

                inner_view1 = torch.einsum("nc,nc->n", q_view1_coarse_feat, k_view1_coarse_feat).view(-1, 1)
                inner_view2 = torch.einsum("nc,nc->n", q_view2_coarse_feat, k_view2_coarse_feat).view(-1, 1)

                neg_coarse_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                                torch.cat((view1_feat_coarse, view2_feat_coarse), dim=1))
                neg_coarse_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                                torch.cat((view1_feat_coarse, view2_feat_coarse), dim=1))

                neg_coarse_view1 = neg_coarse_view1 * neg_mask1
                neg_coarse_view2 = neg_coarse_view2 * neg_mask2

                neg_candidate_view1_coarse_index = \
                    neg_coarse_view1.topk(self.train_cfg.intra_cfg.coarse_pre_select_neg_number, dim=1)[1]
                neg_candidate_view2_coarse_index = \
                    neg_coarse_view2.topk(self.train_cfg.intra_cfg.coarse_pre_select_neg_number, dim=1)[1]

                neg_use_view1_coarse = torch.zeros(
                    (q_view1_coarse_feat.shape[0], self.train_cfg.intra_cfg.coarse_after_select_neg_number))
                neg_use_view2_coarse = torch.zeros(
                    (q_view2_coarse_feat.shape[0], self.train_cfg.intra_cfg.coarse_after_select_neg_number))

                for i in range(q_view1_coarse_feat.shape[0]):
                    use_index = neg_candidate_view1_coarse_index[i,torch.randperm(neg_candidate_view1_coarse_index.shape[1])][
                                :self.train_cfg.intra_cfg.coarse_after_select_neg_number]
                    neg_use_view1_coarse[i, :use_index.shape[0]] = neg_coarse_view1[i, use_index]

                for i in range(q_view2_coarse_feat.shape[0]):
                    use_index = neg_candidate_view2_coarse_index[i,torch.randperm(neg_candidate_view2_coarse_index.shape[1])][
                                :self.train_cfg.intra_cfg.coarse_after_select_neg_number]
                    neg_use_view2_coarse[i, :use_index.shape[0]] = neg_coarse_view2[i, use_index]

                neg_coarse_global_index_1 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_1 = coarse_feat_flatten[:, neg_coarse_global_index_1]
                neg_coarse_global_index_2 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_2 = coarse_feat_flatten[:, neg_coarse_global_index_2]

                neg_coarse_global_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                                       neg_coarse_global_1)
                neg_coarse_global_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                                       neg_coarse_global_2)

                neg_view1_coarse_concate = torch.cat(
                    [neg_use_view1_coarse.to(view1_feat_coarse.device), neg_coarse_global_view1], dim=1)
                neg_view2_coarse_concate = torch.cat(
                    [neg_use_view2_coarse.to(view1_feat_coarse.device), neg_coarse_global_view2], dim=1)

                logits_view1 = torch.cat([inner_view1, neg_view1_coarse_concate], dim=1)
                logits_view2 = torch.cat([inner_view2, neg_view2_coarse_concate], dim=1)
                logits = torch.cat([logits_view1, logits_view2], dim=0)
                logits = logits / self.train_cfg.intra_cfg.coarse_temperature
                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
                loss = self.criterion(logits, labels)
                if out.keys().__len__() == 0:
                    out['loss'] = loss
                else:
                    out['loss'] += loss
            else:
                neg_coarse_global_index_1 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_1 = coarse_feat_flatten[:, neg_coarse_global_index_1]
                neg_coarse_global_index_2 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_2 = coarse_feat_flatten[:, neg_coarse_global_index_2]

                neg_coarse_global_view1 = torch.einsum("nc,ck->nk", view1_feat_coarse.transpose(0, 1),
                                                       neg_coarse_global_1)
                neg_coarse_global_view2 = torch.einsum("nc,ck->nk", view2_feat_coarse.transpose(0, 1),
                                                       neg_coarse_global_2)

                logits_view1 = torch.einsum("nc,ck->nk", view1_feat_coarse.transpose(0, 1),
                                            torch.cat([view1_feat_coarse, view2_feat_coarse], dim=1)).topk(
                    self.train_cfg.intra_cfg.coarse_after_select_neg_number + 1,
                    dim=1)[
                    0]
                logits_view2 = torch.einsum("nc,ck->nk", view2_feat_coarse.transpose(0, 1),
                                            torch.cat([view1_feat_coarse, view2_feat_coarse], dim=1)).topk(
                    self.train_cfg.intra_cfg.coarse_after_select_neg_number + 1,
                    dim=1)[
                    0]

                logits_view1 = torch.cat([logits_view1, neg_coarse_global_view1], dim=1)
                logits_view2 = torch.cat([logits_view2, neg_coarse_global_view2], dim=1)
                logits = torch.cat([logits_view1, logits_view2], dim=0)
                logits = logits / self.train_cfg.intra_cfg.coarse_temperature
                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
                loss = self.criterion(logits, labels)
                if out.keys().__len__() == 0:
                    out['loss'] = loss
                else:
                    out['loss'] += loss

        out['loss'] = out['loss'] / (coarse_feat.shape[0] / 2)
        return out

    def meshgrid3d(self, shape):
        z_ = torch.linspace(0., shape[0] - 1, shape[0])
        y_ = torch.linspace(0., shape[1] - 1, shape[1])
        x_ = torch.linspace(0., shape[2] - 1, shape[2])
        z, y, x = torch.meshgrid(z_, y_, x_)
        return torch.stack((z, y, x), 3)

    async def async_simple_test(self,
                                img,
                                img_meta,
                                proposals=None,
                                rescale=False):
        """Async test without augmentation."""
        assert self.with_bbox, 'Bbox head must be implemented.'
        x = self.extract_feat(img)

        if proposals is None:
            proposal_list = await self.rpn_head.async_simple_test_rpn(
                x, img_meta)
        else:
            proposal_list = proposals

        return await self.roi_head.async_simple_test(
            x, proposal_list, img_meta, rescale=rescale)

    def simple_test(self, img, img_metas, proposals=None, rescale=False,**kwargs):
        """Test without augmentation."""
        is_mri = img_metas[0]['is_mri']
        x = self.extract_feat(img)
        outs = []
        out1 = x[0].data.cpu().numpy()
        out2 = x[1].data.cpu().numpy()
        # out3 = x[2].data.cpu().numpy()
        outs = [out1, out2, img.data.cpu().numpy(), img_metas[0]['filename'].split('.', 1)[0]]
        output_embedding = self.test_cfg.get('output_embedding', True)
        if not output_embedding:
            if not os.path.exists(self.test_cfg.save_path):
                os.mkdir(self.test_cfg.save_path)
            outfilename = self.test_cfg.save_path + \
                          img_metas[0]['filename'].split('.', 1)[0] + '.pkl'
            f = open(outfilename, 'wb')
            pickle.dump(outs, f)
            return [
                x[0][0, 0, 0, 0, 0].data.cpu()]  # we have saved the data into harddisk, this is just for fit the code
        else:
            return outs

    def aug_test(self, imgs, img_metas, **kwargs):
        return self.simple_test(imgs, img_metas, **kwargs)

    def sample_index(self, mask, num_points, num_bg_points=1000):
        labels = torch.unique(mask).tolist()
        labels_organ = [label for label in labels if label > 0]
        if labels_organ.__len__() > 15:
            labels_organ = random.shuffle(labels_organ)
            labels_organ = labels_organ[:15]
        labels_organ.append(0)
        count = 0
        for label in labels_organ:
            all_ind = torch.where(mask == label)[0]
            if label == 0:
                thres = num_bg_points
            else:
                thres = num_points
            if all_ind.shape[0] > thres:
                use_count = thres
            else:
                use_count = all_ind.shape[0]
            pert = torch.randperm(all_ind.shape[0])[:use_count]
            use_ind = all_ind[pert]
            if count == 0:
                sample_index = use_ind
            else:
                sample_index = torch.cat((sample_index, use_ind))
            count = count + 1
        return sample_index


def log_sinkhorn_iterations(Z: torch.Tensor, log_mu: torch.Tensor, log_nu: torch.Tensor, iters: int) -> torch.Tensor:
    """ Perform Sinkhorn Normalization in Log-space for stability"""
    u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
    for _ in range(iters):
        u = log_mu - torch.logsumexp(Z + v.unsqueeze(1), dim=2)
        v = log_nu - torch.logsumexp(Z + u.unsqueeze(2), dim=1)
    return Z + u.unsqueeze(2) + v.unsqueeze(1)


def log_optimal_transport(scores: torch.Tensor, alpha: torch.Tensor, iters: int) -> torch.Tensor:
    """ Perform Differentiable Optimal Transport in Log-space for stability"""
    b, m, n = scores.shape
    one = scores.new_tensor(1)
    ms, ns = (m * one).to(scores), (n * one).to(scores)

    bins0 = alpha.expand(b, m, 1)
    bins1 = alpha.expand(b, 1, n)
    alpha = alpha.expand(b, 1, 1)

    couplings = torch.cat([torch.cat([scores, bins0], -1),
                           torch.cat([bins1, alpha], -1)], 1)

    norm = - (ms + ns).log()
    log_mu = torch.cat([norm.expand(m), ns.log()[None] + norm])
    log_nu = torch.cat([norm.expand(n), ms.log()[None] + norm])
    log_mu, log_nu = log_mu[None].expand(b, -1), log_nu[None].expand(b, -1)

    Z = log_sinkhorn_iterations(couplings, log_mu, log_nu, iters)
    Z = Z - norm  # multiply probabilities by M+N
    return Z


def arange_like(x, dim: int):
    return x.new_ones(x.shape[dim]).cumsum(0) - 1  # traceable in 1.1

sam/models/frameworks/sam_v2_uniform_cross_volume_mask_dual_head_MRI_CT.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai
# Rearrange the ignore computation into useindex, save computation cost
import os.path
import warnings

# import ipdb
import torch
from mmdet.models.builder import DETECTORS, build_backbone, build_head, build_neck, build_loss
from mmdet.models.detectors.base import BaseDetector
import torch.nn.functional as F
from torch import linalg as LA
import time
import pickle
import numpy as np
import math
import matplotlib.pyplot as plt
from ..losses import SupConMeanLoss
import random


@DETECTORS.register_module()
class Sam_uniform_cross_volume_dual_head_mean_vector(BaseDetector):
    def __init__(self,
                 backbone,
                 neck=None,
                 read_out_head=None,
                 sem_neck=None,
                 superloss=None,
                 train_cfg=None,
                 test_cfg=None,
                 init_cfg=None):
        super(Sam_uniform_cross_volume_dual_head_mean_vector, self).__init__(init_cfg)
        self.backbone = build_backbone(backbone)
        self.backbone.init_weights()
        if neck is not None:
            self.neck = build_neck(neck)
        self.criterion = torch.nn.CrossEntropyLoss().cuda()
        self.supcriterion = build_loss(superloss)
        self.read_out_head = build_neck(read_out_head)
        if sem_neck is not None:
            self.semantic_head = build_neck(sem_neck)
        else:
            self.semantic_head = build_neck(neck)
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

    def extract_feat(self, img, normalize=True):
        """Directly extract features from the backbone+neck."""
        x = self.backbone(img)
        out1 = self.neck(x[:self.neck.end_level])[0]
        out2 = self.read_out_head([x[-1]])[0]
        out3 = self.semantic_head(x[:self.semantic_head.end_level])[0]
        if normalize:
            out1 = F.normalize(out1, dim=1)
            out2 = F.normalize(out2, dim=1)
            out3 = F.normalize(out3, dim=1)
        out1 = out1.type(torch.half)
        out2 = out2.type(torch.half)
        out3 = out3.type(torch.half)
        return [out1, out2, out3]

    def forward_dummy(self, img):
        """Used for computing network flops.

        See `mmdetection/tools/analysis_tools/get_flops.py`
        """
        outs = ()
        # backbone
        x = self.extract_feat(img)
        outs = outs + x

        return outs

    def forward_train(self,
                      img,
                      img_metas,
                      **kwargs):

        """
        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.

            img_metas (list[dict]): list of image info dict where each dict
                has: 'img_shape', 'img_name',etc.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        all_intra_info = []
        all_intra_img = []
        all_inter_info = []
        all_inter_img = []

        p = np.random.rand()

        for i in range(img_metas.__len__()):
            im_info = img_metas[i][0]
            if im_info['style'] == 'intra':
                all_intra_info.append(im_info)
                all_intra_img.append(img[i])
            if im_info['style'] == 'inter':
                all_inter_info.append(im_info)
                all_inter_img.append(img[i])

        intra_imgs = torch.stack(all_intra_img)
        mesh_grids = []
        valids = []
        for i in range(all_intra_info.__len__()):
            mesh = all_intra_info[i]['meshgrid']
            mesh_grids.append(mesh)
            valid = all_intra_info[i]['valid']
            valids.append(valid)
            all_intra_info[i].pop('meshgrid')
            all_intra_info[i].pop('valid')
        mesh_grids = torch.stack(mesh_grids).to(intra_imgs.device)
        valids = torch.stack(valids).to(intra_imgs.device)

        losses = dict()
        if p < 0.4:
            all_inter_feats = []
            for i, image in enumerate(all_inter_img):
                feat = self.extract_feat(image.unsqueeze(0))
                all_inter_feats.append(feat)
            loss = self.loss_inter(all_inter_feats, all_inter_info)
        else:
            feats = self.extract_feat(intra_imgs)
            loss = self.loss_intra(feats, mesh_grids, all_intra_info, valids)
        losses['loss'] = loss
        return losses

    def loss_intra(self, feats, meshgrid, img_metas, valid, *args):

        N, C, D, H, W = meshgrid.shape
        output_half = F.interpolate(meshgrid, size=(int(D / 2), int(H / 2), int(W / 2)), mode='trilinear', align_corners=False)
        output_416 = F.interpolate(meshgrid, size=(int(D / 4), int(H / 16), int(W / 16)), mode='trilinear', align_corners=False)
        valid_half = F.interpolate(valid, size=(int(D / 2), int(H / 2), int(W / 2)))
        valid_416 = F.interpolate(valid, size=(int(D / 4), int(H / 16), int(W / 16)))

        output_half = output_half.type(torch.half)
        output_416 = output_416.type(torch.half)
        valid_half = valid_half.type(torch.half)
        valid_416 = valid_416.type(torch.half)

        # fine_time_start = time.time()

        for i in range(int(N / 2)):
            result = self.single_intra_fine_loss(feats[0][2 * i:2 * (i + 1)], feats[1][2 * i:2 * (i + 1)],
                                                 output_half[2 * i:2 * (i + 1)], output_416[2 * i:2 * (i + 1)],
                                                 valid_half[2 * i:2 * (i + 1)], valid_416[2 * i:2 * (i + 1)])
            if i == 0:
                fine_loss = result['loss']
            else:
                fine_loss += result['loss']
        fine_loss = fine_loss / int(N / 2)
        result = self.single_intra_coarse_loss(feats[1], output_416, valid_416)
        coarse_loss = result['loss']

        loss = fine_loss + coarse_loss
        return loss

    def loss_inter(self, feats, img_metas, **kwargs):
        N = feats.__len__()

        for i in range(int(N / 2)):
            result = self.single_inter_loss([feats[2 * i], feats[2 * i + 1]],
                                            [img_metas[2 * i]['mask'], img_metas[2 * i + 1]['mask']])
            if i == 0:
                fine_loss = result['loss']
            else:
                fine_loss += result['loss']
        fine_loss = fine_loss / int(N / 2)
        loss = fine_loss
        return loss

    def single_inter_loss(self, feat, mask):
        out = dict()
        view_1_fine = feat[0][2]
        view_2_fine = feat[1][2]
        view_1_fine = view_1_fine.view(view_1_fine.shape[1], -1).unsqueeze(1).permute(2, 1, 0)
        view_2_fine = view_2_fine.view(view_2_fine.shape[1], -1).unsqueeze(1).permute(2, 1, 0)
        mask_1 = mask[0].to(view_1_fine.device).type(torch.half)
        mask_2 = mask[1].to(view_2_fine.device).type(torch.half)
        mask_1_fine = F.interpolate(mask_1.unsqueeze(0), size=feat[0][2][0, :, :, :, :].shape[1:]).view(-1)
        mask_2_fine = F.interpolate(mask_2.unsqueeze(0), size=feat[1][2][0, :, :, :, :].shape[1:]).view(-1)
        # fine_1_means,fine_1_labels = self.get_mean_vector(mask_1_fine,view_1_fine)
        # fine_2_means,fine_2_labels = self.get_mean_vector(mask_2_fine,view_2_fine)
        fine_feats = torch.cat((view_1_fine, view_2_fine), dim=0)
        fine_labels = torch.cat((mask_1_fine, mask_2_fine), dim=0)
        fine_mean_feat, fine_mean_labels = self.get_mean_vector(fine_labels, fine_feats)
        # fine_mean_feat = torch.cat((fine_1_means,fine_2_means),dim=0)
        # fine_mean_labels = torch.cat((fine_1_labels,fine_2_labels),dim=0)
        loss = self.supcriterion(fine_feats, fine_labels, fine_mean_feat, fine_mean_labels)
        out['loss'] = loss
        return out

    def single_intra_fine_loss(self, fine_feat, coarse_feat, fine_grid, coarse_grid, fine_vaild, coarse_valid, *args):

        view_1_fine = fine_feat[0, :, :, :, :].view(128, -1)
        view_2_fine = fine_feat[1, :, :, :, :].view(128, -1)

        out = dict()

        fine_local_grid = self.meshgrid3d(fine_feat.shape[2:])  # z y x

        p1_y_min = fine_grid[0, 0, fine_vaild[0, 0] > 0].min()
        p1_y_max = fine_grid[0, 0, fine_vaild[0, 0] > 0].max()
        p1_x_min = fine_grid[0, 1, fine_vaild[0, 0] > 0].min()
        p1_x_max = fine_grid[0, 1, fine_vaild[0, 0] > 0].max()
        p1_z_min = fine_grid[0, 2, fine_vaild[0, 0] > 0].min()
        p1_z_max = fine_grid[0, 2, fine_vaild[0, 0] > 0].max()

        p2_y_min = fine_grid[1, 0, fine_vaild[1, 0] > 0].min()
        p2_y_max = fine_grid[1, 0, fine_vaild[1, 0] > 0].max()
        p2_x_min = fine_grid[1, 1, fine_vaild[1, 0] > 0].min()
        p2_x_max = fine_grid[1, 1, fine_vaild[1, 0] > 0].max()
        p2_z_min = fine_grid[1, 2, fine_vaild[1, 0] > 0].min()
        p2_z_max = fine_grid[1, 2, fine_vaild[1, 0] > 0].max()

        intersection_y = [max(p1_y_min, p2_y_min), min(p1_y_max, p2_y_max)]
        intersection_x = [max(p1_x_min, p2_x_min), min(p1_x_max, p2_x_max)]
        intersection_z = [max(p1_z_min, p2_z_min), min(p1_z_max, p2_z_max)]

        if intersection_y[0] > intersection_y[1] or intersection_x[0] > intersection_x[1] or intersection_z[0] > \
                intersection_z[1]:

            # no overlap
            index_use = torch.randperm(fine_feat.shape[2] * fine_feat.shape[3] * fine_feat.shape[4])
            q_index = index_use[:self.train_cfg.intra_cfg.after_select_pos_number]
            fine_local_grid_flatten = fine_local_grid.view(-1, 3)

            q_loc = fine_local_grid_flatten[q_index, :].type(torch.LongTensor)
            pos1_mm_use = fine_grid[0, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]
            pos1_mm_all = fine_grid[0, :, :, :, :].view(3, -1)
            dist1_1 = LA.norm((pos1_mm_use.view(-1, pos1_mm_use.shape[1], 1) - \
                               pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)
            ind_ignores_1_1 = torch.where(dist1_1 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
            use_mask_1_1 = torch.ones_like(dist1_1)
            use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0

            q_fine_feat = fine_feat[0, :, q_loc[:, 0], q_loc[:, 1], q_loc[:, 2]]
            q_fine_feat = q_fine_feat.permute(1, 0)

            neg_fine = torch.einsum("nc,ck->nk", q_fine_feat, view_1_fine)
            neg_fine = neg_fine * use_mask_1_1
            inner = torch.einsum("nc,nc->n", q_fine_feat, q_fine_feat).view(-1, 1)

            neg_fine_use_index = neg_fine.topk(self.train_cfg.intra_cfg.after_select_neg_number)[1]
            neg_fine_use = torch.zeros((neg_fine.shape[0], self.train_cfg.intra_cfg.after_select_neg_number),
                                       device=fine_feat.device)
            for i in range(neg_fine.shape[0]):
                neg_fine_use[i, :] = neg_fine[i, neg_fine_use_index[i, :]]
            logits = torch.cat([inner, neg_fine_use], dim=1)
            logits = logits / self.train_cfg.intra_cfg.temperature
            labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
            loss = self.criterion(logits, labels)
        else:
            # have overlap
            intersection_volume_sub1 = (fine_grid[0, 0] >= intersection_y[0]) * (
                    fine_grid[0, 0] <= intersection_y[1]) \
                                       * (fine_grid[0, 1] >= intersection_x[0]) * (
                                               fine_grid[0, 1] <= intersection_x[1]) \
                                       * (fine_grid[0, 2] >= intersection_z[0]) * (
                                               fine_grid[0, 2] <= intersection_z[1]) * fine_vaild[0, 0]

            intersection_volume_sub2 = (fine_grid[1, 0] >= intersection_y[0]) * (
                    fine_grid[1, 0] <= intersection_y[1]) \
                                       * (fine_grid[1, 1] >= intersection_x[0]) * (
                                               fine_grid[1, 1] <= intersection_x[1]) \
                                       * (fine_grid[1, 2] >= intersection_z[0]) * (
                                               fine_grid[1, 2] <= intersection_z[1]) * fine_vaild[1, 0]
            pos1_index_overlap = fine_local_grid[intersection_volume_sub1 > 0, :]
            pos2_index_overlap = fine_local_grid[intersection_volume_sub2 > 0, :]
            pos_index_all = fine_local_grid.view(-1, 3)
            # yxz pos_mm_overlap      zyx pos_index
            pos1_mm_overlap = fine_grid[0, :, intersection_volume_sub1 > 0]
            pos2_mm_overlap = fine_grid[1, :, intersection_volume_sub2 > 0]
            view1_num_points = pos1_mm_overlap.shape[1]
            view2_num_points = pos2_mm_overlap.shape[1]

            if view1_num_points < self.train_cfg.intra_cfg.pre_select_pos_number:
                points1_select = torch.randperm(view1_num_points)
            else:
                points1_select = torch.randperm(view1_num_points)[:self.train_cfg.intra_cfg.pre_select_pos_number]
            if view2_num_points < self.train_cfg.intra_cfg.pre_select_pos_number:
                points2_select = torch.randperm(view2_num_points)
            else:
                points2_select = torch.randperm(view2_num_points)[:self.train_cfg.intra_cfg.pre_select_pos_number]

            pos1_use = pos1_mm_overlap[:, points1_select]
            pos2_use = pos2_mm_overlap[:, points2_select]

            pos1_mm_all = fine_grid[0, :, :, :, :].view(3, -1)
            pos2_mm_all = fine_grid[1, :, :, :, :].view(3, -1)

            # time1 = time.time()

            with torch.no_grad():

                dist1_1 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                                   pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)

                dist1_2 = LA.norm((pos1_use.view(-1, pos1_use.shape[1], 1) - \
                                   pos2_mm_all.view(-1, 1, pos2_mm_all.shape[1])), dim=0)
                dist2_1 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                                   pos1_mm_all.view(-1, 1, pos1_mm_all.shape[1])), dim=0)
                dist2_2 = LA.norm((pos2_use.view(-1, pos2_use.shape[1], 1) - \
                                   pos2_mm_all.view(-1, 1, pos2_mm_all.shape[1])), dim=0)

            # time2 = time.time()
            # print('dist time:',time2-time1)

            pos1 = torch.unique(torch.where(dist1_2 < self.train_cfg.intra_cfg.positive_distance)[0])
            pos2 = torch.unique(torch.where(dist2_1 < self.train_cfg.intra_cfg.positive_distance)[0])

            if pos1.shape[0] == 0:
                print(dist1_2.min())
                pos1 = torch.unique(torch.where(dist1_2 < dist1_2.min() + 0.5)[0])
                pos2 = torch.unique(torch.where(dist2_1 < dist2_1.min() + 0.5)[0])
            if pos1.shape[0] <= self.train_cfg.intra_cfg.after_select_pos_number:
                pos1 = pos1
            else:
                pos1 = pos1[torch.randperm(pos1.shape[0])[:self.train_cfg.intra_cfg.after_select_pos_number]]

            if pos2.shape[0] <= self.train_cfg.intra_cfg.after_select_pos_number:
                pos2 = pos2
            else:
                pos2 = pos2[torch.randperm(pos2.shape[0])[:self.train_cfg.intra_cfg.after_select_pos_number]]

            dist1_1 = dist1_1[pos1, :]
            dist1_2 = dist1_2[pos1, :]
            dist2_1 = dist2_1[pos2, :]
            dist2_2 = dist2_2[pos2, :]

            pos1_source = points1_select[pos1]
            pos2_source = points2_select[pos2]
            pos1_target = dist1_2.min(dim=1)[1]
            pos2_target = dist2_1.min(dim=1)[1]

            ind_ignores_1_1 = torch.where(dist1_1 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_1_2 = torch.where(dist1_2 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_2_1 = torch.where(dist2_1 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_2_2 = torch.where(dist2_2 < self.train_cfg.intra_cfg.ignore_distance)
            ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
            ind_ignores_1_2 = torch.stack(ind_ignores_1_2)
            ind_ignores_2_1 = torch.stack(ind_ignores_2_1)
            ind_ignores_2_2 = torch.stack(ind_ignores_2_2)

            use_mask_1_1 = torch.ones_like(dist1_1)
            use_mask_1_2 = torch.ones_like(dist1_2)
            use_mask_2_1 = torch.ones_like(dist2_1)
            use_mask_2_2 = torch.ones_like(dist2_2)
            use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0
            use_mask_1_2[ind_ignores_1_2[0, :], ind_ignores_1_2[1, :]] = 0
            use_mask_2_1[ind_ignores_2_1[0, :], ind_ignores_2_1[1, :]] = 0
            use_mask_2_2[ind_ignores_2_2[0, :], ind_ignores_2_2[1, :]] = 0

            neg_mask1 = torch.cat([use_mask_1_1, use_mask_1_2], dim=1)
            neg_mask2 = torch.cat([use_mask_2_1, use_mask_2_2], dim=1)

            coarse_feat_resample = F.interpolate(coarse_feat, fine_feat.shape[2:], mode='trilinear', align_corners=False)
            coarse_feat_resample = F.normalize(coarse_feat_resample, dim=1)
            view_1_coarse = coarse_feat_resample[0, :, :, :, :].view(128, -1)
            view_2_coarse = coarse_feat_resample[1, :, :, :, :].view(128, -1)

            q_view1_location = pos1_index_overlap[pos1_source].type(torch.LongTensor)
            k_view1_location = pos_index_all[pos1_target].type(torch.LongTensor)
            q_view2_location = pos2_index_overlap[pos2_source].type(torch.LongTensor)
            k_view2_location = pos_index_all[pos2_target].type(torch.LongTensor)

            q_view1_fine_feat = fine_feat[0, :, q_view1_location[:, 0], q_view1_location[:, 1],
                                q_view1_location[:, 2]].transpose(0, 1)
            k_view1_fine_feat = fine_feat[1, :, k_view1_location[:, 0], k_view1_location[:, 1],
                                k_view1_location[:, 2]].transpose(0, 1)
            q_view2_fine_feat = fine_feat[1, :, q_view2_location[:, 0], q_view2_location[:, 1],
                                q_view2_location[:, 2]].transpose(0, 1)
            k_view2_fine_feat = fine_feat[0, :, k_view2_location[:, 0], k_view2_location[:, 1],
                                k_view2_location[:, 2]].transpose(0, 1)

            q_view1_coarse_feat = coarse_feat_resample[0, :, q_view1_location[:, 0], q_view1_location[:, 1],
                                  q_view1_location[:, 2]].transpose(0, 1)
            q_view2_coarse_feat = coarse_feat_resample[1, :, q_view2_location[:, 0], q_view2_location[:, 1],
                                  q_view2_location[:, 2]].transpose(0, 1)

            inner_view1 = torch.einsum("nc,nc->n", q_view1_fine_feat, k_view1_fine_feat).view(-1, 1)
            inner_view2 = torch.einsum("nc,nc->n", q_view2_fine_feat, k_view2_fine_feat).view(-1, 1)

            neg_fine_view1 = torch.einsum("nc,ck->nk", q_view1_fine_feat, torch.cat((view_1_fine, view_2_fine), dim=1))
            neg_coarse_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                            torch.cat((view_1_coarse, view_2_coarse), dim=1))
            neg_fine_view2 = torch.einsum("nc,ck->nk", q_view2_fine_feat, torch.cat((view_1_fine, view_2_fine), dim=1))
            neg_coarse_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                            torch.cat((view_1_coarse, view_2_coarse), dim=1))

            # print('compute pos and neg time:',time4-time3)

            neg_all_view1 = neg_fine_view1 + neg_coarse_view1
            neg_all_view2 = neg_fine_view2 + neg_coarse_view2
            neg_all_view1 = neg_all_view1 * neg_mask1
            neg_all_view2 = neg_all_view2 * neg_mask2

            neg_candidate_view1_index = neg_all_view1.topk(self.train_cfg.intra_cfg.pre_select_neg_number, dim=1)[1]
            neg_candidate_view2_index = neg_all_view2.topk(self.train_cfg.intra_cfg.pre_select_neg_number, dim=1)[
                1]  # find diverse and hard negatives

            neg_use_view1 = torch.zeros((q_view1_fine_feat.shape[0], self.train_cfg.intra_cfg.after_select_neg_number),
                                        device=neg_all_view1.device)
            neg_use_view2 = torch.zeros((q_view2_fine_feat.shape[0], self.train_cfg.intra_cfg.after_select_neg_number),
                                        device=neg_all_view2.device)
            for i in range(q_view1_fine_feat.shape[0]):
                use_index = neg_candidate_view1_index[i, torch.randperm(neg_candidate_view1_index[i, :].shape[0])[
                                                         :self.train_cfg.intra_cfg.after_select_neg_number]]
                neg_use_view1[i, :] = neg_fine_view1[i, use_index]

            for i in range(q_view2_fine_feat.shape[0]):
                use_index = neg_candidate_view2_index[i, torch.randperm(neg_candidate_view2_index[i, :].shape[0])[
                                                         :self.train_cfg.intra_cfg.after_select_neg_number]]

                neg_use_view2[i, :] = neg_fine_view2[i, use_index]

            logits_view1 = torch.cat([inner_view1, neg_use_view1], dim=1)
            logits_view2 = torch.cat([inner_view2, neg_use_view2], dim=1)
            logits = torch.cat([logits_view1, logits_view2], dim=0)
            logits = logits / self.train_cfg.intra_cfg.temperature

            labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
            loss = self.criterion(logits, labels)
        out['loss'] = loss
        return out

    def single_intra_coarse_loss(self, coarse_feat, coarse_grid, coarse_valid, *args):  # coarse_grid[B,C,D,H,W] C:y,x,z
        coarse_feat_flatten = coarse_feat.view(coarse_feat.shape[0], 128, -1)
        coarse_feat_flatten = coarse_feat_flatten.permute(1, 0, 2)
        coarse_feat_flatten = coarse_feat_flatten.reshape(128, -1)

        out = dict()

        for i in range(int(coarse_feat.shape[0] / 2)):
            view1_feat_coarse = coarse_feat[2 * i, :, :, :, :].view(128, -1)
            view2_feat_coarse = coarse_feat[2 * i + 1, :, :, :, :].view(128, -1)
            view1_loc_coarse = coarse_grid[2 * i, :, :, :, :].view(3, -1)
            view2_loc_coarse = coarse_grid[2 * i + 1, :, :, :, :].view(3, -1)

            global_use_index_list = [j + 2 * i * view1_feat_coarse.shape[1] for j in
                                     range(0, 2 * view1_feat_coarse.shape[1])]
            global_index_list = [j for j in range(0, coarse_feat.shape[0] * view1_feat_coarse.shape[1])]
            global_search_index_list = list(set(global_index_list) - set(global_use_index_list))
            global_search_index = torch.tensor(global_search_index_list).to(view1_feat_coarse.device)

            with torch.no_grad():
                dist = LA.norm((view1_loc_coarse.view(-1, view1_loc_coarse.shape[1], 1) - \
                                view2_loc_coarse.view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
                dist_z_1 = LA.norm((view1_loc_coarse[2, :].view(-1, view1_loc_coarse.shape[1], 1) - \
                                    view2_loc_coarse[2, :].view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
                dist_z_2 = LA.norm((view2_loc_coarse[2, :].view(-1, view2_loc_coarse.shape[1], 1) - \
                                    view1_loc_coarse[2, :].view(-1, 1, view1_loc_coarse.shape[1])), dim=0)
                dist_1_1 = LA.norm((view1_loc_coarse.view(-1, view1_loc_coarse.shape[1], 1) - \
                                    view1_loc_coarse.view(-1, 1, view1_loc_coarse.shape[1])), dim=0)
                dist_2_2 = LA.norm((view2_loc_coarse.view(-1, view2_loc_coarse.shape[1], 1) - \
                                    view2_loc_coarse.view(-1, 1, view2_loc_coarse.shape[1])), dim=0)
            pos_view1_coarse_can = torch.unique(
                torch.where(dist < self.train_cfg.intra_cfg.coarse_positive_distance)[0])

            key_view1_coarse_can = dist[pos_view1_coarse_can].min(dim=1)[1]
            pos_view1_coarse = pos_view1_coarse_can[dist_z_1[pos_view1_coarse_can, key_view1_coarse_can] < 8]
            key_view1_coarse = key_view1_coarse_can[dist_z_1[pos_view1_coarse_can, key_view1_coarse_can] < 8]
            pos_view2_coarse_can = torch.unique(
                torch.where(dist.transpose(0, 1) < self.train_cfg.intra_cfg.coarse_positive_distance)[0])
            key_view2_coarse_can = dist.transpose(0, 1)[pos_view2_coarse_can].min(dim=1)[1]
            pos_view2_coarse = pos_view2_coarse_can[
                dist_z_2[pos_view2_coarse_can, key_view2_coarse_can] < self.train_cfg.intra_cfg.coarse_z_thres]
            key_view2_coarse = key_view2_coarse_can[
                dist_z_2[pos_view2_coarse_can, key_view2_coarse_can] < self.train_cfg.intra_cfg.coarse_z_thres]

            if pos_view1_coarse.shape[0] > 0:
                q_view1_coarse_feat = view1_feat_coarse[:, pos_view1_coarse].transpose(0, 1)
                k_view1_coarse_feat = view2_feat_coarse[:, key_view1_coarse].transpose(0, 1)
                q_view2_coarse_feat = view2_feat_coarse[:, pos_view2_coarse].transpose(0, 1)
                k_view2_coarse_feat = view1_feat_coarse[:, key_view2_coarse].transpose(0, 1)

                # all_nodes_view1 = []
                # all_nodes_view2 = []
                dist_1_2 = dist.clone()
                dist_2_1 = dist.transpose(0, 1).clone()
                dist_1_2 = dist_1_2[pos_view1_coarse, :]
                dist_1_1 = dist_1_1[pos_view1_coarse, :]
                dist_2_1 = dist_2_1[pos_view2_coarse, :]
                dist_2_2 = dist_2_2[pos_view2_coarse, :]

                ind_ignores_1_1 = torch.where(dist_1_1 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_1_2 = torch.where(dist_1_2 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_2_1 = torch.where(dist_2_1 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_2_2 = torch.where(dist_2_2 < self.train_cfg.intra_cfg.coarse_ignore_distance)
                ind_ignores_1_1 = torch.stack(ind_ignores_1_1)
                ind_ignores_1_2 = torch.stack(ind_ignores_1_2)
                ind_ignores_2_1 = torch.stack(ind_ignores_2_1)
                ind_ignores_2_2 = torch.stack(ind_ignores_2_2)

                use_mask_1_1 = torch.ones_like(dist_1_1)
                use_mask_1_2 = torch.ones_like(dist_1_2)
                use_mask_2_1 = torch.ones_like(dist_2_1)
                use_mask_2_2 = torch.ones_like(dist_2_2)

                use_mask_1_1[ind_ignores_1_1[0, :], ind_ignores_1_1[1, :]] = 0
                use_mask_1_2[ind_ignores_1_2[0, :], ind_ignores_1_2[1, :]] = 0
                use_mask_2_1[ind_ignores_2_1[0, :], ind_ignores_2_1[1, :]] = 0
                use_mask_2_2[ind_ignores_2_2[0, :], ind_ignores_2_2[1, :]] = 0

                neg_mask1 = torch.cat([use_mask_1_1, use_mask_1_2], dim=1)
                neg_mask2 = torch.cat([use_mask_2_1, use_mask_2_2], dim=1)

                inner_view1 = torch.einsum("nc,nc->n", q_view1_coarse_feat, k_view1_coarse_feat).view(-1, 1)
                inner_view2 = torch.einsum("nc,nc->n", q_view2_coarse_feat, k_view2_coarse_feat).view(-1, 1)

                neg_coarse_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                                torch.cat((view1_feat_coarse, view2_feat_coarse), dim=1))
                neg_coarse_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                                torch.cat((view1_feat_coarse, view2_feat_coarse), dim=1))

                neg_coarse_view1 = neg_coarse_view1 * neg_mask1
                neg_coarse_view2 = neg_coarse_view2 * neg_mask2

                neg_candidate_view1_coarse_index = \
                    neg_coarse_view1.topk(self.train_cfg.intra_cfg.coarse_pre_select_neg_number, dim=1)[1]
                neg_candidate_view2_coarse_index = \
                    neg_coarse_view2.topk(self.train_cfg.intra_cfg.coarse_pre_select_neg_number, dim=1)[1]

                neg_use_view1_coarse = torch.zeros(
                    (q_view1_coarse_feat.shape[0], self.train_cfg.intra_cfg.coarse_after_select_neg_number))
                neg_use_view2_coarse = torch.zeros(
                    (q_view2_coarse_feat.shape[0], self.train_cfg.intra_cfg.coarse_after_select_neg_number))

                for i in range(q_view1_coarse_feat.shape[0]):
                    use_index = torch.randperm(neg_candidate_view1_coarse_index.shape[1])[
                                :self.train_cfg.intra_cfg.coarse_after_select_neg_number]
                    neg_use_view1_coarse[i, :use_index.shape[0]] = neg_coarse_view1[i, use_index]

                for i in range(q_view2_coarse_feat.shape[0]):
                    use_index = torch.randperm(neg_candidate_view2_coarse_index.shape[1])[
                                :self.train_cfg.intra_cfg.coarse_after_select_neg_number]
                    neg_use_view2_coarse[i, :use_index.shape[0]] = neg_coarse_view2[i, use_index]

                neg_coarse_global_index_1 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_1 = coarse_feat_flatten[:, neg_coarse_global_index_1]
                neg_coarse_global_index_2 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_2 = coarse_feat_flatten[:, neg_coarse_global_index_2]

                neg_coarse_global_view1 = torch.einsum("nc,ck->nk", q_view1_coarse_feat,
                                                       neg_coarse_global_1)
                neg_coarse_global_view2 = torch.einsum("nc,ck->nk", q_view2_coarse_feat,
                                                       neg_coarse_global_2)

                neg_view1_coarse_concate = torch.cat(
                    [neg_use_view1_coarse.to(view1_feat_coarse.device), neg_coarse_global_view1], dim=1)
                neg_view2_coarse_concate = torch.cat(
                    [neg_use_view2_coarse.to(view1_feat_coarse.device), neg_coarse_global_view2], dim=1)

                logits_view1 = torch.cat([inner_view1, neg_view1_coarse_concate], dim=1)
                logits_view2 = torch.cat([inner_view2, neg_view2_coarse_concate], dim=1)
                logits = torch.cat([logits_view1, logits_view2], dim=0)
                logits = logits / self.train_cfg.intra_cfg.temperature
                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
                loss = self.criterion(logits, labels)
                if out.keys().__len__() == 0:
                    out['loss'] = loss
                else:
                    out['loss'] += loss
            else:
                neg_coarse_global_index_1 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_1 = coarse_feat_flatten[:, neg_coarse_global_index_1]
                neg_coarse_global_index_2 = global_search_index[
                    torch.randperm(global_search_index_list.__len__())[
                    :self.train_cfg.intra_cfg.coarse_global_select_number]]
                neg_coarse_global_2 = coarse_feat_flatten[:, neg_coarse_global_index_2]

                neg_coarse_global_view1 = torch.einsum("nc,ck->nk", view1_feat_coarse.transpose(0, 1),
                                                       neg_coarse_global_1)
                neg_coarse_global_view2 = torch.einsum("nc,ck->nk", view2_feat_coarse.transpose(0, 1),
                                                       neg_coarse_global_2)

                logits_view1 = torch.einsum("nc,ck->nk", view1_feat_coarse.transpose(0, 1),
                                            torch.cat([view1_feat_coarse, view2_feat_coarse], dim=1)).topk(
                    self.train_cfg.intra_cfg.coarse_after_select_neg_number + 1,
                    dim=1)[
                    0]
                logits_view2 = torch.einsum("nc,ck->nk", view2_feat_coarse.transpose(0, 1),
                                            torch.cat([view1_feat_coarse, view2_feat_coarse], dim=1)).topk(
                    self.train_cfg.intra_cfg.coarse_after_select_neg_number + 1,
                    dim=1)[
                    0]

                logits_view1 = torch.cat([logits_view1, neg_coarse_global_view1], dim=1)
                logits_view2 = torch.cat([logits_view2, neg_coarse_global_view2], dim=1)
                logits = torch.cat([logits_view1, logits_view2], dim=0)
                logits = logits / self.train_cfg.intra_cfg.temperature
                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)
                loss = self.criterion(logits, labels)
                if out.keys().__len__() == 0:
                    out['loss'] = loss
                else:
                    out['loss'] += loss

        out['loss'] = out['loss'] / (coarse_feat.shape[0] / 2)
        return out

    def meshgrid3d(self, shape):
        z_ = torch.linspace(0., shape[0] - 1, shape[0])
        y_ = torch.linspace(0., shape[1] - 1, shape[1])
        x_ = torch.linspace(0., shape[2] - 1, shape[2])
        z, y, x = torch.meshgrid(z_, y_, x_)
        return torch.stack((z, y, x), 3)

    async def async_simple_test(self,
                                img,
                                img_meta,
                                proposals=None,
                                rescale=False):
        """Async test without augmentation."""
        assert self.with_bbox, 'Bbox head must be implemented.'
        x = self.extract_feat(img)

        if proposals is None:
            proposal_list = await self.rpn_head.async_simple_test_rpn(
                x, img_meta)
        else:
            proposal_list = proposals

        return await self.roi_head.async_simple_test(
            x, proposal_list, img_meta, rescale=rescale)

    def simple_test(self, img, img_metas, proposals=None, rescale=False):
        """Test without augmentation."""
        x = self.extract_feat(img)
        outs = []
        out1 = x[0].data.cpu().numpy()
        out2 = x[1].data.cpu().numpy()
        out3 = x[2].data.cpu().numpy()
        outs = [out1, out2, out3, img.data.cpu().numpy(), img_metas[0]['filename'].split('.', 1)[0]]
        output_embedding = self.test_cfg.get('output_embedding', True)
        if not output_embedding:
            if not os.path.exists(self.test_cfg.save_path):
                os.mkdir(self.test_cfg.save_path)
            outfilename = self.test_cfg.save_path + \
                          img_metas[0]['filename'].split('.', 1)[0] + '.pkl'
            f = open(outfilename, 'wb')
            pickle.dump(outs, f)
            return [
                x[0][0, 0, 0, 0, 0].data.cpu()]  # we have saved the data into harddisk, this is just for fit the code
        else:
            return outs

    def aug_test(self, imgs, img_metas, **kwargs):
        return self.simple_test(imgs, img_metas, **kwargs)

    def get_mean_vector(self, mask, features):
        labels = torch.unique(mask).tolist()
        labels_organ = [label for label in labels if label > 0]
        mean_vectors = []
        for label in labels_organ:
            all_ind = torch.where(mask == label)[0]
            mean_vector = features[all_ind, 0, :].mean(dim=0)
            mean_vectors.append(mean_vector)
        mean_vectors = torch.stack(mean_vectors)
        labels_organ = torch.tensor(labels_organ)
        return mean_vectors, labels_organ


def log_sinkhorn_iterations(Z: torch.Tensor, log_mu: torch.Tensor, log_nu: torch.Tensor, iters: int) -> torch.Tensor:
    """ Perform Sinkhorn Normalization in Log-space for stability"""
    u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
    for _ in range(iters):
        u = log_mu - torch.logsumexp(Z + v.unsqueeze(1), dim=2)
        v = log_nu - torch.logsumexp(Z + u.unsqueeze(2), dim=1)
    return Z + u.unsqueeze(2) + v.unsqueeze(1)


def log_optimal_transport(scores: torch.Tensor, alpha: torch.Tensor, iters: int) -> torch.Tensor:
    """ Perform Differentiable Optimal Transport in Log-space for stability"""
    b, m, n = scores.shape
    one = scores.new_tensor(1)
    ms, ns = (m * one).to(scores), (n * one).to(scores)

    bins0 = alpha.expand(b, m, 1)
    bins1 = alpha.expand(b, 1, n)
    alpha = alpha.expand(b, 1, 1)

    couplings = torch.cat([torch.cat([scores, bins0], -1),
                           torch.cat([bins1, alpha], -1)], 1)

    norm = - (ms + ns).log()
    log_mu = torch.cat([norm.expand(m), ns.log()[None] + norm])
    log_nu = torch.cat([norm.expand(n), ms.log()[None] + norm])
    log_mu, log_nu = log_mu[None].expand(b, -1), log_nu[None].expand(b, -1)

    Z = log_sinkhorn_iterations(couplings, log_mu, log_nu, iters)
    Z = Z - norm  # multiply probabilities by M+N
    return Z


def arange_like(x, dim: int):
    return x.new_ones(x.shape[dim]).cumsum(0) - 1  # traceable in 1.1

sam/models/frameworks/sam_v2_uniform_cross_volume_mask_dual_head_mean_vector.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from .supconloss import SupConLoss, SupConNegLoss, SupConTopKLoss, SupConMeanLoss, SupConMeanLossCol

sam/models/losses/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
"""
Author: Yonglong Tian (yonglong@mit.edu)
Date: May 07, 2020
"""
from __future__ import print_function

import torch
import torch.nn as nn
from mmdet.models.builder import LOSSES


@LOSSES.register_module()
class SupConLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""

    def __init__(self, temperature=2., contrast_mode='all',
                 base_temperature=0.5):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(features.device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(features.device)
        else:
            mask = mask.float().to(features.device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(features.device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
        mask = mask[(labels > 0)[:, 0], :]
        log_prob = log_prob[(labels > 0)[:, 0], :]
        uses = mask.sum(1) > 0
        mask = mask[uses, :]
        log_prob = log_prob[uses, :]

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, mean_log_prob_pos.shape[0]).mean()

        return loss


@LOSSES.register_module()
class SupConTopKLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""

    def __init__(self, temperature=1., contrast_mode='all',
                 base_temperature=0.5):
        super(SupConTopKLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(features.device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(features.device)
        else:
            mask = mask.float().to(features.device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_feature = anchor_feature[labels.view(-1) > 0, :]
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()
        logits = logits * 10

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(features.device),
            0
        )

        neg_mask = 1 - mask
        mask = mask * logits_mask

        mask = mask[labels.view(-1) > 0, :]
        neg_mask = neg_mask[labels.view(-1) > 0, :]
        logits_mask = logits_mask[labels.view(-1) > 0, :]

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        neg_exp_logits = exp_logits * neg_mask
        v, ind = neg_exp_logits.topk(200)
        mask_topk = torch.zeros_like(mask).scatter_(1, ind, 1)
        all_mask = mask + mask_topk
        all_use_exp_logits = exp_logits * all_mask
        log_prob = logits - torch.log(all_use_exp_logits.sum(1, keepdim=True))
        # mask = mask[(labels>0)[:,0],:]
        # log_prob = log_prob[(labels>0)[:,0],:]
        uses = mask.sum(1) > 0
        mask = mask[uses, :]
        log_prob = log_prob[uses, :]

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, mean_log_prob_pos.shape[0]).mean()

        return loss


@LOSSES.register_module()
class SupConNegLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""

    def __init__(self, temperature=1.0, contrast_mode='all',
                 base_temperature=0.5):
        super(SupConNegLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(features.device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(features.device)
        else:
            mask = mask.float().to(features.device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        # logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        # logits = anchor_dot_contrast - logits_max.detach()
        logits = anchor_dot_contrast

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.zeros_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(features.device),
            1
        )
        # mask = mask * logits_mask

        mask = 1 - mask
        mask = mask + logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits)
        exp_logits = exp_logits * mask
        # v,ind = exp_logits.topk(200)
        # mask_topk = torch.zeros_like(mask).scatter_(1,ind, 1)
        # exp_logits = exp_logits * mask_topk
        # import matplotlib.pyplot as plt
        # plt.figure(figsize=(20, 20))
        # plt.imshow(mask_topk.data.cpu())
        # plt.show()
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
        # mask = mask[(labels>0)[:,0],:]
        # log_prob = log_prob[(labels>0)[:,0],:]
        # uses = mask.sum(1)>0
        # mask = mask[uses,:]
        # log_prob = log_prob[uses,:]

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (logits_mask * log_prob).sum(1) / logits_mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, mean_log_prob_pos.shape[0]).mean()

        return loss


@LOSSES.register_module()
class SupConMeanLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""

    def __init__(self, temperature=0.01, contrast_mode='all',
                 base_temperature=1.):
        super(SupConMeanLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels, prototypes, prototypes_labels):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        labels = labels.view(-1, 1)
        prototypes_labels = prototypes_labels.view(-1, 1).to(labels.device)
        mask_p_f = torch.eq(prototypes_labels, labels.T).float().to(features.device)

        features = features.squeeze(1)
        # compute logits
        prototypes_dot_features = torch.div(
            torch.matmul(prototypes, features.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(prototypes_dot_features, dim=1, keepdim=True)
        logits = prototypes_dot_features - logits_max.max().detach()
        # logits = prototypes_dot_features

        # compute log_prob
        exp_logits = torch.exp(logits)
        # v,ind = exp_logits.topk(200)
        # mask_topk = torch.zeros_like(mask).scatter_(1,ind, 1)
        # exp_logits = exp_logits * mask_topk
        # import matplotlib.pyplot as plt
        # plt.figure(figsize=(20, 20))
        # plt.imshow(mask_topk.data.cpu())
        # plt.show()
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
        # mask = mask[(labels>0)[:,0],:]
        # log_prob = log_prob[(labels>0)[:,0],:]
        uses = mask_p_f.sum(1) > 0
        mask_p_f = mask_p_f[uses, :]
        log_prob = log_prob[uses, :]

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask_p_f * log_prob).sum(1) / mask_p_f.sum(1)

        # loss
        loss = - self.base_temperature * mean_log_prob_pos
        loss = loss.view(1, mean_log_prob_pos.shape[0]).mean()
        return loss


@LOSSES.register_module()
class SupConMeanLossCol(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""

    def __init__(self, temperature=0.1, contrast_mode='all',
                 base_temperature=1.0):
        super(SupConMeanLossCol, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels, prototypes, prototypes_labels):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        labels = labels.view(-1, 1).to(features.device)
        prototypes_labels = prototypes_labels.view(-1, 1).to(features.device)
        # mask_p_f = torch.eq(prototypes_labels,labels.T).float().to(features.device)

        mask_f_p = torch.eq(labels, prototypes_labels.T).float().to(features.device)

        features = features.squeeze(1)
        # compute logits
        features_dot_prototypes = torch.div(
            torch.matmul(features, prototypes.T),
            self.temperature)
        # prototypes_dot_features = torch.div(
        #     torch.matmul(prototypes,features.T),
        #     self.temperature)
        # for numerical stability
        # logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        # logits = anchor_dot_contrast - logits_max.detach()
        logits = features_dot_prototypes

        # compute log_prob
        exp_logits = torch.exp(logits)
        # v,ind = exp_logits.topk(200)
        # mask_topk = torch.zeros_like(mask).scatter_(1,ind, 1)
        # exp_logits = exp_logits * mask_topk
        # import matplotlib.pyplot as plt
        # plt.figure(figsize=(20, 20))
        # plt.imshow(mask_topk.data.cpu())
        # plt.show()
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
        # mask = mask[(labels>0)[:,0],:]
        # log_prob = log_prob[(labels>0)[:,0],:]
        # uses = mask.sum(1)>0
        # mask = mask[uses,:]
        # log_prob = log_prob[uses,:]

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask_f_p * log_prob).sum(1) / mask_f_p.sum(1)

        # loss
        loss = - mean_log_prob_pos
        loss = loss.view(1, mean_log_prob_pos.shape[0]).mean()

        return loss

sam/models/losses/supconloss.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule
from mmcv.runner import BaseModule, auto_fp16

from mmdet.models.builder import NECKS


@NECKS.register_module()
class FPN3d(BaseModule):
    r"""Feature Pyramid Network.

    This is an implementation of paper `Feature Pyramid Networks for Object
    Detection <https://arxiv.org/abs/1612.03144>`_.

    Args:
        in_channels (List[int]): Number of input channels per scale.
        out_channels (int): Number of output channels (used at each scale)
        num_outs (int): Number of output scales.
        start_level (int): Index of the start input backbone level used to
            build the feature pyramid. Default: 0.
        end_level (int): Index of the end input backbone level (exclusive) to
            build the feature pyramid. Default: -1, which means the last level.
        add_extra_convs (bool | str): If bool, it decides whether to add conv
            layers on top of the original feature maps. Default to False.
            If True, it is equivalent to `add_extra_convs='on_input'`.
            If str, it specifies the source feature map of the extra convs.
            Only the following options are allowed

            - 'on_input': Last feat map of neck inputs (i.e. backbone feature).
            - 'on_lateral':  Last feature map after lateral convs.
            - 'on_output': The last output feature map after fpn convs.
        relu_before_extra_convs (bool): Whether to apply relu before the extra
            conv. Default: False.
        no_norm_on_lateral (bool): Whether to apply norm on lateral.
            Default: False.
        conv_cfg (dict): Config dict for convolution layer. Default: None.
        norm_cfg (dict): Config dict for normalization layer. Default: None.
        act_cfg (str): Config dict for activation layer in ConvModule.
            Default: None.
        upsample_cfg (dict): Config dict for interpolate layer.
            Default: `dict(mode='nearest')`
        init_cfg (dict or list[dict], optional): Initialization config dict.

    Example:
        >>> import torch
        >>> in_channels = [2, 3, 5, 7]
        >>> scales = [340, 170, 84, 43]
        >>> inputs = [torch.rand(1, c, s, s)
        ...           for c, s in zip(in_channels, scales)]
        >>> self = FPN(in_channels, 11, len(in_channels)).eval()
        >>> outputs = self.forward(inputs)
        >>> for i in range(len(outputs)):
        ...     print(f'outputs[{i}].shape = {outputs[i].shape}')
        outputs[0].shape = torch.Size([1, 11, 340, 340])
        outputs[1].shape = torch.Size([1, 11, 170, 170])
        outputs[2].shape = torch.Size([1, 11, 84, 84])
        outputs[3].shape = torch.Size([1, 11, 43, 43])
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 num_outs,
                 start_level=0,
                 end_level=-1,
                 add_extra_convs=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'),
                 init_cfg=dict(
                     type='Xavier', layer='Conv3d', distribution='uniform')):
        super(FPN3d, self).__init__(init_cfg)
        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
            self.add_extra_convs = 'on_input'

        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,
                inplace=False)

            self.lateral_convs.append(l_conv)
        fpn_conv = ConvModule(
            out_channels,
            out_channels,
            3,
            padding=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg,
            inplace=False)
        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,
                    inplace=False)
                self.fpn_convs.append(extra_fpn_conv)

    @auto_fp16()
    def forward(self, inputs):
        """Forward function."""
        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] += F.interpolate(laterals[i],
                                                 **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] += F.interpolate(
                    laterals[i], size=prev_shape, **self.upsample_cfg)

        # build outputs
        # part 1: from original levels
        outs = [
            self.fpn_convs[0](laterals[0])
        ]
        # 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(F.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](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)

sam/models/necks/fpn3d.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
import re

from setuptools import find_packages, setup


def get_version():
    version_file = "sam/version.py"
    with open(version_file, "r") as f:
        exec(compile(f.read(), version_file, "exec"))
    return locals()["__version__"]


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

    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__":
    install_requires = parse_requirements()
    setup(
        name="sam",
        version=get_version(),
        description="SAM",
        author="Xiaoyu Bai",
        author_email="not available",
        packages=find_packages(exclude=("configs", "tools")),
        install_requires=install_requires,
        include_package_data=True,
        ext_modules=[],
        zip_safe=False,
    )

setup.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
from demo_semantic_stable_points import fixed_point_iterations
tools/__init__.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import os
import torchio as tio
import torch
import pickle as pkl

# test_path = '/data/sdd/user/rawdata/CT-MRI-Abd/nii_regis_final_test'
test_path = '/data/sde/user/CT-MRI-ABD-resample_padding-test/'
all_cases = os.listdir(test_path)

all_cases = [name.split('_', 1)[0] for name in all_cases if '_c2f_affine_mask.nii.gz' in name]
dis_all = []
dis_x = []
dis_y = []
dis_z = []
all = []
for case in all_cases:
    print(case)
    dis_case = []
    mr_mask = tio.LabelMap(os.path.join(test_path, case + '_0001_mask.nii.gz'))
    ct_mask = tio.LabelMap(os.path.join(test_path, case + '_0000_mask.nii.gz'))
    mr_data = mr_mask.data[0, :, :, :]
    ct_data = ct_mask.data[0, :, :, :]
    # ct_data = torch.round(ct_data)
    ct_labels = torch.unique(ct_data).tolist()
    # mr_data = torch.round(mr_data)
    # mr_labels = torch.unique(mr_data).tolist()
    labels = torch.unique(mr_data).tolist()
    print(labels, ct_labels)
    labels = [label for label in labels if label > 0]
    for label in labels:
        mr_loc = torch.where(mr_data == label)
        mr_loc = torch.stack(mr_loc)
        mr_loc = torch.mean(mr_loc.type(torch.float32), dim=1)
        ct_loc = torch.where(ct_data == label)
        ct_loc = torch.stack(ct_loc)
        if ct_loc.shape[1] == 0:
            continue
        ct_loc = torch.mean(ct_loc.type(torch.float32), dim=1)
        all.append([label, mr_loc, ct_loc])
        dis_all.append(torch.linalg.norm(mr_loc - ct_loc))
        dis_case.append(torch.linalg.norm(mr_loc - ct_loc))
        dis_x.append(torch.linalg.norm(mr_loc[0] - ct_loc[0]))
        dis_y.append(torch.linalg.norm(mr_loc[1] - ct_loc[1]))
        dis_z.append(torch.linalg.norm(mr_loc[2] - ct_loc[2]))
f = open('/data/sde/user/init_abd_results.pkl', 'wb')
pkl.dump(all, f)
dis_all = torch.tensor(dis_all)
dis_x = torch.tensor(dis_x)
dis_y = torch.tensor(dis_y)
dis_z = torch.tensor(dis_z)
print(dis_all.mean(), dis_all.std())
print(dis_x.mean(), dis_x.std())
print(dis_y.mean(), dis_y.std())
print(dis_z.mean(), dis_z.std())

tools/affine_evaluation.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# single process tumor tracking evaluation and visualization

import torch
import torchio as tio
import pickle
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import torch.nn.functional as F
import json
from utils import make_local_grid_patch
from interfaces import get_sim_embed_loc_multi_embedding_space, get_sim_semantic_embed_loc_multi_embedding_space
import statsmodels.api as sm

data_path = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
prediction_path = '/data/sdd/user/results/result-dlt/'
save_path = '/data/sdd/user/visual/visual-dlt_jimmy/'
if not os.path.exists(save_path):
    os.mkdir(save_path)
all_dis = []
all_dis_x = []
all_dis_y = []
all_dis_z = []
draw_results = False

anno_path = '/data/sdd/user/rawdata/DLT-main/data/'
gths = json.load(open(os.path.join(anno_path, 'test.json'), 'r'))
all_keys = list(gths.keys())
all_img = []
for i in range(all_keys.__len__()):
    element = gths[all_keys[i]]
    query_img_name = element['source']
    key_img_name = element['target']
    if query_img_name not in all_img:
        all_img.append(query_img_name)
    if key_img_name not in all_img:
        all_img.append(key_img_name)
import csv

with open("/data/sdd/user/processed_data/deeplesion_dlt_jimmy.csv", "w") as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["name"])
    for nii_name in all_img:
        writer.writerow([nii_name])

normed_spacing = (2., 2., 2.)
gpu_id = 0
cuda_device = torch.device('cuda:' + str(gpu_id))
count = 0
hit = 0
miss = 0
hit_10 = 0
miss_10 = 0
for ii in range(all_keys.__len__()):
    print(ii)
    element = gths[all_keys[ii]]
    count = count + 1
    query_img_name = element['source'].split('.', 1)[0]
    query_bbx = element[
        'source box']  # [x1,y1,z1,x2-x1,y2-y1,z2-z1] z index is the inverse of the sever deeplesion dataset
    query_center = element['source center']
    key_img_name = element['target'].split('.', 1)[0]
    key_bbx = element['target box']
    key_center = element['target center']
    query_img_info = pickle.load(open(prediction_path + query_img_name + '.pkl', 'rb'))
    key_img_info = pickle.load(open(prediction_path + key_img_name + '.pkl', 'rb'))

    query_img = tio.ScalarImage(os.path.join(data_path, query_img_name + '.nii.gz'))
    query_img_origin_spacing = query_img.spacing
    query_img_origin_shape = query_img.shape
    query_spacing_norm_ratio = np.array(query_img_origin_spacing) / np.array(normed_spacing)
    query_bbx_reshaped = np.array(query_bbx).reshape(2, 3)
    query_bbx_reshaped[1, :] = query_bbx_reshaped[1, :] + query_bbx_reshaped[0, :]
    query_bbx_reshaped[:, 2] = query_img.shape[3] - query_bbx_reshaped[:, 2]
    query_bbx_reshaped[:, 2] = query_bbx_reshaped[::-1, 2]
    query_bbx_normed = query_bbx_reshaped * query_spacing_norm_ratio
    query_center[2] = query_img_origin_shape[3] - query_center[2] - 1

    query_point_fine = (query_bbx_reshaped[1, :] + query_bbx_reshaped[0, :]) / 2 * query_spacing_norm_ratio
    query_center_fine = query_center * query_spacing_norm_ratio
    query_point_fine = query_center_fine

    key_img = tio.ScalarImage(os.path.join(data_path, key_img_name + '.nii.gz'))
    key_img_origin_spacing = key_img.spacing
    key_img_origin_shape = key_img.shape
    key_spacing_norm_ratio = np.array(key_img_origin_spacing) / np.array(normed_spacing)
    key_bbx_reshaped = np.array(key_bbx).reshape(2, 3)
    key_bbx_reshaped[1, :] = key_bbx_reshaped[1, :] + key_bbx_reshaped[0, :]

    key_bbx_reshaped[:, 2] = key_img.shape[3] - key_bbx_reshaped[:, 2]
    key_bbx_reshaped[:, 2] = key_bbx_reshaped[::-1, 2]
    key_bbx_normed = key_bbx_reshaped * key_spacing_norm_ratio
    # key_bbx_normed = np.floor(key_bbx_normed).astype(int)
    key_center[2] = key_img_origin_shape[3] - key_center[2] - 1

    key_point_fine_gt = (key_bbx_reshaped[1, :] + key_bbx_reshaped[0, :]) / 2 * key_spacing_norm_ratio
    # key_point_fine_gt = key_point_fine_gt.astype(int)
    key_center_fine = key_center * key_spacing_norm_ratio
    key_point_fine_gt = key_center_fine

    query_point_fine_re = np.floor(query_point_fine / 2.).astype(int)
    query_point_fine_re_matrix = make_local_grid_patch(query_point_fine_re, 1, 1, 1, query_img_info[0].shape[4],
                                                       query_img_info[0].shape[3], query_img_info[0].shape[2])

    iters = 4
    for j in range(iters):
        if j == 0:
            pts_query = query_point_fine_re_matrix
            pts_query_final = pts_query
        if j % 2 == 0:
            pts_key_inter, s_inter, pts_key, s = get_sim_semantic_embed_loc_multi_embedding_space(query_img_info,
                                                                                                  key_img_info,
                                                                                                  pts_query,
                                                                                                  cuda_device,
                                                                                                  return_interploted=True)
            pts_key_final = pts_key
            pts_key_final_inter = pts_key_inter
        else:
            pts_key_inter, s_inter, pts_key, s = get_sim_semantic_embed_loc_multi_embedding_space(key_img_info,
                                                                                                  query_img_info,
                                                                                                  pts_query,
                                                                                                  cuda_device,
                                                                                                  return_interploted=True)
            pts_query_final = pts_key
            pts_query_final_inter = pts_key_inter
        pts_query = pts_key
    dis = np.linalg.norm(pts_query_final_inter - query_point_fine,
                         axis=1)
    final_q = []
    final_k = []
    scores = []
    for i in range(dis.shape[0]):
        if dis[i] < 100 and s[i] > 1.3:
            final_q.append(pts_query_final_inter[i, :])
            final_k.append(pts_key_final_inter[i, :])
            scores.append(s[i])
    final_q = np.array(final_q)
    final_k = np.array(final_k)
    shift = final_q - query_point_fine
    ones = np.ones((final_q.shape[0], 1))

    final_q_ext = np.concatenate((final_q, ones), axis=1)
    m = []
    w = [0, 1, 2]
    for i in w:
        mod = sm.RLM(final_k[:, i], final_q_ext)
        res = mod.fit()
        m.append(res.params)
    m.append(np.array([0, 0, 0, 1]))
    m = np.array(m)

    shift_affine = np.matmul(m[:3, :3], shift.T).transpose(1, 0)
    pt_k_shift = final_k - shift_affine
    # pt_k_shift = pt_k_shift.astype(int)
    pt_k_shift = np.where(pt_k_shift < 0, 0, pt_k_shift)
    max_shape = np.array(key_img_info[3].shape[2:]) - 1
    max_shape = np.array([max_shape[1], max_shape[2], max_shape[0]])
    pt_k_shift = np.where(pt_k_shift > max_shape, max_shape, pt_k_shift)
    pt_k_shift_mean = np.mean(pt_k_shift, axis=0)
    pt_k_shift_std = np.std(pt_k_shift, axis=0)
    use_ind = np.abs((pt_k_shift[:, 2] - pt_k_shift_mean[2]) < 1)
    pt_k_shift_mean_r = np.mean(pt_k_shift[use_ind, :], axis=0)

    pre = pt_k_shift_mean
    pre = pre + 0.5
    scores = np.array(scores)
    score = scores.mean()
    pre_remap = pre / key_spacing_norm_ratio
    final = key_bbx_normed - pre

    dis = np.linalg.norm((key_point_fine_gt - pre)) * 2.
    dis_x = np.linalg.norm((key_point_fine_gt[0] - pre[0])) * 2.
    dis_y = np.linalg.norm((key_point_fine_gt[1] - pre[1])) * 2.
    dis_z = np.linalg.norm((key_point_fine_gt[2] - pre[2])) * 2.
    all_dis.append(dis)
    all_dis_x.append(dis_x)
    all_dis_y.append(dis_y)
    all_dis_z.append(dis_z)
    final = key_bbx_reshaped - pre_remap
    if (final[0, :] <= 0).all() and (final[1, :] > -0).all():
        hit = hit + 1
    else:
        miss = miss + 1

    if dis <= 10:
        hit_10 = hit_10 + 1
    else:
        miss_10 = miss_10 + 1

    if draw_results:  # set True to draw matching results
        query_point_fine = np.floor(query_point_fine).astype(int)
        q_img = query_img_info[3][0, 0, :, :, :]

        q_img = (q_img + 50)
        q_img[q_img < 51] = 51
        q_img[q_img > 76.2] = 76.2
        q_img = (q_img - 51) / (76.2 - 51)

        slice = q_img[query_point_fine[2], :, :]
        # q_img[query_point_fine[1] - 1:query_point_fine[1] + 2,
        # query_point_fine[0] - 1:query_point_fine[0] + 2] = 1
        fig, ax = plt.subplots(3, 3, figsize=(33, 18))
        fig.suptitle('distance:%.2f' % dis + 'score:%.2f' % score)
        # plt.subplot(1,3,1)
        ax[0, 0].set_title('query')
        ax[0, 0].imshow(slice, cmap='gray')
        ax[0, 0].plot((query_point_fine[0]), (query_point_fine[1]), 'o', markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)

        rect = patches.Rectangle((query_bbx_normed[0, 0], query_bbx_normed[0, 1]),
                                 query_bbx_normed[1, 0] - query_bbx_normed[0, 0],
                                 query_bbx_normed[1, 1] - query_bbx_normed[0, 1]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[0, 0].add_patch(rect)

        slice = q_img[:, query_point_fine[1], :]
        slice = slice[::-1, :]
        ax[1, 0].set_title('query')
        ax[1, 0].imshow(slice, cmap='gray')
        ax[1, 0].plot((query_point_fine[0]), (q_img.shape[0] - query_point_fine[2] - 1), 'o',
                      markerfacecolor='none', markeredgecolor="red",
                      markersize=12, markeredgewidth=2)

        rect = patches.Rectangle((query_bbx_normed[0, 0], q_img.shape[0] - query_bbx_normed[1, 2] - 1),
                                 query_bbx_normed[1, 0] - query_bbx_normed[0, 0],
                                 query_bbx_normed[1, 2] - query_bbx_normed[0, 2]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[1, 0].add_patch(rect)

        # plt(x,y...)

        slice = q_img[:, :, query_point_fine[0]]
        slice = slice[::-1, :]
        ax[2, 0].set_title('query')
        ax[2, 0].imshow(slice, cmap='gray')
        ax[2, 0].plot((query_point_fine[1]), (q_img.shape[0] - query_point_fine[2] - 1), 'o',
                      markerfacecolor='none', markeredgecolor="red",
                      markersize=12, markeredgewidth=2)
        rect = patches.Rectangle((query_bbx_normed[0, 1], q_img.shape[0] - query_bbx_normed[1, 2] - 1),
                                 query_bbx_normed[1, 1] - query_bbx_normed[0, 1],
                                 query_bbx_normed[1, 2] - query_bbx_normed[0, 2]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[2, 0].add_patch(rect)

        key_point_fine_gt = np.floor(key_point_fine_gt).astype(int)

        k_img_gt = key_img_info[3][0, 0, :, :, :]

        k_img_gt = (k_img_gt + 50)
        k_img_gt[k_img_gt < 51] = 51
        k_img_gt[k_img_gt > 76.2] = 76.2
        k_img_gt = (k_img_gt - 51) / (76.2 - 51)
        # k_img_gt[key_point_fine_gt[1] - 1:key_point_fine_gt[1] + 2,
        # key_point_fine_gt[0] - 1:key_point_fine_gt[0] + 2] = 1
        # plt.subplot(1,3,2)
        # plt.title('key_gt')
        # plt.imshow(k_img_gt, cmap='gray')
        # plt.savefig(save_path + pair_id + '_' + 'key_gth' + '.png')

        slice = k_img_gt[key_point_fine_gt[2], :, :]

        ax[0, 1].set_title('key_gt')
        ax[0, 1].imshow(slice, cmap='gray')
        ax[0, 1].plot((key_point_fine_gt[0]), (key_point_fine_gt[1]), 'o', markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)
        rect = patches.Rectangle((key_bbx_normed[0, 0], key_bbx_normed[0, 1]),
                                 key_bbx_normed[1, 0] - key_bbx_normed[0, 0],
                                 key_bbx_normed[1, 1] - key_bbx_normed[0, 1]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[0, 1].add_patch(rect)

        slice = k_img_gt[:, key_point_fine_gt[1], :]
        slice = slice[::-1, :]

        ax[1, 1].set_title('key_gt')
        ax[1, 1].imshow(slice, cmap='gray')
        ax[1, 1].plot((key_point_fine_gt[0]), (k_img_gt.shape[0] - key_point_fine_gt[2] - 1), 'o',
                      markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)
        rect = patches.Rectangle((key_bbx_normed[0, 0], k_img_gt.shape[0] - key_bbx_normed[1, 2] - 1),
                                 key_bbx_normed[1, 0] - key_bbx_normed[0, 0],
                                 key_bbx_normed[1, 2] - key_bbx_normed[0, 2]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[1, 1].add_patch(rect)

        slice = k_img_gt[:, :, key_point_fine_gt[0]]
        slice = slice[::-1, :]

        ax[2, 1].set_title('key_gt')
        ax[2, 1].imshow(slice, cmap='gray')
        ax[2, 1].plot((key_point_fine_gt[1]), (k_img_gt.shape[0] - key_point_fine_gt[2] - 1), 'o',
                      markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)
        rect = patches.Rectangle((key_bbx_normed[0, 1], k_img_gt.shape[0] - key_bbx_normed[1, 2] - 1),
                                 key_bbx_normed[1, 1] - key_bbx_normed[0, 1],
                                 key_bbx_normed[1, 2] - key_bbx_normed[0, 2]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[2, 1].add_patch(rect)

        key_img_pre = key_img_info[3][0, 0, :, :, :]
        key_img_pre = (key_img_pre + 50)
        key_img_pre[key_img_pre < 51] = 51
        key_img_pre[key_img_pre > 76.2] = 76.2
        key_img_pre = (key_img_pre - 51) / (76.2 - 51)
        # key_img_pre[ind[1][0] - 1:ind[1][0] + 2,
        # ind[2][0] - 1:ind[2][0] + 2] = 1
        # plt.subplot(1,3,3)
        # plt.title('key_pre')
        # plt.imshow(key_img_pre, cmap='gray')

        pre = np.floor(pre).astype('int')
        slice = key_img_pre[pre[2], :, :]
        ax[0, 2].set_title('key_pre')
        ax[0, 2].imshow(slice, cmap='gray')
        ax[0, 2].plot(pre[0], pre[1], 'o', markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)

        # also draw gt bbox
        rect = patches.Rectangle((key_bbx_normed[0, 0], key_bbx_normed[0, 1]),
                                 key_bbx_normed[1, 0] - key_bbx_normed[0, 0],
                                 key_bbx_normed[1, 1] - key_bbx_normed[0, 1]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[0, 2].add_patch(rect)

        slice = key_img_pre[:, pre[1], :]
        slice = slice[::-1, :]
        ax[1, 2].set_title('key_pre')
        ax[1, 2].imshow(slice, cmap='gray')
        ax[1, 2].plot(pre[0], (key_img_pre.shape[0] - pre[2] - 1),
                      'o', markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)

        # also draw gt bbox
        rect = patches.Rectangle((key_bbx_normed[0, 0], k_img_gt.shape[0] - key_bbx_normed[1, 2] - 1),
                                 key_bbx_normed[1, 0] - key_bbx_normed[0, 0],
                                 key_bbx_normed[1, 2] - key_bbx_normed[0, 2]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[1, 2].add_patch(rect)

        slice = key_img_pre[:, :, pre[0]]
        slice = slice[::-1, :]
        ax[2, 2].set_title('key_pre')
        ax[2, 2].imshow(slice, cmap='gray')
        ax[2, 2].plot(pre[1], (key_img_pre.shape[0] - pre[2] - 1),
                      'o', markerfacecolor='none',
                      markeredgecolor="red",
                      markersize=12, markeredgewidth=2)

        # also draw gt bbox
        rect = patches.Rectangle((key_bbx_normed[0, 1], k_img_gt.shape[0] - key_bbx_normed[1, 2] - 1),
                                 key_bbx_normed[1, 1] - key_bbx_normed[0, 1],
                                 key_bbx_normed[1, 2] - key_bbx_normed[0, 2]
                                 , linewidth=2, edgecolor='g',
                                 facecolor='none')
        ax[2, 2].add_patch(rect)

        plt.tight_layout()

        # plt.show()
        # plt.close()
        plt.savefig(save_path + str(all_keys[ii]) + '_' + str('%.2f' % dis) + '_stable' + '.png')
        plt.close()
print(count, hit, miss)
print(hit / count)
print(count, hit_10, miss_10)
print(hit_10 / count)
print(np.array(all_dis).mean(), np.array(all_dis).std())
print(np.array(all_dis_x).mean(), np.array(all_dis_x).std())
print(np.array(all_dis_y).mean(), np.array(all_dis_y).std())
print(np.array(all_dis_z).mean(), np.array(all_dis_z).std())

tools/deeplesion_eval_metric_visuall_multipoints_dlt_stable.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import numpy as np
import os
import matplotlib.pyplot as plt
import time

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from interfaces import init, get_embedding, get_sim_embed_loc, normalize
from utils import read_image, visualize


os.chdir(os.path.join(os.path.dirname(__file__), os.pardir))  # go to root dir of this project
config_file = 'configs/sam/sam_NIHLN.py'
checkpoint_file = 'checkpoints/SAM.pth'

# assume all input image are in torchio "LPS+" direction which equal to "RAI" orientation in ITK-Snap.
im1_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_001.nii.gz'
im2_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_002.nii.gz'


def get_random_query_point(im):
    im1_shape = list(im['shape'][1:])
    num_try = 0
    bg_th = im['img'].min() + (im['img'].max() - im['img'].min())/10
    while num_try < 20:
        p = np.random.randint(im1_shape)
        if im['img'][p[0], p[1], p[2]] > bg_th:  # not in air
            break
        num_try += 1
    print('query point', p, ', intensity', im['img'][p[0], p[1], p[2]])
    return p


def main():
    time1 = time.time()
    model = init(config_file, checkpoint_file)
    time2 = time.time()
    print('model loading time:', time2 - time1)
    im1, normed_im1, norm_info_1 = read_image(im1_file, is_MRI=False)
    im2, normed_im2, norm_info_2 = read_image(im2_file, is_MRI=False)
    time3 = time.time()
    print('image loading time:', time3 - time2)

    pt1 = get_random_query_point(im1)

    emb1 = get_embedding(normed_im1, model)
    emb2 = get_embedding(normed_im2, model)
    time4 = time.time()
    print('embeddings computing time:', time4 - time3)

    pt1_normed = np.array(pt1) * norm_info_1
    pt2_normed, score = get_sim_embed_loc(emb1, emb2, pt1_normed,
                                          (im2['shape'][3], im2['shape'][1], im2['shape'][2]), norm_info=norm_info_2,
                                          write_sim=False, use_sim_coarse=True)
    pt2 = np.array(pt2_normed).astype(int)
    print(pt2, score)
    time5 = time.time()
    print('matching point computing time:', time5 - time4)
    visualize(im1['img'], im2['img'], norm_info_1, norm_info_2, pt1, pt2, score)


if __name__ == '__main__':
    main()

tools/demo.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import numpy as np
import os

from tools.demo import get_random_query_point

from interfaces import init, get_embedding, get_sim_embed_loc, normalize
from utils import read_image, visualize
import matplotlib.pyplot as plt
import time

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.chdir(os.path.join(os.path.dirname(__file__), os.pardir))  # go to root dir of this project
config_file = 'configs/samv2/samv2_ct_mri_infer.py'
checkpoint_file = 'checkpoints/CrossSAM_iter_20000_ct_mr_inter_1.pth'

# assume all input image are in torchio "LPS+" direction which equal to "RAI" orientation in ITK-Snap.
im1_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_001.nii.gz'; im1_isMRI = False
im2_file = 'data/raw_data/MR/T2WI.nii.gz'; im2_isMRI = True  # T2
# im2_file = 'data/raw_data/MR/AbdomenMRCT_0002_0000.nii.gz'; im2_isMRI = True  # T1
im1_file, im2_file, im1_isMRI, im2_isMRI = im2_file, im1_file, im2_isMRI, im1_isMRI  # swap image 1 and 2

# pt1 = (110, 57, 111)


def main():
    time1 = time.time()
    model = init(config_file, checkpoint_file)
    time2 = time.time()
    print('model loading time:', time2 - time1)
    im1, normed_im1, norm_info_1 = read_image(im1_file, is_MRI=im1_isMRI)
    im2, normed_im2, norm_info_2 = read_image(im2_file, is_MRI=im2_isMRI)
    time3 = time.time()
    print('image loading time:', time3 - time2)
    pt1 = get_random_query_point(im1)

    emb1 = get_embedding(normed_im1, model)
    emb2 = get_embedding(normed_im2, model)
    time4 = time.time()
    print('embeddings computing time:', time4 - time3)

    pt1_normed = np.array(pt1) * norm_info_1
    pt2_normed, score = get_sim_embed_loc(emb1, emb2, pt1_normed,
                                          (im2['shape'][3], im2['shape'][1], im2['shape'][2]),
                                          norm_info=norm_info_2, write_sim=False, use_sim_coarse=True)
    pt2 = np.array(pt2_normed).astype(int)
    print('matched point', pt2, ', similarity', score)
    time5 = time.time()
    print('matching point computing time:', time5 - time4)
    visualize(im1['img'], im2['img'], norm_info_1, norm_info_2, pt1, pt2, score,
              im1_isMRI=im1_isMRI, im2_isMRI=im2_isMRI)


if __name__ == '__main__':
    main()

tools/demo_ct_mr.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import numpy as np
import os
import matplotlib.pyplot as plt
import time

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from interfaces import init, get_embedding, get_sim_embed_semantic_loc, normalize
from utils import read_image, visualize
from demo import get_random_query_point


os.chdir(os.path.join(os.path.dirname(__file__), os.pardir))  # go to root dir of this project
config_file = 'configs/samv2/samv2_NIHLN.py'
checkpoint_file = 'checkpoints/SAMv2_iter_20000.pth'

# assume all input image are in torchio "LPS+" direction which equal to "RAI" orientation in ITK-Snap.
im1_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_001.nii.gz'
im2_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_002.nii.gz'


def main():
    time1 = time.time()
    model = init(config_file, checkpoint_file)
    time2 = time.time()
    print('model loading time:', time2 - time1)
    im1, normed_im1, norm_info_1 = read_image(im1_file, is_MRI=False)
    im2, normed_im2, norm_info_2 = read_image(im2_file, is_MRI=False)
    time3 = time.time()
    print('image loading time:', time3 - time2)

    pt1 = get_random_query_point(im1)

    emb1 = get_embedding(normed_im1, model)
    emb2 = get_embedding(normed_im2, model)
    time4 = time.time()
    print('embeddings computing time:', time4 - time3)

    pt1_normed = np.array(pt1) * norm_info_1
    pt2_normed, score = get_sim_embed_semantic_loc(emb1, emb2, pt1_normed,
                                                   (im2['shape'][3], im2['shape'][1], im2['shape'][2]),
                                                   norm_info=norm_info_2,
                                                   write_sim=False, use_sim_coarse=True)
    pt2 = np.array(pt2_normed).astype(int)
    print(pt2, score)
    time5 = time.time()
    print('matching point computing time:', time5 - time4)
    visualize(im1['img'], im2['img'], norm_info_1, norm_info_2, pt1, pt2, score)


if __name__ == '__main__':
    main()

tools/demo_semantic.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# compute on the original embeddings not interpolate to image size
import os
import numpy as np
import torch
import time
import SimpleITK as sitk
import statsmodels.api as sm
import matplotlib.pyplot as plt

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from interfaces import init, get_embedding, get_sim_embed_loc_multi_embedding_space
from utils import read_image, visualize, make_local_grid_patch
from interfaces import get_sim_semantic_embed_loc_multi_embedding_space
from demo import get_random_query_point

os.chdir(os.path.join(os.path.dirname(__file__), os.pardir))  # go to root dir of this project
config_file = 'configs/samv2/samv2_NIHLN.py'
checkpoint_file = 'checkpoints/SAMv2_iter_20000.pth'

# assume all input image are in torchio "LPS+" direction which equal to "RAI" orientation in ITK-Snap.
im1_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_001.nii.gz'
im2_file = 'data/raw_data/NIH_lymph_node/ABD_LYMPH_002.nii.gz'


def main():
    time1 = time.time()
    model = init(config_file, checkpoint_file)
    time2 = time.time()
    print('model loading time:', time2 - time1)
    im1, normed_im1, norm_info_1 = read_image(im1_file)
    im2, normed_im2, norm_info_2 = read_image(im2_file)
    time3 = time.time()
    print('image loading time:', time3 - time2)

    pt1 = get_random_query_point(im1)

    emb1 = get_embedding(normed_im1, model)
    emb2 = get_embedding(normed_im2, model)
    time4 = time.time()
    print('embeddings computing time:', time4 - time3)

    # fixed_point_iterations is slower but more accurate than simple nearest neighbor matching
    score, pt2 = fixed_point_iterations(emb1, norm_info_1, emb2, norm_info_2, im2['shape'], pt1)
    print(pt2, score)
    visualize(im1['img'], im2['img'], norm_info_1, norm_info_2, pt1, pt2, score)


def fixed_point_iterations(emb1, norm_info_1, emb2, norm_info_2, im2_shape, pt1, x_margin=1, y_margin=1, z_margin=1,
                           iterations=6):
    print('performing fixed point iteration')
    pt_query = np.round(pt1).astype('int')
    pt_query_center = np.floor(pt_query * norm_info_1 * 0.5).astype(int)
    # crop a cube region centered in the template point
    pt_query_normed = make_local_grid_patch(pt_query_center, x_margin, y_margin, z_margin,
                                            emb1[0].shape[4], emb1[0].shape[3], emb1[0].shape[2])
    cuda_device = torch.device('cuda:' + str(0))

    for i in range(iterations):
        if i % 2 == 0:
            if i > 0:
                pt_query_normed = np.floor(pt_query * norm_info_1 * 0.5).astype(int)
            pt_key_normed, score = get_sim_semantic_embed_loc_multi_embedding_space(emb1, emb2, pt_query_normed,
                                                                                    cuda_device)
            pt_key = np.round((pt_key_normed * 2 + 0.5) / norm_info_2).astype(int)
            pt_k_final = pt_key
        else:
            pt_query_normed = np.floor(pt_query * norm_info_2 * 0.5).astype(int)
            pt_key_normed, score = get_sim_semantic_embed_loc_multi_embedding_space(emb2, emb1, pt_query_normed,
                                                                                    cuda_device)
            pt_key = np.round((pt_key_normed * 2 + 0.5) / norm_info_1).astype(int)
            pt_q_final = pt_key
            pt_final_score = score
        print(f'score:{score.mean(), score.max()}')
        pt_query = pt_key
    dis = np.linalg.norm((pt_q_final - np.array(pt1)) * norm_info_1, axis=1)
    final_q_all = []
    final_k_all = []
    final_score_all = []
    for i in range(dis.shape[0]):
        if dis[i] < 100 and pt_final_score[i] > 1.3:
            final_q_all.append(pt_q_final[i, :])
            final_k_all.append(pt_k_final[i, :])
            final_score_all.append(pt_final_score[i])

    final_q_all = np.array(final_q_all)
    final_k_all = np.array(final_k_all)
    final_score_all = np.array(final_score_all)
    print(final_q_all.shape[0])

    final_k_list = list([tuple(t) for t in final_k_all])
    final_k_count = dict((k_ele, final_k_list.count(k_ele)) for k_ele in final_k_list)

    final_q = []
    final_k = []

    for keys in final_k_count:
        ind_k = np.where(np.logical_and(np.logical_and(final_k_all[:, 0] == keys[0], final_k_all[:, 1] == keys[1]),
                                        final_k_all[:, 2] == keys[2]))[0]
        s = final_score_all[ind_k]
        q = final_q_all[ind_k, :]
        final_q.append(q[0, :])
        final_k.append(keys)
    final_q = np.array(final_q)
    final_k = np.array(final_k)

    k_xs = np.unique(final_k[:, 0])
    k_ys = np.unique(final_k[:, 1])
    k_zs = np.unique(final_k[:, 2])
    q_zs = np.unique(final_q[:, 2])

    shift = final_q - pt1
    final_q_mm = final_q * norm_info_1
    final_k_mm = final_k * norm_info_2
    shift_mm = shift * norm_info_1
    # pt_q_shift = final_q - shift

    ones = np.ones((final_q.shape[0], 1))
    final_q_ext = np.concatenate((final_q, ones), axis=1)
    m = []

    if k_zs.shape[0] == 1:
        z_ratro = norm_info_1[2] / norm_info_2[2]
        z_bias = k_zs - q_zs * z_ratro
        k = [0, 1]
        for i in k:
            mod = sm.RLM(final_k[:, i],
                         np.concatenate([final_q_ext[:, 0:1], final_q_ext[:, 1:2], final_q_ext[:, 3:]], axis=1))
            res = mod.fit()
            m.append(res.params)
        m.append([0, 0, z_bias[0]])
        m.append(np.array([0, 0, 1]))
        m = np.array(m)
        m = np.concatenate([m[:, :2], np.array([0, 0, z_ratro, 0]).reshape(-1, 1), m[:, 2:]], axis=1)

    else:
        k = [0, 1, 2]
        for i in k:
            mod = sm.RLM(final_k[:, i], final_q_ext[:, ])
            res = mod.fit()
            m.append(res.params)
        m.append(np.array([0, 0, 0, 1]))
        m = np.array(m)

    shift_affine = np.matmul(m[:3, :3], shift.T).transpose(1, 0)
    pt_k_shift = final_k - shift_affine
    pt_k_shift = pt_k_shift.astype(int)
    pt_k_shift = np.where(pt_k_shift < 0, 0, pt_k_shift)
    max_shape = np.array(im2_shape[1:]) - 1
    max_shape = np.array([max_shape[1], max_shape[0], max_shape[2]])
    pt_k_shift = np.where(pt_k_shift > max_shape, max_shape, pt_k_shift)
    pt_k_shift_mean = np.mean(pt_k_shift, axis=0)
    pt_k_shift_mean = pt_k_shift_mean.astype('int')

    return final_score_all.mean(), pt_k_shift_mean


if __name__ == '__main__':
    main()

tools/demo_semantic_stable_points.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import argparse
import numpy as np
import matplotlib.pyplot as plt
import SimpleITK as sitk
import mmcv
import torch
from mmcv import Config, DictAction
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
                         wrap_fp16_model)

from mmdet.models import build_detector
from sam import *
import torch.nn.functional as F


def init(config, checkpoint):
    cfg = Config.fromfile(config)
    cfg.gpu_ids = range(1)
    model = build_detector(cfg.model, test_cfg=cfg.get('test_cfg'))
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
    model = MMDataParallel(model, device_ids=cfg.gpu_ids)
    model.eval()
    return model


def get_embedding(batched_data, model):
    with torch.no_grad():
        result = model(return_loss=False, rescale=True, **batched_data)
    return result


def normalize(pt, imshape):
    pt = np.array(pt)
    imshape = np.array(imshape)
    unit_space = 2 / np.array(imshape)
    origin = unit_space / 2 - 1
    pt_normed = origin + pt * unit_space
    return pt_normed


# infer matching point of the given (one) query_point (x,y,z). the query_point is on the norm image
# the final output matching point's location is depend on the imshape parameter.
# i.e. you can set the imshape equals to the original pre-normed key image size, then you can get the
# mapping point on the original key image
def get_sim_embed_loc(query_img_info, key_img_info, query_point, imshape, norm_info, write_sim=False, return_sim=False,
                      use_sim_coarse=True):
    query_point = np.array(query_point)
    cuda_device = torch.device('cuda:' + str(0))
    query_point_fine_re = np.floor(query_point / 2.).astype(int)
    fine_query = torch.tensor(query_img_info[0]).to(cuda_device)

    coarse_query = torch.tensor(query_img_info[1]).to(cuda_device)

    coarse_query = F.interpolate(coarse_query, fine_query.shape[2:], mode='trilinear', align_corners=False)
    coarse_query = F.normalize(coarse_query, dim=1)

    fine_key = torch.tensor(key_img_info[0]).to(cuda_device)
    coarse_key = torch.tensor(key_img_info[1]).to(cuda_device)
    coarse_key = F.interpolate(coarse_key, fine_key.shape[2:], mode='trilinear', align_corners=False)
    coarse_key = F.normalize(coarse_key, dim=1)

    query_fine = fine_query[0, :, query_point_fine_re[2], query_point_fine_re[1], query_point_fine_re[0]].view(-1,
                                                                                                               128)
    key_fine = fine_key[0, :, :, :, :].reshape(128, -1)

    query_coarse = coarse_query[0, :, query_point_fine_re[2], query_point_fine_re[1], query_point_fine_re[0]].view(
        -1, 128)
    key_coarse = coarse_key[0, :, :, :, :].reshape(128, -1)

    sim_fine = torch.einsum("nc,ck->nk", query_fine, key_fine)
    sim_coarse = torch.einsum("nc,ck->nk", query_coarse, key_coarse)

    sim_fine = sim_fine.reshape(fine_key.shape[2:])
    sim_coarse = sim_coarse.reshape(coarse_key.shape[2:])
    sim_fine = sim_fine.view(1, 1, sim_fine.shape[0], sim_fine.shape[1], sim_fine.shape[2])
    sim_coarse = sim_coarse.view(1, 1, sim_coarse.shape[0], sim_coarse.shape[1], sim_coarse.shape[2])
    if use_sim_coarse:
        sim = (sim_fine + sim_coarse) / 2
    else:
        sim = sim_fine
        # sim = sim_coarse
    # s im = F.interpolate(sim, tuple(key_img_info[2].shape[2:]), mode='trilinear', align_corners=False)
    sim = F.interpolate(sim, imshape, mode='trilinear', align_corners=False)
    if return_sim:
        return sim
    sim = sim[0, 0, :, :, :]
    ind = torch.where(sim == sim.max())
    z = ind[0][0].data.cpu().numpy()
    y = ind[1][0].data.cpu().numpy()
    x = ind[2][0].data.cpu().numpy()
    if write_sim:
        img = torch.tensor(key_img_info[2]).to(cuda_device)
        img = F.interpolate(img, imshape, mode='trilinear', align_corners=False)
        img_array = img[0, 0, :, :, :].data.cpu().numpy().astype('float32')
        img_sitk = sitk.GetImageFromArray(img_array)
        img_sitk.SetSpacing(norm_info)
        sitk.WriteImage(img_sitk, 'img.nii.gz')

        sim_map_array = sim.data.cpu().numpy().astype('float32')
        sim_map_array = sim_map_array
        sim_map = sitk.GetImageFromArray(sim_map_array)
        sim_map.SetSpacing(norm_info)
        sitk.WriteImage(sim_map, 'sim_map.nii.gz')

    # plt.imshow(torch.flip(sim[:,:,x].data.cpu(),(0,)).numpy())
    # plt.savefig('sim_2.png')
    return [x, y, z], sim.max().data.cpu().numpy()


def get_sim_embed_loc_multi_embedding_space(query_img_info, key_img_info, query_points):
    query_point = np.array(query_points)
    cuda_device = torch.device('cuda:' + str(0))
    fine_query = torch.tensor(query_img_info[0]).to(cuda_device)

    coarse_query = torch.tensor(query_img_info[1]).to(cuda_device)

    coarse_query = F.interpolate(coarse_query, fine_query.shape[2:], mode='trilinear', align_corners=False)
    coarse_query = F.normalize(coarse_query, dim=1)

    fine_key = torch.tensor(key_img_info[0]).to(cuda_device)
    coarse_key = torch.tensor(key_img_info[1]).to(cuda_device)
    coarse_key = F.interpolate(coarse_key, fine_key.shape[2:], mode='trilinear', align_corners=False)
    coarse_key = F.normalize(coarse_key, dim=1)

    query_fine = fine_query[0, :, query_points[:, 2], query_points[:, 1],
                 query_points[:, 0]].transpose(1, 0)

    key_fine = fine_key[0, :, :, :, :].reshape(128, -1)

    query_coarse = coarse_query[0, :, query_points[:, 2], query_points[:, 1],
                   query_points[:, 0]].transpose(1, 0)
    key_coarse = coarse_key[0, :, :, :, :].reshape(128, -1)

    sim_fine = torch.einsum("nc,ck->nk", query_fine, key_fine)
    sim_coarse = torch.einsum("nc,ck->nk", query_coarse, key_coarse)

    sim_fine = sim_fine.view(query_points.shape[0], 1, fine_key.shape[2], fine_key.shape[3],
                             fine_key.shape[4])
    sim_coarse = sim_coarse.view(query_points.shape[0], 1, coarse_key.shape[2], coarse_key.shape[3],
                                 coarse_key.shape[4])
    sim_all = (sim_fine + sim_coarse) / 2
    locs = []
    scores = []

    for i in range(sim_all.shape[0]):
        sim = sim_all[i, 0, :, :, :]
        ind = torch.where(sim == sim.max())
        z = ind[0][0].data.cpu().numpy()
        y = ind[1][0].data.cpu().numpy()
        x = ind[2][0].data.cpu().numpy()
        scores.append(sim.max().data.cpu().numpy())
        locs.append([x, y, z])
    return np.array(locs).astype(int), np.array(scores)


def get_sim_embed_semantic_loc(query_img_info, key_img_info, query_point, imshape, norm_info, write_sim=False,
                               return_sim=False, use_sim_coarse=True):
    query_point = np.array(query_point)
    cuda_device = torch.device('cuda:' + str(0))
    query_point_fine_re = np.floor(query_point / 2.).astype(int)
    fine_query = torch.tensor(query_img_info[0]).to(cuda_device)
    seman_query = torch.tensor(query_img_info[2]).to(cuda_device)

    coarse_query = torch.tensor(query_img_info[1]).to(cuda_device)

    coarse_query = F.interpolate(coarse_query, fine_query.shape[2:], mode='trilinear', align_corners=False)
    coarse_query = F.normalize(coarse_query, dim=1)

    fine_key = torch.tensor(key_img_info[0]).to(cuda_device)
    coarse_key = torch.tensor(key_img_info[1]).to(cuda_device)
    seman_key = torch.tensor(key_img_info[2]).to(cuda_device)
    coarse_key = F.interpolate(coarse_key, fine_key.shape[2:], mode='trilinear', align_corners=False)
    coarse_key = F.normalize(coarse_key, dim=1)

    query_fine = fine_query[0, :, query_point_fine_re[2], query_point_fine_re[1], query_point_fine_re[0]].view(-1,
                                                                                                               128)
    query_sem = seman_query[0, :, query_point_fine_re[2], query_point_fine_re[1], query_point_fine_re[0]].view(-1,
                                                                                                               128)
    key_fine = fine_key[0, :, :, :, :].reshape(128, -1)
    key_sem = seman_key[0, :, :, :, :].reshape(128, -1)

    query_coarse = coarse_query[0, :, query_point_fine_re[2], query_point_fine_re[1], query_point_fine_re[0]].view(
        -1, 128)
    key_coarse = coarse_key[0, :, :, :, :].reshape(128, -1)

    sim_fine = torch.einsum("nc,ck->nk", query_fine, key_fine)
    sim_sem = torch.einsum("nc,ck->nk", query_sem, key_sem)
    sim_coarse = torch.einsum("nc,ck->nk", query_coarse, key_coarse)

    sim_fine = sim_fine.reshape(fine_key.shape[2:])
    sim_sem = sim_sem.reshape(fine_key.shape[2:])
    sim_coarse = sim_coarse.reshape(coarse_key.shape[2:])
    sim_fine = sim_fine.view(1, 1, sim_fine.shape[0], sim_fine.shape[1], sim_fine.shape[2])
    sim_sem = sim_sem.view(1, 1, sim_sem.shape[0], sim_sem.shape[1], sim_sem.shape[2])
    sim_coarse = sim_coarse.view(1, 1, sim_coarse.shape[0], sim_coarse.shape[1], sim_coarse.shape[2])
    if use_sim_coarse:
        sim = (sim_fine + sim_coarse + sim_sem) / 3
        # sim = sim_sem
    else:
        sim = (sim_sem + sim_fine) / 2
        # sim = sim_coarse
    # sim = F.interpolate(sim, tuple(key_img_info[2].shape[2:]), mode='trilinear', align_corners=False)
    sim = F.interpolate(sim, imshape, mode='trilinear', align_corners=False)
    if return_sim:
        return sim
    sim = sim[0, 0, :, :, :]
    ind = torch.where(sim == sim.max())
    z = ind[0][0].data.cpu().numpy()
    y = ind[1][0].data.cpu().numpy()
    x = ind[2][0].data.cpu().numpy()
    if write_sim:
        img = torch.tensor(key_img_info[3]).to(cuda_device)
        img = F.interpolate(img, imshape, mode='trilinear', align_corners=False)
        img_array = img[0, 0, :, :, :].data.cpu().numpy().astype('float32')
        img_sitk = sitk.GetImageFromArray(img_array)
        img_sitk.SetSpacing(norm_info)
        sitk.WriteImage(img_sitk, 'img.nii.gz')

        sim_map_array = sim.data.cpu().numpy().astype('float32')
        sim_map_array = sim_map_array
        sim_map = sitk.GetImageFromArray(sim_map_array)
        sim_map.SetSpacing(norm_info)
        sitk.WriteImage(sim_map, 'sim_map.nii.gz')

    # plt.imshow(torch.flip(sim[:,:,x].data.cpu(),(0,)).numpy())
    # plt.savefig('sim_2.png')
    return [x, y, z], sim.max().data.cpu().numpy()


def get_sim_semantic_embed_loc_multi_embedding_space(query_img_info, key_img_info, query_points, cuda_device,
                                                     return_interploted=False):
    # cuda_device = torch.device('cuda:' + str(0))
    fine_query = torch.tensor(query_img_info[0]).to(cuda_device)

    coarse_query = torch.tensor(query_img_info[1]).to(cuda_device)
    sem_query = torch.tensor(query_img_info[2]).to(cuda_device)

    coarse_query = F.interpolate(coarse_query, fine_query.shape[2:], mode='trilinear', align_corners=False)
    coarse_query = F.normalize(coarse_query, dim=1)

    fine_key = torch.tensor(key_img_info[0]).to(cuda_device)
    coarse_key = torch.tensor(key_img_info[1]).to(cuda_device)
    sem_key = torch.tensor(key_img_info[2]).to(cuda_device)
    coarse_key = F.interpolate(coarse_key, fine_key.shape[2:], mode='trilinear', align_corners=False)
    coarse_key = F.normalize(coarse_key, dim=1)

    query_fine = fine_query[0, :, query_points[:, 2], query_points[:, 1],
                 query_points[:, 0]].transpose(1, 0)

    key_fine = fine_key[0, :, :, :, :].reshape(128, -1)

    query_coarse = coarse_query[0, :, query_points[:, 2], query_points[:, 1],
                   query_points[:, 0]].transpose(1, 0)
    key_coarse = coarse_key[0, :, :, :, :].reshape(128, -1)

    query_sem = sem_query[0, :, query_points[:, 2], query_points[:, 1],
                query_points[:, 0]].transpose(1, 0)
    key_sem = sem_key[0, :, :, :, :].reshape(128, -1)

    sim_fine = torch.einsum("nc,ck->nk", query_fine, key_fine)
    sim_coarse = torch.einsum("nc,ck->nk", query_coarse, key_coarse)
    sim_sem = torch.einsum("nc,ck->nk", query_sem, key_sem)

    sim_fine = sim_fine.view(query_points.shape[0], 1, fine_key.shape[2], fine_key.shape[3],
                             fine_key.shape[4])
    sim_coarse = sim_coarse.view(query_points.shape[0], 1, coarse_key.shape[2], coarse_key.shape[3],
                                 coarse_key.shape[4])
    sim_sem = sim_sem.view(query_points.shape[0], 1, fine_key.shape[2], fine_key.shape[3],
                           fine_key.shape[4])
    sim_all = (sim_fine + sim_coarse + sim_sem) / 3
    # sim = sim_coarse
    if return_interploted:
        sim_inter = F.interpolate(sim_all, tuple(key_img_info[3].shape[2:]), mode='trilinear', align_corners=False)
        # sim_inter = sim_inter[:, 0, :, :, :]
        locs_inter = []
        scores_inter = []
        for i in range(sim_inter.shape[0]):
            sim_here = sim_inter[i, 0, :, :, :]
            ind = torch.where(sim_here == sim_here.max())
            z = ind[0][0].data.cpu().numpy()
            y = ind[1][0].data.cpu().numpy()
            x = ind[2][0].data.cpu().numpy()
            scores_inter.append(sim_here.max().data.cpu().numpy())
            locs_inter.append([x, y, z])
    locs = []
    scores = []

    # ind = torch.argmax(sim, dim=1).cpu().numpy()
    # zyx = np.unravel_index(ind, fine_key_vol.shape[2:])
    # xyz = np.vstack(zyx)[::-1] * cfg.local_emb_stride + .5  # add .5 to closer to stride center
    # xyz = xyz.T / key_data['im_norm_ratio']
    # xyz = np.minimum(np.round(xyz.astype(int)), np.array(key_data['im_shape'])[::-1] - 1)
    for i in range(sim_all.shape[0]):
        sim = sim_all[i, 0, :, :, :]
        ind = torch.where(sim == sim.max())
        z = ind[0][0].data.cpu().numpy()
        y = ind[1][0].data.cpu().numpy()
        x = ind[2][0].data.cpu().numpy()
        scores.append(sim.max().data.cpu().numpy())
        locs.append([x, y, z])
    if return_interploted:
        return np.array(locs_inter).astype(int), np.array(scores_inter), np.array(locs).astype(int), np.array(scores)
    else:
        return np.array(locs).astype(int), np.array(scores)


def get_sim_embed_loc_multi_embedding_space(query_img_info, key_img_info, query_points, cuda_device,
                                            return_interploted=False):
    fine_query = torch.tensor(query_img_info[0]).to(cuda_device)
    coarse_query = torch.tensor(query_img_info[1]).to(cuda_device)

    coarse_query = F.interpolate(coarse_query, fine_query.shape[2:], mode='trilinear', align_corners=False)
    coarse_query = F.normalize(coarse_query, dim=1)

    fine_key = torch.tensor(key_img_info[0]).to(cuda_device)
    coarse_key = torch.tensor(key_img_info[1]).to(cuda_device)
    coarse_key = F.interpolate(coarse_key, fine_key.shape[2:], mode='trilinear', align_corners=False)
    coarse_key = F.normalize(coarse_key, dim=1)

    query_fine = fine_query[0, :, query_points[:, 2], query_points[:, 1],
                 query_points[:, 0]].transpose(1, 0)
    key_fine = fine_key[0, :, :, :, :].reshape(128, -1)

    query_coarse = coarse_query[0, :, query_points[:, 2], query_points[:, 1],
                   query_points[:, 0]].transpose(1, 0)
    key_coarse = coarse_key[0, :, :, :, :].reshape(128, -1)
    sim_fine = torch.einsum("nc,ck->nk", query_fine, key_fine)
    sim_coarse = torch.einsum("nc,ck->nk", query_coarse, key_coarse)

    sim_fine = sim_fine.view(query_points.shape[0], 1, fine_key.shape[2], fine_key.shape[3],
                             fine_key.shape[4])
    sim_coarse = sim_coarse.view(query_points.shape[0], 1, coarse_key.shape[2], coarse_key.shape[3],
                                 coarse_key.shape[4])
    sim_all = (sim_fine + sim_coarse) / 2
    # sim = sim_coarse
    if return_interploted:
        sim_inter = F.interpolate(sim_all, tuple(key_img_info[2].shape[2:]), mode='trilinear', align_corners=False)
        # sim_inter = sim_inter[:, 0, :, :, :]
        locs_inter = []
        scores_inter = []
        for i in range(sim_inter.shape[0]):
            sim_here = sim_inter[i, 0, :, :, :]
            ind = torch.where(sim_here == sim_here.max())
            z = ind[0][0].data.cpu().numpy()
            y = ind[1][0].data.cpu().numpy()
            x = ind[2][0].data.cpu().numpy()
            scores_inter.append(sim_here.max().data.cpu().numpy())
            locs_inter.append([x, y, z])
    locs = []
    scores = []

    # ind = torch.argmax(sim, dim=1).cpu().numpy()
    # zyx = np.unravel_index(ind, fine_key_vol.shape[2:])
    # xyz = np.vstack(zyx)[::-1] * cfg.local_emb_stride + .5  # add .5 to closer to stride center
    # xyz = xyz.T / key_data['im_norm_ratio']
    # xyz = np.minimum(np.round(xyz.astype(int)), np.array(key_data['im_shape'])[::-1] - 1)
    for i in range(sim_all.shape[0]):
        sim = sim_all[i, 0, :, :, :]
        ind = torch.where(sim == sim.max())
        z = ind[0][0].data.cpu().numpy()
        y = ind[1][0].data.cpu().numpy()
        x = ind[2][0].data.cpu().numpy()
        scores.append(sim.max().data.cpu().numpy())
        locs.append([x, y, z])
    if return_interploted:
        return np.array(locs_inter).astype(int), np.array(scores_inter), np.array(locs).astype(int), np.array(scores)
    else:
        return np.array(locs).astype(int), np.array(scores)

tools/interfaces.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torch
import torchio as tio
import pickle
import numpy as np
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
import matplotlib.pyplot as plt
import torch.nn.functional as F
# import ipdb
from time import time, sleep

from multiprocessing import Pool, current_process, Queue, Event
from multiprocessing import set_start_method, current_process
from interfaces import get_sim_semantic_embed_loc_multi_embedding_space, get_sim_embed_loc_multi_embedding_space
from utils import make_local_grid_patch
import statsmodels.api as sm

gth_path = '/data/sdd/user/processed_data/ind_files/DL_lesion_pairs_1022_correct.pth'
count = 0
hit = 0
miss = 0
NUM_GPUS = 2
PROC_PER_GPU = 4
margin_x, margin_y, margin_z = 1, 1, 1
stable_matching = True
semantic = True


def infer_element(element):
    gpu_id = (int(current_process().name.split('-')[-1]) - 1) % NUM_GPUS
    torch.cuda.set_device(f'cuda:{gpu_id}')
    cuda_device = torch.device('cuda:' + str(gpu_id))
    data_path = '/data/sdd/user/rawdata/Deeplesion/Images_nifti/'
    prediction_path = '/data/sdd/user/results/result-dlt/'
    normed_spacing = (2., 2., 2.)
    # ------------------------------------image information------------------------------------
    query_img_name = element[0]
    query_bbx = element[1]
    key_img_name = element[2]
    key_bbx = element[3]
    pair_id = element[4]
    query_img_info = pickle.load(open(prediction_path + query_img_name + '.pkl', 'rb'))
    key_img_info = pickle.load(open(prediction_path + key_img_name + '.pkl', 'rb'))

    query_img = tio.ScalarImage(os.path.join(data_path, query_img_name + '.nii.gz'))
    query_img_origin_spacing = query_img.spacing
    query_spacing_norm_ratio = np.array(query_img_origin_spacing) / np.array(normed_spacing)
    query_bbx_reshaped = np.array(query_bbx).reshape(2, 3)
    query_bbx_reshaped[:, 2] = query_img.shape[3] - query_bbx_reshaped[:, 2]
    query_bbx_reshaped[:, 2] = query_bbx_reshaped[::-1, 2]
    query_point_fine = (query_bbx_reshaped[1, :] + query_bbx_reshaped[0, :]) / 2 * query_spacing_norm_ratio

    key_img = tio.ScalarImage(os.path.join(data_path, key_img_name + '.nii.gz'))
    key_img_origin_spacing = key_img.spacing
    key_spacing_norm_ratio = np.array(key_img_origin_spacing) / np.array(normed_spacing)
    key_bbx_reshaped = np.array(key_bbx).reshape(2, 3)
    key_bbx_reshaped[:, 2] = key_img.shape[3] - key_bbx_reshaped[:, 2]
    key_bbx_reshaped[:, 2] = key_bbx_reshaped[::-1, 2]
    key_bbx_normed = key_bbx_reshaped * key_spacing_norm_ratio
    key_point_fine_gt = (key_bbx_reshaped[1, :] + key_bbx_reshaped[0, :]) / 2 * key_spacing_norm_ratio

    query_point_fine_re = np.floor(query_point_fine / 2.).astype(int)
    query_point_fine_re_matrix = make_local_grid_patch(query_point_fine_re, margin_x, margin_y, margin_z,
                                                       query_img_info[0].shape[4],
                                                       query_img_info[0].shape[3], query_img_info[0].shape[2])

    if stable_matching == False:
        pts_query = query_point_fine_re_matrix
        if semantic:
            pts_key_final_inter, s_inter, pts_key, s = get_sim_semantic_embed_loc_multi_embedding_space(query_img_info,
                                                                                                        key_img_info,
                                                                                                        pts_query,
                                                                                                        cuda_device,
                                                                                                        return_interploted=True)
        else:
            pts_key_final_inter, s_inter, pts_key, s = get_sim_embed_loc_multi_embedding_space(query_img_info,
                                                                                               key_img_info,
                                                                                               pts_query,
                                                                                               cuda_device,
                                                                                               return_interploted=True)
        all = np.array(pts_key_final_inter)
        pre = np.array([all[:, 0].mean(), all[:, 1].mean(), all[:, 2].mean()])

    if stable_matching:
        iters = 4
        for ii in range(iters):
            if ii == 0:
                pts_query = query_point_fine_re_matrix
            if ii % 2 == 0:
                if semantic:
                    pts_key_inter, s_inter, pts_key, s = get_sim_semantic_embed_loc_multi_embedding_space(
                        query_img_info, key_img_info, pts_query, cuda_device, return_interploted=True)
                else:
                    pts_key_inter, s_inter, pts_key, s = get_sim_embed_loc_multi_embedding_space(
                        query_img_info, key_img_info, pts_query, cuda_device, return_interploted=True)
                pts_key_final_inter = pts_key_inter
                if ii == 0:
                    pts_key_final_inter_direct = pts_key_inter
            else:
                if semantic:
                    pts_key_inter, s_inter, pts_key, s = get_sim_semantic_embed_loc_multi_embedding_space(
                        key_img_info, query_img_info, pts_query, cuda_device, return_interploted=True)
                else:
                    pts_key_inter, s_inter, pts_key, s = get_sim_embed_loc_multi_embedding_space(
                        key_img_info, query_img_info, pts_query, cuda_device, return_interploted=True)
                pts_query_final_inter = pts_key_inter
            pts_query = pts_key

        dis = np.linalg.norm(pts_query_final_inter - query_point_fine,
                             axis=1)
        final_q_all = []
        final_k_all = []
        final_score_all = []

        for i in range(dis.shape[0]):
            if dis[i] < 100 and s_inter[i] > 1.3:
                final_q_all.append(pts_query_final_inter[i, :])
                final_k_all.append(pts_key_final_inter[i, :])
                final_score_all.append(s_inter[i])

        final_q_all = np.array(final_q_all)
        final_k_all = np.array(final_k_all)
        final_score_all = np.array(final_score_all)
        final_k_list = list([tuple(t) for t in final_k_all])
        final_k_count = dict((k_ele, final_k_list.count(k_ele)) for k_ele in final_k_list)

        final_q = []
        final_k = []

        for keys in final_k_count:
            ind_k = np.where(np.logical_and(np.logical_and(final_k_all[:, 0] == keys[0], final_k_all[:, 1] == keys[1]),
                                            final_k_all[:, 2] == keys[2]))[0]
            s = final_score_all[ind_k]
            q = final_q_all[ind_k, :]
            final_q.append(q[0, :])
            final_k.append(keys)
        final_q = np.array(final_q).reshape(-1,3)
        final_k = np.array(final_k).reshape(-1,3)

        k_xs = np.unique(final_k[:, 0])
        k_ys = np.unique(final_k[:, 1])
        k_zs = np.unique(final_k[:, 2])
        q_zs = np.unique(final_q[:, 2])

        if  k_zs.shape[0] == 1:
            print('z:',pair_id)
        if  k_ys.shape[0] == 1:
            print('y:',pair_id)
        if  k_xs.shape[0] == 1:
            print('x:',pair_id)

        shift = final_q - query_point_fine
        ones = np.ones((final_q.shape[0], 1))
        final_q_ext = np.concatenate((final_q, ones), axis=1)
        m = []
        if k_zs.shape[0] == 1:
            z_ratro = 1.
            z_bias = k_zs - q_zs * z_ratro
            k = [0, 1]
            for i in k:
                mod = sm.RLM(final_k[:, i],
                             np.concatenate([final_q_ext[:, 0:1], final_q_ext[:, 1:2], final_q_ext[:, 3:]], axis=1))
                res = mod.fit()
                m.append(res.params)
            m.append([0, 0, z_bias[0]])
            m.append(np.array([0, 0, 1]))
            m = np.array(m)
            m = np.concatenate([m[:, :2], np.array([0, 0, z_ratro, 0]).reshape(-1, 1), m[:, 2:]], axis=1)


        else:
            k = [0, 1, 2]
            for i in k:
                mod = sm.RLM(final_k[:, i], final_q_ext[:, ])
                res = mod.fit()
                m.append(res.params)
            m.append(np.array([0, 0, 0, 1]))
            m = np.array(m)
        shift_affine = np.matmul(m[:3, :3], shift.T).transpose(1, 0)
        pt_k_shift = final_k - shift_affine
        pt_k_shift = np.where(pt_k_shift < 0, 0, pt_k_shift)
        if semantic:
            max_shape = np.array(key_img_info[3].shape[2:]) - 1
        else:
            max_shape = np.array(key_img_info[2].shape[2:]) - 1
        max_shape = np.array([max_shape[1], max_shape[2], max_shape[0]])
        pt_k_shift = np.where(pt_k_shift > max_shape, max_shape, pt_k_shift)
        pt_k_shift_mean = np.mean(pt_k_shift, axis=0)
        pre_s = pt_k_shift_mean
        all_direct = np.array(pts_key_final_inter_direct)
        pre_direct = np.array([all_direct[:, 0].mean(), all_direct[:, 1].mean(), all_direct[:, 2].mean()])
        dis_s_d = np.linalg.norm(pre_s - pre_direct)
        # if dis_s_d > 10:
        #     pre = pre_direct
        # else:
        #     pre = pre_s
        pre = pre_s
    pre = pre + 0.5
    pre_remap = pre / key_spacing_norm_ratio
    dis = np.linalg.norm(key_point_fine_gt - pre)
    dis_x = np.linalg.norm(key_point_fine_gt[0] - pre[0])
    dis_y = np.linalg.norm(key_point_fine_gt[1] - pre[1])
    dis_z = np.linalg.norm(key_point_fine_gt[2] - pre[2])
    flag = 0
    final = key_bbx_reshaped - pre_remap
    if (final[0, :] <= 0).all() and (final[1, :] > -0).all():
        flag = 1
    else:
        flag = 0

    final = [flag, dis, dis_x, dis_y, dis_z,pair_id]
    return final


def infer_element_init(q):
    infer_element.q = q


def run_pool():
    set_start_method('spawn')
    queue = Queue()
    for gpu_ids in range(NUM_GPUS):
        for _ in range(PROC_PER_GPU):
            queue.put(gpu_ids)
    gths = torch.load(gth_path)
    all_keys = list(gths.keys())
    all_elements = []
    for i in range(all_keys.__len__()):
        for element in gths[all_keys[i]]:
            all_elements.append(element)
    with Pool(processes=PROC_PER_GPU * NUM_GPUS) as pool:
        allresults = pool.map(infer_element, all_elements)
    return allresults


if __name__ == '__main__':
    start = time()
    results = run_pool()
    hits = 0
    miss = 0
    dis = []
    dis_x = []
    dis_y = []
    dis_z = []
    all_pair = []
    for result in results:
        if result[0] == 1:
            hits += 1
        else:
            miss += 1
        dis.append(result[1])
        dis_x.append(result[2])
        dis_y.append(result[3])
        dis_z.append(result[4])
        all_pair.append(result[5])
    dis = np.array(dis) * 2.
    dis_x = np.array(dis_x) * 2.
    dis_y = np.array(dis_y) * 2.
    dis_z = np.array(dis_z) * 2.
    all_pair = np.array(all_pair)
    print(hits / (hits + miss))
    print(dis.mean(), dis.std(), dis.max())
    print((dis <= 10).sum() / (hits + miss))
    print(dis_x.mean(), dis_x.std(), dis_x.max())
    print(dis_y.mean(), dis_y.std(), dis_y.max())
    print(dis_z.mean(), dis_z.std(), dis_z.max())
    end = time()
    print(end - start)
    print(all_pair[dis>60])
    print(dis[dis > 60])

tools/multi_infer_all_in_one.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
# Written by Xiaoyu Bai
# This code first uses SAM to do rigid registration on two images, and then uses DEEDS for deformable fine-tuning.
# The pipeline is:
# (1) compute SAM embeddings on 2x2x2 mm resampled images and find grid points matchings.
# (2) q,k are template points and their corresponding points on query image. q_o,k_o are points location
# remapped to original image space.
# (3) Using robust rigid estimator to compute the rigid transform matrix on resample images and original images (and do transform).
# (4) Do DEEDS deformation on rigid registered resampled images.
# (5) Apply the deformable field on original rigid registered images.

# for cross-modality training, it contains several stages of registration
# the first registration is based on the intra-training ckpts, and its results  are used for training the next stage SAM


import os

import SimpleITK as sitk
import nibabel as nib
import numpy as np
import torch
import torch.nn.functional as F
import torchio as tio
from numpy.linalg import inv
from scipy import ndimage
from scipy import ndimage as ndi

from interfaces import get_sim_embed_loc_multi_embedding_space
from interfaces import init, get_embedding
from misc.deeds_preprocess_no_mask_origin_image import padding_img
from utils import read_image

np.set_printoptions(suppress=True)
from scipy.spatial.transform import Rotation as ro

config_file = '/data/sdd/user/SAM/configs/samv2/samv2_ct_mri.py'
checkpoint_file = '/data/sdd/user/SAM/work_dirs/samv2_ct_mri/iter_20000_ct_mr_intra_only.pth'

im1_file = '/data/sdd/user/DEEDs/deedsBCV/mri.nii.gz'
im1_mask = '/data/sdd/user/DEEDs/deedsBCV/mri_mask.nii.gz'
im2_file = '/data/sdd/user/DEEDs/deedsBCV/ct.nii.gz'
im2_mask = '/data/sdd/user/DEEDs/deedsBCV/ct_mask.nii.gz'
deeds_path = '/data/sdd/user/DEEDs/deedsBCV/'
s_thres = 1.5


def compute_point(model, normed_im1, norm_info_1, normed_im2, norm_info_2):
    cuda_device = torch.device('cuda:' + str(0))
    emb1 = get_embedding(normed_im1, model)
    emb2 = get_embedding(normed_im2, model)

    shape1 = emb1[0].shape
    shape2 = emb2[0].shape
    img_1 = torch.tensor(emb1[2]).to(cuda_device)
    processed_1_img_downsample = F.interpolate(img_1, shape1[2:], mode='trilinear', align_corners=False)
    body_index = torch.where(processed_1_img_downsample > -25)
    z_min = body_index[2].min()
    z_max = body_index[2].max()
    y_min = body_index[3].min()
    y_max = body_index[3].max()
    x_min = body_index[4].min()
    x_max = body_index[4].max()

    x_strid = 12
    y_strid = 12
    z_strid = 8
    all_points_part1 = []  # 4 mini-batches because of  GPU memory limitation
    all_points_part2 = []
    all_points_part3 = []
    all_points_part4 = []
    for i in range(x_min, x_max, x_strid):
        for j in range(y_min, y_max, y_strid):
            for k in range(z_min+1, z_max, z_strid):
                all_points_part1.append([i, j, k])
    all_points_part1 = np.array(all_points_part1)
    for i in range(x_min + 5, x_max, x_strid):
        for j in range(y_min + 5, y_max, y_strid):
            for k in range(z_min + 4, z_max, z_strid):
                all_points_part2.append([i, j, k])
    all_points_part2 = np.array(all_points_part2)
    for i in range(x_min, x_max, x_strid):
        for j in range(y_min, y_max, y_strid):
            for k in range(z_min + 4, z_max, z_strid):
                all_points_part3.append([i, j, k])
    all_points_part3 = np.array(all_points_part3)
    for i in range(x_min + 5, x_max, x_strid):
        for j in range(y_min + 5, y_max, y_strid):
            for k in range(z_min+1, z_max, z_strid):
                all_points_part4.append([i, j, k])
    all_points_part4 = np.array(all_points_part4)

    its = 1
    for iter in range(its):
        if iter % 2 == 0:
            if iter == 0:
                q_loc_part1 = all_points_part1
            k_loc_part1, k_score_part1 = get_sim_embed_loc_multi_embedding_space(emb1, emb2, q_loc_part1, cuda_device)
            print(k_score_part1.mean())
        else:
            q_loc_part1, q_score_part1 = get_sim_embed_loc_multi_embedding_space(emb2, emb1, k_loc_part1, cuda_device)

    for iter in range(its):
        if iter % 2 == 0:
            if iter == 0:
                q_loc_part2 = all_points_part2
            k_loc_part2, k_score_part2 = get_sim_embed_loc_multi_embedding_space(emb1, emb2, q_loc_part2, cuda_device)
            print(k_score_part2.mean())
        else:
            q_loc_part2, q_score_part2 = get_sim_embed_loc_multi_embedding_space(emb2, emb1, k_loc_part2, cuda_device)

    for iter in range(its):
        if iter % 2 == 0:
            if iter == 0:
                q_loc_part3 = all_points_part3
            k_loc_part3, k_score_part3 = get_sim_embed_loc_multi_embedding_space(emb1, emb2, q_loc_part3, cuda_device)
            print(k_score_part3.mean())
        else:
            q_loc_part3, q_score_part3 = get_sim_embed_loc_multi_embedding_space(emb2, emb1, k_loc_part3, cuda_device)
    for iter in range(its):
        if iter % 2 == 0:
            if iter == 0:
                q_loc_part4 = all_points_part4
            k_loc_part4, k_score_part4 = get_sim_embed_loc_multi_embedding_space(emb1, emb2, q_loc_part4, cuda_device)
            print(k_score_part4.mean())
        else:
            q_loc_part4, q_score_part4 = get_sim_embed_loc_multi_embedding_space(emb2, emb1, k_loc_part4, cuda_device)

    q_loc = np.concatenate([q_loc_part1, q_loc_part2, q_loc_part3, q_loc_part4], axis=0)
    k_loc = np.concatenate([k_loc_part1, k_loc_part2, k_loc_part3, k_loc_part4], axis=0)
    k_score = np.concatenate([k_score_part1, k_score_part2, k_score_part3, k_score_part4], axis=0)
    print(k_score.mean(),k_score.max())
    q_loc = q_loc[k_score > s_thres, :]
    k_loc = k_loc[k_score > s_thres, :]
    q_loc = q_loc[:, ::-1].transpose(1, 0)
    k_loc = k_loc[:, ::-1].transpose(1, 0)
    q_loc = q_loc * 2 + 0.5  # zyx
    k_loc = k_loc * 2 + 0.5
    q_loc_origin = (q_loc.transpose(1, 0) / norm_info_1[::-1]).transpose(1, 0)
    k_loc_origin = (k_loc.transpose(1, 0) / norm_info_2[::-1]).transpose(1, 0)
    return q_loc, k_loc, q_loc_origin, k_loc_origin


def rigid_trans(q, k, im1, normed_im1, im2, normed_im2, DEEDS_fine_tune=True, with_mask=False,
                save_name=None, final_save_path=None, reverse=False):
    saveflag = False
    if not save_name == None:
        saveflag = True
    if reverse:
        querys = np.array(k)
        keys = np.array(q)
    else:
        querys = np.array(q)
        keys = np.array(k)
    querys = querys[::-1, :]  # xyz

    keys = keys[::-1, :]

    # find mean column wise
    centroid_querys = np.mean(querys, axis=1)
    centroid_keys = np.mean(keys, axis=1)

    # ensure centroids are 3x1
    centroid_querys = centroid_querys.reshape(-1, 1)
    centroid_keys = centroid_keys.reshape(-1, 1)

    # subtract mean
    querys_m = querys - centroid_querys
    keys_m = keys - centroid_keys
    querys_m = querys_m.astype(np.float32)
    keys_m = keys_m.astype(np.float32)
    H = querys_m @ np.transpose(keys_m)

    # sanity check
    if np.linalg.matrix_rank(H) < 3:
        raise ValueError("rank of H = {}, expecting 3".format(np.linalg.matrix_rank(H)))

    # find rotation
    U, S, Vt = np.linalg.svd(H)
    R = Vt.T @ U.T

    # special reflection case
    if np.linalg.det(R) < 0:
        print("det(R) < R, reflection detected!, correcting for it ...")
        Vt[2, :] *= -1
        R = Vt.T @ U.T

    rot_thres = 0.2
    r = ro.from_matrix(R)
    rot_est = r.as_rotvec()
    rot_est_flag = rot_est >= 0
    rot_est_binary = np.where(rot_est_flag, 1, -1)
    rot_est_larger_than_thres = np.abs(rot_est) > rot_thres
    rot_est_rect = np.where(np.abs(rot_est) <= rot_thres, rot_est, rot_thres)
    rot_est_rect_all = rot_est_rect * rot_est_binary
    rot_est_rect = np.where(rot_est_larger_than_thres, rot_est_rect_all, rot_est_rect)
    R_rect = ro.from_rotvec(rot_est_rect).as_matrix()
    t = -R_rect @ centroid_querys + centroid_keys
    m = np.zeros((4, 4), dtype=R.dtype)
    m[:3, :3] = R_rect
    m[:3, 3] = t.reshape(-1)
    m[3, :] = np.array([0, 0, 0, 1])
    print(m)
    m_xyz = np.copy(m)
    col0 = m[:, 0].copy()
    col1 = m[:, 1].copy()
    m[:, 0] = col1
    m[:, 1] = col0
    rol0 = m[0, :].copy()
    rol1 = m[1, :].copy()
    m[0, :] = rol1
    m[1, :] = rol0
    m1 = inv(m)
    # ------------------------resampled SAM rotation---------------------------------
    im1_resample_nd = normed_im1['img'][0][0, 0, :, :, :].numpy().transpose(1, 2, 0)
    im2_resample_nd = normed_im2['img'][0][0, 0, :, :, :].numpy().transpose(1, 2, 0)
    im1_resampled_shape = np.array(im1_resample_nd.shape)
    im2_resampled_shape = np.array(im2_resample_nd.shape)
    shapes = np.concatenate([im1_resampled_shape, im2_resampled_shape]).reshape(2, 3)
    covers = shapes.max(0)

    img_use = np.zeros(covers)
    if reverse:
        img_use[:] = im2_resample_nd.min()
        img_use[:im2_resampled_shape[0], :im2_resampled_shape[1], :im2_resampled_shape[2]] = im2_resample_nd
        img_use = img_use[:im1_resampled_shape[0], :im1_resampled_shape[1], :im1_resampled_shape[2]]
    else:
        img_use[:] = im1_resample_nd.min()
        img_use[:im1_resampled_shape[0], :im1_resampled_shape[1], :im1_resampled_shape[2]] = im1_resample_nd
        img_use = img_use[:im2_resampled_shape[0], :im2_resampled_shape[1], :im2_resampled_shape[2]]

    if with_mask == True:
        mask1_resample_nd = im1['processed_mask']
        mask1_resampled_shape = np.array(mask1_resample_nd.shape)
        mask_use = np.zeros(covers)
        mask_use[:mask1_resampled_shape[0], :mask1_resampled_shape[1], :mask1_resampled_shape[2]] = mask1_resample_nd
        mask2_resample_nd = im2['processed_mask']
        mask2_resampled_shape = np.array(mask2_resample_nd.shape)
        mask_use = mask_use[:mask2_resampled_shape[0], :mask2_resampled_shape[1], :mask2_resampled_shape[2]]

    rigid_nd_resampled = ndi.affine_transform(img_use, m1,
                                              cval=img_use.min())
    rigid_mask = rigid_nd_resampled > img_use.min() + 5
    rigid_mask = ndimage.binary_dilation(rigid_mask, structure=np.ones((15, 15, 1)))
    rigid_resampled = nib.Nifti1Image(rigid_nd_resampled, affine=(im2['affine_resampled']))

    fixed_resample_array = normed_im2['img'][0][0, 0, :, :, :].numpy().transpose(1, 2, 0) * rigid_mask
    fixed_resample_array[rigid_mask == 0] = fixed_resample_array.min()

    fixed_resample = nib.Nifti1Image(fixed_resample_array,
                                     affine=(im2['affine_resampled']))
    nib.save(fixed_resample, deeds_path + '/fixed_resample_sam.nii.gz')

    fixed_resample_no_crop = nib.Nifti1Image(normed_im2['img'][0][0, 0, :, :, :].numpy().transpose(1, 2, 0),
                                             affine=(im2['affine_resampled']))
    if saveflag:
        nib.save(rigid_resampled, os.path.join(final_save_path, save_name + '_0001.nii.gz'))
        nib.save(fixed_resample_no_crop, os.path.join(final_save_path, save_name + '_0000.nii.gz'))
    else:
        nib.save(rigid_resampled, deeds_path + 'affined_resample_sam.nii.gz')
        nib.save(fixed_resample_no_crop, deeds_path + 'fixed_resample_nocrop.nii.gz')

    if with_mask == True:
        rigid_mask_nd_resampled = ndi.affine_transform(mask_use, m1, order=0,
                                                       cval=0)
        rigid_mask_nd_resampled = rigid_mask_nd_resampled.astype('int32')
        rigid_mask_resampled = nib.Nifti1Image(rigid_mask_nd_resampled, affine=(im2['affine_resampled']))
        if saveflag:
            nib.save(rigid_mask_resampled, os.path.join(final_save_path, save_name + '_mask_0001.nii.gz'))
        else:
            nib.save(rigid_mask_resampled, deeds_path + 'affined_mask_resample_sam.nii.gz')

        fixed_mask_resample_array = mask2_resample_nd

        fixed_mask_resample = nib.Nifti1Image(fixed_mask_resample_array,
                                              affine=(im2['affine_resampled']))
        if saveflag:
            nib.save(fixed_mask_resample, os.path.join(final_save_path, save_name + '_mask_0000.nii.gz'))
        else:
            nib.save(fixed_mask_resample, deeds_path + 'fixed_mask_resample_sam.nii.gz')
    # -----------------------------resampled DEEDS affine fine tuning------------------------------
    if DEEDS_fine_tune:
        cmd = '%s/linearBCV -F %s/fixed_resample_nocrop.nii.gz -M %s/affined_resample_sam.nii.gz -l 3 -G 5x4x3 -L  3x2x2 -Q 2x2x2 -O %s/mri_ct_affine_resample' % (
            deeds_path, deeds_path, deeds_path, deeds_path)
        os.system(cmd)

        deeds_affine = []
        f = open(os.path.join(deeds_path, 'mri_ct_affine_resample_matrix.txt'), 'r')
        for i in range(4):
            deeds_affine.append(f.readline())
        f.close()

        all_affines_number = []
        for lines in deeds_affine:
            all_in_line = lines.split(' ', -1)
            all_in_line = [part for part in all_in_line if part.__len__() > 0]
            all_in_line.pop(-1)
            all_in_line = [float(part) for part in all_in_line]
            all_affines_number.append(all_in_line)
        all_affines_number = np.array(all_affines_number)
        all_affines_number = all_affines_number.transpose(1, 0)

        shifts = all_affines_number[:3, 3:].copy()
        ro_part = all_affines_number[:3, :3].copy()
        # yxz -> xyz
        ro_part = np.matmul(np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]]), ro_part)
        ro_part = np.matmul(ro_part, np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]]))
        all_affines_number[:3, :3] = ro_part  # xyz
        all_affines_number[0, 3] = shifts[1, 0]
        all_affines_number[1, 3] = shifts[0, 0]
        all_affines_number[2, 3] = shifts[2, 0]
        all_affines_number = np.concatenate([all_affines_number, np.array([0, 0, 0, 1]).reshape(-1, 4)], axis=0)
        all_affines_number = np.linalg.inv(all_affines_number)  # left multiply or right inverse multiply
        m_xyz_1 = np.matmul(m_xyz, all_affines_number)
        m_xyz_1_xyz = m_xyz_1.copy()

        col0 = m_xyz_1[:, 0].copy()
        col1 = m_xyz_1[:, 1].copy()
        m_xyz_1[:, 0] = col1
        m_xyz_1[:, 1] = col0
        rol0 = m_xyz_1[0, :].copy()
        rol1 = m_xyz_1[1, :].copy()
        m_xyz_1[0, :] = rol1
        m_xyz_1[1, :] = rol0
        m_xyz_2 = inv(m_xyz_1)

        rigid_nd_resampled = ndi.affine_transform(img_use, m_xyz_2,
                                                  cval=img_use.min())
        rigid_mask = rigid_nd_resampled > img_use.min() + 5
        rigid_mask = ndimage.binary_dilation(rigid_mask, structure=np.ones((15, 15, 1)))
        rigid_resampled = nib.Nifti1Image(rigid_nd_resampled, affine=(im2['affine_resampled']))
        nib.save(rigid_resampled, deeds_path + '/affined_resample_sam_deeds.nii.gz')
        fixed_resample_array = normed_im2['img'][0][0, 0, :, :, :].numpy().transpose(1, 2, 0) * rigid_mask
        fixed_resample_array[rigid_mask == 0] = fixed_resample_array.min()

        fixed_resample = nib.Nifti1Image(fixed_resample_array,
                                         affine=(im2['affine_resampled']))
        nib.save(fixed_resample, deeds_path + '/fixed_resample_sam_deeds.nii.gz')
    # ----------------------use to original image---------------------------------------------------------------------------
    im1_affine = im1['affine'].copy()
    col0 = im1_affine[:, 0].copy()
    col1 = im1_affine[:, 1].copy()
    im1_affine[:, 0] = col1
    im1_affine[:, 1] = col0

    im1_affine_resampled = im1['affine_resampled'].copy()
    col0 = im1_affine_resampled[:, 0].copy()
    col1 = im1_affine_resampled[:, 1].copy()
    im1_affine_resampled[:, 0] = col1
    im1_affine_resampled[:, 1] = col0

    im2_affine = im2['affine'].copy()
    col0 = im2_affine[:, 0].copy()
    col1 = im2_affine[:, 1].copy()
    im2_affine[:, 0] = col1
    im2_affine[:, 1] = col0

    im2_affine_resampled = im2['affine_resampled'].copy()
    col0 = im2_affine_resampled[:, 0].copy()
    col1 = im2_affine_resampled[:, 1].copy()
    im2_affine_resampled[:, 0] = col1
    im2_affine_resampled[:, 1] = col0

    L = np.matmul(im2_affine, np.linalg.inv(im2_affine_resampled))
    L[:3, 3] = 0
    B = np.matmul(im2_affine_resampled, np.linalg.inv(im2_affine))
    B[:3, 3] = 0
    if DEEDS_fine_tune:
        m_origin = np.matmul(np.linalg.inv(L), m_xyz_1_xyz)
        m_origin = np.matmul(m_origin, np.linalg.inv(B))

        col0 = m_origin[:, 0].copy()
        col1 = m_origin[:, 1].copy()
        m_origin[:, 0] = col1
        m_origin[:, 1] = col0
        rol0 = m_origin[0, :].copy()
        rol1 = m_origin[1, :].copy()
        m_origin[0, :] = rol1
        m_origin[1, :] = rol0
    else:
        m_origin = np.matmul(np.linalg.inv(L), m)  # m has transported the pivots so no change cols and raws
        m_origin = np.matmul(m_origin, np.linalg.inv(B))

    m_origin_sam = np.matmul(np.linalg.inv(L), m)
    m_origin_sam = np.matmul(m_origin_sam, np.linalg.inv(B))

    m_origin_1 = inv(m_origin)

    rigid_nd = ndi.affine_transform(im1['img'], m_origin_1, cval=im1['img'].min())
    fixed = nib.Nifti1Image(im2['img'], affine=im2['affine'])
    nib.save(fixed, deeds_path + 'fixed.nii.gz')
    rigid = nib.Nifti1Image(rigid_nd, affine=(im2['affine']))
    nib.save(rigid, deeds_path + 'affined.nii.gz')
    rigid_nd_sam = ndi.affine_transform(im1['img'], inv(m_origin_sam), cval=im1['img'].min())
    rigid_sam = nib.Nifti1Image(rigid_nd_sam, affine=(im2['affine']))
    nib.save(rigid_sam, deeds_path + 'affined_sam.nii.gz')
    # else:


if __name__ == '__main__':
    print(checkpoint_file)
    print(s_thres)
    model = init(config_file, checkpoint_file)
    root_path = '/data/sdd/user/rawdata/CT-MRI-Abd/nii/'
    print(root_path)
    all_cases = os.listdir(root_path)
    all_cases = [case.split('_', 1)[0] for case in all_cases if '_0000.nii.gz' in case]
    final_save_path = '/data/sde/user/abd_registered/abd_registered_0_'+str(s_thres)+'_double_check'
    print(final_save_path)
    if not os.path.exists(final_save_path):
        os.mkdir(final_save_path)
    s_name = 'ct.nii.gz'
    t_name = 'mri.nii.gz'
    count = 0
    for index, caseinfo in enumerate(all_cases):
        print(count, caseinfo)
        count += 1
        # if count > 1:
        #     break
        # if '1059783' not in caseinfo:
        #     continue
        source_image = os.path.join(root_path, caseinfo + '_0000.nii.gz')
        target_image = os.path.join(root_path, caseinfo + '_0001.nii.gz')
        sou_shape = tio.ScalarImage(source_image).data.shape
        tar_shape = tio.ScalarImage(target_image).data.shape

        padding_img(source_image=source_image, target_image=target_image, save_path=deeds_path, s_name=s_name,
                    t_name=t_name)
        im1, normed_im1, norm_info_1 = read_image(im1_file, is_MRI=True)  # im1 yxz normed_im1 zyx
        im2, normed_im2, norm_info_2 = read_image(im2_file, is_MRI=False, remove_bed=True)
        q, k, q_o, k_o = compute_point(model, normed_im1, norm_info_1, normed_im2, norm_info_2)
        DEEDS_fine_tune = False
        rigid_trans(q, k, im1, normed_im1, im2, normed_im2, DEEDS_fine_tune=DEEDS_fine_tune,
                    with_mask=False, save_name=None, final_save_path=None, reverse=False)

        # # ---------------------------------------------------------------
        if DEEDS_fine_tune:
            cmd = '%s/deedsBCV -F %s/fixed_resample_sam_deeds.nii.gz -M %s/affined_resample_sam_deeds.nii.gz -l 3 -G 5x4x3 -L 5x4x3 -Q 2x2x1 -O %s/MRI_SAM_resample_sam_deeds' % (
                deeds_path, deeds_path, deeds_path, deeds_path)
            os.system(cmd)
            cmd = '%s/applyBCVfloat_aniso -F %s/affined_resample_sam_deeds.nii.gz  -M %s/affined.nii.gz   -O %s/MRI_SAM_resample_sam_deeds -D %s/MRI_SAM_deform.nii.gz' % (
                deeds_path, deeds_path, deeds_path, deeds_path, deeds_path)
            os.system(cmd)
        else:
            cmd = '%s/deedsBCV -F %s/fixed_resample_sam.nii.gz -M %s/affined_resample_sam.nii.gz -l 4 -G 6x5x4x3 -L 6x5x4x3 -Q 3x2x2x1 -O %s/MRI_SAM_resample_sam' % (
                deeds_path, deeds_path, deeds_path, deeds_path)
            os.system(cmd)
            cmd = '%s/applyBCVfloat_aniso -F %s/affined_resample_sam.nii.gz  -M %s/affined_sam.nii.gz   -O %s/MRI_SAM_resample_sam -D %s/MRI_SAM_deform.nii.gz' % (
                deeds_path, deeds_path, deeds_path, deeds_path, deeds_path)
            # cmd = '%s/deedsBCV -F %s/fixed_resample_sam.nii.gz -M %s/affined_resample_sam.nii.gz -l 1 -G 3 -L 0 -Q 3 -O %s/MRI_SAM_resample_sam' % (
            #     deeds_path, deeds_path, deeds_path, deeds_path)
            # os.system(cmd)
            # cmd = '%s/applyBCVfloat_aniso -F %s/affined_resample_sam.nii.gz  -M %s/affined_sam.nii.gz   -O %s/MRI_SAM_resample_sam -D %s/MRI_SAM_deform.nii.gz' % (
            #     deeds_path, deeds_path, deeds_path, deeds_path, deeds_path)
            os.system(cmd)
        # # --------------------------------------------------------------
        ct = tio.ScalarImage(os.path.join(deeds_path, 'fixed.nii.gz'))
        ct_data = ct.data[0, :sou_shape[1], :sou_shape[2], :sou_shape[3]].numpy().transpose(2, 0, 1)
        ct_affine = ct.affine
        col1 = ct_affine[:, 1].copy()
        col0 = ct_affine[:, 0].copy()
        ct_affine[:, 1] = col0
        ct_affine[:, 0] = col1
        mri = tio.ScalarImage(os.path.join(deeds_path, 'MRI_SAM_deform.nii.gz'))
        mri_data = mri.data[0, :sou_shape[1], :sou_shape[2], :sou_shape[3]].numpy().transpose(2, 0, 1)
        mri_affine = mri.affine
        col1 = mri_affine[:, 1].copy()
        col0 = mri_affine[:, 0].copy()
        mri_affine[:, 1] = col0
        mri_affine[:, 0] = col1

        ct_nii = sitk.GetImageFromArray(ct_data)
        ct_nii.SetSpacing(ct.spacing)
        ct_nii.SetOrigin((-ct_affine[0, 3], -ct_affine[1, 3], ct_affine[2, 3]))
        ct_nii.SetDirection((1, 0, 0, 0, 1, 0, 0, 0, 1))
        sitk.WriteImage(ct_nii, os.path.join(final_save_path,
                                             caseinfo + '_0000.nii.gz'))

        mri_nii = sitk.GetImageFromArray(mri_data)
        mri_nii.SetSpacing(mri.spacing)
        mri_nii.SetOrigin((-mri_affine[0, 3], -mri_affine[1, 3], mri_affine[2, 3]))
        mri_nii.SetDirection((1, 0, 0, 0, 1, 0, 0, 0, 1))
        sitk.WriteImage(mri_nii, os.path.join(final_save_path,
                                              caseinfo + '_0001.nii.gz'))
        # # -------------------------------

tools/regis_sam.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
from mmcv import Config, DictAction
from mmcv.cnn import fuse_conv_bn
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
                         wrap_fp16_model)

from mmdet.apis import multi_gpu_test, single_gpu_test
from mmdet.datasets import (build_dataloader, build_dataset,
                            replace_ImageToTensor)
from mmdet.models import build_detector
from sam import *


def parse_args():
    parser = argparse.ArgumentParser(
        description='MMDet test (and eval) a model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='ids of gpus to use '
             '(only applicable to non-distributed testing)')
    parser.add_argument(
        '--gpu-collect',
        action='store_true',
        help='whether to use gpu to collect results.')
    parser.add_argument(
        '--tmpdir',
        help='tmp directory used for collecting results from multiple '
             'workers, available when gpu-collect is not specified')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    samples_per_gpu = 1
    if isinstance(cfg.data.test, dict):
        cfg.data.test.test_mode = True
        samples_per_gpu = cfg.data.test.pop('samples_per_gpu', 1)
        if samples_per_gpu > 1:
            # Replace 'ImageToTensor' to 'DefaultFormatBundle'
            cfg.data.test.pipeline = replace_ImageToTensor(
                cfg.data.test.pipeline)
    elif isinstance(cfg.data.test, list):
        for ds_cfg in cfg.data.test:
            ds_cfg.test_mode = True
        samples_per_gpu = max(
            [ds_cfg.pop('samples_per_gpu', 1) for ds_cfg in cfg.data.test])
        if samples_per_gpu > 1:
            for ds_cfg in cfg.data.test:
                ds_cfg.pipeline = replace_ImageToTensor(ds_cfg.pipeline)

    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids
    else:
        cfg.gpu_ids = range(1)

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
        if len(cfg.gpu_ids) > 1:
            warnings.warn(
                f'We treat {cfg.gpu_ids} as gpu-ids, and reset to '
                f'{cfg.gpu_ids[0:1]} as gpu-ids to avoid potential error in '
                'non-distribute testing time.')
            cfg.gpu_ids = cfg.gpu_ids[0:1]
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    rank, _ = get_dist_info()
    # build the dataloader
    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(
        dataset,
        samples_per_gpu=samples_per_gpu,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=distributed,
        shuffle=False)

    # build the model and load checkpoint
    cfg.model.train_cfg = None
    model = build_detector(cfg.model, test_cfg=cfg.get('test_cfg'))
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
    # for backward compatibility
    if 'CLASSES' in checkpoint.get('meta', {}):
        model.CLASSES = checkpoint['meta']['CLASSES']
    else:
        model.CLASSES = dataset.CLASSES

    if not distributed:
        model = MMDataParallel(model, device_ids=cfg.gpu_ids)
        outputs = single_gpu_test(model, data_loader)
    else:
        model = MMDistributedDataParallel(
            model.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False)
        outputs = multi_gpu_test(model, data_loader)


if __name__ == '__main__':
    main()

tools/test_sam.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import copy
import os
import os.path as osp
import time
import warnings
# os.environ["CUDA_VISIBLE_DEVICES"]="0"
import mmcv
import torch
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist
from mmcv.utils import get_git_hash

from mmdet import __version__
from mmdet.apis import init_random_seed, set_random_seed
from sam.apis.train import train_detector
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.utils import collect_env, get_root_logger
from sam.models import *


def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--auto-resume',
        action='store_true',
        help='resume from the latest checkpoint automatically')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='ids of gpus to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file (deprecate), '
        'change to --cfg-options instead.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    if args.options and args.cfg_options:
        raise ValueError(
            '--options and --cfg-options cannot be both '
            'specified, --options is deprecated in favor of --cfg-options')
    if args.options:
        warnings.warn('--options is deprecated in favor of --cfg-options')
        args.cfg_options = args.options

    return args


def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)
    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    cfg.auto_resume = args.auto_resume
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids
    else:
        cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
        if len(cfg.gpu_ids) > 1:
            warnings.warn(
                f'We treat {cfg.gpu_ids} as gpu-ids, and reset to '
                f'{cfg.gpu_ids[0:1]} as gpu-ids to avoid potential error in '
                'non-distribute training time.')
            cfg.gpu_ids = cfg.gpu_ids[0:1]
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        # re-set gpu_ids with distributed training mode
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info
    meta['config'] = cfg.pretty_text
    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    # set random seeds
    seed = init_random_seed(args.seed)
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = seed
    meta['exp_name'] = osp.basename(args.config)

    model = build_detector(
        cfg.model,
        train_cfg=cfg.get('train_cfg'),
        test_cfg=cfg.get('test_cfg'))
    model.init_weights()

    datasets = [build_dataset(cfg.data.train)]
    if len(cfg.workflow) == 2:
        val_dataset = copy.deepcopy(cfg.data.val)
        val_dataset.pipeline = cfg.data.train.pipeline
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmdet version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmdet_version=__version__ + get_git_hash()[:7],
            CLASSES=datasets[0].CLASSES)
    # add an attribute for visualization convenience
    model.CLASSES = datasets[0].CLASSES
    train_detector(
        model,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        meta=meta)


if __name__ == '__main__':
    main()

tools/train_sam.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2
# Copyright (c) Medical AI Lab, Alibaba DAMO Academy
import torchio as tio
import numpy as np
from sam.datasets.piplines import Resample, RescaleIntensity, GenerateMetaInfo, Collect3d, CropBackground, RemoveCTBed, \
    Crop, DynamicRescaleIntensity
from sam.datasets.collect import collate
import matplotlib.pyplot as plt
import os
from scipy import ndimage
import torch

import SimpleITK as sitk


def read_image(path, mask_path=None, norm_spacing=(2., 2., 2.), crop_loc=None, is_MRI=False, to_canonical=True,
               remove_bed=False):
    data = {}
    origin_info = {}
    image_fn = path.split('/', -1)[-1]
    data_path = path.replace(image_fn, '')

    img_tio = tio.ScalarImage(data_path + image_fn)
    if to_canonical:
        ToCanonical = tio.ToCanonical()
        img_tio = ToCanonical(img_tio)
    # print(img_tio.spacing)
    if img_tio.orientation == ('R', 'A', 'S'):
        img_data = img_tio.data
        img_tio.data = torch.flip(img_data, (1, 2))
        img_tio.affine = np.array(
            [-img_tio.affine[0, :], -img_tio.affine[1, :], img_tio.affine[2, :], img_tio.affine[3, :]])
    assert img_tio.orientation == ('L', 'P', 'S'), print('right now the image orientation need to be LPS+ ')
    img_data = img_tio.data
    img_tio.data = img_data.permute(0, 2, 1, 3)
    img_tio.affine = np.array(
        [img_tio.affine[1, :], img_tio.affine[0, :], img_tio.affine[2, :], img_tio.affine[3, :]])

    if remove_bed:
        remove_bed = RemoveCTBed(thres=-800)
        data = {}
        data['subject'] = tio.Subject(image=img_tio)
        data = remove_bed(data)
        body_mask = data['subject']['body_mask'].data
        img_tio.data[body_mask == 0] = img_tio.data.min()
        data = {}
        print('remove_bed')
    if not mask_path is None:
        # mask_info = mask_path.split('/', -1)[-1]
        mask_tio = tio.LabelMap(mask_path)
        if to_canonical:
            mask_tio = ToCanonical(mask_tio)
        if mask_tio.orientation == ('R', 'A', 'S'):
            mask_data = mask_tio.data
            mask_tio.data = torch.flip(mask_data, (1, 2))
            mask_tio.affine = np.array(
                [-mask_tio.affine[0, :], -mask_tio.affine[1, :], mask_tio.affine[2, :], mask_tio.affine[3, :]])
        assert mask_tio.orientation == ('L', 'P', 'S'), print('right now the image orientation need to be LPS+ ')
        mask_data = mask_tio.data
        # mask_data[mask_data>1] = 1
        mask_tio.data = mask_data.permute(0, 2, 1, 3)
        mask_tio.affine = np.array(
            [mask_tio.affine[1, :], mask_tio.affine[0, :], mask_tio.affine[2, :], mask_tio.affine[3, :]])
        mask_origin_data = mask_tio.data.numpy().copy()

    origin_data = img_tio.data.numpy().copy()
    img_tio_shape = img_tio.data.shape  # tio use lazy loading, so we read image shape from data.shape to force it load image here
    img_tio_spacing = img_tio.spacing
    img_tio_affine = img_tio.affine
    norm_ratio = np.array(img_tio_spacing) / np.array(norm_spacing)
    if not mask_path is None:
        subject = tio.Subject(
            image=img_tio,
            mask=mask_tio
        )
    else:
        subject = tio.Subject(
            image=img_tio
        )
    data['image_fn'] = image_fn
    data['subject'] = subject
    if not crop_loc is None:
        crop_filter = Crop(by_view=False)
        data['crop_matrix_origin'] = crop_loc
        data = crop_filter(data)

    resample_filter = Resample(norm_spacing=norm_spacing, crop_artifacts=False)
    data = resample_filter(data)
    affine_resampled = data['subject'].image.affine
    if is_MRI:
        rescale_filter = DynamicRescaleIntensity()
    else:
        rescale_filter = RescaleIntensity()
    data = rescale_filter(data)
    meta_collect_filter = GenerateMetaInfo(is_train=False)
    data = meta_collect_filter(data)
    data['is_mri'] = is_MRI
    collects = Collect3d(keys=['img'], meta_keys=['filename', 'is_mri'])
    input = collects(data)
    batch = collate([input])
    # batch['img_metas'] =  batch['img_metas'].data
    origin_info['img'] = origin_data[0, :, :, :]  # yxz
    origin_info['shape'] = img_tio_shape
    origin_info['spacing'] = img_tio_spacing
    origin_info['affine'] = img_tio_affine
    origin_info['affine_resampled'] = affine_resampled
    if not mask_path is None:
        origin_info['processed_mask'] = data['subject'].mask.data.numpy()[0, :, :, :]
        origin_info['origin_mask'] = mask_origin_data[0, :, :, :]
    return origin_info, batch, norm_ratio


def read_image_regis(path):
    data = {}
    origin_info = {}
    image_fn = path.split('/', -1)[-1]
    data_path = path.replace(image_fn, '')

    img_tio = tio.ScalarImage(data_path + image_fn)
    if img_tio.orientation == ('R', 'A', 'S'):
        img_data = img_tio.data
        img_tio.data = torch.flip(img_data, (1, 2))
        img_tio.affine = np.array(
            [-img_tio.affine[0, :], -img_tio.affine[1, :], img_tio.affine[2, :], img_tio.affine[3, :]])
    assert img_tio.orientation == ('L', 'P', 'S'), print('right now the image orientation need to be LPS+ ')
    img_data = img_tio.data
    img_tio.data = img_data.permute(0, 2, 1, 3)
    img_tio.affine = np.array(
        [img_tio.affine[1, :], img_tio.affine[0, :], img_tio.affine[2, :], img_tio.affine[3, :]])
    # origin_data = img_tio.data.numpy().copy()

    subject = tio.Subject(
        image=img_tio
    )
    data['image_fn'] = image_fn
    data['subject'] = subject
    remove_bed = RemoveCTBed(thres=-300)
    data = remove_bed(data)
    # mask = data['body_mask'].astype('int16')
    # img = data['subject'].image.data[0,:,:,:].numpy()
    # mask_sitk = sitk.GetImageFromArray(mask)
    # sitk.WriteImage(mask_sitk,'mask.nii.gz')
    # img_sitk = sitk.GetImageFromArray(img)
    # sitk.WriteImage(img_sitk,'img.nii.gz')

    crop = CropBackground(thres=1, by_mask=True)
    data = crop(data)
    img_tio_shape = data[
        'subject'].image.data.shape  # tio use lazy loading, so we read image shape from data.shape to force it load image here
    img_tio_spacing = data['subject'].image.spacing
    img_tio_affine = data['subject'].image.affine
    origin_data = data['subject'].image.data.numpy().copy()
    origin_mask = data['subject']['body_mask'].data[0, :, :, :].numpy().copy()
    norm_ratio = np.array(img_tio_spacing) / np.array([2., 2., 2.])
    resample = Resample()
    data = resample(data)
    rescale = RescaleIntensity()
    data = rescale(data)
    body_mask = data['subject']['body_mask'].data[0, :, :, :]
    meta_collect = GenerateMetaInfo()
    data = meta_collect(data)
    collects = Collect3d(keys=['img'])
    input = collects(data)
    batch = collate([input])
    origin_info['img'] = origin_data[0, :, :, :]  # yxz
    origin_info['shape'] = img_tio_shape
    origin_info['spacing'] = img_tio_spacing
    origin_info['affine'] = img_tio_affine
    origin_info['body_mask'] = body_mask  # this is the processed data body mask not for original image
    origin_info['origin_mask'] = origin_mask
    return origin_info, batch, norm_ratio


def visualize(im1, im2, norm_ratio_1, norm_ratio_2, pt1, pt2, score, savename=None,
              im1_isMRI=False, im2_isMRI=False):
    q_window_low, q_window_high = -300, 300
    k_window_low, k_window_high = -300, 300
    if im1_isMRI:
        q_window_low, q_window_high = im1.min(), im1.max()
    if im2_isMRI:
        k_window_low, k_window_high = im2.min(), im2.max()

    fig, ax = plt.subplots(3, 2, figsize=(20, 30))
    q_img = im1.transpose(2, 0, 1)
    q_img[q_img < q_window_low] = q_window_low
    q_img[q_img > q_window_high] = q_window_high
    q_img = (q_img - q_window_low) / (q_window_high - q_window_low)

    q_img = q_img.astype(np.float32)

    slice = q_img[pt1[2], :, :]
    ax[0, 0].set_title('query')
    ax[0, 0].imshow(slice, cmap='gray')
    ax[0, 0].plot((pt1[0]), (pt1[1]), 'o', markerfacecolor='none',
                  markeredgecolor="red",
                  markersize=12, markeredgewidth=2)
    slice = q_img[:, pt1[1], :]
    slice = slice[::-1, :]
    ax[1, 0].set_title('query')
    ax[1, 0].imshow(slice, cmap='gray', aspect=norm_ratio_1[2] / norm_ratio_1[0])
    ax[1, 0].plot((pt1[0]), (q_img.shape[0] - pt1[2] - 1), 'o',
                  markerfacecolor='none', markeredgecolor="red",
                  markersize=12, markeredgewidth=2)

    slice = q_img[:, :, pt1[0]]
    slice = slice[::-1, :]
    ax[2, 0].set_title('query')
    ax[2, 0].imshow(slice, cmap='gray', aspect=norm_ratio_1[2] / norm_ratio_1[1])
    ax[2, 0].plot((pt1[1]), (q_img.shape[0] - pt1[2] - 1), 'o',
                  markerfacecolor='none', markeredgecolor="red",
                  markersize=12, markeredgewidth=2)
    k_img = im2.transpose(2, 0, 1)

    k_img[k_img < k_window_low] = k_window_low
    k_img[k_img > k_window_high] = k_window_high
    k_img = (k_img - k_window_low) / (k_window_high - k_window_low)
    k_img = k_img.astype(np.float32)

    slice = k_img[pt2[2], :, :]
    ax[0, 1].set_title('key')
    ax[0, 1].imshow(slice, cmap='gray')
    ax[0, 1].plot((pt2[0]), (pt2[1]), 'o', markerfacecolor='none',
                  markeredgecolor="red",
                  markersize=12, markeredgewidth=2)

    slice = k_img[:, pt2[1], :]
    slice = slice[::-1, :]

    ax[1, 1].set_title('key')
    ax[1, 1].imshow(slice, cmap='gray', aspect=norm_ratio_2[2] / norm_ratio_2[0])
    ax[1, 1].plot((pt2[0]), (k_img.shape[0] - pt2[2] - 1), 'o',
                  markerfacecolor='none',
                  markeredgecolor="red",
                  markersize=12, markeredgewidth=2)
    slice = k_img[:, :, pt2[0]]
    slice = slice[::-1, :]
    ax[2, 1].set_title('key')
    ax[2, 1].imshow(slice, cmap='gray', aspect=norm_ratio_2[2] / norm_ratio_2[1])
    ax[2, 1].plot((pt2[1]), (k_img.shape[0] - pt2[2] - 1), 'o',
                  markerfacecolor='none',
                  markeredgecolor="red",
                  markersize=12, markeredgewidth=2)
    plt.suptitle(f'score:{score}')
    plt.tight_layout()
    if not savename is None:
        plt.savefig(
            f'/data/sdd/user/visual/visual-matching/{pt1[0]}_{pt1[1]}_{pt1[2]}_{pt2[0]}_{pt2[1]}_{pt2[2]}_{savename}.png')
        plt.close()
    else:
        plt.show()
        plt.close()


def get_contour_points(mask):
    # mask = ndimage.binary_erosion(mask, structure=np.array(
    #     [[[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 1, 0], [1, 1, 1], [0, 1, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]]]))
    # mask_erosion = ndimage.binary_erosion(mask, structure=np.array(
    #     [[[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 1, 0], [1, 1, 1], [0, 1, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]]]))

    mask = ndimage.binary_erosion(mask, structure=np.ones((3, 3, 3)))
    mask_erosion = ndimage.binary_erosion(mask, structure=np.array(
        [[[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 1, 0], [1, 1, 1], [0, 1, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]]]))

    mask_contour = mask.astype(int) - mask_erosion.astype(int)
    all_points = np.where(mask_contour == 1)
    all_points = np.array(all_points)
    z_range = [all_points[2, :].min(), all_points[2, :].max()]
    all_points = all_points[:, all_points[2, :] < z_range[1]]
    all_points = all_points[:, all_points[2, :] > z_range[0]]
    x = all_points[1, :].copy()
    y = all_points[0, :].copy()
    all_points[0, :] = x
    all_points[1, :] = y
    return all_points


def visualize_seg(img, points, im_fn):
    window_high = 200
    window_low = -200
    if not os.path.exists('/data/sdb/user/mmdet/liver_test/' + str(im_fn) + '/'):
        os.mkdir('/data/sdb/user/mmdet/liver_test/' + str(im_fn) + '/')

    q_img = img.transpose(2, 0, 1)
    q_img[q_img < window_low] = window_low
    q_img[q_img > window_high] = window_high
    q_img = (q_img - window_low) / (window_high - window_low)
    q_img = q_img.astype(np.float32)
    zs = np.unique(points[:, 2])
    for z in zs:
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.imshow(q_img[z, :, :], cmap='gray')
        slice_points = points[points[:, 2] == z, :]
        for i in range(slice_points.shape[0]):
            ax.plot((slice_points[i, 0]), slice_points[i, 1], 'o', markerfacecolor='none',
                    markeredgecolor="red",
                    markersize=12, markeredgewidth=2)
        plt.savefig('/data/sdb/user/mmdet/liver_test/' + str(im_fn) + '/' + str(z) + '.png')
        # plt.show()
        plt.close()


def make_local_grid_patch(pt, margin_x, margin_y, margin_z, imshape_x, imshape_y, imshape_z):
    pt = np.asarray(pt).astype('int')
    pt_patch = np.zeros((2 * margin_z + 1, 2 * margin_y + 1, 2 * margin_x + 1, 3), dtype='int')
    for j in range(-margin_z, margin_z + 1):
        for k in range(-margin_y, margin_y + 1):
            for w in range(-margin_x, margin_x + 1):
                pt_patch[j + margin_z, k + margin_y, w + margin_x, :] = pt + np.array(
                    [w, k, j])
    pt_patch = pt_patch.reshape(-1, 3)

    pt_patch[:, 0] = np.where(pt_patch[:, 0] < imshape_x, pt_patch[:, 0], imshape_x - 1)
    pt_patch[:, 1] = np.where(pt_patch[:, 1] < imshape_y, pt_patch[:, 1], imshape_y - 1)
    pt_patch[:, 2] = np.where(pt_patch[:, 2] < imshape_z, pt_patch[:, 2], imshape_z - 1)
    pt_patch[:, 0] = np.where(pt_patch[:, 0] > 0, pt_patch[:, 0], 0)
    pt_patch[:, 1] = np.where(pt_patch[:, 1] > 0, pt_patch[:, 1], 0)
    pt_patch[:, 2] = np.where(pt_patch[:, 2] > 0, pt_patch[:, 2], 0)
    pt_patch = np.unique(pt_patch, axis=0)
    return pt_patch


def IOX(box1, gts, dim='3D', metric='IOU'):
    """
    Compute the IOU, IOBB, IOGT, or IOBBGT between a box and a list of boxes. We assume the box is a predicted bounding-box and
    the box list is a list of ground-truths of the same image. IOU is intersection over union; IOBB is intersection over
    predicted bounding-box; IOGT is intersection over ground-truths; IOBBGT is max(IOBB, IOGT)
    :param box1: A numpy array of size d. d = 4 for 2D box (x1, y1, x2, y2), and 6 for 3D box (x1, y1, z1, x2, y2, z2).
    :param gts: A numpy array of size n x d. n is the number of gt boxes.
    :param dim: One of the 3 strings: 2D, 3D, P3D. P3D means detections are in 3D but gts are in 2D. If a gt is on one
     slice of the box, we compute the 2D overlap to decide if the box hits the gt. Note that gt is only annotated on one
     slice, but its format should still in 3D (x1, y1, z, x2, y2, z).
    :param metric: One of the 4 strings: IOU, IOBB, IOGT, IOBBGT.
    :return: A numpy array with size n.
    """
    if dim == '3D':
        ixmin = np.maximum(gts[:, 0], box1[0])
        iymin = np.maximum(gts[:, 1], box1[1])
        izmin = np.maximum(gts[:, 2], box1[2])
        ixmax = np.minimum(gts[:, 3], box1[3])
        iymax = np.minimum(gts[:, 4], box1[4])
        izmax = np.minimum(gts[:, 5], box1[5])
        iw = np.maximum(ixmax - ixmin + 1., 0.)
        ih = np.maximum(iymax - iymin + 1., 0.)
        id = np.maximum(izmax - izmin + 1., 0.)
        inters = iw * ih * id
        union = ((box1[3] - box1[0] + 1.) * (box1[4] - box1[1] + 1.) * (box1[5] - box1[2] + 1.) +
                 (gts[:, 3] - gts[:, 0] + 1.) * (gts[:, 4] - gts[:, 1] + 1.) * (gts[:, 5] - gts[:, 2] + 1.) - inters)
        box_size = (box1[3] - box1[0] + 1.) * (box1[4] - box1[1] + 1.) * (box1[5] - box1[2] + 1.)
        gt_sizes = (gts[:, 3] - gts[:, 0] + 1.) * (gts[:, 4] - gts[:, 1] + 1.) * (gts[:, 5] - gts[:, 2] + 1.)

    elif dim == 'P3D':
        assert np.all(gts[:, 2] == gts[:, 5])
        z_inters = (box1[2] <= gts[:, 2]) & (gts[:, 2] <= box1[5])
        ixmin = np.maximum(gts[:, 0], box1[0])
        iymin = np.maximum(gts[:, 1], box1[1])
        ixmax = np.minimum(gts[:, 3], box1[3])
        iymax = np.minimum(gts[:, 4], box1[4])
        iw = np.maximum(ixmax - ixmin + 1., 0.)
        ih = np.maximum(iymax - iymin + 1., 0.)
        inters = iw * ih * z_inters
        union = ((box1[3] - box1[0] + 1.) * (box1[4] - box1[1] + 1.) +
                 (gts[:, 3] - gts[:, 0] + 1.) * (gts[:, 4] - gts[:, 1] + 1.) - inters)
        box_size = (box1[3] - box1[0] + 1.) * (box1[4] - box1[1] + 1.)
        gt_sizes = (gts[:, 3] - gts[:, 0] + 1.) * (gts[:, 4] - gts[:, 1] + 1.)

    elif dim == '2D':
        ixmin = np.maximum(gts[:, 0], box1[0])
        iymin = np.maximum(gts[:, 1], box1[1])
        ixmax = np.minimum(gts[:, 2], box1[2])
        iymax = np.minimum(gts[:, 3], box1[3])
        iw = np.maximum(ixmax - ixmin + 1., 0.)
        ih = np.maximum(iymax - iymin + 1., 0.)
        inters = iw * ih
        union = ((box1[2] - box1[0] + 1.) * (box1[3] - box1[1] + 1.) +
                 (gts[:, 2] - gts[:, 0] + 1.) * (gts[:, 3] - gts[:, 1] + 1.) - inters)
        box_size = (box1[2] - box1[0] + 1.) * (box1[3] - box1[1] + 1.)
        gt_sizes = (gts[:, 2] - gts[:, 0] + 1.) * (gts[:, 3] - gts[:, 1] + 1.)

    if metric == 'IOU':
        return inters / union
    elif metric == 'IOBB':
        return inters / box_size
    elif metric == 'IOGT':
        return inters / gt_sizes
    elif metric == 'IOBBGT':
        return np.maximum(inters / box_size, inters / gt_sizes)


def log_sinkhorn_iterations(Z: torch.Tensor, log_mu: torch.Tensor, log_nu: torch.Tensor, iters: int) -> torch.Tensor:
    """ Perform Sinkhorn Normalization in Log-space for stability"""
    u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
    for _ in range(iters):
        u = log_mu - torch.logsumexp(Z + v.unsqueeze(1), dim=2)
        v = log_nu - torch.logsumexp(Z + u.unsqueeze(2), dim=1)
    return Z + u.unsqueeze(2) + v.unsqueeze(1)


def log_optimal_transport(scores: torch.Tensor, alpha: torch.Tensor, iters: int) -> torch.Tensor:
    """ Perform Differentiable Optimal Transport in Log-space for stability， alpha is the bin_score or Z in paper"""
    b, m, n = scores.shape
    one = scores.new_tensor(1)
    ms, ns = (m * one).to(scores), (n * one).to(scores)

    bins0 = alpha.expand(b, m, 1)
    bins1 = alpha.expand(b, 1, n)
    alpha = alpha.expand(b, 1, 1)

    couplings = torch.cat([torch.cat([scores, bins0], -1),
                           torch.cat([bins1, alpha], -1)], 1)

    norm = - (ms + ns).log()
    log_mu = torch.cat([norm.expand(m), ns.log()[None] + norm])
    log_nu = torch.cat([norm.expand(n), ms.log()[None] + norm])
    log_mu, log_nu = log_mu[None].expand(b, -1), log_nu[None].expand(b, -1)

    Z = log_sinkhorn_iterations(couplings, log_mu, log_nu, iters)
    Z = Z - norm  # multiply probabilities by M+N
    return Z


def arange_like(x, dim: int):
    return x.new_ones(x.shape[dim]).cumsum(0) - 1


if __name__ == '__main__':
    data_path = '/data/sdb/user/Deeplesion/Images_nifti/'
    image_fn = '004450_01_02_089-149.nii.gz'

    _, batch, _ = read_image(data_path + image_fn)

tools/utils.py
alibaba-damo-academy/self-supervised-anatomical-embedding-v2