From 9364bbe9e5efc581fc756b7dcf9739a00f7dde27 Mon Sep 17 00:00:00 2001
From: Steven Palma <steven.palma@huggingface.co>
Date: Tue, 8 Apr 2025 14:10:04 +0200
Subject: [PATCH] refactor(benchmark): use time instead of keyboard for
 capturing video in benchmark

---
 benchmarks/video/capture_camera_feed.py | 17 ++++++++++-------
 pyproject.toml                          |  1 -
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/benchmarks/video/capture_camera_feed.py b/benchmarks/video/capture_camera_feed.py
index 650c05a6..5abba244 100644
--- a/benchmarks/video/capture_camera_feed.py
+++ b/benchmarks/video/capture_camera_feed.py
@@ -18,14 +18,14 @@
 import argparse
 import datetime as dt
 import os
+import time
 from pathlib import Path
 
 import cv2
-import keyboard
 import rerun as rr
 
 
-def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height: int):
+def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height: int, duration: int):
     rr.init("lerobot_capture_camera_feed")
     memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "5%")
     rr.spawn(memory_limit=memory_limit)
@@ -46,7 +46,8 @@ def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height
     cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
 
     frame_index = 0
-    while True:
+    start_time = time.time()
+    while time.time() - start_time < duration:
         ret, frame = cap.read()
 
         if not ret:
@@ -56,10 +57,6 @@ def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height
         cv2.imwrite(str(capture_dir / f"frame_{frame_index:06d}.png"), frame)
         frame_index += 1
 
-        # Break the loop on 'q' key press
-        if keyboard.is_pressed("q"):
-            break
-
     # Release the capture and destroy all windows
     cap.release()
     # TODO(Steven): Find a way to close visualizer: https://github.com/rerun-io/rerun/pull/9400
@@ -93,5 +90,11 @@ if __name__ == "__main__":
         default=720,
         help="Height of the captured images.",
     )
+    parser.add_argument(
+        "--duration",
+        type=int,
+        default=20,
+        help="Duration in seconds for which the video stream should be captured.",
+    )
     args = parser.parse_args()
     display_and_save_video_stream(**vars(args))
diff --git a/pyproject.toml b/pyproject.toml
index 873ce7e7..b39c0720 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,6 @@ dependencies = [
     "numba>=0.59.0",
     "omegaconf>=2.3.0",
     "opencv-python-headless>=4.9.0",
-    "keyboard>=0.13.5",
     "packaging>=24.2",
     "pymunk>=6.6.0",
     "pynput>=1.7.7",