import json
import os
from uuid import uuid4
from ultralytics import YOLO

def yolo_to_labelstudio(results, class_names, model_version="yolov8_1.0"):
    """将 YOLOv8 检测结果转换为 Label Studio 预测格式
    
    Args:
        results: YOLOv8 Results 对象
        class_names: 类别名称列表
        model_version: 模型版本标识
        
    Returns:
        Label Studio 兼容的预测字典
    """
    predictions = []
    
    # 获取原始图像尺寸 (height, width)
    orig_height = results[0].orig_shape[0]
    orig_width = results[0].orig_shape[1]
    
    for box in results[0].boxes:
        # 提取检测框信息
        x_min, y_min, x_max, y_max = box.xyxy[0].tolist()
        conf = box.conf.item()
        cls_idx = int(box.cls.item())
        
        # 计算百分比坐标（保留3位小数）
        x_percent = round((x_min / orig_width) * 100, 3)
        y_percent = round((y_min / orig_height) * 100, 3)
        width_percent = round((x_max - x_min) / orig_width * 100, 3)
        height_percent = round((y_max - y_min) / orig_height * 100, 3)
        
        # 构建预测项
        prediction = {
            "id": str(uuid4()),  # 生成唯一ID
            "type": "rectanglelabels",
            "from_name": "label",
            "to_name": "image",
            "original_width": orig_width,
            "original_height": orig_height,
            "value": {
                "x": x_percent,
                "y": y_percent,
                "width": width_percent,
                "height": height_percent,
                "rotation": 0,
                "rectanglelabels": [class_names[cls_idx]]
            },
            "score": round(conf, 4)  # 保留4位小数
        }
        predictions.append(prediction)
    
    return {
        "model_version": model_version,
        "result": predictions
    }

# 使用示例
if __name__ == "__main__":
    # 初始化模型
    # model_0 = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV1/weights/best.pt')
    
    # 头,安全帽
    model_0 = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV12/weights/best.pt')
    
    # 检测鞋的模型
    model_1 = YOLO('/home/admin-root/haotian/劳保鞋识别/jingzhu_245_1112/tensorrtx-master/yolov8/build_1113/best.pt')
    
    # 定义类别名称 (需与训练时一致)
    CLASS_NAMES_0 = [
        "helmet", "head",
    ]
    
    CLASS_NAMES_1 = [
        "shoe",
    ]
    
    

    images_path = '/home/admin-root/haotian/劳保鞋识别/dataset/new_all/train/images'

    # save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputShoeV2'

    image_names = os.listdir(images_path)

    all_images = [os.path.join(images_path, t) for  t in image_names]
    
    json_list = list()
    
    for i in range(len(all_images)):

        # conf 置信度在.
        results_0 = model_0(all_images[i],
                        conf = 0.5, device=0)
        
        predictions_0 = yolo_to_labelstudio(results_0, CLASS_NAMES_0)
        
        
        results_1 = model_1(all_images[i], 
                            conf=0.5, device=0)
        
        predictions_1 = yolo_to_labelstudio(results_1, CLASS_NAMES_1)
        
        predictions_0["result"] += predictions_1["result"]
        
        
        t_d = dict()
        # "劳保鞋识别/dataset/new_all/train/images/dataset2_-01-15-2-2-2-2-25_jpg.rf.267991c87898d90733d5c7623941a999.jpg"
        t_d["data"] = {"image": f"/data/local-files/?d=劳保鞋识别/dataset/new_all/train/images/{image_names[i]}"}
        t_d["predictions"] = list()
        
        t_d["predictions"].append(predictions_0)
        
        json_list.append(t_d)
        
        # if i == 10:
        #     break
        

with open("label_studio_new_all.json", "w", encoding="utf-8") as f:

    json.dump(json_list,f,
               indent=4,              # 4空格缩进
                sort_keys=True,        # 按键名字典序排序
                ensure_ascii=False     # 支持非ASCII字符
              )