#静态检测，图片检测

from ultralytics import YOLO   # ultralytics是yolo模型的实现库
import cv2                     # cv2是opencv库，用于图像处理
import os                      # os用于文件和目录操作

# 替换为实际的模型路径
model_path = r'\detect\train\weights\best.pt'   # 加载yolo模型
model = YOLO(model_path)                             # model是加载后的yolo模型实例

def infer_and_draw(image_path, output_folder):
    # 检查 image_path 是否是文件夹路径
    if os.path.isdir(image_path):
        # 遍历文件夹中的所有图像文件
        for filename in os.listdir(image_path):
            file_path = os.path.join(image_path, filename)
            if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                process_image(file_path, output_folder)
    else:
        # 处理单张图像
        process_image(image_path, output_folder)

def process_image(image_path, output_folder):
    # 使用模型进行推理
    results = model(image_path)
    original_img = cv2.imread(image_path)  # 读取原始图像

    if results:
        for result in results:
            if result.boxes is not None and len(result.boxes.xyxy) > 0:
                boxes = result.boxes.xyxy.cpu().numpy()  # 获取检测框
                classes = result.boxes.cls.cpu().numpy()  # 获取类别
                confidences = result.boxes.conf.cpu().numpy()  # 获取置信度

                # 确保输出文件夹存在
                if not os.path.exists(output_folder):
                    os.makedirs(output_folder)

                # 存储所有检测框的标签和位置
                detected_boxes = []

                for i, box in enumerate(boxes):
                    x1, y1, x2, y2 = map(int, box[:4])
                    class_id = int(classes[i]) if len(classes) > i else 0
                    # 标签映射（假设0表示人头类别）
                    label = "head" if class_id == 0 else "unknown"
                    conf = confidences[i] if len(confidences) > i else 0.0

                    detected_boxes.append((x1, y1, x2, y2, label, conf))

                # 避免重叠框的问题，通过过滤和合并框
                filtered_boxes = filter_and_merge_boxes(detected_boxes)

                # 生成最终的标注图像
                img_annotated = original_img.copy()
                for (x1, y1, x2, y2, label, conf) in filtered_boxes:
                    cv2.rectangle(img_annotated, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(img_annotated, f'{label} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                # 保存标注后的图像
                output_image_path = os.path.join(output_folder, os.path.basename(image_path))
                cv2.imwrite(output_image_path, img_annotated)
                print(f"Saved annotated image to: {output_image_path}")

                # 显示图像（可选）
                cv2.imshow('Detected Image', img_annotated)
                cv2.waitKey(0)
                cv2.destroyAllWindows()
            else:
                print("No detections found.")
                # 保存原始图像
                output_image_path = os.path.join(output_folder, os.path.basename(image_path))
                cv2.imwrite(output_image_path, original_img)
                print(f"Saved original image to: {output_image_path}")

def filter_and_merge_boxes(boxes):
    """
    简单地合并重叠框的逻辑。可以根据需要调整合并算法。
    """
    filtered_boxes = []
    threshold = 0.5  # 重叠阈值

    def iou(box1, box2):
        x1, y1, x2, y2 = box1
        x1_, y1_, x2_, y2_ = box2
        ix1, iy1 = max(x1, x1_), max(y1, y1_)
        ix2, iy2 = min(x2, x2_), min(y2, y2_)
        iw = max(ix2 - ix1 + 1, 0)
        ih = max(iy2 - iy1 + 1, 0)
        inter = iw * ih
        ua = (x2 - x1 + 1) * (y2 - y1 + 1) + (x2_ - x1_ + 1) * (y2_ - y1_) - inter
        return inter / ua

    for i, box1 in enumerate(boxes):
        keep = True
        for j, box2 in enumerate(filtered_boxes):
            if iou(box1[:4], box2[:4]) > threshold:
                keep = False
                break
        if keep:
            filtered_boxes.append(box1)

    return filtered_boxes

# 使用实际图像路径或文件夹路径进行推理，并指定输出文件夹
infer_and_draw(r'\assets\xx.jpg(图片路径)', r'ultralytics\output（输出路径）')