YOLO 텍스트 라벨을 이용하여 Object Crop 이미지 저장

YOLO 텍스트로 된 라벨 bbox를 이용하여 object들을 crop 하여 이미지로 저장하는 코드이다.
확장자나 경로는 사용자에 맞게 바꿔야 한다.

이미지, 라벨 이름의 짝과 개수가 맞는지는 코드 실행 전 검사하여야 한다.

import cv2
import os
import glob

def get_x_y_points(point1_x, point1_y, point2_x, point2_y):
    xmin, ymin, xmax, ymax = 0, 0, 0, 0
    
    if point1_x < point2_x and point1_y < point2_y:
        xmin = point1_x
        ymin = point1_y
        xmax = point2_x
        ymax = point2_y
    elif point1_x > point2_x and point1_y < point2_y:
        xmin = point2_x
        ymin = point1_y
        xmax = point1_x
        ymax = point2_y 
    elif point1_x < point2_x and point1_y > point2_y:
        xmin = point1_x
        ymin = point2_y
        xmax = point2_x
        ymax = point1_y 
    elif point1_x > point2_x and point1_y > point2_y:
        xmin = point2_x
        ymin = point2_y
        xmax = point1_x
        ymax = point1_y
        
    return xmin, ymin, xmax, ymax

def read_yolo_label(label_path):
    with open(label_path, 'r') as f:
        lines = f.readlines()
    labels = []
    for line in lines:
        parts = line.strip().split()
        labels.append([int(parts[0]), float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4])])
    return labels

def crop_objects(image_path, label_path, output_dir):
    # read image
    image = cv2.imread(image_path)
    h, w, _ = image.shape
    
    # read label
    labels = read_yolo_label(label_path)
    
    for idx, label in enumerate(labels):
        class_id, x_center, y_center, width, height = label

        # real image scale
        x_center *= w
        y_center *= h
        width *= w
        height *= h
        
        # calculate real image points
        x_min = int(x_center - width / 2)
        x_max = int(x_center + width / 2)
        y_min = int(y_center - height / 2)
        y_max = int(y_center + height / 2)
        
        # arrange points
        x_min, y_min, x_max, y_max = get_x_y_points(x_min, y_min, x_max, y_max)
        
        # crop image
        cropped_image = image[y_min:y_max, x_min:x_max]
        image_file_name = os.path.basename(label_path)
        image_file_name = os.path.splitext(image_file_name)[0]
        
        # save
        output_path = os.path.join(output_dir, f'{image_file_name}_class_{class_id}_object_{idx}.jpg')
        if len(cropped_image) != 0:
            cv2.imwrite(output_path, cropped_image)

# 경로 변경
source_dir_path = "./source_dir_path"
output_dir = './output_dir_path'

os.makedirs(output_dir, exist_ok=True)

# 확장자 등 변경
image_path_list = glob.glob(os.path.join(source_dir_path, "**", "*.jpg"), recursive=True)
label_path_list = glob.glob(os.path.join(source_dir_path, "**", "*.txt"), recursive=True)

for i, (image_path, label_path) in enumerate(zip(image_path_list, label_path_list)):
    if ((i + 1) % 10 == 0):
        print(f"{i + 1} / {len(image_path_list)}")
    crop_objects(image_path, label_path, output_dir)