dlt_test.py

import src.tract_feat as tract_feat
import src.nn_model as nn_model

import whitematteranalysis as wma
import numpy as np

import argparse
import h5py

import os
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import Session

from sklearn.metrics import classification_report
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.utils import class_weight
from sklearn.metrics import confusion_matrix

CPU = False
if CPU:
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = ""

    num_cores = 4

    config = ConfigProto(intra_op_parallelism_threads=num_cores,
                            inter_op_parallelism_threads=num_cores,
                            allow_soft_placement=True,
                            device_count={'CPU': 4})

    session = Session(config=config)

    # tf.keras.backend.set_session(session)


else:
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    # Assign GPU usage
    config = ConfigProto()
    # fix the gpu memory usage portion
    # config.gpu_options.per_process_gpu_memory_fraction = 0.75
    # dynamic assign gpu memory
    config.gpu_options.allow_growth = True
    session = Session(config=config)
    # session = tf.InteractiveSession(config=config)

# -----------------
# Parse arguments
# -----------------
parser = argparse.ArgumentParser(
    description="Testing using a CNN model.",
    epilog="Written by Fan Zhang, fzhang@bwh.harvard.edu")

parser.add_argument(
    'inputModel',
    help='Pretrained network model as an h5 file.')
parser.add_argument(
    'inputFeat',
    help='Input tract feature data as an h5 file.')
parser.add_argument(
    'outputDir',
    help='The output directory should be a new empty directory. It will be created if needed.')

parser.add_argument(
    '-modelLabelName', type=str,
    help='Label name in the model as an h5 file.')

parser.add_argument(
    '-inputLabel', type=str,
    help='Input ground truth label as an h5 file.')

parser.add_argument(
    '-outPrefix', type=str,
    help='A prefix string of all output files.')

parser.add_argument(
    '-tractVTKfile', type=str,
    help='Tractography data as a vtkPolyData file. If given, prediction will output tracts')

args = parser.parse_args()

script_name = '<test>'

if not os.path.exists(args.inputModel):
    print(script_name, "Error: Input network model ", args.inputModel, "does not exist.")
    exit()

if not os.path.exists(args.modelLabelName):
    print(script_name, "Error: Input model label name", args.modelLabelName, "does not exist.")
    exit()

if not os.path.exists(args.inputFeat):
    print(script_name, "Error: Input feature ", args.inputFeat, "does not exist.")
    exit()

if args.inputLabel is None:
    print(script_name, "No input label is provided. Will perform prediction only.")
elif not os.path.exists(args.inputLabel):
    print(script_name, "Error: Input label ", args.inputLabel, "does not exist.")
    exit()

if not os.path.exists(args.outputDir):
    print(script_name, "Output directory", args.outputDir, "does not exist, creating it.")
    os.makedirs(args.outputDir)

''' Load data '''

# Load model parameters
print(script_name, 'Load parameters when training the model.')
params = np.load(args.inputModel.replace('_model.h5', '_params.npy'), allow_pickle=True).item(0)

# Load label names in the model
print(script_name, 'Load tracts names along with the model.')
with h5py.File(args.modelLabelName, "r") as f:
    y_names_in_model = f['y_names'].value

# Load test data feature
with h5py.File(args.inputFeat, "r") as f:
    print(script_name, 'Load input feature.')
    x_test = f['feat'].value

# Generate ground truth labels for evaluation
if args.inputLabel is not None:

    print(script_name, 'Load input label.')
    with h5py.File(args.inputLabel, "r") as f:
        y_test = f['label_array'].value.astype(int)
        # y_value = f['label_values'].value
        y_names = f['label_names'].value

        # Used for generate ground truth label
        y_test_orig = y_test.copy()
        y_names_orig = y_names.copy()

    # Generate final ground truth label
    print(script_name, 'Generate FINAL ground truth label for evaluation.')

    print(script_name, ' # Feat Preprocessing - combine subdiviations of some tracts including CBLM, SupT, and Others.')

    # y_test contains index numbers for names of tracts, y_names are names of tracts
    y_test, y_names, _ = tract_feat.combine_tract_subdiviations_and_merge_outliers(y_test, y_names, verbose=False)

    if params['bilateral_feature']:
        y_test, y_names, _ = tract_feat.bilateralize_feature(y_test, y_names, verbose=False)

    y_test_ground_truth_final = tract_feat.update_y_test_based_on_model_y_names(y_test, y_names, y_names_in_model)

else:
    y_test_ground_truth_final = None

if params['bilateral_feature']:
    print(script_name, 'Make a bilateral copy for each fiber.')
    x_test, _ = tract_feat.bilateral_X_data(
        x_test)

# Perform predition of multiple tracts

print('')
print('===================================')
print('')
print(script_name, 'Start multi-tract prediction.')

print(script_name, 'x_test shape:', x_test.shape)
print(script_name, 'tracts to predict:', y_names_in_model)
# print script_name, 'tracts in the input data', y_names


output_multi_tract_predition_mask_path = os.path.join(args.outputDir,
                                                      args.outPrefix + '_multi_tract_specific_prediction_mask.h5')
output_multi_tract_predition_report_path = os.path.join(args.outputDir,
                                                        args.outPrefix + '_multi_tract_prediction_report.h5')
if not os.path.exists(output_multi_tract_predition_mask_path):

    # Load model
    model = load_model(args.inputModel)

    y_prediction, prediction_report, con_matrix = nn_model.predict(model, x_test, y_data=y_test_ground_truth_final,
                                                                   y_name=y_names_in_model, verbose=True)

    if args.inputLabel is not None:
        if prediction_report is not None:
            with h5py.File(output_multi_tract_predition_report_path, "w") as f:
                f.create_dataset('prediction_report', data=prediction_report)
                f.create_dataset('con_matrix', data=con_matrix)

    with h5py.File(output_multi_tract_predition_mask_path, "w") as f:
        f.create_dataset('y_prediction', data=y_prediction)

    del model

else:
    print(script_name, 'Loading prediction result.')
    with h5py.File(output_multi_tract_predition_mask_path, "r") as f:
        y_prediction = f[
            'y_prediction'].value

if args.tractVTKfile is not None:

    print('')
    print('===================================')
    print('')
    print(script_name, 'Output fiber tracts.')
    # Tractography Parcellation
    tract_prediction_mask = y_prediction

    print(script_name, 'Load vtk:', args.tractVTKfile)
    pd_whole_tract = wma.io.read_polydata(args.tractVTKfile)

    print(script_name, ' # labels in mask:', np.unique(tract_prediction_mask))
    print(script_name, ' # y_names:', y_names_in_model)

    number_of_tracts = np.max(tract_prediction_mask) + 1
    pd_t_list = wma.cluster.mask_all_clusters(pd_whole_tract, tract_prediction_mask, number_of_tracts,
                                              preserve_point_data=False, preserve_cell_data=False, verbose=False)

    output_tract_folder = os.path.join(args.outputDir, args.outPrefix + '_prediction_tracts_outlier_removed')
    if not os.path.exists(output_tract_folder):
        os.makedirs(output_tract_folder)

    for t_idx in range(len(pd_t_list)):
        pd_t = pd_t_list[t_idx]

        if y_names_in_model is not None:
            fname_t = os.path.join(output_tract_folder, y_names_in_model[t_idx].decode('UTF-8') + '.vtp')
        else:
            fname_t = os.path.join(output_tract_folder, 'tract_' + str(t_idx) + '.vtp')

        print(script_name, 'output', fname_t)
        wma.io.write_polydata(pd_t, fname_t)

    print(script_name, 'Done! Tracts are in:', output_tract_folder)