File size: 3,495 Bytes
700dd75 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | import os
import queue
import sys
import threading
import time
import av
import numpy as np
class VideoWriter:
def __init__(
self,
output_path: str,
width: int,
height: int,
fps: float,
codec: str = "h264",
buffer_size: int = 50,
):
self.output_path = output_path
self._first_frame = True # Track first frame to suppress x264 info output
# Create output directory if it doesn't exist
output_dir = os.path.dirname(output_path)
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
self.queue = queue.Queue(maxsize=buffer_size)
self.container = av.open(output_path, mode="w")
self.stream = self.container.add_stream(codec, rate=fps)
self.stream.width = width
self.stream.height = height
thread = threading.Thread(target=self._writer_worker, daemon=True)
thread.start()
def _assert_dimensions(self, frame: np.ndarray) -> None:
assert (
frame.shape[1] == self.stream.width and frame.shape[0] == self.stream.height
), f"""Incorrect frame dimensions. Input dimensions: {frame.shape[1]}x{frame.shape[0]}.
Expected dimensions: {self.stream.width}x{self.stream.height}"""
def add_frame(self, frame: np.ndarray) -> None:
self._assert_dimensions(frame)
self.queue.put(frame)
def _writer_worker(self) -> None:
while True:
frame = self.queue.get()
if frame is None:
continue
self._assert_dimensions(frame)
frame = av.VideoFrame.from_ndarray(frame, format="rgb24")
# Suppress stderr for first frame encoding (x264 prints info then)
if self._first_frame:
stderr_fd = sys.stderr.fileno()
old_stderr = os.dup(stderr_fd)
devnull = os.open(os.devnull, os.O_WRONLY)
os.dup2(devnull, stderr_fd)
try:
packets = self.stream.encode(frame)
for packet in packets:
self.container.mux(packet)
finally:
os.dup2(old_stderr, stderr_fd)
os.close(old_stderr)
os.close(devnull)
self._first_frame = False
else:
packets = self.stream.encode(frame)
for packet in packets:
self.container.mux(packet)
def _flush_stream(self) -> None:
packets = self.stream.encode()
for packet in packets:
self.container.mux(packet)
def stop(self) -> str:
"""
Blocking call. Waits until all the frames in the queue have been written to the file
and the video writer has been closed.
"""
if not self.queue.empty():
print("Waiting for video writer queue to empty...")
while not self.queue.empty():
time.sleep(0.1)
print("Video writer queue is empty, flushing stream...")
self._flush_stream()
self.container.close()
return self.output_path
def cancel(self) -> None:
"""Immediately stops writing and deletes the output file"""
if os.path.exists(self.output_path):
os.remove(self.output_path)
self.container.close()
def __del__(self) -> None:
self.container.close()
|