diff --git a/res/4.mp4 b/res/4.mp4 new file mode 100644 index 0000000..074e67d Binary files /dev/null and b/res/4.mp4 differ diff --git a/res/5.mp4 b/res/5.mp4 new file mode 100644 index 0000000..6a85a53 Binary files /dev/null and b/res/5.mp4 differ diff --git a/runs/detect/predict35/4.avi b/runs/detect/predict35/4.avi new file mode 100644 index 0000000..47fc423 Binary files /dev/null and b/runs/detect/predict35/4.avi differ diff --git a/runs/detect/predict36/5.avi b/runs/detect/predict36/5.avi new file mode 100644 index 0000000..75a5614 Binary files /dev/null and b/runs/detect/predict36/5.avi differ diff --git a/src/pic.py b/src/pic.py new file mode 100644 index 0000000..8f65343 --- /dev/null +++ b/src/pic.py @@ -0,0 +1,88 @@ +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:deskAdd_Building_20240616IR" + 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() \ No newline at end of file diff --git a/src/test.py b/src/test.py index 1a5e4fc..406b1d4 100644 --- a/src/test.py +++ b/src/test.py @@ -1,3 +1,3 @@ from ultralytics import YOLO model = YOLO(r"C:\workspace\le-yolo\runs\detect\train26\weights\best.pt") -results = model.predict("../res/3.mp4", show=True, save=True) +results = model.predict("../res/5.mp4", show=True, save=True)