add weights

Author: NotF404

alignshift/models/truncated_densenet3d_alignshift.py

@@ -1,6 +1,4 @@
 # Ke Yan, Imaging Biomarkers and Computer-Aided Diagnosis Laboratory,
-# National Institutes of Health Clinical Center, July 2019
-"""The truncated Densenet-121 with FPN and 3DCE"""
 from collections import namedtuple
 
 import torch
@@ -13,7 +11,7 @@
 import re
 from collections import OrderedDict
 from mmdet.models.registry import BACKBONES
-from convs.operators import AlignShiftConv
+from alignshift.operators import AlignShiftConv
 import torch.utils.checkpoint as cp
 from mmdet.models.utils import build_conv_layer, build_norm_layer
 # mybn = nn.BatchNorm3d
@@ -30,13 +28,13 @@ class _DenseLayer(nn.Sequential):
     def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, n_fold, memory_efficient=False, ref_thickness=None):
         super(_DenseLayer, self).__init__()
         self.add_module('norm1', build_norm_layer(norm_cfg, num_input_features, postfix=1)[1]),
-        self.add_module('relu1', nn.ReLU(inplace=True)),
+        self.add_module('relu1', nn.ReLU(inplace=False)),
         self.add_module('conv1', AlignShiftConv(num_input_features, bn_size *
                                            growth_rate, kernel_size=1, stride=1, alignshift=True,
                                            bias=False, ref_thickness=ref_thickness, n_fold=n_fold,)),
         self.add_module('norm2', build_norm_layer(norm_cfg, bn_size * growth_rate, postfix=1)[1]),
-        self.add_module('relu2', nn.ReLU(inplace=True)),
-        self.add_module('conv2', AlignShiftConv(bn_size * growth_rate, growth_rate,
+        self.add_module('relu2', nn.ReLU(inplace=False)),
+        self.add_module('conv2', AlignShiftConv(bn_size * growth_rate, growth_rate, alignshift=True,
                                            kernel_size=3, stride=1, padding=1,
                                            bias=False, ref_thickness=ref_thickness, n_fold=n_fold)),
         self.drop_rate = drop_rate
@@ -77,11 +75,6 @@ def forward(self, init_features, thickness):
             features.append(new_features)
         return torch.cat(features, 1)
 
-class _StageNormRelu(nn.Sequential):
-    def __init__(self, num_input_features):
-        super().__init__()
-        self.add_module('norm', build_norm_layer(norm_cfg, num_input_features, postfix=1)[1])
-        self.add_module('relu', nn.ReLU(inplace=True))
 
 class _Transition(nn.Sequential):
     def __init__(self, num_input_features, num_output_features):
@@ -101,17 +94,13 @@ def __init__(self, input_features, input_slice):
         # self.add_module('reduction_z_pooling', nn.AvgPool3d(kernel_size=[input_slice, 1, 1], stride=1))
 @BACKBONES.register_module
 class DenseNetCustomTrunc3dAlign(nn.Module):
-    """The truncated Densenet-121 with FPN and 3DCE"""
-    # truncated since transition layer 3 since we find it works better in DeepLesion
-    # We only keep the finest-level feature map after FPN
     def __init__(self, 
                 out_dim=256,
                 n_cts=3,
                 fpn_finest_layer=1,
                 ref_thickness=2.0,
                 n_fold=8, 
-                memory_efficient=True,
-                syncbn=True):
+                memory_efficient=True):
         super().__init__()
         self.depth = 121
         self.feature_upsample = True
@@ -176,8 +165,8 @@ def __init__(self,
                 nn.init.kaiming_uniform_(layer.weight, a=1)
                 nn.init.constant_(layer.bias, 0)
         self.init_weights()
-        if syncbn:
-            self = nn.SyncBatchNorm.convert_sync_batchnorm(self)
+        # if syncbn:
+        #     self = nn.SyncBatchNorm.convert_sync_batchnorm(self)
 
     def forward(self, x, thickness):
         x = self.conv0(x)
@@ -238,8 +227,8 @@ def init_weights(self, pretrained=True):
                 state_dict1[new_key] = state_dict[key]
 
         key = self.load_state_dict(state_dict1, strict=False)
-        print(key)
-
+        #print(key)
+        
     def freeze(self):
         for name, param in self.named_parameters():
             print('freezing', name)

alignshift/models/truncated_densenet3d_tsm.py

@@ -11,7 +11,7 @@
 import re
 from collections import OrderedDict
 from mmdet.models.registry import BACKBONES
-from convs.operators.tsmconv import TSMConv
+from alignshift.operators.tsmconv import TSMConv
 import torch.utils.checkpoint as cp
 from mmdet.models.utils import build_conv_layer, build_norm_layer
 # mybn = nn.BatchNorm3d
@@ -29,12 +29,12 @@ class _DenseLayer(nn.Sequential):
     def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, n_fold, memory_efficient=False):
         super(_DenseLayer, self).__init__()
         self.add_module('norm1', build_norm_layer(norm_cfg, num_input_features, postfix=1)[1]),
-        self.add_module('relu1', nn.ReLU(inplace=True)),
+        self.add_module('relu1', nn.ReLU(inplace=False)),
         self.add_module('conv1', TSMConv(num_input_features, bn_size *
                                            growth_rate, kernel_size=1, stride=1,
                                            bias=False, n_fold=n_fold)),
         self.add_module('norm2', build_norm_layer(norm_cfg, bn_size * growth_rate, postfix=1)[1]),
-        self.add_module('relu2', nn.ReLU(inplace=True)),
+        self.add_module('relu2', nn.ReLU(inplace=False)),
         self.add_module('conv2', TSMConv(bn_size * growth_rate, growth_rate,
                                            kernel_size=3, stride=1, padding=1,
                                            bias=False, n_fold=n_fold)),
@@ -83,7 +83,7 @@ def __init__(self, num_input_features, num_output_features):
         self.add_module('relu', nn.ReLU(inplace=True))
         self.add_module('conv', TSMConv(num_input_features, num_output_features,
                                           kernel_size=1, stride=1, bias=False, tsm=False))
-        self.add_module('pool', nn.AvgPool3d(kernel_size=[1, 2, 2], stride=[1, 2, 2]))
+        self.add_module('pool', nn.AvgPool3d(kernel_size=[1, 2, 2], stride=[1, 2, 2]))#, padding=[0,1,1]
 
 class _Reduction_z(nn.Sequential):
     def __init__(self, input_features, input_slice):
@@ -98,8 +98,7 @@ def __init__(self,
                 n_cts=3,
                 fpn_finest_layer=1,
                 memory_efficient=True,
-                n_fold=8,
-                syncbn=True):
+                n_fold=8,):
         super().__init__()
         self.depth = 121
         self.feature_upsample = True
@@ -126,16 +125,17 @@ def __init__(self,
         # Each denseblock
         num_features = num_init_features
         for i, num_layers in enumerate(block_config):
-            block = _DenseBlock(num_layers=num_layers, num_input_features=num_features,memory_efficient=memory_efficient,
-                                bn_size=bn_size, growth_rate=growth_rate, drop_rate=drop_rate, n_fold=n_fold)
+            block = _DenseBlock(num_layers=num_layers, num_input_features=num_features,
+                                bn_size=bn_size, growth_rate=growth_rate, drop_rate=drop_rate,
+                                n_fold=self.n_fold, memory_efficient=memory_efficient)
             self.add_module('denseblock%d' % (i + 1), block)
             num_features = num_features + num_layers * growth_rate
+            reductionz = _Reduction_z(num_features, self.n_cts)
+            self.add_module('reductionz%d' % (i + 1), reductionz)
             if i != len(block_config) - 1:
                 trans = _Transition(num_input_features=num_features, num_output_features=num_features // 2)
                 self.add_module('transition%d' % (i + 1), trans)
                 num_features = num_features // 2
-            reductionz = _Reduction_z(num_features, self.n_cts)
-            self.add_module('reductionz%d' % (i + 1), reductionz)
 
         # Final batch norm
         # self.add_module('norm5', nn.BatchNorm2d(num_features))
@@ -159,38 +159,40 @@ def __init__(self,
                 nn.init.kaiming_uniform_(layer.weight, a=1)
                 nn.init.constant_(layer.bias, 0)
         self.init_weights()
-        if syncbn:
-            self = nn.SyncBatchNorm.convert_sync_batchnorm(self)
+        # if syncbn:
+        #     self = nn.SyncBatchNorm.convert_sync_batchnorm(self)
 
     def forward(self, x):
         x = self.conv0(x)
         x = self.norm0(x)
-        relu0 = self.relu0(x)
-        pool0 = self.pool0(relu0)
+        x = self.relu0(x)
+        x = self.pool0(x)
 
-        db1 = self.denseblock1(pool0)
-        ts1 = self.transition1(db1)
+        x = self.denseblock1(x)
+        redc1 = self.reductionz1(x)
+        x = self.transition1(x)
 
-        db2 = self.denseblock2(ts1)
-        ts2 = self.transition2(db2)
 
-        db3 = self.denseblock3(ts2)
+        x = self.denseblock2(x)
+        redc2 = self.reductionz2(x)
+        x = self.transition2(x)
 
+
+        x = self.denseblock3(x)
+        redc3 = self.reductionz3(x)
         # truncated since here since we find it works better in DeepLesion
         # ts3 = self.transition3(db3)
         # db4 = self.denseblock4(ts3)
 
-        if self.feature_upsample:
-            ftmaps = [relu0[:,:,self.mid_ct,...], db1[:,:,self.mid_ct,...], db2[:,:,self.mid_ct,...], db3[:,:,self.mid_ct,...]]
-            x = self.lateral4(ftmaps[-1])
-            for p in range(3, self.fpn_finest_layer - 1, -1):
-                x = F.interpolate(x, scale_factor=2, mode="nearest")
-                y = ftmaps[p-1]
-                lateral = getattr(self, 'lateral%d' % p)(y)
-                x += lateral
-            return [x]
-        else:
-            return [db3]
+        # if self.feature_upsample:
+        ftmaps = [None, redc1.squeeze(2), redc2.squeeze(2), redc3.squeeze(2)]
+        x = self.lateral4(ftmaps[-1])
+        for p in range(3, self.fpn_finest_layer - 1, -1):
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+            y = ftmaps[p-1]
+            lateral = getattr(self, 'lateral%d' % p)(y)
+            x += lateral
+        return [x]
 
     def init_weights(self, pretrained=True):
         pattern = re.compile(

alignshift/operators/tsmconv.py

@@ -17,7 +17,7 @@ class TSMConv(_ConvNd):
     def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                  padding=0, dilation=1, groups=1,
                  bias=True, padding_mode='zeros', 
-                 n_fold=8, tsm=True, inplace=True,
+                 n_fold=8, tsm=True, inplace=False,
                  shift_padding_zero=True):
 
         kernel_size = _pair_same(kernel_size)

deeplesion.md

@@ -20,6 +20,7 @@ AlignShift:Bridging the Gap of Imaging Thickness in 3D Anisotropic Volumes ([arX
   the experiment code is base on [mmdetection](https://github.com/open-mmlab/mmdetection)
 ,this directory consists of compounents used in mmdetection.
 * ``mmdet`` 
+  
 ## Convert a 2D model into 3D with a single line of code
 
 ```python

deeplesion/dataset/DeepLesion_25d.py renamed to deeplesion/dataset/DeepLesionDataset_25d.py

@@ -50,8 +50,8 @@ def __call__(self, results):
         if 'gt_semantic_seg' in results:
             results['gt_semantic_seg'] = DC(
                 to_tensor(results['gt_semantic_seg'][None, ...]), stack=True)
-        if 'z_spacing' in results:
-            results['z_spacing'] = DC(to_tensor(results['z_spacing']), stack=True, pad_dims=None)
+        if 'thickness' in results:
+            results['thickness'] = DC(to_tensor(results['thickness']), stack=True, pad_dims=None)
         return results
 
     def __repr__(self):

deeplesion/dataset/transform.py

@@ -1,4 +1,4 @@
-from ..ENVIRON import data_root
+from deeplesion.ENVIRON import data_root
 anchor_scales = [4, 6, 8, 12, 24, 48]#, 64
 # fp16 = dict(loss_scale=96.)
 
@@ -9,7 +9,7 @@
     DATA_AUG_POSITION = False,
     NORM_SPACING = 0.,
     SLICE_INTV = 2.0,
-    NUM_SLICES = 7,
+    NUM_SLICES = 3,
     GROUNP_ZSAPACING = False,
 )
 input_channel = dataset_transform['NUM_SLICES']
@@ -21,6 +21,7 @@
     pretrained= False,
     backbone=dict(
         type='DenseNetCustomTrunc',
+        in_channels=input_channel,
         out_dim=512,
         fpn_finest_layer=2,),
     rpn_head=dict(
@@ -127,7 +128,7 @@
         max_per_img=50,
         mask_thr_binary=0.5))
 # dataset settings
-dataset_type = 'DeepLesionDataset_25d'
+dataset_type = 'DeepLesionDataset25d'
 
 img_norm_cfg = dict(
     mean=[0.] * input_channel, std=[1.] * input_channel, to_rgb=False)
@@ -166,7 +167,7 @@
 ]
 train_pipeline = [
     dict(type='DefaultFormatBundle'),
-    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks', 'z_spacing']),#, 'flage'#, 'img_info'#, 'z_spacing'
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks', 'thickness']),#, 'flage'#, 'img_info'#, 'z_spacing'
 ]
 
 test_pipeline = [
@@ -185,14 +186,14 @@
         dicm2png_cfg=dataset_transform),
         with_mask=True,
         with_label=True,
-    test=dict(
+    val=dict(
         type=dataset_type,
         ann_file=data_root + 'val_ann.pkl',
         image_path=data_root + 'Images_png/',
         pipeline=train_pipeline,
         pre_pipeline = pre_pipeline_test,
         dicm2png_cfg=dataset_transform),
-    val=dict(
+    test=dict(
         type=dataset_type,
         ann_file=data_root + 'test_ann.pkl',
         image_path=data_root + 'Images_png/',

deeplesion/mconfigs/densenet_25d.py

@@ -8,7 +8,7 @@
     DATA_AUG_POSITION = False,
     NORM_SPACING = 0.,
     SLICE_INTV = 2.0,
-    NUM_SLICES = 7,
+    NUM_SLICES = 3,
     GROUNP_ZSAPACING = False,
 )
 input_channel = dataset_transform['NUM_SLICES']

deeplesion/mconfigs/densenet_align.py

@@ -9,7 +9,7 @@
     DATA_AUG_POSITION = False,
     NORM_SPACING = 0.,
     SLICE_INTV = 2.0,
-    NUM_SLICES = 3,
+    NUM_SLICES = 7,
     GROUNP_ZSAPACING = False,
 )
 input_channel = dataset_transform['NUM_SLICES']
@@ -25,12 +25,11 @@
         out_dim=feature_channel,
         fpn_finest_layer=2,
         n_fold=8,
-        memory_efficient=True,
-        syncbn=False),
+        memory_efficient=True),
     rpn_head=dict(
         type='RPNHead',
         in_channels=feature_channel,
-        feat_channels=feature_channel,###原版这俩好像没有conv
+        feat_channels=feature_channel,
         anchor_scales=anchor_scales,
         anchor_ratios=[0.5, 1., 2.0],
         anchor_base_sizes=[4],
@@ -63,7 +62,7 @@
     mask_roi_extractor=dict(
         type='SingleRoIExtractor',
         roi_layer=dict(type='RoIAlign', out_size=14, sample_num=2),
-        finest_scale=196, #按照面积是这个的倍数， 到该层去取feature
+        finest_scale=196, 
         out_channels=feature_channel,
         featmap_strides=[4]),
     mask_head=dict(
@@ -170,7 +169,7 @@
 ]
 train_pipeline = [
     dict(type='DefaultFormatBundle_3d'),
-    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),#, 'flage'#, 'img_info'#, 'z_spacing', , 'thickness'
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks', 'thickness']),#, 'flage'#, 'img_info'#, 'z_spacing', , 'thickness'
 ]
 
 test_pipeline = [
@@ -190,14 +189,14 @@
         dicm2png_cfg=dataset_transform),
         with_mask=True,
         with_label=True,
-    test=dict(
+    val=dict(
         type=dataset_type,
         ann_file=data_root + 'val_ann.pkl',
         image_path=data_root + 'Images_png/',
         pipeline=train_pipeline,
         pre_pipeline = pre_pipeline_test,
         dicm2png_cfg=dataset_transform),
-    val=dict(
+    test=dict(
         type=dataset_type,
         ann_file=data_root + 'test_ann.pkl',
         image_path=data_root + 'Images_png/',

deeplesion/mconfigs/densenet_tsm.py

@@ -16,13 +16,14 @@
 class DenseNetCustomTrunc(nn.Module):
     def __init__(self, 
                 out_dim=256,
+                in_channels=3,
                 fpn_finest_layer=1):
         super().__init__()
         self.depth = 121
         self.feature_upsample = True
         self.fpn_finest_layer = fpn_finest_layer
         self.out_dim = out_dim
-        self.in_channel = 7
+        self.in_channel = in_channels
         assert self.depth in [121]
         if self.depth == 121:
             num_init_features = 64