代碼實現把目標檢測數據集的框框畫在相應圖片上

  衆所周知，目標檢測的數據集由兩部分組成，一部分是圖片，另一部分當然就是圖片對應的標籤了。這裏的標籤就是圖片中每一個物體的邊框了。在數據集中，每一個框框是以五個數字組成的，分別是x、y、w、h，以及類別。

  打開數據集，我不太喜歡這些數字，我想要在圖片上面框框這些框框，下面直接上代碼。

import random
import colorsys
import matplotlib.pyplot as plt
from datetime import datetime
import os
import cv2
import numpy as np
import tensorflow as tf
import shutil


#原圖像需要放在相應路徑下
#生成帶框圖像路徑
out_image_path = '填入你的路徑'
#訓練集txt文件路徑
data_path = '填入你的路徑'
#類別名稱txt文件路徑
class_name_path = '填入你的路徑'

def draw_bbox_new(image, bboxes, classes=None, show_label=True):

    if classes is None:
        classes = read_class_names(class_name_path)
    class_num = len(classes)
    image_h, image_w, _ = image.shape
    hsv_tuples = [(1.0 * x / class_num, 1., 1.) for x in range(class_num)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))
    random.seed(32)
    random.shuffle(colors)
    random.seed(None)
    for i, bbox in enumerate(bboxes):
        coor = np.array(bbox[:4], dtype=np.int32)
        class_ind = int(bbox[4])
        bbox_color = colors[class_ind]
        bbox_thick = 3
        c1 = (coor[0], coor[1])
        c2 = (coor[2], coor[3])
        cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)
        if show_label:
            bbox_mess = '%s: %.2f' % (classes[class_ind], 1)
            t_size = cv2.getTextSize(bbox_mess, 0, 0.5, thickness=1 )[0]
            cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)  # filled
            cv2.putText(image, bbox_mess, (c1[0], c1[1] - 2), cv2.FONT_HERSHEY_SIMPLEX,0.5, (0, 0, 0), 1, lineType=cv2.LINE_AA)

    return image

def read_class_names(class_file_name):
    names = {}
    with open(class_file_name, 'r') as data:
        for name in data:
            names[ID] = name.strip('\n')
    return names

image_list = []
with open(data_path)as file:
    pictures = file.readlines()
    for i in pictures:
        image = i.strip().split()
        image_list.append(image[0])

gt_list = []
with open(data_path)as file:
    pictures = file.readlines()
    for i in pictures:
        gt_list.append(i.strip().split()[1:])

gt = []
for i in gt_list:
    tem = []
    for j in i:
        tem.append(j.split(','))
    gt.append(tem)
    
for index,image_path in enumerate(image_list):
    image_info = cv2.imread(image_path)
    new_image = draw_bbox_new(image_info, gt[index])
    new_image = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)
    plt.imshow(new_image)
    plt.show()
#     img_name = image_path.split('/')[-1]
#     cv2.imwrite(os.path.join(out_image_path,img_name), new_image)

生成圖片如下：