29 lines
964 B
Python
29 lines
964 B
Python
from ultralytics import YOLO
|
|
import os
|
|
import cv2
|
|
|
|
# 创建输出目录
|
|
os.makedirs('output', exist_ok=True)
|
|
|
|
# 加载预训练模型 (YOLOv8会自动下载)
|
|
model = YOLO('models/yolov8n.pt') # 使用nano版本 (最小最快的模型)
|
|
|
|
# 处理photo目录中的所有图片
|
|
for filename in os.listdir('photo'):
|
|
if filename.lower().endswith(('.png', '.jpg', '.jpeg','.HEIC')):
|
|
img_path = os.path.join('photo', filename)
|
|
|
|
# 进行目标检测
|
|
results = model.predict(img_path, save=True, project='output')
|
|
|
|
# 打印检测结果
|
|
print(f"\n检测结果 {filename}:")
|
|
for result in results:
|
|
boxes = result.boxes
|
|
for box in boxes:
|
|
cls_id = int(box.cls[0])
|
|
conf = box.conf[0].item()
|
|
cls_name = model.names[cls_id]
|
|
print(f"- 检测到 {cls_name} ({conf:.2f})")
|
|
|
|
print("\n所有结果已保存到 output 文件夹!") |