131 lines
4.1 KiB
Python
131 lines
4.1 KiB
Python
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字符
|
||
)
|
||
|
||
|
||
|