First upload

59d4c6b verified over 2 years ago

10.2 kB

import os
import os.path
from os import environ
import sys
import json
import subprocess
import time
import nibabel as nib

# +++++++++++++ Conversion imports +++++++++++++++++++++++++
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.abspath(".."))
# +++++++++++++ Conversion imports +++++++++++++++++++++++++

from utils import *
from dicom_to_nii import convert_ct_dicom_to_nii, convert_transform_mr_to_nii, PatientList, save_images
from nii_to_dicom import convert_nii_to_dicom, integer_to_onehot
from predict_nnunet import predictNNUNet

def predict(tempPath, patient_id, regSeriesInstanceUID, runInterpreter):
    
    # Important: Check the input parameters #################
    if not patient_id or patient_id == "":
        sys.exit("No Patient dataset loaded: Load the patient dataset in Study Management.")
        
    if not regSeriesInstanceUID or regSeriesInstanceUID == "":
        sys.exit("No series instance UID for Modality 'REG' file. Check for REG file in your study")
        
    dir_base = os.path.join(tempPath, patient_id)
    createdir(dir_base)
    
    dir_ct_dicom = os.path.join(dir_base, 'ct_dicom')
    createdir(dir_ct_dicom)

    dir_mr_dicom = os.path.join(dir_base, 'mr_dicom')
    createdir(dir_mr_dicom)
    
    dir_reg_dicom = os.path.join(dir_base, 'reg_dicom')
    createdir(dir_reg_dicom)
        
    nnUNet_raw = os.path.join(os.getcwd(), 'nnUNet_raw')    
    nnUNet_preprocessed = os.path.join(os.getcwd(), 'nnUNet_preprocessed')
    RESULTS_FOLDER = os.path.join(os.getcwd(), 'nnUNet_trained_models')
    dataset = "Dataset103_EPTN_T1_CT_all_structures" 
    # IMPORTANT: data set modality: MR or CT ######################
    predictType='MR' 

    # IMPORTANT DOT Remove ########################################
    os.environ['nnUNet_raw'] = nnUNet_raw
    os.environ['nnUNet_preprocessed'] = nnUNet_preprocessed
    os.environ['nnUNet_results'] = RESULTS_FOLDER
    
    # Important ++++++++++++++++++++++++++++++++++++++++++++++++
    # Import the lib after setting environ parameters
    # import nnunet.inference.predict_simple as nnunetpredict
    
    print('** The python enviornment path: ', os.environ["PATH"])
    
    # For nnunet version 2
    import nnunetv2.inference.predict_from_raw_data as nnunetpredict
    # ###########################################################

    # predicted files
    predictedNiiFile = os.path.join(tempPath, patient_id, 'predict_nii')
    createdir(predictedNiiFile)
    
    predictedDicom = os.path.join(tempPath, patient_id, 'predicted_dicom')
    createdir(predictedDicom)
    
    predictedDicomFile = os.path.join(predictedDicom, 'predicted_rtstruct.dcm')
        
    print('** Use python interpreter: ', runInterpreter)
    print('** Patient name: ', patient_id)
    print('** REG series instance UID: ', regSeriesInstanceUID)
    
    # Convert CT image to NII #############
    startTime = time.time()
    
    if predictType == 'CT':
        
        dir_dicom_to_nii = os.path.join(nnUNet_raw, 'nnUNet_raw_data', 'Dataset098_HAN_nodes')
        createdir(dir_dicom_to_nii)
    
        downloadSeriesInstanceByModality(instanceID, dir_ct_dicom, "CT")
        print("Loading CT from Orthanc done: ", time.time()-startTime)
        
        # Convert CT image to NII #############
        refCT= convert_ct_dicom_to_nii(dir_dicom=dir_ct_dicom, dir_nii=dir_dicom_to_nii, outputname='1a_001_0000.nii.gz', newvoxelsize = None)
        print("Convert CT image to NII Done: ",  time.time()-startTime)
        
        # new version 2:  
        cmd = [modelPath, '-i', dir_dicom_to_nii, '-o', predictedNiiFile, '-d', dataset, '-tr', 'nnUNetTrainer_650epochs', '-c', '3d_fullres', '-f', '0']
        
        out = subprocess.check_output(cmd)
        # Important ########################
        sys.argv = cmd
        
        # #### nnunet version 2 #############
        nnunetpredict.predict_entry_point() 
        print("Prediction CT done",  time.time()-startTime)
        
        niiFile = os.path.join(predictedNiiFile, '1a_001.nii.gz')
        
        # POSTPROCESSING TO CONVERT FROM INTEGERS TO 2**i, ADD CONTOURS EXISTS, AND SMOOTH        
        integer_to_onehot(niiFile)
        print("POST processing convert from integers done: ", time.time()-startTime)
        
        startTime = time.time()
        convert_nii_to_dicom(dicomctdir=dir_ct_dicom, predictedNiiFile=niiFile, predictedDicomFile=predictedDicomFile,
                             predicted_structures=predicted_structures, rtstruct_colors=rtstruct_colors, refCT=refCT)
        
        print("Convert CT predicted NII to DICOM done: ", time.time()-startTime)
                
    elif predictType == 'MR':
        
        dir_dicom_to_nii = os.path.join(nnUNet_raw, 'nnUNet_raw_data',dataset)
        createdir(dir_dicom_to_nii)
        
        # Download the REG dicom ##############
        downloadSeriesInstanceByModality(regSeriesInstanceUID, dir_reg_dicom, "REG")
        print("Loading REG from Orthanc done: ", time.time()-startTime)
        
        # Download the MR dicom ###############
        # Read the mr study instance UID from the download REG dicom   
        mrSeriesInstanceUID = getSeriesInstanceUIDFromRegDicom(dir_reg_dicom, regSeriesInstanceUID)
        
        downloadSeriesInstanceByModality(mrSeriesInstanceUID, dir_mr_dicom, "MR")
        print("Loading MR from Orthanc done: ", time.time()-startTime)
        
        # Execute REG tranformation ###########
        ctSeriesInstanceUIDFromRegDicom = getCTSeriesInstanceUIDFromRegDicom(dir_reg_dicom, regSeriesInstanceUID)
        print("CT Series Instance UID referenced by Reg dicom: ", ctSeriesInstanceUIDFromRegDicom)
        
        downloadSeriesInstanceByModality(ctSeriesInstanceUIDFromRegDicom, dir_ct_dicom, "CT")
        
        Patients = PatientList() 
        Patients.list_dicom_files(dir_ct_dicom, 1)
        patient = Patients.list[0]
        patient_name = patient.PatientInfo.PatientName
        patient.import_patient_data(newvoxelsize=None)
        CT = patient.CTimages[0]

        startTime = time.time()
        mr_reg = regMatrixTransformation(dir_mr_dicom, reg_file_path=dir_reg_dicom, regSeriesInstanceUID=regSeriesInstanceUID, CT=CT)
        print("Transforming MR data done (OpenTPS.Core)")
        
        # Convert transform MR image to NII ##################
        refMR = convert_transform_mr_to_nii(dir_mr_dicom=dir_mr_dicom, tranform_mr = mr_reg, dir_nii=dir_dicom_to_nii, outputname='1a_001_0000.nii.gz', CT=CT)
        refCT= convert_ct_dicom_to_nii(dir_dicom=dir_ct_dicom, dir_nii=dir_dicom_to_nii, outputname='1a_001_0001.nii.gz', newvoxelsize = None)
        print("Convert CT image to NII Done: ",  time.time()-startTime)
        print("Convert transform MR image to NII Done: ",  time.time()-startTime)
        
        
        print("## start MR running prediction ###############")
        startTime = time.time()
        # modelPath = '..\\..\\python_environments\\prediction-3.10.9\\Scripts\\nnUNetv2_predict.exe' 
        # cmd = [modelPath, '-i', dir_dicom_to_nii, '-o', predictedNiiFile, '-d', '99', '-c', '3d_fullres' , '--disable_tta', '-tr', 'nnUNetTrainer_650epochs', '-f', '1, 4']
        
        predictNNUNet(os.path.join(RESULTS_FOLDER,dataset, 'nnUNetTrainer_650epochs__nnUNetPlans__3d_fullres'),
                      dir_dicom_to_nii,
                      predictedNiiFile,
                      [1]) 
        
        print("Prediction MR done",  time.time()-startTime)
        
        startTime = time.time()

        predicted_structures = ["background", "BRAIN", "AMYGDALAE", "BRAINSTEM", "CAUDATENUCLEI", "CEREBELLUM", "CHIASM", "COCHLEAS", "CORNEAS", "CORPUSCALLOSUM", "FORNICES", "GLANDPINEAL", "HIPPOCAMPI", "HYPOTHALAMI", "LACRIMALGLANDS", "LENSES", "OPTICNERVES", "ORBITOFRONTALS", "PITUITARY", "RETINAS", "THALAMI", "VSCCs"]
        rtstruct_colors = [[255,0,0]]*len(predicted_structures)
        
        niiFile = os.path.join(predictedNiiFile, '1a_001.nii.gz')
        
        # POSTPROCESSING TO CONVERT FROM INTEGERS TO 2**i, ADD CONTOURS EXISTS, AND SMOOTH        
        integer_to_onehot(niiFile)
        print("POST processing convert from integers done: ", time.time()-startTime)
        
        # Convert CT image to NII #############
        
        
        
        convert_nii_to_dicom(dicomctdir=dir_ct_dicom, predictedNiiFile=niiFile, predictedDicomFile=predictedDicomFile, 
                            predicted_structures=predicted_structures, rtstruct_colors=rtstruct_colors, refCT=refCT)
    else:
        print("Not supported yet")
        
    
    startTime = time.time()
    uploadDicomToOrthanc(predictedDicomFile)
    print("Upload predicted result to Orthanc done: ",  time.time()-startTime)

# tempPath = 'C:\Temp\parrot_prediction'
# regSeriesInstanceUID = '1.2.246.352.205.5029381855449574337.1508502639685232062'
# runInterpreter = 'py3109'
# patientName = 'P0461C0006I7638639'

'''
Prediction parameters provided by the server. Select the parameters to be used for prediction:
    [1] tempPath: The path where the predict.py is stored,
    [2] patientname: python version,
    [3] ctSeriesInstanceUID: Series instance UID for data set with modality = CT. To predict 'MR' modality data, retrieve the CT UID by the code (see Precision Code)
    [4] rtStructSeriesInstanceUID: Series instance UID for modality = RTSTURCT
    [5] regSeriesInstanceUID: Series instance UID for modality = REG,
    [6] runInterpreter: The python version for the python environment
    [7] oarList: only for dose predciton. For contour predicion oarList = []
    [8] tvList:  only for dose prediction. For contour prediction tvList = []
''' 
if __name__ == '__main__':
    predict(tempPath=sys.argv[1], patient_id=sys.argv[2], regSeriesInstanceUID=sys.argv[5], runInterpreter=sys.argv[6])
    # predict(tempPath=tempPath, patient_id=patientName, regSeriesInstanceUID=regSeriesInstanceUID, runInterpreter=runInterpreter)

	import os
	import os.path
	from os import environ
	import sys
	import json
	import subprocess
	import time
	import nibabel as nib

	# +++++++++++++ Conversion imports +++++++++++++++++++++++++
	sys.path.append(os.path.dirname(os.path.abspath(__file__)))
	sys.path.append(os.path.abspath(".."))
	# +++++++++++++ Conversion imports +++++++++++++++++++++++++

	from utils import *
	from dicom_to_nii import convert_ct_dicom_to_nii, convert_transform_mr_to_nii, PatientList, save_images
	from nii_to_dicom import convert_nii_to_dicom, integer_to_onehot
	from predict_nnunet import predictNNUNet

	def predict(tempPath, patient_id, regSeriesInstanceUID, runInterpreter):

	# Important: Check the input parameters #################
	if not patient_id or patient_id == "":
	sys.exit("No Patient dataset loaded: Load the patient dataset in Study Management.")

	if not regSeriesInstanceUID or regSeriesInstanceUID == "":
	sys.exit("No series instance UID for Modality 'REG' file. Check for REG file in your study")

	dir_base = os.path.join(tempPath, patient_id)
	createdir(dir_base)

	dir_ct_dicom = os.path.join(dir_base, 'ct_dicom')
	createdir(dir_ct_dicom)

	dir_mr_dicom = os.path.join(dir_base, 'mr_dicom')
	createdir(dir_mr_dicom)

	dir_reg_dicom = os.path.join(dir_base, 'reg_dicom')
	createdir(dir_reg_dicom)

	nnUNet_raw = os.path.join(os.getcwd(), 'nnUNet_raw')
	nnUNet_preprocessed = os.path.join(os.getcwd(), 'nnUNet_preprocessed')
	RESULTS_FOLDER = os.path.join(os.getcwd(), 'nnUNet_trained_models')
	dataset = "Dataset103_EPTN_T1_CT_all_structures"
	# IMPORTANT: data set modality: MR or CT ######################
	predictType='MR'

	# IMPORTANT DOT Remove ########################################
	os.environ['nnUNet_raw'] = nnUNet_raw
	os.environ['nnUNet_preprocessed'] = nnUNet_preprocessed
	os.environ['nnUNet_results'] = RESULTS_FOLDER

	# Important ++++++++++++++++++++++++++++++++++++++++++++++++
	# Import the lib after setting environ parameters
	# import nnunet.inference.predict_simple as nnunetpredict

	print('** The python enviornment path: ', os.environ["PATH"])

	# For nnunet version 2
	import nnunetv2.inference.predict_from_raw_data as nnunetpredict
	# ###########################################################

	# predicted files
	predictedNiiFile = os.path.join(tempPath, patient_id, 'predict_nii')
	createdir(predictedNiiFile)

	predictedDicom = os.path.join(tempPath, patient_id, 'predicted_dicom')
	createdir(predictedDicom)

	predictedDicomFile = os.path.join(predictedDicom, 'predicted_rtstruct.dcm')

	print('** Use python interpreter: ', runInterpreter)
	print('** Patient name: ', patient_id)
	print('** REG series instance UID: ', regSeriesInstanceUID)

	# Convert CT image to NII #############
	startTime = time.time()

	if predictType == 'CT':

	dir_dicom_to_nii = os.path.join(nnUNet_raw, 'nnUNet_raw_data', 'Dataset098_HAN_nodes')
	createdir(dir_dicom_to_nii)

	downloadSeriesInstanceByModality(instanceID, dir_ct_dicom, "CT")
	print("Loading CT from Orthanc done: ", time.time()-startTime)

	# Convert CT image to NII #############
	refCT= convert_ct_dicom_to_nii(dir_dicom=dir_ct_dicom, dir_nii=dir_dicom_to_nii, outputname='1a_001_0000.nii.gz', newvoxelsize = None)
	print("Convert CT image to NII Done: ", time.time()-startTime)

	# new version 2:
	cmd = [modelPath, '-i', dir_dicom_to_nii, '-o', predictedNiiFile, '-d', dataset, '-tr', 'nnUNetTrainer_650epochs', '-c', '3d_fullres', '-f', '0']

	out = subprocess.check_output(cmd)
	# Important ########################
	sys.argv = cmd

	# #### nnunet version 2 #############
	nnunetpredict.predict_entry_point()
	print("Prediction CT done", time.time()-startTime)

	niiFile = os.path.join(predictedNiiFile, '1a_001.nii.gz')

	# POSTPROCESSING TO CONVERT FROM INTEGERS TO 2**i, ADD CONTOURS EXISTS, AND SMOOTH
	integer_to_onehot(niiFile)
	print("POST processing convert from integers done: ", time.time()-startTime)

	startTime = time.time()
	convert_nii_to_dicom(dicomctdir=dir_ct_dicom, predictedNiiFile=niiFile, predictedDicomFile=predictedDicomFile,
	predicted_structures=predicted_structures, rtstruct_colors=rtstruct_colors, refCT=refCT)

	print("Convert CT predicted NII to DICOM done: ", time.time()-startTime)

	elif predictType == 'MR':

	dir_dicom_to_nii = os.path.join(nnUNet_raw, 'nnUNet_raw_data',dataset)
	createdir(dir_dicom_to_nii)

	# Download the REG dicom ##############
	downloadSeriesInstanceByModality(regSeriesInstanceUID, dir_reg_dicom, "REG")
	print("Loading REG from Orthanc done: ", time.time()-startTime)

	# Download the MR dicom ###############
	# Read the mr study instance UID from the download REG dicom
	mrSeriesInstanceUID = getSeriesInstanceUIDFromRegDicom(dir_reg_dicom, regSeriesInstanceUID)

	downloadSeriesInstanceByModality(mrSeriesInstanceUID, dir_mr_dicom, "MR")
	print("Loading MR from Orthanc done: ", time.time()-startTime)

	# Execute REG tranformation ###########
	ctSeriesInstanceUIDFromRegDicom = getCTSeriesInstanceUIDFromRegDicom(dir_reg_dicom, regSeriesInstanceUID)
	print("CT Series Instance UID referenced by Reg dicom: ", ctSeriesInstanceUIDFromRegDicom)

	downloadSeriesInstanceByModality(ctSeriesInstanceUIDFromRegDicom, dir_ct_dicom, "CT")

	Patients = PatientList()
	Patients.list_dicom_files(dir_ct_dicom, 1)
	patient = Patients.list[0]
	patient_name = patient.PatientInfo.PatientName
	patient.import_patient_data(newvoxelsize=None)
	CT = patient.CTimages[0]

	startTime = time.time()
	mr_reg = regMatrixTransformation(dir_mr_dicom, reg_file_path=dir_reg_dicom, regSeriesInstanceUID=regSeriesInstanceUID, CT=CT)
	print("Transforming MR data done (OpenTPS.Core)")

	# Convert transform MR image to NII ##################
	refMR = convert_transform_mr_to_nii(dir_mr_dicom=dir_mr_dicom, tranform_mr = mr_reg, dir_nii=dir_dicom_to_nii, outputname='1a_001_0000.nii.gz', CT=CT)
	refCT= convert_ct_dicom_to_nii(dir_dicom=dir_ct_dicom, dir_nii=dir_dicom_to_nii, outputname='1a_001_0001.nii.gz', newvoxelsize = None)
	print("Convert CT image to NII Done: ", time.time()-startTime)
	print("Convert transform MR image to NII Done: ", time.time()-startTime)


	print("## start MR running prediction ###############")
	startTime = time.time()
	# modelPath = '..\\..\\python_environments\\prediction-3.10.9\\Scripts\\nnUNetv2_predict.exe'
	# cmd = [modelPath, '-i', dir_dicom_to_nii, '-o', predictedNiiFile, '-d', '99', '-c', '3d_fullres' , '--disable_tta', '-tr', 'nnUNetTrainer_650epochs', '-f', '1, 4']

	predictNNUNet(os.path.join(RESULTS_FOLDER,dataset, 'nnUNetTrainer_650epochs__nnUNetPlans__3d_fullres'),
	dir_dicom_to_nii,
	predictedNiiFile,
	[1])

	print("Prediction MR done", time.time()-startTime)

	startTime = time.time()

	predicted_structures = ["background", "BRAIN", "AMYGDALAE", "BRAINSTEM", "CAUDATENUCLEI", "CEREBELLUM", "CHIASM", "COCHLEAS", "CORNEAS", "CORPUSCALLOSUM", "FORNICES", "GLANDPINEAL", "HIPPOCAMPI", "HYPOTHALAMI", "LACRIMALGLANDS", "LENSES", "OPTICNERVES", "ORBITOFRONTALS", "PITUITARY", "RETINAS", "THALAMI", "VSCCs"]
	rtstruct_colors = [[255,0,0]]*len(predicted_structures)

	niiFile = os.path.join(predictedNiiFile, '1a_001.nii.gz')

	# POSTPROCESSING TO CONVERT FROM INTEGERS TO 2**i, ADD CONTOURS EXISTS, AND SMOOTH
	integer_to_onehot(niiFile)
	print("POST processing convert from integers done: ", time.time()-startTime)

	# Convert CT image to NII #############



	convert_nii_to_dicom(dicomctdir=dir_ct_dicom, predictedNiiFile=niiFile, predictedDicomFile=predictedDicomFile,
	predicted_structures=predicted_structures, rtstruct_colors=rtstruct_colors, refCT=refCT)
	else:
	print("Not supported yet")


	startTime = time.time()
	uploadDicomToOrthanc(predictedDicomFile)
	print("Upload predicted result to Orthanc done: ", time.time()-startTime)

	# tempPath = 'C:\Temp\parrot_prediction'
	# regSeriesInstanceUID = '1.2.246.352.205.5029381855449574337.1508502639685232062'
	# runInterpreter = 'py3109'
	# patientName = 'P0461C0006I7638639'

	'''
	Prediction parameters provided by the server. Select the parameters to be used for prediction:
	[1] tempPath: The path where the predict.py is stored,
	[2] patientname: python version,
	[3] ctSeriesInstanceUID: Series instance UID for data set with modality = CT. To predict 'MR' modality data, retrieve the CT UID by the code (see Precision Code)
	[4] rtStructSeriesInstanceUID: Series instance UID for modality = RTSTURCT
	[5] regSeriesInstanceUID: Series instance UID for modality = REG,
	[6] runInterpreter: The python version for the python environment
	[7] oarList: only for dose predciton. For contour predicion oarList = []
	[8] tvList: only for dose prediction. For contour prediction tvList = []
	'''
	if __name__ == '__main__':
	predict(tempPath=sys.argv[1], patient_id=sys.argv[2], regSeriesInstanceUID=sys.argv[5], runInterpreter=sys.argv[6])
	# predict(tempPath=tempPath, patient_id=patientName, regSeriesInstanceUID=regSeriesInstanceUID, runInterpreter=runInterpreter)