import os
import cv2
import numpy as np

# 总的检测根目录
path_root_labels = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/HardHatsV1/train/labels_helmet_head_shoe'
# 总的检测根目录
path_root_imgs = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/HardHatsV1/train/images'

path_root_save = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/HardHatsV1/train/show'


remove_path = 'fire-8/remove/valid.txt'

type_object = '.txt'

for ii in os.walk(path_root_imgs):
    for j in ii[2]:
        type = j.split(".")[-1]
        if type != 'jpg':
            continue
        path_img = os.path.join(path_root_imgs, j)
        # print(path_img)
        label_name = j[:-4] + type_object
        path_label = os.path.join(path_root_labels, label_name)
        # print(path_label)
        
        if os.path.exists(path_label):
            
            f = open(path_label, 'r+', encoding='utf-8')

            img = cv2.imdecode(np.fromfile(path_img, dtype=np.uint8), -1)
            w = img.shape[1]
            h = img.shape[0]
            new_lines = []
            img_tmp = img.copy()
            while True:
                line = f.readline()
                if line:
                    msg = line.split(" ")
                    # if msg[0] == '1':
                    #     with open(remove_path, 'a') as t:
                    #         t.write(j+'\n')
                        # print(j)
                    # print(x_center,",",y_center,",",width,",",height)
                    x1 = int((float(msg[1]) - float(msg[3]) / 2) * w)  # x_center - width/2
                    y1 = int((float(msg[2]) - float(msg[4]) / 2) * h)  # y_center - height/2
                    x2 = int((float(msg[1]) + float(msg[3]) / 2) * w)  # x_center + width/2
                    y2 = int((float(msg[2]) + float(msg[4]) / 2) * h)  # y_center + height/2
                    # print(x1, ",", y1, ",", x2, ",", y2)
                    cv2.rectangle(img_tmp, (x1, y1), (x2, y2), (0, 0, 255), 5)
                    cv2.putText(img_tmp, msg[0], (x1-5, y1-5), 2, 2, 0, 2)
                else:
                    break
        # cv2.imshow("show", img_tmp)
        # c = cv2.waitKey(0)
        try:
            cv2.imwrite(f'{path_root_save}/{j}', img_tmp)
        except:
            print(f"写入图片失败: {j}")
        # print(f'写入文件{j}')