DeepSarcopenia/main.py at main · AIM-KannLab/DeepSarcopenia · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import argparse
import warnings
import os
import pandas as pd
import SimpleITK as sitk
import numpy as np
import glob as glob
from opts import parse_opts, get_args
from src.infer_segmentation import test_segmentation
from src.infer_slice_selection import test_slice_selection
from src.segmentation_preprocess import preprocess2


def slice_selection(img_dir,slice_model,slice_csv_path):
    """
    Test the Slice Selction Model
    Args:
        Input Scans -- nrrd files
        Model -- C3_Top_Selection_Model_Weight.hdf5
        Output -- C3_Top_Slice_Prediction.csv'
    """


    print('--- slice selection ---')
    test_slice_selection(
        image_dir=img_dir,
        model_weight_path=slice_model,
        csv_write_path=slice_csv_path)


def segmentation(pre_process_dir, seg_model, slice_csv_path, out_dir):

    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    print('--- C3 segmentation ---')
    test_segmentation(
        img_dir=pre_process_dir,
        model_weight_path=seg_model,
        slice_csv_path=slice_csv_path,
        output_dir=out_dir
        )


def preprocess(img_dir, pre_process_dir):

    img_dirs = [i for i in glob.glob(img_dir + '/*nrrd')]

    if not os.path.exists(pre_process_dir):
        os.makedirs(pre_process_dir)

    preprocess2(img_dirs, pre_process_dir)
#def main() :
#    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
#    warnings.filterwarnings('ignore')
#    opt = parse_opts()
#    dict1 = get_args(opt)
#    if (opt.test) :
        # Run Slice Slection
#        slice_selection(dict1["img_dir"], dict1["slice_model"],dict1["slice_csv_path"])
        # Run Processing steps on raw images for preprocessed files needed for segmenation
        #preprocess(dict1["img_dir"], dict1["pre_process_dir"])
        # Run the Segmentation which will generate output segmentations
        #segmentation(dict1["pre_process_dir"], dict1["seg_model"], dict1["slice_csv_path"], dict1["out_dir"])


if __name__ == '__main__':

    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    warnings.filterwarnings('ignore')
    opt = parse_opts()
    dict1 = get_args(opt)
    if (opt.test) :

        if (opt.STEP == 'ALL') :

            # Run Slice Slection, pre-processing and then segmentation
            slice_selection(dict1["img_dir"], dict1["slice_model"],dict1["slice_csv_path"])
            preprocess(dict1["img_dir"], dict1["pre_process_dir"])
            segmentation(dict1["pre_process_dir"], dict1["seg_model"], dict1["slice_csv_path"], dict1["out_dir"])
        elif (opt.STEP == 'SLICE') :
            slice_selection(dict1["img_dir"], dict1["slice_model"],dict1["slice_csv_path"])
        elif (opt.STEP == 'PREPROCESS') :
            preprocess(dict1["img_dir"], dict1["pre_process_dir"])
        elif (opt.STEP == 'SEGMENT') :
            segmentation(dict1["pre_process_dir"], dict1["seg_model"], dict1["slice_csv_path"], dict1["out_dir"])