修改--封装030脚本内容
This commit is contained in:
parent
0047bb87f7
commit
8df4dc1e6a
@ -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)]
|
||||
|
||||
|
||||
@ -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字符
|
||||
# )
|
||||
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user