# -*- coding: utf-8 -*-
# Author: Christoph Berger
# Script for evaluation and bulk segmentation of Brain Tumor Scans
# using the MICCAI BRATS algorithmic repository
#
# Please refer to README.md and LICENSE.md for further documentation
# This software is not certified for clinical use.
__version__ = "0.1"
__author__ = "Christoph Berger"
import errno
import glob
import json
import logging
import os
import os.path as op
import subprocess
import sys
import tempfile
import numpy as np
from . import fusionator
from .util import filemanager as fm
from .util import own_itk as oitk
from .util.citation_reminder import citation_reminder, new_segmentor_note
class Segmentor(object):
"""
Now does it all!
"""
@new_segmentor_note
@citation_reminder
def __init__(
self,
config=None,
fileformats=None,
verbose=True,
tty=False,
newdocker=True,
gpu="0",
):
"""Init the orchestra class with placeholders"""
self.noOfContainers = 0
self.config = []
self.directory = None
self.verbose = verbose
self.tty = tty
self.dockerGPU = newdocker
self.gpu = gpu
self.package_directory = op.dirname(op.abspath(__file__))
# set environment variables to limit GPU usage
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
if config is None:
config = op.join(self.package_directory, "config", "dockers.json")
if fileformats is None:
self.fileformats = op.join(
self.package_directory, "config", "fileformats.json"
)
else:
self.fileformats = fileformats
try:
configfile = open(config, "r")
self.config = json.load(configfile)
self.noOfContainers = len(self.config.keys())
configfile.close()
except IOError as e:
logging.exception("I/O error({0}): {1}".format(e.errno, e.strerror))
raise
except ValueError:
logging.exception("Invalid configuration file")
raise
except:
logging.exception("Unexpected Error!")
raise
def _getFileFormat(self, index):
return self.config[index]["fileformat"]
def _getContainerName(self, index):
return self.config[index]["name"]
def _getNumberOfContainers(self):
return len(self.config)
def _runDummyContainer(self, stop=False):
command = "docker run --rm -it hello-world"
subprocess.check_call(command, shell=True)
def _runContainer(self, id, directory, outputDir, outputName):
"""
Runs one container on one patient folder
"""
logging.info("Now running a segmentation with the Docker {}.".format(id))
logging.info("Output will be in {}.".format(outputDir))
params = self.config[id] # only references, doesn't copy
command = "docker run --rm"
# assemble the rest of the command
flags = params.get("flags", "")
# check if we need to map the user
if params.get("user_mode", False):
user_flags = "--user $(id -u):$(id -g)"
else:
user_flags = ""
# assemble the gpu flags if needed
if params["runtime"] == "nvidia":
if self.dockerGPU:
# TODO clean this up
gpu_flags = "--gpus device=" + str(self.gpu)
else:
gpu_flags = "--runtime=nvidia -e CUDA_VISIBLE_DEVICES=" + str(self.gpu)
else:
gpu_flags = ""
# assemble directory mapping
volume = "-v " + str(directory) + ":" + str(params["mountpoint"])
# assemble execution command
call = str(params["command"])
# stick everything together
command = (
command
+ " "
+ user_flags
+ " "
+ gpu_flags
+ " "
+ flags
+ " "
+ volume
+ " "
+ params["id"]
+ " "
+ call
)
if self.verbose:
print("Executing: {}".format(command))
try:
with open(
op.join(outputDir, "{}_output.log".format(outputName.split(".")[0])),
"w",
) as f:
subprocess.check_call(command, shell=True, stdout=f)
except Exception as e:
logging.error(
"Segmentation failed for case {} with error: {}".format(directory, e)
)
if "exit status 125" in str(e):
logging.error(
"DOCKER DAEMON not running! Please start your Docker runtime."
)
sys.exit(125)
return False
if self.verbose:
logging.info("Container exited without error")
# fileh.close()
return True
def _runIterate(self, dir, cid):
"""Iterates over a directory and runs the segmentation on each patient found"""
logging.info("Looking for BRATS data directories..")
for fn in os.listdir(dir):
if not os.path.isdir(os.path.join(dir, fn)):
continue # Not a directory
if "DE_RI" in fn:
logging.info("Found pat data: {}".format(fn))
try:
os.makedirs(os.path.join(os.path.join(dir, fn), "results"))
except OSError as err:
if err.errno != errno.EEXIST:
raise
logging.info("Calling Container: {}".format(cid))
if not self._runContainer(cid, os.path.join(dir, fn), dir):
logging.info(
"ERROR: Run failed for patient {} with container {}".format(
fn, cid
)
)
return False
# TODO: rename folder and prepend pat_id
# rename_folder(img_id, os.path.join(directory, fn), fn)
return True
def _multiSegment(self, tempDir, inputs, method, outputName, outputDir):
"""
multiSegment [summary]
Args:
tempDir ([type]): [description]
inputs ([type]): [description]
method ([type]): [description]
outputName ([type]): [description]
outputDir ([type]): [description]
"""
logging.debug("CALLED MULTISEGMENT")
fusion = fusionator.Fusionator()
for cid in self.config.keys():
# replace this with a call to single-segment
logging.info("[Orchestra] Segmenting with " + cid)
ff = self._format(self._getFileFormat(cid), self.fileformats)
for key, img in inputs.items():
savepath = op.join(tempDir, ff[key])
img = oitk.get_itk_image(img)
if self.verbose:
logging.info(
"[Weborchestra][Info] Writing to path {}".format(savepath)
)
oitk.write_itk_image(img, savepath)
if self.verbose:
logging.info("[Weborchestra][Info] Images saved correctly")
logging.info(
"[Weborchestra][Info] Starting the Segmentation with container {} now".format(
cid
)
)
status = self._runContainer(cid, tempDir, outputDir)
status = self._runContainer(cid, tempDir, outputDir, outputName)
if status:
if self.verbose:
logging.info("[Weborchestra][Success] Segmentation saved")
resultsDir = op.join(tempDir, "results/")
saveLocation = op.join(outputDir, cid + "_tumor_seg.nii.gz")
self._handleResult(cid, resultsDir, outputPath=saveLocation)
else:
logging.exception("Container run for CID {} failed!".format(cid))
fusion._dirFuse(
outputDir, method=method, outputPath=op.join(outputDir, outputName)
)
def _singleSegment(self, tempDir, inputs, cid, outputName, outputDir):
"""
singleSegment [summary]
Args:
tempDir ([type]): [description]
inputs ([type]): [description]
cid ([type]): [description]
outputName ([type]): [description]
outputDir ([type]): [description]
"""
ff = self._format(self._getFileFormat(cid), self.fileformats)
for key, img in inputs.items():
savepath = op.join(tempDir, ff[key])
img = oitk.get_itk_image(img)
if self.verbose:
logging.info("[Weborchestra][Info] Writing to path {}".format(savepath))
oitk.write_itk_image(img, savepath)
if self.verbose:
logging.info("[Weborchestra][Info] Images saved correctly")
logging.info(
"[Weborchestra][Info] Starting the Segmentation with {} now".format(cid)
)
status = self._runContainer(cid, tempDir, outputDir, outputName)
if status:
if self.verbose:
logging.info("[Weborchestra][Success] Segmentation saved")
resultsDir = op.join(tempDir, "results/")
self._handleResult(
cid, resultsDir, outputPath=op.join(outputDir, outputName)
)
# delete tmp directory if result was saved elsewhere
else:
logging.error("[Weborchestra][Error] Segmentation failed, see output!")
def segment(
self, t1=None, t1c=None, t2=None, fla=None, cid="mocker", outputPath=None
):
"""
segment [summary]
Args:
t1 ([type], optional): [description]. Defaults to None.
t1c ([type], optional): [description]. Defaults to None.
t2 ([type], optional): [description]. Defaults to None.
fla ([type], optional): [description]. Defaults to None.
cid (str, optional): [description]. Defaults to 'mocker'.
outputPath ([type], optional): [description]. Defaults to None.
"""
# Call output method here
outputName, outputDir = self._whereDoesTheFileGo(outputPath, t1, cid)
# set up logging (for all internal functions)
logging.basicConfig(
format="%(asctime)s %(levelname)s:%(message)s",
filename=op.join(outputDir, "segmentor_high_level.log"),
level=logging.DEBUG,
)
logging.getLogger().addHandler(logging.StreamHandler())
logging.debug("DIRNAME is: " + outputDir)
logging.debug("FILENAME is: " + outputName)
logging.info(
"Now running a new set of segmentations on input: {}".format(op.dirname(t1))
)
# switch between
inputs = {"t1": t1, "t2": t2, "t1c": t1c, "fla": fla}
# create temporary directory for storage
storage = tempfile.TemporaryDirectory(dir=self.package_directory)
# TODO this is a potential security hazzard as all users can access the files now, but currently it seems the only way to deal with bad configured docker installations
os.chmod(storage.name, 0o777)
tempDir = op.abspath(storage.name)
resultsDir = op.join(tempDir, "results")
os.mkdir(resultsDir)
# TODO this is a potential security hazzard as all users can access the files now, but currently it seems the only way to deal with bad configured docker installations
os.chmod(resultsDir, 0o777)
logging.debug(tempDir)
logging.debug(resultsDir)
if cid == "mav" or cid == "simple" or cid == "all":
# segment with all containers
logging.info("Called singleSegment with method: " + cid)
self._multiSegment(tempDir, inputs, cid, outputName, outputDir)
else:
# segment only with a single container
logging.info("Called singleSegment with docker: " + cid)
self._singleSegment(tempDir, inputs, cid, outputName, outputDir)
### Private utility methods below ###
def _whereDoesTheFileGo(self, outputPath, t1path, cid):
if outputPath is None:
outputDir = op.join(op.dirname(t1path), "output")
outputName = cid + "_segmentation.nii.gz"
elif outputPath.endswith(".nii.gz"):
if "~" in outputPath:
outputPath = op.expanduser(outputPath)
# valid filename
outputDir = op.dirname(outputPath)
outputName = op.basename(outputPath)
# if only a filename is passed, use the t1 directory
if outputDir == "":
outputDir = op.join(op.dirname(t1path), "output")
else:
outputDir = outputName = None
if outputDir is None or outputName is None:
raise ValueError(
"The outputPath is ambiguous and cannot be determined! path: {}, t1path: {}, cid: {}".format(
outputPath, t1path, cid
)
)
# build abspaths:
outputDir = op.abspath(outputDir)
try:
os.makedirs(outputDir, exist_ok=True)
except Exception as e:
print("could not create target directory: {}".format(outputDir))
raise e
return outputName, outputDir
def _handleResult(self, cid, directory, outputPath):
"""
This function handles the copying and renaming of the
Segmentation result before returning
"""
# Todo: Find segmentation result
contents = glob.glob(op.join(directory, "tumor_" + cid + "_class.nii*"))
if len(contents) == 0:
contents = glob.glob(op.join(directory, "tumor_*_class.nii*"))
if len(contents) == 0:
contents = glob.glob(op.join(directory, cid + "*.nii*"))
if len(contents) == 0:
contents = glob.glob(op.join(directory, "*tumor*.nii*"))
if len(contents) < 1:
logging.error(
"[Weborchestra - Filehandling][Error] No segmentation saved, the container run has most likely failed."
)
elif len(contents) > 1:
logging.warning(
"[Weborchestra - Filehandling][Warning] Multiple Segmentations found"
)
print("found files: {}".format(contents))
img = oitk.get_itk_image(contents[0])
labels = 0
exportImg = None
for _, c in enumerate(contents):
img = oitk.get_itk_image(c)
if labels < len(np.unique(oitk.get_itk_array(img))):
exportImg = img
labels = len(np.unique(oitk.get_itk_array(img)))
oitk.write_itk_image(exportImg, op.join(outputPath))
logging.warning(
"[Weborchestra - Filehandling][Warning] Segmentation with most labels ({}) for cid {} saved".format(
labels, cid
)
)
return
img = oitk.get_itk_image(contents[0])
for c in contents:
os.remove(op.join(directory, c))
oitk.write_itk_image(img, outputPath)
def _format(self, fileformat, configpath, verbose=True):
# load fileformat for a given container
try:
configfile = open(op.abspath(configpath), "r")
config = json.load(configfile)
configfile.close()
except IOError as e:
logging.exception("I/O error({0}): {1}".format(e.errno, e.strerror))
raise
except ValueError:
logging.exception("Invalid configuration file")
raise
except:
logging.exception("Unexpected Error!")
raise
logging.info(
"[Weborchestra][Success]Loaded fileformat: {}".format(config[fileformat])
)
return config[fileformat]
