Repository: fepegar/unet Branch: main Commit: d5f46e0b1e65 Files: 17 Total size: 29.3 KB Directory structure: gitextract_as_2qdiu/ ├── .bumpversion.toml ├── .github/ │ └── ISSUE_TEMPLATE.md ├── .gitignore ├── .zenodo.json ├── CONTRIBUTING.rst ├── LICENSE ├── README.rst ├── justfile ├── pyproject.toml ├── tests/ │ ├── __init__.py │ └── test_unet.py ├── tox.ini └── unet/ ├── __init__.py ├── conv.py ├── decoding.py ├── encoding.py └── unet.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .bumpversion.toml ================================================ [tool.bumpversion] current_version = "0.8.1" tag = true commit = true message = "Bump version: {current_version} → {new_version}" [[tool.bumpversion.files]] filename = "pyproject.toml" search = 'version = "{current_version}"' replace = 'version = "{new_version}"' ================================================ FILE: .github/ISSUE_TEMPLATE.md ================================================ * unet version: * Python version: * Operating System: ### Description Describe what you were trying to get done. Tell us what happened, what went wrong, and what you expected to happen. ### What I Did ``` Paste the command(s) you ran and the output. If there was a crash, please include the traceback here. ``` ================================================ FILE: .gitignore ================================================ # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover .hypothesis/ .pytest_cache/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder target/ # Jupyter Notebook .ipynb_checkpoints # pyenv .python-version # celery beat schedule file celerybeat-schedule # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ # PyCharm .idea # VS Code .vscode # uv uv.lock # ruff .ruff_cache ================================================ FILE: .zenodo.json ================================================ { "creators": [ { "affiliation": "University College London, United Kingdom", "name": "P\u00e9rez-Garc\u00eda, Fernando", "orcid": "0000-0001-9090-3024" } ], "description": "PyTorch implementation of 2D and 3D U-Net", "keywords": [ "pytorch", "medical-image-computing", "deep-learning", "machine-learning", "convolutional-neural-networks" ], "license": "mit-license", "upload_type": "software" } ================================================ FILE: CONTRIBUTING.rst ================================================ .. highlight:: shell ============ Contributing ============ Contributions are welcome, and they are greatly appreciated! Every little bit helps, and credit will always be given. You can contribute in many ways: Types of Contributions ---------------------- Report Bugs ~~~~~~~~~~~ Report bugs at https://github.com/fepegar/unet/issues. If you are reporting a bug, please include: * Your operating system name and version. * Any details about your local setup that might be helpful in troubleshooting. * Detailed steps to reproduce the bug. Fix Bugs ~~~~~~~~ Look through the GitHub issues for bugs. Anything tagged with "bug" and "help wanted" is open to whoever wants to implement it. Implement Features ~~~~~~~~~~~~~~~~~~ Look through the GitHub issues for features. Anything tagged with "enhancement" and "help wanted" is open to whoever wants to implement it. Write Documentation ~~~~~~~~~~~~~~~~~~~ unet could always use more documentation, whether as part of the official unet docs, in docstrings, or even on the web in blog posts, articles, and such. Submit Feedback ~~~~~~~~~~~~~~~ The best way to send feedback is to file an issue at https://github.com/fepegar/unet/issues. If you are proposing a feature: * Explain in detail how it would work. * Keep the scope as narrow as possible, to make it easier to implement. * Remember that this is a volunteer-driven project, and that contributions are welcome :) Get Started! ------------ Ready to contribute? Here's how to set up `unet` for local development. 1. Fork the `unet` repo on GitHub. 2. Clone your fork locally:: $ git clone git@github.com:your_name_here/unet.git 3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: $ mkvirtualenv unet $ cd unet/ $ python setup.py develop 4. Create a branch for local development:: $ git checkout -b name-of-your-bugfix-or-feature Now you can make your changes locally. 5. When you're done making changes, check that your changes pass flake8 and the tests, including testing other Python versions with tox:: $ flake8 unet tests $ python setup.py test or pytest $ tox To get flake8 and tox, just pip install them into your virtualenv. 6. Commit your changes and push your branch to GitHub:: $ git add . $ git commit -m "Your detailed description of your changes." $ git push origin name-of-your-bugfix-or-feature 7. Submit a pull request through the GitHub website. Pull Request Guidelines ----------------------- Before you submit a pull request, check that it meets these guidelines: 1. The pull request should include tests. 2. If the pull request adds functionality, the docs should be updated. Put your new functionality into a function with a docstring, and add the feature to the list in README.rst. 3. The pull request should work for Python 3.6 and 3.7, and for PyPy. Check https://travis-ci.org/fepegar/unet/pull_requests and make sure that the tests pass for all supported Python versions. Tips ---- To run a subset of tests:: $ python -m unittest tests.test_unet Deploying --------- A reminder for the maintainers on how to deploy. Make sure all your changes are committed (including an entry in HISTORY.rst). Then run:: $ bump2version patch # possible: major / minor / patch $ git push $ git push --tags Travis will then deploy to PyPI if tests pass. ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2019 Fernando Perez-Garcia Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: README.rst ================================================ U-Net ===== .. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.3522306.svg :target: https://doi.org/10.5281/zenodo.3522306 :alt: DOI .. image:: https://img.shields.io/badge/License-MIT-yellow.svg :target: https://opensource.org/licenses/MIT :alt: License .. image:: https://img.shields.io/pypi/v/unet.svg :target: https://pypi.python.org/pypi/unet PyTorch implementation of 1D, 2D and 3D U-Net. The U-Net architecture was first described in `Ronneberger et al. 2015, U-Net: Convolutional Networks for Biomedical Image Segmentation `_. The 3D version was described in `Çiçek et al. 2016, 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation `_. Installation ------------ :: pip install unet Credits ------- If you used this code for your research, please cite this repository using the information available on its `Zenodo entry `_: Pérez-García, Fernando. (2020). fepegar/unet: PyTorch implementation of 2D and 3D U-Net (v0.7.5). Zenodo. https://doi.org/10.5281/zenodo.3697931 ================================================ FILE: justfile ================================================ @install_uv: if ! command -v uv >/dev/null 2>&1; then \ echo "uv is not installed. Installing..."; \ curl -LsSf https://astral.sh/uv/install.sh | sh; \ fi setup: install_uv uv sync --all-extras --all-groups bump part='patch': install_uv uv run bump-my-version bump {{part}} --verbose release: install_uv rm -rf dist uv build --no-sources uv publish changelog: install_uv uvx git-changelog --output CHANGELOG.md ruff: install_uv uvx ruff check --fix uvx ruff format test: install_uv uv run tox -p ================================================ FILE: pyproject.toml ================================================ [build-system] requires = ["hatchling"] build-backend = "hatchling.build" [project] name = "unet" version = "0.8.1" description = "PyTorch implementation of 1D, 2D and 3D U-Net." authors = [{ name = "Fernando Perez-Garcia", email = "fepegar@gmail.com" }] readme = { file = "README.rst", content-type = "text/x-rst" } requires-python = ">=3.9" classifiers = [ "Intended Audience :: Science/Research", "License :: OSI Approved :: MIT License", "Natural Language :: English", "Operating System :: OS Independent", ] dependencies = ["torch"] [project.urls] Homepage = "https://github.com/fepegar/unet" Source = "https://github.com/fepegar/unet" "Issue tracker" = "https://github.com/fepegar/unet/issues" [dependency-groups] dev = ["bump-my-version", "pytest", "pytest-sugar", "ruff", "tox-uv"] select = [ # pycodestyle "E", # Pyflakes "F", # pyupgrade "UP", # flake8-bugbear "B", # isort "I", ] ================================================ FILE: tests/__init__.py ================================================ # -*- coding: utf-8 -*- """Unit test package for unet.""" ================================================ FILE: tests/test_unet.py ================================================ import torch from unet import UNet1D, UNet2D, UNet3D residual = False torch.manual_seed(0) torch.set_grad_enabled(False) def run(model, shape): x_sample = torch.rand(*shape) with torch.no_grad(): y = model(x_sample) return y def test_unet_1d(): model = UNet1D( normalization="batch", preactivation=True, residual=False, ).eval() shape = 1, 1, 572 result = 1, 2, 388 y = run(model, shape) assert tuple(y.shape) == result def test_unet_1d_residual(): model = UNet1D( normalization="batch", preactivation=True, residual=True, ).eval() shape = 1, 1, 512 result = 1, 2, 512 y = run(model, shape) assert tuple(y.shape) == result def test_unet_2d(): model = UNet2D( normalization="batch", preactivation=True, residual=False, ).eval() shape = 1, 1, 572, 572 result = 1, 2, 388, 388 y = run(model, shape) assert tuple(y.shape) == result def test_unet_2d_residual(): model = UNet2D( normalization="batch", preactivation=True, residual=True, ).eval() shape = 1, 1, 512, 512 result = 1, 2, 512, 512 y = run(model, shape) assert tuple(y.shape) == result def test_unet_3d(): model = UNet3D( normalization="batch", preactivation=True, residual=False, ).eval() shape = 1, 1, 132, 132, 116 result = 1, 2, 44, 44, 28 y = run(model, shape) assert tuple(y.shape) == result def test_unet_3d_residual(): model = UNet3D( normalization="batch", preactivation=True, residual=True, num_encoding_blocks=2, upsampling_type="trilinear", ).eval() shape = 1, 1, 64, 64, 56 result = 1, 2, 64, 64, 56 y = run(model, shape) assert tuple(y.shape) == result ================================================ FILE: tox.ini ================================================ [tox] envlist = py3{9,10,11,12} [testenv] commands = pytest ================================================ FILE: unet/__init__.py ================================================ __version__ = "0.7.9" from .unet import UNet, UNet1D, UNet2D, UNet3D __all__ = [ "UNet", "UNet1D", "UNet2D", "UNet3D", ] ================================================ FILE: unet/conv.py ================================================ from typing import Optional import torch.nn as nn class ConvolutionalBlock(nn.Module): def __init__( self, dimensions: int, in_channels: int, out_channels: int, normalization: Optional[str] = None, kernel_size: int = 3, activation: Optional[str] = "ReLU", preactivation: bool = False, padding: int = 0, padding_mode: str = "zeros", dilation: Optional[int] = None, dropout: float = 0, ): super().__init__() block = nn.ModuleList() dilation = 1 if dilation is None else dilation if padding: total_padding = kernel_size + 2 * (dilation - 1) - 1 padding = total_padding // 2 class_name = "Conv{}d".format(dimensions) conv_class = getattr(nn, class_name) no_bias = not preactivation and (normalization is not None) conv_layer = conv_class( in_channels, out_channels, kernel_size, padding=padding, padding_mode=padding_mode, dilation=dilation, bias=not no_bias, ) norm_layer = None if normalization is not None: class_name = "{}Norm{}d".format(normalization.capitalize(), dimensions) norm_class = getattr(nn, class_name) num_features = in_channels if preactivation else out_channels norm_layer = norm_class(num_features) activation_layer = None if activation is not None: activation_layer = getattr(nn, activation)() if preactivation: self.add_if_not_none(block, norm_layer) self.add_if_not_none(block, activation_layer) self.add_if_not_none(block, conv_layer) else: self.add_if_not_none(block, conv_layer) self.add_if_not_none(block, norm_layer) self.add_if_not_none(block, activation_layer) dropout_layer = None if dropout: class_name = "Dropout{}d".format(dimensions) dropout_class = getattr(nn, class_name) dropout_layer = dropout_class(p=dropout) self.add_if_not_none(block, dropout_layer) self.conv_layer = conv_layer self.norm_layer = norm_layer self.activation_layer = activation_layer self.dropout_layer = dropout_layer self.block = nn.Sequential(*block) def forward(self, x): return self.block(x) @staticmethod def add_if_not_none(module_list, module): if module is not None: module_list.append(module) ================================================ FILE: unet/decoding.py ================================================ from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from .conv import ConvolutionalBlock CHANNELS_DIMENSION = 1 UPSAMPLING_MODES = ( "nearest", "linear", "bilinear", "bicubic", "trilinear", ) class Decoder(nn.Module): def __init__( self, in_channels_skip_connection: int, dimensions: int, upsampling_type: str, num_decoding_blocks: int, normalization: Optional[str], preactivation: bool = False, residual: bool = False, padding: int = 0, padding_mode: str = "zeros", activation: Optional[str] = "ReLU", initial_dilation: Optional[int] = None, dropout: float = 0, ): super().__init__() upsampling_type = fix_upsampling_type(upsampling_type, dimensions) self.decoding_blocks = nn.ModuleList() self.dilation = initial_dilation for _ in range(num_decoding_blocks): decoding_block = DecodingBlock( in_channels_skip_connection, dimensions, upsampling_type, normalization=normalization, preactivation=preactivation, residual=residual, padding=padding, padding_mode=padding_mode, activation=activation, dilation=self.dilation, dropout=dropout, ) self.decoding_blocks.append(decoding_block) in_channels_skip_connection //= 2 if self.dilation is not None: self.dilation //= 2 def forward(self, skip_connections, x): zipped = zip(reversed(skip_connections), self.decoding_blocks) for skip_connection, decoding_block in zipped: x = decoding_block(skip_connection, x) return x class DecodingBlock(nn.Module): def __init__( self, in_channels_skip_connection: int, dimensions: int, upsampling_type: str, normalization: Optional[str], preactivation: bool = True, residual: bool = False, padding: int = 0, padding_mode: str = "zeros", activation: Optional[str] = "ReLU", dilation: Optional[int] = None, dropout: float = 0, ): super().__init__() self.residual = residual if upsampling_type == "conv": in_channels = out_channels = 2 * in_channels_skip_connection self.upsample = get_conv_transpose_layer( dimensions, in_channels, out_channels ) else: self.upsample = get_upsampling_layer(upsampling_type) in_channels_first = in_channels_skip_connection * (1 + 2) out_channels = in_channels_skip_connection self.conv1 = ConvolutionalBlock( dimensions, in_channels_first, out_channels, normalization=normalization, preactivation=preactivation, padding=padding, padding_mode=padding_mode, activation=activation, dilation=dilation, dropout=dropout, ) in_channels_second = out_channels self.conv2 = ConvolutionalBlock( dimensions, in_channels_second, out_channels, normalization=normalization, preactivation=preactivation, padding=padding, padding_mode=padding_mode, activation=activation, dilation=dilation, dropout=dropout, ) if residual: self.conv_residual = ConvolutionalBlock( dimensions, in_channels_first, out_channels, kernel_size=1, normalization=None, activation=None, ) def forward(self, skip_connection, x): x = self.upsample(x) skip_connection = self.center_crop(skip_connection, x) x = torch.cat((skip_connection, x), dim=CHANNELS_DIMENSION) if self.residual: connection = self.conv_residual(x) x = self.conv1(x) x = self.conv2(x) x = x + connection else: x = self.conv1(x) x = self.conv2(x) return x def center_crop(self, skip_connection, x): skip_shape = torch.tensor(skip_connection.shape) x_shape = torch.tensor(x.shape) crop = skip_shape[2:] - x_shape[2:] half_crop = (crop / 2).int() # If skip_connection is 10, 20, 30 and x is (6, 14, 12) # Then pad will be (-2, -2, -3, -3, -9, -9) pad = -torch.stack((half_crop, half_crop)).t().flatten() skip_connection = F.pad(skip_connection, pad.tolist()) return skip_connection def get_upsampling_layer(upsampling_type: str) -> nn.Upsample: if upsampling_type not in UPSAMPLING_MODES: message = 'Upsampling type is "{}"' " but should be one of the following: {}" message = message.format(upsampling_type, UPSAMPLING_MODES) raise ValueError(message) upsample = nn.Upsample( scale_factor=2, mode=upsampling_type, align_corners=False, ) return upsample def get_conv_transpose_layer(dimensions, in_channels, out_channels): class_name = "ConvTranspose{}d".format(dimensions) conv_class = getattr(nn, class_name) conv_layer = conv_class(in_channels, out_channels, kernel_size=2, stride=2) return conv_layer def fix_upsampling_type(upsampling_type: str, dimensions: int): if upsampling_type == "linear": if dimensions == 2: upsampling_type = "bilinear" elif dimensions == 3: upsampling_type = "trilinear" return upsampling_type ================================================ FILE: unet/encoding.py ================================================ from typing import Optional import torch.nn as nn from .conv import ConvolutionalBlock class Encoder(nn.Module): def __init__( self, in_channels: int, out_channels_first: int, dimensions: int, pooling_type: str, num_encoding_blocks: int, normalization: Optional[str], preactivation: bool = False, residual: bool = False, padding: int = 0, padding_mode: str = "zeros", activation: Optional[str] = "ReLU", initial_dilation: Optional[int] = None, dropout: float = 0, ): super().__init__() self.encoding_blocks = nn.ModuleList() self.dilation = initial_dilation is_first_block = True for _ in range(num_encoding_blocks): encoding_block = EncodingBlock( in_channels, out_channels_first, dimensions, normalization, pooling_type, preactivation, is_first_block=is_first_block, residual=residual, padding=padding, padding_mode=padding_mode, activation=activation, dilation=self.dilation, dropout=dropout, ) is_first_block = False self.encoding_blocks.append(encoding_block) if dimensions in (1, 2): in_channels = out_channels_first out_channels_first = in_channels * 2 elif dimensions == 3: in_channels = 2 * out_channels_first out_channels_first = in_channels if self.dilation is not None: self.dilation *= 2 def forward(self, x): skip_connections = [] for encoding_block in self.encoding_blocks: x, skip_connnection = encoding_block(x) skip_connections.append(skip_connnection) return skip_connections, x @property def out_channels(self): return self.encoding_blocks[-1].out_channels class EncodingBlock(nn.Module): def __init__( self, in_channels: int, out_channels_first: int, dimensions: int, normalization: Optional[str], pooling_type: Optional[str], preactivation: bool = False, is_first_block: bool = False, residual: bool = False, padding: int = 0, padding_mode: str = "zeros", activation: Optional[str] = "ReLU", dilation: Optional[int] = None, dropout: float = 0, ): super().__init__() self.preactivation = preactivation self.normalization = normalization self.residual = residual if is_first_block: normalization = None preactivation = None else: normalization = self.normalization preactivation = self.preactivation self.conv1 = ConvolutionalBlock( dimensions, in_channels, out_channels_first, normalization=normalization, preactivation=preactivation, padding=padding, padding_mode=padding_mode, activation=activation, dilation=dilation, dropout=dropout, ) if dimensions in (1, 2): out_channels_second = out_channels_first elif dimensions == 3: out_channels_second = 2 * out_channels_first self.conv2 = ConvolutionalBlock( dimensions, out_channels_first, out_channels_second, normalization=self.normalization, preactivation=self.preactivation, padding=padding, activation=activation, dilation=dilation, dropout=dropout, ) if residual: self.conv_residual = ConvolutionalBlock( dimensions, in_channels, out_channels_second, kernel_size=1, normalization=None, activation=None, ) self.downsample = None if pooling_type is not None: self.downsample = get_downsampling_layer(dimensions, pooling_type) def forward(self, x): if self.residual: connection = self.conv_residual(x) x = self.conv1(x) x = self.conv2(x) x = x + connection else: x = self.conv1(x) x = self.conv2(x) if self.downsample is None: return x else: skip_connection = x x = self.downsample(x) return x, skip_connection @property def out_channels(self): return self.conv2.conv_layer.out_channels def get_downsampling_layer( dimensions: int, pooling_type: str, kernel_size: int = 2, ) -> nn.Module: class_name = "{}Pool{}d".format(pooling_type.capitalize(), dimensions) class_ = getattr(nn, class_name) return class_(kernel_size) ================================================ FILE: unet/unet.py ================================================ # -*- coding: utf-8 -*- """Main module.""" from typing import Optional import torch.nn as nn from .encoding import Encoder, EncodingBlock from .decoding import Decoder from .conv import ConvolutionalBlock __all__ = ["UNet", "UNet1D", "UNet2D", "UNet3D"] class UNet(nn.Module): def __init__( self, in_channels: int = 1, out_classes: int = 2, dimensions: int = 2, num_encoding_blocks: int = 5, out_channels_first_layer: int = 64, normalization: Optional[str] = None, pooling_type: str = "max", upsampling_type: str = "conv", preactivation: bool = False, residual: bool = False, padding: int = 0, padding_mode: str = "zeros", activation: Optional[str] = "ReLU", initial_dilation: Optional[int] = None, dropout: float = 0, monte_carlo_dropout: float = 0, ): super().__init__() depth = num_encoding_blocks - 1 # Force padding if residual blocks if residual: padding = 1 # Encoder self.encoder = Encoder( in_channels, out_channels_first_layer, dimensions, pooling_type, depth, normalization, preactivation=preactivation, residual=residual, padding=padding, padding_mode=padding_mode, activation=activation, initial_dilation=initial_dilation, dropout=dropout, ) # Bottom (last encoding block) in_channels = self.encoder.out_channels if dimensions in (1, 2): out_channels_first = 2 * in_channels else: out_channels_first = in_channels self.bottom_block = EncodingBlock( in_channels, out_channels_first, dimensions, normalization, pooling_type=None, preactivation=preactivation, residual=residual, padding=padding, padding_mode=padding_mode, activation=activation, dilation=self.encoder.dilation, dropout=dropout, ) # Decoder if dimensions in (1, 2): power = depth - 1 elif dimensions == 3: power = depth in_channels = self.bottom_block.out_channels in_channels_skip_connection = out_channels_first_layer * 2**power num_decoding_blocks = depth self.decoder = Decoder( in_channels_skip_connection, dimensions, upsampling_type, num_decoding_blocks, normalization=normalization, preactivation=preactivation, residual=residual, padding=padding, padding_mode=padding_mode, activation=activation, initial_dilation=self.encoder.dilation, dropout=dropout, ) # Monte Carlo dropout self.monte_carlo_layer = None if monte_carlo_dropout: dropout_class = getattr(nn, "Dropout{}d".format(dimensions)) self.monte_carlo_layer = dropout_class(p=monte_carlo_dropout) # Classifier if dimensions in (1, 2): in_channels = out_channels_first_layer elif dimensions == 3: in_channels = 2 * out_channels_first_layer self.classifier = ConvolutionalBlock( dimensions, in_channels, out_classes, kernel_size=1, activation=None, ) def forward(self, x): skip_connections, encoding = self.encoder(x) encoding = self.bottom_block(encoding) x = self.decoder(skip_connections, encoding) if self.monte_carlo_layer is not None: x = self.monte_carlo_layer(x) return self.classifier(x) class UNet1D(UNet): def __init__(self, *args, **user_kwargs): kwargs = {} kwargs["dimensions"] = 1 kwargs["num_encoding_blocks"] = 5 kwargs["out_channels_first_layer"] = 64 kwargs.update(user_kwargs) super().__init__(*args, **kwargs) class UNet2D(UNet): def __init__(self, *args, **user_kwargs): kwargs = {} kwargs["dimensions"] = 2 kwargs["num_encoding_blocks"] = 5 kwargs["out_channels_first_layer"] = 64 kwargs.update(user_kwargs) super().__init__(*args, **kwargs) class UNet3D(UNet): def __init__(self, *args, **user_kwargs): kwargs = {} kwargs["dimensions"] = 3 kwargs["num_encoding_blocks"] = 4 kwargs["out_channels_first_layer"] = 32 kwargs["normalization"] = "batch" kwargs.update(user_kwargs) super().__init__(*args, **kwargs)