计算机视觉：使用 YOLOv8 创建交通热力图

在一个由数据驱动和导向的世界中，解释、可视化并基于这些数据做出决策的能力变得越来越重要。这意味着，应用正确的工具和技术可能决定一个项目的成败。在计算机视觉领域，有许多技术可以解释从视频（录制、流媒体或实时）中获取的数据。在特定情况下，例如评估交通强度或某些对象（如人、车辆、动物等）的行为区域时，热力图成为一种非常有效的工具。

热力图基于颜色，通常暖色（如红色、橙色和黄色）表示交通流量或密度较大的区域，而冷色（如蓝色和绿色）表示交通流量或密度较低的区域。换句话说，这种技术可以收集大量数据并将其转换为易于理解和直观的可视化效果。如下例所示：

考虑到实际应用，我们可以列举一些场景，例如需要理解交通拥堵中的运动模式，或理解百货商店或购物中心走廊中的消费和运动模式。此外，还可以用于分析牧场中动物的运动模式，尤其是在畜牧业中。

我们以一个包含3条道路的空中视频为例，其中有一些汽车在行驶。我们的目标是创建一条热力图，显示道路上交通最密集的点：

视频链接：https://youtu.be/MNn9qKG2UFI?si=PQWXhDY48dis-Vxs

代码实现

首先，我们需要导入所需的库：

from collections import defaultdict
import cv2
import numpy as np
from ultralytics import YOLO

然后，我们声明一个变量来加载训练好的YOLO模型。在我的例子中，我使用了一个“中等”模型。

model = YOLO('yolov8m.pt')

接下来，我们需要指定要分析的视频路径。

videopath = '/Path/Road traffic video.mp4'
cap = cv2.VideoCapture(videopath)

现在，我们创建一个空字典来存储跟踪位置（track_history），以及另一个字典来存储每个对象的最后推断位置（last_positions）。

track_history = defaultdict(lambda: [])
last_positions = {}

由于我们只考虑移动的物体（例如，停放的汽车可能会干扰分析），我们创建了一个函数来计算两点之间的欧几里得距离（p1和p2，分别是帧A和帧B中的物体点）。该距离使用勾股定理计算，返回两点之间的直线距离。

def calculate_distance(p1, p2):
    return np.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)

当然，我们还需要初始化一个热力图。使用numpy创建一个矩阵，其中所有元素最初都为零。该矩阵还需要包含3个“层”，分别代表3个颜色通道（RGB）。

heatmap = np.zeros((int(cap.get(4)), int(cap.get(3)), 3), dtype=np.float32)

下一步是创建一个while循环，逐帧处理视频。

while cap.isOpened():
    success, frame = cap.read()

如果成功读取帧，则使用YOLO模型进行推理。请注意，我添加了跟踪和持久性算法，这些算法在处理视频（帧序列）时非常理想。此外，由于本例的重点是记录汽车交通，我们只定义了类别2。

if success:
   results = model.track(frame, persist=True, classes=2)

如果有有效的检测结果，我们需要将其记录到历史记录中，同时获取边界框的坐标，并使用以下代码更新热力图：

for box, track_id in zip(boxes, track_ids):
    x_center, y_center, width, height = box
    current_position = (float(x_center), float(y_center))

    top_left_x = int(x_center - width / 2)
    top_left_y = int(y_center - height / 2)
    bottom_right_x = int(x_center + width / 2)
    bottom_right_y = int(y_center + height / 2)

    top_left_x = max(0, top_left_x)
    top_left_y = max(0, top_left_y)
    bottom_right_x = min(heatmap.shape[1], bottom_right_x)
    bottom_right_y = min(heatmap.shape[0], bottom_right_y)

    track = track_history[track_id]
    track.append(current_position)
    if len(track) > 1200:
        track.pop(0)

为了检查物体是否在移动，我们使用calculate_distance函数计算最后一个记录点与当前点之间的距离。如果距离达到最小值，则将其记录到热力图中。

    last_position = last_positions.get(track_id)
    if last_position and calculate_distance(last_position, current_position) > 5:
        heatmap[top_left_y:bottom_right_y, top_left_x:bottom_right_x] += 1
  
    last_positions[track_id] = current_position

随着移动物体的跟踪点被记录，出于美观原因，我们在热力图上添加了高斯模糊滤镜。

 heatmap_blurred = cv2.GaussianBlur(heatmap, (15, 15), 0)

然后，我们对热力图进行归一化处理，并在原始视频帧上叠加颜色，以使其更直观。

  heatmap_norm = cv2.normalize(heatmap_blurred, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
  heatmap_color = cv2.applyColorMap(heatmap_norm, cv2.COLORMAP_JET)

  alpha = 0.7
  overlay = cv2.addWeighted(frame, 1 - alpha, heatmap_color, alpha, 0)

  cv2.imshow("Traffic Heatmap", overlay)

最后，我们为while循环添加一个退出键，并在视频结束时使用cap.release()和cv2.destroyAllWindows()关闭循环。

    if cv2.waitKey(1) & 0xFF == ord("q"):
        break
    else:
        break

cap.release()
cv2.destroyAllWindows()

结果

运行代码后，我们会看到一个窗口打开，屏幕上应用了蓝色背景（热力图滤镜）。当YOLO模型检测到汽车时，滤镜值会根据汽车的强度和数量进行更新，并用更热的颜色突出显示。