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()