Example #1
0
    def segmentation(self, image_in, session_id, model):
        logging.debug('Preparing input data for segmentation')
        self.reportProgress(0)

        result_file = tempfile.NamedTemporaryFile(suffix=self.outputFileExtension(), dir=self.aiaa_tmpdir).name
        aiaaClient = AIAAClient(self.server_url)
        params = aiaaClient.segmentation(model, image_in, result_file, session_id=session_id)

        extreme_points = params.get('points', params.get('extreme_points'))
        logging.debug('Extreme Points: {}'.format(extreme_points))

        self.reportProgress(100)
        return extreme_points, result_file
class AIAALogic():
    def __init__(self,
                 server_url=None,
                 server_version=None,
                 progress_callback=None):

        self.aiaa_tmpdir = slicer.util.tempDirectory('slicer-aiaa')
        self.volumeToImageFiles = dict()
        self.progress_callback = progress_callback
        self.useCompression = True

        # Create Single AIAA Client Instance
        self.aiaaClient = AIAAClient()
        self.setServer(server_url, server_version)

    def __del__(self):
        shutil.rmtree(self.aiaa_tmpdir, ignore_errors=True)

    def inputFileExtension(self):
        return ".nii.gz" if self.useCompression else ".nii"

    def outputFileExtension(self):
        # output is currently always generated as .nii.gz
        return ".nii.gz"

    def setServer(self, server_url=None, server_version=None):
        if not server_url:
            server_url = 'http://0.0.0.0:5000'
        if not server_version:
            server_version = 'v1'
        logging.debug('Using AIAA server {}: {}'.format(
            server_version, server_url))
        self.aiaaClient.server_url = server_url
        self.aiaaClient.api_version = server_version

    def setUseCompression(self, useCompression):
        self.useCompression = useCompression

    def setProgressCallback(self, progress_callback=None):
        self.progress_callback = progress_callback

    def reportProgress(self, progress):
        if self.progress_callback:
            self.progress_callback(progress)

    def list_models(self, label=None):
        #self.reportProgress(0)
        logging.debug('Fetching List of Models for label: {}'.format(label))
        result = self.aiaaClient.model_list(label)
        #self.reportProgress(100)
        return result

    def nodeCacheKey(self, mrmlNode):
        return mrmlNode.GetID() + "*" + str(mrmlNode.GetMTime())

    def segmentation(self, model, inputVolume):
        logging.debug('Preparing input data for segmentation')
        self.reportProgress(0)
        in_file = self.volumeToImageFiles.get(self.nodeCacheKey(inputVolume))
        if in_file is None:
            # No cached file
            in_file = tempfile.NamedTemporaryFile(
                suffix=self.inputFileExtension(), dir=self.aiaa_tmpdir).name
            self.reportProgress(5)
            start = time.time()
            slicer.util.saveNode(inputVolume, in_file)
            logging.info('Saved Input Node into {0} in {1:3.1f}s'.format(
                in_file,
                time.time() - start))
            self.volumeToImageFiles[self.nodeCacheKey(inputVolume)] = in_file
        else:
            logging.debug('Using cached image file: {}'.format(in_file))

        self.reportProgress(30)

        result_file = tempfile.NamedTemporaryFile(
            suffix=self.outputFileExtension(), dir=self.aiaa_tmpdir).name
        params = self.aiaaClient.segmentation(model,
                                              in_file,
                                              result_file,
                                              save_doc=False)

        extreme_points = params.get('points', params.get('extreme_points'))
        logging.debug('Extreme Points: {}'.format(extreme_points))

        self.reportProgress(100)
        return extreme_points, result_file

    def dextr3d(self, model, pointset, inputVolume, modelInfo):
        self.reportProgress(0)

        logging.debug('Preparing for Annotation/Dextr3D Action')

        node_id = inputVolume.GetID()
        in_file = self.volumeToImageFiles.get(self.nodeCacheKey(inputVolume))
        logging.debug('Node Id: {} => {}'.format(node_id, in_file))

        if in_file is None:
            in_file = tempfile.NamedTemporaryFile(
                suffix=self.inputFileExtension(), dir=self.aiaa_tmpdir).name

            self.reportProgress(5)
            start = time.time()
            slicer.util.saveNode(inputVolume, in_file)
            logging.info('Saved Input Node into {0} in {1:3.1f}s'.format(
                in_file,
                time.time() - start))

            self.volumeToImageFiles[self.nodeCacheKey(inputVolume)] = in_file
        else:
            logging.debug('Using Saved Node from: {}'.format(in_file))

        self.reportProgress(30)

        # Pre Process
        pad = 20
        roi_size = '128x128x128'
        if (modelInfo is not None) and ('padding' in modelInfo):
            pad = modelInfo['padding']
            roi_size = 'x'.join(map(str, modelInfo['roi']))

        result_file = tempfile.NamedTemporaryFile(
            suffix=self.outputFileExtension(), dir=self.aiaa_tmpdir).name
        self.aiaaClient.dextr3d(model, pointset, in_file, result_file, pad,
                                roi_size)

        self.reportProgress(100)
        return result_file