Generate a probability segmentation output insted of binary segmentation

[Peregalli]

Hi!
My aim is to have a probability segmentation output of the whole WSI and export it as a numpy array to use it as an image. I have to separate scripts. The first one uses the colon the colon epithelium segmentation HE model with my WSI and saves it in a .tiff file.
The second script imports .tiff segmentation and extracts aligned patches from the original WSI and segmentation inference and saves them as .png.

I would like to get the predictions as a probability (or a heatmap) without output thresholding. My problem is when I import .tiff file is already binary. I have seen that there is an attribute heatmapOutput in SegmentationNetwork class but when I set that parameter as True, I can not save the segmentation correctly. Another thing is that I am using getPatchAsImage instead of getLevelAsImage.

This is my code to get segmentation:

patchSize = 512

magnification = 10
overlapPercent = 0.1
scaleFactor=1.0

importer = fast.WholeSlideImageImporter.create(WSIPath)

imageRenderer = fast.ImageRenderer.create().connect(importer)

tissueSegmentation = fast.TissueSegmentation.create(threshold = 70).connect(importer)

patchGenerator = fast.PatchGenerator.create(patchSize, 
                                            patchSize, 
                                            magnification=magnification,
                                            maskThreshold=0.02, 
                                            overlapPercent=overlapPercent)\
    .connect(0, importer)\
    .connect(1, tissueSegmentation)

maxBatchSize = 1
batchGenerator = fast.ImageToBatchGenerator.create(maxBatchSize=maxBatchSize).connect(0, patchGenerator)

nn = fast.SegmentationNetwork.create(
    scaleFactor=scaleFactor,
    modelFilename=modelPath
    ).connect(0, batchGenerator)

stitcher = fast.PatchStitcher.create().connect(0, nn)

finished = fast.RunUntilFinished.create()\
    .connect(stitcher)

exporter = fast.TIFFImagePyramidExporter.create(os.path.join(output_dir,WSI_fn+'.tiff'))\
    .connect(finished)\
    .run()

And this is how I am extracting patches

level = 0
x = 12000
y = 24000
patch_size = 1024


def get_mask_patch(tiff_path: str, x: int, y : int, level : int,  W : int, H : int, patch_size : int) -> np.ndarray:
    """ Get a patch from a .tiff file scaled in order to be aligned with original WSI

    # Run importer and get data
    mask_pyramid_image = fast.fast.TIFFImagePyramidImporter\
        .create(tiff_path)\
        .runAndGetOutputData()

    w = mask_pyramid_image.getLevelWidth(level)
    h = mask_pyramid_image.getLevelHeight(level)

    access_mask = mask_pyramid_image.getAccess(fast.ACCESS_READ)
    
    #Transform coordinates to mask space
    x = int(x/W*w)
    y = int(y/H*h)

    mask = access_mask.getPatchAsImage(level, x, y, int(patch_size/W*w), int(patch_size/H*h))

    # Convert FAST image to numpy ndarray 
    mask = np.asarray(mask) #Here is already binary
    mask = (mask*255).astype(np.uint8)
    mask = cv.resize(mask,(patch_size,patch_size))

    return mask


def main():
    # Run importer and get data
    wsi_pyramid_image = fast.fast.TIFFImagePyramidImporter\
        .create(WSI_path)\
        .runAndGetOutputData()

    W = wsi_pyramid_image.getLevelWidth(level)
    H = wsi_pyramid_image.getLevelHeight(level)


    # Extract specific patch at highest resolution
    access_wsi_image = wsi_pyramid_image.getAccess(fast.ACCESS_READ)
    wsi_image = access_wsi_image.getPatchAsImage(level, x, y, patch_size,patch_size)
    wsi_image = np.asarray(wsi_image)

    mask = np.zeros((patch_size,patch_size,3))
    mask[:,:,0] = get_mask_patch(TIFF_path_1, x, y, level, W, H, patch_size)
    if TIFF_path_2 is not None:
        mask[:,:,1] = get_mask_patch(TIFF_path_2, x, y, level, W, H, patch_size)

    # This should print: (1024, 1024, 3) uint8 255 1
    plt.imshow(wsi_image)
    plt.imshow(mask,alpha = 0.5)
    plt.savefig(f'image_{level}_{y}_{x}.png')
    plt.show()

    cv.imwrite('Lady_path.png',wsi_image[:,:,::-1])

In this second script I am using two segmentations, one from colon epithelium model and the other from nuclei segmentation, however for both segmentation scripts I have the same problem. I would greatly appreciate your help!

[andreped]

Hello, @Peregalli!

AFAIK, FAST does not have support to store the softmax output as a pyramidal image.

If you wish to make such a pyramidal image, you will have to run your network on the WSI and extract the softmax output as individual image patches and stitch these together yourself. Note that this may be tricky to do if you are doing overlapping inference.

I have modified your first script which should extract patches at the requested patch size, and then you can do whatever you want with these later on:

# hyperparams
patchSize = 512
magnification = 10
overlapPercent = 0.0
scaleFactor=1.0

# pipeline
importer = fast.WholeSlideImageImporter.create(WSIPath)

tissueSegmentation = fast.TissueSegmentation.create(threshold=70).connect(importer)

patchGenerator = fast.PatchGenerator.create(patchSize, patchSize, magnification=magnification,
                                            maskThreshold=0.02, overlapPercent=overlapPercent)\
    .connect(0, importer)\
    .connect(1, tissueSegmentation)

network = fast.NeuralNetwork.create(scaleFactor=scaleFactor, modelFilename=modelPath)\
     .connect(0, patchGenerator)

for i, patch in enumerate(fast.DataStream(network)):
     patch = np.asarray(patch)
     print("patch:", i, patch.shape)

     # [...do whatever post-processing you want here...]

Some notes:

• The reason why the result is "thresholded/argmaxed", is because you are using the SegmentationNetwork PO. Hence, I replaced this PR with the NeuralNetwork PO which should return the raw softmax output.
• I have also removed the ImageToBatchGenerator PO, as it was redundant in your current pipeline.
• Note that to get started, I would avoid doing overlapping inference (set overlapPercent=0), as it might be tricky for you to get the predicted patches aligned with the WSI.
• I noticed in your second script that you wish to extract the patches from the final pyramidal image. I guess you can just do that directly in the for loop I made above.
• If you wish to store the softmax output from individual patches on disk, I recommend storing it in another format than PNG (or similar). HDF5 is a good option and actually what FAST uses to store the output generated by a patch-wise WSI classifier.

Hope it helps! Let me know if you have further questions.

---

EDIT: If you wish to get the image patch and softmax output pairs, you could try to do something like this in the loop after the NeuralNetwork PO:

for i, (image, pred) in enumerate(zip(fast.DataStream(patchGenerator), fast.DataStream(network))):
     image = np.asarray(image)
     pred = np.asarray(pred)
     print("\npatch:", i, image.shape, pred.shape)

     # [...do whatever post-processing you want here...]

But it might be that the extracted image patches will not be on-to-one with the pred patches, but give this a try first.

[andreped]

I made a working example in this gist that can be tested on Google CoLab directly.

I also share the code below:

import fast
import matplotlib.pyplot as plt
import numpy as np

# enable verbose in FAST
fast.Reporter.setGlobalReportMethod(fast.Reporter.COUT)

# Download a nuclei segmentation model from the DataHub
model = fast.DataHub().download('nuclei-segmentation-model')

# hyperparams
patchSize = 512
magnification = 20
overlapPercent = 0.0
scaleFactor=1.0 / 255.0

# pipeline
importer = fast.WholeSlideImageImporter.create("./CMU-1.svs")  # this is a local WSI I downloaded using wget

tissueSegmentation = fast.TissueSegmentation.create(threshold=70).connect(importer)

patchGenerator = fast.PatchGenerator.create(patchSize, patchSize, magnification=magnification,
                                            maskThreshold=0.02, overlapPercent=overlapPercent)\
    .connect(0, importer)\
    .connect(1, tissueSegmentation)

network = fast.NeuralNetwork.create(scaleFactor=scaleFactor, modelFilename=model.paths[0] + '/high_res_nuclei_unet.onnx')\
     .connect(0, patchGenerator)

for i, (image, pred) in enumerate(zip(fast.DataStream(patchGenerator), fast.DataStream(network))):
     image = np.asarray(image)
     pred = np.asarray(pred)
     print("patch:", i, image.shape, pred.shape)
     argmaxed = np.argmax(pred, axis=-1)

     fig, ax = plt.subplots(1, 4)

     ax[0].imshow(image)
     ax[1].imshow(pred[..., 0], vmin=0, vmax=1, cmap="gray")
     ax[2].imshow(pred[..., 1], vmin=0, vmax=1, cmap="gray")
     ax[3].imshow(argmaxed, vmin=0, vmax=1, cmap="gray")
     plt.show()

[Peregalli]

Thank you very much for your response.
I have another question. I would like to perform some postprocessing on the mask but then generate the .tiff file again to be able to visualize it. With my code, I save the patch positions at the image name, Could I generate the pyramidal tiff again?

[andreped]

Could I generate the pyramidal tiff again?

The problem is that you cannot store floats in the pyramidal image, like the softmax output. It needs to be stored as uint8 - one uint per pixel. If you are fine thresholding the post-processed softmax image to get uint8 (each pixel is assigned to a class), then you can simply make a custom python process object, which you can connect to a basic processing pipeline. That way, if you are fine with storing the final post-processed result "thresholded/argmaxed", without raw softmax output, then it should work.

If you really wish to store softmax outputs in a pyramidal image, I guess you could try using tifffile. I have previously used it to build my own pyramidal images in Python, but I have only assigned one value per pixel (uint8). Not sure what you suggest is offered by standard TIFF.

Below is an example how to make your own custom Python PO, which then can modify the NeuralNetwork PO before storing the result in a pyramidal image:

import fast
import numpy as np

class CustomPythonPO(fast.PythonProcessObject):
    def __init__(self, value=2):
        super().__init__()
        self.createInputPort(0)
        self.createOutputPort(0)
        self.value = value  # not actually used in this PO, but now you know how to add arguments

    def execute(self):
        image = self.getInputData()
        new_array = np.asarray(image)
        
        # make some algorithm here, e.g., flips the channels and argmax to get uint8 result (one uint assigned to each pixel)
        new_array = new_array[..., ::-1]
        new_array = np.argmax(new_array, axis=-1)

        # Create new fast image and add as output
        new_output_image = fast.Image.createFromArray(new_array)
        new_output_image.setSpacing(image.getSpacing())
        self.addOutputData(0, new_output_image)

Then you can connect it to your original pipeline script, which should be able to store the result in a pyramidal image:

# hyperparams
patchSize = 256
magnification = 5
overlapPercent = 0.0
scaleFactor = 1.0 / 255.0

# pipeline with custom PO
importer = fast.WholeSlideImageImporter.create("./CMU-1.svs")

tissueSegmentation = fast.TissueSegmentation.create(threshold=70).connect(importer)

patchGenerator = fast.PatchGenerator.create(patchSize, 
                                            patchSize, 
                                            magnification=magnification,
                                            maskThreshold=0.02, 
                                            overlapPercent=overlapPercent)\
    .connect(0, importer)\
    .connect(1, tissueSegmentation)

network = fast.NeuralNetwork.create(scaleFactor=scaleFactor, modelFilename=model.paths[0] + '/high_res_nuclei_unet.onnx')\
     .connect(patchGenerator)

post_process = CustomPythonPO.create().connect(network)

stitcher = fast.PatchStitcher.create().connect(post_process)

finished = fast.RunUntilFinished.create().connect(stitcher)

exporter = fast.TIFFImagePyramidExporter.create("./result.tiff")\
    .connect(finished)\
    .run()

I made a gist here that enables you and others to play around with the script on CoLab easily.

To visualize the result, you can do as you normally would in real-time, or render the generated result after export is finished, like so:

mask_importer = fast.TIFFImagePyramidImporter.create('./result.tiff')

segmentationRenderer = fast.SegmentationRenderer.create()\
    .connect(mask_importer)

fast.SimpleWindow2D.create()\
    .connect(segmentationRenderer)\
    .run()

Note that rendering using fast.SimpleWindow2D does not work in CoLab, so I did not include it in the notebook, but it should work locally.

To render result on top of the WSI as overlay, do something like this:

importer = fast.WholeSlideImageImporter.create("./CMU-1.svs")
mask_importer = fast.TIFFImagePyramidImporter.create('./result.tiff')

# Display the stitched segmentation results on top of the WSI
renderer = fast.ImagePyramidRenderer.create()\
    .connect(importer)

segmentationRenderer = fast.SegmentationRenderer.create()\
    .connect(mask_importer)

fast.SimpleWindow2D.create()\
    .connect(renderer)\
    .connect(segmentationRenderer)\
    .run()

[smistad]

Hi

As @andreped says, the TIFF backend in FAST does not support floats yet, which we need for tensors with giga resolution.
I am working on heatmap overlay on WSIs (and other images) in FAST right now, but not for giga resolution heatmaps.
I will think about this and see if I come up with a solution that can fit.

[Peregalli]

Thanks for your comments! I am using a CustomPostrocessing class as @andreped suggested. But I need to align prediction and image, is there a way to do that?

patchGenerator = fast.PatchGenerator.create(patchSize, 
                                            patchSize, 
                                            magnification=magnification, 
                                            overlapPercent=overlapPercent,
                                            maskThreshold = 0.02)\
    .connect(0, importer)\
    .connect(1, tissueSegmentation)

# Create neural network object ###
nn = fast.NeuralNetwork.create(
    scaleFactor=scaleFactor,
    modelFilename=modelPath,
    ).connect(0, patchGenerator)

post_process = Posprocessing.create(threshold = 0.2).connect(0,nn).connect(1,patchGenerator)

and Postprocessing class:

class Posprocessing(fast.PythonProcessObject):
    def __init__(self, threshold):
        super().__init__()
        self.createInputPort(0)
        self.createInputPort(1)
        self.createOutputPort(0)
        self.threshold = threshold
        self.epithelium_fill_filter = HoleFillFilter()   

    def execute(self):
        #Get prediction 
        prediction = self.getInputData(0)
        image = np.asarray(self.getInputData(1))
        new_prediction = ((np.asarray(prediction)[...,1]>self.threshold)*255).astype(np.uint8)

        #Postprocesing prediction
        mask_fill = self.epithelium_fill_filter.apply(new_prediction)
        
        # Do other this with image and mask_fill

        new_output_image = fast.Image.createFromArray(mask)
        new_output_image.setSpacing(prediction.getSpacing())
        self.addOutputData(0, new_output_image)

But image and prediction are not aligned.

[andreped]

But image and prediction are not aligned.

@Peregalli Not sure if I understand what you mean. Do you mean that the rendered predictions on top of the WSI are not aligned after adding the custom post-processing PO, or do you mean that the pred and image within the custom PO are not aligned/paired?

I just tested the pipeline below on my macOS, and it seems to render fine with the custom PO, but I did not check that the pred and image within the custom PO were correctly paired.

import fast
import numpy as np

# enable verbose in FAST
fast.Reporter.setGlobalReportMethod(fast.Reporter.COUT)

class TwoInputCustomPythonPO(fast.PythonProcessObject):
    def __init__(self, value=2):
        super().__init__()
        self.createInputPort(0)
        self.createInputPort(1)
        self.createOutputPort(0)
        self.value = value  # not actually used in this PO, but now you know how to add arguments

    def execute(self):
        pred = self.getInputData(0)
        new_array = np.asarray(pred)
        patch = self.getInputData(1)  # not actually used in this PO, just as an example

        # make some algorithm here, e.g., thresholds the softmax output to get uint8 result (one uint assigned per pixel)
        new_array = np.argmax(new_array, axis=-1, keepdims=True).astype("uint8")

        # Create new fast image and add as output
        new_output_image = fast.Image.createFromArray(new_array)
        new_output_image.setSpacing(pred.getSpacing())
        self.addOutputData(0, new_output_image)

# Download a nuclei segmentation model from the DataHub
model = fast.DataHub().download('nuclei-segmentation-model')

# hyperparams
patchSize = 256
magnification = 20
overlapPercent = 0.0
scaleFactor = 1.0 / 255.0

# pipeline with custom PO
importer = fast.WholeSlideImageImporter.create("./CMU-1.svs")

tissueSegmentation = fast.TissueSegmentation.create(threshold=70).connect(importer)

patchGenerator = fast.PatchGenerator.create(patchSize,
                                            patchSize,
                                            magnification=magnification,
                                            maskThreshold=0.02,
                                            overlapPercent=overlapPercent)\
    .connect(0, importer)\
    .connect(1, tissueSegmentation)

network = fast.NeuralNetwork.create(scaleFactor=scaleFactor, modelFilename=model.paths[0] + '/high_res_nuclei_unet.onnx')\
     .connect(patchGenerator)

post_process = TwoInputCustomPythonPO.create().connect(0, network).connect(1, patchGenerator)

stitcher = fast.PatchStitcher.create().connect(post_process)

# Display the stitched segmentation results on top of the WSI
renderer = fast.ImagePyramidRenderer.create()\
    .connect(importer)

segmentationRenderer = fast.SegmentationRenderer.create()\
    .connect(stitcher)

fast.SimpleWindow2D.create()\
    .connect(renderer)\
    .connect(segmentationRenderer)\
    .run()

Rendered result below:

[Peregalli]

I mean the pred and image within the custom PO. For example, if I visualize with matplotlib library both image and predicted mask are not aligned :

Another example:

However, when I use SimpleWindow2D I see both are correctly aligned. Is there a way to know the position from where that image is being extracted? In order to align the image with the mask.

[andreped]

I think what may happen is that in your custom PO, when you feed the PatchGenerator PO as input, you get the next output of the PatchGenerator and not the same original patch that was given to the NeuralNetwork PO. This is what I hypothesized could happen earlier in this thread.

My initial thought was to use fast.DataStream to synchronize the patch generation stream, but I do not think it can be connected between POs. Any ideas on how to resolve this issue, @smistad?