refactor(benchmark): use time instead of keyboard for capturing video in benchmark

This commit is contained in:
Steven Palma 2025-04-08 14:10:04 +02:00
parent 9a22c9834d
commit 9364bbe9e5
No known key found for this signature in database
2 changed files with 10 additions and 8 deletions

View File

@ -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))

View File

@ -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",