修改--封装030脚本内容

This commit is contained in:
haotian 2025-06-25 09:31:41 +08:00
parent 0047bb87f7
commit 8df4dc1e6a
2 changed files with 119 additions and 31 deletions

View File

@ -11,9 +11,10 @@ model = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHea
# images_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images'
# images_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images_20250623'
images_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/test_images/duan8_real'
# images_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/test_images/duan8_real'
images_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images_20250624105822'
# save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputHardHatsV1HelmetHeadShoe2_20250623'
save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputHardHatsV1HelmetHeadShoe2_20250623_1'
save_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/output/outputHardHatsV1HelmetHeadShoe2_202506624105822'
all_images = [os.path.join(images_path, t) for t in os.listdir(images_path)]

View File

@ -57,6 +57,62 @@ def yolo_to_labelstudio(results, class_names, model_version="yolov8_1.0"):
"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__":
# 初始化模型
@ -77,55 +133,86 @@ if __name__ == "__main__":
"shoe",
]
model_2 = YOLO('/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV1HelmetHeadShoe2/weights/best.pt')
CLASS_NAMES_2 = [
"helmet",
"head",
"shoe"
]
images_path = '/home/admin-root/haotian/劳保鞋识别/dataset/new_all/train/images'
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'
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字符
)
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字符
# )