from ultralytics import YOLO
from PIL import Image
import cv2
import os

# def draw_boxes(image, detections):
#     for detection in detections:
#         label_name, bbox_list = detection['label'], detection['bbox']
#         confidence = detection['confidence']

#         # 提取边界框坐标
#         x1, y1, x2, y2 = map(int, bbox_list)

#         # 在图像上绘制边界框
#         cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

#         # 在图像上写入标签名和置信度
#         text = f"{label_name}: {confidence}"
#         cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

#     return image

# def analysis_yolov8(frame, model_coco, confidence_set):
#     results_coco = model_coco(frame)
#     re_list = []

#     if results_coco:
#         for r in results_coco:
#             boxes = r.boxes
#             for box in boxes:
#                 b = box.xyxy[0]  # 获取边界框坐标
#                 c = box.cls
#                 labels_name = model_coco.names[int(c)]
#                 confidence = float(box.conf)
#                 confidence = round(confidence, 2)

#                 if confidence < confidence_set:
#                     continue

                
#                 re_dict = {'label': labels_name, 'bbox': b, 'confidence': confidence}
#                 re_list.append(re_dict)

#     return re_list

# # # 创建新文件夹保存结果
# output_folder = "E:/code_files/jy/20240411_test_img_file/t_0"
# os.makedirs(output_folder, exist_ok=True)

# 加载模型和图像
model = YOLO("E:/code_files/jy/20240411_model_test/train/weights/best.onnx")
img_files = r"E:\code_files\jy\20240423"

for filename in os.listdir(img_files):

    img_path = os.path.join(img_files, filename)
    model.predict(img_path, save=True)
    # frame = cv2.imread(img_path)  # 替换为你的图像路径

    # 调用检测函数获取边界框结果
    # detections = analysis_yolov8(frame, model, confidence_set=0.5)

    # # 绘制边界框和标签
    # annotated_image = draw_boxes(frame, detections)

    # # 保存结果图像到新文件夹
    # output_path = os.path.join(output_folder, filename)
    # cv2.imwrite(output_path, annotated_image)

    # print(f"Annotated image saved at: {output_path}")