import os
import cv2
from ultralytics import YOLO

# 加载模型
model = YOLO("best.pt")

# 输入文件夹路径
input_dir = "test_images/"
temp_dir = "temp_rgb_images"

# 创建临时文件夹用于保存转换后的RGB图片
os.makedirs(temp_dir, exist_ok=True)

# 扫描文件夹，将灰度图转换为RGB
for file in os.listdir(input_dir):
    if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
        path = os.path.join(input_dir, file)
        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)

        # 如果是灰度图，转换为3通道
        if len(img.shape) == 2 or img.shape[2] == 1:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

        # 保存到临时目录
        cv2.imwrite(os.path.join(temp_dir, file), img)

# 使用转换后的图像文件夹进行检测
results = model.predict(source=temp_dir, conf=0.2, save=True)

# 输出检测信息
for result in results:
    boxes = result.boxes  # 检测框
    for box in boxes:
        cls = int(box.cls[0])
        conf = float(box.conf[0])
        xyxy = box.xyxy[0].tolist()
        print(f"类别: {model.names[cls]}, 置信度: {conf:.2f}, 坐标: {xyxy}")