import os
import cv2def 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_folderdef 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]) * widthy_center = float(parts[2]) * heightimg_width = float(parts[3]) * widthimg_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 = 1file_path = os.path.join(class_folder, f"{file_number}.jpg")while os.path.exists(file_path):file_number += 1file_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:# 获取类别IDclass_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()