udate coplenet_run.py

taigw · taigw · commit aaa5776fe252 · 2020-08-19T17:00:00.000+08:00
diff --git a/README.md b/README.md
@@ -15,7 +15,7 @@ Based on this project, we have also released a dataset named as `UESTC-COVID-19
 
 # Requirements
 * [Pytorch][torch_link] version >=1.0.1.
-* [PyMIC][pymic_link], a Pytorch-based toolkit for medical image computing. Install it by `pip install PYMIC`.
+* [PyMIC][pymic_link], a Pytorch-based toolkit for medical image computing. Version 0.2 is required. Install it by `pip install PYMIC==0.2`.
 * Some basic python packages such as Numpy, Pandas, SimpleITK.
 
 [data_link]:http://faculty.uestc.edu.cn/HiLab/en/article/379152/list/index.htm
@@ -31,7 +31,7 @@ COPLE-Net has also been implemented in [MONAI][monai_link], a PyTorch-based, ope
 
 # How to use
 1. Download the pretrained model and example CT images from [Google Drive][google_link] or [Baidu Netdisk][baidu_link] (extract code q0ci). Put them into the folder `coplenet_data`.
-2. Run `python net_run.py config/config.cfg`. The results will be saved in `coplenet_data/result`.
+2. Run `python coplenet_run.py test config/config.cfg`. The results will be saved in `coplenet_data/result`.
 3. To segment COVID-19 pneumonia lesions from your own images, make sure that the images have been cropped into the lung region, and the intensity has been normalized into [0, 1] using window width/level of 1500/-650. Open the configure file `config/config.cfg` and edit `root_dir`, `test_csv` and `output_dir` according to the path of your images. Then return to step 2 to obtain the segmentation results.
 
 [google_link]:https://drive.google.com/drive/folders/1K1jbrxWWhG_L7dh6yMyB4DtklBr-bhxH?usp=sharing
diff --git a/config/config.cfg b/config/config.cfg
@@ -19,9 +19,22 @@ Pad_output_size = [1, 32, 32]
 Pad_ceil_mode   = True
 Pad_inverse     = True
 
+[network]
+# this section gives parameters for network
+
+# type of network
+net_type = COPLENet
+
+# number of class, required for segmentation task
+class_num     = 2
+in_chns       = 1
+bilinear      = True
+feature_chns  = [32, 64, 128, 256, 512]
+dropout       = [0.0, 0.0, 0.3, 0.4, 0.5]
+
 [testing]
 # device name" cuda:n or cpu
-device_name = cuda:1
+device_name = cuda:0
 
 checkpoint_name  = coplenet_data/coplenet_pretrain.pt
 output_dir       = coplenet_data/result
diff --git a/coplenet_run.py b/coplenet_run.py
@@ -9,29 +9,29 @@
 from __future__ import print_function, division
 import sys
 import torch
-from coplenet import COPLENet
-from pymic.net_run.net_run import TrainInferAgent
 from pymic.util.parse_config import parse_config
+from pymic.net_run.net_run_agent import  NetRunAgent
+from pymic.net.net_dict import NetDict
+from coplenet import COPLENet
+
+my_net_dict = NetDict
+my_net_dict['COPLENet'] = COPLENet
 
 def main():
-    if(len(sys.argv) < 2):
-        print('Number of arguments should be 2. e.g.')
-        print('    python net_run.py config.cfg')
+    if(len(sys.argv) < 3):
+        print('Number of arguments should be 3. e.g.')
+        print('    python train_infer.py train config.cfg')
         exit()
     cfg_file = str(sys.argv[1])
     config   = parse_config(cfg_file)
 
-    # parameters of COPLENet
-    net_param = {"class_num"   : 2,
-                 "in_chns"     : 1,
-                 "bilinear"    : True,
-                 "feature_chns": [32, 64, 128, 256, 512],
-                 "dropout"     : [0.0, 0.0, 0.3, 0.4, 0.5]}
-    config['network'] = net_param
-    
-    net   = COPLENet(net_param)
-    agent = TrainInferAgent(config, 'test')
-    agent.set_network(net)
+    stage    = str(sys.argv[1])
+    cfg_file = str(sys.argv[2])
+    config   = parse_config(cfg_file)
+
+    # use custormized CNN and loss function
+    agent  = NetRunAgent(config, stage)
+    agent.set_network_dict(my_net_dict)
     agent.run()
 
 if __name__ == "__main__":
