MLScienceChallenge-rh 1.12

ML Science coding challenge for CZI Biohub application.

Raymond Hawkins - CZ Biohub - ML Science Coding Challenge

Installation Instructions

The following installation instructions were used to set up the environment in WSL (Ubuntu 24.04.2).

conda create -y --name 06_image_translation python=3.10
conda activate 06_image_translation
conda install -y ipykernel nbformat nbconvert black jupytext ipywidgets
pip install viscy[metrics,visual]==0.2.0rc1 jupyterlab
pip install torchmetrics==1.4.0
pip install ipyfilechooser ipympl ttach
pip install --upgrade jupyter

Part 1

• I worked through the notebook solution.ipynb in Windows Subsystem for Linux running Ubuntu with a single NVIDIA RTX 6000 Ada generation, changing paths where necessary
• In the notebook, I visualized the trained models outputs following the cells where the pretrained models outputs were visualized.
• To implement an interactive gui, I used ipywidgets to create a user interface that allows one to:

a. Select a .zarr file to load

b. Select a .ckpt VS model to load

c. Select and update the position in the zarr store to use as input to the model

d. Apply the model, and,

e. Visualize input, target, output and all encoder/decoder featuremaps using matplotlib.

The virtual stain viewer (and thus visualizing the trained model feature maps) can be demo'ed in demo_virtual_stain_viewer.ipynb

The interactive gui was inspired by JuNkIE (https://bitbucket.org/rfg_lab/junkie/src/master/junkie.py) a Jupyter Notebook image exploration package developed by our lab. Parts of the code were generated using PyCharm's AI assist autofill.

To run the gui, please ensure that you have IPython and IPyWidgets correctly configured to work with matplotlib. You will also need to install ipyfilechooser (https://pypi.org/project/ipyfilechooser/) and ipympl (https://www.npmjs.com/package/jupyter-matplotlib). Afterwards, you may need to reinstall Jupyter for changes to take effect:

pip install ipyfilechooser ipympl
pip install --upgrade jupyter

Part 2

• I wrote a script (predict_vs.py) based on the solution.ipynb file to infer membrane and nuclei from a phase video in time

• The script works with the associated dataset, although further work would need to be done to generalize it to all data (e.g. the script assumes the XY dimensions are 2046x2003 and pads accordingly, but this could be generalized to any possible dimension)

• To get rid of intensity fluctuations I have implemented naive gaussian smoothing in the temporal dimension and test-time augmentation (TTA). TTA reduces intensity fluctuations by ensembling the predictions from 24 randomly augmented versions of the input image. The result is a reduction in intensity fluctuations without phantom artifacts that are caused by the gaussian filter (at the cost of longer inference time)

  • TTA requires the ttach package (https://github.com/qubvel/ttach/tree/master) which can be installed with

    pip install ttach

• I have implemented the script as a CLI using an argparser (written by copilot and editted by myself) which takes the following arguments:

  • "--model" or "-m": the path to the .ckpt model checkpoint.
  • "--input" or "-i": the path to the ome zarr dataset.
  • "--output" or "-o": the path to save the output ome zarr predictions.
  • "--smoothing": the type of regularization for temporal intensity fluctuations (None, "gaussian", or "tta")