Generating new pipeline with batched input

[naguileraleal]

Hello!

I'm trying to generate a custom .fpl pipeline for a .onnx model I've trained. I've based my pipeline on the "NucleiSegmentation" pipeline from FastPathology's pre-loaded pipelines.

The thing is: My model expects a (8,3,256,256) tensor as its input, but the PatchGenerator process object produces (I'm assuming) a (1,3,256,256) shaped output.

I confirmed this when trying to run the pipeline on a test WSI. The output from FastPathology was

Standard (std) exception caught: Got invalid dimensions for input: input.1 for the following indices index: 0 Got: 1 Expected: 8 Please fix either the inputs or the model.

Looking into the FAST documentation (python tutorial), I found the Batch data object, that enables one to batch several inputs. Now my problem is that I have no clue on how to state this object on the .fpl pipeline, since the inputs to this object are variables and it does not seem to follow the same API as the ProcessObjects did.

Is it possible to achieve what I'm looking for with a .fpl pipeline? Or do I have to move away from FastPathology and implement the pipeline in Python/C++ using FAST?

Thanks in advance!

PD: This is my custom pipeline

PipelineName "Glandulas segmentation"
PipelineDescription "Segmentacion de glandulas con U-Net de QuickAnnotator"
PipelineInputData WSI "Whole-slide image"
PipelineOutputData Segmentation stitcher 0
Attribute classes "Background;Cell nuclei"

### Processing chain

ProcessObject tissueSeg TissueSegmentation
Input 0 WSI

ProcessObject patch PatchGenerator
Attribute patch-size 256 256
Attribute patch-magnification 20
Attribute patch-overlap 0.1
Input 0 WSI
Input 1 tissueSeg 0

ProcessObject network SegmentationNetwork
Attribute scale-factor 0.003921568627451
Attribute model "$CURRENT_PATH$/../models/glandulas.onnx"
Input 0 patch 0

ProcessObject stitcher PatchStitcher
Input 0 network 0

### Renderers
Renderer imgRenderer ImagePyramidRenderer
Input 0 WSI

Renderer segRenderer SegmentationRenderer
Attribute border-opacity 0.5
Input 0 stitcher 0

[andreped]

I assume the model has been trained in QuickAnnotator which uses PyTorch as its backend.

The shape you have (8,3,256,256) means that you have channel first 3 before the image shape 256 x 256. By default, FAST expects channel last. However, it is possible to change this by setting the dimension-order attribute of the NeuralNetwork PO.

Regarding the 8, I would think that corresponds to your batch size, which in this case is hardcoded (not optimal!). Unless you are doing something awfully fancy, the network should support any batch size. In PyTorch, it is possible to define the input shape of the network as -1. In that case, the network can take any batch size.

If you are converting the model to ONNX yourself, it should be possible to set the input shape of the network after conversion, which I would set to (-1,256,256,3).

However, if you just wish to see if the model is working, you could try running this FPL instead:

PipelineName "Glandulas segmentation"
PipelineDescription "Segmentacion de glandulas con U-Net de QuickAnnotator"
PipelineInputData WSI "Whole-slide image"
PipelineOutputData Segmentation stitcher 0
Attribute classes "Background;Cell nuclei"

### Processing chain

ProcessObject tissueSeg TissueSegmentation
Input 0 WSI

ProcessObject patch PatchGenerator
Attribute patch-size 256 256
Attribute patch-magnification 20
Attribute patch-overlap 0.1
Input 0 WSI
Input 1 tissueSeg 0

ProcessObject batch ImageToBatchGenerator
Attribute max-batch-size 8
Input 0 patch 0

ProcessObject network NeuralNetwork
Attribute scale-factor 0.003921568627451
Attribute model "$CURRENT_PATH$/../models/glandulas.onnx"
Attribute dimension-ordering "channel-first"
Input 0 batch 0

ProcessObject converter TensorToSegmentation
Input 0 network 0

ProcessObject stitcher PatchStitcher
Input 0 converter 0

### Renderers
Renderer imgRenderer ImagePyramidRenderer
Input 0 WSI

Renderer segRenderer SegmentationRenderer
Attribute border-opacity 0.5
Input 0 stitcher 0

Note that I have made some modifications. I have introduced a batch generator, which is stacking up 8 patches at a time before they are fed to the NeuralNetwork PO. Note that I had to change the inference PO, as the SegmentationNetwork did not support the dimension-ordering attribute. For the NeuralNetwork PO, you also need to convert the output tensor to a segmentation, which the SegmentationNetwork does by default.

However, note that this might crash at the last patches, as it might fail to build up a batch of 8 patches, which your network does not support.

All this can be resolved by properly setting the input size during model conversion to ONNX. Give this a go first, and if you are interested, I could assist you in the model conversion to get the appropriate model input shape.

[MarkusDrange]

@andreped and I just managed to convert a PyTorch model to ONNX and deploy it in FastPathology.

However, the dimension-ordering attribute was not available in FastPathology v1.0.0. After upgrading FP to the latest artifact (see here), it worked as expected. Might be that the latest release v1.1.0 also works.

Regarding the batch size, this was handled by setting the dynamic_axis during conversion. As you are using QuickAnnotator, this might be tricky to change for you, but I'm sharing an example on how to do it below:

import torch
from models import SomeModel

dummy_input = torch.randn(1, 3, 299, 299, device="cuda")
model = SomeModel()
model.load_from_checkpoint("/path/to/pretrained/model.ckpt")

model.to_onnx("model_lightning_export.onnx", dummy_input, export_params=True,
    input_names=['input'], output_names=['output'],  
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})

[smistad]

By default, FAST expects channel last. However, it is possible to change this by setting the dimension-order attribute of the NeuralNetwork PO.

This is not entirely true. FAST uses channel last, but it can handle networks which are trained with channel first. FAST will try to detect this automatically, and then convert the input and output tensors accordingly. If it fails to detect it you can override it with the dimension-order attribute.

I would also encourage you to change the batch dimension to 1 or even better: -1. If you really want to do batch processing you can do that by using the ImageToBatchGenerator

[naguileraleal]

Thank you all for your responses!

I changed the batch size of the .onnx model to 1 and used the pipeline above by @andreped, changing the max-batch-size parameter to 1, and the model is now segmenting my WSI!

With respect to this line

model.to_onnx("model_lightning_export.onnx", dummy_input, export_params=True, input_names=['input'], output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})

I did not understand if the dynamic_axes argument was supposed to solve the variable-size batch problem. Nevertheless, I used it to get a new .onnx model. When changing back max-batch-size to 8, and processing the image with this new .onnx model, FastPathology crashes, and closes without error.

[andreped]

If you were able to set the batch size to 1, then you can remove the ImageToBatchGenerator PO, and then just connect the output of the PatchGenerator to the NeuralNetwork like you did before.

As @smistad mentioned, might be that your original FPL file works fine now, if you use the updated model.

Regarding batch processing, by setting the batch size to dynamic, you essentially say that it can take any size, which both for PyTorch and TF should be -1, AFAIK. If you set batch size to 1, it only supports batch size 1 and thus fails when using max-batch-size equals to 8.

Anyways, batch processing is not required for inference. As you now have a working solution, I believe this issue can be closed.