refactor(benchmark): use time instead of keyboard for capturing video in benchmark
This commit is contained in:
parent
9a22c9834d
commit
9364bbe9e5
|
@ -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))
|
||||||
|
|
|
@ -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",
|
||||||
|
|
Loading…
Reference in New Issue