DeepLabV3+：对预测处理的详解

相信大家对于这一部分才是最感兴趣的，能够实实在在的看到效果。这里我们就只需要两个.py文件（deeplab.py、predict_img.py）。

创建DeeplabV3类

deeplab.py的作用是为了创建一个DeeplabV3类，提供一个检测图片的方法，而predict_img.py则是为了单独检测图片的效果。

在这里我需要一个defaults字典用来包含我在这个类要使用的变量，而需要把数据类型转换成字典数据再做存储，这时候就需要用到类的内置属性__dict__。

这里简单说一下字典defaults的key和value。

1、model_path=model_date/deeplab_mobilenetv2.pth

此文件是基于VOC拓展数据集训练的权重，放心使用，附上下载地址的权值文件。

2、num_classes=2

对于需要区分的类数+1，比如我这里是识别裂缝，所以我的num_classes为1+1，再比如经典的猫狗分类问题，那么它们的num_classes为2+1＝3。

3、backbone=mobilenet

这里是使用的主干网络，有mobilenet和xception可供选择。

4、input_shape=[512,512]

输入图片的大小

5、downsample_factor=16

下采样的倍数，可选的有8和16，但8训练要求更大的内存，这里要与训练时相同。

6、mix_type=0

0代表原图与生成的图进行混合；1代表仅保留生成的图；2代表扣去背景，仅保留原图中的目标。

7、cude=False

有cuda就是Ture，没有就用cpu。

在这里，请看detect_image函数里面，首先要用cvtColor函数对图片进行一个转化，因为RGB图像才有权重。

在对图像的大小修改时，需要增添一个灰度框，想想这样的请况，如果图像比输入大小小就会使图像被强行放大，可能会伸长也可能会扩展。那么为了避免这种请况，所以要添加这个灰度边界，后期因为要与原图大小匹配，会将这部分去掉。

然后，对图像的每个像素点进行分类。

# deeplab.pyimport colorsys
import copyimport cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nnfrom Deeplabv3_plus.deeplabv3plus import DeepLab
from utily.utils import cvtColor, preprocess_input, resize_image, show_configclass DeeplabV3(object):defaults = {"model_path": 'model_data/deeplab_mobilenetv2.pth',"num_classes": 2,"backbone": "mobilenet","input_shape": [512, 512],  "downsample_factor": 16,  "mix_type": 0,"cuda": False}def __init__(self, **kwargs):self.__dict__.update(self.defaults)for name, value in kwargs.items():setattr(self, name, value)if self.num_classes <= 21:self.colors = [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128), (128, 64, 12)]# 画框设置不同的颜色else:hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))# 获得模型self.generate()show_config(**self.defaults)# 获得所有的分类def generate(self, onnx=False):# 载入模型与权值self.net = DeepLab(num_classes=self.num_classes, backbone=self.backbone, downsample_factor=self.downsample_factor, pretrained=False)device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')self.net.load_state_dict(torch.load(self.model_path, map_location=device))self.net = self.net.eval()print('{} model, and classes loaded.'.format(self.model_path))if not onnx:if self.cuda:self.net = nn.DataParallel(self.net)self.net = self.net.cuda()def detect_image(self, image, count=False, name_classes=None):"""* 在这里将图像转换成RGB图像，防止灰度图在预测时报错。* 代码仅仅支持RGB图像的预测，所以其它类型的图像都会转化成RGB:param image: 图片:param count: 计数:param name_classes::return:"""image = cvtColor(image)# 对输入图像进行一个备份，后面用于绘图old_img = copy.deepcopy(image)orininal_h = np.array(image).shape[0]orininal_w = np.array(image).shape[1]# 给图像增加灰条，实现不失真的resize# 也可以直接resize进行识别image_data, nw, nh  = resize_image(image, (self.input_shape[1],self.input_shape[0]))# 添加上batch_size维度image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)with torch.no_grad():images = torch.from_numpy(image_data)if self.cuda:images = images.cuda()# 图片传入网络进行预测pr = self.net(images)[0]# 取出每一个像素点的种类pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy()# 将灰条部分截取掉pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]# 进行图片的resizepr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR)# 取出每一个像素点的种类pr = pr.argmax(axis=-1)if count:classes_nums        = np.zeros([self.num_classes])total_points_num    = orininal_h * orininal_wprint('-' * 63)print("|%25s | %15s | %15s|"%("Key", "Value", "Ratio"))print('-' * 63)for i in range(self.num_classes):num     = np.sum(pr == i)ratio   = num / total_points_num * 100if num > 0:print("|%25s | %15s | %14.2f%%|"%(str(name_classes[i]), str(num), ratio))print('-' * 63)classes_nums[i] = numprint("classes_nums:", classes_nums)if self.mix_type == 0:# seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))# for c in range(self.num_classes):#     seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')#     seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')#     seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])# 将新图片转换成Image的形式image   = Image.fromarray(np.uint8(seg_img))# 将新图与原图及进行混合image   = Image.blend(old_img, image, 0.7)elif self.mix_type == 1:# seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))# for c in range(self.num_classes):#     seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')#     seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')#     seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])# 将新图片转换成Image的形式image   = Image.fromarray(np.uint8(seg_img))elif self.mix_type == 2:seg_img = (np.expand_dims(pr != 0, -1) * np.array(old_img, np.float32)).astype('uint8')# 将新图片转换成Image的形式image = Image.fromarray(np.uint8(seg_img))return image

单张图片的预测

由于我想要将图片放在PyQt5设计的ui中，所以要单张单张的显示。

from PIL import Image
from deeplab import DeeplabV3if __name__ == "__main__":deeplab = DeeplabV3()mode = "predict" count = False    #指定了是否进行目标的像素点计数（即面积）与比例计算# name_classes    = ["background","aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]name_classes = ["background","crack"]if mode == "predict":while True:img = input('Input image filename:')try:image = Image.open(img)except:print('Open Error! Try again!')continueelse:r_image = deeplab.detect_image(image, count=count, name_classes=name_classes)r_image.show()

我们来看看效果：

原图

效果图