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 } def main(model_list, class_name_list, images_path, save_file, pre_root, conf=0.3): """_summary_ Args: model_list (_type_): 模型列表 class_name_list (_type_): 模型输出类名称 images_path (_type_): 图像路径 save_file (_type_): 保存json文件路径 pre_root (_type_): json中文件路径前缀 """ 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)): predictions = None for j in range(len(model_list)): results = model_list[j](all_images[i], conf = conf, device=0) predictions_t = yolo_to_labelstudio(results, class_name_list[j]) if predictions is None: predictions = predictions_t else: predictions["result"] += predictions_t["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["data"] = {"image": f"/data/local-files/?d=images_20250624105811/{image_names[i]}"} t_d["data"] = {"image": f"/data/local-files/?d={pre_root}/{image_names[i]}"} t_d["predictions"] = list() t_d["predictions"].append(predictions) json_list.append(t_d) print(len(json_list)) # return with open(save_file, "w", encoding="utf-8") as f: json.dump(json_list,f, indent=4, # 4空格缩进 sort_keys=True, # 按键名字典序排序 ensure_ascii=False # 支持非ASCII字符 ) # 使用示例 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", ] model_2 = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV1HelmetHeadShoe2/weights/best.pt') CLASS_NAMES_2 = [ "helmet", "head", "shoe" ] model_list = list() model_list.append(model_2) class_name_list = list() class_name_list.append(CLASS_NAMES_2) # images_path = '/home/admin-root/haotian/劳保鞋识别/dataset/new_all/train/images' # images_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images_20250624105811' images_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images_20250624105822' # save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputShoeV2' main(model_list, class_name_list, images_path, "duan8_22.json", "images_20250624105822") #-----------------------------------------------------------------旧版代码------------------------------------------------------------------------------- # 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["data"] = {"image": f"/data/local-files/?d=images_20250624105811/{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字符 # )