yolov5热力图代码

YOLOv5是一种目标检测算法，它是YOLO（You Only Look Once）系列算法的最新版本。YOLOv5的热力图代码可以用于生成目标检测结果的热力图，以可视化目标在图像中的分布情况。

以下是一个简单的YOLOv5热力图代码示例：

import torch
from PIL import Image
from torchvision.transforms import functional as F
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box

# 加载YOLOv5模型
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))

# 加载图像
image = Image.open('image.jpg')

# 图像预处理
img = F.to_tensor(image)
img = img.unsqueeze(0)

# 模型推理
pred = model(img)

# 非最大抑制
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)

# 绘制热力图
heatmap = torch.zeros_like(img.squeeze(0))
for det in pred:
    x1, y1, x2, y2, conf, cls = det
    plot_one_box((x1, y1, x2, y2), heatmap, label=int(cls))

# 可视化热力图
heatmap_image = Image.fromarray(heatmap.numpy())
heatmap_image.show()

这段代码首先加载了YOLOv5模型，然后加载了一张图像。接下来，对图像进行预处理，并通过模型进行推理得到目标检测结果。然后，使用非最大抑制方法对结果进行筛选，得到最终的目标框。最后，根据目标框的位置信息，在热力图上绘制相应的热力点。最终，将生成的热力图可视化展示出来。