import colorsys import copy import time import cv2 import numpy as np import torch import torch.nn.functional as F from PIL import Image from torch import nn from nets.deeplabv3_plus import DeepLab from utils.utils import cvtColor, preprocess_input, resize_image, show_config #-----------------------------------------------------------------------------------# # 使用自己训练好的模型预测需要修改3个参数 # model_path、backbone和num_classes都需要修改！ # 如果出现shape不匹配，一定要注意训练时的model_path、backbone和num_classes的修改 #-----------------------------------------------------------------------------------# class DeeplabV3(object): _defaults = { #-------------------------------------------------------------------# # model_path指向logs文件夹下的权值文件 # 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。 # 验证集损失较低不代表miou较高，仅代表该权值在验证集上泛化性能较好。 #-------------------------------------------------------------------# "model_path" : 'model_data/last_epoch_weights1.pth', #----------------------------------------# # 所需要区分的类的个数+1 #----------------------------------------# "num_classes" : 46, #----------------------------------------# # 所使用的的主干网络： # mobilenet # xception #----------------------------------------# "backbone" : "mobilenet", #----------------------------------------# # 输入图片的大小 #----------------------------------------# "input_shape" : [1024, 1042], #----------------------------------------# # 下采样的倍数，一般可选的为8和16 # 与训练时设置的一样即可 #----------------------------------------# "downsample_factor" : 16, #-------------------------------------------------# # mix_type参数用于控制检测结果的可视化方式 # # mix_type = 0的时候代表原图与生成的图进行混合 # mix_type = 1的时候代表仅保留生成的图 # mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标 #-------------------------------------------------# "mix_type" : 0, #-------------------------------# # 是否使用Cuda # 没有GPU可以设置成False #-------------------------------# "cuda" : True, } #---------------------------------------------------# # 初始化Deeplab #---------------------------------------------------# def __init__(self, **kwargs): self.__dict__.update(self._defaults) for name, value in kwargs.items(): setattr(self, name, value) #---------------------------------------------------# # 画框设置不同的颜色 #---------------------------------------------------# if self.num_classes <= 46: 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), (0, 0, 142), (119, 11, 32), (244,164,140), (188,143,143), (64,224,205), (127,255,0), (199,97,20), (189,252,201), (0,255,127), (160,32,240), (138,42,226), (255,97,0), (255,215,0), (255,128,0), (189,252,201), (240,255,240), (0, 130, 180), (152, 251, 152), (107, 142, 35), (153, 153, 153), (190, 153, 153), (250, 170, 30), (220, 220, 0), (107, 142, 35), ] 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 #---------------------------------------------------------# 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)] #---------------------------------------------------# # 进行图片的resize #---------------------------------------------------# pr = 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_w print('-' * 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 * 100 if num > 0: print("|%25s | %15s | %14.2f%%|"%(str(name_classes[i]), str(num), ratio)) print('-' * 63) classes_nums[i] = num print("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 def get_FPS(self, image, test_interval): #---------------------------------------------------------# # 在这里将图像转换成RGB图像，防止灰度图在预测时报错。 # 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB #---------------------------------------------------------# image = cvtColor(image) #---------------------------------------------------------# # 给图像增加灰条，实现不失真的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().argmax(axis=-1) #--------------------------------------# # 将灰条部分截取掉 #--------------------------------------# 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)] t1 = time.time() for _ in range(test_interval): with torch.no_grad(): #---------------------------------------------------# # 图片传入网络进行预测 #---------------------------------------------------# pr = self.net(images)[0] #---------------------------------------------------# # 取出每一个像素点的种类 #---------------------------------------------------# pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy().argmax(axis=-1) #--------------------------------------# # 将灰条部分截取掉 #--------------------------------------# 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)] t2 = time.time() tact_time = (t2 - t1) / test_interval return tact_time def convert_to_onnx(self, simplify, model_path): import onnx self.generate(onnx=True) im = torch.zeros(1, 3, *self.input_shape).to('cpu') # image size(1, 3, 512, 512) BCHW input_layer_names = ["images"] output_layer_names = ["output"] # Export the model print(f'Starting export with onnx {onnx.__version__}.') torch.onnx.export(self.net, im, f = model_path, verbose = False, opset_version = 12, training = torch.onnx.TrainingMode.EVAL, do_constant_folding = True, input_names = input_layer_names, output_names = output_layer_names, dynamic_axes = None) # Checks model_onnx = onnx.load(model_path) # load onnx model onnx.checker.check_model(model_onnx) # check onnx model # Simplify onnx if simplify: import onnxsim print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.') model_onnx, check = onnxsim.simplify( model_onnx, dynamic_input_shape=False, input_shapes=None) assert check, 'assert check failed' onnx.save(model_onnx, model_path) print('Onnx model save as {}'.format(model_path)) def get_miou_png(self, image): #---------------------------------------------------------# # 在这里将图像转换成RGB图像，防止灰度图在预测时报错。 # 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB #---------------------------------------------------------# image = cvtColor(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)] #---------------------------------------------------# # 进行图片的resize #---------------------------------------------------# pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR) #---------------------------------------------------# # 取出每一个像素点的种类 #---------------------------------------------------# pr = pr.argmax(axis=-1) image = Image.fromarray(np.uint8(pr)) return image