Eric Chen's Blog Semantic Segmentation using FCN and DeepLabV3¶
Semantic Segmentation is an image analysis task in which we classify each pixel in the image into a class.
In this post, we will perform semantic segmentation using pre-trained models built in Pytorch. They are FCN and DeepLabV3.
Understanding model inputs and outputs¶
Now before we get started, we need to know about the inputs and outputs of these semantic segmentation models. So, let's start!
These models expect a 3-channled image which is normalized with the Imagenet mean and standard deviation, i.e.,
mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]
So, the input is [Ni x Ci x Hi x Wi]
where,
Ni-> the batch sizeCi-> the number of channels (which is 3)Hi-> the height of the imageWi-> the width of the image
And the output of the model is [No x Co x Ho x Wo]
where,
No-> is the batch size (same asNi)Co-> is the number of classes that the dataset have! So each class will have a map in the output.Ho-> the height of the image (which is the same asHiin almost all cases)Wo-> the width of the image (which is the same asWiin almost all cases)
Note that torchvision models output is an OrderedDict and not a torch.Tensor
And in .eval() mode it just has one key out and thus to get the output we need to get the value
stored in that key.
The out key of this OrderedDict is the key that holds the output.
FCN with Resnet-101 backbone¶
FCN - Fully Convolutional Netowrks, are among the most early invented Neural Networks for the task of Semantic Segmentation.
Let's load one up!
from torchvision import models
fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()
And that's it we have a pretrained model of FCN (which stands for Fully Convolutional Neural Networks) with a Resnet101 backbone :)
Now, let's get an image!
from PIL import Image
import matplotlib.pyplot as plt
import torch
!wget -nv https://www.goodfreephotos.com/cache/other-photos/car-and-traffic-on-the-road-coming-towards-me.jpg -O car.png
img = Image.open('./car.png')
plt.imshow(img)
plt.axis('off')
plt.show()
Now, that we have the image we need to preprocess it and normalize it! These preprocess is following the ImageNet training data.
# Apply the transformations needed
import torchvision.transforms as T
trf = T.Compose([T.Resize(256),
T.CenterCrop(224),
T.ToTensor(),
T.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225])])
inp = trf(img)
print(inp.shape)
inp = inp.unsqueeze(0)
print(inp.shape)
Forward pass the input tensor. The output of the model is a OrderedDict so, we need to take the out key from that to get the output of the model.
# Pass the input through the net
out = fcn(inp)['out']
print (out.shape)
Alright! So, out is the final output of the model. And as we can see, its shape is [1 x 21 x H x W] as discussed earlier. So, the model was trained on 21 classes and thus our output have 21 channels!
Since the pretrained FCN segmentation model is trained on PASCAL VOC dataset, which have 20 class labels + 1 background class. Therefore, there are 21 classes in total. Now we need to compress the 21 channels into a single channel which contains all the class labels for each pixel.
import numpy as np
om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()
print (om.shape)
print (np.unique(om))
So, we as we can see now have a 2D image. Where each pixel corresponds to a class!
The last thing is to take this 2D image where each pixel corresponds to a class label and convert this
into a segmentation map where each class label is converted into a RGB color and thus helping in easy visualization.
We will use the following function to convert this 2D image to an RGB image wheree each label is mapped to its
corresponding color.
# Define the helper function
def decode_segmap(image, nc=21):
label_colors = np.array([(0, 0, 0), # 0=background
# 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
(128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
# 6=bus, 7=car, 8=cat, 9=chair, 10=cow
(0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),
# 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
(192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
# 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
(0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])
r = np.zeros_like(image).astype(np.uint8)
g = np.zeros_like(image).astype(np.uint8)
b = np.zeros_like(image).astype(np.uint8)
for l in range(0, nc):
idx = image == l
r[idx] = label_colors[l, 0]
g[idx] = label_colors[l, 1]
b[idx] = label_colors[l, 2]
rgb = np.stack([r, g, b], axis=2)
return rgb
rgb = decode_segmap(om)
plt.imshow(rgb); plt.show()
Because we are doing semantic segmentation. So the two cars in the middle are assigned to the same class. And that class is car
Also, Do note that the image after segmentation is smaller than the original image as in the preprocessing step the image is resized and cropped.
Next, let's move all this under one function and play with a few more images!
def segment(net, path, show_orig=True, dev='cuda'):
img = Image.open(path)
if show_orig: plt.imshow(img); plt.axis('off'); plt.show()
# Comment the Resize and CenterCrop for better inference results
trf = T.Compose([T.Resize(640),
#T.CenterCrop(224),
T.ToTensor(),
T.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225])])
inp = trf(img).unsqueeze(0).to(dev)
out = net.to(dev)(inp)['out']
om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()
rgb = decode_segmap(om)
plt.imshow(rgb); plt.axis('off'); plt.show()
And let's get a new image!
!wget -nv https://upload.wikimedia.org/wikipedia/commons/thumb/f/f4/Sussex_cow_4.JPG/1200px-Sussex_cow_4.JPG -O cow.jpg
segment(fcn, './cow.jpg')
!wget -nv https://storage.needpix.com/rsynced_images/pedestrian-zone-456909_1280.jpg -O pedestrian.jpg
segment(fcn, './pedestrian.jpg')
DeepLabv3 with Resnet-101 backbone¶
dlab = models.segmentation.deeplabv3_resnet101(pretrained=1).eval()
Alright! Now we have god-level segmentation model!
Let's see how we perform with the same image on this model!
segment(dlab, './car.png')
Wow. It looks like DeepLabV3 is able to find the third car which is very dark in the original picture.
Let's see how it performs on other pictures.
segment(dlab, './cow.jpg')
segment(dlab, './pedestrian.jpg')
Comparision¶
For, now we will see how these two models compare with each other in 3 metrics
- Inference time
- Size of the model
- GPU memory used by the model
Inference Time¶
import time
def infer_time(net, path='./pedestrian.jpg', dev='cuda'):
img = Image.open(path)
trf = T.Compose([T.Resize(256),
T.CenterCrop(224),
T.ToTensor(),
T.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225])])
inp = trf(img).unsqueeze(0).to(dev)
st = time.time()
out1 = net.to(dev)(inp)
et = time.time()
return et - st
On CPU
avg_over = 100
fcn_infer_time_list_cpu = [infer_time(fcn, dev='cpu') for _ in range(avg_over)]
fcn_infer_time_avg_cpu = sum(fcn_infer_time_list_cpu) / avg_over
dlab_infer_time_list_cpu = [infer_time(dlab, dev='cpu') for _ in range(avg_over)]
dlab_infer_time_avg_cpu = sum(dlab_infer_time_list_cpu) / avg_over
print ('Inference time for first few calls for FCN : {}'.format(fcn_infer_time_list_cpu[:10]))
print ('Inference time for first few calls for DeepLabv3: {}'.format(dlab_infer_time_list_cpu[:10]))
print ('The Average Inference time on FCN is: {:.2f}s'.format(fcn_infer_time_avg_cpu))
print ('The Average Inference time on DeepLab is: {:.2f}s'.format(dlab_infer_time_avg_cpu))
On GPU
avg_over = 100
fcn_infer_time_list_gpu = [infer_time(fcn) for _ in range(avg_over)]
fcn_infer_time_avg_gpu = sum(fcn_infer_time_list_gpu) / avg_over
dlab_infer_time_list_gpu = [infer_time(dlab) for _ in range(avg_over)]
dlab_infer_time_avg_gpu = sum(dlab_infer_time_list_gpu) / avg_over
print ('Inference time for first few calls for FCN : {}'.format(fcn_infer_time_list_gpu[:10]))
print ('Inference time for first few calls for DeepLabv3: {}'.format(dlab_infer_time_list_gpu[:10]))
print ('The Average Inference time on FCN is: {:.3f}s'.format(fcn_infer_time_avg_gpu))
print ('The Average Inference time on DeepLab is: {:.3f}s'.format(dlab_infer_time_avg_gpu))
We can see that in both cases (for GPU and CPU) its taking longer for the DeepLabv3 model, as its a much deeper model as compared to FCN.
plt.bar([0.1, 0.2], [fcn_infer_time_avg_cpu, dlab_infer_time_avg_cpu], width=0.08)
plt.ylabel('Time/s')
plt.xticks([0.1, 0.2], ['FCN', 'DeepLabv3'])
plt.title('Inference time of FCN and DeepLabv3 with Resnet-101 backbone on CPU')
plt.show()
plt.bar([0.1, 0.2], [fcn_infer_time_avg_gpu, dlab_infer_time_avg_gpu], width=0.08)
plt.ylabel('Time/s')
plt.xticks([0.1, 0.2], ['FCN', 'DeepLabv3'])
plt.title('Inference time of FCN and DeepLabv3 with Resnet-101 backbone on GPU')
plt.show()
Okay! Now, let's move on to the next comparison, where we will compare the model sizes for both the models.
Model Size¶
import os
resnet101_size = os.path.getsize('/root/.cache/torch/checkpoints/resnet101-5d3b4d8f.pth')
fcn_size = os.path.getsize('/root/.cache/torch/checkpoints/fcn_resnet101_coco-7ecb50ca.pth')
dlab_size = os.path.getsize('/root/.cache/torch/checkpoints/deeplabv3_resnet101_coco-586e9e4e.pth')
fcn_total = fcn_size + resnet101_size
dlab_total = dlab_size + resnet101_size
print ('Size of the FCN model with Resnet101 backbone is: {:.2f} MB'.format(fcn_total / (1024 * 1024)))
print ('Size of the DeepLabv3 model with Resnet101 backbone is: {:.2f} MB'.format(dlab_total / (1024 * 1024)))
plt.bar([0, 1], [fcn_total / (1024 * 1024), dlab_total / (1024 * 1024)])
plt.ylabel('Size of the model in MegaBytes')
plt.xticks([0, 1], ['FCN', 'DeepLabv3'])
plt.title('Comparison of the model size of FCN and DeepLabv3 with Resnet-101 backbone')
plt.show()
Comments
comments powered by Disqus