Skip to content

Commit

Permalink
Fix writers so RecordAction works
Browse files Browse the repository at this point in the history
  • Loading branch information
daniilpastukhov committed Aug 2, 2023
1 parent 952e33b commit 908468f
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 8 deletions.
29 changes: 24 additions & 5 deletions depthai_sdk/src/depthai_sdk/recorders/video_writers/base_writer.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import time
from abc import ABC
from collections import deque
from collections import deque, defaultdict
from pathlib import Path
from typing import Dict

Expand All @@ -10,27 +11,45 @@ def __init__(self, path: Path, name: str):
self.name = name

self._buffers: Dict[str, deque] = {}
self._buffers_max_seconds: Dict[str, int] = {} # in seconds
self._buffers_timestamps = defaultdict(list)
self._buffers_approx_fps: Dict[str, float] = {}
self._file = None
self._fps = None

def create_file_for_buffer(self, subfolder: str, bufname: str):
raise NotImplementedError()

def init_buffer(self, name: str, max_seconds: int):
if max_seconds > 0:
self._buffers[name] = deque(maxlen=int(max_seconds * self._fps))
self._buffers[name] = deque()
self._buffers_max_seconds[name] = max_seconds

def add_to_buffer(self, name: str, frame):
if self._buffers[name] is None:
return

if len(self._buffers[name]) == self._buffers[name].maxlen:
timestamp = time.time()
self._buffers_timestamps[name].append(timestamp)

# Calculate time window based on max_seconds
time_window = self._buffers_max_seconds[name]

# Remove frames that fall outside the time window
while self._buffers_timestamps[name] and (timestamp - self._buffers_timestamps[name][0] > time_window):
self._buffers[name].popleft()
self._buffers_timestamps[name].pop(0)

self._buffers[name].append(frame)

def is_buffer_full(self, name: str) -> bool:
return len(self._buffers[name]) == self._buffers[name].maxlen
if self._buffers[name].maxlen:
return len(self._buffers[name]) == self._buffers[name].maxlen

if not self._buffers_timestamps[name]:
return False

diff = self._buffers_timestamps[name][0] + self._buffers_max_seconds[name] - self._buffers_timestamps[name][-1]
return diff < 0.1

def is_buffer_empty(self, name: str) -> bool:
return len(self._buffers[name]) == 0
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,7 @@ def __init__(self, path: Path, name: str, lossless: bool = False):
Args:
path: Path to save the output. Either a folder or a file.
name: Name of the stream.
fourcc: FourCC code of the codec used to compress the frames.
fps: Frames per second.
lossless: If True, save the stream without compression.
"""

super().__init__(path, name)
Expand All @@ -33,6 +32,16 @@ def __init__(self, path: Path, name: str, lossless: bool = False):
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()

def create_file_for_buffer(self, subfolder: str, buf_name: str):
if self._buffers[buf_name] is None:
raise RuntimeError(f"Buffer {buf_name} is not enabled")

if len(self._buffers[buf_name]) == 0:
return None

frame = self._buffers[buf_name][0]
self.create_file(subfolder, frame)

def create_file(self, subfolder: str, frame: dai.ImgFrame):
if self._lossless or frame.getType() == dai.ImgFrame.Type.RAW16:
extension = 'avi'
Expand All @@ -58,7 +67,6 @@ def _create_file(self, path_to_file: str, frame: dai.ImgFrame):
self._fourcc = 'h264'
options['crf'] = '15'


self._file = av.open(path_to_file, 'w')
self._stream = self._file.add_stream(self._fourcc)
self._stream.options = options
Expand Down

0 comments on commit 908468f

Please sign in to comment.