Skip to content

Commit

Permalink
feat: ✨ Add models used in CellMap Segmentation Challenge.
Browse files Browse the repository at this point in the history
  • Loading branch information
rhoadesScholar committed Aug 12, 2024
1 parent 9f29e90 commit 125894e
Show file tree
Hide file tree
Showing 8 changed files with 1,177 additions and 2 deletions.
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,8 @@ dependencies = [
'torchvision',
'numpy',
'tqdm',
'cellpose'
'cellpose',
'ml-collections',
]

[project.optional-dependencies]
Expand Down
2 changes: 1 addition & 1 deletion src/cellmap_models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@
"""

from .utils import download_url_to_file
from .pytorch import cosem, cellpose
from .pytorch import cosem, cellpose, untrained_models
22 changes: 22 additions & 0 deletions src/cellmap_models/pytorch/untrained_models/README.md
Original file line number Diff line number Diff line change
@@ -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)
```
4 changes: 4 additions & 0 deletions src/cellmap_models/pytorch/untrained_models/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
from .resnet import ResNet
from .vitnet import ViTVNet
from .unet_2D import UNet2D
from .unet_3D import UNet3D
Loading

0 comments on commit 125894e

Please sign in to comment.