TX2之多线程读取视频及深度学习推理

参考：https://jkjung-avt.github.io/camera-caffe-threaded/

1 背景

一般在TX2上部署深度学习模型时，都是读取摄像头视频或传入视频文件进行推理，从视频中抽取帧进行目标检测等任务。对于大点的模型，推理的速度是赶不上摄像头或视频的帧率的，如果我们使用单线程进行处理，即读取一帧检测一帧，推理会堵塞视频的正常传输，表现出来就是摄像头视频有很大的延迟，如果是对实时性要求较高，这种延迟是难以接受的。因此，采用多线程的方法，将视频读取与深度学习推理放在两个线程里，互不影响，达到实时的效果。

2 实现方法

将摄像头的视频读取放入子线程，充当一个生产者的角色，将推理放入主线程，充当消费者的角色，主线程推理完一帧后从子线程提数据，继续推理，下图是原博文的一幅图片，描述了两个线程的关系

3 程序实现

3.1 子线程

子线程程序参考https://github.com/jkjung-avt/tf_trt_models/blob/tf-1.8/utils/camera.py，实现摄像头和视频的读取

"""camera.py

This code implements the Camera class, which encapsulates code to
handle IP CAM, USB webcam or the Jetson onboard camera.  The Camera
class is further extend to take either a video or an image file as
input.
"""

import time
import logging
import threading

import numpy as np
import cv2

def open_cam_rtsp(uri, width, height, latency):
    """Open an RTSP URI (IP CAM)."""
    gst_str = ('rtspsrc location={} latency={} ! '
               'rtph264depay ! h264parse ! omxh264dec ! '
               'nvvidconv ! '
               'video/x-raw, width=(int){}, height=(int){}, '
               'format=(string)BGRx ! videoconvert ! '
               'appsink').format(uri, latency, width, height)
    return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

def open_cam_usb(dev, width, height):
    """Open a USB webcam.

    We want to set width and height here, otherwise we could just do:
        return cv2.VideoCapture(dev)
    """
    gst_str = ('v4l2src device=/dev/video{} ! '
               'video/x-raw, width=(int){}, height=(int){}, '
               'format=(string)RGB ! videoconvert ! '
               'appsink').format(dev, width, height)
    return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

def open_cam_onboard(width, height):
    """Open the Jetson onboard camera.

    On versions of L4T prior to 28.1, you might need to add
    'flip-method=2' into gst_str.
    """
    gst_str = ('nvcamerasrc ! '
               'video/x-raw(memory:NVMM), '
               'width=(int)2592, height=(int)1458, '
               'format=(string)I420, framerate=(fraction)30/1 ! '
               'nvvidconv ! '
               'video/x-raw, width=(int){}, height=(int){}, '
               'format=(string)BGRx ! videoconvert ! '
               'appsink').format(width, height)
    return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

def grab_img(cam):
    """This 'grab_img' function is designed to be run in the sub-thread.
    Once started, this thread continues to grab a new image and put it
    into the global 'img_handle', until 'thread_running' is set to False.
    """
    while cam.thread_running:
        if cam.args.use_image:
            assert cam.img_handle is not None, 'img_handle is empty in use_image case!'
            # keep using the same img, no need to update it
            time.sleep(0.01)  # yield CPU to other threads
        else:
            _, cam.img_handle = cam.cap.read()
            fps = cam.cap.get(cv2.CAP_PROP_FPS)
            time.sleep(1/fps)  # fps = 20hz
            print('time sleep ', 1/fps)
            if cam.img_handle is None:
                logging.warning('grab_img(): cap.read() returns None...')
                break
    cam.thread_running = False

class Camera():
    """Camera class which supports reading images from theses video sources:

    1. Video file
    2. Image (jpg, png, etc.) file, repeating indefinitely
    3. RTSP (IP CAM)
    4. USB webcam
    5. Jetson onboard camera
    """

    def __init__(self, args):
        self.args = args
        self.is_opened = False
        self.thread_running = False
        self.img_handle = None
        self.img_width = 0
        self.img_height = 0
        self.cap = None
        self.thread = None

    def open(self):
        """Open camera based on command line arguments."""
        assert self.cap is None, 'Camera is already opened!'
        args = self.args
        if args.use_file:
            self.cap = cv2.VideoCapture(args.filename)
            # ignore image width/height settings here
        elif args.use_image:
            self.cap = 'OK'
            self.img_handle = cv2.imread(args.filename)
            # ignore image width/height settings here
            if self.img_handle is not None:
                self.is_opened = True
                self.img_height, self.img_width, _ = self.img_handle.shape
        elif args.use_rtsp:
            self.cap = open_cam_rtsp(
                args.rtsp_uri,
                args.image_width,
                args.image_height,
                args.rtsp_latency
            )
        elif args.use_usb:
            self.cap = open_cam_usb(
                args.video_dev,
                args.image_width,
                args.image_height
            )
        else:  # by default, use the jetson onboard camera
            self.cap = open_cam_onboard(
                args.image_width,
                args.image_height
            )
        if self.cap != 'OK':
            if self.cap.isOpened():
                # Try to grab the 1st image and determine width and height
                _, img = self.cap.read()
                if img is not None:
                    self.img_height, self.img_width, _ = img.shape
                    self.is_opened = True

    def start(self):
        assert not self.thread_running
        self.thread_running = True
        self.thread = threading.Thread(target=grab_img, args=(self,))
        self.thread.start()

    def stop(self):
        self.thread_running = False
        self.thread.join()

    def read(self):
        if self.args.use_image:
            return np.copy(self.img_handle)
        else:
            return self.img_handle

    def release(self):
        assert not self.thread_running
        if self.cap != 'OK':
            self.cap.release()

3.2 主线程

主线程程序调用camera类，然后进行open、start、read等操作，read函数返回的是最新一帧的frame，这样可以起到实时的效果，而多余的帧通过python的垃圾回收机制自动回收，不会对系统内存造成影响

def detect_in_video():
    '''其它初始化等操作'''

    cam = Camera(args)
    cam.open()
    cam.start()

    while cam.thread_running:
        frame = cam.read()
        '''后续识别操作'''

def main():
  detect_in_video()

if __name__ =='__main__':
  main()