feat: ✨ Add models used in CellMap Segmentation Challenge.

pyproject.toml

@@ -21,7 +21,8 @@ dependencies = [
   'torchvision',
   'numpy',
   'tqdm',
-  'cellpose'
+  'cellpose',
+  'ml-collections',
 ]
 
 [project.optional-dependencies]

src/cellmap_models/__init__.py

@@ -3,4 +3,4 @@
 """
 
 from .utils import download_url_to_file
-from .pytorch import cosem, cellpose
+from .pytorch import cosem, cellpose, untrained_models

src/cellmap_models/pytorch/untrained_models/README.md

@@ -0,0 +1,22 @@
+<img src="https://raw.githubusercontent.com/janelia-cellmap/cellmap-models/main/assets/CellMapLogo2.png" alt="CellMap logo" width="85%">
+
+## This directory contains various untrained PyTorch model architectures.
+
+## Models
+
+***ResNet***: Parameterizable 2D and 3D ResNet models with a variable number of layers and channels. This model is based on the original ResNet architecture with the addition of a decoding path, which mirrors the encoder, after the bottleneck, to produce an image output.
+
+***UNet2D***: A simple 2D UNet model with a variable number of output channels.
+
+***UNet3D***: A simple 3D UNet model with a variable number of output channels.
+
+***ViTVNet***: A 3D VNet model with a Vision Transformer (ViT) encoder. This model is based on the original VNet architecture with the addition of a ViT encoder in place of the original convolutional encoder.
+
+## Usage
+
+To use these models, you can import them directly from the `cellmap_models.pytorch.untrained_models` module. For example, to import the ResNet model, you can use the following code:
+
+```python
+from cellmap_models.pytorch.untrained_models import ResNet
+model = ResNet(ndim=2, input_nc=1, output_nc=3, n_blocks=18)
+```

src/cellmap_models/pytorch/untrained_models/__init__.py

@@ -0,0 +1,4 @@
+from .resnet import ResNet
+from .vitnet import ViTVNet
+from .unet_2D import UNet2D
+from .unet_3D import UNet3D