Implementation of Beta VAE for benchmarking (#273)

* simualte different embeddings

* update the msd calculation to re-use cdist functions in the repo

* adding a test for the msd

* removing unused msd functions

* renaming msd to compute_track_displacement

* default to cosine distance

* adding the gradient attribution video.

* extend to training ratios

* demo beta_vae 2.5D

* improving the logging for readability and drop pythae baseclasses

* condense the logging to have less tabs.

* fix disentagle metrics

* fixing beta warmup bug

* renaming to loss

* updating architecture to flatten vs spatial VAE with convs

* chaning to use mse with mean reduction and normalizing the kl loss by batch size.

* optunea proof of concept

* add normalized sampled into the transforms so we can use it with MONAIs vae

* update loss debugging code

* adding sync for disentaglement metrics

* adding the dataloader for rpe1 dataset and plotting utils

* cleanup the vae  and add the monai to lightning. adding configs

* add saving hyperparameters

* fix hyperparameter logging

* add embedding logging to the CLIP version

* test and plot of monaivae

* handle monai_vae 2d

* redifining rotation agumentsations

* adding optional scaling to phate

* adding alias and output 2d

* normalizing by also the latent dim and swapping to FP32 for forward pass to avoid overflow with log and exp

* update test for magnitudes

* expose the normalization for vae

* add sam 2 test

* refactor smoothness metrics

* rever to normalalize kl wrt to batch size  and removing the the beta min value

* commit dtwembeddings w sam

* added a clamp to logvar, switch to mse loss sum reduction like the original formulation.

* remove unecessary vae logging losses.

* add a way to handle when using 'mean' reduction for proper scaling

* adding optional config for middle slice index for computing sam2 embeddings and dinov3

* converting latent stats active_dimensions parameter to float to remove warning

* ruff

* removing the optuna config

* numpy docstring

* fix compute smoothness script

* archiving old scripts

* re org the pc features scripts

* embeddings for phase

* add smoothness (mean rand vs adj frame) to the csv

* archiving old beta vae code

* ruff

* fix format

* fix typo

* remove the archived unecessary files

* remove the test run archived file

* adding normalizeintensity

* fixing the vae_logging typing and removing PC plotting from here

* fixing the compute_embedding_smoothness docstring

* simplify the distance metrics and removing deprecated functions and scripts

* remove deprecated functions from clustering.py

* add timelpase to grad_attr.py script

* refactoring the betavaemodule. removing the hyperparamter logging, adding the nn.Module as input for typing purposes  and removing the fp32 custom fwd

* remove the optuna dependency

* deleting old msd test

* ruff format

* fix to explicitly stratify on fov level

* adding reference to dataset for rpe1

* fix pyproject.toml dev

* format and lint

* restore no-augmentation flag effect

* format tests

* rename the sam2 file

* removing unused arguments for logging embeddings.

* removing duplication in the lca

* remove disentaglement metrics

* vectorized the anchor filtering for celldivisiontriplet dataset

* map the channels to the rpe dataset convention

* fix logistic regresion standardization

* update rpe classifier to include mitosis

* ruff

* remove unused logging

* datamodule agnostic

* cleaning up duplicated code in the benchmarking

* cleanup vae

* keeping it consistent and using residual units

* fix typings betavaemonai

* update smoothness to handle adata

* update clustering method and add test

* pre-commit

* Update viscy/data/cell_division_triplet.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/benchmarking/DynaCLR/SAM2/sam2_visualizations.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/pseudotime_analysis/evaluation/compare_dtw_embeddings_sam2.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/pseudotime_analysis/evaluation/compare_dtw_embeddings_sam2.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/contrastive_phenotyping/evaluation/smoothness/compute_smoothness.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/pseudotime_analysis/evaluation/compare_dtw_embeddings_sam2.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/contrastive_phenotyping/evaluation/smoothness/compute_smoothness.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update applications/contrastive_phenotyping/evaluation/archive/ALFI_MSD_v2.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* valuerror on the fidn peaks function

* add literal to the betavae25d normalization

* clipping similarity that was breaking the tests

---------

Co-authored-by: Ziwen Liu <67518483+ziw-liu@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

applications/benchmarking/DynaCLR/DINOV3/config_dinov3_convnext_tiny.yml

@@ -0,0 +1,65 @@
+datamodule_class: viscy.data.triplet.TripletDataModule
+datamodule:
+  data_path: /hpc/projects/intracellular_dashboard/organelle_dynamics/2024_02_04_A549_DENV_ZIKV_timelapse/8-train-test-split/registered_test.zarr
+  tracks_path: /hpc/projects/intracellular_dashboard/organelle_dynamics/2024_02_04_A549_DENV_ZIKV_timelapse/8-train-test-split/track_test.zarr
+  batch_size: 32
+  final_yx_patch_size:
+  - 256
+  - 256
+  include_fov_names: null
+  include_track_ids: null
+  initial_yx_patch_size:
+  - 256
+  - 256
+  normalizations:
+  - class_path: viscy.transforms.ScaleIntensityRangePercentilesd
+    init_args:
+      b_max: 1.0
+      b_min: 0.0
+      keys:
+      - RFP
+      lower: 50
+      upper: 99
+  - class_path: viscy.transforms.NormalizeIntensityd
+    init_args:
+      keys:
+      - Phase3D
+  num_workers: 10
+  source_channel:
+  - RFP
+  - Phase3D
+  z_range:
+  - 15
+  - 45
+
+embedding:
+  pca_kwargs:
+    n_components: 8
+  phate_kwargs:
+    decay: 40
+    knn: 5
+    n_components: 2
+    n_jobs: -1
+    random_state: 42
+  reductions:
+  - PHATE
+  - PCA
+
+execution:
+  overwrite: false
+  save_config: true
+  show_config: true
+
+model:
+  model_name: facebook/dinov3-convnext-tiny-pretrain-lvd1689m
+  pooling_method: mean  # Options: "mean", "max", "cls_token"
+  middle_slice_index: 18  # Specific z-slice index (if null, uses D//2)
+  channel_reduction_methods:
+    Phase3D: middle_slice
+    RFP: max
+  channel_names:
+  - RFP
+  - Phase3D
+
+paths:
+  output_path: /hpc/mydata/eduardo.hirata/repos/viscy/applications/benchmarking/DynaCLR/DINOV3/embeddings_convnext_tiny_mean.zarr

applications/benchmarking/DynaCLR/DINOV3/dinov3_embeddings.py

@@ -0,0 +1,170 @@
+import sys
+from pathlib import Path
+from typing import Dict, List, Literal, Optional
+
+import numpy as np
+import torch
+from PIL import Image
+from skimage.exposure import rescale_intensity
+from transformers import AutoImageProcessor, AutoModel
+
+sys.path.append(str(Path(__file__).parent.parent))
+
+from base_embedding_module import BaseEmbeddingModule, create_embedding_cli
+
+
+class DINOv3Module(BaseEmbeddingModule):
+    def __init__(
+        self,
+        model_name: str = "facebook/dinov3-vitb16-pretrain-lvd1689m",
+        channel_reduction_methods: Optional[
+            Dict[str, Literal["middle_slice", "mean", "max"]]
+        ] = None,
+        channel_names: Optional[List[str]] = None,
+        pooling_method: Literal["mean", "max", "cls_token"] = "mean",
+        middle_slice_index: Optional[int] = None,
+    ):
+        super().__init__(channel_reduction_methods, channel_names, middle_slice_index)
+        self.model_name = model_name
+        self.pooling_method = pooling_method
+
+        self.model = None
+        self.processor = None
+
+    @classmethod
+    def from_config(cls, cfg):
+        """Create model instance from configuration."""
+        model_config = cfg.get("model", {})
+        return cls(
+            model_name=model_config.get(
+                "model_name", "facebook/dinov3-vitb16-pretrain-lvd1689m"
+            ),
+            pooling_method=model_config.get("pooling_method", "mean"),
+            channel_reduction_methods=model_config.get("channel_reduction_methods", {}),
+            channel_names=model_config.get("channel_names", []),
+            middle_slice_index=model_config.get("middle_slice_index", None),
+        )
+
+    def on_predict_start(self):
+        if self.model is None:
+            self.processor = AutoImageProcessor.from_pretrained(self.model_name)
+            self.model = AutoModel.from_pretrained(self.model_name)
+            self.model.eval()
+            self.model.to(self.device)
+
+    def _process_input(self, x: torch.Tensor):
+        """Convert tensor to PIL Images for DINOv3 processing."""
+        return self._convert_to_pil_images(x)
+
+    def _extract_features(self, pil_images):
+        """Extract features using DINOv3 model."""
+        inputs = self.processor(pil_images, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            token_features = outputs.last_hidden_state
+            features = self._pool_features(token_features)
+
+        return features
+
+    def _convert_to_pil_images(self, x: torch.Tensor) -> List[Image.Image]:
+        """
+        Convert tensor to list of PIL Images for DINOv3 processing.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor with shape (B, C, H, W).
+
+        Returns
+        -------
+        list of PIL.Image.Image
+            List of PIL Images ready for DINOv3 processing.
+        """
+        images = []
+
+        for b in range(x.shape[0]):
+            img_tensor = x[b]  # (C, H, W)
+
+            if img_tensor.shape[0] == 1:
+                # Single channel - convert to grayscale PIL
+                img_array = img_tensor[0].cpu().numpy()
+                # Normalize to 0-255
+                img_normalized = (
+                    (img_array - img_array.min())
+                    / (img_array.max() - img_array.min())
+                    * 255
+                ).astype(np.uint8)
+                pil_img = Image.fromarray(img_normalized, mode="L")
+
+            elif img_tensor.shape[0] == 2:
+                img_array = img_tensor.cpu().numpy()
+                rgb_array = np.zeros(
+                    (img_array.shape[1], img_array.shape[2], 3), dtype=np.uint8
+                )
+
+                ch0_norm = rescale_intensity(img_array[0], out_range=(0, 255)).astype(
+                    np.uint8
+                )
+                ch1_norm = rescale_intensity(img_array[1], out_range=(0, 255)).astype(
+                    np.uint8
+                )
+
+                rgb_array[:, :, 0] = ch0_norm  # Red
+                rgb_array[:, :, 1] = ch1_norm  # Green
+                rgb_array[:, :, 2] = (ch0_norm + ch1_norm) // 2  # Blue
+
+                pil_img = Image.fromarray(rgb_array, mode="RGB")
+
+            elif img_tensor.shape[0] == 3:
+                # Three channels - direct RGB
+                img_array = img_tensor.cpu().numpy().transpose(1, 2, 0)  # HWC
+                img_normalized = rescale_intensity(
+                    img_array, out_range=(0, 255)
+                ).astype(np.uint8)
+                pil_img = Image.fromarray(img_normalized, mode="RGB")
+
+            else:
+                # More than 3 channels - use first 3
+                img_array = img_tensor[:3].cpu().numpy().transpose(1, 2, 0)  # HWC
+                img_normalized = rescale_intensity(
+                    img_array, out_range=(0, 255)
+                ).astype(np.uint8)
+                pil_img = Image.fromarray(img_normalized, mode="RGB")
+
+            images.append(pil_img)
+
+        return images
+
+    def _pool_features(self, features: torch.Tensor) -> torch.Tensor:
+        """
+        Pool spatial features from DINOv3 tokens.
+
+        Parameters
+        ----------
+        features : torch.Tensor
+            Token features with shape (B, num_tokens, hidden_dim).
+
+        Returns
+        -------
+        torch.Tensor
+            Pooled features with shape (B, hidden_dim).
+        """
+        if self.pooling_method == "cls_token":
+            # For ViT models, first token is usually CLS token
+            if "vit" in self.model_name.lower():
+                return features[:, 0, :]  # CLS token
+            else:
+                # For ConvNeXt, no CLS token, fall back to mean
+                return features.mean(dim=1)
+
+        elif self.pooling_method == "max":
+            return features.max(dim=1)[0]
+        else:  # mean pooling
+            return features.mean(dim=1)
+
+
+if __name__ == "__main__":
+    main = create_embedding_cli(DINOv3Module, "DINOv3")
+    main()

applications/benchmarking/DynaCLR/OpenPhenom/config_template.yml

@@ -2,8 +2,6 @@
 
 # Paths section
 paths:
-  data_path: /hpc/projects/intracellular_dashboard/organelle_dynamics/2024_11_07_A549_SEC61_ZIKV_DENV/2-assemble/2024_11_07_A549_SEC61_DENV.zarr
-  tracks_path: /hpc/projects/intracellular_dashboard/organelle_dynamics/2024_11_07_A549_SEC61_ZIKV_DENV/1-preprocess/label-free/4-track-gt/2024_11_07_A549_SEC61_ZIKV_DENV_2_cropped.zarr
   output_path: "/home/eduardo.hirata/repos/viscy/applications/benchmarking/DynaCLR/OpenPhenom/openphenom_sec61b_n_phase_3.zarr"
 
 # Model configuration
@@ -16,7 +14,10 @@ model:
     "raw GFP EX488 EM525-45": "max"
 
 # Data module configuration
+datamodule_class: viscy.data.triplet.TripletDataModule
 datamodule:
+  data_path: /hpc/projects/intracellular_dashboard/organelle_dynamics/2024_11_07_A549_SEC61_ZIKV_DENV/2-assemble/2024_11_07_A549_SEC61_DENV.zarr
+  tracks_path: /hpc/projects/intracellular_dashboard/organelle_dynamics/2024_11_07_A549_SEC61_ZIKV_DENV/1-preprocess/label-free/4-track-gt/2024_11_07_A549_SEC61_ZIKV_DENV_2_cropped.zarr
   source_channel:
     - Phase3D
     - "raw GFP EX488 EM525-45"