File size: 1,039 Bytes
2568013 |
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 |
from einops import rearrange
from .projection import sample_image_grid, get_local_rays
from ..misc.sht import rsh_cart_2, rsh_cart_4, rsh_cart_6, rsh_cart_8
def get_intrinsic_embedding(context, degree=0, downsample=1, merge_hw=False):
assert degree in [0, 2, 4, 8]
b, v, _, h, w = context["image"].shape
device = context["image"].device
tgt_h, tgt_w = h // downsample, w // downsample
xy_ray, _ = sample_image_grid((tgt_h, tgt_w), device)
xy_ray = xy_ray[None, None, ...].expand(b, v, -1, -1, -1) # [b, v, h, w, 2]
directions = get_local_rays(xy_ray, rearrange(context["intrinsics"], "b v i j -> b v () () i j"),)
if degree == 2:
directions = rsh_cart_2(directions)
elif degree == 4:
directions = rsh_cart_4(directions)
elif degree == 8:
directions = rsh_cart_8(directions)
if merge_hw:
directions = rearrange(directions, "b v h w d -> b v (h w) d")
else:
directions = rearrange(directions, "b v h w d -> b v d h w")
return directions
|