【目标检测】绘图——将目标框绘制在图像（PIL）

前言

近期在做目标检测的项目，除了算法流程本身，我发现可视化也能极大的加深对整个pipeline的理解，因此想借这篇博客记录一下在目标检测项目中的绘图注意点。

本文使用PIL，能够在原图上绘制目标框以及label_txt，并且框与文字的大小能够自适应原图像的大小。

准备

1. 图像集
2. 标注txt文件，
   格式：img_path box1(x1,y1,x2,y2,cls_id) box2(x1,y1,x2,y2,cls_id)...
   
3. class的txt文件，每一行为class_name
   

绘制代码

PIL

def get_classes(classes_path):
    with open(classes_path, 'r', encoding='utf-8') as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names, len(class_names)
  
def draw_annotation(img_path, annot_path, save_path, classes_path, batch=False):
"""
	batch=True时表示批量处理图像，False表示绘制单张图像并且显示
	img_path: 如果batch=True，则为图像集的dir_path；否则表示单张图像的path
	annot_path： 标注`txt`文件的路径
	save_path：当batch=True才有效，表示绘制框后的图像保存的dir_path
	classes_path: `class`文件的path
"""
    class_names, num_classes = get_classes(classes_path)
    
    # 设置不同类别的框的颜色（r,g,b）
    hsv_tuples = [(x / num_classes, 1., 1.) for x in range(num_classes)]
    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))
    if batch:
        print('start draw ground truth in batch way')
        annot = open(annot_path, 'r', encoding='UTF-8').readlines()
        num = 0

        for line in tqdm(annot):
            line = line.split()

            image = Image.open(line[0],)
            if np.shape(image) != 3 or np.shape(image)[2] != 3:
                image = image.convert('RGB')
            img_name = os.path.basename(line[0])
			
			# 设置字体和框的厚度，会根据原始图像的大小改变
            font = ImageFont.truetype(font=r'C:\Windows\Fonts/Arial.ttf',
                                      size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
            thickness = int(
                max((image.size[0] + image.size[1]) // np.mean(np.array(image.size[:2])), 1))

            for box in line[1:]:
                left, top, right, bottom, cls_id= box.split(',')

                top     = max(0, int(top))
                left    = max(0, int(left))
                bottom  = min(image.size[1], int(bottom))
                right   = min(image.size[0], int(right))

                label = '{}'.format(class_names[int(cls_id)])
                draw = ImageDraw.Draw(image)
                label_size = draw.textsize(label, font)
                label = label.encode('utf-8')
                print(label, left, top, right, bottom)

                if top - label_size[1] >= 0:
                    text_origin = np.array([left, top - label_size[1]])
                else:
                    text_origin = np.array([left, top + 1])
                
                for i in range(thickness):
                    draw.rectangle([left + i, top + i, right - i, bottom - i], outline=colors[int(cls_id)])
                draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=colors[int(cls_id)])
                draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
                del draw
            if not os.path.exists(save_path):
                os.makedirs(os.path.join(save_path), exist_ok=True)
            image.save(os.path.join(save_path, img_name))
            num += 1
        print('draw {} ground truth in batch way done!'.format(num))
    else:
        img_name = os.path.basename(img_path)
        image = Image.open(img_path)
        annot = open(annot_path, 'r').readlines()

        font = ImageFont.truetype(font='model_data/simhei.ttf',
                                    size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness = int(
            max((image.size[0] + image.size[1]) // np.mean(np.array(image.size[:2])), 1))

        for line in annot:
            line = line.split()
            if os.path.basename(line[0]) == img_name:
                for box in line[1:]:
                    left, top, right, bottom, cls_id= box.split(',')

                    top     = max(0, int(top))
                    left    = max(0, int(left))
                    bottom  = min(image.size[1], int(bottom))
                    right   = min(image.size[0], int(right))

                    label = '{}'.format(class_names[int(cls_id)])
                    draw = ImageDraw.Draw(image)
                    label_size = draw.textsize(label, font)
                    label = label.encode('utf-8')
                    print(label, top, left, bottom, right)

                    if top - label_size[1] >= 0:
                        text_origin = np.array([left, top - label_size[1]])
                    else:
                        text_origin = np.array([left, top + 1])

                    for i in range(thickness):
                        draw.rectangle([left + i, top + i, right - i, bottom - i], outline=colors[int(cls_id)])
                    draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=colors[int(cls_id)])
                    draw.text(text_origin, str(label,'UTF-8'), fill=(0,0,0), font=font)
                    del draw
                
                image.show()

font = ImageFont.truetype(font=r'C:\Windows\Fonts\Arial.ttf',
      size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))

对于windows系统，C:\Windows\Fonts里很多系统自带的字体可供选择。

另外python可以识别的颜色，比如可以：

draw.rectangle([left + i, top + i, right - i, bottom - i], 
outline='red') # colors[int(cls_id)]

效果图