259 lines
7.6 KiB
Python
259 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pyrealsense2 as rs
|
|
|
|
|
|
SAVE_DIR = Path("d405_recordings")
|
|
WINDOW_NAME = "RealSense D405 Video Recorder"
|
|
PREFERRED_COLOR_FORMATS = (rs.format.bgr8, rs.format.rgb8)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ColorStreamConfig:
|
|
width: int
|
|
height: int
|
|
fps: int
|
|
stream_format: rs.format
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Record color video from an Intel RealSense D405 camera.")
|
|
parser.add_argument("--width", type=int, default=1280, help="Color stream width. Default: 1280")
|
|
parser.add_argument("--height", type=int, default=720, help="Color stream height. Default: 720")
|
|
parser.add_argument("--fps", type=int, default=30, help="Color stream FPS. Default: 30")
|
|
parser.add_argument(
|
|
"--list-profiles",
|
|
action="store_true",
|
|
help="List supported color video profiles for the connected camera and exit.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def get_color_stream_candidates() -> list[ColorStreamConfig]:
|
|
context = rs.context()
|
|
devices = context.query_devices()
|
|
if len(devices) == 0:
|
|
raise RuntimeError("No RealSense device detected.")
|
|
|
|
device = devices[0]
|
|
stream_candidates: list[ColorStreamConfig] = []
|
|
|
|
for sensor in device.sensors:
|
|
for profile in sensor.get_stream_profiles():
|
|
if profile.stream_type() != rs.stream.color:
|
|
continue
|
|
if profile.format() not in PREFERRED_COLOR_FORMATS:
|
|
continue
|
|
|
|
try:
|
|
video_profile = profile.as_video_stream_profile()
|
|
except RuntimeError:
|
|
continue
|
|
|
|
stream_candidates.append(
|
|
ColorStreamConfig(
|
|
width=video_profile.width(),
|
|
height=video_profile.height(),
|
|
fps=profile.fps(),
|
|
stream_format=profile.format(),
|
|
)
|
|
)
|
|
|
|
if not stream_candidates:
|
|
raise RuntimeError("No usable color stream profile found for the RealSense device.")
|
|
|
|
return stream_candidates
|
|
|
|
|
|
def list_profiles() -> None:
|
|
try:
|
|
stream_candidates = get_color_stream_candidates()
|
|
except RuntimeError as exc:
|
|
print(exc)
|
|
return
|
|
|
|
print("Supported color stream profiles:")
|
|
for candidate in sorted(
|
|
set(stream_candidates),
|
|
key=lambda item: (item.width, item.height, item.fps, str(item.stream_format)),
|
|
):
|
|
print(f" {candidate.width}x{candidate.height}@{candidate.fps} {candidate.stream_format}")
|
|
|
|
|
|
def select_color_stream(width: int, height: int, fps: int) -> ColorStreamConfig:
|
|
stream_candidates = get_color_stream_candidates()
|
|
matching_candidates = [
|
|
candidate
|
|
for candidate in stream_candidates
|
|
if candidate.width == width and candidate.height == height and candidate.fps == fps
|
|
]
|
|
|
|
if not matching_candidates:
|
|
raise RuntimeError(
|
|
f"No supported color stream profile matches {width}x{height}@{fps}. "
|
|
"Run with --list-profiles to see available profiles."
|
|
)
|
|
|
|
def sort_key(candidate: ColorStreamConfig) -> int:
|
|
return 1 if candidate.stream_format == rs.format.bgr8 else 0
|
|
|
|
return max(matching_candidates, key=sort_key)
|
|
|
|
|
|
def init_camera(stream_config: ColorStreamConfig) -> rs.pipeline:
|
|
pipeline = rs.pipeline()
|
|
config = rs.config()
|
|
config.enable_stream(
|
|
rs.stream.color,
|
|
stream_config.width,
|
|
stream_config.height,
|
|
stream_config.stream_format,
|
|
stream_config.fps,
|
|
)
|
|
pipeline.start(config)
|
|
return pipeline
|
|
|
|
|
|
def build_video_writer(output_path: Path, stream_config: ColorStreamConfig) -> cv2.VideoWriter:
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
|
writer = cv2.VideoWriter(
|
|
str(output_path),
|
|
fourcc,
|
|
stream_config.fps,
|
|
(stream_config.width, stream_config.height),
|
|
)
|
|
if not writer.isOpened():
|
|
raise RuntimeError(f"Failed to create video writer: {output_path}")
|
|
return writer
|
|
|
|
|
|
def draw_status(frame: np.ndarray, is_recording: bool, output_path: Path | None) -> np.ndarray:
|
|
preview = frame.copy()
|
|
frame_height = preview.shape[0]
|
|
if is_recording:
|
|
cv2.circle(preview, (25, 30), 8, (0, 0, 255), -1)
|
|
cv2.putText(
|
|
preview,
|
|
"REC",
|
|
(40, 36),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.8,
|
|
(0, 0, 255),
|
|
2,
|
|
)
|
|
else:
|
|
cv2.putText(
|
|
preview,
|
|
"Press 's' to start recording",
|
|
(20, 36),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.7,
|
|
(0, 255, 255),
|
|
2,
|
|
)
|
|
|
|
cv2.putText(
|
|
preview,
|
|
"Press 'q' to stop and quit",
|
|
(20, 70),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.7,
|
|
(255, 255, 255),
|
|
2,
|
|
)
|
|
|
|
if output_path is not None:
|
|
cv2.putText(
|
|
preview,
|
|
f"Saving: {output_path.name}",
|
|
(20, frame_height - 20),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
(255, 255, 255),
|
|
2,
|
|
)
|
|
|
|
return preview
|
|
|
|
|
|
def frame_to_bgr(color_frame: rs.video_frame, stream_config: ColorStreamConfig) -> np.ndarray:
|
|
color_image = np.asanyarray(color_frame.get_data())
|
|
if stream_config.stream_format == rs.format.rgb8:
|
|
return cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
|
|
return color_image
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.list_profiles:
|
|
list_profiles()
|
|
return
|
|
|
|
SAVE_DIR.mkdir(parents=True, exist_ok=True)
|
|
stream_config = select_color_stream(args.width, args.height, args.fps)
|
|
pipeline = init_camera(stream_config)
|
|
writer: cv2.VideoWriter | None = None
|
|
output_path: Path | None = None
|
|
is_recording = False
|
|
|
|
print("-" * 40)
|
|
print("RealSense D405 video recorder")
|
|
print(
|
|
"Selected color stream: "
|
|
f"{stream_config.width}x{stream_config.height} @ {stream_config.fps} FPS "
|
|
f"({stream_config.stream_format})"
|
|
)
|
|
print("Press 's' to start recording")
|
|
print("Press 'q' to stop recording and quit")
|
|
print(f"Video files will be saved to: {SAVE_DIR.resolve()}")
|
|
print("-" * 40)
|
|
|
|
try:
|
|
while True:
|
|
frames = pipeline.wait_for_frames(timeout_ms=2000)
|
|
color_frame = frames.get_color_frame()
|
|
if not color_frame:
|
|
continue
|
|
|
|
color_image = frame_to_bgr(color_frame, stream_config)
|
|
|
|
if is_recording and writer is not None:
|
|
writer.write(color_image)
|
|
|
|
preview = draw_status(color_image, is_recording, output_path)
|
|
cv2.imshow(WINDOW_NAME, preview)
|
|
|
|
key = cv2.waitKey(1) & 0xFF
|
|
|
|
if key == ord("s") and not is_recording:
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
output_path = SAVE_DIR / f"d405_record_{timestamp}.mp4"
|
|
writer = build_video_writer(output_path, stream_config)
|
|
is_recording = True
|
|
print(f"Recording started: {output_path}")
|
|
|
|
if key == ord("q"):
|
|
if is_recording:
|
|
print("Recording stopped.")
|
|
else:
|
|
print("Exit without recording.")
|
|
break
|
|
finally:
|
|
if writer is not None:
|
|
writer.release()
|
|
pipeline.stop()
|
|
cv2.destroyAllWindows()
|
|
if output_path is not None and output_path.exists():
|
|
print(f"Saved video: {output_path.resolve()}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|