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 argparse
import datetime as dt import datetime as dt
import os import os
import time
from pathlib import Path from pathlib import Path
import cv2 import cv2
import keyboard
import rerun as rr 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") rr.init("lerobot_capture_camera_feed")
memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "5%") memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "5%")
rr.spawn(memory_limit=memory_limit) 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) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
frame_index = 0 frame_index = 0
while True: start_time = time.time()
while time.time() - start_time < duration:
ret, frame = cap.read() ret, frame = cap.read()
if not ret: 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) cv2.imwrite(str(capture_dir / f"frame_{frame_index:06d}.png"), frame)
frame_index += 1 frame_index += 1
# Break the loop on 'q' key press
if keyboard.is_pressed("q"):
break
# Release the capture and destroy all windows # Release the capture and destroy all windows
cap.release() cap.release()
# TODO(Steven): Find a way to close visualizer: https://github.com/rerun-io/rerun/pull/9400 # 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, default=720,
help="Height of the captured images.", 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() args = parser.parse_args()
display_and_save_video_stream(**vars(args)) display_and_save_video_stream(**vars(args))

View File

@ -61,7 +61,6 @@ dependencies = [
"numba>=0.59.0", "numba>=0.59.0",
"omegaconf>=2.3.0", "omegaconf>=2.3.0",
"opencv-python-headless>=4.9.0", "opencv-python-headless>=4.9.0",
"keyboard>=0.13.5",
"packaging>=24.2", "packaging>=24.2",
"pymunk>=6.6.0", "pymunk>=6.6.0",
"pynput>=1.7.7", "pynput>=1.7.7",