medusa.render.image#

Module with a renderer class based on pytorch3d.

Module Contents#

class medusa.render.image.PytorchRenderer(viewport, cam_mat, cam_type, shading='flat', lights=None, background=(0, 0, 0), device=DEVICE)[source]#

A pytorch3d-based renderer.

Parameters:

  • viewport (tuple[int]) – Desired output image size (width, height), in pixels; should match the original image (before cropping) that was reconstructed

  • cam_mat (torch.tensor) – A camera matrix to set the position/angle of the camera

  • cam_type (str) – Either ‘orthographic’ (for Flame-based reconstructions) or ‘perpective’ (for mediapipe reconstructions)

  • shading (str) – Type of shading (‘flat’, ‘smooth’)

  • wireframe_opts (None, dict) – Dictionary with extra options for wireframe rendering (options: ‘width’, ‘color’)

  • device (str) – Device to store the image on (‘cuda’ or ‘cpu’)

draw_landmarks(img, lms, radius=3, colors='green')[source]#
alpha_blend(img, background, face_alpha=None)[source]#

Simple alpha blend of a rendered image and a background. The image (img) is assumed to be an RGBA image and the background (background) is assumed to be a RGB image. The alpha channel of the image is used to blend them together. The optional threshold parameter can be used to impose a sharp cutoff.

Parameters:

  • img (torch.tensor) – A 3D or 4D tensor of shape (batch size) x height x width x 4 (RGBA)

  • background (np.ndarray) – A 3D or 4D tensor shape height x width x 3 (RGB[A])

Returns:

img – A blended image

Return type:

torch.tensor

static save_image(f_out, img)[source]#

Saves a single image (using PIL) to disk.

Parameters:

f_out (str, Path) – Path where the image should be saved

forward(v, tris, overlay=None, single_image=True)[source]#

Performs the actual rendering for a given (batch of) mesh(es).

Parameters:

  • v (torch.tensor) – A 3D (batch size x vertices x 3) tensor with vertices

  • tris (torch.tensor) – A 3D (batch size x vertices x 3) tensor with triangles

  • overlay (torch.tensor) – A tensor with shape (batch size x vertices) with vertex colors

  • single_image (bool) – Whether a single image with (potentially) multiple faces should be renderer (True) or multiple images with a single face should be renderered (False)

Returns:

img – A 4D tensor with uint8 values of shape batch size x h x w x 3 (RGB), where h and w are defined in the viewport

Return type:

torch.tensor