From 908468f74d5c4186f1c9e27761a6af9ba9df0cca Mon Sep 17 00:00:00 2001
From: Daniil Pastukhov <pastudan@fit.cvut.cz>
Date: Wed, 2 Aug 2023 17:00:51 +0200
Subject: [PATCH] Fix writers so RecordAction works

---
 .../recorders/video_writers/base_writer.py    | 29 +++++++++++++++----
 .../recorders/video_writers/video_writer.py   | 14 +++++++--
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/depthai_sdk/src/depthai_sdk/recorders/video_writers/base_writer.py b/depthai_sdk/src/depthai_sdk/recorders/video_writers/base_writer.py
index d35bb4344..1bcf51d5b 100644
--- a/depthai_sdk/src/depthai_sdk/recorders/video_writers/base_writer.py
+++ b/depthai_sdk/src/depthai_sdk/recorders/video_writers/base_writer.py
@@ -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
 
@@ -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
diff --git a/depthai_sdk/src/depthai_sdk/recorders/video_writers/video_writer.py b/depthai_sdk/src/depthai_sdk/recorders/video_writers/video_writer.py
index 4321df7ce..30ac0af8d 100644
--- a/depthai_sdk/src/depthai_sdk/recorders/video_writers/video_writer.py
+++ b/depthai_sdk/src/depthai_sdk/recorders/video_writers/video_writer.py
@@ -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)
@@ -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'
@@ -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