天天看點

【目标檢測】小腳本:YOLO标簽可視化

需求分析

在下載下傳别人标注好的目标檢測資料集時,我突然想到一個問題:怎麼直覺得看别人标注的是否正确呢?于是我想到了可以利用​

​opencv​

​​将标注資料還原到原圖上。

更具體的說,指定圖檔和标簽檔案夾,批量輸出還原後的圖檔。

需求實作

import os
import numpy as np
import cv2

# 修改輸入圖檔檔案夾
img_folder = "images/"
img_list = os.listdir(img_folder)
img_list.sort()
# 修改輸入标簽檔案夾
label_folder = "labels/"
label_list = os.listdir(label_folder)
label_list.sort()
# 輸出圖檔檔案夾位置
path = os.getcwd()
output_folder = path + '/' + str("output")
os.mkdir(output_folder)
# labels = ['truck', 'panzer', 'tank', 'SUV', 'cam_net', 'cam_tar']
colormap = [(0, 255, 0), (132, 112, 255), (0, 191, 255)]  # 色盤,可根據類别添加新顔色


# 坐标轉換
def xywh2xyxy(x, w1, h1, img):
    label, x, y, w, h = x
    # print("原圖寬高:\nw1={}\nh1={}".format(w1, h1))
    # 邊界框反歸一化
    x_t = x * w1
    y_t = y * h1
    w_t = w * w1
    h_t = h * h1
    # print("反歸一化後輸出:\n第一個:{}\t第二個:{}\t第三個:{}\t第四個:{}\t\n\n".format(x_t, y_t, w_t, h_t))
    # 計算坐标
    top_left_x = x_t - w_t / 2
    top_left_y = y_t - h_t / 2
    bottom_right_x = x_t + w_t / 2
    bottom_right_y = y_t + h_t / 2

    # print('标簽:{}'.format(labels[int(label)]))
    # print("左上x坐标:{}".format(top_left_x))
    # print("左上y坐标:{}".format(top_left_y))
    # print("右下x坐标:{}".format(bottom_right_x))
    # print("右下y坐标:{}".format(bottom_right_y))
    # 繪制矩形框
    cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), colormap[1], 2)
    """
    # (可選)給不同目标繪制不同的顔色框
    if int(label) == 0:
       cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (0, 255, 0), 2)
    elif int(label) == 1:
       cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (255, 0, 0), 2)
    """
    return img


if __name__ == '__main__':
    for i in range(len(img_list)):
        image_path = img_folder + "/" + img_list[i]
        label_path = label_folder + "/" + label_list[i]
        # 讀取圖像檔案
        img = cv2.imread(str(image_path))
        h, w = img.shape[:2]
        # 讀取 labels
        with open(label_path, 'r') as f:
            lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)
        # 繪制每一個目标
        for x in lb:
            # 反歸一化并得到左上和右下坐标,畫出矩形框
            img = xywh2xyxy(x, w, h, img)
        """
        # 直接檢視生成結果圖
        cv2.imshow('show', img)
        cv2.waitKey(0)
        """
        cv2.imwrite(output_folder + '/' + '{}.png'.format(image_path.split('/')[-1][:-4]), img)      

檔案結構