From f427b6bb3bd01659559bd059234ee2dcb8fb55e0 Mon Sep 17 00:00:00 2001
From: Ruben van de Ven <git@rubenvandeven.com>
Date: Tue, 27 May 2025 17:37:51 +0200
Subject: [PATCH] Video recorder in frame emitter (for gige video)

---
 trap/cv_renderer.py      |  6 +++---
 trap/frame_emitter.py    | 42 ++++++++++++++++++++++++++--------------
 trap/preview_renderer.py |  2 +-
 trap/video_sources.py    |  2 --
 4 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/trap/cv_renderer.py b/trap/cv_renderer.py
index 399d288..3c76f8a 100644
--- a/trap/cv_renderer.py
+++ b/trap/cv_renderer.py
@@ -371,9 +371,9 @@ def decorate_frame(frame: Frame, tracker_frame: Frame, prediction_frame: Frame,
         for option, value in prediction_frame.log['predictor'].items():
             options.append(f"{option}: {value}")
         
-
-    cv2.putText(img, options.pop(-1), (20,img.shape[0]-30),  cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
-    cv2.putText(img, " | ".join(options), (20,img.shape[0]-10),  cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
+    if len(options):
+        cv2.putText(img, options.pop(-1), (20,img.shape[0]-30),  cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
+        cv2.putText(img, " | ".join(options), (20,img.shape[0]-10),  cv2.FONT_HERSHEY_PLAIN, 1, base_color, 1)
 
     return img
 
diff --git a/trap/frame_emitter.py b/trap/frame_emitter.py
index 21e676f..7e46895 100644
--- a/trap/frame_emitter.py
+++ b/trap/frame_emitter.py
@@ -1,18 +1,15 @@
 from __future__ import annotations
 
 import logging
-import multiprocessing
 import pickle
 from argparse import ArgumentParser, Namespace
 from multiprocessing import Event
 from pathlib import Path
 
-import zmq
-
 from trap import node
 from trap.base import *
 from trap.base import LambdaParser
-from trap.timer import Timer
+from trap.preview_renderer import FrameWriter
 from trap.video_sources import get_video_source
 
 logger = logging.getLogger('trap.frame_emitter')
@@ -39,20 +36,31 @@ class FrameEmitter(node.Node):
         offset = int(self.config.video_offset or 0)
         source = get_video_source(self.video_srcs, self.config.camera, offset, self.config.video_end, self.config.video_loop)
         video_gen = enumerate(source, start = offset)
-        while self.run_loop():
 
-            try:
-                i, img = next(video_gen)
-            except StopIteration as e:
-                logger.info("Video source ended")
-                break
+        # writer = FrameWriter(self.config.record, None, None) if self.config.record else nullcontext
+        print(self.config.record)
+        writer = FrameWriter(str(self.config.record), None, None) if self.config.record else None
+        try:
+            while self.run_loop():
 
-            frame = Frame(i, img=img, H=self.config.camera.H, camera=self.config.camera)
+                try:
+                    i, img = next(video_gen)
+                except StopIteration as e:
+                    logger.info("Video source ended")
+                    break
 
-            # TODO: this is very dirty, need to find another way.
-            # perhaps multiprocessing Array?
-            self.frame_noimg_sock.send(pickle.dumps(frame.without_img()))
-            self.frame_sock.send(pickle.dumps(frame))
+                frame = Frame(i, img=img, H=self.config.camera.H, camera=self.config.camera)
+
+                # TODO: this is very dirty, need to find another way.
+                # perhaps multiprocessing Array?
+                self.frame_noimg_sock.send(pickle.dumps(frame.without_img()))
+                self.frame_sock.send(pickle.dumps(frame))
+
+                if writer:
+                    writer.write(frame.img)
+        finally:
+            if writer:
+                writer.release()
 
 
         logger.info("Stopping")
@@ -84,6 +92,10 @@ class FrameEmitter(node.Node):
                     help="End (or loop) playback at given frame.",
                     default=None,
                     type=int)
+        argparser.add_argument("--record",
+                    help="Record source video to given filename",
+                    default=None,
+                    type=Path)
 
         argparser.add_argument("--video-loop",
                             help="By default it emitter will run only once. This allows it to loop the video file to keep testing.",
diff --git a/trap/preview_renderer.py b/trap/preview_renderer.py
index 8ff71e4..bc7ce9c 100644
--- a/trap/preview_renderer.py
+++ b/trap/preview_renderer.py
@@ -300,7 +300,7 @@ class FrameWriter:
     """
     def __init__(self, filename: str, fps: float, frame_size: Optional[tuple] = None) -> None:
         self.filename = filename
-        self.fps = fps
+        self._fps = fps
         self.frame_size = frame_size
 
         self.tmp_dir = tempfile.TemporaryDirectory(prefix="trap-output-")
diff --git a/trap/video_sources.py b/trap/video_sources.py
index fe893c7..37ad555 100644
--- a/trap/video_sources.py
+++ b/trap/video_sources.py
@@ -107,8 +107,6 @@ class GigE(VideoSource):
         br = self.config.post_crop_br or (img.shape[1], img.shape[0])
         
         return img[tl[1]:br[1],tl[0]:br[0],:]
-                
-    
 
 class SingleCvVideoSource(VideoSource):
     def recv(self):