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