-
Notifications
You must be signed in to change notification settings - Fork 16
/
video_output.py
75 lines (59 loc) · 2.71 KB
/
video_output.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#!/usr/bin/env python3
import time
import logging
from nano_llm import Plugin
from nano_llm.utils import cuda_image
from jetson_utils import videoOutput
class VideoOutput(Plugin):
"""
Saves images to a compressed video or directory of individual images, the display, or a network stream.
https://github.com/dusty-nv/jetson-inference/blob/master/docs/aux-streaming.md
"""
def __init__(self, video_output: str="webrtc://@:8554/output",
video_output_codec: str=None, video_output_bitrate: int=None,
video_output_save: str=None, **kwargs):
"""
Output video to a network stream (RTP/RTSP/WebRTC), video file, or display.
Args:
video_output (str): Stream URL, path to video file, or a directory of images.
video_output_codec (str): Force a particular codec (H264, H265, VP8, VP9, MJPEG)
video_output_bitrate (int): The desired bitrate in bits per second (default is 4 Mbps)
video_output_save (str): Save the encoded video to disk in MP4, MKV, AVI, or FLV format.
"""
super().__init__(outputs=0, **kwargs)
options = {}
if video_output_codec:
options['codec'] = video_output_codec
if video_output_bitrate:
options['bitrate'] = video_output_bitrate
if video_output_save:
options['save'] = video_output_save
if video_output is None:
video_output = ''
args = None if 'display://' in video_output else ['--headless']
self.stream = videoOutput(video_output, options=options, argv=args)
self.resource = video_output
self.time_last = time.perf_counter()
self.framerate = 0
def process(self, image, **kwargs):
"""
Input should be a jetson_utils.cudaImage, np.ndarray, torch.Tensor, or have __cuda_array_interface__
"""
image = cuda_image(image)
shape = image.shape
self.stream.Render(image, stream=image.stream)
curr_time = time.perf_counter()
self.framerate = self.framerate * 0.9 + (1.0 / (curr_time - self.time_last)) * 0.1
self.time_last = curr_time
self.send_stats(summary=[f"{shape[1]}x{shape[0]}", f"{self.framerate:.1f} FPS"])
@classmethod
def type_hints(cls):
"""
Return static metadata about the plugin settings.
"""
return dict(
video_output = dict(display_name='Output'),
video_output_codec = dict(display_name='Codec'),
video_output_bitrate = dict(display_name='Bitrate'),
video_output_save = dict(display_name='Save'),
)