125 lines
3.3 KiB
Python
125 lines
3.3 KiB
Python
import os.path
|
||
import string
|
||
import json
|
||
from ultralytics import YOLO
|
||
|
||
'''
|
||
注意修改配置文件中的文件夹路径
|
||
|
||
|
||
'''
|
||
|
||
# 使用secrets模块（更安全但更慢）
|
||
import secrets
|
||
|
||
def secure_generate(num: int):
|
||
chars = string.ascii_letters + string.digits + '_'
|
||
candidate = ''.join(secrets.choice(chars) for _ in range(num))
|
||
|
||
return candidate
|
||
|
||
|
||
model = YOLO('/home/admin-root/haotian/劳保鞋识别/jingzhu_245_1112/tensorrtx-master/yolov8/build_1113/best.pt')
|
||
|
||
images_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/HardHatsV1/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 = model(all_images[i],
|
||
conf = 0.2, device=0)
|
||
|
||
# # Process results list
|
||
# = 0
|
||
# for result in results:
|
||
# boxes = result.boxes # Boxes object for bounding box outputs
|
||
# # masks = result.masks # Masks object for segmentation masks outputs
|
||
# # keypoints = result.keypoints # Keypoints object for pose outputs
|
||
|
||
# # 属于这个类的置信度
|
||
# probs = result.probs # Probs object for classification outputs
|
||
# # obb = result.obb # Oriented boxes object for OBB outputs
|
||
# print(boxes,' ',probs)
|
||
# # result.show() # display to screen
|
||
# results[0].save(filename=f"{save_path}/result_{i}.jpg") # save to disk
|
||
|
||
|
||
|
||
names = results[0].names
|
||
|
||
xywhns = results[0].boxes.xywhn
|
||
clss = results[0].boxes.cls
|
||
|
||
t_d = dict()
|
||
t_d["data"] = {"image": f"/data/local-files/?d=HardHatsV1/train/images/{image_names[i]}"}
|
||
t_d["predictions"] = list()
|
||
|
||
t_d_p = dict()
|
||
t_d_p["model_version"] = "version 1"
|
||
t_d_p["score"] = 0.8
|
||
t_d_p["result"] = list()
|
||
|
||
|
||
|
||
orig_img = results[0].orig_shape
|
||
|
||
for idx,boxe in enumerate(xywhns):
|
||
|
||
|
||
|
||
t_d_p_b = dict()
|
||
t_d_p_b["original_width"] = orig_img[1]
|
||
t_d_p_b["original_height"] = orig_img[0]
|
||
t_d_p_b["image_rotation"] = 0
|
||
|
||
#随机字符串
|
||
t_d_p_b["id"] = secure_generate(9)
|
||
t_d_p_b["from_name"] = "label"
|
||
t_d_p_b["to_name"] = "image"
|
||
t_d_p_b["type"] = "rectanglelabels"
|
||
t_d_p_b["origin"] = "manual"
|
||
|
||
|
||
x, y, w, h = boxe
|
||
class_name = names[clss[idx].item()]
|
||
|
||
t_d_p_b["value"]={
|
||
"x": x.item() * 100,
|
||
"y": y.item() * 100,
|
||
"width" : w.item() * 100,
|
||
"height" : h.item() * 100,
|
||
"rotation": 0,
|
||
"rectanglelabels": [ class_name]
|
||
}
|
||
|
||
t_d_p["result"].append(t_d_p_b)
|
||
|
||
t_d["predictions"].append(t_d_p)
|
||
|
||
json_list.append(t_d)
|
||
|
||
|
||
break
|
||
|
||
with open("test_json.json", "w", encoding="utf-8") as f:
|
||
|
||
json.dump(json_list,f,
|
||
indent=4, # 4空格缩进
|
||
sort_keys=True, # 按键名字典序排序
|
||
ensure_ascii=False # 支持非ASCII字符
|
||
)
|
||
|
||
|
||
|
||
|
||
# i += 1 |