即插即用 | YOLOv8热力图可视化方法详解,揭秘AI如何「看」世界!【附完整源码】

简介: 即插即用 | YOLOv8热力图可视化方法详解,揭秘AI如何「看」世界!【附完整源码】

效果展示

先展示一下热力图显示效果:

引言

计算机视觉中,热力图(Heatmap)经常被用作一种可视化工具,用于表现网络模型的关注点或预测结果。在一些任务中,例如目标检测,关键点检测或语义分割等,模型可以通过生成一个“热力图”来表达不同位置的预测概率。这个“热力图”通常与原始输入图像的大小相同,或者稍微小一些,每一个位置(像素)的值都代表了该位置匹配特定类别或者特性的预测概率。然后,热力图还可以用于生成预测框,或者定位关键点等。由于其形式直观,便于理解,因此热力图也经常用来帮助理解和分析模型的预测逻辑和行为。

本文将介绍如何使用YOLOv8模型结合Grad-CAM(梯度加权类激活映射)技术生成YOLOv8不同网络层中的图像热力图。这是一种基于梯度的可视化方法,通过计算特征图的梯度来生成热力图。这种方法可以帮助您更直观地理解模型在不同区域的关注程度,从而提高对模型的理解和应用能力。

安装环境

在进行热力图计算前,我们需要先安装好需要的环境ultralyticsgrad-cam,命令如下:

pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install grad-cam -i https://pypi.tuna.tsinghua.edu.cn/simple

完整源码

新建一个YOLOv8HeatMap.py文件,写入如下代码。然后放入YOLOv8源码中,如下图所示:

import warnings
warnings.filterwarnings('ignore')
warnings.simplefilter('ignore')
import torch, yaml, cv2, os, shutil, sys
import numpy as np
np.random.seed(0)
import matplotlib.pyplot as plt
from tqdm import trange
from PIL import Image
from ultralytics.nn.tasks import attempt_load_weights
from ultralytics.utils.torch_utils import intersect_dicts
from ultralytics.utils.ops import xywh2xyxy, non_max_suppression
from pytorch_grad_cam import GradCAMPlusPlus, GradCAM, XGradCAM, EigenCAM, HiResCAM, LayerCAM, RandomCAM, EigenGradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)

class ActivationsAndGradients:
    """ Class for extracting activations and
    registering gradients from targetted intermediate layers """

    def __init__(self, model, target_layers, reshape_transform):
        self.model = model
        self.gradients = []
        self.activations = []
        self.reshape_transform = reshape_transform
        self.handles = []
        for target_layer in target_layers:
            self.handles.append(
                target_layer.register_forward_hook(self.save_activation))
            # Because of https://github.com/pytorch/pytorch/issues/61519,
            # we don't use backward hook to record gradients.
            self.handles.append(
                target_layer.register_forward_hook(self.save_gradient))

    def save_activation(self, module, input, output):
        activation = output

        if self.reshape_transform is not None:
            activation = self.reshape_transform(activation)
        self.activations.append(activation.cpu().detach())

    def save_gradient(self, module, input, output):
        if not hasattr(output, "requires_grad") or not output.requires_grad:
            # You can only register hooks on tensor requires grad.
            return

        # Gradients are computed in reverse order
        def _store_grad(grad):
            if self.reshape_transform is not None:
                grad = self.reshape_transform(grad)
            self.gradients = [grad.cpu().detach()] + self.gradients

        output.register_hook(_store_grad)

    def post_process(self, result):
        logits_ = result[:, 4:]
        boxes_ = result[:, :4]
        sorted, indices = torch.sort(logits_.max(1)[0], descending=True)
        return torch.transpose(logits_[0], dim0=0, dim1=1)[indices[0]], torch.transpose(boxes_[0], dim0=0, dim1=1)[indices[0]], xywh2xyxy(torch.transpose(boxes_[0], dim0=0, dim1=1)[indices[0]]).cpu().detach().numpy()
  
    def __call__(self, x):
        self.gradients = []
        self.activations = []
        model_output = self.model(x)
        post_result, pre_post_boxes, post_boxes = self.post_process(model_output[0])
        return [[post_result, pre_post_boxes]]

    def release(self):
        for handle in self.handles:
            handle.remove()

class yolov8_target(torch.nn.Module):
    def __init__(self, ouput_type, conf, ratio) -> None:
        super().__init__()
        self.ouput_type = ouput_type
        self.conf = conf
        self.ratio = ratio
    
    def forward(self, data):
        post_result, pre_post_boxes = data
        result = []
        for i in trange(int(post_result.size(0) * self.ratio)):
            if float(post_result[i].max()) < self.conf:
                break
            if self.ouput_type == 'class' or self.ouput_type == 'all':
                result.append(post_result[i].max())
            elif self.ouput_type == 'box' or self.ouput_type == 'all':
                for j in range(4):
                    result.append(pre_post_boxes[i, j])
        return sum(result)

class yolov8_heatmap:
    def __init__(self, weight, device, method, layer, backward_type, conf_threshold, ratio, show_box, renormalize):
        device = torch.device(device)
        ckpt = torch.load(weight)
        model_names = ckpt['model'].names
        model = attempt_load_weights(weight, device)
        model.info()
        for p in model.parameters():
            p.requires_grad_(True)
        model.eval()
        
        target = yolov8_target(backward_type, conf_threshold, ratio)
        target_layers = [model.model[l] for l in layer]
        method = eval(method)(model, target_layers, use_cuda=device.type == 'cuda')
        method.activations_and_grads = ActivationsAndGradients(model, target_layers, None)
        
        colors = np.random.uniform(0, 255, size=(len(model_names), 3)).astype(np.uint8)
        self.__dict__.update(locals())
    
    def post_process(self, result):
        result = non_max_suppression(result, conf_thres=self.conf_threshold, iou_thres=0.65)[0]
        return result

    def draw_detections(self, box, color, name, img):
        xmin, ymin, xmax, ymax = list(map(int, list(box)))
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), tuple(int(x) for x in color), 2)
        cv2.putText(img, str(name), (xmin, ymin - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.8, tuple(int(x) for x in color), 2, lineType=cv2.LINE_AA)
        return img

    def renormalize_cam_in_bounding_boxes(self, boxes, image_float_np, grayscale_cam):
        """Normalize the CAM to be in the range [0, 1] 
        inside every bounding boxes, and zero outside of the bounding boxes. """
        renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
        for x1, y1, x2, y2 in boxes:
            x1, y1 = max(x1, 0), max(y1, 0)
            x2, y2 = min(grayscale_cam.shape[1] - 1, x2), min(grayscale_cam.shape[0] - 1, y2)
            renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())    
        renormalized_cam = scale_cam_image(renormalized_cam)
        eigencam_image_renormalized = show_cam_on_image(image_float_np, renormalized_cam, use_rgb=True)
        return eigencam_image_renormalized
    
    def process(self, img_path, save_path):
        # img process
        img = cv2.imread(img_path)
        img = letterbox(img)[0]
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.float32(img) / 255.0
        tensor = torch.from_numpy(np.transpose(img, axes=[2, 0, 1])).unsqueeze(0).to(self.device)
        
        try:
            grayscale_cam = self.method(tensor, [self.target])
        except AttributeError as e:
            return
        
        grayscale_cam = grayscale_cam[0, :]
        cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True)
        
        pred = self.model(tensor)[0]
        pred = self.post_process(pred)
        if self.renormalize:
            cam_image = self.renormalize_cam_in_bounding_boxes(pred[:, :4].cpu().detach().numpy().astype(np.int32), img, grayscale_cam)
        if self.show_box:
            for data in pred:
                data = data.cpu().detach().numpy()
                cam_image = self.draw_detections(data[:4], self.colors[int(data[4:].argmax())], f'{self.model_names[int(data[4:].argmax())]} {float(data[4:].max()):.2f}', cam_image)
        
        cam_image = Image.fromarray(cam_image)
        cam_image.save(save_path)
    
    def __call__(self, img_path, save_path, grad_name):
        # remove dir if exist
        # if os.path.exists(save_path):
        #     shutil.rmtree(save_path)
        # make dir if not exist
        if not os.path.exists(save_path):
            os.makedirs(save_path, exist_ok=True)

        if os.path.isdir(img_path):
            for img_path_ in os.listdir(img_path):
                name = img_path_.rsplit('.')[0]
                end_name = img_path_.rsplit('.')[-1]
                self.process(f'{img_path}/{img_path_}', f'{save_path}/{name}_{grad_name}.{end_name}')
        else:
            self.process(img_path, f'{save_path}/result_{grad_name}.png')
        
def get_params():
    # 绘制热力图方法列表
    grad_list = [
        'GradCAM',
        'GradCAMPlusPlus',
        'XGradCAM',
        'EigenCAM',
        'HiResCAM',
        'LayerCAM',
        'RandomCAM',
        'EigenGradCAM'
    ]
    # 自定义需要绘制热力图的层索引,可以用列表绘制不同层的热力图,如[10, 12, 14, 16, 18],将多层的话会将结果进行汇总到一张图上
    layers = [10, 12, 14, 16, 18]
    for grad_name in grad_list:
        params = {
            'weight': 'yolov8n.pt', # 训练好的权重路径
            'device': 'cuda:0',  # cpu或者cuda:0
            'method': grad_name, # GradCAMPlusPlus, GradCAM, XGradCAM, EigenCAM, HiResCAM, LayerCAM, RandomCAM, EigenGradCAM
            'layer': layers,  # 计算梯度的层, 指定层的索引
            'backward_type': 'class', # class, box, all
            'conf_threshold': 0.2, # 置信度阈值默认0.2, 根据情况调节
            'ratio': 0.02, # 建议0.02-0.1,取前多少数据,默认是0.02,只取置信度排序后的前百分之2的目标进行计算热力图。
            'show_box': False,  #是否显示检测框
            'renormalize': True   #是否优化热力图显示结果
        }
        yield params

if __name__ == '__main__':
    for each in get_params():
        model = yolov8_heatmap(**each)
        # model第一个参数:单张图片路径或者图片文件夹路径; 第二个参数:保存路径; 第三个参数:绘制热力图方法
        # model(r'images/00052.jpg', 'result', each['method'])
        model(r'images', 'result', each['method'])

源码详解

上述代码实现了一个用于生成YOLOv8网络热力图的工具,它使用了pytorch_grad_cam库来计算梯度激活图(GradCAM)。主要包含以下几个部分:

导入需要的库

import warnings
warnings.filterwarnings('ignore')
warnings.simplefilter('ignore')
import torch, yaml, cv2, os, shutil, sys
import numpy as np
np.random.seed(0)
import matplotlib.pyplot as plt
from tqdm import trange
from PIL import Image
from ultralytics.nn.tasks import attempt_load_weights
from ultralytics.utils.torch_utils import intersect_dicts
from ultralytics.utils.ops import xywh2xyxy, non_max_suppression
from pytorch_grad_cam import GradCAMPlusPlus, GradCAM, XGradCAM, EigenCAM, HiResCAM, LayerCAM, RandomCAM, EigenGradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients

letterbox 函数

对输入图像进行预处理,包括缩放和填充,以适应特定的尺寸要求,同时保持宽高比。

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
• 1
  • im: 输入图像。
  • new_shape: 目标图像的新尺寸,默认为 640x640。
  • color: 用于边缘填充的颜色,默认为灰色(114, 114, 114)。
  • auto: 是否自动调整填充大小以满足步长约束。
  • scaleFill: 是否拉伸图像以填充新形状。
  • scaleup: 是否允许放大图像。
  • stride: 模型步长,用于确保输出尺寸是该数值的倍数。

ActivationsAndGradients 类

一个辅助类,用于从模型的中间层提取激活和梯度。它注册了前向传播钩子以在每个目标层捕获这些信息。

yolov8_target 类

根据给定的输出类型(‘class’、‘box’ 或 ‘all’)和置信度阈值,从模型输出中选择感兴趣的检测结果。

yolov8_heatmap 类

class yolov8_heatmap:
    def __init__(self, weight, cfg, device, method, layer, backward_type, conf_threshold, ratio):

初始化时加载模型、设置热力图计算方法、指定要计算梯度的层,以及其它相关参数。

process 方法:处理单个图像或图像文件夹,生成热力图并保存结果。

draw_detections 方法:在图像上绘制检测框。

renormalize_cam_in_bounding_boxes 方法:将热力图归一化到边界框内。

get_params 函数:生成参数字典,用于遍历不同的热力图方法。

get_params函数详解

def get_params():
    # 绘制热力图方法列表
    grad_list = [
        'GradCAM',
        'GradCAMPlusPlus',
        'XGradCAM',
        'EigenCAM',
        'HiResCAM',
        'LayerCAM',
        'RandomCAM',
        'EigenGradCAM'
    ]
    # 自定义需要绘制热力图的层索引,可以用列表绘制不同层的热力图,如[10, 12, 14, 16, 18],将多层的话会将结果进行汇总到一张图上
    layers = [10, 12, 14, 16, 18]
    for grad_name in grad_list:
        params = {
            'weight': 'yolov8n.pt', # 训练好的权重路径
            'device': 'cuda:0',  # cpu或者cuda:0
            'method': grad_name, # GradCAMPlusPlus, GradCAM, XGradCAM, EigenCAM, HiResCAM, LayerCAM, RandomCAM, EigenGradCAM
            'layer': layers,  # 计算梯度的层, 指定层的索引
            'backward_type': 'class', # class, box, all
            'conf_threshold': 0.2, # 置信度阈值默认0.2, 根据情况调节
            'ratio': 0.02, # 建议0.02-0.1,取前多少数据,默认是0.02,只取置信度排序后的前百分之2的目标进行计算热力图。
            'show_box': False,  #是否显示检测框
            'renormalize': True   #是否优化热力图显示结果
        }
        yield params

这段代码定义了一个名为 get_params 的函数,它用于设置并返回一个字典格式的生成器,包含了一系列参数,这些参数是用于配置和运行yolov8_heatmap 类的。下面我将逐个解释这些参数的含义和作用。

参数说明:

weight
  • 'rtdetr-l.pt': 这个参数指定了模型的权重文件。在这种情况下,它是一个预先训练的YOLOv8模型的权重文件。这个文件包含了模型的所有训练参数,是模型运行的基础。
device
  • 'cuda:0': 这个参数指定了模型运行的设备。在这里,'cuda:0' 表明模型将在第一个NVIDIA GPU上运行。如果没有GPU或希望在CPU上运行,可以将其更改为 'cpu'
method
  • 有8种热力图生成算法可选:GradCAMPlusPlus, GradCAM, XGradCAM, EigenCAM, HiResCAM, LayerCAM, RandomCAM, EigenGradCAM
layer
  • '[10]': 这个参数指定了用于生成热图的网络层。可以自行选择单一层,如[10]或者多个层列表[10, 12, 14, 16, 18],最终会对多层结果汇总到一张热力图上。
backward_type
  • 'all': 这个参数决定了反向传播的类型。它可以是 'class''box''all''class' 仅关注类别预测,'box' 仅关注边界框预测,而 'all' 结合了两者。
conf_threshold
  • 0.2: 这是一个置信度阈值,用于过滤模型的预测。仅当模型对其预测的置信度高于0.2时,这些预测才会被考虑。
ratio

取前多少数据,默认是0.02,就是只取置信度由大到小排序后的前百分之2的目标进行计算热力图。

这个可能比较难理解,一般0.02就可以了,这个值不是越大越好,最大建议是0.1。

主函数部分

创建yolov8_heatmap实例,并使用get_params生成的参数遍历不同的热力图方法,为指定的图像或图像文件夹生成热力图。

if __name__ == '__main__':
    for each in get_params():
        model = yolov8_heatmap(**each)
        # model第一个参数:单张图片路径或者图片文件夹路径; 第二个参数:保存路径; 第三个参数:绘制热力图方法
        # model(r'images/00052.jpg', 'result', each['method'])
        model(r'images', 'result', each['method'])

参数说明

第一个参数:单张图片路径或者图片文件夹路径;

第二个参数:‘result’,为保存路径;

第三个参数:绘制热力图方法。

通过循环可一次性保存8种热力图计算结果。可自行挑选结果较好的热力图计算方式进行展示,图片保存的命名规则为图片名+热力图计算方法名称.jpg。保存的结果如下:

整体而言,这个工具可以用来分析YOLOv8模型的特征可视化,特别是在目标检测任务中,帮助理解模型是如何关注图像的特定区域来做出预测的,可以增加论文的结果展示对比,丰富论文不同模型的对比角度,还是十分实用的。

目录
打赏
0
0
0
0
115
分享
相关文章
写歌词的技巧和方法:塑造完美歌词结构的艺术,妙笔生词AI智能写歌词软件
歌词是音乐的灵魂,其结构艺术至关重要。开头需引人入胜,主体部分无论是叙事还是抒情,都应层次分明、情感丰富,结尾则需升华或留白,给人以深刻印象。《妙笔生词智能写歌词软件》提供多种AI辅助功能,助你轻松创作完美歌词,成为音乐创作的得力助手。
MILS:无需对LLM进行额外训练就能处理多模态任务,Meta AI提出零样本生成多模态描述方法
MILS 是 Meta AI 推出的零样本生成高质量多模态描述方法,支持图像、视频和音频的描述生成,无需额外训练。
109 34
MILS:无需对LLM进行额外训练就能处理多模态任务,Meta AI提出零样本生成多模态描述方法
|
16天前
|
AI对话网站一键生成系统源码
可以添加进自己的工具箱,也可以嵌入自己博客的页面中,引流效果杠杠的,新拟态设计风格,有能力的大佬可以进行二开,仅提供学习,用户可输入网站名称、AI默认的开场白、AI头像昵称、AI网站中引流的你的网站等等内容,所有生成的网页全部保存到你的服务器上
57 27
AI对话网站一键生成系统源码
AI大模型发展对语音直播交友系统源码开发搭建的影响
近年来,AI大模型技术的迅猛发展深刻影响了语音直播交友系统的开发与应用。本文探讨了AI大模型如何提升语音交互的自然流畅性、内容审核的精准度、个性化推荐的智能性以及虚拟主播的表现力,并分析其对开发流程和用户体验的变革。同时,展望了多模态交互、情感陪伴及元宇宙社交等未来发展方向,指出在把握机遇的同时需应对数据安全、算法偏见等挑战,以实现更智能、安全、有趣的语音直播交友平台。
TimesFM 2.0:用 AI 预测流量、销量和金融市场等走势!谷歌开源超越统计方法的预测模型
TimesFM 2.0 是谷歌研究团队开源的时间序列预测模型,支持长达2048个时间点的单变量预测,具备零样本学习能力,适用于零售、金融、交通等多个领域。
234 23
TimesFM 2.0:用 AI 预测流量、销量和金融市场等走势!谷歌开源超越统计方法的预测模型
使用tree命令把自己的代码归类文件目录的方法-优雅草央千澈以优雅草AI智能功能为例给大家展示tree命令实际用法
使用tree命令把自己的代码归类文件目录的方法-优雅草央千澈以优雅草AI智能功能为例给大家展示tree命令实际用法
115 29
使用tree命令把自己的代码归类文件目录的方法-优雅草央千澈以优雅草AI智能功能为例给大家展示tree命令实际用法
Midscene.js:AI 驱动的 UI 自动化测试框架,支持自然语言交互,生成可视化报告
Midscene.js 是一款基于 AI 技术的 UI 自动化测试框架,通过自然语言交互简化测试流程,支持动作执行、数据查询和页面断言,提供可视化报告,适用于多种应用场景。
614 1
Midscene.js:AI 驱动的 UI 自动化测试框架,支持自然语言交互,生成可视化报告
新手指南:人工智能poe ai 怎么用?国内使用poe记住这个方法就够了!
由于国内网络限制,许多用户在尝试访问Poe AI时面临障碍。幸运的是,现在国内用户也能轻松畅玩Poe AI,告别繁琐的设置,直接开启AI创作之旅!🎉
288 13
已解决:Poe AI国内能用吗?国内用户如何使用Poe AI?亲测有效的方法来了!
人工智能正在重塑我们的世界,而Poe AI作为AI聊天机器人平台的佼者,更是引领着这场变革。它集成了众多顶尖AI模型,如OpenAI的GPT系列、Anthropic的Claude系列以及Google的PaLM等,为用户提供了一个探索AI无限可能的开放平台
148 12
AI Agents Loop异步执行可视化Tutorial 借助AgentBoard工具可视化工作流
本文介绍了AI Agent的异步执行循环(Agent Loop),并展示了如何利用开源框架agentboard可视化这一过程。通过分析不同框架(如AutoGen、LangGraph、AutoAgent)对Agent Loop的抽象,文章详细说明了从简单的功能调用到复杂的多阶段执行流程的设计。此外,还提供了使用agentboard进行日志记录与流程可视化的具体示例,包括安装步骤、代码实现及运行方法,帮助开发者更高效地调试和优化AI Agent的应用。
AI Agents Loop异步执行可视化Tutorial 借助AgentBoard工具可视化工作流

热门文章

最新文章