YOLOv8 detect person only

  • demo.py
from io import BytesIO
from pathlib import Path

# pip install opencv-python pillow ultralytics
import cv2
from PIL import Image
from ultralytics import YOLO  # type:ignore[import-untyped]

filename = "test.jpg"
model = YOLO("yolov8s.pt")
person_cls = 0
img = Image.open(BytesIO(Path(filename).read_bytes()))
# 通过指定classes参数,限定要检测的类别
result = model(img, classes=[person_cls], save=False, show=False, verbose=False)[0]
has_person = bool(result.summary())
if has_person:
    print(f"file: {filename}, result: {has_person}")
    content = cv2.imencode(".jpg", result.plot())[1].tobytes()  # 检测结果转成jpg图片
    content2 = cv2.imencode(".jpg", result.plot(conf=False))[1].tobytes()  # 隐藏置信度
    content3 = cv2.imencode(".jpg", result.plot(labels=False))[1].tobytes()  # 隐藏标签
    Path("test_result.jpg").write_bytes(content)  # 保存到本地
    Path("test_result_without_conf.jpg").write_bytes(content2)
    Path("test_result_without_labels.jpg").write_bytes(content3)
    result.show()  # 打开画框了的图片
else:
    print(f"{filename=}, result: {result.verbose()}")

搭配摄像头识别人

  • camera.py
#!/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)
USE_THREAD = False


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.")
posted @ 2024-09-05 18:49  waketzheng  阅读(12)  评论(0编辑  收藏  举报