detect_vedio.py

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