返回汇报

RGB ViPE:共享 DA3 深度,关闭内部深度

原路径:/SSD2/yangzijing/egoview-v9-foundationhand-s9l10-20260901/egoview/v5_pipeline/vipe_worker.py
SHA256:75c0132a4c1ee7e526192bb9f1b2ef13452f60bf5d21846eb4e410c3d0240043

   1 """Run ViPE once with the normalized V5 canonical camera held fixed."""
   2 
   3 from __future__ import annotations
   4 
   5 import argparse
   6 import json
   7 import math
   8 from pathlib import Path
   9 from typing import Any, Mapping, Sequence
  10 
  11 import numpy as np
  12 
  13 from egoview.v5_pipeline.state import (
  14     CAMERA_CONVENTION,
  15     VIPE_REQUEST_SCHEMA_VERSION,
  16     VIPE_WORKER_SCHEMA_VERSION,
  17     scene_depth_digest,
  18 )
  19 
  20 
  21 CANONICAL_VIPE_INTRINSICS_SOURCE = "normalized canonical K fixed"
  22 SCENE_CAMERA_VIPE_INTRINSICS_SOURCE = "scene camera K fixed"
  23 SCENE_CAMERA_ATTESTATIONS = {
  24     "scene_only": True,
  25     "affects_formal_hand_input": False,
  26     "affects_training_labels": False,
  27     "v5_hand_camera_is_calibrated": False,
  28     "v5_hand_camera_contract_preserved": True,
  29 }
  30 SCENE_CAMERA_REQUEST_FIELD = "scene_camera_contract_json"
  31 
  32 
  33 def _scalar_string(value: Any, *, label: str) -> str:
  34     array = np.asarray(value)
  35     if array.shape != () or array.dtype.kind not in {"S", "U"}:
  36         raise ValueError(f"{label} must be a scalar string")
  37     item = array.item()
  38     return item.decode("utf-8") if isinstance(item, bytes) else str(item)
  39 
  40 
  41 def _request_scene_camera_contract(
  42     payload: Mapping[str, np.ndarray],
  43     intrinsics: np.ndarray,
  44     image_size_hw: np.ndarray,
  45 ) -> dict[str, Any] | None:
  46     raw = payload.get(SCENE_CAMERA_REQUEST_FIELD)
  47     if raw is None:
  48         return None
  49     try:
  50         contract = json.loads(
  51             _scalar_string(raw, label=SCENE_CAMERA_REQUEST_FIELD)
  52         )
  53     except json.JSONDecodeError as error:
  54         raise ValueError("scene camera contract JSON is invalid") from error
  55     if not isinstance(contract, Mapping):
  56         raise ValueError("scene camera contract must be a JSON object")
  57     if contract.get("real_calibration_used") is not True:
  58         raise ValueError("scene camera contract must use real calibration")
  59     for name, expected in SCENE_CAMERA_ATTESTATIONS.items():
  60         if contract.get(name) != expected:
  61             raise ValueError(
  62                 f"scene camera contract must declare {name}={expected!r}"
  63             )
  64     try:
  65         contract_size = np.asarray(contract.get("image_size_hw"), dtype=np.int64)
  66         contract_intrinsics = np.asarray(
  67             contract.get("intrinsics"), dtype=np.float32
  68         )
  69     except (TypeError, ValueError):
  70         contract_size = np.empty(0, dtype=np.int64)
  71         contract_intrinsics = np.empty(0, dtype=np.float32)
  72     if (
  73         contract.get("camera_model") != "pinhole"
  74         or contract.get("coordinate_convention") != CAMERA_CONVENTION
  75         or not np.array_equal(
  76             contract_size,
  77             np.asarray(image_size_hw, dtype=np.int64),
  78         )
  79         or not np.array_equal(
  80             contract_intrinsics,
  81             np.asarray(intrinsics, dtype=np.float32),
  82         )
  83     ):
  84         raise ValueError("scene camera contract differs from the request camera arrays")
  85     return dict(contract)
  86 
  87 
  88 def _load_request(path: Path) -> dict[str, np.ndarray]:
  89     with np.load(path, allow_pickle=False) as archive:
  90         payload = {name: np.asarray(archive[name]) for name in archive.files}
  91     required = {
  92         "schema_version",
  93         "frame_ids",
  94         "frame_keys",
  95         "intrinsics",
  96         "image_size_hw",
  97         "fps",
  98         "da3_depth_camera_z_m",
  99         "da3_valid_mask",
 100         "da3_intrinsics",
 101         "da3_prediction_digest",
 102     }
 103     field_names = set(payload)
 104     if field_names not in (required, required | {SCENE_CAMERA_REQUEST_FIELD}):
 105         raise ValueError(
 106             f"ViPE request fields differ from the sealed contract: {sorted(payload)}"
 107         )
 108     if _scalar_string(payload["schema_version"], label="schema_version") != (
 109         VIPE_REQUEST_SCHEMA_VERSION
 110     ):
 111         raise ValueError("ViPE request schema is unsupported")
 112 
 113     frame_ids = payload["frame_ids"]
 114     frame_keys = payload["frame_keys"]
 115     intrinsics = payload["intrinsics"]
 116     image_size = payload["image_size_hw"]
 117     fps = payload["fps"]
 118     da3_depth = payload["da3_depth_camera_z_m"]
 119     da3_valid = payload["da3_valid_mask"]
 120     da3_intrinsics = payload["da3_intrinsics"]
 121     if (
 122         frame_ids.dtype.kind not in {"i", "u"}
 123         or frame_ids.ndim != 1
 124         or len(frame_ids) == 0
 125         or not np.array_equal(frame_ids, np.arange(len(frame_ids)))
 126     ):
 127         raise ValueError("ViPE request frame_ids must be consecutive decode positions")
 128     if frame_keys.dtype.kind not in {"S", "U"} or frame_keys.shape != frame_ids.shape:
 129         raise ValueError("ViPE request frame_keys must align with frame_ids")
 130     if intrinsics.dtype.kind != "f" or intrinsics.shape != (3, 3):
 131         raise ValueError("ViPE request intrinsics must be one floating 3x3 matrix")
 132     if (
 133         image_size.dtype.kind not in {"i", "u"}
 134         or image_size.shape != (2,)
 135         or np.any(image_size < 2)
 136     ):
 137         raise ValueError("ViPE request image_size_hw must be an integer HxW vector")
 138     if fps.shape != () or not math.isfinite(float(fps)) or float(fps) <= 0.0:
 139         raise ValueError("ViPE request FPS must be finite and positive")
 140     if not np.isfinite(intrinsics).all() or not np.array_equal(
 141         intrinsics[2], np.asarray([0.0, 0.0, 1.0], dtype=intrinsics.dtype)
 142     ):
 143         raise ValueError("ViPE request intrinsics are not a finite pinhole K")
 144     if (
 145         intrinsics[0, 0] <= 0.0
 146         or intrinsics[1, 1] <= 0.0
 147         or intrinsics[0, 1] != 0.0
 148         or intrinsics[1, 0] != 0.0
 149     ):
 150         raise ValueError("ViPE request intrinsics are not a canonical pinhole K")
 151     scene_camera_contract = _request_scene_camera_contract(
 152         payload,
 153         intrinsics,
 154         image_size,
 155     )
 156     count = len(frame_ids)
 157     if (
 158         da3_depth.dtype.kind != "f"
 159         or da3_depth.ndim != 3
 160         or da3_depth.shape[0] != count
 161         or min(da3_depth.shape[1:]) < 2
 162     ):
 163         raise ValueError("shared DA3 depth must be floating [frames,height,width]")
 164     if da3_valid.dtype.kind != "b" or da3_valid.shape != da3_depth.shape:
 165         raise ValueError("shared DA3 valid mask must match depth and be boolean")
 166     if da3_valid.any() and (
 167         not np.isfinite(da3_depth[da3_valid]).all()
 168         or np.any(da3_depth[da3_valid] <= 0.0)
 169     ):
 170         raise ValueError("shared DA3 valid depth must be finite and positive")
 171     if (
 172         da3_intrinsics.dtype.kind != "f"
 173         or da3_intrinsics.shape != (count, 3, 3)
 174         or not np.isfinite(da3_intrinsics).all()
 175         or np.any(da3_intrinsics[:, 0, 0] <= 0.0)
 176         or np.any(da3_intrinsics[:, 1, 1] <= 0.0)
 177         or np.any(da3_intrinsics[:, 0, 1] != 0.0)
 178         or np.any(da3_intrinsics[:, 1, 0] != 0.0)
 179         or not np.array_equal(
 180             da3_intrinsics[:, 2],
 181             np.broadcast_to(np.asarray([0.0, 0.0, 1.0], dtype=da3_intrinsics.dtype), (count, 3)),
 182         )
 183     ):
 184         raise ValueError("shared DA3 intrinsics are not finite canonical pinhole K")
 185     source_height, source_width = (int(value) for value in image_size)
 186     depth_height, depth_width = da3_depth.shape[1:]
 187     expected_da3_k = np.broadcast_to(
 188         np.asarray(intrinsics, dtype=np.float32), (count, 3, 3)
 189     ).copy()
 190     expected_da3_k[:, 0, [0, 2]] *= float(depth_width) / float(source_width)
 191     expected_da3_k[:, 1, [1, 2]] *= float(depth_height) / float(source_height)
 192     if not np.allclose(da3_intrinsics, expected_da3_k, atol=2e-3, rtol=2e-5):
 193         raise ValueError(
 194             "shared DA3 intrinsics/raster do not match "
 195             + (
 196                 "scene RGB camera scaling"
 197                 if scene_camera_contract is not None
 198                 else "canonical RGB camera scaling"
 199             )
 200         )
 201     digest = _scalar_string(
 202         payload["da3_prediction_digest"], label="da3_prediction_digest"
 203     )
 204     expected_digest = scene_depth_digest(
 205         {
 206             "frame_ids": frame_ids,
 207             "frame_keys": frame_keys,
 208             "depth_camera_z_m": da3_depth,
 209             "valid_mask": da3_valid,
 210             "intrinsics": da3_intrinsics,
 211         }
 212     )
 213     if digest != expected_digest:
 214         raise ValueError("shared DA3 prediction digest differs from request arrays")
 215     return payload
 216 
 217 
 218 def _resize_nearest_raster(values: np.ndarray, target_hw: tuple[int, int]) -> np.ndarray:
 219     """Resize frame-first rasters without interpolating invalid depth holes."""
 220 
 221     source = np.asarray(values)
 222     if source.ndim != 3:
 223         raise ValueError("shared raster must be [frames,height,width]")
 224     target_height, target_width = target_hw
 225     rows = np.minimum(
 226         (np.arange(target_height, dtype=np.int64) * source.shape[1]) // target_height,
 227         source.shape[1] - 1,
 228     )
 229     cols = np.minimum(
 230         (np.arange(target_width, dtype=np.int64) * source.shape[2]) // target_width,
 231         source.shape[2] - 1,
 232     )
 233     return source[:, rows[:, None], cols[None, :]]
 234 
 235 
 236 def run_worker(
 237     request_path: Path,
 238     rgb_dir: Path,
 239     output_path: Path,
 240     *,
 241     device: str,
 242     scene_mask_process_res: int,
 243     person_mask_enabled: bool,
 244     person_mask_keyframe_gap_sec: float,
 245     person_mask_expand_px: int,
 246 ) -> None:
 247     """Execute exactly one ViPE pipeline call and publish a sealed pose payload."""
 248 
 249     request = _load_request(request_path)
 250     rgb_dir = rgb_dir.expanduser().resolve()
 251     output_path = output_path.expanduser().resolve()
 252     if not rgb_dir.is_dir():
 253         raise FileNotFoundError(rgb_dir)
 254 
 255     # These imports intentionally live inside the external-runtime worker.  The
 256     # main EgoView environment does not need ViPE's torch extension stack.
 257     import torch
 258 
 259     from vipe.config import parse_typed_config
 260     from vipe.pipeline.processors import TrackAnythingProcessor
 261     from vipe.slam.system import SLAMSystem
 262     from vipe.streams.base import (
 263         FrameAttribute,
 264         ProcessedVideoStream,
 265         StreamProcessor,
 266         VideoFrame,
 267     )
 268     from vipe.streams.frame_dir_stream import FrameDirStream
 269     from vipe.utils.cameras import CameraType
 270     from vipe.utils.model_cache import ModelCache
 271 
 272     torch_device = torch.device(device)
 273     if torch_device.type != "cuda":
 274         raise ValueError("ViPE V5 trajectory inference requires a CUDA device")
 275     if person_mask_enabled and torch_device.index not in {None, 0}:
 276         raise ValueError(
 277             "ViPE TrackAnything scene masking currently requires logical cuda:0"
 278         )
 279     if scene_mask_process_res < 2:
 280         raise ValueError("scene mask process resolution must be at least two pixels")
 281     if (
 282         not math.isfinite(person_mask_keyframe_gap_sec)
 283         or person_mask_keyframe_gap_sec <= 0.0
 284     ):
 285         raise ValueError("person-mask keyframe gap must be finite and positive")
 286     if person_mask_expand_px < 0:
 287         raise ValueError("person-mask expansion cannot be negative")
 288     torch.cuda.set_device(torch_device)
 289 
 290     frame_ids = np.asarray(request["frame_ids"], dtype=np.int64)
 291     frame_keys = np.asarray(request["frame_keys"]).astype(str)
 292     image_size = np.asarray(request["image_size_hw"], dtype=np.int64)
 293     matrix = np.asarray(request["intrinsics"], dtype=np.float32)
 294     scene_camera_contract = _request_scene_camera_contract(
 295         request,
 296         matrix,
 297         image_size,
 298     )
 299     intrinsics_source = (
 300         SCENE_CAMERA_VIPE_INTRINSICS_SOURCE
 301         if scene_camera_contract is not None
 302         else CANONICAL_VIPE_INTRINSICS_SOURCE
 303     )
 304     fps = float(np.asarray(request["fps"]).item())
 305     shared_da3_digest = _scalar_string(
 306         request["da3_prediction_digest"], label="da3_prediction_digest"
 307     )
 308     intrinsics4 = torch.as_tensor(
 309         [matrix[0, 0], matrix[1, 1], matrix[0, 2], matrix[1, 2]],
 310         dtype=torch.float32,
 311     )
 312     source_height, source_width = (int(value) for value in image_size)
 313     da3_depth = np.asarray(request["da3_depth_camera_z_m"], dtype=np.float32)
 314     da3_valid = np.asarray(request["da3_valid_mask"], dtype=bool)
 315     metric_depth_rgb = _resize_nearest_raster(
 316         np.where(
 317             da3_valid & np.isfinite(da3_depth) & (da3_depth > 0.0),
 318             da3_depth,
 319             0.0,
 320         ).astype(np.float32, copy=False),
 321         (source_height, source_width),
 322     )
 323     metric_valid_rgb = _resize_nearest_raster(
 324         da3_valid.astype(bool, copy=False),
 325         (source_height, source_width),
 326     )
 327     metric_depth_rgb[~metric_valid_rgb] = 0.0
 328     mask_scale = min(
 329         1.0,
 330         float(scene_mask_process_res) / float(min(source_height, source_width)),
 331     )
 332     scene_mask_size = (
 333         max(2, int(round(source_height * mask_scale))),
 334         max(2, int(round(source_width * mask_scale))),
 335     )
 336 
 337     class CanonicalCameraProcessor:
 338         n_passes_required = 1
 339 
 340         @staticmethod
 341         def update_frame_size(previous: tuple[int, int]) -> tuple[int, int]:
 342             return previous
 343 
 344         @staticmethod
 345         def update_attributes(previous: set[FrameAttribute]) -> set[FrameAttribute]:
 346             return previous.union(
 347                 {FrameAttribute.INTRINSICS, FrameAttribute.CAMERA_TYPE}
 348             )
 349 
 350         @staticmethod
 351         def update_iterator(previous: Any, pass_idx: int) -> Any:
 352             del pass_idx
 353             for frame_index, frame in enumerate(previous):
 354                 yield CanonicalCameraProcessor()(frame_index, frame)
 355 
 356         @staticmethod
 357         def update_fps(_previous: float) -> float:
 358             return fps
 359 
 360         def __call__(self, frame_index: int, frame: VideoFrame) -> VideoFrame:
 361             if frame_index >= len(frame_ids):
 362                 raise ValueError("ViPE RGB stream has more frames than its request")
 363             if tuple(int(value) for value in frame.size()) != tuple(image_size):
 364                 raise ValueError("ViPE RGB raster differs from the canonical request")
 365             frame.intrinsics = intrinsics4.to(device=frame.device)
 366             frame.camera_type = CameraType.PINHOLE
 367             return frame
 368 
 369     class SharedDA3MetricDepthProcessor:
 370         n_passes_required = 1
 371 
 372         @staticmethod
 373         def update_frame_size(previous: tuple[int, int]) -> tuple[int, int]:
 374             return previous
 375 
 376         @staticmethod
 377         def update_attributes(previous: set[FrameAttribute]) -> set[FrameAttribute]:
 378             return previous.union({FrameAttribute.METRIC_DEPTH})
 379 
 380         def __call__(self, frame_index: int, frame: VideoFrame) -> VideoFrame:
 381             if frame_index >= len(metric_depth_rgb):
 382                 raise ValueError("shared DA3 depth has fewer frames than RGB")
 383             if tuple(int(value) for value in frame.size()) != tuple(image_size):
 384                 raise ValueError("shared DA3 depth cannot be attached to a different raster")
 385             frame.metric_depth = torch.as_tensor(
 386                 metric_depth_rgb[frame_index],
 387                 dtype=torch.float32,
 388                 device=frame.device,
 389             )
 390             return frame
 391 
 392     class PersonMaskCollectorProcessor(StreamProcessor):
 393         n_passes_required = 1
 394 
 395         def __init__(self) -> None:
 396             self.masks: list[np.ndarray] = []
 397 
 398         def __call__(self, frame_index: int, frame: VideoFrame) -> VideoFrame:
 399             if frame_index != len(self.masks):
 400                 raise ValueError("ViPE scene masks were not produced in frame order")
 401             if not person_mask_enabled:
 402                 person = torch.zeros(frame.size(), dtype=torch.bool, device=frame.device)
 403             else:
 404                 if frame.instance is None or frame.instance_phrases is None:
 405                     raise ValueError("ViPE person segmentation did not populate a mask")
 406                 labels = [
 407                     int(label)
 408                     for label, phrase in frame.instance_phrases.items()
 409                     if str(phrase).strip().casefold() == "person"
 410                 ]
 411                 person = torch.zeros_like(frame.instance, dtype=torch.bool)
 412                 for label in labels:
 413                     person |= frame.instance == label
 414                 if person_mask_expand_px:
 415                     kernel = 2 * person_mask_expand_px + 1
 416                     person = (
 417                         torch.nn.functional.max_pool2d(
 418                             person[None, None].float(),
 419                             kernel_size=kernel,
 420                             stride=1,
 421                             padding=person_mask_expand_px,
 422                         )[0, 0]
 423                         > 0.0
 424                     )
 425             if tuple(int(value) for value in person.shape) != scene_mask_size:
 426                 person = torch.nn.functional.adaptive_max_pool2d(
 427                     person[None, None].float(),
 428                     output_size=scene_mask_size,
 429                 )[0, 0] > 0.0
 430             self.masks.append(person.detach().cpu().numpy().astype(bool, copy=False))
 431             return frame
 432 
 433     raw_stream = FrameDirStream(rgb_dir, seek_range=range(0, -1, 1), name="sequence")
 434     if len(raw_stream) != len(frame_ids):
 435         raise ValueError("ViPE RGB stream frame count differs from its request")
 436     expected_names = [f"{index:06d}.png" for index in frame_ids]
 437     if [path.name for path in raw_stream.frame_files] != expected_names:
 438         raise ValueError("ViPE RGB stream files differ from normalized decode order")
 439     model_cache = ModelCache()
 440     processors: list[StreamProcessor] = [CanonicalCameraProcessor()]
 441     processors.append(SharedDA3MetricDepthProcessor())
 442     if person_mask_enabled:
 443         processors.append(
 444             TrackAnythingProcessor(
 445                 ["person"],
 446                 add_sky=False,
 447                 sam_run_gap=max(1, int(round(fps * person_mask_keyframe_gap_sec))),
 448                 mask_expand=2 * person_mask_expand_px + 1,
 449                 model_cache=model_cache,
 450             )
 451         )
 452     mask_collector = PersonMaskCollectorProcessor()
 453     processors.append(mask_collector)
 454     stream = ProcessedVideoStream(raw_stream, processors).cache(
 455         "ViPE V5 RGB person masking",
 456         online=True,
 457         async_prefetch=False,
 458     )
 459 
 460     config = parse_typed_config(
 461         hydra_args=(
 462             "pipeline=default",
 463             "streams=frame_dir_stream",
 464             "streams.base_path=.",
 465             "pipeline.init.camera_type=pinhole",
 466             "pipeline.init.intrinsics=gt",
 467             "pipeline.init.instance=null",
 468             "pipeline.init.async_prefetch=false",
 469             "pipeline.slam.optimize_intrinsics=false",
 470             "pipeline.slam.keyframe_depth=null",
 471             "pipeline.post.depth_align_model=null",
 472             "pipeline.output.save_artifacts=false",
 473             "pipeline.output.save_slam_map=false",
 474             "pipeline.output.save_viz=false",
 475         )
 476     )
 477     if (
 478         config.pipeline.slam.optimize_intrinsics is not False
 479         or config.pipeline.slam.keyframe_depth is not None
 480     ):
 481         raise ValueError(
 482             "ViPE resolved SLAM config must use fixed K and no internal metric-depth model"
 483         )
 484     slam = SLAMSystem(
 485         device=torch_device,
 486         config=config.pipeline.slam.to_dictconfig(),
 487         model_cache=model_cache,
 488     )
 489     slam_output = slam.run([stream], camera_type=CameraType.PINHOLE)
 490     # Complete the same cached preprocessing pass if SLAM stopped requesting
 491     # frames early. Already-cached frames never execute TrackAnything again.
 492     for _ in stream:
 493         pass
 494     if len(mask_collector.masks) != len(frame_ids):
 495         raise ValueError("ViPE scene exclusion does not cover every normalized frame")
 496     scene_exclusion_mask = np.stack(mask_collector.masks, axis=0)
 497 
 498     poses = (
 499         slam_output.get_view_trajectory(0)
 500         .matrix()
 501         .detach()
 502         .cpu()
 503         .numpy()
 504         .astype(np.float32)
 505     )
 506     returned4 = (
 507         slam_output.intrinsics[0].detach().cpu().numpy().astype(np.float32)
 508     )
 509     if poses.shape != (len(frame_ids), 4, 4):
 510         raise ValueError("ViPE trajectory does not cover every normalized frame")
 511     if returned4.shape != (4,) or not np.allclose(
 512         returned4, intrinsics4.numpy(), atol=1e-4, rtol=1e-6
 513     ):
 514         raise ValueError(
 515             "ViPE changed the fixed "
 516             + (
 517                 "scene camera intrinsics"
 518                 if scene_camera_contract is not None
 519                 else "normalized canonical intrinsics"
 520             )
 521         )
 522     returned_matrix = np.asarray(
 523         [
 524             [returned4[0], 0.0, returned4[2]],
 525             [0.0, returned4[1], returned4[3]],
 526             [0.0, 0.0, 1.0],
 527         ],
 528         dtype=np.float32,
 529     )
 530     metadata = {
 531         "schema_version": VIPE_WORKER_SCHEMA_VERSION,
 532         "provider": "ViPE",
 533         "vipe_run_calls": 1,
 534         "vipe_entrypoint": "vipe.slam.system.SLAMSystem.run",
 535         "input_camera_type": "pinhole",
 536         "init_intrinsics": "gt",
 537         "intrinsics_source": intrinsics_source,
 538         "intrinsics_optimization_enabled": False,
 539         "geocalib_invocations": 0,
 540         "post_depth_alignment_enabled": False,
 541         "keyframe_depth_model": None,
 542         "unidepth_used": False,
 543         "precomputed_metric_depth_provider": "depth-anything/DA3METRIC-LARGE",
 544         "shared_da3_prediction_digest": shared_da3_digest,
 545         "da3_model_invocations": 0,
 546         "translation_scale_source": "shared DA3Metric camera-Z",
 547         "metric_depth_injected_frames": int(len(metric_depth_rgb)),
 548         "metric_depth_injected_raster_hw": [source_height, source_width],
 549         "confidence_available": False,
 550         "ba_residual_available": False,
 551         "ba_residual": None,
 552         "valid_policy": "trajectory_returned_for_every_normalized_frame",
 553         "per_frame_tracking_quality_available": False,
 554         "requested_device": device,
 555         "resolved_device": str(torch_device),
 556         "input_frame_count": int(len(frame_ids)),
 557         "scene_person_mask_enabled": person_mask_enabled,
 558         "scene_person_mask_passes": int(person_mask_enabled),
 559         "scene_person_mask_prompt": "person",
 560         "scene_person_mask_keyframe_gap_sec": person_mask_keyframe_gap_sec,
 561         "scene_person_mask_expand_px": person_mask_expand_px,
 562         "scene_person_mask_raster_hw": list(scene_mask_size),
 563         "scene_person_mask_excluded_pixels": int(scene_exclusion_mask.sum()),
 564         "scene_exclusion_role": "scene_point_exclusion_only",
 565         "scene_exclusion_hand_prediction_consumed": False,
 566         "scene_exclusion_hand_root_depth_target_used": False,
 567     }
 568     if scene_camera_contract is not None:
 569         metadata.update(SCENE_CAMERA_ATTESTATIONS)
 570     output_path.parent.mkdir(parents=True, exist_ok=True)
 571     np.savez_compressed(
 572         output_path,
 573         frame_ids=frame_ids,
 574         frame_keys=frame_keys,
 575         camera_to_world=poses,
 576         input_intrinsics=matrix,
 577         vipe_intrinsics=returned_matrix,
 578         image_size_hw=image_size,
 579         scene_exclusion_mask=scene_exclusion_mask,
 580         scene_exclusion_size_hw=np.asarray(scene_mask_size, dtype=np.int64),
 581         metadata_json=np.asarray(
 582             json.dumps(metadata, sort_keys=True, separators=(",", ":"))
 583         ),
 584     )
 585 
 586 
 587 def build_parser() -> argparse.ArgumentParser:
 588     parser = argparse.ArgumentParser()
 589     parser.add_argument("--request", required=True, type=Path)
 590     parser.add_argument("--rgb-dir", required=True, type=Path)
 591     parser.add_argument("--output", required=True, type=Path)
 592     parser.add_argument("--device", required=True)
 593     parser.add_argument("--scene-mask-process-res", required=True, type=int)
 594     parser.add_argument("--person-mask-keyframe-gap-sec", required=True, type=float)
 595     parser.add_argument("--person-mask-expand-px", required=True, type=int)
 596     parser.add_argument("--disable-person-mask", action="store_true")
 597     return parser
 598 
 599 
 600 def main(argv: Sequence[str] | None = None) -> int:
 601     args = build_parser().parse_args(argv)
 602     run_worker(
 603         args.request,
 604         args.rgb_dir,
 605         args.output,
 606         device=args.device,
 607         scene_mask_process_res=args.scene_mask_process_res,
 608         person_mask_enabled=not args.disable_person_mask,
 609         person_mask_keyframe_gap_sec=args.person_mask_keyframe_gap_sec,
 610         person_mask_expand_px=args.person_mask_expand_px,
 611     )
 612     return 0
 613 
 614 
 615 if __name__ == "__main__":
 616     raise SystemExit(main())
 617 
 618 
 619 __all__ = [
 620     "VIPE_REQUEST_SCHEMA_VERSION",
 621     "VIPE_WORKER_SCHEMA_VERSION",
 622     "build_parser",
 623     "main",
 624     "run_worker",
 625 ]