目录
1. OpenCV
2. matplotlib
这里我们用YOLOv3COCO的模型来给照片做检测。我们用tensornets导入预训练的YOLO模型,用OpenCV给照片添加方框和文字,用matplotlib显示图片。
tensornets提供对各种预训练模型的高层接口,包括YOLOv2、YOLOv3、Faster RCNN。
OpenCV中的cv2.rectangle和cv2.putText两个函数分别给图片添加方框和文字。请点击超链接查看两个函数的文档。
我们输入这样的图片
用下面这些代码处理上图。
import tensorflow as tf
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2 as cv
import tensornets as nets
from tensornets.datasets.coco import classnames
# 输入占位符
input_img = tf.placeholder(shape=(None, 416, 416, 3), dtype=tf.float32)
# 实例化YOLO模型
yolov3 = nets.YOLOv3COCO(input_img)
img_path = './road.bmp'
# 读入图片
image = Image.open(img_path)
# 放缩图片到指定大小
image_resize = image.resize((416, 416))
image_resize = np.expand_dims(np.asarray(image_resize), 0)
# TensorFlow的配置
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# 运行YOLO模型获得各个物体类别的方框及其可信度
with tf.Session(config=config) as sess:
sess.run(yolov3.pretrained())
res = sess.run(yolov3.preds, {input_img: yolov3.preprocess(image_resize)})
boxes = yolov3.get_boxes(res, image_resize.shape[1:3])
# MSCOCO中有80种物体类别
print(len(boxes))
# 图片的第一维是数组里的y坐标,第二维是数组里的x坐标
width, height = image.size
# 复制图片
image2 = np.copy(np.asarray(image))
for idx, box in enumerate(boxes):
if box.shape[0] > 0:
for line in box:
x1, y1, x2, y2, score = line
# 只显示可信度超过0.4的方框
if score > 0.4:
x1 *= width / 416
y1 *= height / 416
x2 *= width / 416
y2 *= height / 416
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
# 画框
image2 = cv.rectangle(image2, (x1, y1), (x2, y2), (255, 0, 0), thickness=3)
# 上文字
image2 = cv.putText(image2, classnames[idx], (x1, y1-10), cv. FONT_HERSHEY_TRIPLEX, (x2-x1)/120, (0, 255, 0), 2)
# 显示图片
plt.figure(figsize=(20,20))
plt.imshow(image2)
最终获得的图片类似于这样:
可以看到YOLO找到了两辆car、一个stop sign和一辆truck,其中一辆car和truck重叠。
matplotlib也可以给图片加上方框。这里我们要使用matplotlib.pyplot.text函数和matplotlib.patches.Rectangle类。
class
matplotlib.patches.
Rectangle
(xy, width, height, angle=0.0, **kwargs)[source]Bases:
matplotlib.patches.Patch
A rectangle with lower left at xy = (x, y) with specified width, height and rotation angle.
Parameters: xy : (float, float)
The bottom and left rectangle coordinates
width : float
Rectangle width
height : float
Rectangle height
angle : float, optional
rotation in degrees anti-clockwise about xy (default is 0.0)
fill : bool, optional
Whether to fill the rectangle (default is
True
)
Rectangle是继承了Patch的类。它初始化的第一个参数是矩形左下角的坐标,这里的xy是一般数学意义上的xy,也就是横坐标在前,纵坐标在后,这与计算机领域数组的表示恰好相反。它第二个参数是宽度,第三个参数是高度。另外你还可以给它传递关键词参数,比如指定线宽的linewidth,边缘颜色的edgecolor和表面颜色facecolor。
matplotlib.pyplot.
text
(x, y, s, fontdict=None, withdash=, **kwargs)[source] Add text to the axes.
Add the text s to the axes at location x, y in data coordinates.
Parameters: x, y : scalars
The position to place the text. By default, this is in data coordinates. The coordinate system can be changed using the transform parameter.
s : str
The text.
fontdict : dictionary, optional, default: None
A dictionary to override the default text properties. If fontdict is None, the defaults are determined by your rc parameters.
withdash : boolean, optional, default: False
Creates a
TextWithDash
instance instead of aText
instance.Returns: text :
Text
The created
Text
instance.Other Parameters: **kwargs :
Text
properties.Other misc
text是个函数,用于在当前axes上添加文字。matplotlib用figure表示单独一张图,axes分别表示图中若干个子图。作图操作直接作用于axes上。text的第一个参数是横坐标,第二个参数是纵坐标,第三个参数是要添加的字符串,第四个参数是指定字体。除此以外你还可以指定关键词参数,比如文字颜色color,垂直对齐verticalalignment,文字框参数bbox。
下面是一个例子
for cls, score, box in zip(d['detection_classes'], d['detection_scores'], d['detection_boxes']):
bbox = patches.Rectangle((box[1]*w, box[0]*h),
(box[3]-box[1])*w, (box[2]-box[0])*h,
linewidth=1,
edgecolor=idx2color[cls],
facecolor='none')
plt.gca().add_patch(bbox)
plt.text(box[1]*w,
box[0]*h, s='%s %.2f'%(idx2cls[cls], score),
color='white',
verticalalignment='bottom',
bbox={'color':idx2color[cls], 'pad': 0})
在一次检测后我们给图像这样添加方框可以得到这样的效果: