Function compute_uv_space_data

data_loader/mesh_utils.py:186–242

Compute a UV-space barycentric map where each texel contains barycentric coordinates for the closest point on a UV triangle. Args: vt: torch.Tensor Texture coordinates. Shape = [n_texcoords, 2] vti: torch.Tensor Face texture coordinate indices. Shape = [n_face

(
    vt: torch.Tensor,
    vti: torch.Tensor,
    uv_shape: Union[Tuple[int, int], int],
    flip_uv: bool = True,
)

#incurs a copy to cpu
#generates uv space map where each pixel has the index of the triangle, the 3 vertex indices of the triangle and the barycentric weights
def compute_uv_space_data(
    vt: torch.Tensor,
    vti: torch.Tensor,
    uv_shape: Union[Tuple[int, int], int],
    flip_uv: bool = True,
):
    """Compute a UV-space barycentric map where each texel contains barycentric
    coordinates for the closest point on a UV triangle.
    Args:
        vt: torch.Tensor
        Texture coordinates. Shape = [n_texcoords, 2]
        vti: torch.Tensor
        Face texture coordinate indices. Shape = [n_faces, 3]
        uv_shape: Tuple[int, int] or int
        Shape of the texture map. (HxW)
        flip_uv: bool
        Whether or not to flip UV coordinates along the V axis (OpenGL -> numpy/pytorch convention).
    Returns:
        torch.Tensor: index_img: Face index image, shape [uv_shape[0], uv_shape[1]]
        torch.Tensor: Barycentric coordinate map, shape [uv_shape[0], uv_shape[1], 3]Â
    """

    if isinstance(uv_shape, int):
        uv_shape = (uv_shape, uv_shape)

    if flip_uv:
        # Flip here because texture coordinates in some of our topo files are
        # stored in OpenGL convention with Y=0 on the bottom of the texture
        # unlike numpy/torch arrays/tensors.
        vt = vt.clone()
        vt[:, 1] = 1 - vt[:, 1]

    # Texel to UV mapping (as per OpenGL linear filtering)
    # https://www.khronos.org/registry/OpenGL/specs/gl/glspec46.core.pdf
    # Sect. 8.14, page 261
    # uv=(0.5,0.5)/w is at the center of texel [0,0]
    # uv=(w-0.5, w-0.5)/w is the center of texel [w-1,w-1]
    # texel = floor(u*w - 0.5)
    # u = (texel+0.5)/w

    uv_grid = torch.meshgrid(
        torch.linspace(0.5, uv_shape[0] - 1 + 0.5, uv_shape[0]) / uv_shape[0],
        torch.linspace(0.5, uv_shape[1] - 1 + 0.5, uv_shape[1]) / uv_shape[1],
        indexing="ij")  # HxW, v,u
    uv_grid = torch.stack(uv_grid[::-1], dim=2)  # HxW, u, v
    uv = uv_grid.reshape(-1, 2).data.to("cpu").numpy()
    vth = np.hstack((vt, vt[:, 0:1] * 0 + 1))
    uvh = np.hstack((uv, uv[:, 0:1] * 0 + 1))

    # print("vth",vth)

    closest_points, barys, vertex_idxs, face_idxs = closest_point_barycentrics(uvh, vth, vti.numpy())

    index_img = torch.from_numpy(face_idxs.reshape(uv_shape[0], uv_shape[1])).long()
    vert_idxs_img = torch.from_numpy(vertex_idxs.reshape(uv_shape[0], uv_shape[1],3)).long()
    bary_img = torch.from_numpy(barys.reshape(uv_shape[0], uv_shape[1], 3)).float()
    return index_img, vert_idxs_img, bary_img