update train and test files

Author: taigw

pymic/io/image_read_write.py

@@ -112,6 +112,7 @@ def save_nd_array_as_image(data, image_name, reference_name = None):
 
     elif(image_name.endswith(".jpg") or image_name.endswith(".jpeg") or
          image_name.endswith(".tif") or image_name.endswith(".png")):
+         assert(data_dim == 2)
          save_array_as_rgb_image(data, image_name)
     else:
         raise ValueError("unsupported image format {0:}".format(

pymic/io/nifty_dataset.py

@@ -363,11 +363,13 @@ class RandomCrop(object):
 
     def __init__(self, output_size, fg_focus = False, fg_ratio = 0.0, mask_label = None,  inverse = True):
         assert isinstance(output_size, (list, tuple))
+        assert isinstance(mask_label, (list, tuple))
         self.output_size = output_size
         self.inverse  = inverse
         self.fg_focus = fg_focus
         self.fg_ratio = fg_ratio
         self.mask_label = mask_label
+        
 
     def __call__(self, sample):
         image = sample['image']

pymic/io/transform3d.py

@@ -99,12 +99,13 @@ def forward(self, x):
     params = {'in_chns':4,
               'feature_chns':[2, 8, 32, 48, 64],
               'class_num': 2,
+              'dropout': True,
               'acti_func': 'leakyrelu',
               'leakyrelu_alpha': 0.01}
     Net = UNet2D5(params)
     Net = Net.double()
 
-    x  = np.random.rand(4, 4, 32, 96, 96)
+    x  = np.random.rand(1, 4, 16, 128, 128)
     xt = torch.from_numpy(x)
     xt = torch.tensor(xt)
 

pymic/net2d/unet2d.py

@@ -12,6 +12,7 @@
 import torch.optim as optim
 import torch.nn.functional as F
 
+from scipy import special
 from datetime import datetime
 from tensorboardX import SummaryWriter
 from pymic.io.image_read_write import save_nd_array_as_image
@@ -221,7 +222,7 @@ def __infer(self):
         device = torch.device(self.config['testing']['device_name'])
         self.net.to(device)
         # laod network parameters and set the network as evaluation mode
-        self.checkpoint = torch.load(self.config['testing']['checkpoint_name'])
+        self.checkpoint = torch.load(self.config['testing']['checkpoint_name'], map_location = device)
         self.net.load_state_dict(self.checkpoint['model_state_dict'])
 
         if(self.config['testing']['evaluation_mode'] == True):
@@ -294,6 +295,8 @@ def test_time_dropout(m):
                     for c in range(0, class_num):
                         temp_prob = prob[c]
                         prob_save_name = "{0:}_prob_{1:}.{2:}".format(save_prefix, c, save_format)
+                        if(len(temp_prob.shape) == 2):
+                            temp_prob = np.asarray(temp_prob * 255, np.uint8)
                         save_nd_array_as_image(temp_prob, prob_save_name, root_dir + '/' + names[0])
 
         avg_time = (time.time() - start_time) / len(self.test_loder)

pymic/net3d/unet2d5.py

@@ -0,0 +1,22 @@
+import torch
+
+checkpoint_name1 = "/home/guotai/projects/PyMIC/examples/brats/model/casecade/wt/unet3d_4_8000.pt"
+checkpoint1 = torch.load(checkpoint_name1)
+state_dict1 = checkpoint1['model_state_dict']
+
+checkpoint_name2 = "/home/guotai/projects/PyMIC/examples/brats/model/casecade/wt/unet3d_4_10000.pt"
+checkpoint2 = torch.load(checkpoint_name2)
+state_dict2 = checkpoint2['model_state_dict']
+
+checkpoint_name3 = "/home/guotai/projects/PyMIC/examples/brats/model/casecade/wt/unet3d_4_12000.pt"
+checkpoint3 = torch.load(checkpoint_name3)
+state_dict3 = checkpoint3['model_state_dict']
+
+state_dict = {}
+for item in state_dict1:
+    print(item)
+    state_dict[item] = (state_dict1[item] + state_dict2[item] + state_dict3[item])/3
+
+save_dict = {'model_state_dict': state_dict}
+save_name = "/home/guotai/projects/PyMIC/examples/brats/model/casecade/wt/unet3d_4_avg.pt"
+torch.save(save_dict, save_name)