【LogoDetection 数据集处理】（2）画出训练集图片的标注框

数据集介绍参考：【LogoDetection 数据集处理】（1）将数据集切分为训练集和验证集

为了查看训练集标注的情况，我们需要画出训练集图片的标注框。代码如下：

"""
画出训练集图片的标注框
"""""

import os
import json
import cv2
import numpy as np
from tqdm import tqdm


# _COLORS有70个元素，如果类别数大于70，则_COLORS需要增加
_COLORS = np.array(
    [
        0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.929, 0.694, 0.125,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.300, 0.300, 0.300,
        0.600, 0.600, 0.600,
        1.000, 0.000, 0.000,
        1.000, 0.500, 0.000,
        0.749, 0.749, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 1.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        1.000, 0.667, 0.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.286, 0.286, 0.286,
        0.429, 0.429, 0.429,
        0.571, 0.571, 0.571,
        0.714, 0.714, 0.714,
        0.857, 0.857, 0.857,
        0.000, 0.447, 0.741,
        0.314, 0.717, 0.741,
        0.50, 0.5, 0
    ]
).astype(np.float32).reshape(-1, 3)


# 在训练集图片上标注bbox
def visual(img_path, bbox, cls_id, class_names):
    img = cv2.imread(img_path)
    color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
    text = '{}'.format(class_names[cls_id - 1])
    txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
    font = cv2.FONT_HERSHEY_SIMPLEX
    txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
    txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()

    x_min = bbox[0]
    y_min = bbox[1]
    width = bbox[2]
    height = bbox[3]
    cv2.rectangle(img, (x_min, y_min), (x_min + width, y_min + height), color, thickness=2)
    cv2.rectangle(img, (x_min, y_min - int(1.5 * txt_size[1])), (x_min + txt_size[0] + 1, y_min), txt_bk_color, -1)
    cv2.putText(img, text, (x_min, y_min - 3), font, 0.4, txt_color, thickness=1)

    cv2.imwrite(img_path, img)


if __name__=='__main__':
    # 原始数据路径
    data_path = "dataset/fewshotlogodetection_round1_train_202204/train"
    annoations_path = os.path.join(data_path, "annotations/instances_train2017.json")
    images_path = os.path.join(data_path, "images")

    # 读取annoations的json文件
    with open(annoations_path, 'r', encoding='utf-8') as f:
        annoations_dict = json.load(f)
    images_list = annoations_dict["images"]
    annotations_list = annoations_dict["annotations"]
    categories_list = annoations_dict["categories"]

    # 可视化
    images_name = os.listdir(images_path)
    for image_name in tqdm(images_name,desc="visual process"):
        image_path=os.path.join(images_path,image_name)
        for i in range(len(images_list)):
            pic_file_name = os.path.basename(image_path)
            if pic_file_name == images_list[i]["file_name"]:
                image_id = images_list[i]["id"]
                for j in range(len(annotations_list)):
                    if image_id == annotations_list[j]["image_id"]:
                        bbox = annotations_list[j]["bbox"]
                        cls_id = annotations_list[j]["category_id"]
                        visual(image_path, bbox, cls_id, categories_list)

看一下可视化的结果：