Course
PyTorch模型推理及多任务通用范式 课程4:
- 对语义分割任务和DeepLabV3模型做了简单介绍;
- 根据 pytorch模型推理的三板斧:数据预处理、数据进网络、数据后处理,逐行实现了DeepLabV3的推理代码;
- 对模型输入大小的选取,做了详细介绍。
Assignment
必做题
- 对 "./images/car.jpg" 做语义分割,提取出里面的车辆,模仿上课时,对“可视化推理结果”和“BGRA四通道图”进行保存。
- 自己找2张其他图,对图中某个类别进行分割,并保存“BGRA四通道图”。
思考题
- 用time模块和for循环,对”./images/car.jpg”连续推理100次,统计时间开销。有CUDA的同学,改下代码:self.device=torch.device('cuda'),统计时间开销。
- 以0.5为阈值,计算”./images/car.jpg”图中车辆的面积(单位:像素)。
Solutions
Code
import torch
import torchvision.models as models
import numpy as np
import cv2
import torch.nn.functional as F
import time
class ModelPipline(object):
def __init__(self, device=torch.device('cuda')):
# 进入模型的图片大小:为数据预处理和后处理做准备
self.inputs_size = (520, 520)
# CPU or CUDA:为数据预处理和模型加载做准备
self.device = device
# 载入模型结构和模型权重
self.model = self.get_model()
def predict(self, id, image):
# 数据预处理
inputs, image_h, image_w = self.preprocess(image)
# 数据进网路
outputs = self.model(inputs)
# 数据后处理
results = self.postprocess(id, outputs, image_h, image_w)
return results
def get_model(self):
# 上一节课的内容
model = models.segmentation.deeplabv3_resnet50(num_classes=21, pretrained_backbone=False, aux_loss=True)
pretrained_state_dict = torch.load('./weights/deeplabv3_resnet50_coco-cd0a2569.pth',
map_location=lambda storage, loc: storage)
model.load_state_dict(pretrained_state_dict, strict=True)
model.to(self.device)
model.eval()
return model
def preprocess(self, image):
# opencv默认读入是BGR,需要转为RGB,和训练时保持一致
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 提取原图大小
image_h, image_w = image.shape[:2]
# resize成模型输入的大小,和训练时保持一致
image = cv2.resize(image, dsize=self.inputs_size)
# 归一化和标准化,和训练时保持一致
inputs = image / 255
inputs = (inputs - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])
##以下是图像任务的通用处理
# (H,W,C) ——> (C,H,W)
inputs = inputs.transpose(2, 0, 1)
# (C,H,W) ——> (1,C,H,W)
inputs = inputs[np.newaxis, :, :, :]
# NumpyArray ——> Tensor
inputs = torch.from_numpy(inputs)
# dtype float32
inputs = inputs.type(torch.float32)
# 与self.model放在相同硬件上
inputs = inputs.to(self.device)
return inputs, image_h, image_w
def postprocess(self, id, outputs, image_h, image_w):
# 获取模型输出output
outputs = outputs['out']
# 取softmax得到每个类别的置信度
outputs = torch.softmax(outputs, dim=1)
# 取出目标标签(比如:人体)的那一层置信度
outputs = outputs[:, id : id + 1, :, :]
# 将结果图resize回原图大小
outputs = F.interpolate(outputs, size=(image_h, image_w), mode='bilinear', align_corners=True)
# 数据类型转换:torch.autograd.Variable ——> torch.Tensor ——> numpy.ndarray
mask_person = outputs.data.cpu().numpy().squeeze()
return mask_person
if __name__ == '__main__':
def write_output(image, result, output_1, output_2):
cv2.imwrite(output_1, (result * 255).astype(np.uint8))
mask = result.copy()
mask[mask >= 0.5] = 255
mask[mask < 0.5] = 0
# count pixel for area
unique, counts = np.unique(mask, return_counts=True)
# print area on image
text = 'area = ' + str(counts[1])
font = cv2.FONT_HERSHEY_SIMPLEX
text_size = cv2.getTextSize(text, font, 1, 2)[0]
text_x = round((mask.shape[1] - text_size[0]) / 2)
text_y = mask.sum(axis=1).argmax()
cv2.putText(image, text, (text_x, text_y), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
image_mask = np.concatenate([image, mask[:, :, np.newaxis]], axis=2)
cv2.imwrite(output_2, image_mask)
# 实例化
model_segment = ModelPipline(device=torch.device('cpu'))
# get label list
with open('./labels/pascalvoc_label.txt') as f:
labels = f.read().splitlines()
# 第一张图
image = cv2.imread('./images/car.jpg')
id = labels.index('car')
result = model_segment.predict(id, image)
write_output(image, result, './results/car.jpg', './results/car.png')
# 第二张图
image = cv2.imread('./images/j20.jpg')
id = labels.index('aeroplane')
result = model_segment.predict(id, image)
write_output(image, result, './results/j20.jpg', './results/j20.png')
# 第三张图
image = cv2.imread('./images/hr.jpg')
id = labels.index('train')
result = model_segment.predict(id, image)
write_output(image, result, './results/hr.jpg', './results/hr.png')
# CPU run 100 times
image = cv2.imread('./images/car.jpg')
id = labels.index('car')
model_segment = ModelPipline(torch.device('cpu'))
t_all=0
for i in range(100):
t_start = time.time()
result = model_segment.predict(id, image)
t_end = time.time()
t_all += t_end - t_start
print('CPU 100 time lapse: {:.4f} seconds.'.format(t_all))
# GPU run 100 times
model_segment = ModelPipline(torch.device('cuda'))
t_all = 0
for i in range(100):
t_start = time.time()
result = model_segment.predict(id, image)
t_end = time.time()
t_all += t_end - t_start
print('GPU 100 time lapse: {:.4f} seconds.'.format(t_all))
必做题
思考题
- CPU推理和CUDA推理,各自的时间开销。
CPU 100 time lapse: 93.3184 seconds.
GPU 100 time lapse: 9.1371 seconds.
时间开销明显大于图像分类。 - 面积(单位:像素)。
通过统计模型输出中大于0.5的值的个数,得到面积,并打印在BGRA四通道图上。
其中,图一车辆的面积为102781。
学习心得
对语义分割模型的理解还需要学习,没理解清楚模型如何做到基于像素得到分类概率的。