diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 0000000..9a74a56 --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,7 @@ +{ + "permissions": { + "allow": [ + "Bash(git log:*)" + ] + } +} diff --git a/pyproject.toml b/pyproject.toml index acdba93..aaa1add 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "matplotlib>=3.10.7", "numpy>=2.2.6", "opencv-contrib-python-headless>=4.13.0.92", + "py-gpmf-parser>=0.1.1", "pydantic>=2.12.5", "rich>=14.1.0", "tqdm>=4.67.1", @@ -67,3 +68,6 @@ lint.ignore = [ "PLW2901", # For loop variable overwritten. "PLW0642", # Reassigned self in instance method. ] + +[tool.setuptools.package-data] +opai = ["configs/*"] diff --git a/src/opai/__init__.py b/src/opai/__init__.py index 28c0e65..ace59ee 100644 --- a/src/opai/__init__.py +++ b/src/opai/__init__.py @@ -1,6 +1,6 @@ from opai.presentation.facade import ( add_demos, - add_mapping, + add_mapping_video, browse_session, calibrate, calibrate_with_video, @@ -11,12 +11,14 @@ main, plot_video_frames, register_gopro, + run_extract_trajectories_batch, + run_mapping, verify_calibrated_parameters, ) __all__ = [ "add_demos", - "add_mapping", + "add_mapping_video", "browse_session", "calibrate", "calibrate_with_video", @@ -26,6 +28,8 @@ "list_sessions", "main", "plot_video_frames", + "run_extract_trajectories_batch", + "run_mapping", "verify_calibrated_parameters", "register_gopro", ] diff --git a/src/opai/application/__init__.py b/src/opai/application/__init__.py index a47c163..ce8c0aa 100644 --- a/src/opai/application/__init__.py +++ b/src/opai/application/__init__.py @@ -11,17 +11,22 @@ from opai.application.session import ( add_demos, add_mapping, + add_mapping_video, browse_session, list_sessions, ) +from opai.application.slam import run_extract_trajectories_batch, run_mapping __all__ = [ "add_demos", "add_mapping", + "add_mapping_video", "browse_session", "calibrate", "generate_charuco_board", "list_sessions", + "run_extract_trajectories_batch", + "run_mapping", "verify_calibrated_parameters", "get_media_list", "list_downloaded_thumbnails", diff --git a/src/opai/application/gopro.py b/src/opai/application/gopro.py index af8f8ae..7c9b65c 100644 --- a/src/opai/application/gopro.py +++ b/src/opai/application/gopro.py @@ -13,6 +13,7 @@ ) from opai.domain.context import Context from opai.domain.gopro import GPMedia, GPMediaList, GPThumbnail, GPThumbnailIndex +from opai.infrastructure.logger import get_logger from opai.infrastructure.persistence import ( load_gopro_thumbnail_index, write_gopro_thumbnail_index, @@ -23,6 +24,8 @@ DOWNLOAD_CHUNK_SIZE = 1024 * 1024 THUMBNAIL_INDEX_FILENAME = "gopro_thumbnail_index.json" +logger = get_logger(__name__) + def ensure_gopro_connected(ctx: Context) -> None: _run_async(_ensure_gopro_connected(ctx)) @@ -69,7 +72,13 @@ def register_gopro( socket_address_template = "172.2{}.1{}{}.51:8080" ctx.set_gopro_socket_address(socket_address_template.format(*serial_number[-3:])) + logger.info( + "Registered GoPro for session %s at %s", + ctx.name, + ctx.gopro_socket_address, + ) if download_thumbnails: + logger.info("Downloading GoPro thumbnails for session %s", ctx.name) _download_thumbnails(ctx) @@ -93,6 +102,12 @@ def download_file_from_gopro( url = f"http://{ctx.gopro_socket_address}/videos/DCIM/{source_directory}/{source_filename}" output_file = destination / (output_filename or source_filename) + logger.info( + "Downloading GoPro file %s/%s to %s", + source_directory, + source_filename, + output_file, + ) _run_async( _download_stream_to_file(ctx, url, output_file, progress_label=output_file.name) ) @@ -110,10 +125,13 @@ async def _get_media_list(ctx: Context) -> GPMediaList: url = f"http://{ctx.gopro_socket_address}/gopro/media/list" try: + logger.info("Fetching GoPro media list from %s", ctx.gopro_socket_address) async with _create_async_client() as client: response = await client.get(url) response.raise_for_status() - return GPMediaList(**response.json()) + media_list = GPMediaList(**response.json()) + logger.info("Fetched %s GoPro media directories", len(media_list.media)) + return media_list except httpx.HTTPError as exc: raise OPAIGoProRegistrationError( "Failed to fetch the GoPro media list. Verify the camera connection and try opai.register_gopro(...) again.", @@ -135,14 +153,20 @@ def list_downloaded_thumbnails(ctx: Context) -> list[GPThumbnail]: raise OPAIValidationError( "No GoPro thumbnails found. Call opai.register_gopro(..., download_thumbnails=True) first." ) + logger.info("Found %s downloaded GoPro thumbnails", len(items)) return items def _download_thumbnails(ctx: Context) -> None: if _thumbnail_index_path(ctx).exists(): + logger.info( + "Skipping GoPro thumbnail download because index already exists at %s", + _thumbnail_index_path(ctx), + ) return destination_root = ctx.session_directory / "gopro_thumbnails" + logger.info("Downloading GoPro thumbnails into %s", destination_root) thumbnails: list[GPThumbnail] = [] for media in get_media_list(ctx).media: thumbnails.extend( @@ -163,6 +187,9 @@ def _download_thumbnails_for_directory( destination = destination_root / media.d destination.mkdir(parents=True, exist_ok=True) thumbnails: list[GPThumbnail] = [] + logger.info( + "Downloading %s thumbnails from GoPro directory %s", len(media.fs), media.d + ) for file in media.fs: thumbnail_filename = Path(file.n).with_suffix(".jpg").name @@ -204,6 +231,7 @@ def _download_thumbnail_from_gopro( url = f"http://{ctx.gopro_socket_address}/gopro/media/thumbnail?path={media_path}" output_file = destination / output_filename + logger.info("Downloading GoPro thumbnail %s to %s", media_path, output_file) _run_async( _download_stream_to_file(ctx, url, output_file, progress_label=output_file.name) ) @@ -243,7 +271,14 @@ async def _download_stream_to_file( handle.write(chunk) progress_bar.update(len(chunk)) temp_file.replace(output_file) + logger.info("Downloaded GoPro asset %s to %s", progress_label, output_file) except httpx.HTTPError as exc: + logger.error( + "Failed to download GoPro asset %s from %s: %s", + progress_label, + url, + exc, + ) raise OPAIGoProRegistrationError( f"Failed to download '{progress_label}' from the GoPro.", details={ diff --git a/src/opai/application/imu.py b/src/opai/application/imu.py new file mode 100644 index 0000000..3e2239a --- /dev/null +++ b/src/opai/application/imu.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from py_gpmf_parser.gopro_telemetry_extractor import GoProTelemetryExtractor + +from opai.core.exceptions import OPAIWorkflowError +from opai.domain.imu import ( + DEFAULT_IMU_RECORDING_ID, + IMUPayload, + IMURecording, + IMUSample, + IMUStream, +) +from opai.infrastructure.logger import get_logger + +SECS_TO_MS = 1e3 +DEFAULT_STREAM_TYPES = [ + "ACCL", + "GYRO", + "GPS5", + "GPSP", + "GPSU", + "GPSF", + "GRAV", + "MAGN", + "CORI", + "IORI", + "TMPC", +] + +logger = get_logger(__name__) + + +def extract_imu_from_video( + video_path: Path, + dest: Path, + stream_types: list[str] = DEFAULT_STREAM_TYPES, +) -> Path: + """Extract IMU data from a video using py_gpmf_parser.""" + + logger.info("Extracting IMU streams from %s to %s", video_path, dest) + extractor = GoProTelemetryExtractor(str(video_path)) + try: + extractor.open_source() + + streams: dict[str, IMUStream] = {} + for stream in stream_types: + payload = extractor.extract_data(stream) + if payload and len(payload[0]) > 0: + logger.info( + "Extracted %s IMU samples for stream %s", + len(payload[0]), + stream, + ) + streams[stream] = IMUStream( + samples=tuple( + IMUSample( + value=data.tolist(), + cts=_coerce_timestamp_millis(ts), + ) + for data, ts in zip(*payload) + ) + ) + + output = IMUPayload( + recording_id=DEFAULT_IMU_RECORDING_ID, + recording=IMURecording(streams=streams), + frames_per_second=0.0, + ) + dest.write_text(json.dumps(output.to_dict(), indent=2), encoding="utf-8") + logger.info("Wrote IMU payload to %s", dest) + + return dest + + except Exception as exc: + logger.exception("Failed to extract IMU data from %s", video_path) + raise OPAIWorkflowError(f"Error processing {video_path}: {exc}") from exc + + finally: + extractor.close_source() + + +def _coerce_timestamp_millis(timestamp: object) -> int | float: + millis = timestamp * SECS_TO_MS + if hasattr(millis, "tolist"): + millis = millis.tolist() + if isinstance(millis, list): + if len(millis) != 1: + raise OPAIWorkflowError( + "IMU timestamp payload must contain exactly one scalar value.", + details={"timestamp": str(millis)}, + ) + return millis[0] + return millis diff --git a/src/opai/application/session.py b/src/opai/application/session.py index 66ef4b5..0cc0072 100644 --- a/src/opai/application/session.py +++ b/src/opai/application/session.py @@ -85,6 +85,10 @@ def add_mapping(ctx: Context, video_path: str | Path) -> MappingAsset: return mapping_asset +def add_mapping_video(ctx: Context, video_path: str | Path) -> MappingAsset: + return add_mapping(ctx, video_path) + + def list_sessions() -> list[str]: return [session.name for session in describe_sessions().sessions] diff --git a/src/opai/application/slam.py b/src/opai/application/slam.py new file mode 100644 index 0000000..3ee20c2 --- /dev/null +++ b/src/opai/application/slam.py @@ -0,0 +1,405 @@ +from __future__ import annotations + +import subprocess +from pathlib import Path + +import cv2 +import numpy as np + +from opai.application.imu import extract_imu_from_video +from opai.core.exceptions import OPAIWorkflowError +from opai.domain.context import Context +from opai.domain.slam import MappingRunResult +from opai.infrastructure.context_store import get_demo_videos, get_mapping_video +from opai.infrastructure.docker import ( + ALLOWED_RETURN_CODES, + pull_docker_image, + run_mapping_container, + run_trajectory_extraction_container, +) +from opai.infrastructure.logger import get_logger +from opai.infrastructure.video import ( + convert_video_to_fps, + get_video_duration_seconds, + get_video_fps, + get_video_resolution, +) + +DEFAULT_SLAM_SETTING_FILENAME = "gopro13_fisheye_ratio_4-3_2-7k.yaml" +DEFAULT_RESOLUTION = (2704, 2028) # (width, height) +DEFAULT_DOCKER_IMAGE = "authoree13/orb-slam3:latest" +TARGET_SLAM_FPS = 60 +HIGH_SPEED_SOURCE_FPS = 120.0 +FPS_TOLERANCE = 1.0 +PREPARED_MAPPING_VIDEO_FILENAME = "mapping_input_60fps.mp4" +PREPARED_TRAJECTORY_VIDEO_FILENAME = "trajectory_input_60fps.mp4" +DEFAULT_TRAJECTORY_TIMEOUT_MULTIPLE = 10.0 + +logger = get_logger(__name__) + + +def run_mapping(ctx: Context, slam_setting_path: str) -> MappingRunResult: + logger.info("Starting SLAM mapping for session %s", ctx.name) + original_video_path = get_mapping_video(ctx) + if original_video_path is None: + raise OPAIWorkflowError( + "No mapping video found. Call opai.add_mapping_video(...) or opai.add_mapping(...) before opai.run_mapping(...)." + ) + + work_directory = ctx.session_directory / "slam" / "mapping" + work_directory.mkdir(parents=True, exist_ok=True) + _clear_mapping_outputs(work_directory) + + prepared_video_path = original_video_path + source_fps = get_video_fps(original_video_path) + if abs(source_fps - HIGH_SPEED_SOURCE_FPS) < FPS_TOLERANCE: + logger.info( + "Converting mapping video from %.1f fps to %s fps", + source_fps, + TARGET_SLAM_FPS, + ) + prepared_video_path = convert_video_to_fps( + original_video_path, + work_directory / PREPARED_MAPPING_VIDEO_FILENAME, + TARGET_SLAM_FPS, + ) + else: + logger.info( + "Using mapping video without FPS conversion: %s", prepared_video_path + ) + + video_resolution = get_video_resolution(prepared_video_path) + logger.info( + "Prepared mapping video resolution is %sx%s", + video_resolution[0], + video_resolution[1], + ) + if video_resolution != DEFAULT_RESOLUTION: + raise OPAIWorkflowError( + "Mapping requires a " + f"{DEFAULT_RESOLUTION[0]}x{DEFAULT_RESOLUTION[1]} resolution video. " + f"Video resolution is {video_resolution}." + ) + + imu_json_path = extract_imu_from_video( + original_video_path, work_directory / "imu_data.json" + ) + logger.info("Using IMU payload at %s", imu_json_path) + slam_mask = np.zeros((DEFAULT_RESOLUTION[1], DEFAULT_RESOLUTION[0]), dtype=np.uint8) + slam_mask = draw_predefined_mask( + slam_mask, color=255, mirror=True, gripper=False, finger=True + ) + mask_path = work_directory / "slam_mask.png" + cv2.imwrite(str(mask_path), slam_mask) + logger.info("Wrote SLAM mask to %s", mask_path) + + map_path = work_directory / "map_atlas.osa" + trajectory_csv_path = work_directory / "mapping_camera_trajectory.csv" + stdout_log_path = work_directory / "slam_stdout.txt" + stderr_log_path = work_directory / "slam_stderr.txt" + + logger.info("Pulling SLAM Docker image %s", DEFAULT_DOCKER_IMAGE) + pull_docker_image(DEFAULT_DOCKER_IMAGE) + logger.info("Running SLAM container with settings file %s", slam_setting_path) + run_mapping_container( + docker_image=DEFAULT_DOCKER_IMAGE, + prepared_video_path=prepared_video_path, + work_directory=work_directory, + slam_setting_path=slam_setting_path, + stdout_path=stdout_log_path, + stderr_path=stderr_log_path, + enable_gui=True, + ) + + logger.info("Completed SLAM mapping for session %s", ctx.name) + return MappingRunResult( + input_video_path=original_video_path, + imu_json_path=imu_json_path, + mask_path=mask_path, + map_path=map_path, + trajectory_csv_path=trajectory_csv_path, + stdout_log_path=stdout_log_path, + stderr_log_path=stderr_log_path, + ) + + +def run_extract_trajectories_batch(ctx: Context) -> dict[str, object]: + atlas_path = ctx.session_directory / "slam" / "mapping" / "map_atlas.osa" + if not atlas_path.is_file(): + raise OPAIWorkflowError( + "Mapping trajectory extraction requires a map atlas file. Call opai.run_mapping(...) before calling opai.run_extract_trajectories_batch(...)." + ) + slam_setting_path = ctx.session_directory / "slam_settings.yaml" + if not slam_setting_path.is_file(): + raise OPAIWorkflowError( + "Trajectory extraction requires a session slam_settings.yaml file. Reinitialize the session with opai.init(...) before calling opai.run_extract_trajectories_batch(...)." + ) + + demo_video_paths = get_demo_videos(ctx) + if not demo_video_paths: + logger.info( + "No demo videos registered for session %s. Nothing to extract.", + ctx.name, + ) + return { + "processed_videos": [], + "total_processed": 0, + } + + logger.info( + "Starting trajectory extraction batch for session %s with %s demo videos", + ctx.name, + len(demo_video_paths), + ) + pull_docker_image(DEFAULT_DOCKER_IMAGE) + processed_videos: list[dict] = [] + + for demo_video_path in demo_video_paths: + demo_directory = demo_video_path.parent + trajectory_csv_path = demo_directory / "camera_trajectory.csv" + stdout_log_path = demo_directory / "slam_stdout.txt" + stderr_log_path = demo_directory / "slam_stderr.txt" + imu_json_path = demo_directory / "imu_data.json" + mask_path = demo_directory / "slam_mask.png" + prepared_video_path = demo_video_path + + if trajectory_csv_path.is_file(): + logger.warning( + "camera_trajectory.csv already exists, skipping %s", + demo_directory.name, + ) + continue + + logger.info( + "Preparing trajectory extraction inputs for %s", demo_directory.name + ) + + try: + source_fps = get_video_fps(demo_video_path) + if abs(source_fps - HIGH_SPEED_SOURCE_FPS) < FPS_TOLERANCE: + logger.info( + "Converting demo video %s from %.1f fps to %s fps", + demo_video_path, + source_fps, + TARGET_SLAM_FPS, + ) + prepared_video_path = convert_video_to_fps( + demo_video_path, + demo_directory / PREPARED_TRAJECTORY_VIDEO_FILENAME, + TARGET_SLAM_FPS, + ) + else: + logger.info( + "Using demo video without FPS conversion: %s", + prepared_video_path, + ) + + video_resolution = get_video_resolution(prepared_video_path) + if video_resolution != DEFAULT_RESOLUTION: + raise OPAIWorkflowError( + "Trajectory extraction requires a " + f"{DEFAULT_RESOLUTION[0]}x{DEFAULT_RESOLUTION[1]} resolution video. " + f"Video resolution is {video_resolution}.", + details={ + "video_path": str(prepared_video_path), + "demo_directory": str(demo_directory), + }, + ) + + extract_imu_from_video(demo_video_path, imu_json_path) + slam_mask = np.zeros( + (DEFAULT_RESOLUTION[1], DEFAULT_RESOLUTION[0]), + dtype=np.uint8, + ) + slam_mask = draw_predefined_mask( + slam_mask, color=255, mirror=True, gripper=False, finger=True + ) + wrote_mask = cv2.imwrite(str(mask_path), slam_mask) + if not wrote_mask: + raise OPAIWorkflowError( + "Failed to write the SLAM mask image.", + details={"path": str(mask_path)}, + ) + + duration_seconds = get_video_duration_seconds(prepared_video_path) + timeout_seconds = None + if duration_seconds > 0: + timeout_seconds = duration_seconds * DEFAULT_TRAJECTORY_TIMEOUT_MULTIPLE + + result = run_trajectory_extraction_container( + docker_image=DEFAULT_DOCKER_IMAGE, + prepared_video_path=prepared_video_path, + work_directory=demo_directory, + atlas_path=atlas_path, + slam_setting_path=slam_setting_path, + imu_json_path=imu_json_path, + trajectory_csv_path=trajectory_csv_path, + mask_path=mask_path, + stdout_path=stdout_log_path, + stderr_path=stderr_log_path, + timeout_seconds=timeout_seconds, + ) + status = ( + "success" if result.returncode in ALLOWED_RETURN_CODES else "failed" + ) + error_message = None + if status == "failed": + error_message = ( + f"SLAM extraction exited with code {result.returncode}. " + "Inspect the stderr log for details." + ) + except subprocess.TimeoutExpired: + logger.exception( + "Trajectory extraction timed out for demo %s", + demo_directory.name, + ) + status = "timeout" + error_message = "SLAM extraction timed out." + except OPAIWorkflowError as exc: + logger.exception( + "Trajectory extraction failed for demo %s", + demo_directory.name, + ) + status = "failed" + error_message = exc.message + + processed_videos.append( + { + "demo_id": demo_directory.name, + "video_path": str(demo_video_path), + "prepared_video_path": str(prepared_video_path), + "trajectory_csv": str(trajectory_csv_path), + "stdout_log": str(stdout_log_path), + "stderr_log": str(stderr_log_path), + "status": status, + "error_message": error_message, + } + ) + + logger.info( + "Completed trajectory extraction batch for session %s. Processed %s demo videos", + ctx.name, + len(processed_videos), + ) + return { + "processed_videos": processed_videos, + "total_processed": len(processed_videos), + } + + +def _clear_mapping_outputs(work_directory: Path) -> None: + logger.info("Clearing mapping outputs in %s", work_directory) + for filename in ( + "imu_data.json", + "slam_mask.png", + PREPARED_MAPPING_VIDEO_FILENAME, + "map_atlas.osa", + "mapping_camera_trajectory.csv", + "slam_stdout.txt", + "slam_stderr.txt", + ): + (work_directory / filename).unlink(missing_ok=True) + + +################### +### NOTE: The following functions are adapted directly from the official universal-manipulation-interface repository with minimal modifications. +### Will be updated to a more generalized version in the near-future. +################### + + +def draw_predefined_mask( + img, color=(0, 0, 0), mirror=True, gripper=True, finger=True, use_aa=False +): + all_coords = list() + if mirror: + all_coords.extend(get_mirror_canonical_polygon()) + if gripper: + all_coords.extend(get_gripper_canonical_polygon()) + if finger: + all_coords.extend(get_finger_canonical_polygon()) + + for coords in all_coords: + image_size = (img.shape[1], img.shape[0]) + pts = canonical_to_pixel_coords(coords, image_size) + pts = np.round(pts).astype(np.int32) + flag = cv2.LINE_AA if use_aa else cv2.LINE_8 + cv2.fillPoly(img, [pts], color=color, lineType=flag) + return img + + +def get_mirror_canonical_polygon(): + left_pts = [ + [540, 1700], + [680, 1450], + [590, 1070], + [290, 1130], + [290, 1770], + [550, 1770], + ] + left_coords = pixel_coords_to_canonical(left_pts, DEFAULT_RESOLUTION) + right_coords = left_coords.copy() + right_coords[:, 0] *= -1 + coords = np.stack([left_coords, right_coords]) + return coords + + +def get_gripper_canonical_polygon(): + left_pts = [ + [1352, 1730], + [1100, 1700], + [650, 1500], + [0, 1350], + [0, 2028], + [1352, 2028], + ] + left_coords = pixel_coords_to_canonical(left_pts, DEFAULT_RESOLUTION) + right_coords = left_coords.copy() + right_coords[:, 0] *= -1 + coords = np.stack([left_coords, right_coords]) + return coords + + +def get_finger_canonical_polygon( + height=0.37, top_width=0.25, bottom_width=1.4, resolution=DEFAULT_RESOLUTION +): + # image size + img_w, img_h = resolution + + # calculate coordinates + top_y = 1.0 - height + bottom_y = 1.0 + width = img_w / img_h + middle_x = width / 2.0 + top_left_x = middle_x - top_width / 2.0 + top_right_x = middle_x + top_width / 2.0 + bottom_left_x = middle_x - bottom_width / 2.0 + bottom_right_x = middle_x + bottom_width / 2.0 + + top_y *= img_h + bottom_y *= img_h + top_left_x *= img_h + top_right_x *= img_h + bottom_left_x *= img_h + bottom_right_x *= img_h + + # create polygon points for opencv API + points = [ + [ + [bottom_left_x, bottom_y], + [top_left_x, top_y], + [top_right_x, top_y], + [bottom_right_x, bottom_y], + ] + ] + coords = pixel_coords_to_canonical(points, img_shape=resolution) + return coords + + +def pixel_coords_to_canonical(pts, img_shape=DEFAULT_RESOLUTION): + coords = (np.asarray(pts) - np.array(img_shape) * 0.5) / img_shape[1] + return coords + + +def canonical_to_pixel_coords(coords, img_shape=DEFAULT_RESOLUTION): + pts = np.asarray(coords) * img_shape[1] + np.array(img_shape) * 0.5 + return pts diff --git a/src/opai/configs/__init__.py b/src/opai/configs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/opai/configs/slam/__init__.py b/src/opai/configs/slam/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/opai/configs/slam/gopro13_fisheye_ratio_4-3_2-7k.yaml b/src/opai/configs/slam/gopro13_fisheye_ratio_4-3_2-7k.yaml new file mode 100644 index 0000000..d4d2af8 --- /dev/null +++ b/src/opai/configs/slam/gopro13_fisheye_ratio_4-3_2-7k.yaml @@ -0,0 +1,89 @@ +%YAML:1.0 + +#-------------------------------------------------------------------------------------------- +# Camera Parameters. Adjust them! +#-------------------------------------------------------------------------------------------- +File.version: "1.0" +Camera.type: "KannalaBrandt8" +# Camera calibration and distortion parameters (OpenCV) +Camera1.fx: 282.906909765 # 190.97847715128717 +Camera1.fy: 282.906909765 # 190.9733070521226 +Camera1.cx: 480.0 # 254.93170605935475 +Camera1.cy: 360.0 # 256.8974428996504 + +Camera1.k1: 0.0813337 +Camera1.k2: -0.0852681 +Camera1.k3: 0.07741 +Camera1.k4: -0.0296763 + +# Camera resolution +Camera.width: 960 +Camera.height: 720 + +# Camera frames per second +Camera.fps: 60 + +# Color order of the images (0: BGR, 1: RGB. It is ignored if images are grayscale) +Camera.RGB: 1 + +# Transformation from body-frame (imu) to camera +IMU.T_b_c1: !!opencv-matrix + rows: 4 + cols: 4 + dt: f +# permuted calibration +# perm_mat = np.array([ +# [0,0,1,0], +# [1,0,0,0], +# [0,1,0,0], +# [0,0,0,1] +# ]) +# mat = perm_mat @ open_cam_imu_tbc + data: [ 0.00156717, -0.99997878, 0.00632289, -0.01321271, -0.99996531, + -0.00161881, -0.00817069, -0.00330095, 0.00818075, -0.00630987, + -0.99994663, -0.05175258, 0. , 0. , 0. , + 1. ] + +# IMU noise +IMU.NoiseGyro: 0.033863333333333336 +IMU.NoiseAcc: 0.15026666666666666 +IMU.GyroWalk: 0.25133333333333335 +IMU.AccWalk: 0.0016666666666666668 +IMU.Frequency: 197.577 + +#-------------------------------------------------------------------------------------------- +# ORB Parameters +#-------------------------------------------------------------------------------------------- + +# ORB Extractor: Number of features per image +ORBextractor.nFeatures: 1250 + +# ORB Extractor: Scale factor between levels in the scale pyramid +ORBextractor.scaleFactor: 1.2 + +# ORB Extractor: Number of levels in the scale pyramid +ORBextractor.nLevels: 8 + +# ORB Extractor: Fast threshold +# Image is divided in a grid. At each cell FAST are extracted imposing a minimum response. +# Firstly we impose iniThFAST. If no corners are detected we impose a lower value minThFAST +# You can lower these values if your images have low contrast +# ORBextractor.iniThFAST: 20 +# ORBextractor.minThFAST: 7 +ORBextractor.iniThFAST: 20 +ORBextractor.minThFAST: 7 + +#-------------------------------------------------------------------------------------------- +# Viewer Parameters +#-------------------------------------------------------------------------------------------- +Viewer.KeyFrameSize: 0.05 +Viewer.KeyFrameLineWidth: 1.0 +Viewer.GraphLineWidth: 0.9 +Viewer.PointSize: 2.0 +Viewer.CameraSize: 0.08 +Viewer.CameraLineWidth: 3.0 +Viewer.ViewpointX: 0.0 +Viewer.ViewpointY: -0.7 +Viewer.ViewpointZ: -3.5 # -1.8 +Viewer.ViewpointF: 500.0 +Viewer.imageViewScale: 1.0 diff --git a/src/opai/core/exceptions.py b/src/opai/core/exceptions.py index cbc1753..e88b768 100644 --- a/src/opai/core/exceptions.py +++ b/src/opai/core/exceptions.py @@ -61,3 +61,9 @@ class OPAIGoProNotConnectedError(OPAIError, RuntimeError): """Raised when a workflow fails after input validation.""" default_error_code = "gopro_not_connected_error" + + +class OPAIPackageResourceError(OPAIError, RuntimeError): + """Raised when package resources cannot be loaded.""" + + default_error_code = "package_resource_error" diff --git a/src/opai/domain/__init__.py b/src/opai/domain/__init__.py index 78d08ea..f4043cf 100644 --- a/src/opai/domain/__init__.py +++ b/src/opai/domain/__init__.py @@ -7,6 +7,7 @@ ) from opai.domain.context import Context from opai.domain.gopro import GPThumbnail, GPThumbnailIndex +from opai.domain.imu import IMUPayload, IMURecording, IMUSample, IMUStream from opai.domain.session import DemoAsset, MappingAsset, SessionManifest __all__ = [ @@ -18,6 +19,10 @@ "Context", "GPThumbnail", "GPThumbnailIndex", + "IMUPayload", + "IMURecording", + "IMUSample", + "IMUStream", "DemoAsset", "MappingAsset", "SessionManifest", diff --git a/src/opai/domain/gopro.py b/src/opai/domain/gopro.py index fb19c7b..9f63912 100644 --- a/src/opai/domain/gopro.py +++ b/src/opai/domain/gopro.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import datetime from pydantic import BaseModel diff --git a/src/opai/domain/imu.py b/src/opai/domain/imu.py new file mode 100644 index 0000000..5a3c86e --- /dev/null +++ b/src/opai/domain/imu.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from dataclasses import dataclass + +DEFAULT_IMU_RECORDING_ID = "1" +IMUNumeric = int | float +IMUValue = IMUNumeric | list[IMUNumeric] + + +@dataclass(frozen=True) +class IMUSample: + value: IMUValue + cts: IMUNumeric + + def to_dict(self) -> dict[str, IMUValue | IMUNumeric]: + return { + "value": self.value, + "cts": self.cts, + } + + +@dataclass(frozen=True) +class IMUStream: + samples: tuple[IMUSample, ...] + + def to_dict(self) -> dict[str, list[dict[str, IMUValue | IMUNumeric]]]: + return { + "samples": [sample.to_dict() for sample in self.samples], + } + + +@dataclass(frozen=True) +class IMURecording: + streams: dict[str, IMUStream] + + def to_dict(self) -> dict[str, dict[str, dict[str, list[dict[str, object]]]]]: + return { + "streams": { + stream_name: stream.to_dict() + for stream_name, stream in self.streams.items() + } + } + + +@dataclass(frozen=True) +class IMUPayload: + recording: IMURecording + frames_per_second: float = 0.0 + recording_id: str = DEFAULT_IMU_RECORDING_ID + + def to_dict(self) -> dict[str, object]: + return { + self.recording_id: self.recording.to_dict(), + "frames/second": self.frames_per_second, + } diff --git a/src/opai/domain/slam.py b/src/opai/domain/slam.py new file mode 100644 index 0000000..a37ea0e --- /dev/null +++ b/src/opai/domain/slam.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class MappingRunResult: + input_video_path: Path + imu_json_path: Path + mask_path: Path + map_path: Path + trajectory_csv_path: Path + stdout_log_path: Path + stderr_log_path: Path diff --git a/src/opai/infrastructure/__init__.py b/src/opai/infrastructure/__init__.py index a163694..3392417 100644 --- a/src/opai/infrastructure/__init__.py +++ b/src/opai/infrastructure/__init__.py @@ -1,9 +1,12 @@ from opai.infrastructure.context_store import ( get_active_context, + get_demo_videos, + get_mapping_video, get_session_directory, init_context, list_session_names, ) +from opai.infrastructure.logger import configure_logging, get_logger from opai.infrastructure.persistence import ( load_gopro_thumbnail_index, load_session_manifest, @@ -18,11 +21,15 @@ __all__ = [ "get_active_context", + "get_demo_videos", + "get_logger", + "get_mapping_video", "get_session_directory", "init_context", "load_gopro_thumbnail_index", "list_session_names", "load_session_manifest", + "configure_logging", "write_calibration_result", "write_calibration_verification_result", "write_charuco_board_config", diff --git a/src/opai/infrastructure/context_store.py b/src/opai/infrastructure/context_store.py index bf1ccb6..3604a25 100644 --- a/src/opai/infrastructure/context_store.py +++ b/src/opai/infrastructure/context_store.py @@ -1,7 +1,10 @@ from __future__ import annotations +import importlib.resources +import shutil from pathlib import Path +from opai.core.exceptions import OPAIPackageResourceError, OPAIWorkflowError from opai.domain.context import Context from opai.domain.session import SessionManifest from opai.infrastructure.persistence import ( @@ -12,6 +15,7 @@ _ACTIVE_CONTEXT: Context | None = None SESSION_ROOT_DIRNAME = "sessions" SESSION_MANIFEST_FILENAME = "session.json" +DEFAULT_SLAM_SETTING_FILENAME = "gopro13_fisheye_ratio_4-3_2-7k.yaml" def init_context(name: str) -> Context: @@ -27,9 +31,18 @@ def init_context(name: str) -> Context: session_directory=session_directory, manifest_path=manifest_path, ) + slam_setting_full_path = get_slam_config_dir() / DEFAULT_SLAM_SETTING_FILENAME + shutil.copy(slam_setting_full_path, session_directory / "slam_settings.yaml") return _ACTIVE_CONTEXT +def get_slam_config_dir() -> Path: + config_dir = importlib.resources.files("opai.configs") + if not isinstance(config_dir, Path): + raise OPAIPackageResourceError("Failed to load slam configs directory") + return config_dir / "slam" + + def get_active_context() -> Context | None: return _ACTIVE_CONTEXT @@ -55,6 +68,47 @@ def persist_manifest_for_context(ctx: Context, manifest: SessionManifest) -> Pat return write_session_manifest(manifest_path, manifest) +def get_mapping_video(ctx: Context) -> Path | None: + manifest = load_manifest_for_context(ctx) + if manifest.mapping is None: + return None + absolute_path = (ctx.session_directory / manifest.mapping.stored_path).resolve() + if absolute_path.exists() and absolute_path.is_file(): + return absolute_path + + raise OPAIWorkflowError( + f"Session mapping video is missing from stored session artifacts: {absolute_path}", + details={ + "session_name": ctx.name, + "asset_kind": "mapping", + "stored_path": manifest.mapping.stored_path, + "resolved_path": str(absolute_path), + }, + ) + + +def get_demo_videos(ctx: Context) -> list[Path]: + manifest = load_manifest_for_context(ctx) + demo_video_paths: list[Path] = [] + for demo in manifest.demos: + absolute_path = (ctx.session_directory / demo.stored_path).resolve() + if absolute_path.exists() and absolute_path.is_file(): + demo_video_paths.append(absolute_path) + continue + + raise OPAIWorkflowError( + f"Session demo video is missing from stored session artifacts: {absolute_path}", + details={ + "session_name": ctx.name, + "asset_kind": "demo", + "stored_path": demo.stored_path, + "resolved_path": str(absolute_path), + "asset_id": demo.demo_id, + }, + ) + return demo_video_paths + + def get_session_directory(name: str) -> Path: return session_root() / name diff --git a/src/opai/infrastructure/docker.py b/src/opai/infrastructure/docker.py new file mode 100644 index 0000000..23b7b2a --- /dev/null +++ b/src/opai/infrastructure/docker.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +from opai.core.exceptions import OPAIWorkflowError +from opai.infrastructure.logger import get_logger + +ALLOWED_RETURN_CODES = [ + 0, + 139, +] # NOTE: This is a noticible segfault, returncode of 139, in the chicheng/orb_slam3 image. We'll need to investigate this further. +DEFAULT_SLAM_CONTAINER_SETTINGS_PATH = Path("/slam_settings.yaml") +DEFAULT_SLAM_CONTAINER_VOCABULARY_PATH = Path("/ORB_SLAM3/Vocabulary/ORBvoc.txt") +DEFAULT_SLAM_MAPPING_COMMAND = "/ORB_SLAM3/Examples/Monocular-Inertial/gopro_slam" + +logger = get_logger(__name__) + + +def pull_docker_image(docker_image: str) -> None: + logger.info("Pulling Docker image: %s", docker_image) + result = subprocess.run( + ["docker", "pull", docker_image], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise OPAIWorkflowError( + f"Failed to pull Docker image: {docker_image}", + details={ + "docker_image": docker_image, + "stderr": result.stderr.strip(), + }, + ) + + +def run_mapping_container( + *, + docker_image: str, + prepared_video_path: Path, + work_directory: Path, + slam_setting_path: str, + stdout_path: Path, + stderr_path: Path, + enable_gui: bool = True, +) -> None: + workspace_target = Path("/workspace") + video_mount_target = workspace_target + command = ["docker", "run"] + command.extend( + [ + "--volume", + f"{work_directory.resolve()}:{workspace_target}", + ] + ) + command.extend( + [ + "--volume", + f"{slam_setting_path}:{DEFAULT_SLAM_CONTAINER_SETTINGS_PATH}:ro", + ] + ) + + if prepared_video_path.parent.resolve() != work_directory.resolve(): + video_mount_target = Path("/input") + command.extend( + [ + "--volume", + f"{prepared_video_path.parent.resolve()}:{video_mount_target}:ro", + ] + ) + + if enable_gui: + display_env = os.environ.get("DISPLAY") + if not display_env: + raise OPAIWorkflowError( + "DISPLAY environment variable is not set. GUI forwarding requires X11." + ) + command.extend(["--volume", "/tmp/.X11-unix:/tmp/.X11-unix"]) + command.extend(["--env", f"DISPLAY={display_env}"]) + xdg_runtime_dir = os.environ.get("XDG_RUNTIME_DIR") + if xdg_runtime_dir and os.path.exists(xdg_runtime_dir): + command.extend(["--volume", f"{xdg_runtime_dir}:{xdg_runtime_dir}"]) + command.extend(["--env", f"XDG_RUNTIME_DIR={xdg_runtime_dir}"]) + + xauthority_file = os.environ.get( + "XAUTHORITY", + os.path.expanduser("~/.Xauthority"), + ) + if os.path.exists(xauthority_file): + command.extend(["--volume", f"{xauthority_file}:{xauthority_file}"]) + command.extend(["--env", f"XAUTHORITY={xauthority_file}"]) + command.extend(["--env", "LIBGL_ALWAYS_SOFTWARE=1"]) + command.extend(["--ipc", "host"]) + + command.extend( + [ + docker_image, + DEFAULT_SLAM_MAPPING_COMMAND, + "--vocabulary", + str(DEFAULT_SLAM_CONTAINER_VOCABULARY_PATH), + "--setting", + str(DEFAULT_SLAM_CONTAINER_SETTINGS_PATH), + "--input_video", + str(video_mount_target / prepared_video_path.name), + "--input_imu_json", + str(workspace_target / "imu_data.json"), + "--output_trajectory_csv", + str(workspace_target / "mapping_camera_trajectory.csv"), + "--save_map", + str(workspace_target / "map_atlas.osa"), + "--mask_img", + str(workspace_target / "slam_mask.png"), + ] + ) + + if enable_gui: + command.extend(["--enable_gui"]) + + logger.info("Running Docker command: %s", " ".join(command)) + + with ( + stdout_path.open("w", encoding="utf-8") as stdout_file, + stderr_path.open("w", encoding="utf-8") as stderr_file, + ): + process = subprocess.Popen( + command, + cwd=str(work_directory.resolve()), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for line in iter(process.stdout.readline, ""): + logger.info("SUBPROCESS STDOUT: %s", line.strip()) + stdout_file.write(line) + stdout_file.flush() + for line in iter(process.stderr.readline, ""): + logger.warning("SUBPROCESS STDERR: %s", line.strip()) + stderr_file.write(line) + stderr_file.flush() + + process.wait() + + if process.returncode not in ALLOWED_RETURN_CODES: + raise OPAIWorkflowError( + "SLAM mapping failed. Inspect the generated stdout/stderr logs for details.", + details={ + "docker_image": docker_image, + "return_code": process.returncode, + "stdout_log": str(stdout_path), + "stderr_log": str(stderr_path), + "work_directory": str(work_directory), + }, + ) + + +def run_trajectory_extraction_container( + *, + docker_image: str, + prepared_video_path: Path, + work_directory: Path, + atlas_path: Path, + slam_setting_path: Path, + imu_json_path: Path, + trajectory_csv_path: Path, + mask_path: Path, + stdout_path: Path, + stderr_path: Path, + timeout_seconds: float | None = None, +) -> subprocess.CompletedProcess: + workspace_target = Path("/workspace") + atlas_mount_target = Path("/map") + video_mount_target = workspace_target + command = ["docker", "run", "--rm"] + command.extend( + [ + "--volume", + f"{work_directory.resolve()}:{workspace_target}", + ] + ) + command.extend( + [ + "--volume", + f"{atlas_path.parent.resolve()}:{atlas_mount_target}:ro", + ] + ) + command.extend( + [ + "--volume", + f"{slam_setting_path.resolve()}:{DEFAULT_SLAM_CONTAINER_SETTINGS_PATH}:ro", + ] + ) + + if prepared_video_path.parent.resolve() != work_directory.resolve(): + video_mount_target = Path("/input") + command.extend( + [ + "--volume", + f"{prepared_video_path.parent.resolve()}:{video_mount_target}:ro", + ] + ) + + command.extend( + [ + docker_image, + DEFAULT_SLAM_MAPPING_COMMAND, + "--vocabulary", + str(DEFAULT_SLAM_CONTAINER_VOCABULARY_PATH), + "--setting", + str(DEFAULT_SLAM_CONTAINER_SETTINGS_PATH), + "--input_video", + str(video_mount_target / prepared_video_path.name), + "--input_imu_json", + str(workspace_target / imu_json_path.name), + "--output_trajectory_csv", + str(workspace_target / trajectory_csv_path.name), + "--load_map", + str(atlas_mount_target / atlas_path.name), + "--mask_img", + str(workspace_target / mask_path.name), + ] + ) + + logger.info("Running Docker command: %s", " ".join(command)) + + with ( + stdout_path.open("w", encoding="utf-8") as stdout_file, + stderr_path.open("w", encoding="utf-8") as stderr_file, + ): + return subprocess.run( + command, + cwd=str(work_directory.resolve()), + stdout=stdout_file, + stderr=stderr_file, + timeout=timeout_seconds, + check=False, + ) diff --git a/src/opai/infrastructure/logger.py b/src/opai/infrastructure/logger.py new file mode 100644 index 0000000..73426a2 --- /dev/null +++ b/src/opai/infrastructure/logger.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import logging + +PACKAGE_LOGGER_NAME = "opai" +DEFAULT_LOG_FORMAT = "%(levelname)s %(name)s: %(message)s" + +_package_logger = logging.getLogger(PACKAGE_LOGGER_NAME) +if not any( + isinstance(handler, logging.NullHandler) for handler in _package_logger.handlers +): + _package_logger.addHandler(logging.NullHandler()) + + +def get_logger(name: str | None = None) -> logging.Logger: + if name is None or not name.strip(): + return logging.getLogger(PACKAGE_LOGGER_NAME) + if name == PACKAGE_LOGGER_NAME or name.startswith(f"{PACKAGE_LOGGER_NAME}."): + return logging.getLogger(name) + return logging.getLogger(f"{PACKAGE_LOGGER_NAME}.{name.lstrip('.')}") + + +def configure_logging(level: int | str = logging.INFO) -> logging.Logger: + logger = logging.getLogger(PACKAGE_LOGGER_NAME) + default_handler = next( + ( + handler + for handler in logger.handlers + if getattr(handler, "_opai_default_handler", False) + ), + None, + ) + if default_handler is None: + default_handler = logging.StreamHandler() + default_handler.setFormatter(logging.Formatter(DEFAULT_LOG_FORMAT)) + setattr(default_handler, "_opai_default_handler", True) + logger.addHandler(default_handler) + + logger.setLevel(level) + logger.propagate = False + return logger diff --git a/src/opai/infrastructure/persistence.py b/src/opai/infrastructure/persistence.py index d52e6ba..6e3e506 100644 --- a/src/opai/infrastructure/persistence.py +++ b/src/opai/infrastructure/persistence.py @@ -16,6 +16,9 @@ ) from opai.domain.gopro import GPThumbnail, GPThumbnailIndex from opai.domain.session import DemoAsset, MappingAsset, SessionManifest +from opai.infrastructure.logger import get_logger + +logger = get_logger(__name__) def write_calibration_result( @@ -43,7 +46,7 @@ def write_calibration_result( output_path = session_directory / filename output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") - print(f"Wrote calibration verification result to {output_path}") + logger.info("Wrote calibration result to %s", output_path) return output_path @@ -71,7 +74,7 @@ def write_calibration_verification_result( output_path = session_directory / filename output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") - print(f"Wrote calibration verification result to {output_path}") + logger.info("Wrote calibration verification result to %s", output_path) return output_path diff --git a/src/opai/infrastructure/video.py b/src/opai/infrastructure/video.py index 883fcbd..6b1bdc7 100644 --- a/src/opai/infrastructure/video.py +++ b/src/opai/infrastructure/video.py @@ -1,5 +1,6 @@ from __future__ import annotations +import subprocess from pathlib import Path import cv2 @@ -36,3 +37,120 @@ def sample_video_frames( capture.release() return tuple(sampled_frames) + + +def get_video_fps(video_path: str | Path) -> float: + path = Path(video_path).expanduser() + capture = cv2.VideoCapture(str(path)) + if not capture.isOpened(): + capture.release() + raise OPAIWorkflowError( + f"Unable to open video for FPS detection: {path}", + details={"path": str(path)}, + ) + + try: + fps = float(capture.get(cv2.CAP_PROP_FPS)) + finally: + capture.release() + + if fps <= 0: + raise OPAIWorkflowError( + f"Unable to detect a positive FPS value for video: {path}", + details={"path": str(path), "fps": fps}, + ) + return fps + + +def get_video_resolution(video_path: str | Path) -> tuple[int, int]: + path = Path(video_path).expanduser() + capture = cv2.VideoCapture(str(path)) + if not capture.isOpened(): + capture.release() + raise OPAIWorkflowError( + f"Unable to open video for resolution detection: {path}", + details={"path": str(path)}, + ) + + try: + return ( + int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), + int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)), + ) + finally: + capture.release() + + +def get_video_duration_seconds(video_path: str | Path) -> float: + path = Path(video_path).expanduser() + capture = cv2.VideoCapture(str(path)) + if not capture.isOpened(): + capture.release() + raise OPAIWorkflowError( + f"Unable to open video for duration detection: {path}", + details={"path": str(path)}, + ) + + try: + fps = float(capture.get(cv2.CAP_PROP_FPS)) + frame_count = float(capture.get(cv2.CAP_PROP_FRAME_COUNT)) + finally: + capture.release() + + if fps <= 0: + raise OPAIWorkflowError( + f"Unable to detect a positive FPS value for video: {path}", + details={"path": str(path), "fps": fps}, + ) + if frame_count <= 0: + return 0.0 + return frame_count / fps + + +def convert_video_to_fps( + video_path: str | Path, + output_path: str | Path, + target_fps: int, +) -> Path: + input_path = Path(video_path).expanduser() + destination_path = Path(output_path).expanduser() + destination_path.parent.mkdir(parents=True, exist_ok=True) + + command = [ + "ffmpeg", + "-i", + str(input_path), + "-map_metadata", + "0", + "-movflags", + "+faststart+use_metadata_tags", + "-vf", + f"fps={target_fps}", + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "23", + "-c:a", + "copy", + "-y", + str(destination_path), + ] + result = subprocess.run( + command, + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise OPAIWorkflowError( + f"Failed to convert video to {target_fps} fps: {input_path}", + details={ + "input_path": str(input_path), + "output_path": str(destination_path), + "target_fps": target_fps, + "stderr": result.stderr.strip(), + }, + ) + return destination_path diff --git a/src/opai/presentation/__init__.py b/src/opai/presentation/__init__.py index 139a4ee..4282067 100644 --- a/src/opai/presentation/__init__.py +++ b/src/opai/presentation/__init__.py @@ -1,6 +1,5 @@ from opai.presentation.facade import ( add_demos, - add_mapping, browse_session, calibrate, calibrate_with_video, @@ -10,12 +9,12 @@ list_sessions, main, plot_video_frames, + run_extract_trajectories_batch, verify_calibrated_parameters, ) __all__ = [ "add_demos", - "add_mapping", "browse_session", "calibrate", "calibrate_with_video", @@ -25,5 +24,6 @@ "list_sessions", "main", "plot_video_frames", + "run_extract_trajectories_batch", "verify_calibrated_parameters", ] diff --git a/src/opai/presentation/facade.py b/src/opai/presentation/facade.py index 76ebd79..848c68c 100644 --- a/src/opai/presentation/facade.py +++ b/src/opai/presentation/facade.py @@ -1,6 +1,5 @@ from __future__ import annotations -import html import re from collections.abc import Sequence from pathlib import Path @@ -27,9 +26,9 @@ CharucoBoardConfig, ) from opai.domain.context import Context -from opai.domain.gopro import GPThumbnail from opai.domain.plot import plot_frames from opai.domain.session import DemoAsset, MappingAsset +from opai.domain.slam import MappingRunResult from opai.infrastructure.context_store import get_active_context, init_context SESSION_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") @@ -237,13 +236,34 @@ def add_demos(video_paths: Sequence[str | Path]) -> tuple[DemoAsset, ...]: return add_demos_with_context(ctx, video_paths) -def add_mapping(video_path: str | Path) -> MappingAsset: +def add_mapping_video(video_path: str | Path) -> MappingAsset: ctx = get_context() from opai.application.session import add_mapping as add_mapping_with_context return add_mapping_with_context(ctx, video_path) +def run_mapping( + slam_setting_path: str, +) -> MappingRunResult: + ctx = get_context() + from opai.application.slam import run_mapping as run_mapping_with_context + + return run_mapping_with_context( + ctx=ctx, + slam_setting_path=slam_setting_path, + ) + + +def run_extract_trajectories_batch() -> dict[str, object]: + ctx = get_context() + from opai.application.slam import ( + run_extract_trajectories_batch as run_extract_trajectories_batch_with_context, + ) + + return run_extract_trajectories_batch_with_context(ctx=ctx) + + def list_sessions() -> list[str]: try: from rich.console import Console @@ -347,7 +367,8 @@ def browse_session(name: str) -> list[str]: def main() -> None: print( - "Use opai.init(name), opai.add_demos(...), opai.add_mapping(...), " + "Use opai.init(name), opai.add_demos(...), opai.add_mapping_video(...), " + "opai.add_mapping(...), opai.run_mapping(...), opai.run_extract_trajectories_batch(...), " "opai.generate_charuco_board(...), opai.calibrate(...), " "opai.calibrate_with_video(...), opai.verify_calibrated_parameters(...), " "and opai.plot_video_frames(...) from Python." diff --git a/tests/conftest.py b/tests/conftest.py index e69de29..868ef54 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import pytest + +from opai.infrastructure import context_store + + +@pytest.fixture(autouse=True) +def reset_active_context() -> None: + context_store._ACTIVE_CONTEXT = None + yield + context_store._ACTIVE_CONTEXT = None diff --git a/tests/test_context_store.py b/tests/test_context_store.py new file mode 100644 index 0000000..aa6476f --- /dev/null +++ b/tests/test_context_store.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import pytest + +from opai.core.exceptions import OPAIWorkflowError +from opai.domain.context import Context +from opai.domain.session import DemoAsset, MappingAsset, SessionManifest +from opai.infrastructure.context_store import get_demo_videos, get_mapping_video +from opai.infrastructure.persistence import write_session_manifest + + +def test_get_mapping_video_returns_absolute_stored_path(tmp_path) -> None: + ctx = _build_context(tmp_path) + mapping_path = ( + ctx.session_directory / "captures" / "mapping" / "current" / "map.mp4" + ) + mapping_path.parent.mkdir(parents=True) + mapping_path.write_bytes(b"mapping") + write_session_manifest( + ctx.manifest_path, + SessionManifest( + session_name=ctx.name, + demos=(), + mapping=MappingAsset( + source_path="/tmp/source-map.mp4", + stored_path="captures/mapping/current/map.mp4", + original_filename="map.mp4", + ), + ), + ) + + assert get_mapping_video(ctx) == mapping_path.resolve() + + +def test_get_mapping_video_returns_none_when_manifest_has_no_mapping(tmp_path) -> None: + ctx = _build_context(tmp_path) + write_session_manifest( + ctx.manifest_path, + SessionManifest(session_name=ctx.name, demos=(), mapping=None), + ) + + assert get_mapping_video(ctx) is None + + +def test_get_mapping_video_raises_when_stored_file_is_missing(tmp_path) -> None: + ctx = _build_context(tmp_path) + write_session_manifest( + ctx.manifest_path, + SessionManifest( + session_name=ctx.name, + demos=(), + mapping=MappingAsset( + source_path="/tmp/source-map.mp4", + stored_path="captures/mapping/current/missing.mp4", + original_filename="missing.mp4", + ), + ), + ) + + with pytest.raises(OPAIWorkflowError, match="mapping video is missing") as exc_info: + get_mapping_video(ctx) + + assert exc_info.value.details == { + "session_name": "session-001", + "asset_kind": "mapping", + "stored_path": "captures/mapping/current/missing.mp4", + "resolved_path": str( + ( + ctx.session_directory + / "captures" + / "mapping" + / "current" + / "missing.mp4" + ).resolve() + ), + } + + +def test_get_demo_videos_returns_absolute_paths_in_manifest_order(tmp_path) -> None: + ctx = _build_context(tmp_path) + demo_a = ctx.session_directory / "captures" / "demos" / "demo-0002" / "b.mp4" + demo_b = ctx.session_directory / "captures" / "demos" / "demo-0001" / "a.mp4" + demo_a.parent.mkdir(parents=True) + demo_b.parent.mkdir(parents=True) + demo_a.write_bytes(b"b") + demo_b.write_bytes(b"a") + write_session_manifest( + ctx.manifest_path, + SessionManifest( + session_name=ctx.name, + demos=( + DemoAsset( + demo_id="demo-0002", + source_path="/tmp/b.mp4", + stored_path="captures/demos/demo-0002/b.mp4", + original_filename="b.mp4", + ), + DemoAsset( + demo_id="demo-0001", + source_path="/tmp/a.mp4", + stored_path="captures/demos/demo-0001/a.mp4", + original_filename="a.mp4", + ), + ), + mapping=None, + ), + ) + + assert get_demo_videos(ctx) == [demo_a.resolve(), demo_b.resolve()] + + +def test_get_demo_videos_returns_empty_list_when_manifest_has_no_demos( + tmp_path, +) -> None: + ctx = _build_context(tmp_path) + write_session_manifest( + ctx.manifest_path, + SessionManifest(session_name=ctx.name, demos=(), mapping=None), + ) + + assert get_demo_videos(ctx) == [] + + +def test_get_demo_videos_raises_when_any_stored_file_is_missing(tmp_path) -> None: + ctx = _build_context(tmp_path) + demo_path = ctx.session_directory / "captures" / "demos" / "demo-0001" / "a.mp4" + demo_path.parent.mkdir(parents=True) + demo_path.write_bytes(b"a") + write_session_manifest( + ctx.manifest_path, + SessionManifest( + session_name=ctx.name, + demos=( + DemoAsset( + demo_id="demo-0001", + source_path="/tmp/a.mp4", + stored_path="captures/demos/demo-0001/a.mp4", + original_filename="a.mp4", + ), + DemoAsset( + demo_id="demo-0002", + source_path="/tmp/missing.mp4", + stored_path="captures/demos/demo-0002/missing.mp4", + original_filename="missing.mp4", + ), + ), + mapping=None, + ), + ) + + with pytest.raises(OPAIWorkflowError, match="demo video is missing") as exc_info: + get_demo_videos(ctx) + + assert exc_info.value.details == { + "session_name": "session-001", + "asset_kind": "demo", + "stored_path": "captures/demos/demo-0002/missing.mp4", + "resolved_path": str( + ( + ctx.session_directory + / "captures" + / "demos" + / "demo-0002" + / "missing.mp4" + ).resolve() + ), + "asset_id": "demo-0002", + } + + +def _build_context(tmp_path) -> Context: + session_directory = tmp_path / "sessions" / "session-001" + session_directory.mkdir(parents=True) + return Context( + name="session-001", + session_directory=session_directory, + manifest_path=session_directory / "session.json", + ) diff --git a/tests/test_docker.py b/tests/test_docker.py new file mode 100644 index 0000000..0399f0e --- /dev/null +++ b/tests/test_docker.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from opai.core.exceptions import OPAIWorkflowError +from opai.infrastructure import docker as docker_module + + +def test_run_mapping_container_places_docker_flags_before_image( + tmp_path, + monkeypatch, +) -> None: + work_directory = tmp_path / "work" + input_directory = tmp_path / "input" + settings_file = tmp_path / "settings.yaml" + stdout_path = tmp_path / "stdout.txt" + stderr_path = tmp_path / "stderr.txt" + work_directory.mkdir() + input_directory.mkdir() + settings_file.write_text("settings", encoding="utf-8") + prepared_video_path = input_directory / "raw_video.mp4" + prepared_video_path.write_bytes(b"video") + + xdg_runtime_dir = tmp_path / "xdg" + xauthority_file = tmp_path / ".Xauthority" + xdg_runtime_dir.mkdir() + xauthority_file.write_text("cookie", encoding="utf-8") + monkeypatch.setenv("DISPLAY", ":0") + monkeypatch.setenv("XDG_RUNTIME_DIR", str(xdg_runtime_dir)) + monkeypatch.setenv("XAUTHORITY", str(xauthority_file)) + monkeypatch.setattr(docker_module.os, "getuid", lambda: 1000) + monkeypatch.setattr(docker_module.os, "getgid", lambda: 1000) + + captured_command: list[str] = [] + + def fake_run(command: list[str], **_kwargs) -> SimpleNamespace: + captured_command[:] = command + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(docker_module.subprocess, "run", fake_run) + + docker_module.run_mapping_container( + docker_image="chicheng/orb_slam3:latest", + prepared_video_path=prepared_video_path, + work_directory=work_directory, + settings_file=settings_file, + stdout_path=stdout_path, + stderr_path=stderr_path, + ) + + image_index = captured_command.index("chicheng/orb_slam3:latest") + assert captured_command[:3] == ["docker", "run", "--rm"] + assert ( + captured_command[image_index + 1] == docker_module.DEFAULT_SLAM_MAPPING_COMMAND + ) + user_index = captured_command.index("--user") + assert captured_command[user_index : user_index + 2] == ["--user", "1000:1000"] + assert user_index < image_index + ipc_index = captured_command.index("--ipc") + assert captured_command[ipc_index : ipc_index + 2] == ["--ipc", "host"] + assert ipc_index < image_index + assert "/tmp/.X11-unix:/tmp/.X11-unix" in captured_command + assert f"{input_directory.resolve()}:/input:ro" in captured_command + assert "DISPLAY=:0" in captured_command + assert f"XDG_RUNTIME_DIR={xdg_runtime_dir}" in captured_command + assert f"XAUTHORITY={xauthority_file}" in captured_command + + +def test_run_mapping_container_requires_display_for_gui( + tmp_path, + monkeypatch, +) -> None: + work_directory = tmp_path / "work" + settings_file = tmp_path / "settings.yaml" + stdout_path = tmp_path / "stdout.txt" + stderr_path = tmp_path / "stderr.txt" + work_directory.mkdir() + settings_file.write_text("settings", encoding="utf-8") + prepared_video_path = work_directory / "raw_video.mp4" + prepared_video_path.write_bytes(b"video") + monkeypatch.delenv("DISPLAY", raising=False) + + with pytest.raises(OPAIWorkflowError, match="DISPLAY"): + docker_module.run_mapping_container( + docker_image="chicheng/orb_slam3:latest", + prepared_video_path=prepared_video_path, + work_directory=work_directory, + settings_file=settings_file, + stdout_path=stdout_path, + stderr_path=stderr_path, + ) diff --git a/tests/test_facade.py b/tests/test_facade.py index aef85b3..bc832bc 100644 --- a/tests/test_facade.py +++ b/tests/test_facade.py @@ -10,6 +10,7 @@ import opai from opai.application import calibration as calibration_module +from opai.application import slam as slam_module from opai.core.exceptions import ( OPAIContextError, OPAIDependencyError, @@ -48,6 +49,39 @@ def test_verify_calibrated_parameters_requires_context(tmp_path, monkeypatch) -> ) +def test_run_extract_trajectories_batch_requires_context( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(context_store, "_ACTIVE_CONTEXT", None) + + with pytest.raises(OPAIContextError, match="Call opai.init"): + opai.run_extract_trajectories_batch() + + +def test_run_extract_trajectories_batch_uses_active_context( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + ctx = opai.init("session-001") + expected = {"processed_videos": [], "total_processed": 0} + + def fake_run_extract_trajectories_batch(*, ctx): + assert ctx == opai.get_context() + return expected + + monkeypatch.setattr( + slam_module, + "run_extract_trajectories_batch", + fake_run_extract_trajectories_batch, + ) + + assert opai.run_extract_trajectories_batch() == expected + assert opai.get_context() == ctx + + def test_init_creates_context_directory(tmp_path, monkeypatch) -> None: monkeypatch.chdir(tmp_path) ctx = opai.init("session-001") diff --git a/tests/test_gopro.py b/tests/test_gopro.py index b11779d..ffa0cdc 100644 --- a/tests/test_gopro.py +++ b/tests/test_gopro.py @@ -1,4 +1,5 @@ import json +import logging from pathlib import Path import httpx @@ -259,6 +260,16 @@ def test_register_gopro_skips_thumbnail_download_when_index_exists( ) == {"items": []} +def test_register_gopro_logs_socket_address(tmp_path, caplog) -> None: + ctx = Context(name="session-001", session_directory=tmp_path) + caplog.set_level(logging.INFO, logger="opai") + + register_gopro(ctx, "C3441320092154", download_thumbnails=False) + + assert "Registered GoPro for session session-001" in caplog.text + assert "172.21.154.51:8080" in caplog.text + + def test_download_thumbnail_from_gopro_uses_thumbnail_endpoint( tmp_path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -394,7 +405,6 @@ async def __aenter__(self) -> "_FakeAsyncResponse": async def __aexit__(self, exc_type, exc, tb) -> None: del exc_type, exc, tb - return None def raise_for_status(self) -> None: return None @@ -422,7 +432,6 @@ async def __aenter__(self) -> "_FakeAsyncClient": async def __aexit__(self, exc_type, exc, tb) -> None: del exc_type, exc, tb - return None async def get(self, url: str): response = self._get_responses[url] @@ -449,7 +458,6 @@ def __enter__(self) -> "_FakeProgressBar": def __exit__(self, exc_type, exc, tb) -> None: del exc_type, exc, tb - return None def update(self, amount: int) -> None: self.updates.append(amount) diff --git a/tests/test_imu.py b/tests/test_imu.py new file mode 100644 index 0000000..0c50a83 --- /dev/null +++ b/tests/test_imu.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import json +import logging + +import numpy as np +import pytest + +import opai.application.imu as imu_module +from opai.core.exceptions import OPAIWorkflowError +from opai.domain.imu import IMUPayload, IMURecording, IMUSample, IMUStream + + +def test_imu_payload_to_dict_matches_wire_shape() -> None: + payload = IMUPayload( + recording=IMURecording( + streams={ + "ACCL": IMUStream( + samples=( + IMUSample(value=[1.0, 2.0, 3.0], cts=125.0), + IMUSample(value=4.0, cts=250.0), + ) + ) + } + ), + frames_per_second=30.0, + ) + + assert payload.to_dict() == { + "1": { + "streams": { + "ACCL": { + "samples": [ + {"value": [1.0, 2.0, 3.0], "cts": 125.0}, + {"value": 4.0, "cts": 250.0}, + ] + } + } + }, + "frames/second": 30.0, + } + + +def test_extract_imu_from_video_writes_typed_payload_json( + tmp_path, monkeypatch +) -> None: + telemetry = { + "ACCL": ( + [np.array([1.0, 2.0, 3.0]), np.array([4.0, 5.0, 6.0])], + np.array([0.125, 0.25]), + ), + "GYRO": ( + [np.array([7.0, 8.0, 9.0])], + np.array([0.5]), + ), + } + monkeypatch.setattr( + imu_module, + "GoProTelemetryExtractor", + lambda video_path: _FakeExtractor(video_path, telemetry=telemetry), + ) + + output_path = imu_module.extract_imu_from_video( + tmp_path / "input.mp4", + tmp_path / "imu_data.json", + stream_types=["ACCL", "GYRO"], + ) + + assert output_path == tmp_path / "imu_data.json" + assert json.loads(output_path.read_text(encoding="utf-8")) == { + "1": { + "streams": { + "ACCL": { + "samples": [ + {"value": [1.0, 2.0, 3.0], "cts": 125.0}, + {"value": [4.0, 5.0, 6.0], "cts": 250.0}, + ] + }, + "GYRO": { + "samples": [ + {"value": [7.0, 8.0, 9.0], "cts": 500.0}, + ] + }, + } + }, + "frames/second": 0.0, + } + + +def test_extract_imu_from_video_omits_empty_streams(tmp_path, monkeypatch) -> None: + telemetry = { + "ACCL": ( + [np.array([1.0, 2.0, 3.0])], + np.array([0.125]), + ), + "GYRO": ( + [], + np.array([]), + ), + } + monkeypatch.setattr( + imu_module, + "GoProTelemetryExtractor", + lambda video_path: _FakeExtractor(video_path, telemetry=telemetry), + ) + + output_path = imu_module.extract_imu_from_video( + tmp_path / "input.mp4", + tmp_path / "imu_data.json", + stream_types=["ACCL", "GYRO"], + ) + + assert json.loads(output_path.read_text(encoding="utf-8")) == { + "1": { + "streams": { + "ACCL": { + "samples": [ + {"value": [1.0, 2.0, 3.0], "cts": 125.0}, + ] + } + } + }, + "frames/second": 0.0, + } + + +def test_extract_imu_from_video_logs_output_path( + tmp_path, + monkeypatch, + caplog, +) -> None: + telemetry = { + "ACCL": ( + [np.array([1.0, 2.0, 3.0])], + np.array([0.125]), + ) + } + monkeypatch.setattr( + imu_module, + "GoProTelemetryExtractor", + lambda video_path: _FakeExtractor(video_path, telemetry=telemetry), + ) + caplog.set_level(logging.INFO, logger="opai") + + output_path = imu_module.extract_imu_from_video( + tmp_path / "input.mp4", + tmp_path / "imu_data.json", + stream_types=["ACCL"], + ) + + assert f"Wrote IMU payload to {output_path}" in caplog.text + + +def test_extract_imu_from_video_wraps_extractor_failures(tmp_path, monkeypatch) -> None: + monkeypatch.setattr( + imu_module, + "GoProTelemetryExtractor", + lambda video_path: _RaisingExtractor(video_path, "explode"), + ) + video_path = tmp_path / "input.mp4" + output_path = tmp_path / "imu_data.json" + + with pytest.raises(OPAIWorkflowError, match="Error processing") as exc_info: + imu_module.extract_imu_from_video(video_path, output_path) + + assert str(video_path) in str(exc_info.value) + assert "explode" in str(exc_info.value) + + +class _FakeExtractor: + def __init__( + self, + video_path: str, + *, + telemetry: dict[str, tuple[list[np.ndarray], np.ndarray]], + ) -> None: + self.video_path = video_path + self.telemetry = telemetry + self.opened = False + self.closed = False + + def open_source(self) -> None: + self.opened = True + + def extract_data( + self, stream_type: str + ) -> tuple[list[np.ndarray], np.ndarray] | None: + return self.telemetry.get(stream_type) + + def close_source(self) -> None: + self.closed = True + + +class _RaisingExtractor: + def __init__(self, video_path: str, message: str) -> None: + self.video_path = video_path + self.message = message + + def open_source(self) -> None: + raise RuntimeError(self.message) + + def close_source(self) -> None: + return None diff --git a/tests/test_logger.py b/tests/test_logger.py new file mode 100644 index 0000000..b3788ed --- /dev/null +++ b/tests/test_logger.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import logging + +import numpy as np + +from opai.domain.calibration import CalibrationIntrinsics, CalibrationResult +from opai.infrastructure.logger import get_logger +from opai.infrastructure.persistence import write_calibration_result + + +def test_get_logger_uses_opai_namespace() -> None: + logger = get_logger("infrastructure.persistence") + + assert logger.name == "opai.infrastructure.persistence" + + +def test_write_calibration_result_emits_native_log_record(tmp_path, caplog) -> None: + caplog.set_level(logging.INFO, logger="opai") + + result = CalibrationResult( + mse_reproj_error=0.12, + image_height=1080, + image_width=1920, + intrinsic_type="fisheye", + intrinsics=CalibrationIntrinsics( + aspect_ratio=1.0, + focal_length=500.0, + principal_pt_x=960.0, + principal_pt_y=540.0, + radial_distortion_1=0.1, + radial_distortion_2=0.01, + radial_distortion_3=0.001, + radial_distortion_4=0.0001, + skew=0.0, + ), + camera_matrix=np.eye(3), + dist_coeffs=np.zeros((4, 1)), + ) + + output_path = write_calibration_result(tmp_path, result) + + assert output_path.exists() + assert f"Wrote calibration result to {output_path}" in caplog.text diff --git a/tests/test_mapping.py b/tests/test_mapping.py new file mode 100644 index 0000000..3e431ba --- /dev/null +++ b/tests/test_mapping.py @@ -0,0 +1,617 @@ +from __future__ import annotations + +import logging +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +import opai +from opai.application import slam as slam_module +from opai.core.exceptions import OPAIContextError, OPAIWorkflowError +from opai.domain.context import Context +from opai.infrastructure import context_store + + +def test_run_mapping_requires_context(tmp_path, monkeypatch) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(context_store, "_ACTIVE_CONTEXT", None) + + with pytest.raises(OPAIContextError, match="Call opai.init"): + opai.run_mapping(slam_settings_file=tmp_path / "slam_settings.yaml") + + +def test_add_mapping_video_replaces_active_mapping(tmp_path, monkeypatch) -> None: + monkeypatch.chdir(tmp_path) + opai.init("session-001") + mapping_a = tmp_path / "mapping-a.mp4" + mapping_b = tmp_path / "mapping-b.mp4" + mapping_a.write_bytes(b"first") + mapping_b.write_bytes(b"second") + + first = opai.add_mapping_video(mapping_a) + second = opai.add_mapping_video(mapping_b) + + assert first.original_filename == "mapping-a.mp4" + assert second.original_filename == "mapping-b.mp4" + mapping_dir = ( + tmp_path / "sessions" / "session-001" / "captures" / "mapping" / "current" + ) + assert sorted(path.name for path in mapping_dir.iterdir()) == ["mapping-b.mp4"] + + +def test_run_mapping_requires_registered_mapping_video(tmp_path) -> None: + session_directory = tmp_path / "sessions" / "session-001" + session_directory.mkdir(parents=True) + ctx = Context( + name="session-001", + session_directory=session_directory, + manifest_path=session_directory / "session.json", + ) + settings_path = tmp_path / "slam_settings.yaml" + settings_path.write_text("settings", encoding="utf-8") + + with pytest.raises(OPAIWorkflowError, match="No mapping video found"): + slam_module.run_mapping(ctx, slam_settings_file=settings_path) + + +def test_clear_mapping_outputs_logs_directory(tmp_path, caplog) -> None: + work_directory = tmp_path / "slam" / "mapping" + work_directory.mkdir(parents=True) + stale_output = work_directory / "imu_data.json" + stale_output.write_text("stale", encoding="utf-8") + caplog.set_level(logging.INFO, logger="opai") + + slam_module._clear_mapping_outputs(work_directory) + + assert not stale_output.exists() + assert f"Clearing mapping outputs in {work_directory}" in caplog.text + + +def test_run_mapping_always_reruns_and_overwrites_outputs( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + ctx = opai.init("session-001") + mapping_video = tmp_path / "mapping.mp4" + mapping_video.write_bytes(b"mapping") + opai.add_mapping_video(mapping_video) + settings_path = tmp_path / "slam_settings.yaml" + settings_path.write_text("settings", encoding="utf-8") + + events: list[str] = [] + container_runs: list[int] = [] + prepared_video_existed_before_convert: list[bool] = [] + config_dir = tmp_path / "configs" / "slam" + config_dir.mkdir(parents=True) + config_path = config_dir / settings_path.name + config_path.write_text("settings", encoding="utf-8") + + def fake_get_slam_config_dir() -> Path: + return config_dir + + def fake_get_video_fps(video_path: Path) -> float: + assert video_path.name == "mapping.mp4" + return 120.0 + + def fake_convert_video_to_fps( + video_path: Path, + output_path: Path, + target_fps: int, + ) -> Path: + events.append(f"convert:{target_fps}") + prepared_video_existed_before_convert.append(output_path.exists()) + assert video_path.name == "mapping.mp4" + assert target_fps == slam_module.TARGET_SLAM_FPS + output_path.write_text("prepared", encoding="utf-8") + return output_path + + def fake_get_video_resolution(video_path: Path) -> tuple[int, int]: + assert video_path.name == slam_module.PREPARED_MAPPING_VIDEO_FILENAME + return slam_module.DEFAULT_RESOLUTION + + def fake_extract(video_path: Path, dest: Path) -> Path: + events.append("extract") + assert video_path.name == "mapping.mp4" + dest.write_text("imu", encoding="utf-8") + return dest + + def fake_draw_mask(mask, **_kwargs): + events.append("mask") + return mask + + def fake_imwrite(path: str, _image) -> bool: + Path(path).write_text("mask", encoding="utf-8") + return True + + def fake_pull(docker_image: str) -> None: + events.append(f"pull:{docker_image}") + + def fake_run( + *, + docker_image: str, + prepared_video_path: Path, + work_directory: Path, + settings_file: Path, + stdout_path: Path, + stderr_path: Path, + enable_gui: bool, + ) -> None: + events.append(f"run:{docker_image}") + container_runs.append(len(container_runs) + 1) + assert prepared_video_path.name == slam_module.PREPARED_MAPPING_VIDEO_FILENAME + assert settings_file == config_path + assert enable_gui is False + (work_directory / "map_atlas.osa").write_text( + f"map-{container_runs[-1]}", + encoding="utf-8", + ) + (work_directory / "mapping_camera_trajectory.csv").write_text( + f"trajectory-{container_runs[-1]}", + encoding="utf-8", + ) + stdout_path.write_text(f"stdout-{container_runs[-1]}", encoding="utf-8") + stderr_path.write_text(f"stderr-{container_runs[-1]}", encoding="utf-8") + + monkeypatch.setattr(slam_module, "get_slam_config_dir", fake_get_slam_config_dir) + monkeypatch.setattr(slam_module, "get_video_fps", fake_get_video_fps) + monkeypatch.setattr(slam_module, "convert_video_to_fps", fake_convert_video_to_fps) + monkeypatch.setattr(slam_module, "get_video_resolution", fake_get_video_resolution) + monkeypatch.setattr(slam_module, "extract_imu_from_video", fake_extract) + monkeypatch.setattr(slam_module, "draw_predefined_mask", fake_draw_mask) + monkeypatch.setattr(slam_module.cv2, "imwrite", fake_imwrite) + monkeypatch.setattr(slam_module, "pull_docker_image", fake_pull) + monkeypatch.setattr(slam_module, "run_mapping_container", fake_run) + + first = slam_module.run_mapping( + ctx, + slam_settings_file=settings_path.name, + ) + second = slam_module.run_mapping( + ctx, + slam_settings_file=settings_path.name, + ) + + assert len(container_runs) == 2 + assert prepared_video_existed_before_convert == [False, False] + assert events.count(f"convert:{slam_module.TARGET_SLAM_FPS}") == 2 + assert events.count("extract") == 2 + assert events.count("mask") == 2 + assert sum(event.startswith("pull:") for event in events) == 2 + assert sum(event.startswith("run:") for event in events) == 2 + assert first.input_video_path.name == "mapping.mp4" + assert first.map_path.read_text(encoding="utf-8") == "map-2" + assert second.map_path.read_text(encoding="utf-8") == "map-2" + assert second.trajectory_csv_path.read_text(encoding="utf-8") == "trajectory-2" + + +def test_run_mapping_uses_original_input_when_video_is_already_60fps( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + ctx = opai.init("session-001") + mapping_video = tmp_path / "mapping.mp4" + mapping_video.write_bytes(b"mapping") + opai.add_mapping_video(mapping_video) + settings_path = tmp_path / "slam_settings.yaml" + settings_path.write_text("settings", encoding="utf-8") + + config_dir = tmp_path / "configs" / "slam" + config_dir.mkdir(parents=True) + config_path = config_dir / settings_path.name + config_path.write_text("settings", encoding="utf-8") + events: list[str] = [] + + def fake_get_slam_config_dir() -> Path: + return config_dir + + def fake_get_video_fps(video_path: Path) -> float: + assert video_path.name == "mapping.mp4" + return 60.0 + + def fake_convert_video_to_fps(*_args, **_kwargs) -> Path: + raise AssertionError("60 fps input should not be converted") + + def fake_get_video_resolution(video_path: Path) -> tuple[int, int]: + assert video_path.name == "mapping.mp4" + return slam_module.DEFAULT_RESOLUTION + + def fake_extract(video_path: Path, dest: Path) -> Path: + events.append("extract") + assert video_path.name == "mapping.mp4" + dest.write_text("imu", encoding="utf-8") + return dest + + def fake_draw_mask(mask, **_kwargs): + events.append("mask") + return mask + + def fake_imwrite(path: str, _image) -> bool: + Path(path).write_text("mask", encoding="utf-8") + return True + + def fake_pull(_docker_image: str) -> None: + events.append("pull") + + def fake_run( + *, + docker_image: str, + prepared_video_path: Path, + work_directory: Path, + settings_file: Path, + stdout_path: Path, + stderr_path: Path, + enable_gui: bool, + ) -> None: + events.append("run") + assert docker_image == slam_module.DEFAULT_DOCKER_IMAGE + assert prepared_video_path.name == "mapping.mp4" + assert settings_file == config_path + assert work_directory == ctx.session_directory / "slam" / "mapping" + assert enable_gui is False + stdout_path.write_text("stdout", encoding="utf-8") + stderr_path.write_text("stderr", encoding="utf-8") + + monkeypatch.setattr(slam_module, "get_slam_config_dir", fake_get_slam_config_dir) + monkeypatch.setattr(slam_module, "get_video_fps", fake_get_video_fps) + monkeypatch.setattr(slam_module, "convert_video_to_fps", fake_convert_video_to_fps) + monkeypatch.setattr(slam_module, "get_video_resolution", fake_get_video_resolution) + monkeypatch.setattr(slam_module, "extract_imu_from_video", fake_extract) + monkeypatch.setattr(slam_module, "draw_predefined_mask", fake_draw_mask) + monkeypatch.setattr(slam_module.cv2, "imwrite", fake_imwrite) + monkeypatch.setattr(slam_module, "pull_docker_image", fake_pull) + monkeypatch.setattr(slam_module, "run_mapping_container", fake_run) + + result = slam_module.run_mapping( + ctx, + slam_settings_file=settings_path.name, + ) + + assert result.input_video_path.name == "mapping.mp4" + assert events == ["extract", "mask", "pull", "run"] + + +def test_run_mapping_raises_if_120fps_conversion_fails( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + ctx = opai.init("session-001") + mapping_video = tmp_path / "mapping.mp4" + mapping_video.write_bytes(b"mapping") + opai.add_mapping_video(mapping_video) + settings_path = tmp_path / "slam_settings.yaml" + settings_path.write_text("settings", encoding="utf-8") + + config_dir = tmp_path / "configs" / "slam" + config_dir.mkdir(parents=True) + config_path = config_dir / settings_path.name + config_path.write_text("settings", encoding="utf-8") + container_started = False + image_pulled = False + + def fake_get_slam_config_dir() -> Path: + return config_dir + + def fake_get_video_fps(video_path: Path) -> float: + assert video_path.name == "mapping.mp4" + return 120.0 + + def fake_convert_video_to_fps(*_args, **_kwargs) -> Path: + raise OPAIWorkflowError("conversion failed") + + def fake_pull(_docker_image: str) -> None: + nonlocal image_pulled + image_pulled = True + + def fake_run(**_kwargs) -> None: + nonlocal container_started + container_started = True + + monkeypatch.setattr(slam_module, "get_slam_config_dir", fake_get_slam_config_dir) + monkeypatch.setattr(slam_module, "get_video_fps", fake_get_video_fps) + monkeypatch.setattr(slam_module, "convert_video_to_fps", fake_convert_video_to_fps) + monkeypatch.setattr(slam_module, "pull_docker_image", fake_pull) + monkeypatch.setattr(slam_module, "run_mapping_container", fake_run) + + with pytest.raises(OPAIWorkflowError, match="conversion failed"): + slam_module.run_mapping( + ctx, + slam_settings_file=config_path.name, + ) + + assert image_pulled is False + assert container_started is False + + +def test_draw_predefined_mask_uses_width_height_image_order() -> None: + image_width, image_height = slam_module.DEFAULT_RESOLUTION + mask = np.zeros((image_height, image_width), dtype=np.uint8) + + slam_module.draw_predefined_mask( + mask, color=255, mirror=False, gripper=False, finger=True + ) + + probe_row = int(image_height * 0.7) + probe_col = image_width // 2 + assert mask[probe_row, probe_col] == 255 + + +def test_gripper_polygon_round_trips_within_default_resolution() -> None: + image_width, image_height = slam_module.DEFAULT_RESOLUTION + left_gripper = slam_module.get_gripper_canonical_polygon()[0] + points = slam_module.canonical_to_pixel_coords( + left_gripper, slam_module.DEFAULT_RESOLUTION + ) + + assert np.all(points[:, 0] >= 0) + assert np.all(points[:, 0] <= image_width) + assert np.all(points[:, 1] >= 0) + assert np.all(points[:, 1] <= image_height) + + +def test_run_extract_trajectories_batch_requires_map_atlas(tmp_path) -> None: + session_directory = tmp_path / "sessions" / "session-001" + session_directory.mkdir(parents=True) + ctx = Context( + name="session-001", + session_directory=session_directory, + manifest_path=session_directory / "session.json", + ) + + with pytest.raises(OPAIWorkflowError, match="map atlas file"): + slam_module.run_extract_trajectories_batch(ctx) + + +def test_run_extract_trajectories_batch_processes_registered_demos( + tmp_path, + monkeypatch, +) -> None: + session_directory = tmp_path / "sessions" / "session-001" + atlas_path = session_directory / "slam" / "mapping" / "map_atlas.osa" + atlas_path.parent.mkdir(parents=True, exist_ok=True) + atlas_path.write_text("atlas", encoding="utf-8") + settings_path = session_directory / "slam_settings.yaml" + settings_path.parent.mkdir(parents=True, exist_ok=True) + settings_path.write_text("settings", encoding="utf-8") + + demo_a_directory = session_directory / "captures" / "demos" / "demo-0001" + demo_b_directory = session_directory / "captures" / "demos" / "demo-0002" + demo_a_directory.mkdir(parents=True, exist_ok=True) + demo_b_directory.mkdir(parents=True, exist_ok=True) + demo_a_video = demo_a_directory / "demo-a.mp4" + demo_b_video = demo_b_directory / "demo-b.mp4" + demo_a_video.write_bytes(b"demo-a") + demo_b_video.write_bytes(b"demo-b") + (demo_b_directory / "camera_trajectory.csv").write_text( + "existing", + encoding="utf-8", + ) + ctx = Context( + name="session-001", + session_directory=session_directory, + manifest_path=session_directory / "session.json", + ) + + events: list[str] = [] + + def fake_get_demo_videos(_ctx: Context) -> list[Path]: + return [demo_a_video, demo_b_video] + + def fake_pull(docker_image: str) -> None: + events.append(f"pull:{docker_image}") + + def fake_get_video_fps(video_path: Path) -> float: + assert video_path == demo_a_video + return 120.0 + + def fake_convert_video_to_fps( + video_path: Path, + output_path: Path, + target_fps: int, + ) -> Path: + events.append(f"convert:{target_fps}") + assert video_path == demo_a_video + output_path.write_text("prepared", encoding="utf-8") + return output_path + + def fake_get_video_resolution(video_path: Path) -> tuple[int, int]: + assert video_path.name == slam_module.PREPARED_TRAJECTORY_VIDEO_FILENAME + return slam_module.DEFAULT_RESOLUTION + + def fake_extract(video_path: Path, destination: Path) -> Path: + events.append("extract") + assert video_path == demo_a_video + destination.write_text("imu", encoding="utf-8") + return destination + + def fake_draw_mask(mask, **_kwargs): + events.append("mask") + return mask + + def fake_imwrite(path: str, _image) -> bool: + Path(path).write_text("mask", encoding="utf-8") + return True + + def fake_get_video_duration_seconds(video_path: Path) -> float: + assert video_path.name == slam_module.PREPARED_TRAJECTORY_VIDEO_FILENAME + return 12.5 + + def fake_run_trajectory_extraction_container( + *, + docker_image: str, + prepared_video_path: Path, + work_directory: Path, + atlas_path: Path, + slam_setting_path: Path, + imu_json_path: Path, + trajectory_csv_path: Path, + mask_path: Path, + stdout_path: Path, + stderr_path: Path, + timeout_seconds: float | None, + ) -> SimpleNamespace: + events.append("run") + assert docker_image == slam_module.DEFAULT_DOCKER_IMAGE + assert prepared_video_path == ( + demo_a_directory / slam_module.PREPARED_TRAJECTORY_VIDEO_FILENAME + ) + assert work_directory == demo_a_directory + assert atlas_path == session_directory / "slam" / "mapping" / "map_atlas.osa" + assert slam_setting_path == settings_path + assert imu_json_path == demo_a_directory / "imu_data.json" + assert trajectory_csv_path == demo_a_directory / "camera_trajectory.csv" + assert mask_path == demo_a_directory / "slam_mask.png" + assert timeout_seconds == ( + 12.5 * slam_module.DEFAULT_TRAJECTORY_TIMEOUT_MULTIPLE + ) + trajectory_csv_path.write_text("trajectory", encoding="utf-8") + stdout_path.write_text("stdout", encoding="utf-8") + stderr_path.write_text("stderr", encoding="utf-8") + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(slam_module, "get_demo_videos", fake_get_demo_videos) + monkeypatch.setattr(slam_module, "pull_docker_image", fake_pull) + monkeypatch.setattr(slam_module, "get_video_fps", fake_get_video_fps) + monkeypatch.setattr(slam_module, "convert_video_to_fps", fake_convert_video_to_fps) + monkeypatch.setattr(slam_module, "get_video_resolution", fake_get_video_resolution) + monkeypatch.setattr(slam_module, "extract_imu_from_video", fake_extract) + monkeypatch.setattr(slam_module, "draw_predefined_mask", fake_draw_mask) + monkeypatch.setattr(slam_module.cv2, "imwrite", fake_imwrite) + monkeypatch.setattr( + slam_module, + "get_video_duration_seconds", + fake_get_video_duration_seconds, + ) + monkeypatch.setattr( + slam_module, + "run_trajectory_extraction_container", + fake_run_trajectory_extraction_container, + ) + + result = slam_module.run_extract_trajectories_batch(ctx) + + assert events == [ + f"pull:{slam_module.DEFAULT_DOCKER_IMAGE}", + f"convert:{slam_module.TARGET_SLAM_FPS}", + "extract", + "mask", + "run", + ] + assert result["total_processed"] == 1 + assert result["processed_videos"] == [ + { + "demo_id": "demo-0001", + "video_path": str(demo_a_video), + "prepared_video_path": str( + demo_a_directory / slam_module.PREPARED_TRAJECTORY_VIDEO_FILENAME + ), + "trajectory_csv": str(demo_a_directory / "camera_trajectory.csv"), + "stdout_log": str(demo_a_directory / "slam_stdout.txt"), + "stderr_log": str(demo_a_directory / "slam_stderr.txt"), + "status": "success", + "error_message": None, + } + ] + + +def test_run_extract_trajectories_batch_records_timeouts_and_failures( + tmp_path, + monkeypatch, +) -> None: + session_directory = tmp_path / "sessions" / "session-001" + atlas_path = session_directory / "slam" / "mapping" / "map_atlas.osa" + atlas_path.parent.mkdir(parents=True, exist_ok=True) + atlas_path.write_text("atlas", encoding="utf-8") + settings_path = session_directory / "slam_settings.yaml" + settings_path.parent.mkdir(parents=True, exist_ok=True) + settings_path.write_text("settings", encoding="utf-8") + + demo_a_directory = session_directory / "captures" / "demos" / "demo-0001" + demo_b_directory = session_directory / "captures" / "demos" / "demo-0002" + demo_a_directory.mkdir(parents=True, exist_ok=True) + demo_b_directory.mkdir(parents=True, exist_ok=True) + demo_a_video = demo_a_directory / "demo-a.mp4" + demo_b_video = demo_b_directory / "demo-b.mp4" + demo_a_video.write_bytes(b"demo-a") + demo_b_video.write_bytes(b"demo-b") + ctx = Context( + name="session-001", + session_directory=session_directory, + manifest_path=session_directory / "session.json", + ) + + def fake_get_demo_videos(_ctx: Context) -> list[Path]: + return [demo_a_video, demo_b_video] + + def fake_pull(_docker_image: str) -> None: + return None + + def fake_get_video_fps(_video_path: Path) -> float: + return 60.0 + + def fake_get_video_resolution(_video_path: Path) -> tuple[int, int]: + return slam_module.DEFAULT_RESOLUTION + + def fake_extract(_video_path: Path, destination: Path) -> Path: + destination.write_text("imu", encoding="utf-8") + return destination + + def fake_draw_mask(mask, **_kwargs): + return mask + + def fake_imwrite(path: str, _image) -> bool: + Path(path).write_text("mask", encoding="utf-8") + return True + + def fake_get_video_duration_seconds(_video_path: Path) -> float: + return 5.0 + + def fake_run_trajectory_extraction_container( + *, + work_directory: Path, + stdout_path: Path, + stderr_path: Path, + **_kwargs, + ) -> SimpleNamespace: + stdout_path.write_text("stdout", encoding="utf-8") + stderr_path.write_text("stderr", encoding="utf-8") + if work_directory == demo_a_directory: + raise subprocess.TimeoutExpired(cmd=["docker"], timeout=50.0) + return SimpleNamespace(returncode=2) + + monkeypatch.setattr(slam_module, "get_demo_videos", fake_get_demo_videos) + monkeypatch.setattr(slam_module, "pull_docker_image", fake_pull) + monkeypatch.setattr(slam_module, "get_video_fps", fake_get_video_fps) + monkeypatch.setattr(slam_module, "get_video_resolution", fake_get_video_resolution) + monkeypatch.setattr(slam_module, "extract_imu_from_video", fake_extract) + monkeypatch.setattr(slam_module, "draw_predefined_mask", fake_draw_mask) + monkeypatch.setattr(slam_module.cv2, "imwrite", fake_imwrite) + monkeypatch.setattr( + slam_module, + "get_video_duration_seconds", + fake_get_video_duration_seconds, + ) + monkeypatch.setattr( + slam_module, + "run_trajectory_extraction_container", + fake_run_trajectory_extraction_container, + ) + + result = slam_module.run_extract_trajectories_batch(ctx) + + assert result["total_processed"] == 2 + assert [entry["status"] for entry in result["processed_videos"]] == [ + "timeout", + "failed", + ] + assert ( + result["processed_videos"][0]["error_message"] == "SLAM extraction timed out." + ) + assert result["processed_videos"][1]["error_message"] == ( + "SLAM extraction exited with code 2. Inspect the stderr log for details." + ) diff --git a/uv.lock b/uv.lock index 8fd5187..a3e5378 100644 --- a/uv.lock +++ b/uv.lock @@ -1757,6 +1757,7 @@ dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "opencv-contrib-python-headless" }, + { name = "py-gpmf-parser" }, { name = "pydantic" }, { name = "rich" }, { name = "tqdm" }, @@ -1779,6 +1780,7 @@ requires-dist = [ { name = "numpy", specifier = ">=2.2.6" }, { name = "opencv-contrib-python-headless", specifier = ">=4.13.0.92" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = "==3.3.2" }, + { name = "py-gpmf-parser", specifier = ">=0.1.1" }, { name = "pydantic", specifier = ">=2.12.5" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.308" }, { name = "pytest", marker = "extra == 'dev'" }, @@ -2054,6 +2056,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, ] +[[package]] +name = "py-gpmf-parser" +version = "0.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/0b/2a00e75a281c0de0839a1e210eebee566cc62c35d59264c8fe00b6c75df3/py_gpmf_parser-0.1.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:993836967dc540f8419e69b94e23599944f71f70ffdc0b2c8ee1db037bd4cfbd", size = 2224375, upload-time = "2025-08-16T13:46:48.987Z" }, + { url = "https://files.pythonhosted.org/packages/45/55/b7d06b5fb83652479f95bd41468642b18b8c0614dd1791dd4f585fcf1af3/py_gpmf_parser-0.1.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1d874a1c4925ab371ccc2a5375257478f39ed84a7c8cc3eb8994a2933d4b8693", size = 2245060, upload-time = "2025-08-16T13:46:39.145Z" }, + { url = "https://files.pythonhosted.org/packages/9f/8b/c3d5ea5ff587bb84f5b2f6a896b0a7c7b42e49db763f3b4f882aa686e2f7/py_gpmf_parser-0.1.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4388b321e5d3d85614bb31650b17bddcf858ebae0f86c7c4a0ea81ccdf1dd085", size = 2276941, upload-time = "2025-08-16T13:46:44.481Z" }, +] + [[package]] name = "pycparser" version = "3.0"