Purpose and Goal
Mitoception is a web application that provides a user interface for a machine learning model developed by Haining Zhang in the Meng Wang lab at Howard Hughes Medical Institute's Janelia Research Campus. The application enables researchers to upload >C. elegans mitochondria images and receive automated classifications.
My first implementation of the application used a FastAPI backend with PyTorch to perform inference on uploaded images. The new architecture gives the user the choice of using the server-based inference, or running inference entirely in the browser using WebAssembly, eliminating concerns external users might have of submitting their data to Janelia servers and reducing latency when there are many application users, while maintaining model accuracy.
Spotlight: Implementing persistent model storage with IndexedDB
One of the primary challenges in migrating to client-side inference was managing the 28MB ONNX model file. Intially, the file was downloaded on every page refresh, resulting in delays of several seconds to over a minute depending on network conditions. This created an unacceptable user experience.
I implemented a caching solution using the browser's IndexedDB API to store the model file locally after the initial download. The implementation checks IndexedDB for an existing model before attempting to fetch from the server, and stores the model blob after successful download. This approach reduced model loading time from several seconds down to approximately 25 milliseconds on subsequent visits, dramatically improving the user experience.
The solution required careful coordination between React hooks and
the ONNX Runtime Web initialization sequence. I used a
useEffect hook to trigger model retrieval and inference
session creation when users navigate to the analysis page, storing
the session in React state for reuse across multiple image analyses.
To handle the case where users might submit images before the model
finishes loading, I implemented a file queuing system that holds
uploaded files in state and automatically processes them once the
inference session is ready.
Spotlight: Troubleshooting tensor preprocessing for ResNet models
A critical challenge emerged during testing when the ONNX model produced different results than the original PyTorch implementation for identical input images. After extensive debugging, I discovered the issue was in the image preprocessing pipeline—specifically, the normalization step that converts pixel values into the tensor format expected by the ResNet architecture.
The ResNet model family expects input images to be normalized using specific mean and standard deviation values for each RGB channel: mean values of [0.485, 0.456, 0.406] and standard deviations of [0.229, 0.224, 0.225]. These values were derived from ImageNet statistics and are critical for proper model performance. The tutorial code I started from used a simplified normalization approach that didn't match these requirements.
I modified the imageDataToTensor function in
imageHelper.ts to apply the correct channel-specific
normalization. For each pixel, I converted the RGB values from 0-255
to 0-1 range, then applied the ImageNet normalization formula:
(pixel_value / 255.0 - mean) / std_dev for each
channel. This required restructuring the tensor layout to organize
pixel data by channel rather than by pixel, ensuring the normalized
values aligned with the model's expectations. After implementing
these changes, the ONNX model outputs matched the PyTorch
implementation exactly.
Current status
Mitoception is deployed in production but awaiting peer-review of the associated scientific publication before it is released for public use. The application supports batch uploads of up to 500 images and provides both visual results with color-coded probability overlays and downloadable CSV files for further analysis.
