Posts

Showing posts from February, 2022

PyTorch and Albumentations for Image Segmentation

Image
This example illustrates PyTorch how to use Albumentations for binary semantic segmentation. We will use The Oxford-IIIT Pet Dataset . The task is to classify each pixel of the input image as a pet or a background. Installing the required libraries We will use TernausNet, a library that provides a pre-trained UNet model for the semantic segmentation task. pip install ternausnet ternaus/TernausNet: UNet model with VGG11 encoder pre-trained on Kaggle Carvana dataset (github.com) Importing the required libraries from collections import defaultdictimport copyimport randomimport osimport shutilfrom urllib.request import urlretrieveimport albumentations as Aimport albumentations.augmentations.functional as Ffrom albumentations.pytorch import ToTensorV2import cv2import matplotlib.pyplot as pltimport numpy as npimport ternausnet.modelsfrom tqdm import tqdmimport torchimport torch.backends.cudnn as cudnnimport torch.nn as nnimport torch.optimfrom torch.utils.data import Dataset, DataLoadercud