forked from YikaiFu-cart/acAubo
106 lines
2.9 KiB
Python
106 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
|
|
|
|
DEFAULT_VIDEO_DIR = Path("d405_recordings")
|
|
DEFAULT_OUTPUT_DIR = Path("d405_rgb_frames")
|
|
DEFAULT_TARGET_FPS = 8.0
|
|
|
|
|
|
def find_latest_video(video_dir: Path) -> Path:
|
|
candidates = sorted(video_dir.glob("*.mp4"), key=lambda path: path.stat().st_mtime, reverse=True)
|
|
if not candidates:
|
|
raise FileNotFoundError(f"No .mp4 videos found in {video_dir.resolve()}")
|
|
return candidates[0]
|
|
|
|
|
|
def extract_frames(video_path: Path, output_dir: Path, target_fps: float) -> None:
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
if not cap.isOpened():
|
|
raise RuntimeError(f"Failed to open video: {video_path}")
|
|
|
|
source_fps = cap.get(cv2.CAP_PROP_FPS)
|
|
if source_fps <= 0:
|
|
source_fps = 30.0
|
|
|
|
frame_interval = max(source_fps / target_fps, 1.0)
|
|
next_frame_to_save = 0.0
|
|
frame_index = 0
|
|
saved_count = 0
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
print("-" * 40)
|
|
print(f"Input video : {video_path.resolve()}")
|
|
print(f"Output dir : {output_dir.resolve()}")
|
|
print(f"Source FPS : {source_fps:.2f}")
|
|
print(f"Target FPS : {target_fps:.2f}")
|
|
print("-" * 40)
|
|
|
|
try:
|
|
while True:
|
|
success, frame = cap.read()
|
|
if not success:
|
|
break
|
|
|
|
if frame_index + 1e-6 >= next_frame_to_save:
|
|
saved_count += 1
|
|
image_path = output_dir / f"{saved_count:06d}.jpg"
|
|
cv2.imwrite(str(image_path), frame)
|
|
next_frame_to_save += frame_interval
|
|
|
|
frame_index += 1
|
|
finally:
|
|
cap.release()
|
|
|
|
print(f"Saved {saved_count} RGB frames.")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Extract RGB frames from a recorded D405 video at 8 FPS.",
|
|
)
|
|
parser.add_argument(
|
|
"--video",
|
|
type=Path,
|
|
default=None,
|
|
help="Path to the input video. If omitted, the latest video in d405_recordings is used.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
type=Path,
|
|
default=None,
|
|
help="Directory to save extracted RGB images.",
|
|
)
|
|
parser.add_argument(
|
|
"--target-fps",
|
|
type=float,
|
|
default=DEFAULT_TARGET_FPS,
|
|
help="Frame extraction rate. Default is 8 FPS.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
video_path = args.video if args.video is not None else find_latest_video(DEFAULT_VIDEO_DIR)
|
|
if not video_path.exists():
|
|
raise FileNotFoundError(f"Video file does not exist: {video_path.resolve()}")
|
|
|
|
output_dir = args.output_dir
|
|
if output_dir is None:
|
|
output_dir = DEFAULT_OUTPUT_DIR / video_path.stem
|
|
|
|
if args.target_fps <= 0:
|
|
raise ValueError("--target-fps must be greater than 0.")
|
|
|
|
extract_frames(video_path, output_dir, args.target_fps)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|