File size: 763 Bytes
2f28ec8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from __future__ import annotations

from typing import Any

import torch


def scene_encode_kwargs(batch: dict[str, Any], device: torch.device, model_uses_rgbd: bool) -> dict[str, Any]:
    kwargs: dict[str, Any] = {}
    if not model_uses_rgbd:
        return kwargs
    kwargs["depths"] = batch.get("depths").to(device) if batch.get("depths") is not None else None
    kwargs["depth_valid"] = batch.get("depth_valid").to(device) if batch.get("depth_valid") is not None else None
    kwargs["camera_intrinsics"] = batch.get("camera_intrinsics").to(device) if batch.get("camera_intrinsics") is not None else None
    kwargs["camera_extrinsics"] = batch.get("camera_extrinsics").to(device) if batch.get("camera_extrinsics") is not None else None
    return kwargs