textures.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-unsafe

import itertools
import warnings
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union

import torch
import torch.nn.functional as F
from pytorch3d.ops import interpolate_face_attributes
from pytorch3d.structures.utils import list_to_packed, list_to_padded, padded_to_list
from torch.nn.functional import interpolate

from .utils import pack_unique_rectangles, PackedRectangle, Rectangle


# This file contains classes and helper functions for texturing.
# There are three types of textures: TexturesVertex, TexturesAtlas
# and TexturesUV which inherit from a base textures class TexturesBase.
#
# Each texture class has a method 'sample_textures' to sample a
# value given barycentric coordinates.
#
# All the textures accept either list or padded inputs. The values
# are stored as either per face values (TexturesAtlas, TexturesUV),
# or per face vertex features (TexturesVertex).


def _list_to_padded_wrapper(
    x: List[torch.Tensor],
    pad_size: Union[list, tuple, None] = None,
    pad_value: float = 0.0,
) -> torch.Tensor:
    r"""
    This is a wrapper function for
    pytorch3d.structures.utils.list_to_padded function which only accepts
    3-dimensional inputs.

    For this use case, the input x is of shape (F, 3, ...) where only F
    is different for each element in the list

    Transforms a list of N tensors each of shape (Mi, ...) into a single tensor
    of shape (N, pad_size, ...), or (N, max(Mi), ...)
    if pad_size is None.

    Args:
      x: list of Tensors
      pad_size: int specifying the size of the first dimension
        of the padded tensor
      pad_value: float value to be used to fill the padded tensor

    Returns:
      x_padded: tensor consisting of padded input tensors
    """
    N = len(x)
    dims = x[0].ndim
    reshape_dims = x[0].shape[1:]
    D = torch.prod(torch.tensor(reshape_dims)).item()
    x_reshaped = []
    for y in x:
        if y.ndim != dims and y.shape[1:] != reshape_dims:
            msg = (
                "list_to_padded requires tensors to have the same number of dimensions"
            )
            raise ValueError(msg)
        # pyre-fixme[6]: For 2nd param expected `int` but got `Union[bool, float, int]`.
        x_reshaped.append(y.reshape(-1, D))
    x_padded = list_to_padded(x_reshaped, pad_size=pad_size, pad_value=pad_value)
    # pyre-fixme[58]: `+` is not supported for operand types `Tuple[int, int]` and
    #  `Size`.
    return x_padded.reshape((N, -1) + reshape_dims)


def _padded_to_list_wrapper(
    x: torch.Tensor, split_size: Union[list, tuple, None] = None
) -> List[torch.Tensor]:
    r"""
    This is a wrapper function for pytorch3d.structures.utils.padded_to_list
    which only accepts 3-dimensional inputs.

    For this use case, the input x is of shape (N, F, ...) where F
    is the number of faces which is different for each tensor in the batch.

    This function transforms a padded tensor of shape (N, M, ...) into a
    list of N tensors of shape (Mi, ...) where (Mi) is specified in
    split_size(i), or of shape (M,) if split_size is None.

    Args:
      x: padded Tensor
      split_size: list of ints defining the number of items for each tensor
        in the output list.

    Returns:
      x_list: a list of tensors
    """
    N, M = x.shape[:2]
    reshape_dims = x.shape[2:]
    D = torch.prod(torch.tensor(reshape_dims)).item()
    # pyre-fixme[6]: For 3rd param expected `int` but got `Union[bool, float, int]`.
    x_reshaped = x.reshape(N, M, D)
    x_list = padded_to_list(x_reshaped, split_size=split_size)
    # pyre-fixme[58]: `+` is not supported for operand types `Tuple[typing.Any]` and
    #  `Size`.
    x_list = [xl.reshape((xl.shape[0],) + reshape_dims) for xl in x_list]
    return x_list


def _pad_texture_maps(
    images: Union[Tuple[torch.Tensor], List[torch.Tensor]], align_corners: bool
) -> torch.Tensor:
    """
    Pad all texture images so they have the same height and width.

    Args:
        images: list of N tensors of shape (H_i, W_i, C)
        align_corners: used for interpolation

    Returns:
        tex_maps: Tensor of shape (N, max_H, max_W, C)
    """
    tex_maps = []
    max_H = 0
    max_W = 0
    for im in images:
        h, w, _C = im.shape
        if h > max_H:
            max_H = h
        if w > max_W:
            max_W = w
        tex_maps.append(im)
    max_shape = (max_H, max_W)

    for i, image in enumerate(tex_maps):
        if image.shape[:2] != max_shape:
            image_BCHW = image.permute(2, 0, 1)[None]
            new_image_BCHW = interpolate(
                image_BCHW,
                size=max_shape,
                mode="bilinear",
                align_corners=align_corners,
            )
            tex_maps[i] = new_image_BCHW[0].permute(1, 2, 0)
    tex_maps = torch.stack(tex_maps, dim=0)  # (num_tex_maps, max_H, max_W, C)
    return tex_maps


def _pad_texture_multiple_maps(
    multiple_texture_maps: Union[Tuple[torch.Tensor], List[torch.Tensor]],
    align_corners: bool,
) -> torch.Tensor:
    """
    Pad all texture images so they have the same height and width.

    Args:
        images: list of N tensors of shape (M_i, H_i, W_i, C)
        M_i : Number of texture maps:w

        align_corners: used for interpolation

    Returns:
        tex_maps: Tensor of shape (N, max_M, max_H, max_W, C)
    """
    tex_maps = []
    max_M = 0
    max_H = 0
    max_W = 0
    C = 0
    for im in multiple_texture_maps:
        m, h, w, C = im.shape
        if m > max_M:
            max_M = m
        if h > max_H:
            max_H = h
        if w > max_W:
            max_W = w
        tex_maps.append(im)
    max_shape = (max_M, max_H, max_W, C)
    max_im_shape = (max_H, max_W)
    for i, tms in enumerate(tex_maps):
        new_tex_maps = torch.zeros(max_shape)
        for j in range(tms.shape[0]):
            im = tms[j]
            if im.shape[:2] != max_im_shape:
                image_BCHW = im.permute(2, 0, 1)[None]
                new_image_BCHW = interpolate(
                    image_BCHW,
                    size=max_im_shape,
                    mode="bilinear",
                    align_corners=align_corners,
                )
                new_tex_maps[j] = new_image_BCHW[0].permute(1, 2, 0)
            else:
                new_tex_maps[j] = im
        tex_maps[i] = new_tex_maps
    tex_maps = torch.stack(tex_maps, dim=0)  # (num_tex_maps, max_H, max_W, C)
    return tex_maps


# A base class for defining a batch of textures
# with helper methods.
# This is also useful to have so that inside `Meshes`
# we can allow the input textures to be any texture
# type which is an instance of the base class.
class TexturesBase:
    def isempty(self):
        if self._N is not None and self.valid is not None:
            return self._N == 0 or self.valid.eq(False).all()
        return False

    def to(self, device):
        for k in dir(self):
            v = getattr(self, k)
            if isinstance(v, (list, tuple)) and all(
                torch.is_tensor(elem) for elem in v
            ):
                v = [elem.to(device) for elem in v]
                setattr(self, k, v)
            if torch.is_tensor(v) and v.device != device:
                setattr(self, k, v.to(device))
        self.device = device
        return self

    def _extend(self, N: int, props: List[str]) -> Dict[str, Union[torch.Tensor, List]]:
        """
        Create a dict with the specified properties
        repeated N times per batch element.

        Args:
            N: number of new copies of each texture
                in the batch.
            props: a List of strings which refer to either
                class attributes or class methods which
                return tensors or lists.

        Returns:
            Dict with the same keys as props. The values are the
            extended properties.
        """
        if not isinstance(N, int):
            raise ValueError("N must be an integer.")
        if N <= 0:
            raise ValueError("N must be > 0.")

        new_props = {}
        for p in props:
            t = getattr(self, p)
            if callable(t):
                t = t()  # class method
            if t is None:
                new_props[p] = None
            elif isinstance(t, list):
                if not all(isinstance(elem, (int, float)) for elem in t):
                    raise ValueError("Extend only supports lists of scalars")
                t = [[ti] * N for ti in t]
                new_props[p] = list(itertools.chain(*t))
            elif torch.is_tensor(t):
                new_props[p] = t.repeat_interleave(N, dim=0)
            else:
                raise ValueError(
                    f"Property {p} has unsupported type {type(t)}."
                    "Only tensors and lists are supported."
                )
        return new_props

    def _getitem(self, index: Union[int, slice], props: List[str]):
        """
        Helper function for __getitem__
        """
        new_props = {}
        if isinstance(index, (int, slice)):
            for p in props:
                t = getattr(self, p)
                if callable(t):
                    t = t()  # class method
                new_props[p] = t[index] if t is not None else None
        elif isinstance(index, list):
            index = torch.tensor(index)
        if isinstance(index, torch.Tensor):
            if index.dtype == torch.bool:
                index = index.nonzero()
                index = index.squeeze(1) if index.numel() > 0 else index
                index = index.tolist()
            for p in props:
                t = getattr(self, p)
                if callable(t):
                    t = t()  # class method
                new_props[p] = [t[i] for i in index] if t is not None else None
        return new_props

    def sample_textures(self) -> torch.Tensor:
        """
        Different texture classes sample textures in different ways
        e.g. for vertex textures, the values at each vertex
        are interpolated across the face using the barycentric
        coordinates.
        Each texture class should implement a sample_textures
        method to take the `fragments` from rasterization.
        Using `fragments.pix_to_face` and `fragments.bary_coords`
        this function should return the sampled texture values for
        each pixel in the output image.
        """
        raise NotImplementedError()

    def submeshes(
        self,
        vertex_ids_list: List[List[torch.LongTensor]],
        faces_ids_list: List[List[torch.LongTensor]],
    ) -> "TexturesBase":
        """
        Extract sub-textures used for submeshing.
        """
        raise NotImplementedError(f"{self.__class__} does not support submeshes")

    def faces_verts_textures_packed(self) -> torch.Tensor:
        """
        Returns the texture for each vertex for each face in the mesh.
        For N meshes, this function returns sum(Fi)x3xC where Fi is the
        number of faces in the i-th mesh and C is the dimensional of
        the feature (C = 3 for RGB textures).
        You can use the utils function in structures.utils to convert the
        packed representation to a list or padded.
        """
        raise NotImplementedError()

    def clone(self) -> "TexturesBase":
        """
        Each texture class should implement a method
        to clone all necessary internal tensors.
        """
        raise NotImplementedError()

    def detach(self) -> "TexturesBase":
        """
        Each texture class should implement a method
        to detach all necessary internal tensors.
        """
        raise NotImplementedError()

    def __getitem__(self, index) -> "TexturesBase":
        """
        Each texture class should implement a method
        to get the texture properties for the
        specified elements in the batch.
        The TexturesBase._getitem(i) method
        can be used as a helper function to retrieve the
        class attributes for item i. Then, a new
        instance of the child class can be created with
        the attributes.
        """
        raise NotImplementedError()


def Textures(
    maps: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
    faces_uvs: Optional[torch.Tensor] = None,
    verts_uvs: Optional[torch.Tensor] = None,
    verts_rgb: Optional[torch.Tensor] = None,
) -> TexturesBase:
    """
    Textures class has been DEPRECATED.
    Preserving Textures as a function for backwards compatibility.

    Args:
        maps: texture map per mesh. This can either be a list of maps
          [(H, W, C)] or a padded tensor of shape (N, H, W, C).
        faces_uvs: (N, F, 3) tensor giving the index into verts_uvs for each
            vertex in the face. Padding value is assumed to be -1.
        verts_uvs: (N, V, 2) tensor giving the uv coordinate per vertex.
        verts_rgb: (N, V, C) tensor giving the color per vertex. Padding
            value is assumed to be -1. (C=3 for RGB.)


    Returns:
        a Textures class which is an instance of TexturesBase e.g. TexturesUV,
        TexturesAtlas, TexturesVertex

    """

    warnings.warn(
        """Textures class is deprecated,
        use TexturesUV, TexturesAtlas, TexturesVertex instead.
        Textures class will be removed in future releases.""",
        PendingDeprecationWarning,
    )

    if faces_uvs is not None and verts_uvs is not None and maps is not None:
        return TexturesUV(maps=maps, faces_uvs=faces_uvs, verts_uvs=verts_uvs)

    if verts_rgb is not None:
        return TexturesVertex(verts_features=verts_rgb)

    raise ValueError(
        "Textures either requires all three of (faces uvs, verts uvs, maps) or verts rgb"
    )


class TexturesAtlas(TexturesBase):
    def __init__(self, atlas: Union[torch.Tensor, List[torch.Tensor]]) -> None:
        """
        A texture representation where each face has a square texture map.
        This is based on the implementation from SoftRasterizer [1].

        Args:
            atlas: (N, F, R, R, C) tensor giving the per face texture map.
                The atlas can be created during obj loading with the
                pytorch3d.io.load_obj function - in the input arguments
                set `create_texture_atlas=True`. The atlas will be
                returned in aux.texture_atlas.


        The padded and list representations of the textures are stored
        and the packed representations is computed on the fly and
        not cached.

        [1] Liu et al, 'Soft Rasterizer: A Differentiable Renderer for Image-based
            3D Reasoning', ICCV 2019
            See also https://github.com/ShichenLiu/SoftRas/issues/21
        """
        if isinstance(atlas, (list, tuple)):
            correct_format = all(
                (
                    torch.is_tensor(elem)
                    and elem.ndim == 4
                    and elem.shape[1] == elem.shape[2]
                    and elem.shape[1] == atlas[0].shape[1]
                )
                for elem in atlas
            )
            if not correct_format:
                msg = (
                    "Expected atlas to be a list of tensors of shape (F, R, R, C) "
                    "with the same value of R."
                )
                raise ValueError(msg)
            self._atlas_list = atlas
            self._atlas_padded = None
            self.device = torch.device("cpu")

            # These values may be overridden when textures is
            # passed into the Meshes constructor. For more details
            # refer to the __init__ of Meshes.
            self._N = len(atlas)
            self._num_faces_per_mesh = [len(a) for a in atlas]

            if self._N > 0:
                self.device = atlas[0].device

        elif torch.is_tensor(atlas):
            if atlas.ndim != 5:
                msg = "Expected atlas to be of shape (N, F, R, R, C); got %r"
                raise ValueError(msg % repr(atlas.ndim))
            self._atlas_padded = atlas
            self._atlas_list = None
            self.device = atlas.device

            # These values may be overridden when textures is
            # passed into the Meshes constructor. For more details
            # refer to the __init__ of Meshes.
            self._N = len(atlas)
            max_F = atlas.shape[1]
            self._num_faces_per_mesh = [max_F] * self._N
        else:
            raise ValueError("Expected atlas to be a tensor or list")

        # The num_faces_per_mesh, N and valid
        # are reset inside the Meshes object when textures is
        # passed into the Meshes constructor. For more details
        # refer to the __init__ of Meshes.
        self.valid = torch.ones((self._N,), dtype=torch.bool, device=self.device)

    def clone(self) -> "TexturesAtlas":
        tex = self.__class__(atlas=self.atlas_padded().clone())
        if self._atlas_list is not None:
            tex._atlas_list = [atlas.clone() for atlas in self._atlas_list]
        num_faces = (
            self._num_faces_per_mesh.clone()
            if torch.is_tensor(self._num_faces_per_mesh)
            else self._num_faces_per_mesh
        )
        tex.valid = self.valid.clone()
        tex._num_faces_per_mesh = num_faces
        return tex

    def detach(self) -> "TexturesAtlas":
        tex = self.__class__(atlas=self.atlas_padded().detach())
        if self._atlas_list is not None:
            tex._atlas_list = [atlas.detach() for atlas in self._atlas_list]
        num_faces = (
            self._num_faces_per_mesh.detach()
            if torch.is_tensor(self._num_faces_per_mesh)
            else self._num_faces_per_mesh
        )
        tex.valid = self.valid.detach()
        tex._num_faces_per_mesh = num_faces
        return tex

    def __getitem__(self, index) -> "TexturesAtlas":
        props = ["atlas_list", "_num_faces_per_mesh"]
        new_props = self._getitem(index, props=props)
        atlas = new_props["atlas_list"]
        if isinstance(atlas, list):
            # multiple batch elements
            new_tex = self.__class__(atlas=atlas)
        elif torch.is_tensor(atlas):
            # single element
            new_tex = self.__class__(atlas=[atlas])
        else:
            raise ValueError("Not all values are provided in the correct format")
        new_tex._num_faces_per_mesh = new_props["_num_faces_per_mesh"]
        return new_tex

    def atlas_padded(self) -> torch.Tensor:
        if self._atlas_padded is None:
            if self.isempty():
                self._atlas_padded = torch.zeros(
                    (self._N, 0, 0, 0, 3), dtype=torch.float32, device=self.device
                )
            else:
                self._atlas_padded = _list_to_padded_wrapper(
                    self._atlas_list, pad_value=0.0
                )
        return self._atlas_padded

    def atlas_list(self) -> List[torch.Tensor]:
        if self._atlas_list is None:
            if self.isempty():
                self._atlas_padded = [
                    torch.empty((0, 0, 0, 3), dtype=torch.float32, device=self.device)
                ] * self._N
            self._atlas_list = _padded_to_list_wrapper(
                self._atlas_padded, split_size=self._num_faces_per_mesh
            )
        return self._atlas_list

    def atlas_packed(self) -> torch.Tensor:
        if self.isempty():
            return torch.zeros(
                (self._N, 0, 0, 3), dtype=torch.float32, device=self.device
            )
        atlas_list = self.atlas_list()
        return list_to_packed(atlas_list)[0]

    def extend(self, N: int) -> "TexturesAtlas":
        new_props = self._extend(N, ["atlas_padded", "_num_faces_per_mesh"])
        new_tex = self.__class__(atlas=new_props["atlas_padded"])
        new_tex._num_faces_per_mesh = new_props["_num_faces_per_mesh"]
        return new_tex

    # pyre-fixme[14]: `sample_textures` overrides method defined in `TexturesBase`
    #  inconsistently.
    def sample_textures(self, fragments, **kwargs) -> torch.Tensor:
        """
        This is similar to a nearest neighbor sampling and involves a
        discretization step. The barycentric coordinates from
        rasterization are used to find the nearest grid cell in the texture
        atlas and the RGB is returned as the color.
        This means that this step is differentiable with respect to the RGB
        values of the texture atlas but not differentiable with respect to the
        barycentric coordinates.

        TODO: Add a different sampling mode which interpolates the barycentric
        coordinates to sample the texture and will be differentiable w.r.t
        the barycentric coordinates.

        Args:
            fragments:
                The outputs of rasterization. From this we use

                - pix_to_face: LongTensor of shape (N, H, W, K) specifying the indices
                of the faces (in the packed representation) which
                overlap each pixel in the image.
                - barycentric_coords: FloatTensor of shape (N, H, W, K, 3) specifying
                the barycentric coordinates of each pixel
                relative to the faces (in the packed
                representation) which overlap the pixel.

        Returns:
            texels: (N, H, W, K, C)
        """
        N, H, W, K = fragments.pix_to_face.shape
        atlas_packed = self.atlas_packed()
        R = atlas_packed.shape[1]
        bary = fragments.bary_coords
        pix_to_face = fragments.pix_to_face

        bary_w01 = bary[..., :2]
        # pyre-fixme[16]: `bool` has no attribute `__getitem__`.
        mask = (pix_to_face < 0)[..., None]
        bary_w01 = torch.where(mask, torch.zeros_like(bary_w01), bary_w01)
        # If barycentric coordinates are > 1.0 (in the case of
        # blur_radius > 0.0), wxy might be > R. We need to clamp this
        # index to R-1 to index into the texture atlas.
        w_xy = (bary_w01 * R).to(torch.int64).clamp(max=R - 1)  # (N, H, W, K, 2)

        below_diag = (
            bary_w01.sum(dim=-1) * R - w_xy.float().sum(dim=-1)
        ) <= 1.0  # (N, H, W, K)
        w_x, w_y = w_xy.unbind(-1)
        w_x = torch.where(below_diag, w_x, (R - 1 - w_x))
        w_y = torch.where(below_diag, w_y, (R - 1 - w_y))

        texels = atlas_packed[pix_to_face, w_y, w_x]
        texels = texels * (pix_to_face >= 0)[..., None].float()

        return texels

    def submeshes(
        self,
        vertex_ids_list: List[List[torch.LongTensor]],
        faces_ids_list: List[List[torch.LongTensor]],
    ) -> "TexturesAtlas":
        """
        Extract a sub-texture for use in a submesh.

        If the meshes batch corresponding to this TextureAtlas contains
        `n = len(faces_ids_list)` meshes, then self.atlas_list()
        will be of length n. After submeshing, we obtain a batch of
        `k = sum(len(v) for v in atlas_list` submeshes (see Meshes.submeshes). This
        function creates a corresponding TexturesAtlas object with `atlas_list`
        of length `k`.
        """
        if len(faces_ids_list) != len(self.atlas_list()):
            raise IndexError(
                "faces_ids_list must be of " "the same length as atlas_list."
            )

        sub_features = []
        for atlas, faces_ids in zip(self.atlas_list(), faces_ids_list):
            for faces_ids_submesh in faces_ids:
                sub_features.append(atlas[faces_ids_submesh])

        return self.__class__(sub_features)

    def faces_verts_textures_packed(self) -> torch.Tensor:
        """
        Samples texture from each vertex for each face in the mesh.
        For N meshes with {Fi} number of faces, it returns a
        tensor of shape sum(Fi)x3xC (C = 3 for RGB).
        You can use the utils function in structures.utils to convert the
        packed representation to a list or padded.
        """
        atlas_packed = self.atlas_packed()
        # assume each face consists of (v0, v1, v2).
        # to sample from the atlas we only need the first two barycentric coordinates.
        # for details on how this texture sample works refer to the sample_textures function.
        t0 = atlas_packed[:, 0, -1]  # corresponding to v0  with bary = (1, 0)
        t1 = atlas_packed[:, -1, 0]  # corresponding to v1 with bary = (0, 1)
        t2 = atlas_packed[:, 0, 0]  # corresponding to v2 with bary = (0, 0)
        return torch.stack((t0, t1, t2), dim=1)

    def join_batch(self, textures: List["TexturesAtlas"]) -> "TexturesAtlas":
        """
        Join the list of textures given by `textures` to
        self to create a batch of textures. Return a new
        TexturesAtlas object with the combined textures.

        Args:
            textures: List of TexturesAtlas objects

        Returns:
            new_tex: TexturesAtlas object with the combined
            textures from self and the list `textures`.
        """
        tex_types_same = all(isinstance(tex, TexturesAtlas) for tex in textures)
        if not tex_types_same:
            raise ValueError("All textures must be of type TexturesAtlas.")

        atlas_list = []
        atlas_list += self.atlas_list()
        num_faces_per_mesh = self._num_faces_per_mesh.copy()
        for tex in textures:
            atlas_list += tex.atlas_list()
            num_faces_per_mesh += tex._num_faces_per_mesh
        new_tex = self.__class__(atlas=atlas_list)
        new_tex._num_faces_per_mesh = num_faces_per_mesh
        return new_tex

    def join_scene(self) -> "TexturesAtlas":
        """
        Return a new TexturesAtlas amalgamating the batch.
        """
        return self.__class__(atlas=[torch.cat(self.atlas_list())])

    def check_shapes(
        self, batch_size: int, max_num_verts: int, max_num_faces: int
    ) -> bool:
        """
        Check if the dimensions of the atlas match that of the mesh faces
        """
        # (N, F) should be the same
        return self.atlas_padded().shape[0:2] == (batch_size, max_num_faces)


class TexturesUV(TexturesBase):
    def __init__(
        self,
        maps: Union[torch.Tensor, List[torch.Tensor]],
        faces_uvs: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]],
        verts_uvs: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]],
        *,
        maps_ids: Optional[
            Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]
        ] = None,
        padding_mode: str = "border",
        align_corners: bool = True,
        sampling_mode: str = "bilinear",
    ) -> None:
        """
        Textures are represented as a per mesh texture map and uv coordinates for each
        vertex in each face. NOTE: this class only supports one texture map per mesh.

        Args:
            maps: Either (1) a texture map per mesh. This can either be a list of maps
                    [(H, W, C)] or a padded tensor of shape (N, H, W, C).
                    For RGB, C = 3. In this case maps_ids must be None.
                Or (2) a set of M texture maps per mesh. This can either be a list of sets
                    [(M, H, W, C)] or a padded tensor of shape (N, M, H, W, C).
                    For RGB, C = 3. In this case maps_ids must be provided to
                    identify which is relevant to each face.
            faces_uvs: (N, F, 3) LongTensor giving the index into verts_uvs
                    for each face
            verts_uvs: (N, V, 2) tensor giving the uv coordinates per vertex
                    (a FloatTensor with values between 0 and 1).
            maps_ids: Used if there are to be multiple maps per face.
                    This can be either a list of map_ids [(F,)]
                    or a long tensor of shape (N, F) giving the id of the texture map
                    for each face. If maps_ids is present, the maps has an extra dimension M
                    (so maps_padded is (N, M, H, W, C) and maps_list has elements of
                    shape (M, H, W, C)).
                    Specifically, the color
                    of a vertex V is given by an average of
                       maps_padded[i, maps_ids[i, f], u, v, :]
                    over u and v integers adjacent to
                       _verts_uvs_padded[i, _faces_uvs_padded[i, f, 0], :] .
            align_corners: If true, the extreme values 0 and 1 for verts_uvs
                    indicate the centers of the edge pixels in the maps.
            padding_mode: padding mode for outside grid values
                    ("zeros", "border" or "reflection").
            sampling_mode: type of interpolation used to sample the texture.
                    Corresponds to the mode parameter in PyTorch's
                    grid_sample ("nearest" or "bilinear").

        The align_corners and padding_mode arguments correspond to the arguments
        of the `grid_sample` torch function. There is an informative illustration of
        the two align_corners options at
        https://discuss.pytorch.org/t/22663/9 .

        An example of how the indexing into the maps, with align_corners=True,
        works is as follows.
        If maps[i] has shape [1001, 101] and the value of verts_uvs[i][j]
        is [0.4, 0.3], then a value of j in faces_uvs[i] means a vertex
        whose color is given by maps[i][700, 40]. padding_mode affects what
        happens if a value in verts_uvs is less than 0 or greater than 1.
        Note that increasing a value in verts_uvs[..., 0] increases an index
        in maps, whereas increasing a value in verts_uvs[..., 1] _decreases_
        an _earlier_ index in maps.

        If align_corners=False, an example would be as follows.
        If maps[i] has shape [1000, 100] and the value of verts_uvs[i][j]
        is [0.405, 0.2995], then a value of j in faces_uvs[i] means a vertex
        whose color is given by maps[i][700, 40].
        When align_corners=False, padding_mode even matters for values in
        verts_uvs slightly above 0 or slightly below 1. In this case, the
        padding_mode matters if the first value is outside the interval
        [0.0005, 0.9995] or if the second is outside the interval
        [0.005, 0.995].
        """
        self.padding_mode = padding_mode
        self.align_corners = align_corners
        self.sampling_mode = sampling_mode
        if isinstance(faces_uvs, (list, tuple)):
            for fv in faces_uvs:
                if fv.ndim != 2 or fv.shape[-1] != 3:
                    msg = "Expected faces_uvs to be of shape (F, 3); got %r"
                    raise ValueError(msg % repr(fv.shape))
            self._faces_uvs_list = faces_uvs
            self._faces_uvs_padded = None
            self.device = torch.device("cpu")

            # These values may be overridden when textures is
            # passed into the Meshes constructor. For more details
            # refer to the __init__ of Meshes.
            self._N = len(faces_uvs)
            self._num_faces_per_mesh = [len(fv) for fv in faces_uvs]

            if self._N > 0:
                self.device = faces_uvs[0].device

        elif torch.is_tensor(faces_uvs):
            if faces_uvs.ndim != 3 or faces_uvs.shape[-1] != 3:
                msg = "Expected faces_uvs to be of shape (N, F, 3); got %r"
                raise ValueError(msg % repr(faces_uvs.shape))
            self._faces_uvs_padded = faces_uvs
            self._faces_uvs_list = None
            self.device = faces_uvs.device

            # These values may be overridden when textures is
            # passed into the Meshes constructor. For more details
            # refer to the __init__ of Meshes.
            self._N = len(faces_uvs)
            max_F = faces_uvs.shape[1]
            self._num_faces_per_mesh = [max_F] * self._N
        else:
            raise ValueError("Expected faces_uvs to be a tensor or list")

        if isinstance(verts_uvs, (list, tuple)):
            for fv in verts_uvs:
                if fv.ndim != 2 or fv.shape[-1] != 2:
                    msg = "Expected verts_uvs to be of shape (V, 2); got %r"
                    raise ValueError(msg % repr(fv.shape))
            self._verts_uvs_list = verts_uvs
            self._verts_uvs_padded = None

            if len(verts_uvs) != self._N:
                raise ValueError(
                    "verts_uvs and faces_uvs must have the same batch dimension"
                )
            if not all(v.device == self.device for v in verts_uvs):
                raise ValueError("verts_uvs and faces_uvs must be on the same device")

        elif torch.is_tensor(verts_uvs):
            if (
                verts_uvs.ndim != 3
                or verts_uvs.shape[-1] != 2
                or verts_uvs.shape[0] != self._N
            ):
                msg = "Expected verts_uvs to be of shape (N, V, 2); got %r"
                raise ValueError(msg % repr(verts_uvs.shape))
            self._verts_uvs_padded = verts_uvs
            self._verts_uvs_list = None

            if verts_uvs.device != self.device:
                raise ValueError("verts_uvs and faces_uvs must be on the same device")
        else:
            raise ValueError("Expected verts_uvs to be a tensor or list")

        self._maps_ids_padded, self._maps_ids_list = self._format_maps_ids(maps_ids)

        if isinstance(maps, (list, tuple)):
            self._maps_list = maps
        else:
            self._maps_list = None
        self._maps_padded = self._format_maps_padded(maps)

        if self._maps_padded.device != self.device:
            raise ValueError("maps must be on the same device as verts/faces uvs.")
        self.valid = torch.ones((self._N,), dtype=torch.bool, device=self.device)

    def _format_maps_ids(
        self,
        maps_ids: Optional[
            Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]
        ],
    ) -> Tuple[
        Optional[torch.Tensor], Optional[Union[List[torch.Tensor], Tuple[torch.Tensor]]]
    ]:
        if maps_ids is None:
            return None, None
        elif isinstance(maps_ids, (list, tuple)):
            for mid in maps_ids:
                if mid.ndim != 1:
                    msg = "Expected maps_ids to be of shape (F,); got %r"
                    raise ValueError(msg % repr(mid.shape))
            if len(maps_ids) != self._N:
                raise ValueError(
                    "map_ids, faces_uvs and verts_uvs must have the same batch dimension"
                )
            if not all(mid.device == self.device for mid in maps_ids):
                raise ValueError(
                    "maps_ids and verts/faces uvs must be on the same device"
                )

            if not all(
                mid.shape[0] == nfm
                for mid, nfm in zip(maps_ids, self._num_faces_per_mesh)
            ):
                raise ValueError(
                    "map_ids and faces_uvs must have the same number of faces per mesh"
                )
            if not all(mid.device == self.device for mid in maps_ids):
                raise ValueError(
                    "maps_ids and verts/faces uvs must be on the same device"
                )
            if not self._num_faces_per_mesh:
                return torch.Tensor(), maps_ids
            return list_to_padded(maps_ids, pad_value=0), maps_ids
        elif isinstance(maps_ids, torch.Tensor):
            if maps_ids.ndim != 2 or maps_ids.shape[0] != self._N:
                msg = "Expected maps_ids to be of shape (N, F); got %r"
                raise ValueError(msg % repr(maps_ids.shape))
            maps_ids_padded = maps_ids
            max_F = max(self._num_faces_per_mesh)
            if not maps_ids.shape[1] == max_F:
                raise ValueError(
                    "map_ids and faces_uvs must have the same number of faces per mesh"
                )
            if maps_ids.device != self.device:
                raise ValueError(
                    "maps_ids and verts/faces uvs must be on the same device"
                )
            return maps_ids_padded, None
        raise ValueError("Expected maps_ids to be a tensor or list")

    def _format_maps_padded(
        self, maps: Union[torch.Tensor, List[torch.Tensor]]
    ) -> torch.Tensor:
        maps_ids_none = self._maps_ids_padded is None
        if isinstance(maps, torch.Tensor):
            if not maps_ids_none:
                if maps.ndim != 5 or maps.shape[0] != self._N:
                    msg = "Expected maps to be of shape (N, M, H, W, C); got %r"
                    raise ValueError(msg % repr(maps.shape))
            elif maps.ndim != 4 or maps.shape[0] != self._N:
                msg = "Expected maps to be of shape (N, H, W, C); got %r"
                raise ValueError(msg % repr(maps.shape))
            return maps

        if isinstance(maps, (list, tuple)):
            if len(maps) != self._N:
                raise ValueError("Expected one texture map per mesh in the batch.")
            if self._N > 0:
                ndim = 3 if maps_ids_none else 4
                if not all(map.ndim == ndim for map in maps):
                    raise ValueError("Invalid number of dimensions in texture maps")
                if not all(map.shape[-1] == maps[0].shape[-1] for map in maps):
                    raise ValueError("Inconsistent number of channels in maps")
                maps_padded = (
                    _pad_texture_maps(maps, align_corners=self.align_corners)
                    if maps_ids_none
                    else _pad_texture_multiple_maps(
                        maps, align_corners=self.align_corners
                    )
                )
            else:
                if maps_ids_none:
                    maps_padded = torch.empty(
                        (self._N, 0, 0, 3), dtype=torch.float32, device=self.device
                    )
                else:
                    maps_padded = torch.empty(
                        (self._N, 0, 0, 0, 3), dtype=torch.float32, device=self.device
                    )
            return maps_padded

        raise ValueError("Expected maps to be a tensor or list of tensors.")

    def clone(self) -> "TexturesUV":
        tex = self.__class__(
            self.maps_padded().clone(),
            self.faces_uvs_padded().clone(),
            self.verts_uvs_padded().clone(),
            maps_ids=(
                self._maps_ids_padded.clone()
                if self._maps_ids_padded is not None
                else None
            ),
            align_corners=self.align_corners,
            padding_mode=self.padding_mode,
            sampling_mode=self.sampling_mode,
        )
        if self._maps_list is not None:
            tex._maps_list = [m.clone() for m in self._maps_list]
        if self._verts_uvs_list is not None:
            tex._verts_uvs_list = [v.clone() for v in self._verts_uvs_list]
        if self._faces_uvs_list is not None:
            tex._faces_uvs_list = [f.clone() for f in self._faces_uvs_list]
        if self._maps_ids_list is not None:
            tex._maps_ids_list = [f.clone() for f in self._maps_ids_list]
        num_faces = (
            self._num_faces_per_mesh.clone()
            if torch.is_tensor(self._num_faces_per_mesh)
            else self._num_faces_per_mesh
        )
        tex._num_faces_per_mesh = num_faces
        tex.valid = self.valid.clone()
        return tex

    def detach(self) -> "TexturesUV":
        tex = self.__class__(
            self.maps_padded().detach(),
            self.faces_uvs_padded().detach(),
            self.verts_uvs_padded().detach(),
            maps_ids=(
                self._maps_ids_padded.detach()
                if self._maps_ids_padded is not None
                else None
            ),
            align_corners=self.align_corners,
            padding_mode=self.padding_mode,
            sampling_mode=self.sampling_mode,
        )
        if self._maps_list is not None:
            tex._maps_list = [m.detach() for m in self._maps_list]
        if self._verts_uvs_list is not None:
            tex._verts_uvs_list = [v.detach() for v in self._verts_uvs_list]
        if self._faces_uvs_list is not None:
            tex._faces_uvs_list = [f.detach() for f in self._faces_uvs_list]
        if self._maps_ids_list is not None:
            tex._maps_ids_list = [mi.detach() for mi in self._maps_ids_list]
        num_faces = (
            self._num_faces_per_mesh.detach()
            if torch.is_tensor(self._num_faces_per_mesh)
            else self._num_faces_per_mesh
        )
        tex._num_faces_per_mesh = num_faces
        tex.valid = self.valid.detach()
        return tex

    def __getitem__(self, index) -> "TexturesUV":
        props = [
            "faces_uvs_list",
            "verts_uvs_list",
            "maps_list",
            "maps_ids_list",
            "_num_faces_per_mesh",
        ]
        new_props = self._getitem(index, props)
        faces_uvs = new_props["faces_uvs_list"]
        verts_uvs = new_props["verts_uvs_list"]
        maps = new_props["maps_list"]
        maps_ids = new_props["maps_ids_list"]

        # if index has multiple values then faces/verts/maps may be a list of tensors
        if all(isinstance(f, (list, tuple)) for f in [faces_uvs, verts_uvs, maps]):
            if maps_ids is not None and not isinstance(maps_ids, (list, tuple)):
                raise ValueError(
                    "Maps ids are  not in the correct format expected list or tuple"
                )
            new_tex = self.__class__(
                faces_uvs=faces_uvs,
                verts_uvs=verts_uvs,
                maps=maps,
                maps_ids=maps_ids,
                padding_mode=self.padding_mode,
                align_corners=self.align_corners,
                sampling_mode=self.sampling_mode,
            )
        elif all(torch.is_tensor(f) for f in [faces_uvs, verts_uvs, maps]):
            if maps_ids is not None and not torch.is_tensor(maps_ids):
                raise ValueError(
                    "Maps ids are not in the correct format expected tensor"
                )
            new_tex = self.__class__(
                faces_uvs=[faces_uvs],
                verts_uvs=[verts_uvs],
                maps=[maps],
                maps_ids=[maps_ids] if maps_ids is not None else None,
                padding_mode=self.padding_mode,
                align_corners=self.align_corners,
                sampling_mode=self.sampling_mode,
            )
        else:
            raise ValueError("Not all values are provided in the correct format")
        new_tex._num_faces_per_mesh = new_props["_num_faces_per_mesh"]
        return new_tex

    def faces_uvs_padded(self) -> torch.Tensor:
        if self._faces_uvs_padded is None:
            if self.isempty():
                self._faces_uvs_padded = torch.zeros(
                    (self._N, 0, 3), dtype=torch.float32, device=self.device
                )
            else:
                self._faces_uvs_padded = list_to_padded(
                    self._faces_uvs_list, pad_value=0.0
                )
        return self._faces_uvs_padded

    def faces_uvs_list(self) -> List[torch.Tensor]:
        if self._faces_uvs_list is None:
            if self.isempty():
                self._faces_uvs_list = [
                    torch.empty((0, 3), dtype=torch.float32, device=self.device)
                ] * self._N
            else:
                self._faces_uvs_list = padded_to_list(
                    self._faces_uvs_padded, split_size=self._num_faces_per_mesh
                )
        return self._faces_uvs_list

    def verts_uvs_padded(self) -> torch.Tensor:
        if self._verts_uvs_padded is None:
            if self.isempty():
                self._verts_uvs_padded = torch.zeros(
                    (self._N, 0, 2), dtype=torch.float32, device=self.device
                )
            else:
                self._verts_uvs_padded = list_to_padded(
                    self._verts_uvs_list, pad_value=0.0
                )
        return self._verts_uvs_padded

    def verts_uvs_list(self) -> List[torch.Tensor]:
        if self._verts_uvs_list is None:
            if self.isempty():
                self._verts_uvs_list = [
                    torch.empty((0, 2), dtype=torch.float32, device=self.device)
                ] * self._N
            else:
                # The number of vertices in the mesh and in verts_uvs can differ
                # e.g. if a vertex is shared between 3 faces, it can
                # have up to 3 different uv coordinates.
                self._verts_uvs_list = list(self._verts_uvs_padded.unbind(0))
        return self._verts_uvs_list

    def maps_ids_padded(self) -> Optional[torch.Tensor]:
        return self._maps_ids_padded

    def maps_ids_list(self) -> Optional[List[torch.Tensor]]:
        if self._maps_ids_list is not None:
            return self._maps_ids_list
        elif self._maps_ids_padded is not None:
            return self._maps_ids_padded.unbind(0)
        else:
            return None

    # Currently only the padded maps are used.
    def maps_padded(self) -> torch.Tensor:
        return self._maps_padded

    def maps_list(self) -> List[torch.Tensor]:
        if self._maps_list is not None:
            return self._maps_list
        return self._maps_padded.unbind(0)

    def extend(self, N: int) -> "TexturesUV":
        new_props = self._extend(
            N,
            [
                "maps_padded",
                "verts_uvs_padded",
                "faces_uvs_padded",
                "maps_ids_padded",
                "_num_faces_per_mesh",
            ],
        )
        new_tex = self.__class__(
            maps=new_props["maps_padded"],
            faces_uvs=new_props["faces_uvs_padded"],
            verts_uvs=new_props["verts_uvs_padded"],
            maps_ids=new_props["maps_ids_padded"],
            padding_mode=self.padding_mode,
            align_corners=self.align_corners,
            sampling_mode=self.sampling_mode,
        )

        new_tex._num_faces_per_mesh = new_props["_num_faces_per_mesh"]
        return new_tex

    # pyre-fixme[14]: `sample_textures` overrides method defined in `TexturesBase`
    #  inconsistently.
    def sample_textures(self, fragments, **kwargs) -> torch.Tensor:
        """
        Interpolate a 2D texture map using uv vertex texture coordinates for each
        face in the mesh. First interpolate the vertex uvs using barycentric coordinates
        for each pixel in the rasterized output. Then interpolate the texture map
        using the uv coordinate for each pixel.

        Args:
            fragments:
                The outputs of rasterization. From this we use

                - pix_to_face: LongTensor of shape (N, H, W, K) specifying the indices
                of the faces (in the packed representation) which
                overlap each pixel in the image.
                - barycentric_coords: FloatTensor of shape (N, H, W, K, 3) specifying
                the barycentric coordinates of each pixel
                relative to the faces (in the packed
                representation) which overlap the pixel.

        Returns:
            texels: tensor of shape (N, H, W, K, C) giving the interpolated
            texture for each pixel in the rasterized image.
        """
        if self.isempty():
            faces_verts_uvs = torch.zeros(
                (self._N, 3, 2), dtype=torch.float32, device=self.device
            )
        else:
            packing_list = [
                i[j] for i, j in zip(self.verts_uvs_list(), self.faces_uvs_list())
            ]
            faces_verts_uvs = torch.cat(packing_list)

        # pixel_uvs: (N, H, W, K, 2)
        pixel_uvs = interpolate_face_attributes(
            fragments.pix_to_face, fragments.bary_coords, faces_verts_uvs
        )

        N, H_out, W_out, K = fragments.pix_to_face.shape

        texture_maps = self.maps_padded()
        maps_ids_padded = self.maps_ids_padded()
        if maps_ids_padded is None:
            # pixel_uvs: (N, H, W, K, 2) -> (N, K, H, W, 2) -> (NK, H, W, 2)
            pixel_uvs = pixel_uvs.permute(0, 3, 1, 2, 4).reshape(N * K, H_out, W_out, 2)
            N, H_in, W_in, C = texture_maps.shape  # 3 for RGB

            # textures.map:
            #   (N, H, W, C) -> (N, C, H, W) -> (1, N, C, H, W)
            #   -> expand (K, N, C, H, W) -> reshape (N*K, C, H, W)
            texture_maps = (
                texture_maps.permute(0, 3, 1, 2)[None, ...]
                .expand(K, -1, -1, -1, -1)
                .transpose(0, 1)
                .reshape(N * K, C, H_in, W_in)
            )
            # Textures: (N*K, C, H, W), pixel_uvs: (N*K, H, W, 2)
            # Now need to format the pixel uvs and the texture map correctly!
            # From pytorch docs, grid_sample takes `grid` and `input`:
            #   grid specifies the sampling pixel locations normalized by
            #   the input spatial dimensions It should have most
            #   values in the range of [-1, 1]. Values x = -1, y = -1
            #   is the left-top pixel of input, and values x = 1, y = 1 is the
            #   right-bottom pixel of input.

            # map to a range of [-1, 1] and flip the y axis
            pixel_uvs = torch.lerp(
                pixel_uvs.new_tensor([-1.0, 1.0]),
                pixel_uvs.new_tensor([1.0, -1.0]),
                pixel_uvs,
            )

            if texture_maps.device != pixel_uvs.device:
                texture_maps = texture_maps.to(pixel_uvs.device)
            texels = F.grid_sample(
                texture_maps,
                pixel_uvs,
                mode=self.sampling_mode,
                align_corners=self.align_corners,
                padding_mode=self.padding_mode,
            )
            # texels now has shape (NK, C, H_out, W_out)
            texels = texels.reshape(N, K, C, H_out, W_out).permute(0, 3, 4, 1, 2)
            return texels
        else:
            # We have maps_ids_padded: (N, F), textures_map: (N, M, Hi, Wi, C),
            # fragments.pix_to_face: (N, Ho, Wo, K)
            # Get pixel_to_map_ids: (N, K, Ho, Wo) by indexing pix_to_face into maps_ids
            N, M, H_in, W_in, C = texture_maps.shape  # 3 for RGB

            mask = fragments.pix_to_face < 0
            pix_to_face = fragments.pix_to_face.clone()
            pix_to_face[mask] = 0

            pixel_to_map_ids = (
                maps_ids_padded.flatten()
                .gather(0, pix_to_face.flatten())
                .view(N, H_out, W_out, K, 1)
                .permute(0, 3, 1, 2, 4)
            )  # N x H_out x W_out x K x 1

            # Normalize between -1 and 1 with M (number of maps)
            pixel_to_map_ids = (2.0 * pixel_to_map_ids.float() / float(M - 1)) - 1
            pixel_uvs = pixel_uvs.permute(0, 3, 1, 2, 4)
            pixel_uvs = torch.lerp(
                pixel_uvs.new_tensor([-1.0, 1.0]),
                pixel_uvs.new_tensor([1.0, -1.0]),
                pixel_uvs,
            )  # N x H_out x W_out x K x 2

            # N x H_out x W_out x K x 3
            pixel_uvms = torch.cat((pixel_uvs, pixel_to_map_ids), dim=4)
            # (N, M, H, W, C) -> (N, C, M, H, W)
            texture_maps = texture_maps.permute(0, 4, 1, 2, 3)
            if texture_maps.device != pixel_uvs.device:
                texture_maps = texture_maps.to(pixel_uvs.device)
            texels = F.grid_sample(
                texture_maps,
                pixel_uvms,
                mode=self.sampling_mode,
                align_corners=self.align_corners,
                padding_mode=self.padding_mode,
            )
            # (N, C, K, H_out, W_out) -> (N, H_out, W_out, K, C)
            texels = texels.permute(0, 3, 4, 2, 1).contiguous()
            return texels

    def faces_verts_textures_packed(self) -> torch.Tensor:
        """
        Samples texture from each vertex and for each face in the mesh.
        For N meshes with {Fi} number of faces, it returns a
        tensor of shape sum(Fi)x3xC (C = 3 for RGB).
        You can use the utils function in structures.utils to convert the
        packed representation to a list or padded.
        """
        if self.isempty():
            return torch.zeros(
                (0, 3, self.maps_padded().shape[-1]),
                dtype=torch.float32,
                device=self.device,
            )
        else:
            packing_list = [
                i[j] for i, j in zip(self.verts_uvs_list(), self.faces_uvs_list())
            ]
            faces_verts_uvs = _list_to_padded_wrapper(
                packing_list, pad_value=0.0
            )  # Nxmax(Fi)x3x2
        # map to a range of [-1, 1] and flip the y axis
        faces_verts_uvs = torch.lerp(
            faces_verts_uvs.new_tensor([-1.0, 1.0]),
            faces_verts_uvs.new_tensor([1.0, -1.0]),
            faces_verts_uvs,
        )
        texture_maps = self.maps_padded()  # NxHxWxC or NxMxHxWxC
        maps_ids_padded = self.maps_ids_padded()
        if maps_ids_padded is None:
            texture_maps = texture_maps.permute(0, 3, 1, 2)  # NxCxHxW
        else:
            M = texture_maps.shape[1]
            # (N, M, H, W, C) -> (N, C, M, H, W)
            texture_maps = texture_maps.permute(0, 4, 1, 2, 3)
            # expand maps_ids to (N, F, 3, 1)
            maps_ids_padded = maps_ids_padded[:, :, None, None].expand(-1, -1, 3, -1)
            maps_ids_padded = (2.0 * maps_ids_padded.float() / float(M - 1)) - 1.0

            # (N, F, 3, 2+1) -> (N, 1, F, 3, 3)
            faces_verts_uvs = torch.cat(
                (faces_verts_uvs, maps_ids_padded), dim=3
            ).unsqueeze(1)
            # (N, M, H, W, C) -> (N, C, H, W, M)
            # texture_maps = texture_maps.permute(0, 4, 2, 3, 1)
        textures = F.grid_sample(
            texture_maps,
            faces_verts_uvs,
            mode=self.sampling_mode,
            align_corners=self.align_corners,
            padding_mode=self.padding_mode,
        )  # (N, C, max(Fi), 3)
        if maps_ids_padded is not None:
            textures = textures.squeeze(dim=2)
        # (N, C, max(Fi), 3) -> (N, max(Fi), 3, C)
        textures = textures.permute(0, 2, 3, 1)
        textures = _padded_to_list_wrapper(
            textures, split_size=self._num_faces_per_mesh
        )  # list of N {Fix3xC} tensors
        return list_to_packed(textures)[0]

    def join_batch(self, textures: List["TexturesUV"]) -> "TexturesUV":
        """
        Join the list of textures given by `textures` to
        self to create a batch of textures. Return a new
        TexturesUV object with the combined textures.

        Args:
            textures: List of TexturesUV objects

        Returns:
            new_tex: TexturesUV object with the combined
            textures from self and the list `textures`.
        """
        if self.maps_ids_padded() is not None:
            # TODO
            raise NotImplementedError(
                "join_batch does not support TexturesUV with multiple maps"
            )
        tex_types_same = all(isinstance(tex, TexturesUV) for tex in textures)
        if not tex_types_same:
            raise ValueError("All textures must be of type TexturesUV.")

        padding_modes_same = all(
            tex.padding_mode == self.padding_mode for tex in textures
        )
        if not padding_modes_same:
            raise ValueError("All textures must have the same padding_mode.")
        align_corners_same = all(
            tex.align_corners == self.align_corners for tex in textures
        )
        if not align_corners_same:
            raise ValueError("All textures must have the same align_corners value.")
        sampling_mode_same = all(
            tex.sampling_mode == self.sampling_mode for tex in textures
        )
        if not sampling_mode_same:
            raise ValueError("All textures must have the same sampling_mode.")

        verts_uvs_list = []
        faces_uvs_list = []
        maps_list = []
        faces_uvs_list += self.faces_uvs_list()
        verts_uvs_list += self.verts_uvs_list()
        maps_list += self.maps_list()
        num_faces_per_mesh = self._num_faces_per_mesh.copy()
        for tex in textures:
            verts_uvs_list += tex.verts_uvs_list()
            faces_uvs_list += tex.faces_uvs_list()
            num_faces_per_mesh += tex._num_faces_per_mesh
            maps_list += tex.maps_list()

        new_tex = self.__class__(
            maps=maps_list,
            faces_uvs=faces_uvs_list,
            verts_uvs=verts_uvs_list,
            padding_mode=self.padding_mode,
            align_corners=self.align_corners,
            sampling_mode=self.sampling_mode,
        )
        new_tex._num_faces_per_mesh = num_faces_per_mesh
        return new_tex

    def _place_map_into_single_map(
        self, single_map: torch.Tensor, map_: torch.Tensor, location: PackedRectangle
    ) -> None:
        """
        Copy map into a larger tensor single_map at the destination specified by location.
        If align_corners is False, we add the needed border around the destination.

        Used by join_scene.

        Args:
            single_map: (total_H, total_W, C)
            map_: (H, W, C) source data
            location: where to place map
        """
        do_flip = location.flipped
        source = map_.transpose(0, 1) if do_flip else map_
        border_width = 0 if self.align_corners else 1
        lower_u = location.x + border_width
        lower_v = location.y + border_width
        upper_u = lower_u + source.shape[0]
        upper_v = lower_v + source.shape[1]
        single_map[lower_u:upper_u, lower_v:upper_v] = source

        if self.padding_mode != "zeros" and not self.align_corners:
            single_map[lower_u - 1, lower_v:upper_v] = single_map[
                lower_u, lower_v:upper_v
            ]
            single_map[upper_u, lower_v:upper_v] = single_map[
                upper_u - 1, lower_v:upper_v
            ]
            single_map[lower_u:upper_u, lower_v - 1] = single_map[
                lower_u:upper_u, lower_v
            ]
            single_map[lower_u:upper_u, upper_v] = single_map[
                lower_u:upper_u, upper_v - 1
            ]
            single_map[lower_u - 1, lower_v - 1] = single_map[lower_u, lower_v]
            single_map[lower_u - 1, upper_v] = single_map[lower_u, upper_v - 1]
            single_map[upper_u, lower_v - 1] = single_map[upper_u - 1, lower_v]
            single_map[upper_u, upper_v] = single_map[upper_u - 1, upper_v - 1]

    def join_scene(self) -> "TexturesUV":
        """
        Return a new TexturesUV amalgamating the batch.

        We calculate a large single map which contains the original maps,
        and find verts_uvs to point into it. This will not replicate
        behavior of padding for verts_uvs values outside [0,1].

        If align_corners=False, we need to add an artificial border around
        every map.

        We use the function `pack_unique_rectangles` to provide a layout for
        the single map. This means that if self was created with a list of maps,
        and to() has not been called, and there were two maps which were exactly
        the same tensor object, then they will become the same data in the unified map.
        _place_map_into_single_map is used to copy the maps into the single map.
        The merging of verts_uvs and faces_uvs is handled locally in this function.
        """
        if self.maps_ids_padded() is not None:
            # TODO
            raise NotImplementedError("join_scene does not support multiple maps.")
        maps = self.maps_list()
        heights_and_widths = []
        extra_border = 0 if self.align_corners else 2
        for map_ in maps:
            heights_and_widths.append(
                Rectangle(
                    map_.shape[0] + extra_border, map_.shape[1] + extra_border, id(map_)
                )
            )
        merging_plan = pack_unique_rectangles(heights_and_widths)
        C = maps[0].shape[-1]
        single_map = maps[0].new_zeros((*merging_plan.total_size, C))
        verts_uvs = self.verts_uvs_list()
        verts_uvs_merged = []

        for map_, loc, uvs in zip(maps, merging_plan.locations, verts_uvs):
            new_uvs = uvs.clone()
            if loc.is_first:
                self._place_map_into_single_map(single_map, map_, loc)
            do_flip = loc.flipped
            x_shape = map_.shape[1] if do_flip else map_.shape[0]
            y_shape = map_.shape[0] if do_flip else map_.shape[1]

            if do_flip:
                # Here we have flipped / transposed the map.
                # In uvs, the y values are decreasing from 1 to 0 and the x
                # values increase from 0 to 1. We subtract all values from 1
                # as the x's become y's and the y's become x's.
                new_uvs = 1.0 - new_uvs[:, [1, 0]]
                if TYPE_CHECKING:
                    new_uvs = torch.Tensor(new_uvs)

            # If align_corners is True, then an index of x (where x is in
            # the range 0 .. map_.shape[1]-1) in one of the input maps
            # was hit by a u of x/(map_.shape[1]-1).
            # That x is located at the index loc[1] + x in the single_map, and
            # to hit that we need u to equal (loc[1] + x) / (total_size[1]-1)
            # so the old u should be mapped to
            #   { u*(map_.shape[1]-1) + loc[1] } / (total_size[1]-1)

            # Also, an index of y (where y is in
            # the range 0 .. map_.shape[0]-1) in one of the input maps
            # was hit by a v of 1 - y/(map_.shape[0]-1).
            # That y is located at the index loc[0] + y in the single_map, and
            # to hit that we need v to equal 1 - (loc[0] + y) / (total_size[0]-1)
            # so the old v should be mapped to
            #   1 - { (1-v)*(map_.shape[0]-1) + loc[0] } / (total_size[0]-1)
            # =
            # { v*(map_.shape[0]-1) + total_size[0] - map.shape[0] - loc[0] }
            #        / (total_size[0]-1)

            # If align_corners is False, then an index of x (where x is in
            # the range 1 .. map_.shape[1]-2) in one of the input maps
            # was hit by a u of (x+0.5)/(map_.shape[1]).
            # That x is located at the index loc[1] + 1 + x in the single_map,
            # (where the 1 is for the border)
            # and to hit that we need u to equal (loc[1] + 1 + x + 0.5) / (total_size[1])
            # so the old u should be mapped to
            #   { loc[1] + 1 + u*map_.shape[1]-0.5 + 0.5 } / (total_size[1])
            #  = { loc[1] + 1 + u*map_.shape[1] } / (total_size[1])

            # Also, an index of y (where y is in
            # the range 1 .. map_.shape[0]-2) in one of the input maps
            # was hit by a v of 1 - (y+0.5)/(map_.shape[0]).
            # That y is located at the index loc[0] + 1 + y in the single_map,
            # (where the 1 is for the border)
            # and to hit that we need v to equal 1 - (loc[0] + 1 + y + 0.5) / (total_size[0])
            # so the old v should be mapped to
            #   1 - { loc[0] + 1 + (1-v)*map_.shape[0]-0.5 + 0.5 } / (total_size[0])
            #  = { total_size[0] - loc[0] -1 - (1-v)*map_.shape[0]  }
            #         / (total_size[0])
            #  = { total_size[0] - loc[0] - map.shape[0] - 1 + v*map_.shape[0] }
            #         / (total_size[0])

            # We change the y's in new_uvs for the scaling of height,
            # and the x's for the scaling of width.
            # That is why the 1's and 0's are mismatched in these lines.
            one_if_align = 1 if self.align_corners else 0
            one_if_not_align = 1 - one_if_align
            denom_x = merging_plan.total_size[0] - one_if_align
            scale_x = x_shape - one_if_align
            denom_y = merging_plan.total_size[1] - one_if_align
            scale_y = y_shape - one_if_align
            new_uvs[:, 1] *= scale_x / denom_x
            new_uvs[:, 1] += (
                merging_plan.total_size[0] - x_shape - loc.x - one_if_not_align
            ) / denom_x
            new_uvs[:, 0] *= scale_y / denom_y
            new_uvs[:, 0] += (loc.y + one_if_not_align) / denom_y

            verts_uvs_merged.append(new_uvs)

        faces_uvs_merged = []
        offset = 0
        for faces_uvs_, verts_uvs_ in zip(self.faces_uvs_list(), verts_uvs):
            faces_uvs_merged.append(offset + faces_uvs_)
            offset += verts_uvs_.shape[0]

        return self.__class__(
            maps=[single_map],
            faces_uvs=[torch.cat(faces_uvs_merged)],
            verts_uvs=[torch.cat(verts_uvs_merged)],
            align_corners=self.align_corners,
            padding_mode=self.padding_mode,
            sampling_mode=self.sampling_mode,
        )

    def centers_for_image(self, index: int) -> torch.Tensor:
        """
        Return the locations in the texture map which correspond to the given
        verts_uvs, for one of the meshes. This is potentially useful for
        visualizing the data. See the texturesuv_image_matplotlib and
        texturesuv_image_PIL functions.

        Args:
            index: batch index of the mesh whose centers to return.

        Returns:
            centers: coordinates of points in the texture image
                - a FloatTensor of shape (V,2)
        """
        if self.maps_ids_padded() is not None:
            # TODO: invent a visualization for the multiple maps case
            raise NotImplementedError("This function does not support multiple maps.")
        if self._N != 1:
            raise ValueError(
                "This function only supports plotting textures for one mesh."
            )
        texture_image = self.maps_padded()
        verts_uvs = self.verts_uvs_list()[index][None]
        _, H, W, _3 = texture_image.shape
        coord1 = torch.arange(W).expand(H, W)
        coord2 = torch.arange(H)[:, None].expand(H, W)
        coords = torch.stack([coord1, coord2])[None]
        with torch.no_grad():
            # Get xy cartesian coordinates based on the uv coordinates
            centers = F.grid_sample(
                torch.flip(coords.to(texture_image), [2]),
                # Convert from [0, 1] -> [-1, 1] range expected by grid sample
                verts_uvs[:, None] * 2.0 - 1,
                mode=self.sampling_mode,
                align_corners=self.align_corners,
                padding_mode=self.padding_mode,
            ).cpu()
            centers = centers[0, :, 0].T
        return centers

    def check_shapes(
        self, batch_size: int, max_num_verts: int, max_num_faces: int
    ) -> bool:
        """
        Check if the dimensions of the verts/faces uvs match that of the mesh
        """
        # (N, F) should be the same
        # (N, V) is not guaranteed to be the same
        return (self.faces_uvs_padded().shape[0:2] == (batch_size, max_num_faces)) and (
            self.verts_uvs_padded().shape[0] == batch_size
        )

    def submeshes(
        self,
        vertex_ids_list: List[List[torch.LongTensor]],
        faces_ids_list: List[List[torch.LongTensor]],
    ) -> "TexturesUV":
        """
        Extract a sub-texture for use in a submesh.

        If the meshes batch corresponding to this  TexturesUV contains
        `n = len(faces_ids_list)` meshes, then self.faces_uvs_padded()
        will be of length n. After submeshing, we obtain a batch of
        `k = sum(len(f) for f in faces_ids_list` submeshes (see Meshes.submeshes). This
        function creates a corresponding  TexturesUV object with `faces_uvs_padded`
        of length `k`.

        Args:
            vertex_ids_list: Not used when submeshing TexturesUV.

            face_ids_list: A list of length equal to self.faces_uvs_padded. Each
                element is a LongTensor listing the face ids that the submesh keeps in
                each respective mesh.


        Returns:
            A  "TexturesUV in which faces_uvs_padded, verts_uvs_padded, and maps_padded
            have length sum(len(faces) for faces in faces_ids_list)
        """
        if self.maps_ids_padded() is not None:
            # TODO
            raise NotImplementedError("This function does not support multiple maps.")
        if len(faces_ids_list) != len(self.faces_uvs_padded()):
            raise IndexError(
                "faces_uvs_padded must be of " "the same length as face_ids_list."
            )

        sub_faces_uvs, sub_verts_uvs, sub_maps = [], [], []
        for faces_ids, faces_uvs, verts_uvs, map_ in zip(
            faces_ids_list,
            self.faces_uvs_padded(),
            self.verts_uvs_padded(),
            self.maps_padded(),
        ):
            for faces_ids_submesh in faces_ids:
                sub_faces_uvs.append(faces_uvs[faces_ids_submesh])
                sub_verts_uvs.append(verts_uvs)
                sub_maps.append(map_)

        return self.__class__(
            maps=sub_maps,
            faces_uvs=sub_faces_uvs,
            verts_uvs=sub_verts_uvs,
            padding_mode=self.padding_mode,
            align_corners=self.align_corners,
            sampling_mode=self.sampling_mode,
        )


class TexturesVertex(TexturesBase):
    def __init__(
        self,
        verts_features: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]],
    ) -> None:
        """
        Batched texture representation where each vertex in a mesh
        has a C dimensional feature vector.

        Args:
            verts_features: list of (Vi, C) or (N, V, C) tensor giving a feature
                vector with arbitrary dimensions for each vertex.
        """
        if isinstance(verts_features, (tuple, list)):
            correct_shape = all(
                (torch.is_tensor(v) and v.ndim == 2) for v in verts_features
            )
            if not correct_shape:
                raise ValueError(
                    "Expected verts_features to be a list of tensors of shape (V, C)."
                )

            self._verts_features_list = verts_features
            self._verts_features_padded = None
            self.device = torch.device("cpu")

            # These values may be overridden when textures is
            # passed into the Meshes constructor. For more details
            # refer to the __init__ of Meshes.
            self._N = len(verts_features)
            self._num_verts_per_mesh = [len(fv) for fv in verts_features]

            if self._N > 0:
                self.device = verts_features[0].device

        elif torch.is_tensor(verts_features):
            if verts_features.ndim != 3:
                msg = "Expected verts_features to be of shape (N, V, C); got %r"
                raise ValueError(msg % repr(verts_features.shape))
            self._verts_features_padded = verts_features
            self._verts_features_list = None
            self.device = verts_features.device

            # These values may be overridden when textures is
            # passed into the Meshes constructor. For more details
            # refer to the __init__ of Meshes.
            self._N = len(verts_features)
            max_F = verts_features.shape[1]
            self._num_verts_per_mesh = [max_F] * self._N
        else:
            raise ValueError("verts_features must be a tensor or list of tensors")

        # This is set inside the Meshes object when textures is
        # passed into the Meshes constructor. For more details
        # refer to the __init__ of Meshes.
        self.valid = torch.ones((self._N,), dtype=torch.bool, device=self.device)

    def clone(self) -> "TexturesVertex":
        tex = self.__class__(self.verts_features_padded().clone())
        if self._verts_features_list is not None:
            tex._verts_features_list = [f.clone() for f in self._verts_features_list]
        tex._num_verts_per_mesh = self._num_verts_per_mesh.copy()
        tex.valid = self.valid.clone()
        return tex

    def detach(self) -> "TexturesVertex":
        tex = self.__class__(self.verts_features_padded().detach())
        if self._verts_features_list is not None:
            tex._verts_features_list = [f.detach() for f in self._verts_features_list]
        tex._num_verts_per_mesh = self._num_verts_per_mesh.copy()
        tex.valid = self.valid.detach()
        return tex

    def __getitem__(self, index) -> "TexturesVertex":
        props = ["verts_features_list", "_num_verts_per_mesh"]
        new_props = self._getitem(index, props)
        verts_features = new_props["verts_features_list"]
        if isinstance(verts_features, list):
            # Handle the case of an empty list
            if len(verts_features) == 0:
                verts_features = torch.empty(
                    size=(0, 0, 3),
                    dtype=torch.float32,
                    device=self.verts_features_padded().device,
                )
            new_tex = self.__class__(verts_features=verts_features)
        elif torch.is_tensor(verts_features):
            new_tex = self.__class__(verts_features=[verts_features])
        else:
            raise ValueError("Not all values are provided in the correct format")
        new_tex._num_verts_per_mesh = new_props["_num_verts_per_mesh"]
        return new_tex

    def verts_features_padded(self) -> torch.Tensor:
        if self._verts_features_padded is None:
            if self.isempty():
                self._verts_features_padded = torch.zeros(
                    (self._N, 0, 3, 0), dtype=torch.float32, device=self.device
                )
            else:
                self._verts_features_padded = list_to_padded(
                    self._verts_features_list, pad_value=0.0
                )
        return self._verts_features_padded

    def verts_features_list(self) -> List[torch.Tensor]:
        if self._verts_features_list is None:
            if self.isempty():
                self._verts_features_list = [
                    torch.empty((0, 3), dtype=torch.float32, device=self.device)
                ] * self._N
            else:
                self._verts_features_list = padded_to_list(
                    self._verts_features_padded, split_size=self._num_verts_per_mesh
                )
        return self._verts_features_list

    def verts_features_packed(self) -> torch.Tensor:
        if self.isempty():
            return torch.zeros((self._N, 3, 0), dtype=torch.float32, device=self.device)
        verts_features_list = self.verts_features_list()
        return list_to_packed(verts_features_list)[0]

    def extend(self, N: int) -> "TexturesVertex":
        new_props = self._extend(N, ["verts_features_padded", "_num_verts_per_mesh"])
        new_tex = self.__class__(verts_features=new_props["verts_features_padded"])
        new_tex._num_verts_per_mesh = new_props["_num_verts_per_mesh"]
        return new_tex

    # pyre-fixme[14]: `sample_textures` overrides method defined in `TexturesBase`
    #  inconsistently.
    def sample_textures(self, fragments, faces_packed=None) -> torch.Tensor:
        """
        Determine the color for each rasterized face. Interpolate the colors for
        vertices which form the face using the barycentric coordinates.
        Args:
            fragments:
                The outputs of rasterization. From this we use

                - pix_to_face: LongTensor of shape (N, H, W, K) specifying the indices
                of the faces (in the packed representation) which
                overlap each pixel in the image.
                - barycentric_coords: FloatTensor of shape (N, H, W, K, 3) specifying
                the barycentric coordinates of each pixel
                relative to the faces (in the packed
                representation) which overlap the pixel.

        Returns:
            texels: A texture per pixel of shape (N, H, W, K, C).
            There will be one C dimensional value for each element in
            fragments.pix_to_face.
        """
        verts_features_packed = self.verts_features_packed()
        faces_verts_features = verts_features_packed[faces_packed]

        texels = interpolate_face_attributes(
            fragments.pix_to_face, fragments.bary_coords, faces_verts_features
        )
        return texels

    def submeshes(
        self,
        vertex_ids_list: List[List[torch.LongTensor]],
        faces_ids_list: List[List[torch.LongTensor]],
    ) -> "TexturesVertex":
        """
        Extract a sub-texture for use in a submesh.

        If the meshes batch corresponding to this TexturesVertex contains
        `n = len(vertex_ids_list)` meshes, then self.verts_features_list()
        will be of length n. After submeshing, we obtain a batch of
        `k = sum(len(v) for v in vertex_ids_list` submeshes (see Meshes.submeshes). This
        function creates a corresponding TexturesVertex object with `verts_features_list`
        of length `k`.

        Args:
            vertex_ids_list: A list of length equal to self.verts_features_list. Each
                element is a LongTensor listing the vertices that the submesh keeps in
                each respective mesh.

            face_ids_list: Not used when submeshing TexturesVertex.

        Returns:
            A TexturesVertex in which verts_features_list has length
            sum(len(vertices) for vertices in vertex_ids_list). Each element contains
            vertex features corresponding to the subset of vertices in that submesh.
        """
        if len(vertex_ids_list) != len(self.verts_features_list()):
            raise IndexError(
                "verts_features_list must be of " "the same length as vertex_ids_list."
            )

        sub_features = []
        for vertex_ids, features in zip(vertex_ids_list, self.verts_features_list()):
            for vertex_ids_submesh in vertex_ids:
                sub_features.append(features[vertex_ids_submesh])

        return self.__class__(sub_features)

    def faces_verts_textures_packed(self, faces_packed=None) -> torch.Tensor:
        """
        Samples texture from each vertex and for each face in the mesh.
        For N meshes with {Fi} number of faces, it returns a
        tensor of shape sum(Fi)x3xC (C = 3 for RGB).
        You can use the utils function in structures.utils to convert the
        packed representation to a list or padded.
        """
        verts_features_packed = self.verts_features_packed()
        faces_verts_features = verts_features_packed[faces_packed]
        return faces_verts_features

    def join_batch(self, textures: List["TexturesVertex"]) -> "TexturesVertex":
        """
        Join the list of textures given by `textures` to
        self to create a batch of textures. Return a new
        TexturesVertex object with the combined textures.

        Args:
            textures: List of TexturesVertex objects

        Returns:
            new_tex: TexturesVertex object with the combined
            textures from self and the list `textures`.
        """
        tex_types_same = all(isinstance(tex, TexturesVertex) for tex in textures)
        if not tex_types_same:
            raise ValueError("All textures must be of type TexturesVertex.")

        verts_features_list = []
        verts_features_list += self.verts_features_list()
        num_verts_per_mesh = self._num_verts_per_mesh.copy()
        for tex in textures:
            verts_features_list += tex.verts_features_list()
            num_verts_per_mesh += tex._num_verts_per_mesh

        new_tex = self.__class__(verts_features=verts_features_list)
        new_tex._num_verts_per_mesh = num_verts_per_mesh
        return new_tex

    def join_scene(self) -> "TexturesVertex":
        """
        Return a new TexturesVertex amalgamating the batch.
        """
        return self.__class__(verts_features=[torch.cat(self.verts_features_list())])

    def check_shapes(
        self, batch_size: int, max_num_verts: int, max_num_faces: int
    ) -> bool:
        """
        Check if the dimensions of the verts features match that of the mesh verts
        """
        # (N, V) should be the same
        return self.verts_features_padded().shape[:-1] == (batch_size, max_num_verts)