Closed ZJDATY closed 1 year ago
非常感谢，还有个问题请教，例如一个deeplab v3+语义分割模型，训练出来的20类语义模型，格式是 1x21x512x512 。它的后处理基本都是python实现的，我们一般部署都用C++来实现，我想在onnx模型里就处理好，先softmax 再 argmax，最终输出 1x1x512x512 的结果。我想知道是否可以用该工具实现，可以的话，大佬能教一下怎么修改吗？ 用的这个开源示例。 https://github.com/bubbliiiing/deeplabv3-plus-pytorch
下面是后处理的实现
#---------------------------------------------------# # 检测图片 #---------------------------------------------------# 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
模型在这里：链接: https://pan.baidu.com/s/1cP3VmAy_CSG5L1k6eRx_4A?pwd=mdsw 提取码: mdsw
@ZJDATY 谢谢提醒，百度网盘下载链接已更新。
我想在onnx模型里就处理好，先softmax 再 argmax，最终输出 1x1x512x512 的结果。
可以尝试一下新增节点功能，用ONNX算子把后处理的语义组合出来。
ZhangGe6
commented
1 year ago
非常感谢，还有个问题请教，例如一个deeplab v3+语义分割模型，训练出来的20类语义模型，格式是 1x21x512x512 。它的后处理基本都是python实现的，我们一般部署都用C++来实现，我想在onnx模型里就处理好，先softmax 再 argmax，最终输出 1x1x512x512 的结果。我想知道是否可以用该工具实现，可以的话，大佬能教一下怎么修改吗？ 用的这个开源示例。 https://github.com/bubbliiiing/deeplabv3-plus-pytorch
下面是后处理的实现
模型在这里：链接: https://pan.baidu.com/s/1cP3VmAy_CSG5L1k6eRx_4A?pwd=mdsw 提取码: mdsw