計算機視覺:使用 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模型檢測到汽車時,濾鏡值會根據汽車的強度和數量進行更新,并用更熱的顏色突出顯示。