The architecture in UNET consists of two paths. Contracting path (encoder network) followed by an expansive path (decoder network). U-Net gets its name from the architecture, like the letter U, as depicted in the image below.
The encoder (left side) typically consists of a pre-trained classification network, such as ResNet, using convolution blocks followed by maxpool downsampling to encode the input image into feature representations at various levels. In this architecture, there is repeated application of two 3x3 convolutions. Each convolution is followed by a ReLU and batch normalization. Then a 2x2 max pooling operation is applied to reduce dimensions by half. Again, at each downsampling step, we double the number of feature channels, while we cut in half the spatial dimensions.
The decoder path (right side) consists of upsampling of the feature map followed by a 2x2 transpose convolution, which cuts the number of feature channels in half. The concatenation with the matching feature map from the contracting path, as well as a 3x3 convolutional filter (each followed by a ReLU) is also implemented. A 1x1 convolution is employed in the final layer to map the channels to the required number of classes.
Hello, thank you for using the code provided by Hasty. Please note that some code blocks might not be 100% complete and ready to be run as is. This is done intentionally as we focus on implementing only the most challenging parts that might be tough to pick up from scratch. View our code block as a LEGO block - you can’t use it as a standalone solution, but you can take it and add to your system to complement it. If you have questions about using the tool, please get in touch with us to get direct help from the Hasty team.
import torch
#Loading a U-Net model pre-trained for abnormality segmentation
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
in_channels=3, out_channels=1, init_features=32, pretrained=True)
# Download an example input image
import urllib
url, filename = ("https://github.com/mateuszbuda/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png", "TCGA_CS_4944.png")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)
import numpy as np
from PIL import Image
from torchvision import transforms
input_image = Image.open(filename)
m, s = np.mean(input_image, axis=(0, 1)), np.std(input_image, axis=(0, 1))
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=m, std=s),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model = model.to('cuda')
with torch.no_grad():
output = model(input_batch)
print(torch.round(output[0]))
In Model Playground, we can select feature extraction (encoding) network to use as either ResNet or EffiecientNet. The depth of the ResNet model/ EfficientNet sub type has to be specified as well.
After selecting the desired encoder network, the weights to be used for model initialization can be chosen.
U-Net: Convolutional Networks for Biomedical Image Segmentation
Automate 90% of the work, reduce your time to deployment by 40%, and replace your whole ML software stack with our platform.
Start for free Request a demo