参考了这个:YOLOv8只检测人(或其他一种或者多种类别)_yolov8只检测指定类别-CSDN博客
1. 只检测人:predict的时候指定参数classes=[0]
2. 只画框不要标签:plot的时候传入labels=False
3. 标签中去掉置信度:result.plot(conf=False)
示例代码如下:
#!/usr/bin/env python
import contextlib
from datetime import datetime
from pathlib import Path
# pip install opencv-python pillow ultralytics rich
import cv2
from ultralytics import YOLO # type:ignore[import-untyped]
with contextlib.suppress(ImportError):
from rich import print
model = YOLO("yolov8s.pt")
person_cls = 0 # 人员标签的类别ID
total = 200 # 总共截多少桢
interval = 20 # 每隔几桢推理一次
folder = Path(__file__).parent / "images" # 未检测到人的图片存在这里
folder.mkdir(exist_ok=True)
cap = cv2.VideoCapture(0) # 打开电脑摄像头
count = 0
while True:
if not cap.isOpened():
print(f"Failed to open video capture with {cap=}")
break
ok, frame = cap.read()
if not ok:
continue
count += 1
if count > total:
print(f"Success to capture {total} frames~")
break
elif count % interval != 0:
continue
# 通过指定classes参数,限定要检测的类别
# result = model(frame, classes=[person_cls], save=False, show=False, verbose=False)[0]
result = model(frame, classes=[person_cls], save=False, verbose=False)[0]
has_person = bool(result.summary())
if has_person:
print(f"{count=}, result: {result.verbose()}")
result.show()
else:
file = folder / f"{count}.jpg"
content = cv2.imencode(".jpg", result.plot())[1].tobytes()
file.write_bytes(content)
print(datetime.now(), f"{count=}, result: {has_person}, save to {file}")
cap.release()
print("✨ Done.")