import os
import cv2

def crop_image(image_path, label_path, output_folder):
    # 读取图片
    img = cv2.imread(image_path)
    height, width, _ = img.shape

    # 读取标签文件
    with open(label_path, 'r') as file:
        labels = file.readlines()
    img_id = 1
    # 遍历每个标签
    for label in labels:
        parts = label.strip().split()
        class_id = int(parts[0])
        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)]

        # 保存裁剪后的图片
        cv2.imwrite(os.path.join(output_folder, f"{img_id}.jpg"), cropped_img)
        img_id+=1

def main():
    images_folder = r'E:\BaiduNetdiskDownload\20240616\Add_Dataset\Add_Building_20240616\images'
    labels_folder = r'E:\BaiduNetdiskDownload\20240616\Add_Dataset\Add_Building_20240616\labels'
    output_base_folder = r'E:\BaiduNetdiskDownload\20240616\Add_Dataset\Add_Building_20240616\cut_imgs'

    # 确保输出目录存在
    if not os.path.exists(output_base_folder):
        os.makedirs(output_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_filename = image_filename.replace('.jpg', '.txt').replace('.jpeg', '.txt').replace('.png', '.txt')
            label_path = os.path.join(labels_folder, label_filename)

            # 确保标签文件存在
            if os.path.exists(label_path):
                # 以图片的名字为文件夹名创建输出目录
                output_folder = os.path.join(output_base_folder, os.path.splitext(image_filename)[0])
                if not os.path.exists(output_folder):
                    os.makedirs(output_folder)

                # 裁剪图片
                crop_image(image_path, label_path, output_folder)

if __name__ == '__main__':
    main()