From ec31810d807b1c027ec3e5ad61fb25c1507208f2 Mon Sep 17 00:00:00 2001
From: Ruben van de Ven <git@rubenvandeven.com>
Date: Tue, 17 Oct 2023 12:04:53 +0200
Subject: [PATCH] Render directly to video

---
 trap/frame_emitter.py     |  4 +++-
 trap/plumber.py           | 24 ++++++++++++++++++++++--
 trap/prediction_server.py |  5 +++++
 trap/renderer.py          | 14 +++++++++++++-
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/trap/frame_emitter.py b/trap/frame_emitter.py
index 12579d8..dcd477c 100644
--- a/trap/frame_emitter.py
+++ b/trap/frame_emitter.py
@@ -37,6 +37,7 @@ class FrameEmitter:
         video = cv2.VideoCapture(str(self.config.video_src))
         fps = video.get(cv2.CAP_PROP_FPS)
         frame_duration = 1./fps
+        logger.info(f"Emit frames at {fps} fps")
 
         prev_time = time.time()
         while self.is_running.is_set():
@@ -65,4 +66,5 @@ class FrameEmitter:
 
 def run_frame_emitter(config: Namespace, is_running: Event):
     router = FrameEmitter(config, is_running)
-    router.emit_video()
\ No newline at end of file
+    router.emit_video()
+    is_running.clear()
\ No newline at end of file
diff --git a/trap/plumber.py b/trap/plumber.py
index 2f9a6b7..da37a8f 100644
--- a/trap/plumber.py
+++ b/trap/plumber.py
@@ -1,7 +1,11 @@
+import atexit
 import logging
 from logging.handlers import SocketHandler
 from multiprocessing import Event, Process, Queue
+import multiprocessing
+import signal
 import sys
+import time
 from trap.config import parser
 from trap.frame_emitter import run_frame_emitter
 from trap.prediction_server import run_prediction_server
@@ -12,17 +16,28 @@ from trap.tracker import run_tracker
 
 logger = logging.getLogger("trap.plumbing")
 
+
 class ExceptionHandlingProcess(Process):
-    
+            
     def run(self):
         assert 'is_running' in self._kwargs
+
+        # exit handler to make sure that on many kinds of crashes and kills the
+        # suprocess warns parent/siblings
+        # TODO: Does not work with OOM kill. Would need a watchdog process for that
+        def exit_handler(*args):
+            self._kwargs['is_running'].clear()
+
+        atexit.register(exit_handler)
+        signal.signal(signal.SIGTERM, exit_handler)
+        signal.signal(signal.SIGINT, exit_handler)
+
         try:
             super(Process, self).run()
         except Exception as e:
             logger.exception(e)
             self._kwargs['is_running'].clear()
 
-
 def start():
     args = parser.parse_args()
     loglevel = logging.NOTSET if args.verbose > 1 else logging.DEBUG if args.verbose > 0  else logging.INFO
@@ -40,6 +55,8 @@ def start():
     isRunning.set()
     
 
+
+
     if args.remote_log_addr:
         logging.captureWarnings(True)
         root_logger.setLevel(logging.NOTSET)  # to send all records to cutelog
@@ -64,6 +81,9 @@ def start():
     for proc in procs:
         proc.start()
 
+    # wait for processes to clean up
     for proc in procs:
         proc.join()
+    
+    logger.info('Stop')
 
diff --git a/trap/prediction_server.py b/trap/prediction_server.py
index ba68b68..b9467a5 100644
--- a/trap/prediction_server.py
+++ b/trap/prediction_server.py
@@ -207,8 +207,13 @@ class PredictionServer:
         trajectron.set_environment(online_env, init_timestep)
 
         timestep = init_timestep + 1
+        prev_run_time = 0
         while self.is_running.is_set():
             timestep += 1
+            this_run_time = time.time()
+            logger.debug('test')
+            time.sleep(max(0, prev_run_time - this_run_time + .5))
+            prev_run_time = this_run_time
             # for timestep in range(init_timestep + 1, eval_scene.timesteps):
 
             # input_dict = eval_scene.get_clipped_input_dict(timestep, hyperparams['state'])
diff --git a/trap/renderer.py b/trap/renderer.py
index c185fcc..d2b64cc 100644
--- a/trap/renderer.py
+++ b/trap/renderer.py
@@ -1,5 +1,6 @@
 
 from argparse import Namespace
+import datetime
 import logging
 from multiprocessing import Event
 import cv2
@@ -36,6 +37,15 @@ class Renderer:
         if not self.config.output_dir.exists():
             raise FileNotFoundError("Path does not exist")
     
+        date_str = datetime.datetime.now().isoformat(timespec="minutes")
+        filename = self.config.output_dir / f"render_predictions-{date_str}.mp4"
+        logger.info(f"Write to {filename}")
+
+        fourcc = cv2.VideoWriter_fourcc(*'vp09')
+        # TODO: get FPS from frame_emitter
+        self.out = cv2.VideoWriter(str(filename), fourcc, 23.97, (1280,720))
+
+    
     def run(self):
         predictions = {}
         i=0
@@ -81,8 +91,10 @@ class Renderer:
 
             img_path = (self.config.output_dir / f"{i:05d}.png").resolve()
 
-            cv2.imwrite(str(img_path), img)
+            # cv2.imwrite(str(img_path), img)
+            self.out.write(img)
         logger.info('Stopping')
+        self.out.release()