// An highlighted block
import os
from collections import defaultdict
import cv2

# VisDrone类别ID到YOLO类别ID的映射
# VisDrone类别ID: “ignored regions”、“pedestrian”、“people”、“bicycle”、“car”、
# “van”、“truck”、“tricycle”、“awning-tricycle”、“bus”、“motor”、“others”。
# 这里假设YOLO类别ID从0开始，忽略类别0和11
category_mapping = {
    1: 0,  # pedestrian -> 0
    2: 1,  # people -> 1
    3: 2,  # bicycle -> 2
    4: 3,  # car -> 3
    5: 4,  # van -> 4
    6: 5,  # truck -> 5
    7: 6,  # tricycle -> 6
    8: 7,  # awning-tricycle -> 7
    9: 8,  # bus -> 8
    10: 9,  # motor -> 8
}

def convert_visdrone_to_yolo(visdrone_annotation_path, yolo_annotation_dir, image_width, image_height):
    # 按帧索引分组
    frame_annotations = defaultdict(list)

    with open(visdrone_annotation_path, 'r') as f:
        lines = f.readlines()

    for line in lines:
        data = line.strip().split(',')
        if len(data) < 8:
            continue

        # 解析VisDrone格式
        frame_index = int(data[0])  # 帧索引
        bbox_left = float(data[2])
        bbox_top = float(data[3])
        bbox_width = float(data[4])
        bbox_height = float(data[5])
        category_id = int(data[7])

        # 忽略类别0或无效类别
        if category_id == 0 or category_id not in category_mapping:
            continue

        # 计算归一化坐标
        x_center = (bbox_left + bbox_width / 2) / image_width
        y_center = (bbox_top + bbox_height / 2) / image_height
        width = bbox_width / image_width
        height = bbox_height / image_height

        # 映射类别ID
        yolo_class_id = category_mapping[category_id]

        # 添加到对应帧的标注列表
        frame_annotations[frame_index].append(f"{yolo_class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")

    # 保存每帧的YOLO格式标注文件
    for frame_index, annotations in frame_annotations.items():
        yolo_annotation_path = os.path.join(yolo_annotation_dir, f"{frame_index:06d}.txt")
        with open(yolo_annotation_path, 'w') as f:
            f.writelines(annotations)


# 示例：转换单个视频的标注文件
def batch_convert_visdrone_to_yolo(visdrone_annotation_dir, yolo_annotation_dir, visdrone_sequences_dir):
    if not os.path.exists(yolo_annotation_dir):
        os.makedirs(yolo_annotation_dir)

    for annotation_file in os.listdir(visdrone_annotation_dir):
        if not annotation_file.endswith(".txt"):
            continue
        visdrone_sequences_path = os.path.join(visdrone_sequences_dir, str(annotation_file.split('.')[0]))
        print('visdrone_sequences_path',visdrone_sequences_path)
        image = cv2.imread(visdrone_sequences_path + "/0000001.jpg")
        image_height,image_width = image.shape[:2]
        print('shape', image_width, image_height)
        visdrone_annotation_path = os.path.join(visdrone_annotation_dir, annotation_file)
        video_name = os.path.splitext(annotation_file)[0]
        video_yolo_annotation_dir = os.path.join(yolo_annotation_dir, video_name)

        if not os.path.exists(video_yolo_annotation_dir):
            os.makedirs(video_yolo_annotation_dir)

        convert_visdrone_to_yolo(visdrone_annotation_path, video_yolo_annotation_dir, image_width, image_height)


# 示例：批量转换
visdrone_annotation_dir = "./VisDrone2019-VID-val/annotations"
visdrone_sequences_dir = "./VisDrone2019-VID-val/sequences"
yolo_annotation_dir = "./VisDrone2019-VID-val/annotations_yolo"

batch_convert_visdrone_to_yolo(visdrone_annotation_dir, yolo_annotation_dir, visdrone_sequences_dir)