自用代码 | YOLOv5 特征图可视化代码

1. 在项目/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

2. 在/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

3. 在/home/xxx/code/yolov5-feature/train.py中，修改pred = model(imgs)为 pred = model(imgs, paths[0][-12:-4])。