mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-30 07:59:43 +00:00
Martino Russi
134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Camera recorder for MuJoCo simulator
|
|
Subscribes to ZMQ camera feed and saves images to disk
|
|
"""
|
|
import argparse
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# Add sim module to path
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from sim.sensor_utils import SensorClient
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Save camera streams from MuJoCo simulator")
|
|
parser.add_argument("--host", type=str, default="localhost",
|
|
help="Simulator host address (default: localhost)")
|
|
parser.add_argument("--port", type=int, default=5555,
|
|
help="ZMQ port (default: 5555)")
|
|
parser.add_argument("--save-dir", type=str, default="./camera_recordings",
|
|
help="Directory to save images (default: ./camera_recordings)")
|
|
parser.add_argument("--save-rate", type=int, default=5,
|
|
help="Save every Nth frame (default: 5 = ~6Hz at 30fps stream)")
|
|
parser.add_argument("--max-frames", type=int, default=None,
|
|
help="Maximum frames to save before exiting (default: unlimited)")
|
|
args = parser.parse_args()
|
|
|
|
# Create save directory
|
|
save_dir = Path(args.save_dir)
|
|
save_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
print("="*60)
|
|
print("📷 MuJoCo Camera Recorder")
|
|
print("="*60)
|
|
print(f"🌐 Connecting to: tcp://{args.host}:{args.port}")
|
|
print(f"💾 Saving to: {save_dir.absolute()}")
|
|
print(f"⏬ Save rate: Every {args.save_rate} frames")
|
|
if args.max_frames:
|
|
print(f"🎬 Max frames: {args.max_frames}")
|
|
print("="*60)
|
|
print("\nWaiting for camera data...")
|
|
print("Press Ctrl+C to stop recording\n")
|
|
|
|
# Connect to sensor server
|
|
client = SensorClient()
|
|
client.start_client(server_ip=args.host, port=args.port)
|
|
|
|
frame_count = 0
|
|
saved_count = 0
|
|
last_time = time.time()
|
|
fps_display = 0
|
|
|
|
# Track per-camera frame counts
|
|
camera_frame_counts = {}
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
# Receive image data
|
|
data = client.receive_message()
|
|
|
|
# Calculate FPS
|
|
frame_count += 1
|
|
current_time = time.time()
|
|
if current_time - last_time >= 1.0:
|
|
fps_display = frame_count / (current_time - last_time)
|
|
frame_count = 0
|
|
last_time = current_time
|
|
|
|
# Process each camera
|
|
for key, value in data.items():
|
|
if key == "timestamps":
|
|
continue
|
|
|
|
# Initialize counter for this camera
|
|
if key not in camera_frame_counts:
|
|
camera_frame_counts[key] = 0
|
|
# Create subdirectory for each camera
|
|
(save_dir / key).mkdir(exist_ok=True)
|
|
|
|
camera_frame_counts[key] += 1
|
|
|
|
# Only save every Nth frame
|
|
if camera_frame_counts[key] % args.save_rate != 0:
|
|
continue
|
|
|
|
# Decode image if it's a string (base64 encoded)
|
|
if isinstance(value, str):
|
|
from sim.sensor_utils import ImageUtils
|
|
img = ImageUtils.decode_image(value)
|
|
elif isinstance(value, np.ndarray):
|
|
img = value
|
|
else:
|
|
continue
|
|
|
|
# Save image
|
|
filename = f"{key}_{saved_count:06d}.jpg"
|
|
save_path = save_dir / key / filename
|
|
cv2.imwrite(str(save_path), img)
|
|
|
|
# Print progress
|
|
if saved_count % 10 == 0:
|
|
print(f"[{fps_display:5.1f} FPS] Saved {saved_count:4d} frames to {key}/")
|
|
|
|
saved_count += 1
|
|
|
|
# Check if we've reached max frames
|
|
if args.max_frames and saved_count >= args.max_frames:
|
|
print(f"\n✓ Reached max frames ({args.max_frames}). Stopping...")
|
|
return
|
|
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error receiving/saving frame: {e}")
|
|
time.sleep(0.1)
|
|
|
|
except KeyboardInterrupt:
|
|
print(f"\n\n✓ Recording stopped")
|
|
print(f"📊 Total frames saved: {saved_count}")
|
|
print(f"📁 Location: {save_dir.absolute()}")
|
|
finally:
|
|
client.stop_client()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|