import os
import json

# 修改为你的标签目录路径
LABEL_DIR = "D:\ChromeDownload\***"
# 输出标签目录，和图像文件名对应的txt文件存这里
OUTPUT_DIR = "D:\ChromeDownload\labels"

# 假设图像尺寸，若不是1920x1080，修改这里
IMAGE_WIDTH = 1920
IMAGE_HEIGHT = 1080

# 类别映射，没在这里的类别会跳过
CLASS_MAP = {
    "Car": 0,           # 小汽车
    "Truck": 1,         # 卡车/大货车
    "Van": 2,           # 面包车/厢式货车
    "Bus": 3,           # 公交车
    "Pedestrian": 4,    # 行人
    "Cyclist": 5,       # 自行车
    "Tricyclist": 6,    # 三轮车
    "Motorcyclist": 7,  # 摩托车
    "Barrowlist": 8     # 手推车
}


def convert_bbox_to_yolo(xmin, ymin, xmax, ymax, img_w, img_h):
    x_center = (xmin + xmax) / 2.0 / img_w
    y_center = (ymin + ymax) / 2.0 / img_h
    width = (xmax - xmin) / img_w
    height = (ymax - ymin) / img_h
    return x_center, y_center, width, height
# 务必根据json中的标签做适当调整
def convert_json_to_txt(json_path, txt_path):
    with open(json_path, 'r') as f:
        data = json.load(f)

    lines = []
    for obj in data:
        cls = obj.get("type")
        if cls not in CLASS_MAP:
            continue  # 不在映射表中类别跳过

        bbox = obj.get("2d_box")
        xmin = bbox["xmin"]
        ymin = bbox["ymin"]
        xmax = bbox["xmax"]
        ymax = bbox["ymax"]

        x_c, y_c, w, h = convert_bbox_to_yolo(xmin, ymin, xmax, ymax, IMAGE_WIDTH, IMAGE_HEIGHT)

        # 限制范围在0~1内（可选）
        x_c = max(min(x_c, 1), 0)
        y_c = max(min(y_c, 1), 0)
        w = max(min(w, 1), 0)
        h = max(min(h, 1), 0)

        cls_id = CLASS_MAP[cls]
        lines.append(f"{cls_id} {x_c:.6f} {y_c:.6f} {w:.6f} {h:.6f}")

    if lines:
        with open(txt_path, "w") as f:
            f.write("\n".join(lines))
    else:
        # 如果没有有效标签，写空文件或者不写都可以，这里写空文件
        with open(txt_path, "w") as f:
            pass

def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    for filename in os.listdir(LABEL_DIR):
        if not filename.endswith(".json"):
            continue
        json_path = os.path.join(LABEL_DIR, filename)
        txt_filename = os.path.splitext(filename)[0] + ".txt"
        txt_path = os.path.join(OUTPUT_DIR, txt_filename)
        convert_json_to_txt(json_path, txt_path)
        print(f"Converted {filename} -> {txt_filename}")

if __name__ == "__main__":
    main()