def GetContours(organ, patientFileName, threshold, withReal=True, tryLoad=True): #with real loads pre=existing DICOM roi to compare the prediction with model = Model.UNet() model.load_state_dict( torch.load( os.path.join( pathlib.Path(__file__).parent.absolute(), "Models/Model_" + organ.replace(" ", "") + ".pt"))) model = model.cuda() model.eval() #Make a list for all the contour images try: CTs = pickle.load( open( os.path.join( pathlib.Path(__file__).parent.absolute(), str("Prediction_Patients/" + patientFileName + "_Processed.txt")), 'rb')) except: CTs = DicomParsing.GetPredictionData(patientFileName, organ) if tryLoad: try: contourImages, contours = pickle.load( open( os.path.join( pathlib.Path(__file__).parent.absolute(), str("Prediction_Patients/" + organ + "/" + patientFileName + "_predictedContours.txt")), 'rb')) except: contours = [] zValues = [] #keep track of z position to add to contours after contourImages = [] exportNum = 0 for CT in CTs: ipp = CT[1] zValues.append(float(ipp[2])) pixelSpacing = CT[2] sliceThickness = CT[3] x = torch.from_numpy(CT[0]).cuda() xLen, yLen = x.shape #need to reshape x = torch.reshape(x, (1, 1, xLen, yLen)).float() predictionRaw = (model(x)).cpu().detach().numpy() prediction = PostProcessing.Process(predictionRaw, threshold) contourImage, contourPoints = PostProcessing.MaskToContour( prediction[0, 0, :, :]) contourImages.append(contourImage) contours.append(contourPoints) contours = PostProcessing.FixContours(contours) contours = PostProcessing.AddZToContours(contours, zValues) contours = DicomParsing.PixelToContourCoordinates( contours, ipp, zValues, pixelSpacing, sliceThickness) with open( os.path.join( pathlib.Path(__file__).parent.absolute(), str("Prediction_Patients/" + organ + "/" + patientFileName + "_predictedContours.txt")), "wb") as fp: pickle.dump([contourImages, contours], fp) else: contours = [] zValues = [] #keep track of z position to add to contours after contourImages = [] for CT in CTs: ipp = CT[1] zValues.append(float(ipp[2])) pixelSpacing = CT[2] sliceThickness = CT[3] x = torch.from_numpy(CT[0]).cuda() xLen, yLen = x.shape #need to reshape x = torch.reshape(x, (1, 1, xLen, yLen)).float() predictionRaw = (model(x)).cpu().detach().numpy() prediction = PostProcessing.Process(predictionRaw, threshold) contourImage, contourPoints = PostProcessing.MaskToContour( prediction[0, 0, :, :]) contourImages.append(contourImage) contours.append(contourPoints) contours = PostProcessing.FixContours(contours) contours = PostProcessing.AddZToContours(contours, zValues) contours = DicomParsing.PixelToContourCoordinates( contours, ipp, zValues, pixelSpacing, sliceThickness) with open( os.path.join( pathlib.Path(__file__).parent.absolute(), str("Prediction_Patients/" + organ + "/" + patientFileName + "_predictedContours.txt")), "wb") as fp: pickle.dump([contourImages, contours], fp) existingContours = [] if withReal: try: existingContours = pickle.load( open( os.path.join( pathlib.Path(__file__).parent.absolute(), str("Prediction_Patients/" + organ + "/" + patientFileName + "_ExistingContours.txt")), "rb")) except: pass Test.PlotPatientContours(contours, existingContours)
def GetContours(organ, patientFileName, path, threshold, withReal=True, tryLoad=True, plot=True): #with real loads pre=existing DICOM roi to compare the prediction with if path == None: #if no path supplied, assume that data folders are set up as default in the working directory. path = pathlib.Path(__file__).parent.absolute() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print("Device being used for predicting: " + device.type) model = Model.UNet() model.load_state_dict( torch.load( os.path.join(path, "Models/Model_" + organ.replace(" ", "") + ".pt"))) model = model.to(device) model.eval() contoursList = [] #The 1d contours list to be returned existingContoursList = [] #Make a list for all the contour images try: CTs = pickle.load( open( os.path.join( path, str("Predictions_Patients/" + patientFileName + "_Processed.txt")), 'rb')) except: CTs = DicomParsing.GetPredictionCTs(patientFileName, path) if tryLoad: try: contourImages, contours = pickle.load( open( os.path.join( path, str("Predictions_Patients/" + organ + "/" + patientFileName + "_predictedContours.txt")), 'rb')) except: contours = [] zValues = [] #keep track of z position to add to contours after contourImages = [] exportNum = 0 for CT in CTs: ipp = CT[1] zValues.append(float(ipp[2])) pixelSpacing = CT[2] sliceThickness = CT[3] x = torch.from_numpy(CT[0]).to(device) xLen, yLen = x.shape #need to reshape x = torch.reshape(x, (1, 1, xLen, yLen)).float() predictionRaw = (model(x)).cpu().detach().numpy() prediction = PostProcessing.Process(predictionRaw, threshold) contourImage, contourPoints = PostProcessing.MaskToContour( prediction[0, 0, :, :]) contourImages.append(contourImage) contours.append(contourPoints) contours = PostProcessing.FixContours(contours) contours = PostProcessing.AddZToContours(contours, zValues) contours = DicomParsing.PixelToContourCoordinates( contours, ipp, zValues, pixelSpacing, sliceThickness) for layer_idx in range(len(contours)): if len(contours[layer_idx]) > 0: for point_idx in range(len(contours[layer_idx])): x = contours[layer_idx][point_idx][0] y = contours[layer_idx][point_idx][1] z = contours[layer_idx][point_idx][2] contoursList.append(x) contoursList.append(y) contoursList.append(z) with open( os.path.join( path, str("Predictions_Patients/" + organ + "/" + patientFileName + "_predictedContours.txt")), "wb") as fp: pickle.dump([contourImages, contours], fp) else: contours = [] zValues = [] #keep track of z position to add to contours after contourImages = [] for CT in CTs: ipp = CT[1] zValues.append(float(ipp[2])) pixelSpacing = CT[2] sliceThickness = CT[3] x = torch.from_numpy(CT[0]).to(device) xLen, yLen = x.shape #need to reshape x = torch.reshape(x, (1, 1, xLen, yLen)).float() predictionRaw = (model(x)).cpu().detach().numpy() prediction = predictionRaw prediction[0, 0, :, :] = PostProcessing.Process( predictionRaw[0, 0, :, :], threshold) contourImage, contourPoints = PostProcessing.MaskToContour( prediction[0, 0, :, :]) contourImages.append(contourImage) contours.append(contourPoints) contours = PostProcessing.FixContours(contours) contours = PostProcessing.AddZToContours(contours, zValues) contours = DicomParsing.PixelToContourCoordinates( contours, ipp, zValues, pixelSpacing, sliceThickness) for layer_idx in range(len(contours)): if len(contours[layer_idx]) > 0: for point_idx in range(len(contours[layer_idx])): x = contours[layer_idx][point_idx][0] y = contours[layer_idx][point_idx][1] z = contours[layer_idx][point_idx][2] contoursList.append(x) contoursList.append(y) contoursList.append(z) with open( os.path.join( path, str("Predictions_Patients/" + organ + "/" + patientFileName + "_predictedContours.txt")), "wb") as fp: pickle.dump([contourImages, contours], fp) existingContours = [] if withReal: try: existingContours = pickle.load( open( os.path.join( path, str("Predictions_Patients/" + organ + "/" + patientFileName + "_ExistingContours.txt")), "rb")) for layer_idx in range(len(existingContours)): if len(existingContours[layer_idx]) > 0: for point_idx in range(len(existingContours[layer_idx])): x = existingContours[layer_idx][point_idx][0] y = existingContours[layer_idx][point_idx][1] z = existingContours[layer_idx][point_idx][2] existingContoursList.append(x) existingContoursList.append(y) existingContoursList.append(z) except: existingContours = DicomParsing.GetDICOMContours( patientFileName, organ, path) for layer_idx in range(len(existingContours)): if len(existingContours[layer_idx]) > 0: for point_idx in range(len(existingContours[layer_idx])): x = existingContours[layer_idx][point_idx][0] y = existingContours[layer_idx][point_idx][1] z = existingContours[layer_idx][point_idx][2] existingContoursList.append(x) existingContoursList.append(y) existingContoursList.append(z) if plot == True: Test.PlotPatientContours(contours, existingContours) return contoursList, existingContoursList