国家企业信用网查询系统杭州网站建设优化

目标检测是计算机视觉领域的重要研究方向，而YOLO（You Only Look Once）系列算法因其高效性和准确性成为该领域的代表性方法。YOLOv8作为YOLO系列的最新版本，在目标检测任务中表现出色。然而，传统的目标检测结果通常以边界框和类别标签的形式呈现，缺乏对模型决策过程的直观解释。热力图作为一种可视化工具，能够直观地展示模型在图像中的关注区域，为模型的可解释性提供了重要支持。本文将探讨基于YOLOv8的热力图生成与可视化方法，并支持自定义模型与置信度阈值的多维度分析。

1. YOLOv8与热力图生成技术概述

1.1 YOLOv8的核心特性

YOLOv8在YOLO系列的基础上进行了多项优化，包括：

1. 更高效的网络结构：采用CSP（Cross Stage Partial）架构，提升特征提取能力。
2. 更精确的检测性能：通过改进的损失函数和训练策略，提高目标检测的准确性。
3. 更灵活的部署支持：支持多种硬件平台和推理框架，便于实际应用。

1.2 热力图生成的基本原理

热力图是一种通过颜色梯度表示数据分布的可视化方法。在目标检测中，热力图通常用于展示模型对图像中不同区域的关注程度。生成热力图的核心步骤包括：

1. 特征提取：从YOLOv8的中间层提取特征图。
2. 权重计算：根据模型的输出计算每个像素的权重。
3. 可视化映射：将权重映射为颜色梯度，生成热力图。
4. 后处理：对热力图进行归一化和增强处理，确保其清晰直观。

1.3 实现细节

1. 模型加载与初始化

   • 权重加载 ：通过torch.load加载预训练权重文件（weight），并根据配置文件（cfg）初始化YOLOv8模型。
   • 设备适配 ：支持GPU（如cuda:0）和CPU运行，用户可以根据硬件条件灵活选择。
   • 权重适配：使用intersect_dicts方法将预训练权重适配到当前模型结构中，确保兼容性。

2. 图像预处理

   • 图像缩放与填充 ：使用letterbox函数对输入图像进行缩放和填充，以满足YOLOv8的输入要求。
   • 归一化与转换 ：将图像转换为RGB格式并归一化到[0, 1]范围，便于后续处理。

3. 激活值与梯度提取

   • 激活值捕获 ：通过ActivationsAndGradients类捕获指定层（layer）的激活值。
   • 梯度计算 ：根据反向传播类型（backward_type）计算类别分数或边界框坐标的梯度。

4. 热力图生成与叠加

   • 梯度加权 ：根据激活值和梯度计算权重，生成加权热力图。
   • 归一化与增强 ：将热力图归一化到[0, 1]范围，并通过颜色映射增强其可读性。
   • 叠加显示 ：使用show_cam_on_image函数将热力图叠加到原图上，同时绘制检测结果（边界框和类别标签）。

5. 结果保存

   • 将生成的热力图保存为PNG文件，路径由save_path参数指定。每张热力图对应一个检测目标，便于用户逐一查看。

2. 项目源代码

import warnings
warnings.filterwarnings('ignore')
warnings.simplefilter('ignore')
import torch, yaml, cv2, os, shutil
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 DetectionModel as Model
from ultralytics.utils.torch_utils import intersect_dicts
from ultralytics.utils.ops import xywh2xyxy
from pytorch_grad_cam import GradCAMPlusPlus, GradCAM, XGradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradientsdef 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 constraintsshape = 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 paddingratio = r, r  # width, height ratiosnew_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 paddingif auto:  # minimum rectangledw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh paddingelif scaleFill:  # stretchdw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratiosdw /= 2  # divide padding into 2 sidesdh /= 2if shape[::-1] != new_unpad:  # resizeim = 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 borderreturn im, ratio, (dw, dh)class yolov8_heatmap:def __init__(self, weight, cfg, device, method, layer, backward_type, conf_threshold, ratio):device = torch.device(device)ckpt = torch.load(weight)model_names = ckpt['model'].namescsd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32model = Model(cfg, ch=3, nc=len(model_names)).to(device)csd = intersect_dicts(csd, model.state_dict(), exclude=['anchor'])  # intersectmodel.load_state_dict(csd, strict=False)  # loadmodel.eval()print(f'Transferred {len(csd)}/{len(model.state_dict())} items')target_layers = [eval(layer)]method = eval(method)colors = np.random.uniform(0, 255, size=(len(model_names), 3)).astype(np.int32)self.__dict__.update(locals())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 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 imgdef __call__(self, img_path, save_path):# remove dir if existif os.path.exists(save_path):shutil.rmtree(save_path)# make dir if not existos.makedirs(save_path, exist_ok=True)# img processimg = cv2.imread(img_path)img = letterbox(img)[0]img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)img = np.float32(img) / 255.0tensor = torch.from_numpy(np.transpose(img, axes=[2, 0, 1])).unsqueeze(0).to(self.device)# init ActivationsAndGradientsgrads = ActivationsAndGradients(self.model, self.target_layers, reshape_transform=None)# get ActivationsAndResultresult = grads(tensor)activations = grads.activations[0].cpu().detach().numpy()# postprocess to yolo outputpost_result, pre_post_boxes, post_boxes = self.post_process(result[0])for i in trange(int(post_result.size(0) * self.ratio)):if float(post_result[i].max()) < self.conf_threshold:breakself.model.zero_grad()# get max probability for this predictionif self.backward_type == 'class' or self.backward_type == 'all':score = post_result[i].max()score.backward(retain_graph=True)if self.backward_type == 'box' or self.backward_type == 'all':for j in range(4):score = pre_post_boxes[i, j]score.backward(retain_graph=True)# process heatmapif self.backward_type == 'class':gradients = grads.gradients[0]elif self.backward_type == 'box':gradients = grads.gradients[0] + grads.gradients[1] + grads.gradients[2] + grads.gradients[3]else:gradients = grads.gradients[0] + grads.gradients[1] + grads.gradients[2] + grads.gradients[3] + grads.gradients[4]b, k, u, v = gradients.size()weights = self.method.get_cam_weights(self.method, None, None, None, activations, gradients.detach().numpy())weights = weights.reshape((b, k, 1, 1))saliency_map = np.sum(weights * activations, axis=1)saliency_map = np.squeeze(np.maximum(saliency_map, 0))saliency_map = cv2.resize(saliency_map, (tensor.size(3), tensor.size(2)))saliency_map_min, saliency_map_max = saliency_map.min(), saliency_map.max()if (saliency_map_max - saliency_map_min) == 0:continuesaliency_map = (saliency_map - saliency_map_min) / (saliency_map_max - saliency_map_min)# add heatmap and box to imagecam_image = show_cam_on_image(img.copy(), saliency_map, use_rgb=True)"不想在图片中绘画出边界框和置信度，注释下面的一行代码即可"cam_image = self.draw_detections(post_boxes[i], self.colors[int(post_result[i, :].argmax())], f'{self.model_names[int(post_result[i, :].argmax())]} {float(post_result[i].max()):.2f}', cam_image)cam_image = Image.fromarray(cam_image)cam_image.save(f'{save_path}/{i}.png')def get_params():params = {'weight': 'yolov8n.pt',   # 训练出来的权重文件'cfg': 'ultralytics/cfg/models/v8/yolov8n.yaml',  # 训练权重对应的yaml配置文件'device': 'cuda:0','method': 'GradCAM', # GradCAMPlusPlus, GradCAM, XGradCAM , 使用的热力图库文件不同的效果不一样可以多尝试'layer': 'model.model[9]',  # 想要检测的对应层'backward_type': 'all', # class, box, all'conf_threshold': 0.01, # 0.6  # 置信度阈值，有的时候你的进度条到一半就停止了就是因为没有高于此值的了'ratio': 0.02 # 0.02-0.1}return paramsif __name__ == '__main__':model = yolov8_heatmap(**get_params())model(r'ultralytics/assets/10.jpg', 'result')  # 第一个是检测的文件, 第二个是保存的路径

参数名	参数类型	解释
weight	str	训练出来的权重文件路径，用于加载训练好的模型。
cfg	str	训练权重对应的YAML配置文件路径，定义了模型的结构和参数。
device	str	指定运行设备，例如 cuda:0 表示使用第一个GPU，cpu 表示使用CPU。
method	str	使用的热力图生成方法，例如 GradCAM、GradCAMPlusPlus、XGradCAM。
layer	str	想要检测的模型层，例如 model.model[11]，表示检测第11层。
backward_type	str	反向传播类型，可选 class（类别）、box（边界框）、all（全部）。
conf_threshold	float	置信度阈值，过滤低于该值的检测结果，默认值为 0.9。
ratio	float	热力图与原图的叠加比例，取值范围为 0.02-0.1，默认值为 0.3。