predict.py

import os
import subprocess
import tempfile

import h5py  # HDF5 is ilastik's preferred file format
import logging
import skimage

from cellprofiler_core.image import Image
from cellprofiler_core.module import Module
import cellprofiler_core.setting
from cellprofiler_core.setting.choice import Choice
from cellprofiler_core.setting.text import Pathname


logger = logging.getLogger(__name__)

__doc__ = """\
Predict
=======

Use an ilastik pixel classifier to generate a probability image. Each
channel represents the probability of the pixels in the image belong to
a particular class. Use **ColorToGray** to separate channels for further
processing. For example, use **IdentifyPrimaryObjects** on a
(single-channel) probability map to generate a segmentation. The order
of the channels in **ColorToGray** is the same as the order of the
labels within the ilastik project.

CellProfiler automatically scales grayscale and color images to the
[0.0, 1.0] range on load. Your ilastik classifier should be trained on
images with the same scale as the prediction images. You can ensure
consistent scales by:

-  using **ImageMath** to convert the images loaded by CellProfiler back
   to their original scale. Use these settings to rescale an image:

   -  **Operation**: *None*
   -  **Multiply the first image by**: *RESCALE_VALUE*
   -  **Set values greater than 1 equal to 1?**: *No*

   where *RESCALE_VALUE* is determined by your image data and the value
   of *Set intensity range from* in **NamesAndTypes**. For example, the
   *RESCALE_VALUE* for 32-bit images rescaled by "*Image bit-depth*" is
   65535 (the maximum value allowed by this data type). Please refer to
   the help for the setting *Set intensity range from* in
   **NamesAndTypes** for more information.

   This option is best when your training and prediction images do not
   require any preprocessing by CellProfiler.

-  preprocessing any training images with CellProfiler (e.g.,
   **RescaleIntensity**) and applying the same pre-processing steps to
   your analysis pipeline. You can use **SaveImages** to export training
   images as 32-bit TIFFs.

   This option requires two CellProfiler pipelines, but is effective
   when your training and prediction images require preprocessing by
   CellProfiler.

Additionally, please ensure CellProfiler is configured to load images in
the same format as ilastik. For example, if your ilastik classifier is
trained on RGB images, use **NamesAndTypes** to load images as RGB by
selecting "*Color image*" from the *Select the image type* dropdown. If
your classifier expects grayscale images, use **NamesAndTypes** to load
images as "*Grayscale image*".
"""


class Predict(cellprofiler_core.module.ImageProcessing):
    module_name = "Predict"

    variable_revision_number = 2

    def create_settings(self):
        super(Predict, self).create_settings()

        self.executable = Pathname(
            "Executable",
            doc="ilastik command line executable name, or location if it is not on your path."
        )

        self.project_file = Pathname(
            "Project file",
            doc="Path to the project file (\*.ilp)."
        )

        self.project_type = Choice(
            "Select the project type",
            [
                "Pixel Classification",
                "Autocontext (2-stage)"
            ],
            "Pixel Classification",
            doc="""\
Select the project type which matches the project file specified by
*Project file*. CellProfiler supports two types of ilastik projects:

-  *Pixel Classification*: Classify the pixels of an image given user
   annotations. `Read more`_.

-  *Autocontext (2-stage)*: Perform pixel classification in multiple
   stages, sharing predictions between stages to improve results. `Read
   more <http://ilastik.org/documentation/autocontext/autocontext>`__.

.. _Read more: http://ilastik.org/documentation/pixelclassification/pixelclassification
"""
        )

    def settings(self):
        settings = super(Predict, self).settings()

        settings += [
            self.executable,
            self.project_file,
            self.project_type
        ]

        return settings

    def visible_settings(self):
        visible_settings = super(Predict, self).visible_settings()

        visible_settings += [
            self.executable,
            self.project_file,
            self.project_type
        ]

        return visible_settings

    def run(self, workspace):
        image = workspace.image_set.get_image(self.x_name.value)

        x_data = image.pixel_data

        fin = tempfile.NamedTemporaryFile(suffix=".h5", delete=False)

        fout = tempfile.NamedTemporaryFile(suffix=".h5", delete=False)

        cmd = [
            self.executable.value,
            "--headless",
            "--project", self.project_file.value,
            "--output_format", "hdf5"
        ]

        if self.project_type.value in ["Pixel Classification"]:
            cmd += ["--export_source", "Probabilities"]
        elif self.project_type.value in ["Autocontext (2-stage)"]:
            x_data = skimage.img_as_ubyte(x_data)  # ilastik requires UINT8. Might be relaxed in future.

            cmd += ["--export_source", "probabilities stage 2"]
            #cmd += ["--export_source", "probabilities all stages"]

        cmd += [
            "--output_filename_format", fout.name,
            fin.name
        ]

        try:
            with h5py.File(fin.name, "w") as f:
                shape = x_data.shape

                
                f.create_dataset("data", shape, data=x_data)

            fin.close()

            fout.close()

            subprocess.check_call(cmd)

            with h5py.File(fout.name, "r") as f:
                y_data = f["exported_data"][()]

            y = Image(y_data)

            workspace.image_set.add(self.y_name.value, y)

            if self.show_window:
                workspace.display_data.x_data = x_data

                workspace.display_data.y_data = y_data

                workspace.display_data.dimensions = image.dimensions
        except subprocess.CalledProcessError as cpe:
            logger.error("Command {} exited with status {}".format(cpe.output, cpe.returncode), cpe)

            raise cpe
        except IOError as ioe:
            raise ioe
        finally:
            os.unlink(fin.name)

            os.unlink(fout.name)