如何用Python在图片上绘制BoundingBox
参考资料:
https://blog.csdn.net/weixin_41735859/article/details/106599903
在目标检测等CV领域的任务里,经常会涉及到在图片上绘制BBox,也就是画一个矩形框把目标框起来,并且还可能会打上标签。这篇随笔记录一下在Python里如何在图片上绘制BoundingBox。
我的主要参考对象是参考资料里的方法,参考资料里给出了opencv和PIL来绘制的方法,考虑到大多数机器学习环境里都会有PIL/pillow(这玩意儿都快和python绑定了吧,是python原生的),所以我就记录一下用PIL绘制的方法(不用装新的包嘛!)废话不多说,直接放代码和效果图,并且附上详细的注释:
from PIL import Image, ImageDraw, ImageFont # 1.读取图片 real_im = T.ToPILImage()(XXX_1) bbox_list = XXX_2 label_list = XXX_3 # 2.设置字体格式及大小 font = ImageFont.truetype(font='./Courier-12.ttf', size=np.floor(1.5e-2 * np.shape(real_im)[1] + 10).astype('int32')) # 3.给画笔添加句柄(我就是要在real_im上画画!) draw = ImageDraw.Draw(real_im) for (label_, bbox_) in zip(label_list, bbox_list): # 4.获取边框坐标 # 最终边框格式 bbox = [xl, yl, xr, yr] bbox_[0], bbox_[1], bbox_[2], bbox_[3] = bbox_[0], bbox_[1], bbox_[0] + bbox_[2], bbox_[1] + bbox_[3] bbox_[0] = int(bbox_[0] * real_im.size[0]) bbox_[1] = int(bbox_[1] * real_im.size[1]) bbox_[2] = int(bbox_[2] * real_im.size[0]) bbox_[3] = int(bbox_[3] * real_im.size[1]) # 5.获取label长宽 label_size = draw.textsize(label_, font) # 6.设置label起点 text_origin = np.array([bbox_[0], bbox_[1] + 0.2 * label_size[1]]) # 7.随机一个RGB值作颜色 color_for_draw = tuple(np.random.randint(0, 255, size=[3])) # 8. 绘制矩形框,加入label文本 draw.rectangle([bbox_[0], bbox_[1], bbox_[2], bbox_[3]],outline=color_for_draw, width=2) draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=color_for_draw) draw.text(text_origin, str(label_), fill=(255, 255, 255), font=font) # 9.删除画笔句柄 del draw
下面说一下几个要点:
- XXX_1这个地方需要是一个PIL.Image格式的图片,无论你是用torchvision里面的ToPIL搞出来,还是用Image.open打开图片文件,都需要是PIL.Image格式
- XXX_2这个地方是一个BBox的列表,众所周知需要四个值才能确定一个矩形,所以它的形状可能是(n, 4)的array。
- XXX_3这个地方是Label的列表,内容就是str,也就是我们想在对应BBox旁边打上的文字,它的形状可能是(n,)的array或者就是list。
- 这段代码最终生效的BBox是xl, yl, xr, yr格式,坐标系参考资料里面有,如果你的BBox做了归一化,或者是其他格式,你需要相应的调整。(我展示的代码里bbox值一开始是归一化的,并且是COCO的格式,[xl,yl,w,h])
- 上述代码可能不好直接跑通,包括ttf字体还得自己下载,但是希望你能读到并理解这个绘制的方法,就好。
最终的效果如下两张图所示,因为是128的分辨率所以看着别扭,label框的大小需要调整,但是可以看到确实能够很方便地把BoundingBox在图片上画出来。