--- 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() + +