Universal implementation of the UNet architecture for image segmentation.
Project description
UNET SEGMENTATION PYTORCH
Installation
pip install segment-torch
Usage
from segment_torch.unet import UNet
from torch import nn
device = "cuda"
config = dict(
in_channels=3,
out_channels=1,
hiddens=[4, 8, 16, 32],
dropouts=[0, 0.15, 0.15, 0.15], # hiddens
maxpools=2, # hiddens - 1
kernel_sizes=3, # 2*hiddens + 3*hiddens + 2
paddings='same', # 2*hiddens + 3*hiddens + 2
strides=1, # 2*hiddens + 3*hiddens
dilation=1,
criterion=nn.BCELoss(),
output_activation=nn.Sigmoid(),
activation=nn.ReLU(),
dimensions=2,
device=device
)
unet = UNet(**config)
Different ways to define configs
# 0. None: default values are used
kernel_sizes=None
# 1. Single value or tuple: all layers have the same value
kernel_sizes = 3
kernel_sizes = (3, 3)
# 2. Lists of values
encooder_kernel_sizes = [3, 3, 3, 3]
decoder_kernel_sizes = [3, 3, 3, 3, 3]
kernel_sizes = [encooder_kernel_sizes, decoder_kernel_sizes]
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
segment_torch-0.0.6.tar.gz
(8.5 kB
view hashes)
Built Distribution
Close
Hashes for segment_torch-0.0.6-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 5f3f83642ba842eb290b466474109a5f71e48024e150afc944f15bea9d506e30 |
|
MD5 | 4457003957fedbacef0f222dd848e60b |
|
BLAKE2b-256 | d7b70341a9a2534c92cff038b536322f751aff4f379e3031e0b95e6ff7c3b3c9 |