TensorRT-Demo/utils/writer.py

31 lines
942 B
Python
Raw Normal View History

2023-03-06 20:44:29 +08:00
"""writer.py
"""
import subprocess
import cv2
def get_video_writer(name, width, height, fps=30):
"""Get a VideoWriter object for saving output video.
This function tries to use Jetson's hardware H.264 encoder (omxh264enc)
if available, in which case the output video would be a MPEG-2 TS file.
Otherwise, it uses cv2's built-in encoding mechanism and saves a MP4
file.
"""
gst_elements = str(subprocess.check_output('gst-inspect-1.0'))
if 'omxh264dec' in gst_elements:
filename = name + '.ts' # Transport Stream
gst_str = ('appsrc ! videoconvert ! omxh264enc ! mpegtsmux ! '
'filesink location=%s') % filename
return cv2.VideoWriter(
gst_str, cv2.CAP_GSTREAMER, 0, fps, (width, height))
else:
filename = name + '.mp4' # MP4
return cv2.VideoWriter(
filename, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))