WIP: refactor contrastive learning code with virtual staining code

[mattersoflight]

This issue tracks our progress toward integration of the contrastive learning code with virtual staining code.

Our preprocessing code is currently in good shape and consists of:

• fluorescence deskewing/deconvolution: shrimPy
• phase deconvolution: recOrder
• Registration with fluorescence: shrimPy
• virtual staining of nuclei and membrane: VisCy
• Segmentation: VisCy (Adding segmentation utility functions  #108)
• Tracking: ultrack

We are still improving the tracking to capture cell division and cells near the boundary of the FOV: @tayllatheodoro

Training
It works well via pytorch lightning CLI and configs. The dataloader also works well with the HCS data format.

Pending improvements:
Architecture :

• Hard-coded stem stride in ContrastiveEncoder #138 @ziw-liu
• Unify stem implementation #120 @ziw-liu
• Update the projection head (normalization and size). #139 @ziw-liu

Data loader and loss functions:

• Implement NT-Xent loss (implicit negative samples) #136
• Time sampling for positive pair  #123
• data loader that pools multiple datasets.

Prediction and evaluation
It works well via pytorch lightning CLI and configs.

• Analyze embeddings with PCA and UMAP to generate a standard report: produce a report of useful visualizations to assess the dimensionality and features learned by embeddings #140 . @Soorya19Pradeep
• Refine the napari visualization tool for annotation of cell states in latent space: Visualize tracked cells and their features czbiohub-sf/napari-iohub#13. @ziw-liu

In this round, we should make any changes in the code path that can affect the architecture.

[ziw-liu]

Currently (Slurm) Cellpose segmentation is implemented in shrimPy, and will be moved to the new bioimage analysis repo.

@Soorya19Pradeep I'm outdated! It's now in #108

[ziw-liu]

For the napari UI I think we should first try interacting with the plugin through standardized data files so we don't have to maintain our own interface.

[ziw-liu]

The napari-clusters-plotter plugin does not implement readers. So it relies on what's available in the napari layer list (features are stored as attribute of the labels layer).

I now think a workable way is to implement a custom reader in napari-iohub for the images and tracks so the visualization is easier (handle mixed dimensions and scales etc). The ultrack plugin does load the extra columns in its output CSVs as layer features so it can be used by the cluster plotter.

As for clustering, I think dimensionality reduction should be done beforehand on all the cells, instead of on the limited number of cells in each FOV.

[mattersoflight]

The ultrack plugin does load the extra columns in its output CSVs as layer features so it can be used by the cluster plotter.

That's interesting. Can this work?

• Write projected embeddings to the same table as the output of ultrack
• load the tracks using ultrack plugin
• use napari-cluster-plotter to visualize labels that match projected embeddings.
  If yes, we don't need to create a new widget.

@ziw-liu please go ahead and decide on a useful and low-maintenance solution.

[mattersoflight]

@ziw-liu @alishbaimran Given our offline discussion, here is the prioritization of features:

1. update data format and data module for efficiency + allow selection of channels to encode,
2. define positive pairs based on temporal closeness, and
3. pool different datasets.

You could partition the refactor in 3 PRs, each of which implements the above and is tested with the corresponding training run.

We will train contrastive phenotyping models via the python scripts and CLIs that wrap these scripts. We don't have to prioritize integration with lighting CLI yet.

[mattersoflight]

@ziw-liu and @alishbaimran I think we can bypass the patchification step by chunking the Zarr store in C*Y*X sized chunks. The data chunked like this can be loaded fast enough on VAST since we just have to fetch the images within specific Z and T ranges. This will also simplify the processing pipeline and avoid the need to track one more data format.

I got this idea while exploring the data we are preparing for release with the paper on mantis.
Take a look at:
/hpc/projects/comp.micro/mantis/mantis_paper_data_release/figure_1.zarr