YOLACT时2019新出的用于实例分割的深度学习模型,在一番环境配置后成功复现了yolact,针对改进版的YOLACT++我还未尝试,下面简单的测试了yolact模型中对mask的处理,当然限于时间精力的有限,本文是针对一个博主的一篇文章(点此查看)出现的错误进行的小小改动,避免后来学习yolact的人在这个地方又踩坑,仅此提供一个参考
出错代码:
#查看预测结果
for key in preds[0].keys():
if key=='class'or key=='score':
print(key,':',preds[0][key].shape,'\t',preds[0][key])
else:
print(key,':',preds[0][key].shape)
原因分析:
由pycharm的debug功能可调试查看到一下信息:
从图中可以看出preds是一个list,里面包含两个dict:键分别为detection和net,我们需要的是detection该键对应的值,这个值是包含了mask , proto,score ,class,box这5个键的字典。所以preds[0].keys()取得的是detection和net,我改动后的代码如下:
d = preds[0].get('detection')
for key in d.keys():
if key == 'class' or key == 'score':
print(key, ':', d[key].shape, '\t', d[key])
else:
print(key, ':', d[key].shape)
到此问题解决,希望对路人有所帮助
(懒人福利)完整代码,复制即可运行:
from model.yolact import Yolact
from model.utils.augmentations import BaseTransform, FastBaseTransform, Resize
from model.data import cfg, set_cfg, set_dataset
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import os
import matplotlib.pyplot as plt
import cv2
from pylab import *
import matplotlib.patches as patches
def evalimage(net: Yolact, path: str, save_path: str = None):
"""
net:即yolact网络
path:给定图片路径
savepath:该参数该函数暂时不用
preds:模型正向跑完一张图的结果
"""
frame = torch.from_numpy(cv2.imread(path)).cuda().float()
batch = FastBaseTransform()(frame.unsqueeze(0))
preds = net(batch)
return preds
# img_numpy = prep_display(preds, frame, None, None, undo_transform=False)
#
# if save_path is None:
# img_numpy = img_numpy[:, :, (2, 1, 0)]
#
# if save_path is None:
# plt.imshow(img_numpy)
# plt.title(path)
# plt.show()
# else:
# cv2.imwrite(save_path, img_numpy)
if __name__ == '__main__':
CONFIG = 'yolact_base_config' # 导入yolact模型指定resnet50的backbone
MODEL_PATH = 'weights/yolact_base_54_800000.pth' # 下载下来的预训练模型
PIC_PATH = 'dog.png' # 测试图片路径
set_cfg(CONFIG) # yolact项目指定的导入config的函数
# 我们先查看一下这张图
image = plt.imread(PIC_PATH)
plt.imshow(image)
plt.axis('off') # 不显示坐标轴
plt.show()
# 预测过程,preds即模型计算结果
with torch.no_grad():
cudnn.benchmark = True
cudnn.fastest = True
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net = Yolact()
net.load_weights(MODEL_PATH)
net.eval()
preds = evalimage(net, PIC_PATH)
# 查看预测结果
d = preds[0].get('detection')
for key in d.keys():
if key == 'class' or key == 'score':
print(key, ':', d[key].shape, '\t', d[key])
else:
print(key, ':', d[key].shape)
proto = d['proto'].cpu().numpy()
plt.figure(figsize=(4 * 1.38 * 2, 8 * 1.38 * 2)) # 8行4列
# 这里可以看见32个特征图的效果,我们可以发现这些有的是加强前景(1排2列,5排3列),
# 有的是加强背景(2排4列,8排4列),有的是加强左边(一二排3列),有的是加强右边(4排3列)
# yolact所有的mask都是根据这32(k)张特征图进行不同的加权方式产生的。而决定加权方式的就是mask系数。
for i in range(32):
plt.subplot(8, 4, i + 1)
plt.imshow(proto[:, :, i])
plt.axis('off')
plt.show()
# 模型数据转为numpy格式
box = d['box'].cpu().numpy()
score = d['score'].cpu().numpy()
class_id = d['class'].cpu().numpy()
mask = d['mask'].cpu().numpy()
##############################
de_num = mask.shape[0] # 确定检测到目标的个数
col = 4 # 展示4列
row = (de_num / col + 0.5) # 展示行数
plt.figure(figsize=(col * 1.38 * 2, row * 1.38 * 2))
for j in range(de_num):
result = proto * np.transpose(mask[j]) # proto乘以mask系数
result = 1 / (1 + np.exp(-result)) # sigmoid处理
result = np.sum(result, 2) # 累加
plt.subplot(row, col, j + 1)
title = 'class_id:' + str(class_id[j]) + ' ' + str(score[j])
plt.title(title, color='red', fontsize='large', fontweight='bold')
plt.imshow(result)
axis('off')
# 画框处理,从box中读取位置信息
currentAxis = plt.gca()
x_min = int(box[j][0] * 138)
y_min = int(box[j][1] * 138)
x_max = int(box[j][2] * 138)
y_max = int(box[j][3] * 138)
rect = patches.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, linewidth=1, edgecolor='r',
facecolor='none')
currentAxis.add_patch(rect)
plt.show()