Product backlog

[JLrumberger]

Here we pool all the potential features and improvements for the self-supervised adaptation model. You can all edit the first post and add stuff directly to the list.

Milestones

• MVP: Adapt cellpose on one out of domain dataset, better after self-supervised adaptation (compare with cellpose v2 finetuning)
• Do actual science: try different uncertainty quantification methods (MC dropout, conformal prediction), regularization methods, etc.
• Provide a easy-to-use notebook that runs domain adaptation of cellpose on user provided data
• Make it work on 3D data
• Expand it to panoptic segmentation
• In-house use case: Julia Mellert's dataset of neurons in cell culture from alzheimer's model organism. Initial cellpose results weren't good so Julia curated annotations.

Data loading

• Data loader for cellpose data
• Data loader for tissuenet data <- Jannik
• Data loader for livecell data <- Dinesh

Image augmentation

• Image augmentation function that does inverse transformations and can be parameterized explicitly

Model architecture

• add cellpose to requirements and check dependencies errors -> Jannik
• Specify a format / class and adapt cellpose including post-processing
• Subclass the cellpose class for deepcell mesmer including post-processing

Model training

• Implement a default test_time_augmentation class which consists of the actual TTA functionality and a fusing functionality (that may need to be implemented individually for every model)
• Create StudentTeacher class, which inits with an instance of model, student augmentations and teacher augmentations. It ingests an unaugmented batch and outputs fused teacher prediction, student prediction, teacher prediction uncertainty estimate. It should be easily adaptable to other uncertainty measures, model architectures (mesmer?) and fusing methods (stardist).
• Prepare a training class that takes a model class, a dataloader and training procedure (i.e. StudentTeacher class) and starts the training. This class should be extendable, so that it works with other training procedures, dataloaders and models <- Fabian
• Add validation script for use during training. Should write artefacts to W&B or tensorboard <- Jannik
• Loss function probably like ACTIS with moving-average of the quantiles of some uncertainty metrics?
• Elastic regularization (i.e. not deviate too much from initialization weights)
• Uncertainty-based loss weighting comparable to EDAPS

Error Analysis

• Prepare an automatic way to quantify model performance to assess performance metrics of experiments, maybe use peters library
• visualize true/false positives and negatives and log the filter maps from the uncertainty-based filtering
• split validation loss into mse and ce loss for wandb logging

Misc

• Multi GPU (DDP) and mixed precision training for faster iterations
• Hyperparameter search with W&B
• Add pypi publish toolchain via github actions.

Experiments

• try multiple different augmentations on the same sample instead of bigger batches (https://arxiv.org/pdf/2203.03304.pdf)
• Crazy idea: use a second loss predictor model that takes input imgs and predictions as input and predicts loss value. Supervised: MSE(predicted loss, true loss), Unsupervised: fix weights of loss predictor and do gradient descent on predicted loss (this is kind of a reward model then..)
• Optimize hyperparameters of the post-processing (diameter, cellprob_threshold, flow_threshold) of cellpose on the tissuenet validation dataset and see if no test-time-augmentations vs. flips only vs. 90° rotations vs. full test-time-augmentations yields best results. Check if flows and uncertainties of the flow maps look better with/out test-time-augmentations. Should one do mean/max/median to combine flows? <- Fabian & Jannik (keep us updated)
• Visualize how individual teacher tta predictions look like pre-transformation, post-transformation and visualize differences and means.

Publication

• Write down planned experiments and how you want to present results
• Add github actions for pypi and readthedocs

Information

• lucidrains repo as a template for class design
• git book
• Nimbus test code as a reference