众所周知,目标检测的数据集由两部分组成,一部分是图片,另一部分当然就是图片对应的标签了。这里的标签就是图片中每一个物体的边框了。在数据集中,每一个框框是以五个数字组成的,分别是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)
生成图片如下: