|
|
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) |
|
|
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 |
|
|
|
