自用代码 | YOLOv5 特征图可视化代码
YOLOv5 特征图可视化代码
在yolov5的项目中，添加一些代码就可以实现特征图可视化。
- 在项目
/home/xxx/code/yolov5-feature/utils/utils.py中添加如下的代码：
outDir是你的特征图可视化目录
outDir = "/home/xxx/code/yolov5-feature/output/feature_new/"
def visualize_feature(features, imgid, i, num, sigmod=False):
outDir = "/home/xxx/code/yolov5-feature/output/feature_new/"
'''
features:[batch, 256, 128, 128]
'''
# 只取一张图的第一个通道
# [1, 1, 128, 128]
out_f_path = os.path.join(outDir, "feature_" + imgid + num + ".png")
out_i_path = os.path.join(outDir, "img_" + imgid + num + ".png")
feature = features[:1, 0, :, :]
# [128, 128]
feature = feature.view(feature.shape[1],feature.shape[2])
#to numpy (128, 128)
feature = feature.detach().cpu().numpy()
if sigmod:
#use sigmod to [0,1]
feature = 1.0/(1+np.exp(-1*feature))
# to [0,255]
feature = np.round(feature*255)
feature = feature.astype(np.uint8)
feature = cv2.applyColorMap(feature, cv2.COLORMAP_JET)
cv2.imwrite(out_f_path, feature)
feature = _gen_maskmap(feature)
if i == 21:
out_f_38_path = os.path.join(outDir, "feature_" + imgid + num + "38_" + ".png")
features = F.interpolate(features, size=[38, 38])
features_38 = F.interpolate(features, size=[76, 76])
feature_38 = features_38[:1, 0, :, :]
feature = features[:1, 0, :, :]
feature_38 = feature_38.view(feature_38.shape[1],feature_38.shape[2])
feature = feature.view(feature.shape[1],feature.shape[2])
feature_38 = feature_38.detach().cpu().numpy()
feature = feature.detach().cpu().numpy()
feature_38 = np.round(feature_38*255)
feature_38 = feature_38.astype(np.uint8)
feature_38 = cv2.applyColorMap(feature_38, cv2.COLORMAP_JET)
feature = np.round(feature*255)
feature = feature.astype(np.uint8)
feature = cv2.applyColorMap(feature, cv2.COLORMAP_JET)
cv2.imwrite(out_f_38_path, feature_38)
cv2.imwrite(out_f_path, feature)
feature_38 = _gen_maskmap(feature_38)
feature = _gen_maskmap(feature)
if i == 25:
out_f_19_path = os.path.join(outDir, "feature_" + imgid + num + "19_" + ".png")
features = F.interpolate(features, size=[19, 19])
features_19 = F.interpolate(features, size=[76, 76])
feature_19 = features_19[:1, 0, :, :]
feature = features[:1, 0, :, :]
feature_19 = feature_19.view(feature_19.shape[1],feature_19.shape[2])
feature = feature.view(feature.shape[1],feature.shape[2])
feature_19 = feature_19.detach().cpu().numpy()
feature = feature.detach().cpu().numpy()
feature_19 = np.round(feature_19*255)
feature_19 = feature_19.astype(np.uint8)
feature_19 = cv2.applyColorMap(feature_19, cv2.COLORMAP_JET)
feature = np.round(feature*255)
feature = feature.astype(np.uint8)
feature = cv2.applyColorMap(feature, cv2.COLORMAP_JET)
cv2.imwrite(out_f_19_path, feature_19)
cv2.imwrite(out_f_path, feature)
feature_19 = _gen_maskmap(feature_19)
feature = _gen_maskmap(feature)
# mean = np.array([0.40789654, 0.44719302, 0.47026115],
# dtype=np.float32).reshape(3, 1, 1)
# std = np.array([0.28863828, 0.27408164, 0.27809835],
# dtype=np.float32).reshape(3, 1, 1)
# image = image * 255
# image = (image * std + mean) * 255
# 再把图片transpose成标准的样子
# image = image.astype(np.uint8)
# outImg = _blend_img(image, feature)
# cv2.imwrite(out_i_path,outImg)
def _gen_maskmap(feature):
# feature[64, 64]
h, w = feature.shape[0], feature.shape[1]
color_map = np.zeros((h, w, 3), dtype=np.uint8)
for i in range(h):
for j in range(w):
if any(feature[i][j]>128):
color_map[i][j] = [255,255,255]
return color_map
def _blend_img(back, fore, trans=0.7):
'''
back = img-->[h*4, w*4, 3]
fore = tl_hm-->[h, w, 3]
'''
if fore.shape[0] != back.shape[0] or fore.shape[0] != back.shape[1]:
fore = cv2.resize(fore, (back.shape[1], back.shape[0]))
if len(fore.shape) == 2:
fore = fore.reshape(fore.shape[0], fore.shape[1], 1)
# 两幅图像进行合并时，按公式：blended_img = img1 * (1 – alpha) + img2* alpha 进行
ret = (back * (1. - trans) + fore * trans).astype(np.uint8)
# 别越界了,ret的大小就是原图的大小
ret[ret > 255] = 255
return ret
- 在
/home/xxx/code/yolov5-feature/models/yolo.py中修改并添加下面的代码：
class Model(nn.Module)
def forward(self, x, imgid="", augment=False, profile=False):
if augment:
img_size = x.shape[-2:] # height, width
s = [0.83, 0.67] # scales
y = []
for i, xi in enumerate((x,
torch_utils.scale_img(x.flip(3), s[0]), # flip-lr and scale
torch_utils.scale_img(x, s[1]), # scale
)):
# cv2.imwrite('img%g.jpg' % i, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1])
y.append(self.forward_once(xi)[0])
y[1][..., :4] /= s[0] # scale
y[1][..., 0] = img_size[1] - y[1][..., 0] # flip lr
y[2][..., :4] /= s[1] # scale
return torch.cat(y, 1), None # augmented inference, train
else:
return self.forward_once(x, imgid, profile) # single-scale inference, train
def forward_once(self, x, imgid="", profile=False):
y, dt = [], [] # outputs
for ii, m in enumerate(self.model):
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if profile:
try:
import thop
o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS
except:
o = 0
t = torch_utils.time_synchronized()
for _ in range(10):
_ = m(x)
dt.append((torch_utils.time_synchronized() - t) * 100)
print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
x = m(x) # run, 网络各个层(backbone + neck + output)
y.append(x if m.i in self.save else None) # save output
# by ljj
if m.i == 17: # [128,76,76]
visualize_feature(x, imgid, m.i, "_1_")
if m.i == 21: # [256,38,38]
visualize_feature(x, imgid, m.i, "_2_")
if m.i == 25: # [512,19,19]
visualize_feature(x, imgid, m.i, "_3_")
if profile:
print('%.1fms total' % sum(dt))
return x
- 在/home/xxx/code/yolov5-feature/train.py中，修改
pred = model(imgs)为pred = model(imgs, paths[0][-12:-4])。