YOLO 텍스트 라벨을 이미지에 bbox 표시, 시각화

YOLO 텍스트로 된 라벨 bbox를 이미지에 표시하는 샘플 코드이다.

다른 용도로 변형 가능하며, 경로는 사용자에 맞게 바꿔야 한다.

import os
import glob
import cv2
import shutil
import numpy as np


# opencv 한글 경로 읽을 수 있도록
def imread(file_path):
    f = open(file_path.encode("utf-8"), "rb")
    bytes = bytearray(f.read())
    npArr = np.asarray(bytes, dtype=np.uint8)

    return cv2.imdecode(npArr, cv2.IMREAD_UNCHANGED)


# image, label 짝 체크
def pair_img_label_check(img_path, label_path, log_path="./", log_name="none_pair_data.txt"):
    img_ext = ["jpg", "jpeg", "png"]
    img_path_list = []
    check = True

    for ext in img_ext:
        temp_img_path_list = glob.glob(os.path.join(img_path, "*." + ext))
        for temp_img_path in temp_img_path_list:
            img_path_list.append(temp_img_path)

    label_path_list = glob.glob(os.path.join(label_path, "*.txt"))

    img_name_list = [os.path.split(os.path.splitext(img_path)[0])[1] for img_path in img_path_list]
    label_name_list = [os.path.split(os.path.splitext(label_path)[0])[1] for label_path in label_path_list]

    img_name_dict = {img_name: 1 for img_name in img_name_list}
    label_name_dict = {label_name: 1 for label_name in label_name_list}

    for img_name in img_name_list:
        if not img_name in label_name_dict:
            print("none pair img:", img_name)
            with open(log_path + log_name, "a") as f:
                f.write(f"img: {img_name}\n")
            check = False

    for label_name in label_name_list:
        if not label_name in img_name_dict:
            print("none pair label:", label_name)
            with open(log_path + log_name, "a") as f:
                f.write(f"label: {label_name}\n")
            check = False

    return check, img_path_list, label_path_list


# YOLO text label xmin, ymin, xmax, ymax로 변환
def calculate_points(image_width: int, image_height: int, x_center_scaling: float, y_center_scaling:float,
                     w_scaling:float, h_scaling:float):
    w = w_scaling * image_width
    h = h_scaling * image_height

    xmin = (x_center_scaling * image_width) - int(w / 2)
    ymin = (y_center_scaling * image_height) - int(h / 2)
    xmax = (x_center_scaling * image_width) + int(w / 2)
    ymax = (y_center_scaling * image_height) + int(h / 2)

    return xmin, ymin, xmax, ymax

# 경로 변경 필요
origin_img_path = "/image_path"
origin_label_path = "/label_path"

log_path = "/log_path"

# dest_path_sub_dir = ["images", "labels"]

pair_check, img_path_list, label_path_list = pair_img_label_check(origin_img_path, origin_label_path, log_path=log_path)
if pair_check != True:
    exit()

img_path_list.sort()
label_path_list.sort()

count = 0
for img_path, label_path in zip(img_path_list, label_path_list):
    count += 1
    if count % 100 == 0:
        print(f"{count} / {len(img_path_list)}")

    img = imread(img_path)

    with open (label_path, "r") as f:
        lines = f.readlines()

        if len(lines) == 0 or lines[0].replace("\n", "") == "":
            print("Empty label data: ", label_path)

        else:
            h, w, c = img.shape

            for line in lines:
                object_data = line.replace("\n", "").split(" ")
                xmin, ymin, xmax, ymax = calculate_points(w, h, float(object_data[1]), float(object_data[2]), float(object_data[3]), float(object_data[4]))
                img = cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 255, 0), 3)

    cv2.imshow(label_path, img)
    key = cv2.waitKey(0)

    if key == ord('q'):
        exit()

    cv2.destroyAllWindows()