# -*- coding: utf-8 -*-
import os
import yaml
from PIL import Image
from yolov.lanrenyolov import LanrenYolov
def convert_to_labels_txt(datas, output_file, classes, img_width, img_height):
    with open(output_file, 'w') as f:
        for data in datas:
            # 遍历字典的键值对，查找值为 2 的键名
            for key, value in classes["names"].items():
                if value == data["label"]:
                    class_id=key
                    break  # 如果只需要找到第一个匹配的键名，则可以使用 break 来提前结束循环
            x = data["x"] / img_width
            y = data["y"] / img_height
            width =( data["points"][1][0] - data["points"][0][0])/ img_width
            height =(  data["points"][1][1] - data["points"][0][1])/ img_height

            # 保留小数点6位
            x = "{:.6f}".format(x)
            y = "{:.6f}".format(y)
            width = "{:.6f}".format(width)
            height = "{:.6f}".format(height)

            label_line = f"{class_id} {x} {y} {width} {height}\n"
            f.write(label_line)
# 这个文件主要用来调用已经预先训练了一部分数据的模型,然后让它自动标注,最后你可以将数据导入到标注工具中进行手动检查
# 检测和手动校准后,再将它加入到你的数据集中,这样就能省一大波力气
weights = "./datas/yolov5l_best.pt"
yaml_file_path = "D:\pythonProjects\yuanshenAI\datasets\yuanshen\yuanshen.yaml"
device = "0"
yolov = LanrenYolov(weights=weights,data=yaml_file_path, device=device,   conf_thres=0.8,iou_thres=0.5)
data_dir = "./autolabel/"  # 总目录
img_dir = "./autolabel/images"  # 图片目录
label_dir = "./autolabel/labels"  # 标签目录
num_ok=0
# 读取 YAML 文件并解析为 Python 对象
with open(yaml_file_path, "r", encoding="utf-8", errors="ignore") as f:
    classes = yaml.safe_load(f)
# 如果没有上方目录则创建
if not os.path.exists(data_dir):
    os.makedirs(data_dir)
if not os.path.exists(img_dir):
    os.makedirs(img_dir)
if not os.path.exists(label_dir):
    os.makedirs(label_dir)
print("准备自动标注...")
# 循环遍历imag_dir文件夹中的所有图片
for filename in os.listdir(img_dir):
    # 所有格式为png和jpg的图片
    if filename.endswith('.png') or filename.endswith('.jpg'):
        #如果文件格式是.png转成jpg
        if filename.endswith('.png'):  # 找到所有以.png结尾的文件
            img_path = os.path.join(img_dir, filename)  # 获取图片的完整路径
            img = Image.open(img_path)  # 打开图片
            jpg_filename = filename[:-4] + '.jpg'  # 将文件名的后缀改为.jpg
            jpg_path = os.path.join(img_dir, jpg_filename)  # 新的jpg文件路径
            img=img.convert("RGB")
            img.save(jpg_path)  # 保存为jpg格式
            img.close()  # 关闭图片对象
            os.remove(img_path)  # 删除原来的png文件
            filename = os.path.splitext(filename)[0] + '.jpg'  # 更改文件名后缀为.jpg

        # 取图片的大小
        img_width, img_height = Image.open(os.path.join(img_dir, filename)).size
        output_file = os.path.join(label_dir, filename.split(".")[0] + ".txt")
        datas = yolov.detect(img_dir + "/" + filename, view_img=False)
        if datas!=[]:
            # 将结果生成标签  找出id
            convert_to_labels_txt(datas[0], output_file, classes, img_width, img_height)
            num_ok += 1
            print("完成",num_ok,"张图片","保存在",output_file)
print("自动标注完成...总数",num_ok)