Diff of /coplenet_run.py [000000] .. [5b2be7]

Switch to side-by-side view

--- a
+++ b/coplenet_run.py
@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Author: Guotai Wang
+# Date:   12 June, 2020
+# Implementation of of COPLENet for COVID-19 pneumonia lesion segmentation from CT images.
+# Reference: 
+#     G. Wang et al. A Noise-robust Framework for Automatic Segmentation of COVID-19 Pneumonia Lesions 
+#     from CT Images. IEEE Transactions on Medical Imaging, 2020. DOI:10.1109/TMI.2020.3000314.
+
+from __future__ import print_function, division
+import sys
+import torch
+from pymic.util.parse_config import parse_config
+from pymic.net_run.agent_seg import  SegmentationAgent
+from pymic.net.net_dict_seg import SegNetDict
+from coplenet import COPLENet
+
+net_dict = SegNetDict
+net_dict['COPLENet'] = COPLENet
+
+def main():
+    if(len(sys.argv) < 3):
+        print('Number of arguments should be 3. e.g.')
+        print('    python coplenet_run.py test config.cfg')
+        exit()
+    cfg_file = str(sys.argv[1])
+    config   = parse_config(cfg_file)
+
+    stage    = str(sys.argv[1])
+    cfg_file = str(sys.argv[2])
+    config   = parse_config(cfg_file)
+
+    # use custormized CNN and loss function
+    agent  = SegmentationAgent(config, stage)
+    net_name = config['network']['net_type']
+    if(net_name in net_dict):
+        net = net_dict[net_name](config['network'])
+        agent.set_network(net)
+        agent.run()
+    else:
+        raise ValueError("undefined network {0:}".format(net_name))
+
+if __name__ == "__main__":
+    main()
+    
+