update net dict and loss dictg

Author: taigw

examples/JSRT2/jsrt_net_run.py

@@ -4,38 +4,15 @@
 import sys
 from pymic.util.parse_config import parse_config
 from pymic.net_run.net_run_agent import  NetRunAgent
-from pymic.net.net_factory import net_dict
-from pymic.loss.loss_factory import loss_dict
+from pymic.net.net_dict import NetDict
+from pymic.loss.loss_dict import LossDict
 from my_net2d import MyUNet2D 
 from my_loss  import MyFocalDiceLoss
 
-my_net_dict = {
-    "MyUNet2D": MyUNet2D
-}
-
-def get_network(params):
-    net_type = params["net_type"]
-    if(net_type in my_net_dict):
-        net = my_net_dict[net_type](params)
-    elif(net_type in net_dict):
-        net = net_dict[net_type](params)
-    else:
-        raise ValueError("Undefined network: {0:}".format(net_type))
-    return net 
-
-my_loss_dict = {
-    "MyFocalDiceLoss": MyFocalDiceLoss
-}
-
-def get_loss(params):
-    loss_type = params["loss_type"]
-    if(loss_type in my_loss_dict):
-        loss_obj = my_loss_dict[loss_type](params)
-    elif(loss_type in net_dict):
-        loss_obj = loss_dict[loss_type](params)
-    else:
-        raise ValueError("Undefined loss: {0:}".format(loss_type))
-    return loss_obj
+my_net_dict = NetDict
+my_net_dict["MyUNet2D"] = MyUNet2D
+my_loss_dict = LossDict
+my_loss_dict["MyFocalDiceLoss"] = MyFocalDiceLoss
 
 def main():
     if(len(sys.argv) < 3):
@@ -47,11 +24,9 @@ def main():
     config   = parse_config(cfg_file)
 
     # use custormized CNN and loss function
-    net      = get_network(config['network'])
-    loss_obj = get_loss(config['training'])
     agent  = NetRunAgent(config, stage)
-    agent.set_network(net)
-    agent.set_loss_calculater(loss_obj)
+    agent.set_network_dict(my_net_dict)
+    agent.set_loss_dict(my_loss_dict)
     agent.run()
 
 if __name__ == "__main__":

pymic/loss/loss_factory.py pymic/loss/loss_dict.pypymic/loss/loss_factory.py renamed to pymic/loss/loss_dict.py

@@ -5,18 +5,10 @@
 from pymic.loss.dice import DiceWithCrossEntropyLoss, NoiseRobustDiceLoss
 from pymic.loss.exp_log import ExpLogLoss
 
-loss_dict = {'CrossEntropyLoss': CrossEntropyLoss,
+LossDict = {'CrossEntropyLoss': CrossEntropyLoss,
     'GeneralizedCrossEntropyLoss': GeneralizedCrossEntropyLoss,
     'DiceLoss': DiceLoss,
     'MultiScaleDiceLoss': MultiScaleDiceLoss,
     'DiceWithCrossEntropyLoss': DiceWithCrossEntropyLoss,
     'NoiseRobustDiceLoss': NoiseRobustDiceLoss,
     'ExpLogLoss': ExpLogLoss}
-
-def get_loss(params):
-    loss_type = params['loss_type']
-    if(loss_type in loss_dict):
-        loss_obj = loss_dict[loss_type](params)
-    else:
-        raise ValueError("Undefined loss type {0:}".format(loss_type))
-    return loss_obj

pymic/net/net_factory.py pymic/net/net_dict.pypymic/net/net_factory.py renamed to pymic/net/net_dict.py

@@ -6,18 +6,10 @@
 from pymic.net.net3d.unet2d5 import UNet2D5
 from pymic.net.net3d.unet3d import UNet3D
 
-net_dict = {
+NetDict = {
 	'UNet2D': UNet2D,
 	'COPLENet': COPLENet,
 	'UNet2D_ScSE': UNet2D_ScSE,
 	'UNet2D5': UNet2D5,
 	'UNet3D': UNet3D
-	}
-	
-def get_network(params):
-	net_type = params['net_type']
-	if(net_type in net_dict):
-		net_obj = net_dict[net_type](params)
-	else:
-		raise ValueError("Undefined network type {0:}".format(net_type))
-	return net_obj
+	}

pymic/net_run/net_run_agent.py

@@ -17,11 +17,11 @@
 from tensorboardX import SummaryWriter
 from pymic.io.image_read_write import save_nd_array_as_image
 from pymic.io.nifty_dataset import NiftyDataset
-from pymic.io.transform3d import get_transform
-from pymic.net.net_factory import get_network
+from pymic.transform.trans_dict import TransformDict
+from pymic.net.net_dict import NetDict
 from pymic.net_run.infer_func import volume_infer
 from pymic.net_run.get_optimizer import get_optimiser
-from pymic.loss.loss_factory import get_loss
+from pymic.loss.loss_dict import LossDict
 from pymic.loss.util import get_soft_label
 from pymic.loss.util import reshape_prediction_and_ground_truth
 from pymic.loss.util import get_classwise_dice
@@ -35,23 +35,29 @@ def __init__(self, config, stage = 'train'):
         self.stage  = stage
         if(stage == 'inference'):
             self.stage = 'test'
-        self.net    = None
         self.train_set = None 
         self.valid_set = None 
         self.test_set  = None
+        self.net       = None
         self.loss_calculater = None 
+        self.transform_dict  = TransformDict
+        self.loss_dict       = LossDict
+        self.net_dict        = NetDict
         self.tensor_type = config['dataset']['tensor_type']
 
     def set_datasets(self, train_set, valid_set, test_set):
         self.train_set = train_set
         self.valid_set = valid_set
         self.test_set  = test_set
 
-    def set_network(self, net):
-        self.net = net 
+    def set_transform_dict(self, custom_transform_dict):
+        self.transform_dict = custom_transform_dict
 
-    def set_loss_calculater(self, loss_calculater):
-        self.loss_calculater = loss_calculater
+    def set_network_dict(self, custom_net_dict):
+        self.net_dict = custom_net_dict 
+
+    def set_loss_dict(self, custom_loss_dict):
+        self.loss_dict = custom_loss_dict
 
     def get_stage_dataset_from_config(self, stage):
         assert(stage in ['train', 'valid', 'test'])
@@ -66,11 +72,14 @@ def get_stage_dataset_from_config(self, stage):
             with_weight = False 
         else:
             raise ValueError("Incorrect value for stage: {0:}".format(stage))
+        self.transform_list  = []
+        for name in transform_names:
+            if(name not in self.transform_dict):
+                raise(ValueError("Undefined transform {0:}".format(name))) 
+            one_transform = self.transform_dict[name](self.config['dataset'])
+            self.transform_list.append(one_transform)
 
-        self.transform_list = [get_transform(name, self.config['dataset']) \
-                    for name in transform_names ]    
         csv_file = self.config['dataset'].get(stage + '_csv', None)
-    
         dataset  = NiftyDataset(root_dir=root_dir,
                                 csv_file  = csv_file,
                                 modal_num = modal_num,
@@ -99,8 +108,10 @@ def create_dataset(self):
                 batch_size=batch_size, shuffle=False, num_workers=batch_size)
 
     def create_network(self):
-        if(self.net is None):
-            self.net = get_network(self.config['network'])
+        net_name = self.config['network']['net_type']
+        if(net_name not in self.net_dict):
+            raise ValueError("Undefined network {0:}".format(net_name))
+        self.net = self.net_dict[net_name](self.config['network'])
         if(self.tensor_type == 'float'):
             self.net.float()
         else:
@@ -156,11 +167,14 @@ def train(self):
             self.checkpoint = None
         self.create_optimizer()
 
-        train_loss      = 0
+        loss_name = self.config['training']['loss_type']
+        if(loss_name not in self.loss_dict):
+            raise ValueError("Undefined loss function {0:}".format(loss_name))
+        self.loss_calculater = self.loss_dict[loss_name](self.config['training'])
+
+        trainIter  = iter(self.train_loader)
+        train_loss = 0
         train_dice_list = []
-        if(self.loss_calculater is None):
-            self.loss_calculater = get_loss(self.config['training'])
-        trainIter = iter(self.train_loader)
         print("{0:} training start".format(str(datetime.now())[:-7]))
         for it in range(iter_start, iter_max):
             try:

pymic/transform/__init__.py

@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+from pymic.transform.transform3d import * 
+
+TransformDict = {
+    'ChannelWiseGammaCorrection': ChannelWiseGammaCorrection,
+    'ChannelWiseNormalize': ChannelWiseNormalize,
+    'ChannelWiseThreshold': ChannelWiseThreshold,
+    'ChannelWiseThresholdWithNormalize': ChannelWiseThresholdWithNormalize,
+    'CropWithBoundingBox': CropWithBoundingBox,
+    'LabelConvert': LabelConvert,
+    'LabelConvertNonzero': LabelConvertNonzero,
+    'LabelToProbability': LabelToProbability,
+    'RandomCrop': RandomCrop,
+    'RandomFlip': RandomFlip,
+    'RandomRotate': RandomRotate,
+    'ReduceLabelDim': ReduceLabelDim,
+    'Rescale': Rescale,
+    'Pad': Pad,
+}