假设自己的yolov8模型已经训练完成,名称为“D:\yolov8\runs\detect\train5\weights\best.pt”,请以此模型编写一个实时识别当前屏幕并打印出当前界面中存在的对象
时间: 2024-12-10 08:40:09 浏览: 15
在Python环境中,我们可以使用PyTorch及其YOLOv8库来实现这个功能。首先确保安装了相关的库,如torch、torchvision、opencv-python等。以下是一个简单的示例,说明如何加载预训练模型并在摄像头视频流上做实时检测:
```python
import torch
from models.common import Detect
from utils.datasets import LoadImagesAndLabels, LoadStreams, collate_fn
from utils.general import check_img_size, non_max_suppression, scale_coords
from PIL import Image
# 指定模型路径
model_path = "D:/yolov8/runs/detect/train5/weights/best.pt"
# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 使用GPU加速,如果不可用则用CPU
model = Detect(weights=model_path, device=device)
# 初始化摄像头
cam = cv2.VideoCapture(0) # 如果需要读取文件,则改为cv2.VideoCapture('video.mp4')
while True:
ret, frame = cam.read()
if not ret:
break
# 调整图像大小以便模型接受
img = Image.fromarray(frame)
img, _, _ = check_img_size(img, model.stride)
# 将PIL图像转为Tensor并进行预测
img = torch.from_numpy(img).unsqueeze(0).float().to(device)
pred = model(img)[0]
# 非极大抑制,获取最终的物体框
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.45)
# 打印出检测到的对象及它们的位置
for det in pred:
label, score, left, top, right, bottom = det
print(f"Detected object: {label.decode()}, Score: {score:.2f}, Coordinates: [{left:.0f}, {top:.0f}] to [{right:.0f}, {bottom:.0f}]")
# 显示原图以及标注
cv2.imshow("Object Detection", frame)
key = cv2.waitKey(1)
if key == ord('q'): # 按下'q'键退出
break
# 关闭摄像头
cam.release()
cv2.destroyAllWindows()
```
阅读全文