YOLOv5怎么實現在圖片上顯示統計出單一檢測目標的個數

YOLOv5怎么實現在圖片上顯示統計出單一檢測目標的個數

YOLOv5是一種流行的目標檢測算法，廣泛應用于計算機視覺任務中。它能夠快速準確地檢測出圖像中的多個目標，并為每個目標生成邊界框和類別標簽。然而，在某些應用場景中，我們不僅需要檢測目標，還需要統計出圖像中某一特定類別的目標數量，并在圖像上顯示出來。本文將詳細介紹如何使用YOLOv5實現這一功能。

1. 環境準備

在開始之前，確保你已經安裝了YOLOv5所需的依賴環境。你可以通過以下命令克隆YOLOv5的官方倉庫并安裝依賴：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

2. 加載預訓練模型

YOLOv5提供了多種預訓練模型，你可以根據需要選擇合適的模型。以下代碼展示了如何加載一個預訓練的YOLOv5模型：

import torch

# 加載預訓練模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # 使用YOLOv5s模型

3. 進行目標檢測

接下來，我們可以使用加載的模型對圖像進行目標檢測。以下代碼展示了如何對一張圖片進行檢測，并獲取檢測結果：

from PIL import Image

# 加載圖片
img_path = 'path_to_your_image.jpg'
img = Image.open(img_path)

# 進行目標檢測
results = model(img)

# 獲取檢測結果
detections = results.pandas().xyxy[0]  # 獲取檢測結果的DataFrame

4. 統計單一目標數量

假設我們需要統計圖像中某一特定類別的目標數量，例如“person”。我們可以通過以下代碼實現：

# 統計特定類別的目標數量
target_class = 'person'
target_count = len(detections[detections['name'] == target_class])

print(f'Number of {target_class}: {target_count}')

5. 在圖片上顯示統計結果

為了在圖片上顯示統計結果，我們可以使用OpenCV庫來繪制文本。以下代碼展示了如何在圖片上繪制統計結果：

import cv2
import numpy as np

# 將PIL圖像轉換為OpenCV格式
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

# 在圖片上繪制統計結果
text = f'Number of {target_class}: {target_count}'
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 1
font_color = (255, 0, 0)  # 紅色
thickness = 2

# 獲取文本大小
(text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)

# 計算文本位置
text_x = 10
text_y = 30

# 繪制文本
cv2.putText(img_cv, text, (text_x, text_y), font, font_scale, font_color, thickness)

# 顯示圖片
cv2.imshow('Image with Count', img_cv)
cv2.waitKey(0)
cv2.destroyAllWindows()

6. 保存結果

如果你希望將結果保存為新的圖片文件，可以使用以下代碼：

# 保存結果
output_path = 'output_image.jpg'
cv2.imwrite(output_path, img_cv)

7. 完整代碼示例

以下是完整的代碼示例，展示了如何加載模型、檢測目標、統計數量并在圖片上顯示結果：

import torch
from PIL import Image
import cv2
import numpy as np

# 加載預訓練模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 加載圖片
img_path = 'path_to_your_image.jpg'
img = Image.open(img_path)

# 進行目標檢測
results = model(img)

# 獲取檢測結果
detections = results.pandas().xyxy[0]

# 統計特定類別的目標數量
target_class = 'person'
target_count = len(detections[detections['name'] == target_class])

# 將PIL圖像轉換為OpenCV格式
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

# 在圖片上繪制統計結果
text = f'Number of {target_class}: {target_count}'
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 1
font_color = (255, 0, 0)  # 紅色
thickness = 2

# 獲取文本大小
(text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)

# 計算文本位置
text_x = 10
text_y = 30

# 繪制文本
cv2.putText(img_cv, text, (text_x, text_y), font, font_scale, font_color, thickness)

# 顯示圖片
cv2.imshow('Image with Count', img_cv)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 保存結果
output_path = 'output_image.jpg'
cv2.imwrite(output_path, img_cv)

8. 總結

通過以上步驟，我們成功地使用YOLOv5實現了在圖片上顯示統計出單一檢測目標的個數。這一功能在許多實際應用中非常有用，例如監控系統中的行人計數、交通流量統計等。希望本文對你有所幫助，祝你在使用YOLOv5進行目標檢測時取得更多成果！