This commit is contained in:
2026-05-06 10:47:14 +08:00
commit e3be8861b8
9 changed files with 4380 additions and 0 deletions

16
.gitignore vendored Normal file
View File

@ -0,0 +1,16 @@
# Python 缓存文件
__pycache__/
*.py[cod]
*.class
*.pyc
*.cache
# IDE 配置文件(可选,若不需要提交 IDE 配置)
.idea/
*.csv
ui_settings.json
# VSCode 配置
.vscode/

205
README.md Normal file
View File

@ -0,0 +1,205 @@
# TomatoPick 番茄采摘系统
TomatoPick 是一个面向番茄自动采摘实验的 Python 项目,集成了 Tkinter 参数界面、Intel RealSense RGB-D 相机、YOLO pose 目标检测、AUBO 机械臂控制和 AGV 移动平台控制。
当前仓库是运行端工程,核心代码集中在 `main.py``control.py``control_core.py`,辅助脚本位于 `tools/`
## 项目结构
```text
TomatoPick/
├── main.py # Tkinter 图形界面入口
├── control.py # UI 与核心逻辑之间的配置转发/兼容入口
├── control_core.py # 相机、YOLO、机械臂、AGV 与采摘状态机
├── ui_settings.json # UI 参数缓存,运行后自动保存/更新
├── tools/
│ ├── realsense_record_video.py # RealSense 录制辅助脚本
│ ├── video_to_rgb_frames.py # 视频抽帧辅助脚本
│ ├── aubo_joint_position.py # AUBO 关节位置辅助脚本
│ └── 1.png # UI 相机区域背景图
└── README.md
```
## 功能概览
- 通过图形界面配置机械臂 IP、AGV IP、运行时长、YOLO 模型路径、放置位置等参数。
- 支持完整采摘流程AGV 巡检、视觉检测、停车、精确定位、机械臂抓取、放置、回 Home、AGV 继续前进。
- 支持“视觉测试”模式,只启动 RealSense 和 YOLO 推理,不启动机械臂和 AGV适合先调试模型与相机画面。
- 将运行日志、终端输出和相机检测画面实时显示在 UI 中。
- 通过 `ui_settings.json` 保存上一次 UI 参数,方便下次启动恢复。
## 核心文件说明
### `main.py`
图形界面入口,直接运行即可打开参数配置界面。
主要职责:
- 构建 Tkinter 主界面,包括参数区、运行日志、相机画面和终端输出。
- 读取和保存 `ui_settings.json`
- 选择 YOLO `.pt` 模型文件和相机背景图。
- 启动完整采摘流程,内部调用 `control.main()`
- 启动视觉测试流程,只初始化 RealSense 与 YOLO。
- 接收 `control_core.py` 推送的相机帧并叠加检测信息。
### `control.py`
兼容入口和配置桥接层。UI 或外部脚本会先修改这里的模块级参数,然后由 `_sync_config()` 同步到 `control_core.configure()`
常用参数包括:
| 参数 | 说明 |
| --- | --- |
| `ROBOT_IP` / `ROBOT_PORT` | AUBO 机械臂连接地址 |
| `AGV_IP` / `AGV_PORT` | AGV 连接地址 |
| `AGV_SPEED_FORWARD` | AGV 前进速度 |
| `TOTAL_DURATION` | 系统总运行时长,单位秒 |
| `AGV_STOP_TIMEOUT` | AGV 停车后等待采摘的超时时间 |
| `YOLO_MODEL_PATH` | YOLO pose 模型路径,默认 `best.pt` |
| `YOLO_DETECT_CONF` | YOLO 推理置信度阈值 |
| `PICK_CONFIDENCE_THRESHOLD` | 进入采摘逻辑的目标置信度阈值 |
| `HOME_JOINTS` | 机械臂 Home 关节位 |
| `place_positions` | 番茄放置位 |
| `SCISSORS_ENABLED` | 是否启用末端剪刀/夹爪动作 |
### `control_core.py`
核心业务逻辑负责设备连接、视觉检测、坐标转换、机械臂动作、AGV 控制和资源释放。
主要模块:
- `PlacementManager`:管理放置位。
- `DetectedTomato` / `TomatoCandidate`:保存检测目标、关键点、深度和坐标信息。
- `RobotArmController`:封装机械臂回 Home、移动到预抓取点、抓取点和放置点等动作。
- `VisionController`:封装 RealSense 采集、YOLO pose 检测、关键点筛选、深度读取、坐标转换和抓取触发。
- `AgvController`:封装 AGV 获取控制权、前进、停止和清理。
核心状态由 3 个 `threading.Event` 协调:
| 状态 | 说明 |
| --- | --- |
| `running` | 系统总运行开关 |
| `has_tomato` | 视觉线程发现候选番茄后置位,通知 AGV 停车 |
| `picking_done` | 本轮抓取放置完成后置位,允许 AGV 继续前进 |
简化流程:
```text
连接机械臂
-> 连接 AGV 并获取控制权
-> 初始化 RealSense
-> 加载 YOLO 模型
-> 启动视觉线程和 AGV 线程
-> AGV 前进巡检
-> 视觉发现成熟番茄
-> AGV 停止
-> 视觉计算采摘点和抓取角度
-> 机械臂抓取并放置
-> 机械臂返回 Home
-> AGV 继续前进
```
## 环境依赖
建议在 Windows 环境下运行,并提前安装设备 SDK。
主要 Python 依赖:
```text
opencv-python
numpy
Pillow
ultralytics
pyrealsense2
pyaubo_sdk
pyaubo_agvc_sdk
```
`tkinter``threading``ctypes``signal` 等为 Python 标准库或随 Python 发行版提供。
可参考安装方式:
```bash
pip install opencv-python numpy Pillow ultralytics pyrealsense2
```
`pyaubo_sdk``pyaubo_agvc_sdk` 通常需要按设备厂商提供的 SDK 包安装。
## 启动方式
### 1. 启动图形界面
推荐使用 UI 启动:
```bash
python main.py
```
界面中可执行:
- `启动程序`运行完整采摘流程会连接机械臂、AGV、RealSense 并加载 YOLO。
- `视觉测试`:只测试相机采集和 YOLO 检测,不启动机械臂与 AGV。
- `停止程序`:停止当前任务并释放资源。
- `保存参数`:将当前 UI 参数写入 `ui_settings.json`
### 2. 直接启动核心流程
如果不需要 UI也可以直接运行
```bash
python control.py
```
这种方式会使用 `control.py` 顶部定义的默认参数,不读取 UI 输入。
## 模型文件
默认模型路径为:
```text
best.pt
```
可以将模型文件放在项目根目录,也可以在 UI 中选择任意 `.pt` 文件。使用绝对路径时,请确认路径在当前机器上存在。
当前 `ui_settings.json` 中保存的模型路径和背景图路径可能来自其他电脑或旧目录,首次运行前建议在 UI 中重新选择。
## 参数保存与传递
UI 参数传递流程:
```text
main.py
-> 保存到 ui_settings.json
-> 写入 control.py 模块变量
-> control.py 调用 _sync_config()
-> control_core.configure()
-> control_core.main()
```
因此:
- 从 UI 启动时,以界面中的参数为准。
- 直接运行 `control.py` 时,以 `control.py` 中写死的默认参数为准。
- 如果修改了 `control.py` 的默认值,需要重新启动程序才会生效。
## 运行前检查
完整流程会控制真实硬件,启动前请确认:
- AUBO 机械臂、AGV、RealSense 相机均已连接,并能被对应 SDK 访问。
- `ROBOT_IP``AGV_IP`、端口号与现场设备一致。
- `HOME_JOINTS`、手眼标定矩阵 `R_tc` / `T_tc`、放置位 `place_positions` 已按现场设备校准。
- YOLO 模型路径正确,且模型类别与关键点定义符合当前采摘逻辑。
- 调试视觉效果时优先使用 `视觉测试`,确认画面和检测结果稳定后再运行完整采摘流程。
## 辅助脚本
`tools/` 目录提供了一些调试脚本:
- `realsense_record_video.py`:录制 RealSense 视频。
- `video_to_rgb_frames.py`:从视频中导出 RGB 图片帧。
- `aubo_joint_position.py`:辅助读取/查看 AUBO 机械臂关节位置。
这些脚本主要用于采集数据、调试相机和校准机械臂位姿。

243
control.py Normal file
View File

@ -0,0 +1,243 @@
import signal
from typing import Any, Sequence
import control_core as _core
"""
对外兼容入口。
这个文件主要做两件事:
1. 暴露和旧版本一致的函数名,方便 UI 或外部脚本继续调用;
2. 在真正进入核心逻辑前,把这里可修改的配置同步到 `control_core.py`。
"""
# -------------------- 设备连接参数 --------------------
ROBOT_IP = "192.168.192.100"
ROBOT_PORT = 30004
M_PI = _core.M_PI
AGV_IP = "192.168.192.100"
AGV_PORT = 30104
AGV_SPEED_FORWARD = 0
AGV_SPEED_STOP = 0.0
AGV_RUN_DISTANCE = 5.0
# -------------------- 运行时长与检测参数 --------------------
TOTAL_DURATION = 300
AGV_STOP_TIMEOUT = 10
AGV_PICK_SETTLE_DELAY = 1.0
YOLO_MODEL_PATH = "best.pt"
YOLO_DETECT_CONF = 0.5
PICK_CONFIDENCE_THRESHOLD = 0.7
# -------------------- 位姿与轨迹参数 --------------------
# Home 是机械臂的初始/复位关节位。
HOME_JOINTS = [-1.5247, 1.0899, 2.4671, 1.2761, 1.5113, 0.0001]
# 视觉算出的目标点会再叠加这两个偏移:
# APPROACH_Y_OFFSET 用于先到预抓取点;
# LIFT_Z_OFFSET 用于修正真正执行抓取时的高度。
APPROACH_Y_OFFSET = 0.15
LIFT_Z_OFFSET = 0.035
PICKUP_X_OFFSET = 0.0
PICKUP_Y_OFFSET = -0.015
PICK_ANGLE_ORIENTATION_INDEX = 4
PICK_ANGLE_SIGN = 1.0
MAX_PICK_ANGLE_DEG = 30.0
# 当前项目只使用一个放置位,仍保留列表结构以兼容 control_core。
place_positions = [[-0.6283, 0.0354, -0.0889, 2.8798, 0.02, -1.4346]]
# -------------------- 夹爪与机械臂运动参数 --------------------
GRIPPER_CLOSE_IO = 1
GRIPPER_OPEN_IO = 0
GRIPPER_ACTION_DELAY = 2.0
SCISSORS_ENABLED = True
# 这里沿用底层 SDK 的弧度制。
ARM_SPEED = 100 * (M_PI / 180)
ARM_ACCEL = 100 * (M_PI / 180)
# -------------------- 参数同步清单 --------------------
_CONFIG_NAMES = (
"ROBOT_IP",
"ROBOT_PORT",
"AGV_IP",
"AGV_PORT",
"AGV_SPEED_FORWARD",
"AGV_SPEED_STOP",
"AGV_RUN_DISTANCE",
"TOTAL_DURATION",
"AGV_STOP_TIMEOUT",
"AGV_PICK_SETTLE_DELAY",
"YOLO_MODEL_PATH",
"YOLO_DETECT_CONF",
"PICK_CONFIDENCE_THRESHOLD",
"HOME_JOINTS",
"APPROACH_Y_OFFSET",
"LIFT_Z_OFFSET",
"PICKUP_X_OFFSET",
"PICKUP_Y_OFFSET",
"PICK_ANGLE_ORIENTATION_INDEX",
"PICK_ANGLE_SIGN",
"MAX_PICK_ANGLE_DEG",
"GRIPPER_CLOSE_IO",
"GRIPPER_OPEN_IO",
"GRIPPER_ACTION_DELAY",
"SCISSORS_ENABLED",
"ARM_SPEED",
"ARM_ACCEL",
"place_positions",
)
# -------------------- 运行状态别名 --------------------
running = _core.running
has_tomato = _core.has_tomato
picking_done = _core.picking_done
robot_rpc_client = _core.robot_rpc_client
ui_callback = _core.ui_callback
def _sync_config() -> None:
# 如果 UI 或外部脚本只改了 control.py这里会在进入核心逻辑前同步到 control_core。
config = {name: globals()[name] for name in _CONFIG_NAMES}
if place_positions and isinstance(place_positions[0], (list, tuple)):
config["place_positions"] = [list(position) for position in place_positions]
else:
config["place_positions"] = [list(place_positions)]
config["HOME_JOINTS"] = list(HOME_JOINTS)
_core.configure(config)
def set_ui_callback(callback):
global ui_callback
ui_callback = callback
_core.set_ui_callback(callback)
def exampleState(robot_name: str):
return _core.exampleState(robot_name)
def get_robot_end_effector_pose(robot_name: str):
return _core.get_robot_end_effector_pose(robot_name)
def exampleInverseK(robot_name: str, pose: Sequence[float], reference_q: Sequence[float]):
return _core.exampleInverseK(robot_name, pose, reference_q)
def exampleStartup():
_sync_config()
return _core.exampleStartup()
def quaternion_to_rotation_matrix(q: Sequence[float]):
return _core.quaternion_to_rotation_matrix(q)
def euler_to_rotation_matrix(roll: float, pitch: float, yaw: float):
return _core.euler_to_rotation_matrix(roll, pitch, yaw)
def get_robot_end_matrix(robot_name: str):
return _core.get_robot_end_matrix(robot_name)
def camera_to_base(point_cam: Sequence[float], robot_name: str):
return _core.camera_to_base(point_cam, robot_name)
def check_joint_position(robot_name: str, target_joints: Sequence[float]):
return _core.check_joint_position(robot_name, target_joints)
def check_tcp_pose(robot_name: str, target_pose: Sequence[float]):
return _core.check_tcp_pose(robot_name, target_pose)
def waitArrival(impl: Any, target_joints=None, target_pose=None):
return _core.waitArrival(impl, target_joints=target_joints, target_pose=target_pose)
def control_tool_io(robot_name: str, io_index: int, state: bool):
return _core.control_tool_io(robot_name, io_index, state)
def get_next_placement():
_sync_config()
return _core.get_next_placement()
def get_placement_index():
return _core.get_placement_index()
def create_placement_manager():
_sync_config()
return _core.create_placement_manager()
def return_to_home(robot_name: str):
_sync_config()
return _core.return_to_home(robot_name)
def control_robot(robot_name: str, pose: Sequence[float], angle_rad: float):
_sync_config()
return _core.control_robot(robot_name, pose, angle_rad)
def init_camera():
_sync_config()
return _core.init_camera()
def init_tomato_detector(model_path=None):
_sync_config()
return _core.init_tomato_detector(model_path)
def filter_ripe_tomatoes(results, confidence_threshold=None):
_sync_config()
return _core.filter_ripe_tomatoes(results, confidence_threshold)
def draw_tomato_tilt_line(color_image, x1, y1, x2, y2, pixel_x):
return _core.draw_tomato_tilt_line(color_image, x1, y1, x2, y2, pixel_x)
def fit_red_line(crop_image, x1, y1):
return _core.fit_red_line(crop_image, x1, y1)
def vision_detection_thread(pipeline, align, depth_intrinsics, model, robot_name: str):
_sync_config()
return _core.vision_detection_thread(pipeline, align, depth_intrinsics, model, robot_name)
def agv_control_thread(agv_client):
_sync_config()
return _core.agv_control_thread(agv_client)
def main():
# control.py 本身只做兼容和配置转发,真正业务流程在 control_core.main()。
_sync_config()
return _core.main()
def signal_handler(sig: int, frame: Any):
_sync_config()
return _core.signal_handler(sig, frame)
if __name__ == "__main__":
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
main()

1775
control_core.py Normal file

File diff suppressed because it is too large Load Diff

1672
main.py Normal file

File diff suppressed because it is too large Load Diff

BIN
tools/1.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.2 MiB

View File

@ -0,0 +1,106 @@
from pyaubo_sdk import RpcClient, RtdeClient
import time
import threading
ROBOT_IP = "192.168.192.100"
RPC_PORT = 30004
RTDE_PORT = 30010
USERNAME = "aubo"
PASSWORD = "123456"
PRINT_INTERVAL = 2.0
latest_joint_positions = None
latest_tcp_pose = None
lock = threading.Lock()
def main():
global latest_joint_positions, latest_tcp_pose
rpc = RpcClient()
rtde = RtdeClient()
topic_id = None
try:
# 1. RPC 连接
rpc.connect(ROBOT_IP, RPC_PORT)
rpc.login(USERNAME, PASSWORD)
robot_names = rpc.getRobotNames()
if not robot_names:
print("未找到机器人")
return
print(f"已连接机器人: {robot_names[0]}")
# 2. RTDE 连接 + 登录
rtde.connect(ROBOT_IP, RTDE_PORT)
rtde.login(USERNAME, PASSWORD)
# 3. 设置 RTDE 订阅项
names = ["R1_actual_q", "R1_actual_TCP_pose"]
topic_id = rtde.setTopic(False, names, 50, 0)
# 4. 回调函数
def callback(parser):
global latest_joint_positions, latest_tcp_pose
joint_positions = parser.popVectorDouble()
tcp_pose = parser.popVectorDouble()
with lock:
latest_joint_positions = joint_positions
latest_tcp_pose = tcp_pose
# 5. 订阅
rtde.subscribe(topic_id, callback)
print(f"开始监听(每 {PRINT_INTERVAL:.0f}s 输出一次)... Ctrl+C 退出")
while True:
time.sleep(PRINT_INTERVAL)
with lock:
joints = latest_joint_positions[:] if latest_joint_positions is not None else None
tcp = latest_tcp_pose[:] if latest_tcp_pose is not None else None
if joints is not None and tcp is not None:
joints_fmt = [round(j, 4) for j in joints]
tcp_fmt = [round(t, 4) for t in tcp]
print("\n" + "=" * 60)
print(f"关节角 (rad): {joints_fmt}")
print(f"末端位姿 : {tcp_fmt}")
print("=" * 60)
else:
print("尚未收到完整数据...")
except KeyboardInterrupt:
print("\n程序退出")
except Exception as e:
print(f"\n运行出错: {e}")
finally:
try:
if topic_id is not None:
rtde.removeTopic(False, topic_id)
except Exception:
pass
try:
rtde.disconnect()
except Exception:
pass
try:
rpc.logout()
except Exception:
pass
try:
rpc.disconnect()
except Exception:
pass
if __name__ == "__main__":
main()

View File

@ -0,0 +1,258 @@
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()

View File

@ -0,0 +1,105 @@
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()