import os
import cv2


def get_class_folder(catagetory,class_id, base_folder):
    # 根据类别ID创建文件夹路径
    class_folder = os.path.join(base_folder, catagetory[int(class_id)])
    if not os.path.exists(class_folder):
        os.makedirs(class_folder)
    return class_folder


def crop_and_save(image_path, label, class_folder):
    # 读取图片
    img = cv2.imread(image_path)
    height, width = img.shape[:2]

    # 解析标签
    parts = label.strip().split()

    x_center = float(parts[1]) * width
    y_center = float(parts[2]) * height
    img_width = float(parts[3]) * width
    img_height = float(parts[4]) * height

    # 计算左上角坐标
    x1 = int(x_center - (img_width / 2))
    y1 = int(y_center - (img_height / 2))

    # 确保裁剪区域在图片范围内
    x1 = max(0, x1)
    y1 = max(0, y1)
    x2 = min(width, x1 + img_width)
    y2 = min(height, y1 + img_height)

    # 裁剪图片
    cropped_img = img[int(y1):int(y2), int(x1):int(x2)]

    # 保存裁剪后的图片，文件名为数字递增
    file_number = 1
    file_path = os.path.join(class_folder, f"{file_number}.jpg")
    while os.path.exists(file_path):
        file_number += 1
        file_path = os.path.join(class_folder, f"{file_number}.jpg")

    cv2.imwrite(file_path, cropped_img)


def main():
    root = r"D:\desk\Add_Building_20240616\IR"
    images_folder = os.path.join(root, 'images')
    labels_folder = os.path.join(root,'labels')
    cropped_images_base_folder = os.path.join(root,'cropped_images')
    class_path = os.path.join(labels_folder, r"classes.txt")
    category ={}
    with open(class_path, 'r', encoding='utf-8') as f:
        cls = f.readlines()
        for i, c in enumerate(cls):
            category[i] = c.strip()

    # 确保裁剪图片的基文件夹存在
    if not os.path.exists(cropped_images_base_folder):
        os.makedirs(cropped_images_base_folder)

    # 遍历images文件夹中的所有图片
    for image_filename in os.listdir(images_folder):
        if image_filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            # 构建图片和标签的完整路径
            image_path = os.path.join(images_folder, image_filename)
            label_path = os.path.join(labels_folder, os.path.splitext(image_filename)[0] + '.txt')

            # 确保标签文件存在
            if os.path.exists(label_path):
                # 读取标签文件
                with open(label_path, 'r') as file:
                    labels = file.readlines()

                # 遍历每个标签
                for label in labels:
                    # 获取类别ID
                    class_id = int(label.strip().split()[0])

                    # 获取类别文件夹
                    class_folder = get_class_folder(category,class_id, cropped_images_base_folder)

                    # 裁剪并保存图片
                    crop_and_save(image_path, label, class_folder)


if __name__ == '__main__':
    main()