feat: added mae3d (#652)

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>

Author: anna-grim

src/neuron_proofreader/machine_learning/vision_models.py

@@ -5,10 +5,11 @@
 @email: anna.grim@alleninstitute.org
 
 Code for vision models that perform image classification tasks within
-NeuronProofreading pipelines.
+NeuronProofreader pipelines.
 
 """
 
+#from neurobase.finetune import finetune_model
 from einops import rearrange
 
 import torch
@@ -130,6 +131,36 @@ def forward(self, x):
 
 
 # --- Transformers ---
+class MAE3D(nn.Module):
+
+    def __init__(self, checkpoint_path, model_config):
+        # Call parent closs
+        super().__init__()
+
+        # Load model
+        full_model = finetune_model(
+            checkpoint_path=checkpoint_path,
+            model_config=model_config,
+            task_head_config="binary_classifier",
+            freeze_encoder=True
+        )
+
+        # Instance attributes
+        self.encoder = full_model.encoder
+        self.output = ml_util.init_feedforward(384, 1, 2)
+
+    def forward(self, x):
+        latent0 = self.encoder(x[:, 0:1, ...])
+        latent1 = self.encoder(x[:, 1:2, ...])
+
+        x0 = latent0["latents"][:, 0, :]
+        x1 = latent1["latents"][:, 0, :]
+
+        x = torch.cat((x0, x1), dim=1)
+        x = self.output(x)
+        return x
+
+
 class ViT3D(nn.Module):
     """
     A class that implements a 3D Vision transformer.