Bringing ANNs and Intellesis together
Training artificial neural networks (ANNs) to tackle semantic segmentation problems has become a very popular task and an increasing number of solutions have become available that require little technical understanding to train ANNs. Of course, this statement does not hold for models highly optimized for specific use cases where solving the problem at hand requires a high amount of experience, creativity and some “magic”. However, in many cases it is possible to achieve decent results with widely used ANNs model architectures, a sufficient amount of data and enough compute power.
While about a decade ago solving computer vision tasks required an expert to design algorithms and understand the math behind them, ANNs and simple-to-use frameworks like TensorFlow and PyTorch allow learning those algorithms from data and therefore can be treated as a black-box even achieving super-human performance in some cases. This simplicity and the high expressiveness of neural network models are presumably the reason that ANNs have become extremely popular and there is no way around this technology in research as well as industry. It should, however, be mentioned here that it is still extremely helpful to understand what happens behind the scenes.
The challenges of deploying a trained model to large images
While research mostly focuses on training and evaluating models, deployment is a different issue that poses different challenges. Especially in microscopy there usually are not only captured small images but stacks of images, time series of images or tiled scenes that are stitched to one large image of multiple or even hundreds of gigabytes. The training process, in contrast, is typically performed on small image crops (e.g. 256 x 256) such that the data for the calculations fit into the GPU or main memory of the employed machine. Applying a trained model on considerably larger images as mentioned above brings about several pitfalls that need to be handled.
Input image with nuclei
After Segmentation
The ZEN Intellesis infrastructure provides a solution for that problem. It applies a sophisticated tiling strategy that transforms multi-dimensional image captures – as they are generated by, for example, ZEISS microscopes – to image tiles of shape (height x width x color channels) where height and width are defined by the applied ANNs model. Finally, it performs inference with the model and reassembles the resulting segmentation masks to fit the initially captured image with intact metadata. It, therefore, offers a solution to apply a trained ANN to large images without worrying too much about the required memory and the size of the image.
czmodel – From TensorFlow to Intellesis
Starting with ZEN Blue3.2 it will be possible to run prediction with custom segmentation models implemented in TensorFlow 2. To be compatible with the Intellesis infrastructure, a TensorFlow/Keras model has to comply with our specification and some metadata must be provided that is described in detail in the official ANN model specification in the project description of our czmodel package. This blogpost introduces a python library called czmodel that offers a set of tools to easily convert a trained TensorFlow/Keras model to a CZMODEL file that can be directly loaded into ZEN Intellesis. Generating a CZMODEL file this way is just a matter of a few lines of code and can therefore easily be integrated into any training pipeline as well as custom Apeer modules.
In this blog post, we assume that there is some kind of training pipeline in place that generates trained TensorFlow/Keras models as instances of tensorflow.keras.Model. For a simple example of such a pipeline see the attached Jupyter notebook in our demo project. Note that the notebook is not to be understood as a best practice guide for training ANN models but rather illustrates the use of the czmodel library with a very simple model generation process that will usually be more sophisticated.
The entire czmodel package operates on so-called ModelSpec objects that contain all information needed to convert a TensorFlow/Keras model to CZMODEL. Specifically, it contains the trained model itself and a metadata object called ModelMetadata that reflects the JSON metadata described in the ANN model specification as a python class. Thus, the typical workflow of converting a TensorFlow/Keras model to CZMODEL is to create a ModelMetadata object, wrap it into a ModelSpec object together with the corresponding TensorFlow/Keras model and, finally, provide the ModelSpec object to the conversion function that generates the CZMODEL file.
The czmodel package supports three different ways to generate a CZMODEL.
Convert a load TensorFlow/Keras model to CZMODEL using a ModelSpec
This scenario reflects exactly the workflow described above. The ModelMetadata and ModelSpec objects for a loaded TensorFlow/Keras model named “model” are created as follows:
Convert a Saved Model on disk to CZMODEL using a ModelSpec
In this case, we assume the TensorFlow/Keras model not to reside in the main memory but on disk in the saved model format. Let “path” be the path to the SavedModel folder. We can then apply the same procedure as shown above with loaded models but instantiate the model spec object with the path of the model instead of the loaded model:
The czmodel library will automatically load the model from the given path before wrapping it into a CZMODEL.
Convert a Saved Model on disk to CZMODEL by providing a JSON meta data file
This functionality replicates the import functionality of ZEN Intellesis to load models via a JSON specification. If models should e.g. be exchanged with third parties, it is usually easier to deliver CZMODEL files that contain everything needed to run the model within Intellesis than exchanging JSON files with the corresponding SavedModel folders that would, nonetheless, need to be bundled another way.
The czmodel library provides a simple conversion routine that does all this in one line:
Note that the path to the model is implicitly defined by the provided JSON file.
Adding-processing and setting the tile size
Most trained models require some kind of pre-processing (e.g. normalization/standardization) to transform the input data to the distribution the model was trained on. The Intellesis inference engine requires this to be part of the model while – depending on the training infrastructure - this is not necessarily the case during training.
All conversion functions in the model library wrap the provided model into a deployment model if pre-processing is applied or the spatial dimensions should be reset. They additionally accept a list of TensorFlow/Keras layers as a preprocessing argument that will be prepended to the provided model. Additionally, setting the parameter spatial_dims allows redefining the spatial dimensions of the expected model input that implicitly defines the size of the tiles provided to the inference engine. Note that this is only possible for models that are invariant w.r.t the spatial dimensions.
While the czmodel.util.preprocessing module provides a set of useful pre-processinglayers, any TensorFlow/Keras layer can in theory be used for pre-processing. Thefollowing call generates a CZMODEL from the above defined ModelSpec object, appends an RgbToBgr pre-processing layer to transform the input images to a different color space and sets the expected size of the input images to 512 x512:
Remarks
The current implementation of the czmodel heavily relies on the SavedModel import/export functionality in TensorFlow 2. We, therefore, recommend to always use the latest version of TensorFlow 2 compatible with czmodel as there are continuously released improvements and bug fixes in the TensorFlow code base.
