Spaces:
Running on Zero
Running on Zero
File size: 3,003 Bytes
87608ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Timed raw PXDepth forward path used by evaluation tools.
This module applies the selected fixed-size or equal-area preprocessing,
measures only model-forward latency, and restores predictions to input
resolution. It intentionally returns only genuine network outputs rather than
performing metric alignment or estimating camera intrinsics.
"""
import time
from typing import Dict, Optional, Tuple
import torch
from ..model import PXDepth
from .resize import resize_image, resize_map
def synchronize(device: torch.device) -> None:
"""Synchronize pending CUDA work before or after timing model forward.
Args:
device: Model device. CPU and other devices require no action.
Returns:
``None``.
"""
if device.type == "cuda":
torch.cuda.synchronize(device)
@torch.inference_mode()
def predict_raw(
model: PXDepth,
image: torch.Tensor,
input_size: Optional[Tuple[int, int]] = (1022, 770),
resize_by_area: bool = True,
use_fp16: bool = False,
use_fp32: bool = False,
) -> Dict[str, torch.Tensor]:
"""Run raw model forward and return outputs at the original image size.
Only ``model.forward`` is included in ``inference_time``. Resizing,
synchronization overhead, and output packaging are excluded. No GT depth
alignment or camera-intrinsics prediction is performed.
Args:
model: Evaluation-mode :class:`PXDepth` model.
image: RGB tensor ``[3,H,W]`` or ``[B,3,H,W]`` in ``[0,1]``.
input_size: Exact/reference tuple ``(width,height)``. Defaults to
``(1022,770)``.
resize_by_area: Preserve aspect ratio at ``input_size`` area. Enabled
by default.
use_fp16: Use FP16 for attention-heavy model regions.
use_fp32: Force full-precision model execution.
Returns:
Dictionary with raw normalized log-depth ``depth_affine_invariant``
``[B,H,W]``, ``depth_affine_space='log'``, boolean ``mask`` ``[B,H,W]``,
and scalar forward time. The leading batch dimension is removed for
unbatched input.
"""
image, original_size = resize_image(image, input_size, resize_by_area, model.patch_size)
batched = image.ndim == 4
model_input = image if batched else image.unsqueeze(0)
model_input = model_input.to(device=model.device, dtype=torch.float32)
synchronize(model.device)
start = time.perf_counter()
output = model.forward(model_input, use_fp16=use_fp16, use_fp32=use_fp32)
synchronize(model.device)
elapsed = time.perf_counter() - start
depth = resize_map(output["depth"], original_size)
mask = resize_map(output["mask"], original_size, is_mask=True)
pred = {
"depth_affine_invariant": depth,
"depth_affine_space": "log",
"mask": mask,
"inference_time": elapsed,
}
if not batched:
pred = {key: value[0] if isinstance(value, torch.Tensor) and value.ndim > 0 else value for key, value in pred.items()}
return pred
|