Mutex Watershed in the Web (as an Imjoy Plugin)
The goal of this project is to make the mutex watershed algorithm and our trained networks as accessible as possible. What is the most accessible you ask? Running everything in the browser.
The mutex watershed segmentation pipeline consists of two steps. We need to take an input image and feed it through a trained neural network and then take its output and feed that into our C++ backend that runs the MWS.
For the network predictions, we are converting our trained pytorch networks into the ONNX format which then lets us run the same model in javascript. Unfortunately, this is not straightforward for the UNet architecture. ONNX does not support all the needed operators, at least not in a consistent version. To get support for 2D pooling we use this special ONNX branch. It gives us access to a 2D average pooling layer and lets us build a modified UNet in ONNX (with average pooling instead of max pooling).
To access our backend, Constantine has compiled our C++ Module to WebAssembly so we can call it in javascript.
The project is still in its infancy but since sharing is a breeze you get to experience this super early prototype with me. When you click on the EM microscopy image below, it kicks of the pipeline; the neural network predicts affinities and the Mutex Watershed turns those affinities into a segmentation. After a few seconds you can see both the network output and the segmentation on the right. Note that this is just a tech demo and the network has not been trained to optimality.
Clearly processing a static image is not the end of the road. The next step is to integrate this prototype into a functional imjoy plugin to make use of imjoys user interface and ease of deployment.