File size: 499 Bytes
41ff959
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from typing import TypedDict

from jaxtyping import Float
from torch import Tensor


class BatchedViews(TypedDict, total=False):
    """Single- or multi-view tensors consumed by the inference encoders."""

    extrinsics: Float[Tensor, "batch view 4 4"]
    intrinsics: Float[Tensor, "batch view 3 3"]
    image: Float[Tensor, "batch view channel height width"]
    prompt_disparity: Float[Tensor, "batch view channel height width"]
    prompt_mask: Float[Tensor, "batch view channel height width"]