Add video recording code

Author: rdesarz

biped_walking_controller/simulation.py

@@ -1,3 +1,5 @@
+import os
+import subprocess
 import typing
 from pathlib import Path
 from typing import Tuple, Dict
@@ -340,6 +342,50 @@ def _compute_force(
     return total_force, [mean_x, mean_y, mean_z]
 
 
+def _process_video(filename: str, simulation_duration: float, start_time: float = 0.0):
+    cmd = [
+        "ffprobe",
+        "-i",
+        filename,
+        "-show_entries",
+        "format=duration",
+        "-v",
+        "quiet",
+        "-of",
+        "csv=p=0",
+    ]
+
+    out = subprocess.check_output(cmd).decode().strip()
+    video_duration = float(out)
+
+    ratio = simulation_duration / video_duration
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-ss",
+        str(start_time / ratio),
+        "-i",
+        filename,
+        "-filter:v",
+        f"setpts={ratio}*PTS",
+        "-r",
+        "60",
+        "-c:v",
+        "libx264",
+        "-crf",
+        "18",
+        "-preset",
+        "slow",
+        "output.mp4",
+    ]
+
+    subprocess.run(cmd, check=True)
+
+    # Remove the original file and rename the output to the original file
+    os.remove(filename)
+    os.rename("output.mp4", filename)
+
+
 class Simulator:
     """
     Thin PyBullet wrapper for loading a robot URDF and driving it from
@@ -422,6 +468,8 @@ def __init__(self, dt, path_to_robot_urdf: Path, model, launch_gui=True):
 
         self._displayed_lines = None
         self._displayed_points = None
+        self.log_id = None
+        self.filename = None
 
     def step(self):
         """
@@ -723,3 +771,12 @@ def get_zmp_pose(self):
         px = -M[1] / F[2]
         py = M[0] / F[2]
         return np.array([px, py, 0.0])
+
+    def start_video_record(self, filename: str = "recording.mp4"):
+        self.filename = filename
+        self.log_id = pb.startStateLogging(pb.STATE_LOGGING_VIDEO_MP4, self.filename)
+
+    def stop_video_record(self, duration: float, t_start: float = 0.0):
+        pb.stopStateLogging(self.log_id)
+
+        _process_video(self.filename, duration, t_start)

examples/example_4_physics_simulation.py

@@ -1,6 +1,5 @@
 import math
 import argparse
-import subprocess
 from pathlib import Path
 
 import numpy as np
@@ -34,51 +33,12 @@
 )
 
 
-def process_video(simulation_duration: float, start_time: float = 0.0):
-    cmd = [
-        "ffprobe",
-        "-i",
-        "talos_run.mp4",
-        "-show_entries",
-        "format=duration",
-        "-v",
-        "quiet",
-        "-of",
-        "csv=p=0",
-    ]
-
-    out = subprocess.check_output(cmd).decode().strip()
-    video_duration = float(out)
-
-    ratio = simulation_duration / video_duration
-
-    cmd = [
-        "ffmpeg",
-        "-ss",
-        str(start_time / ratio),
-        "-i",
-        "talos_run.mp4",
-        "-filter:v",
-        f"setpts={ratio}*PTS",
-        "-r",
-        "60",
-        "-c:v",
-        "libx264",
-        "-crf",
-        "18",
-        "-preset",
-        "slow",
-        "talos_run_fixed.mp4",
-    ]
-
-    subprocess.run(cmd, check=True)
-
-
 def main():
     p = argparse.ArgumentParser()
     p.add_argument("--path-talos-data", type=Path, help="Path to talos_data root")
     p.add_argument("--plot-results", action="store_true")
     p.add_argument("--launch-gui", action="store_true")
+    p.add_argument("--record-video", action="store_true")
     args = p.parse_args()
 
     np.set_printoptions(suppress=True, precision=3)
@@ -262,7 +222,8 @@ def main():
     rf_forces = np.zeros((len(phases), 1))
     lf_forces = np.zeros((len(phases), 1))
 
-    log_id = pb.startStateLogging(pb.STATE_LOGGING_VIDEO_MP4, "talos_run.mp4")
+    if args.record_video:
+        simulator.start_video_record()
 
     # We start the walking phase
     for k, _ in enumerate(phases[:-2]):
@@ -344,9 +305,8 @@ def main():
 
             rf_forces[k], lf_forces[k] = simulator.get_contact_forces()
 
-    pb.stopStateLogging(log_id)
-
-    process_video(t[-1], t_init + 2.0)
+    if args.record_video:
+        simulator.stop_video_record(duration=t[-1])
 
     if args.plot_results: