O2vae integration rebase

bioimage_embed/augmentations.py

@@ -32,6 +32,39 @@
 DEFAULT_AUGMENTATION = A.Compose(DEFAULT_AUGMENTATION_LIST)
 DEFAULT_ALBUMENTATION = A.Compose(DEFAULT_AUGMENTATION_LIST)
 
+DEFAULT_AUGMENTATION_LIST = [
+    # Flip the images horizontally or vertically with a 50% chance
+    A.OneOf(
+        [
+            A.HorizontalFlip(p=0.5),
+            A.VerticalFlip(p=0.5),
+        ],
+        p=0.5,
+    ),
+    # Rotate the images by a random angle within a specified range
+    A.Rotate(limit=45, p=0.5),
+    # Randomly scale the image intensity to adjust brightness and contrast
+    A.RandomGamma(gamma_limit=(80, 120), p=0.5),
+    # Apply random elastic transformations to the images
+    A.ElasticTransform(
+        alpha=1,
+        sigma=50,
+        alpha_affine=50,
+        p=0.5,
+    ),
+    # Shift the image channels along the intensity axis
+    A.ChannelShuffle(p=0.5),
+    # Add a small amount of noise to the images
+    A.GaussNoise(var_limit=(10.0, 50.0), p=0.5),
+    # Crop a random part of the image and resize it back to the original size
+    A.RandomResizedCrop(
+        height=512, width=512, scale=(0.9, 1.0), ratio=(0.9, 1.1), p=0.5
+    ),
+    # Adjust image intensity with a specified range for individual channels
+    A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
+]
+
+DEFAULT_AUGMENTATION = A.Compose(DEFAULT_AUGMENTATION_LIST)
 
 class VisionWrapper:
     def __init__(self, transform_dict, *args, **kwargs):

bioimage_embed/lightning/torch.py

@@ -58,6 +58,8 @@ def __init__(self, model, args=SimpleNamespace()):
         # TODO update all models to use this for export to onxx
         # self.example_input_array = torch.randn(1, *self.model.input_dim)
         # self.model.train()
+        # keep a handle on metrics logged by the model
+        self.metrics = {}
 
     def forward(self, x: torch.Tensor) -> ModelOutput:
         """

bioimage_embed/models/factory.py

@@ -18,7 +18,6 @@
 from . import bolts
 from functools import partial
 
-
 class ModelFactory:
     def __init__(
         self, input_dim, latent_dim, pretrained=False, progress=True, **kwargs
@@ -144,6 +143,19 @@ def resnet18_beta_vae(self):
             bolts.ResNet18VAEDecoder,
         )
 
+    def resnet50_vae(self):
+        return self.create_model(
+            partial(
+                pythae.models.VAEConfig,
+                use_default_encoder=False,
+                use_default_decoder=False,
+                **self.kwargs
+            ),
+            pythae.models.VAE,
+            bolts.ResNet50VAEEncoder,
+            bolts.ResNet50VAEDecoder,
+        )
+
     def resnet50_vqvae(self):
         return self.create_model(
             partial(

bioimage_embed/models/o2vae_shapeembed_integration.diff

@@ -0,0 +1,97 @@
+diff --git a/models/align_reconstructions.py b/models/align_reconstructions.py
+index d07d1ab..c52b40d 100644
+--- a/models/align_reconstructions.py
++++ b/models/align_reconstructions.py
+@@ -6,7 +6,7 @@ import torch
+ import torchgeometry as tgm
+ import torchvision.transforms.functional as T_f
+
+-from registration import registration
++from ..registration import registration
+
+
+ def loss_reconstruction_fourier_batch(x, y, recon_loss_type="bce", mask=None):
+diff --git a/models/decoders/cnn_decoder.py b/models/decoders/cnn_decoder.py
+index ba3a1cc..1740945 100644
+--- a/models/decoders/cnn_decoder.py
++++ b/models/decoders/cnn_decoder.py
+@@ -58,7 +58,7 @@ class CnnDecoder(nn.Module):
+
+         self.dec_conv = nn.Sequential(*layers)
+
+-    def forward(self, x):
++    def forward(self, x, epoch = None):
+         bs = x.size(0)
+         x = self.fc(x)
+         dim = x.size(1)
+diff --git a/models/encoders_o2/e2scnn.py b/models/encoders_o2/e2scnn.py
+index 9c4f47f..e292b1e 100644
+--- a/models/encoders_o2/e2scnn.py
++++ b/models/encoders_o2/e2scnn.py
+@@ -219,14 +219,20 @@ class E2SFCNN(torch.nn.Module):
+                 repr += f"\t{i: <3} - {name: <70} | {params: <8} |\n"
+         return repr
+
+-    def forward(self, input: torch.tensor):
++    def forward(self, input: torch.tensor, epoch = None):
++        #print(f"DEBUG: e2scnn forward: input.shape: {input.shape}")
+         x = GeometricTensor(input, self.in_repr)
++        #print(f"DEBUG: e2scnn forward: pre layers x.shape: {x.shape}")
+
+         for layer in self.eq_layers:
+             x = layer(x)
+
++        #print(f"DEBUG: e2scnn forward: pre fully_net x.shape: {x.shape}")
++
+         x = self.fully_net(x.tensor.reshape(x.tensor.shape[0], -1))
+
++        #print(f"DEBUG: e2scnn forward: pre final x.shape: {x.shape}")
++
+         return x
+
+     def build_layer_regular(
+diff --git a/models/vae.py b/models/vae.py
+index 3af262b..af1a2dc 100644
+--- a/models/vae.py
++++ b/models/vae.py
+@@ -3,8 +3,9 @@ import importlib
+ import numpy as np
+ import torch
+ import torchvision
++from pythae.models.base.base_utils import ModelOutput
+
+-from models import align_reconstructions
++from . import align_reconstructions
+
+ from . import model_utils as mut
+
+@@ -273,10 +274,11 @@ class VAE(torch.nn.Module):
+
+         return y
+
+-    def forward(self, x):
++    def forward(self, x, epoch = None):
++        x = x["data"]
+         in_shape = x.shape
+         bs = in_shape[0]
+-        assert x.ndim == 4
++        assert len(in_shape) == 4
+
+         # inference and sample
+         z = self.q_net(x)
+@@ -290,8 +292,12 @@ class VAE(torch.nn.Module):
+             y = torch.sigmoid(y)
+         # check the spatial dimensions are good (if doing multiclass prediction per pixel, the `c` dim may be different)
+         assert in_shape[-2:] == y.shape[-2:], (
+-            "output image different dimension to "
+-            "input image ... probably change the number of layers (cnn_dims) in the decoder"
++            f"output image different dimension {y.shape[-2:]} to "
++            f"input image {in_shape[-2:]} ... probably change the number of layers (cnn_dims) in the decoder"
+         )
+
+-        return x, y, mu, logvar
++        # gather losses
++        losses = self.loss(x, y, mu, logvar)
++
++        return ModelOutput(recon_x=y, z=z_sample, loss=losses['loss'], recon_loss=losses['loss_recon'])
++        #return ModelOutput(recon_x=y, z=z_sample)

bioimage_embed/shapes/lightning.py

@@ -41,8 +41,8 @@ def eval_step(self, batch, batch_idx):
                 [
                     loss_ops.diagonal_loss(),
                     loss_ops.symmetry_loss(),
-                    # loss_ops.triangle_inequality(),
                     loss_ops.non_negative_loss(),
+                    # loss_ops.triangle_inequality(),
                     # loss_ops.clockwise_order_loss(),
                 ]
             )

bioimage_embed/shapes/mds.py

@@ -9,11 +9,12 @@ def mds(d):
     :return: A matrix of x, y coordinates.
     """
     n = d.size(0)
-    I = torch.eye(n)
+    I = torch.eye(n, dtype=torch.float64)
     H = I - torch.ones((n, n)) / n
 
     S = -0.5 * H @ d @ H
-    eigvals, eigvecs = S.symeig(eigenvectors=True)
+    #eigvals, eigvecs = S.symeig(eigenvectors=True)
+    eigvals, eigvecs = torch.linalg.eigh(S)
 
     # Sort the eigenvalues and eigenvectors in decreasing order
     idx = eigvals.argsort(descending=True)

scripts/shapeembed/__init__.py

@@ -0,0 +1,2 @@
+from .dataset_transformations import mask2distmatrix
+from .evaluation import *

scripts/shapeembed/common_helpers.py

@@ -0,0 +1,42 @@
+import re
+import os
+import glob
+import types
+import logging
+
+def compressed_n_features(dist_mat_size, comp_fact):
+  return dist_mat_size*(dist_mat_size-1)//(2**comp_fact)
+
+def model_str(params):
+  s = f'{params.model_name}'
+  if hasattr(params, 'model_args'):
+    s += f"-{'_'.join([f'{k}{v}' for k, v in vars(params.model_args).items()])}"
+  return s
+
+def job_str(params):
+  return f"{params.dataset.name}-{model_str(params)}-{params.compression_factor}-{params.latent_dim}-{params.batch_size}"
+
+def job_str_re():
+  return re.compile("(.*)-(.*)-(\d+)-(\d+)-(\d+)")
+
+def params_from_job_str(jobstr):
+  raw = jobstr.split('-')
+  ps = types.SimpleNamespace()
+  ps.batch_size = int(raw.pop())
+  ps.latent_dim = int(raw.pop())
+  ps.compression_factor = int(raw.pop())
+  if len(raw) == 3:
+    ps.model_args = types.SimpleNamespace()
+    for p in raw.pop().split('-'):
+      if p[0:4] == 'beta': ps.model_args.beta = float(p[4:])
+  ps.model_name = raw.pop()
+  ps.dataset = types.SimpleNamespace(name=raw.pop())
+  return ps
+
+def find_existing_run_scores(dirname, logger=logging.getLogger(__name__)):
+  ps = []
+  for f in glob.glob(f'{dirname}/*-shapeembed-score_df.csv'):
+    p = params_from_job_str(os.path.basename(f)[:-24])
+    p.csv_file = f
+    ps.append(p)
+  return ps