accessibilityamodal/depth.py

"""Generate a depth map for the accessibility geometry pipeline.

This script keeps depth estimation separate from reconstruction so the core
pipeline can run in lightweight environments and can optionally use heavier
models such as Depth Anything V2 when their dependencies are installed.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageOps


DEFAULT_DEPTH_MODEL = 'depth-anything/Depth-Anything-V2-Small-hf'
DEFAULT_VGGT_MODEL = 'facebook/VGGT-1B'


def depth_model_identity(
    *,
    engine: str,
    model: str,
    checkpoint: str | Path | None,
    encoder: str,
    metric_depth: bool,
    vggt_model: str,
) -> dict[str, Any]:
    """Return an auditable third-party model identity without guessing revision."""

    if engine == 'depth_anything_v2':
        size = {
            'vits': 'Small',
            'vitb': 'Base',
            'vitl': 'Large',
            'vitg': 'Giant',
        }.get(encoder, encoder)
        checkpoint_name = Path(checkpoint).name.lower() if checkpoint else ''
        dataset = (
            'Hypersim'
            if 'hypersim' in checkpoint_name
            else 'VKITTI'
            if 'vkitti' in checkpoint_name
            else None
        )
        if metric_depth and dataset:
            model_id = f'depth-anything/Depth-Anything-V2-Metric-{dataset}-{size}'
        else:
            model_id = f'depth-anything/Depth-Anything-V2-{size}'
        return {
            'model_id': model_id,
            'family': 'Depth Anything V2',
            'encoder': encoder,
            'metric_training_dataset': dataset,
            'license': 'Apache-2.0' if encoder == 'vits' else 'CC-BY-NC-4.0',
            'third_party': True,
            'revision': 'local_checkpoint_revision_unrecorded',
            'revision_review_required_before_public_release': True,
        }
    if engine == 'transformers':
        return {
            'model_id': model,
            'family': 'Transformers depth-estimation pipeline',
            'license': 'see_upstream_model_card',
            'third_party': True,
            'revision': 'runtime_default_or_local_cache',
            'revision_review_required_before_public_release': True,
        }
    if engine == 'vggt':
        return {
            'model_id': vggt_model,
            'family': 'VGGT',
            'license': 'see_upstream_model_card',
            'third_party': True,
            'revision': 'local_checkpoint_or_runtime_default',
            'revision_review_required_before_public_release': True,
        }
    return {
        'model_id': 'opencv_depth_fallback',
        'family': 'nonlearned_fallback',
        'third_party': True,
    }


def read_rgb(path: str | Path, max_size: int | None = None) -> Image.Image:
    image = ImageOps.exif_transpose(Image.open(path)).convert('RGB')
    if max_size and max(image.size) > max_size:
        scale = max_size / max(image.size)
        size = (round(image.size[0] * scale), round(image.size[1] * scale))
        image = image.resize(size, Image.Resampling.LANCZOS)
    return image


def normalize_depth(depth: np.ndarray, near: float, far: float, invert: bool) -> np.ndarray:
    depth = depth.astype(np.float32)
    valid = np.isfinite(depth)
    if not np.any(valid):
        raise ValueError('Depth prediction contains no finite values')

    values = depth[valid]
    lo, hi = np.percentile(values, [1, 99])
    if hi <= lo:
        hi = lo + 1.0
    norm = np.clip((depth - lo) / (hi - lo), 0.0, 1.0)
    if invert:
        norm = 1.0 - norm
    return near + norm * (far - near)


def resize_float_map(values: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    image = Image.fromarray(values.astype(np.float32), mode='F')
    image = image.resize(size, Image.Resampling.BILINEAR)
    return np.asarray(image, dtype=np.float32)


def colorize_confidence(confidence: np.ndarray) -> Image.Image:
    confidence = confidence.astype(np.float32)
    valid = np.isfinite(confidence)
    if not np.any(valid):
        return Image.fromarray(np.zeros(confidence.shape, dtype=np.uint8), mode='L')
    lo, hi = np.percentile(confidence[valid], [2, 98])
    if hi <= lo:
        hi = lo + 1.0
    vis = np.clip((confidence - lo) / (hi - lo), 0.0, 1.0)
    vis[~valid] = 0
    return Image.fromarray((vis * 255).astype(np.uint8), mode='L')


def preprocess_square_tensor(image: Image.Image, target_size: int):
    import torch

    width, height = image.size
    max_dim = max(width, height)
    left = (max_dim - width) // 2
    top = (max_dim - height) // 2
    scale = target_size / max_dim
    crop = (
        left * scale,
        top * scale,
        (left + width) * scale,
        (top + height) * scale,
    )
    square = Image.new('RGB', (max_dim, max_dim), (0, 0, 0))
    square.paste(image, (left, top))
    square = square.resize((target_size, target_size), Image.Resampling.BICUBIC)
    array = np.asarray(square, dtype=np.float32) / 255.0
    tensor = torch.from_numpy(array).permute(2, 0, 1).contiguous()
    return tensor[None], crop


def crop_and_resize_prediction(values: np.ndarray, crop: tuple[float, float, float, float], size: tuple[int, int]) -> np.ndarray:
    left, top, right, bottom = crop
    h, w = values.shape[:2]
    left_i = max(0, min(w - 1, int(np.floor(left))))
    top_i = max(0, min(h - 1, int(np.floor(top))))
    right_i = max(left_i + 1, min(w, int(np.ceil(right))))
    bottom_i = max(top_i + 1, min(h, int(np.ceil(bottom))))
    cropped = values[top_i:bottom_i, left_i:right_i]
    if values.ndim == 2:
        return resize_float_map(cropped, size)
    channels = [resize_float_map(cropped[..., channel], size) for channel in range(values.shape[-1])]
    return np.stack(channels, axis=-1).astype(np.float32)


def write_point_cloud_ply(
    path: Path,
    points: np.ndarray,
    colors: np.ndarray,
    confidence: np.ndarray | None,
    confidence_threshold: float,
    max_points: int,
) -> int:
    valid = np.all(np.isfinite(points), axis=-1) & (points[..., 2] > 0)
    if confidence is not None:
        valid &= np.isfinite(confidence) & (confidence >= confidence_threshold)
    ys, xs = np.where(valid)
    if ys.size == 0:
        selected = np.array([], dtype=np.int64)
    elif max_points > 0 and ys.size > max_points:
        selected = np.linspace(0, ys.size - 1, max_points, dtype=np.int64)
    else:
        selected = np.arange(ys.size, dtype=np.int64)
    pts = points[ys[selected], xs[selected]]
    cols = colors[ys[selected], xs[selected]]
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open('w', encoding='ascii') as handle:
        handle.write('ply\nformat ascii 1.0\n')
        handle.write(f'element vertex {len(pts)}\n')
        handle.write('property float x\nproperty float y\nproperty float z\n')
        handle.write('property uchar red\nproperty uchar green\nproperty uchar blue\n')
        handle.write('end_header\n')
        for point, color in zip(pts, cols):
            handle.write(
                f'{point[0]:.6f} {point[1]:.6f} {point[2]:.6f} '
                f'{int(color[0])} {int(color[1])} {int(color[2])}\n'
            )
    return int(len(pts))


def heuristic_depth(image: Image.Image, near: float, far: float) -> np.ndarray:
    """Build a deterministic perspective prior from image position and edges."""
    rgb = np.asarray(image, dtype=np.float32) / 255.0
    h, w = rgb.shape[:2]
    yy = np.linspace(0.0, 1.0, h, dtype=np.float32)[:, None]
    depth = far - yy * (far - near)
    depth = np.repeat(depth, w, axis=1)

    gray = 0.299 * rgb[..., 0] + 0.587 * rgb[..., 1] + 0.114 * rgb[..., 2]
    grad_y = np.abs(np.gradient(gray, axis=0))
    if np.isfinite(grad_y).any() and float(grad_y.max()) > 0:
        grad_y = grad_y / float(grad_y.max())
        depth -= 0.08 * (far - near) * grad_y.astype(np.float32)
    return np.clip(depth, min(near, far), max(near, far)).astype(np.float32)


def resolve_device(device: str) -> int | str:
    if device == 'cpu':
        return -1
    if device == 'auto':
        import torch

        return 0 if torch.cuda.is_available() else -1
    if device.startswith('cuda'):
        if ':' in device:
            return int(device.split(':', 1)[1])
        return 0
    return device


def transformers_depth(image: Image.Image, model: str, device: str) -> np.ndarray:
    try:
        from transformers import pipeline
    except ImportError as exc:
        raise RuntimeError('transformers is not installed; install optional depth dependencies first') from exc

    estimator = pipeline('depth-estimation', model=model, device=resolve_device(device))
    result: dict[str, Any] = estimator(image)
    if 'predicted_depth' in result:
        predicted = result['predicted_depth']
        if hasattr(predicted, 'detach'):
            predicted = predicted.detach().cpu().numpy()
        depth = np.asarray(predicted, dtype=np.float32)
        if depth.ndim == 3:
            depth = depth.squeeze()
    elif 'depth' in result:
        depth = np.asarray(result['depth'], dtype=np.float32)
    else:
        raise RuntimeError(f'Unexpected depth-estimation result keys: {sorted(result)}')

    if depth.shape[:2] != (image.height, image.width):
        depth_img = Image.fromarray(depth.astype(np.float32), mode='F')
        depth_img = depth_img.resize(image.size, Image.Resampling.BILINEAR)
        depth = np.asarray(depth_img, dtype=np.float32)
    return depth.astype(np.float32)


def resolve_torch_device(device: str) -> str:
    if device != 'auto':
        return device
    import torch

    return 'cuda' if torch.cuda.is_available() else 'cpu'


def depth_anything_v2_depth(
    image: Image.Image,
    repo: str | Path,
    checkpoint: str | Path,
    encoder: str,
    device: str,
    input_size: int,
    metric: bool,
    max_depth: float,
) -> np.ndarray:
    repo = Path(repo).resolve()
    checkpoint = Path(checkpoint).resolve()
    if not checkpoint.exists():
        raise FileNotFoundError(f'Depth Anything V2 checkpoint not found: {checkpoint}')
    if not repo.exists():
        raise FileNotFoundError(f'Depth Anything V2 repo not found: {repo}')
    source_root = repo / 'metric_depth' if metric else repo
    if not source_root.exists():
        raise FileNotFoundError(f'Depth Anything V2 source root not found: {source_root}')

    sys.path.insert(0, str(source_root))
    import torch
    from depth_anything_v2.dpt import DepthAnythingV2
    import depth_anything_v2.dinov2_layers.attention as da_attention
    import depth_anything_v2.dinov2_layers.block as da_block

    model_configs = {
        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
        'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]},
    }
    if encoder not in model_configs:
        raise ValueError(f'Unsupported Depth Anything V2 encoder: {encoder}')
    config = dict(model_configs[encoder])
    if metric:
        config['max_depth'] = max_depth

    torch_device = resolve_torch_device(device)
    # The vendored Depth Anything V2 DINOv2 blocks prefer xFormers when it is
    # importable, but the available xFormers kernels do not support every
    # server GPU/dtype combination, including RTX 5090 sm_120 with float32.
    # Use plain PyTorch attention for reproducible CPU and Slurm inference.
    da_attention.XFORMERS_AVAILABLE = False
    da_block.XFORMERS_AVAILABLE = False

    model = DepthAnythingV2(**config)
    model.load_state_dict(torch.load(str(checkpoint), map_location='cpu'))
    model = model.to(torch_device).eval()

    # ``image`` was EXIF-normalized and optionally downscaled by ``read_rgb``.
    # Depth Anything's infer_image receives BGR when loaded via OpenCV, so
    # preserve that channel convention without reopening the raw JPEG.
    raw = np.asarray(image.convert('RGB'))[:, :, ::-1].copy()

    with torch.inference_mode():
        depth = model.infer_image(raw, input_size)
    return depth.astype(np.float32)


def load_vggt_model(
    repo: str | Path,
    model_id: str,
    checkpoint: str | Path | None,
    device: str,
):
    repo = Path(repo).resolve()
    if not repo.exists():
        raise FileNotFoundError(f'VGGT repo not found: {repo}')
    if str(repo) not in sys.path:
        sys.path.insert(0, str(repo))

    import torch
    from vggt.models.vggt import VGGT

    torch_device = resolve_torch_device(device)
    # We only need camera/depth/point for accessibility geometry. Disabling
    # track avoids loading the point-tracking branch and lowers memory use.
    model = VGGT(enable_track=False)
    if checkpoint:
        checkpoint = Path(checkpoint).resolve()
        if not checkpoint.exists():
            raise FileNotFoundError(f'VGGT checkpoint not found: {checkpoint}')
        state = torch.load(str(checkpoint), map_location='cpu')
        if isinstance(state, dict) and 'model' in state:
            state = state['model']
    else:
        local_model = Path(model_id)
        if local_model.exists():
            model_path = local_model / 'model.pt' if local_model.is_dir() else local_model
            if not model_path.exists():
                raise FileNotFoundError(f'VGGT local model checkpoint not found: {model_path}')
        else:
            from huggingface_hub import hf_hub_download

            model_path = Path(hf_hub_download(repo_id=model_id, filename='model.pt'))
        state = torch.load(str(model_path), map_location='cpu')
        if isinstance(state, dict) and 'model' in state:
            state = state['model']
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        print(f'VGGT checkpoint missing keys: {len(missing)}', file=sys.stderr)
    if unexpected:
        print(f'VGGT checkpoint unexpected keys: {len(unexpected)}', file=sys.stderr)
    model = model.to(torch_device).eval()
    return model, torch_device


def vggt_predict_depth(
    image: Image.Image,
    model,
    device: str,
    input_size: int,
    near: float,
    far: float,
    invert: bool,
    keep_raw_depth: bool,
) -> dict[str, np.ndarray]:
    import torch

    tensor, crop = preprocess_square_tensor(image, input_size)
    tensor = tensor.to(device)
    if device.startswith('cuda'):
        major = torch.cuda.get_device_capability()[0]
        dtype = torch.bfloat16 if major >= 8 else torch.float16
        autocast = torch.cuda.amp.autocast(dtype=dtype)
    else:
        autocast = torch.autocast(device_type='cpu', enabled=False)

    with torch.inference_mode():
        with autocast:
            predictions = model(tensor)

    raw_depth = predictions['depth'][0, 0, ..., 0].detach().float().cpu().numpy()
    depth_conf = predictions['depth_conf'][0, 0].detach().float().cpu().numpy()
    point_map = predictions['world_points'][0, 0].detach().float().cpu().numpy()
    point_conf = predictions['world_points_conf'][0, 0].detach().float().cpu().numpy()

    raw_depth = crop_and_resize_prediction(raw_depth, crop, image.size)
    depth_conf = crop_and_resize_prediction(depth_conf, crop, image.size)
    point_map = crop_and_resize_prediction(point_map, crop, image.size)
    point_conf = crop_and_resize_prediction(point_conf, crop, image.size)
    depth = raw_depth if keep_raw_depth else normalize_depth(raw_depth, near, far, invert)
    return {
        'depth': depth.astype(np.float32),
        'raw_depth': raw_depth.astype(np.float32),
        'depth_conf': depth_conf.astype(np.float32),
        'world_points': point_map.astype(np.float32),
        'world_points_conf': point_conf.astype(np.float32),
    }


def save_depth_vis(path: str | Path, depth: np.ndarray) -> None:
    valid = np.isfinite(depth)
    values = depth[valid]
    if values.size == 0:
        vis = np.zeros(depth.shape, dtype=np.uint8)
    else:
        lo, hi = np.percentile(values, [2, 98])
        if hi <= lo:
            hi = lo + 1.0
        vis = np.clip((depth - lo) / (hi - lo), 0.0, 1.0)
        vis = (vis * 255).astype(np.uint8)
    Image.fromarray(vis, mode='L').save(path)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description='Generate relative depth for accessibility geometry completion.')
    parser.add_argument('--image', required=True, help='Input RGB image.')
    parser.add_argument('--output-depth', required=True, help='Output .npy depth path.')
    parser.add_argument('--output-vis', default=None, help='Optional grayscale depth visualization path.')
    parser.add_argument('--manifest', default=None, help='Optional JSON manifest path.')
    parser.add_argument('--engine', choices=['heuristic', 'transformers', 'depth_anything_v2', 'vggt'], default='heuristic')
    parser.add_argument('--model', default=DEFAULT_DEPTH_MODEL, help='Transformers depth-estimation model id or local path.')
    parser.add_argument('--device', default='auto', help='auto, cpu, cuda, cuda:0, or a transformers device string.')
    parser.add_argument('--depth-anything-repo', default='../diffusion-vas/models/Depth_Anything_V2', help='Local Depth Anything V2 source repo.')
    parser.add_argument('--checkpoint', default=None, help='Depth Anything V2 .pth checkpoint for --engine depth_anything_v2.')
    parser.add_argument('--encoder', choices=['vits', 'vitb', 'vitl', 'vitg'], default='vitl')
    parser.add_argument('--metric-depth', action='store_true', help='Treat the Depth Anything V2 checkpoint as a metric-depth model and do not near/far normalize.')
    parser.add_argument('--max-depth', type=float, default=20.0, help='Metric Depth Anything max depth in meters.')
    parser.add_argument('--input-size', type=int, default=518, help='Depth Anything V2 inference input size.')
    parser.add_argument('--near', type=float, default=1.0, help='Depth value assigned to the near end after normalization.')
    parser.add_argument('--far', type=float, default=6.0, help='Depth value assigned to the far end after normalization.')
    parser.add_argument('--invert-depth', action='store_true', help='Invert predicted relative depth before near/far normalization.')
    parser.add_argument('--max-size', type=int, default=1280, help='Resize longest side before inference. Use 0 to keep original size.')
    parser.add_argument('--vggt-repo', default='vggt', help='Local facebookresearch/VGGT checkout.')
    parser.add_argument('--vggt-model', default=DEFAULT_VGGT_MODEL, help='Hugging Face model id or local VGGT model directory.')
    parser.add_argument('--vggt-checkpoint', default=None, help='Optional local VGGT model.pt checkpoint.')
    parser.add_argument('--vggt-input-size', type=int, default=518, help='Square VGGT inference resolution.')
    parser.add_argument('--vggt-keep-raw-depth', action='store_true', help='Do not near/far normalize VGGT depth before saving --output-depth.')
    parser.add_argument('--output-conf', default=None, help='Optional output .npy confidence map path.')
    parser.add_argument('--output-conf-vis', default=None, help='Optional output confidence visualization path.')
    parser.add_argument('--output-raw-depth', default=None, help='Optional output raw VGGT depth .npy path.')
    parser.add_argument('--output-world-points', default=None, help='Optional output VGGT world point map .npy path.')
    parser.add_argument('--output-point-conf', default=None, help='Optional output VGGT world point confidence .npy path.')
    parser.add_argument('--output-point-cloud', default=None, help='Optional output VGGT point cloud .ply path.')
    parser.add_argument('--point-conf-threshold', type=float, default=1.0, help='VGGT point confidence threshold for --output-point-cloud.')
    parser.add_argument('--max-point-cloud-points', type=int, default=120000, help='Maximum VGGT point-cloud vertices to export.')
    return parser


def main() -> None:
    args = build_parser().parse_args()
    max_size = None if args.max_size == 0 else args.max_size
    image = read_rgb(args.image, max_size=max_size)

    if args.engine == 'heuristic':
        raw_depth = heuristic_depth(image, args.near, args.far)
        depth = raw_depth
    elif args.engine == 'transformers':
        raw_depth = transformers_depth(image, args.model, args.device)
        depth = normalize_depth(raw_depth, args.near, args.far, args.invert_depth)
    elif args.engine == 'depth_anything_v2':
        raw_depth = depth_anything_v2_depth(
            image,
            args.depth_anything_repo,
            args.checkpoint,
            args.encoder,
            args.device,
            args.input_size,
            args.metric_depth,
            args.max_depth,
        )
        depth = raw_depth if args.metric_depth else normalize_depth(raw_depth, args.near, args.far, args.invert_depth)
    else:
        model, torch_device = load_vggt_model(
            args.vggt_repo,
            args.vggt_model,
            args.vggt_checkpoint,
            args.device,
        )
        prediction = vggt_predict_depth(
            image,
            model,
            torch_device,
            args.vggt_input_size,
            args.near,
            args.far,
            args.invert_depth,
            args.vggt_keep_raw_depth,
        )
        depth = prediction['depth']
        raw_depth = prediction['raw_depth']
        if args.output_conf:
            Path(args.output_conf).parent.mkdir(parents=True, exist_ok=True)
            np.save(Path(args.output_conf), prediction['depth_conf'].astype(np.float32))
        if args.output_conf_vis:
            Path(args.output_conf_vis).parent.mkdir(parents=True, exist_ok=True)
            colorize_confidence(prediction['depth_conf']).save(args.output_conf_vis)
        if args.output_raw_depth:
            Path(args.output_raw_depth).parent.mkdir(parents=True, exist_ok=True)
            np.save(Path(args.output_raw_depth), prediction['raw_depth'].astype(np.float32))
        if args.output_world_points:
            Path(args.output_world_points).parent.mkdir(parents=True, exist_ok=True)
            np.save(Path(args.output_world_points), prediction['world_points'].astype(np.float32))
        if args.output_point_conf:
            Path(args.output_point_conf).parent.mkdir(parents=True, exist_ok=True)
            np.save(Path(args.output_point_conf), prediction['world_points_conf'].astype(np.float32))
        if args.output_point_cloud:
            rgb = np.asarray(image.convert('RGB'), dtype=np.uint8)
            point_count = write_point_cloud_ply(
                Path(args.output_point_cloud),
                prediction['world_points'],
                rgb,
                prediction['world_points_conf'],
                args.point_conf_threshold,
                args.max_point_cloud_points,
            )
            print(f'Wrote VGGT point cloud to {args.output_point_cloud} ({point_count} points)')

    output_depth = Path(args.output_depth)
    output_depth.parent.mkdir(parents=True, exist_ok=True)
    np.save(output_depth, depth.astype(np.float32))

    output_vis = Path(args.output_vis) if args.output_vis else output_depth.with_suffix('.png')
    output_vis.parent.mkdir(parents=True, exist_ok=True)
    save_depth_vis(output_vis, depth)

    note = 'Metric monocular depth estimate; calibrate intrinsics/scale before path-planning use.'
    if args.engine == 'vggt':
        note = 'VGGT single-view/few-view geometry estimate; scale is not calibrated metric ground truth.'
    elif args.engine != 'depth_anything_v2' or not args.metric_depth:
        note = 'Relative monocular depth; use calibrated metric depth for path-planning truth.'

    model_identity = depth_model_identity(
        engine=args.engine,
        model=args.model,
        checkpoint=args.checkpoint,
        encoder=args.encoder,
        metric_depth=args.metric_depth,
        vggt_model=args.vggt_model,
    )
    manifest = {
        'image': args.image,
        'raster_orientation_policy': 'RGB is decoded with PIL ImageOps.exif_transpose before depth inference.',
        'display_raster_size': {'width': image.width, 'height': image.height},
        'engine': args.engine,
        'model': model_identity['model_id'],
        'model_identity': model_identity,
        'depth_anything_repo': args.depth_anything_repo if args.engine == 'depth_anything_v2' else None,
        'checkpoint': args.checkpoint if args.engine == 'depth_anything_v2' else None,
        'encoder': args.encoder if args.engine == 'depth_anything_v2' else None,
        'vggt_repo': args.vggt_repo if args.engine == 'vggt' else None,
        'vggt_model': args.vggt_model if args.engine == 'vggt' else None,
        'vggt_checkpoint': args.vggt_checkpoint if args.engine == 'vggt' else None,
        'vggt_input_size': args.vggt_input_size if args.engine == 'vggt' else None,
        'vggt_keep_raw_depth': args.vggt_keep_raw_depth if args.engine == 'vggt' else None,
        'metric_depth': args.metric_depth if args.engine == 'depth_anything_v2' else None,
        'max_depth': args.max_depth if args.metric_depth else None,
        'device': args.device if args.engine in {'transformers', 'depth_anything_v2', 'vggt'} else None,
        'output_depth': str(output_depth),
        'output_vis': str(output_vis),
        'output_conf': args.output_conf if args.engine == 'vggt' else None,
        'output_raw_depth': args.output_raw_depth if args.engine == 'vggt' else None,
        'output_world_points': args.output_world_points if args.engine == 'vggt' else None,
        'output_point_cloud': args.output_point_cloud if args.engine == 'vggt' else None,
        'near': args.near,
        'far': args.far,
        'invert_depth': args.invert_depth,
        'shape': list(depth.shape),
        'note': note,
    }
    manifest_path = Path(args.manifest) if args.manifest else output_depth.with_name('depth_manifest.json')
    manifest_path.write_text(json.dumps(manifest, indent=2), encoding='utf-8')
    print(f'Wrote depth to {output_depth}')
    print(f'Wrote depth visualization to {output_vis}')


if __name__ == '__main__':
    main()