Refactor

.github/workflows/test.yaml

@@ -33,6 +33,8 @@ jobs:
         run: poetry --help
       - name: Install poetry
         uses: abatilo/actions-poetry@v2
+      - name: Dynamic versioning
+        uses: poetry self add poetry-dynamic-versioning
       - name: Setup a local virtual environment (if no poetry.toml file)
         run: |
           poetry config virtualenvs.create true --local

bioimage_embed/__init__.py

@@ -1,8 +1,12 @@
-import torch
-
-torch.cuda.empty_cache()
 # from . import models, lightning, cli, export, config
-from .lightning import AESupervised, AEUnsupervised, AE,  AutoEncoderSupervised, AutoEncoderUnsupervised, AutoEncoder
+from .lightning import (
+    AESupervised,
+    AEUnsupervised,
+    AE,
+    AutoEncoderSupervised,
+    AutoEncoderUnsupervised,
+    AutoEncoder,
+)
 
 # TODO: Fix this import as it currently produces too many warnings
 from .models import ModelFactory, create_model
@@ -13,14 +17,19 @@
 import logging
 logging.captureWarnings(True)
 
+import torch
+
+torch.cuda.empty_cache()
 __all__ = [
     "AESupervised",
     "AutoEncoderUnsupervised",
     "AEUnsupervised",
     "AutoEncoderSupervised",
-    "AutoEncoder"
-    "AE"
+    "AutoEncoder",
+    "AE",
     "BioImageEmbed",
     "Config",
     "augmentations",
+    "ModelFactory",
+    "create_model",
 ]

bioimage_embed/lightning/__init__.py

@@ -1,5 +1,21 @@
 from .pyro import LitAutoEncoderPyro
-from .torch import AESupervised, AEUnsupervised, AutoEncoder, AE, AutoEncoderSupervised, AutoEncoderUnsupervised
+from .torch import (
+    AESupervised,
+    AEUnsupervised,
+    AutoEncoder,
+    AE,
+    AutoEncoderSupervised,
+    AutoEncoderUnsupervised,
+)
 from .dataloader import DataModule
 
-__all__ = ["LitAutoEncoderPyro", "AESupervised", "AEUnsupervised", "DataModule", "AutoEncoder","AE","AutoEncoderUnsupervised","AutoEncoderSupervised"]
+__all__ = [
+    "LitAutoEncoderPyro",
+    "AESupervised",
+    "AEUnsupervised",
+    "DataModule",
+    "AutoEncoder",
+    "AE",
+    "AutoEncoderUnsupervised",
+    "AutoEncoderSupervised",
+]

bioimage_embed/lightning/pyro.py

@@ -7,6 +7,10 @@
 
 
 class LitAutoEncoderPyro(pl.LightningModule):
+    """
+    WIP Unsupported
+    """
+
     def __init__(self, model, batch_size=1, learning_rate=1e-3):
         super().__init__()
         # self.autoencoder = AutoEncoder(batch_size, 1)
@@ -59,6 +63,3 @@ def pyro_training_step(self, train_batch, batch_idx):
 
     def training_step(self, train_batch, batch_idx):
         return self.torch_training_step(train_batch, batch_idx)
-
-    def training_step(self, train_batch, batch_idx):
-        return self.pyro_training_step(train_batch, batch_idx)

bioimage_embed/lightning/tests/test_ndims.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+
 # Fixture for batch sizes
 @pytest.fixture(params=[1, 16])
 def batch_size(request):

bioimage_embed/lightning/torch.py

@@ -143,16 +143,6 @@ def eval_step(self, batch, batch_idx):
         """
         return self.predict_step(batch, batch_idx)
 
-    # def lr_scheduler_step(self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None):
-    #     # Implement your own logic for updating the lr scheduler
-    #     # This method will be called at each training step
-    #     # Update the lr scheduler based on the provided arguments
-    #     # You can access the lr scheduler using `self.lr_schedulers()`
-
-    #     # Example:
-    #     for lr_scheduler in self.lr_schedulers():
-    #         lr_scheduler.step()
-
     def timm_optimizers(self, model):
         optimizer = optim.create_optimizer(self.args, model.parameters())
         lr_scheduler = scheduler.create_scheduler(self.args, optimizer)[0]
@@ -168,8 +158,6 @@ def timm_to_lightning(self, optimizer, lr_scheduler):
         }
 
     def configure_optimizers(self):
-        # optimizer = optim.create_optimizer(self.args, self.model.parameters())
-        # lr_scheduler = scheduler.create_scheduler(self.args, optimizer)[0]
         optimizer, lr_scheduler = self.timm_optimizers(self.model)
         return self.timm_to_lightning(optimizer, lr_scheduler)
 

bioimage_embed/models/__init__.py

@@ -1,18 +1,19 @@
-import torch
-import torch.nn.functional as F
+# Description: This file is the main entry point for the models module. It imports all the necessary modules and classes for the models module to function properly.
 
 # Note - you must have torchvision installed for this example
-from torch.utils.data import DataLoader
-
-# from .ae import AutoEncoder
-
-# from .vae_bio import Mask_VAE, Image_VAE
-# from .utils import BaseVAE
-# from .legacy.vae import VAE
-# from .vq_vae import VQ_VAE
 
 from .bolts import ResNet18VAEEncoder, ResNet18VAEDecoder
 
 from . import bolts
 from . import pythae
-from .factory import ModelFactory, create_model, __all_models__
+from .factory import ModelFactory, create_model, __all_models__
+
+__all__ = [
+    "ModelFactory",
+    "create_model",
+    "__all_models__",
+    "ResNet18VAEEncoder",
+    "ResNet18VAEDecoder",
+    "bolts",
+    "pythae",
+]

bioimage_embed/models/tests/test_factory.py

@@ -14,7 +14,9 @@
 latent_dim = [64, 16]
 pretrained_options = [True, False]
 progress_options = [True, False]
-batch = [1,]
+batch = [
+    1,
+]
 
 
 @pytest.mark.parametrize("model", __all_models__)

bioimage_embed/tests/test_bioimage_embed.py

@@ -2,13 +2,14 @@
 import torch
 from ..bie import BioImageEmbed
 
+
 @pytest.fixture()
 def test_bioimage_embed():
     bie = BioImageEmbed()
     bie.train()
     bie.infer()
     bie.validate()
-    model_output = bie(torch.tensor([1, 2, 3, 4, 5]))
-    tensor = bie.model(torch.tensor([1, 2, 3, 4, 5]))
+    assert bie(torch.tensor([1, 2, 3, 4, 5]))
+    assert bie.model(torch.tensor([1, 2, 3, 4, 5]))
 
     bie.model(torch.tensor([1, 2, 3, 4, 5]))