Function rendering_batches

motion_rep/motion_checker.py:185–239

Render depth maps for all frames in batches.

(
    all_verts: torch.Tensor,
    faces: torch.Tensor,
    width: int,
    height: int,
    focal_length: float,
    R: Optional[torch.Tensor] = None,
    T: Optional[torch.Tensor] = None,
    batch_size: int = 8,
    render_multiple: bool = True,
    reverse_axis: bool = False,
)

def rendering_batches(
    all_verts: torch.Tensor,
    faces: torch.Tensor,
    width: int,
    height: int,
    focal_length: float,
    R: Optional[torch.Tensor] = None,
    T: Optional[torch.Tensor] = None,
    batch_size: int = 8,
    render_multiple: bool = True,
    reverse_axis: bool = False,
) -> torch.Tensor:
    """Render depth maps for all frames in batches."""
    person_num = all_verts.shape[0]
    frame_num = all_verts.shape[1]
    verts_num = all_verts.shape[2]

    if render_multiple:
        multiple_person_faces = []
        vert_offset = 0
        for _ in range(person_num):
            faces_offset = faces + vert_offset
            multiple_person_faces.append(faces_offset)
            vert_offset += verts_num
        multiple_person_faces = torch.cat(multiple_person_faces, dim=0)

    all_depth_maps = []
    for i in range(0, frame_num, batch_size):
        batch_start = i
        batch_end = min(i + batch_size, frame_num)
        batch_verts = all_verts[:, batch_start:batch_end].transpose(1, 0)
        actual_batch_size = batch_end - batch_start
        batch_R = R[batch_start:batch_end] if R is not None else None
        batch_T = T[batch_start:batch_end] if T is not None else None

        if render_multiple:
            verts = batch_verts.reshape(actual_batch_size, person_num * verts_num, 3)
            faces_tensor = multiple_person_faces.unsqueeze(0).repeat(actual_batch_size, 1, 1)
        else:
            verts = batch_verts
            faces_tensor = faces.unsqueeze(0).repeat(actual_batch_size, 1, 1)

        depth_maps = render_depth_maps(
            vertices=verts,
            faces=faces_tensor,
            image_size=(height, width),
            focal_length=focal_length,
            reverse_axis=reverse_axis,
            R=batch_R,
            T=batch_T,
        )
        all_depth_maps.append(depth_maps)

    depth_maps = torch.cat(all_depth_maps, dim=0)
    return depth_maps


def render_and_save_overlay(