import time
from absl import app, flags, logging
from absl.flags import FLAGS
import cv2
import tensorflow as tf
from yolov3_tf2.models import (
YoloV3, YoloV3Tiny
)
from yolov3_tf2.dataset import transform_images
from yolov3_tf2.utils import draw_outputs
#flags.DEFINE_XXX 表示增加为命令行可选参数,括号内三个参数分别为变量名、默认值、提示信息。
flags.DEFINE_string('classes', './data/coco.names', 'path to classes file') #确定类别名classes的路径
flags.DEFINE_string('weights', './checkpoints/yolov3.tf', #确定权重文件yolov3.tf的路径
'path to weights file')
flags.DEFINE_boolean('tiny', False, 'yolov3 or yolov3-tiny')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_string('video', '/home/yaoqing/下载/数据集/vehicle.mp4',
'path to video file or number for webcam)')
# flags.DEFINE_string('video', './data/video.mp4',
# 'path to video file or number for webcam)')
flags.DEFINE_string('output', '/home/yaoqing/下载/数据集/vehicle_ed.mp4', 'path to output video')
flags.DEFINE_string('output_format', 'XVID', 'codec used in VideoWriter when saving video to file')
flags.DEFINE_integer('num_classes', 80, 'number of classes in the model')
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True) #来打开内存增长,它试图只分配运行时所需的GPU内存:它开始分配非常少的内存,随着程序运行和更多的GPU内存需要
if FLAGS.tiny:
yolo = YoloV3Tiny(classes=FLAGS.num_classes)
else:
yolo = YoloV3(classes=FLAGS.num_classes)
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded')
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
times = []
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
out = None
if FLAGS.output:
# by default VideoCapture returns float instead of int
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height)) #cv2.VideoWriter函数是保存视频,四个参数分别为:保存路径、编码器(把视频或者音频按照编码格式,编码成特定文件格式)、帧率、保存画面大小。
while True:
_, img = vid.read()#读入每一帧画面
if img is None: #空画面则输出“Empty Frame”,睡眠0.1s,然后退出循环。
logging.warning("Empty Frame")
time.sleep(0.1)
break#continue
img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #将读进来的BGR格式的图片转化为RGB
img_in = tf.expand_dims(img_in, 0) #0代表的第一维度,1代表第二维度
img_in = transform_images(img_in, FLAGS.size)
t1 = time.time() #当前时间的时间戳
boxes, scores, classes, nums = yolo.predict(img_in)
t2 = time.time()
times.append(t2-t1)
times = times[-20:]
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
img = cv2.putText(img, "Time: {:.2f}ms".format(sum(times)/len(times)*1000), (0, 30),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2) #cv2.putText()中七个参数分别为:图片,添加的文字,左上角坐标,字体,字体大小,颜色,字体粗细
if FLAGS.output:
out.write(img)
cv2.imshow('output', img)#显示图像,两个参数分别是:1.窗口的名字,2.显示的图像
if cv2.waitKey(1) == ord('q'):#键盘输入'q',退出循环
break
cv2.destroyAllWindows()#销毁所有窗口
if __name__ == '__main__':
try:
app.run(main)
except SystemExit:
pass