added resnet50_beta_vae to the factory

uhlmanngroup · Oct 1, 2024 · 0235109 · 0235109
1 parent 931d485
commit 0235109
Showing 1 changed file with 27 additions and 0 deletions.
diff --git a/bioimage_embed/models/factory.py b/bioimage_embed/models/factory.py
@@ -144,6 +144,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(
@@ -157,6 +170,19 @@ def resnet50_vqvae(self):
             bolts.ResNet50VQVAEDecoder,
         )
 
+    def resnet50_beta_vae(self):
+        return self.create_model(
+            partial(
+                pythae.models.BetaVAEConfig,
+                use_default_encoder=False,
+                use_default_decoder=False,
+                **self.kwargs
+            ),
+            pythae.models.BetaVAE,
+            bolts.ResNet50VAEEncoder,
+            bolts.ResNet50VAEDecoder,
+        )
+
     def resnet_vae_legacy(self, depth):
         return self.create_model(
             pythae.models.VAEConfig,
@@ -210,6 +236,7 @@ def __call__(self, model):
     "resnet18_vae",
     "resnet18_beta_vae",
     "resnet50_vae",
+    "resnet50_beta_vae",
     "resnet18_vae_bolt",
     "resnet50_vae_bolt",
     "resnet18_vqvae",