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 文件夹!")