【目标检测】使用YOLO格式标注将目标框中的图片截取保存

将目标框中的图片截取保存到文件夹中用于分类网络训练。

Python代码：

import os
import cv2

"""
|-----------------------------------------------|
| index |           class           | new label |
|-----------------------------------------------|
|   0   | Class 1                   |  label 1  |
|   1   | Class 2                   |  label 2  |
|   2   | Class 3                   |  label 3  |
|-----------------------------------------------|
"""

# 指定图片文件夹路径和txt文件夹路径
img_folder = "images"
txt_folder = "labels"

# 遍历图片文件夹中的所有文件
for filename in os.listdir(img_folder):
    # 确保文件是一个图像文件
    if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"):
        # 构造图像文件路径
        image_file_path = os.path.join(img_folder, filename)

        # 构造对应的txt文件路径
        txt_filename = os.path.splitext(filename)[0] + ".txt"
        label_file_path = os.path.join(txt_folder, txt_filename)

        # 读取标注文件和原始图片
        with open(label_file_path, 'r') as f:
            lines = f.readlines()
        image = cv2.imread(image_file_path)

        # 循环处理每个标注框
        i = 0
        for line in lines:
            i = i + 1
            class_id, x_center, y_center, width, height = map(float, line.strip().split())

            # 将YOLOv5格式的坐标转换为常规坐标
            left = int((x_center - width / 2) * image.shape[1])
            top = int((y_center - height / 2) * image.shape[0])
            right = int((x_center + width / 2) * image.shape[1])
            bottom = int((y_center + height / 2) * image.shape[0])

            class_id = int(class_id)

            if class_id == 0:
                folder = 'label 1'
            elif class_id == 1:
                folder = 'label 2'
            elif class_id == 2:
                folder = 'label 3'

            # 截取标注框内的内容并保存为新图片
            object_image = image[top:bottom, left:right]
            object_file_path = image_file_path.split("\\")
            object_file_path = f"{folder}/" + os.path.splitext(object_file_path[1])[0] + f"_{int(class_id)}_{i}.jpg"
            print(object_file_path)
            cv2.imwrite(object_file_path, object_image)