imaginaire.utils.visualization package

Submodules

imaginaire.utils.visualization.common module

class imaginaire.utils.visualization.common.Colorize(n=35)[source]

Bases: object

Class to colorize segmentation maps.

imaginaire.utils.visualization.common.labelcolormap(N)[source]

Create colors for segmentation label ids.

Parameters

N (int) – Number of labels.

imaginaire.utils.visualization.common.plot_keypoints(image, keypoints, normalize=True)[source]

Plot keypoints on image.

Parameters

  • image (PIL.Image, or numpy.ndarray, or torch.Tensor) – Input image.

  • keypoints (np.ndarray or torch.Tensor, Nx2) – Keypoint locations.

  • normalize (bool) – Whether to normalize the image or not.

imaginaire.utils.visualization.common.plot_keypoints_on_black(resize_h, resize_w, crop_h, crop_w, is_flipped, cfgdata, keypoints)[source]

Plot keypoints on black image.

Parameters

  • resize_h (int) – Height to be resized to.

  • resize_w (int) – Width to be resized to.

  • crop_h (int) – Height of the cropping.

  • crop_w (int) – Width of the cropping.

  • is_flipped (bool) – If image is a flipped version.

  • cfgdata (obj) – Data configuration object.

  • keypoints (np.ndarray) – Keypoint locations. Shape of (Nx2) or (TxNx2).

Returns

List of images (output_h, output_w, 3).

Return type

(list of np.ndarray)

imaginaire.utils.visualization.common.save_tensor_image(filename, image, minus1to1_normalized=False)[source]

Convert a 3 dimensional torch tensor to a PIL image with the desired width and height.

Parameters

  • filename (str) – Image filename to be saved to.

  • image (3 x W1 x H1 tensor) – Image tensor

  • minus1to1_normalized (bool) – True if the tensor values are in [-1,

  • Otherwise, we assume the values are in [0, 1] (1]) –

Returns

The resulting PIL image.

Return type

(PIL image)

imaginaire.utils.visualization.common.tensor2flow(tensor, imtype=<class 'numpy.uint8'>)[source]

Convert flow tensor to color image.

Parameters

  • tensor (tensor) –

  • tensor then (If) –

  • imtype (np.dtype) – Type of output image.

Returns

(numpy.ndarray or normalized torch image).

imaginaire.utils.visualization.common.tensor2im(image_tensor, imtype=<class 'numpy.uint8'>, normalize=True, three_channel_output=True)[source]

Convert tensor to image.

Parameters

  • image_tensor (torch.tensor or list of torch.tensor) – If tensor then (NxCxHxW) or (NxTxCxHxW) or (CxHxW).

  • imtype (np.dtype) – Type of output image.

  • normalize (bool) – Is the input image normalized or not? three_channel_output (bool): Should single channel images be made 3 channel in output?

Returns

(numpy.ndarray, list if case 1, 2 above).

imaginaire.utils.visualization.common.tensor2label(segmap, n_label=None, imtype=<class 'numpy.uint8'>, colorize=True, output_normalized_tensor=False)[source]

Convert segmentation mask tensor to color image. :param segmap: :type segmap: tensor :param If tensor then: :type If tensor then: NxCxHxW) or (NxTxCxHxW) or (CxHxW :param n_label: If None, then segmap.size(0). :type n_label: int :param imtype: Type of output image. :type imtype: np.dtype :param colorize: Put colors in. :type colorize: bool

Returns

(numpy.ndarray or normalized torch image).

imaginaire.utils.visualization.common.tensor2pilimage(image, width=None, height=None, minus1to1_normalized=False)[source]

Convert a 3 dimensional torch tensor to a PIL image with the desired width and height.

Parameters

  • image (3 x W1 x H1 tensor) – Image tensor

  • width (int) – Desired width for the result PIL image.

  • height (int) – Desired height for the result PIL image.

  • minus1to1_normalized (bool) – True if the tensor values are in [-1,

  • Otherwise, we assume the values are in [0, 1] (1]) –

Returns

The resulting PIL image.

Return type

(PIL image)

imaginaire.utils.visualization.face module

imaginaire.utils.visualization.face.add_face_keypoints(label_map, image, keypoints)[source]

Add additional keypoints to label map.

Parameters

  • label_map (Nx1xHxW tensor or None) –

  • image (Nx3xHxW tensor) –

  • keypoints (NxKx2 tensor) –

imaginaire.utils.visualization.face.connect_face_keypoints(resize_h, resize_w, crop_h, crop_w, original_h, original_w, is_flipped, cfgdata, keypoints)[source]

Connect the face keypoints to edges and draw the sketch.

Parameters

  • resize_h (int) – Height the input image was resized to.

  • resize_w (int) – Width the input image was resized to.

  • crop_h (int) – Height the input image was cropped.

  • crop_w (int) – Width the input image was cropped.

  • original_h (int) – Original height of the input image.

  • original_w (int) – Original width of the input image.

  • is_flipped (bool) – Is the input image flipped.

  • cfgdata (obj) – Data configuration.

  • keypoints (NxKx2 numpy array) – Facial landmarks (with K keypoints).

Returns

Drawn label map.

Return type

(list of HxWxC numpy array)

imaginaire.utils.visualization.face.convert_face_landmarks_to_image(cfgdata, landmarks, output_size, output_tensor=True, cpu_only=False)[source]

Convert the facial landmarks to a label map.

Parameters

  • cfgdata (obj) – Data configuration.

  • landmarks –

  • output_size (tuple of int) – H, W of output label map.

  • output_tensor (bool) – Output tensors instead of numpy arrays.

  • cpu_only (bool) – Output CPU tensor only.

Returns

Label maps.

Return type

(NxCxHxW tensor or list of HxWxC numpy arrays)

imaginaire.utils.visualization.face.draw_edge(im, x, y, bw=1, color=(255, 255, 255), draw_end_points=False)[source]

Set colors given a list of x and y coordinates for the edge.

Parameters

  • im (HxWxC numpy array) – Canvas to draw.

  • x (1D numpy array) – x coordinates of the edge.

  • y (1D numpy array) – y coordinates of the edge.

  • bw (int) – Width of the stroke.

  • color (list or tuple of int) – Color to draw.

  • draw_end_points (bool) – Whether to draw end points of the edge.

imaginaire.utils.visualization.face.func(x, a, b, c)[source]

Quadratic fitting function.

imaginaire.utils.visualization.face.get_126_landmarks_from_image(imgs, landmarks_network)[source]

Get face keypoints from an image.

Parameters

  • imgs (Nx3xHxW tensor or NxHxWx3 numpy array) –

  • images. (Input) –

  • landmarks_network (obj) – The landmark detection network.

Returns

Predicted landmarks.

Return type

(Nx126x2 numpy array)

imaginaire.utils.visualization.face.get_dlib_landmarks_from_image(imgs, predictor_path='shape_predictor_68_face_landmarks.dat')[source]

Get face keypoints from an image.

Parameters

  • imgs (N x 3 x H x W tensor or N x H x W x 3 numpy array) – Input images.

  • predictor_path (str) – Path to the predictor model.

imaginaire.utils.visualization.face.interp_points(x, y)[source]

Given the start and end points, interpolate to get a curve/line.

Parameters

  • x (1D array) – x coordinates of the points to interpolate.

  • y (1D array) – y coordinates of the points to interpolate.

Returns

  • curve_x (1D array): x coordinates of the interpolated points.

  • curve_y (1D array): y coordinates of the interpolated points.

Return type

(dict)

imaginaire.utils.visualization.face.linear(x, a, b)[source]

Linear fitting function.

imaginaire.utils.visualization.face.normalize_and_connect_face_keypoints(cfg, is_inference, data)[source]

Normalize face keypoints w.r.t. reference face keypoints and connect keypoints to form 2D images.

Parameters

  • cfg (obj) – Data configuration.

  • is_inference (bool) – Is doing inference or not.

  • data (dict) – Input data.

Returns

Output data.

Return type

(dict)

imaginaire.utils.visualization.face.normalize_face_keypoints(keypoints, ref_keypoints, dist_scales=None, momentum=0.9)[source]

Normalize face keypoints w.r.t. the reference face keypoints.

Parameters

  • keypoints (Kx2 numpy array) – Target facial keypoints to be normalized.

  • ref_keypoints (Kx2 numpy array) – Reference facial keypoints.

  • dist_scales (list of list of floats) – Normalization params.

  • momentum (float) – Temporal momentum for the normalization params.

Returns

Normalized facial keypoints.

Return type

(Kx2 numpy array)

imaginaire.utils.visualization.face.npy_to_tensor(keypoints)[source]

Convert numpy array to pytorch tensor.

imaginaire.utils.visualization.face.set_color(im, yy, xx, color)[source]

Set pixels of the image to the given color.

Parameters

  • im (HxWxC numpy array) – Canvas to draw.

  • xx (1D numpy array) – x coordinates of the pixels.

  • yy (1D numpy array) – y coordinates of the pixels.

  • color (list or tuple of int) – Color to draw.

imaginaire.utils.visualization.face.smooth_face_keypoints(concat_keypoints, ks)[source]

Temporally smooth the face keypoints by median filtering.

Parameters

  • concat_keypoints (TxKx2 numpy array) – Face keypoints to be filtered.

  • ks (int) – Filter kernel size.

Returns

Output face keypoints.

Return type

(1xKx2 numpy array)

imaginaire.utils.visualization.pose module

imaginaire.utils.visualization.pose.base_openpose_to_npy(inputs, return_largest_only=False)[source]

Convert OpenPose dicts to numpy arrays of keypoints.

Parameters

  • inputs (list of dicts) – List of OpenPose dicts.

  • return_largest_only (bool) – Whether to return only the largest person.

Returns

Keypoints.

Return type

(list of numpy arrays)

imaginaire.utils.visualization.pose.connect_pose_keypoints(pts, edge_lists, size, basic_points_only, remove_face_labels, random_drop_prob)[source]

Draw edges by connecting the keypoints onto the label map.

Parameters

  • pts (Px3 numpy array) – Keypoint xy coordinates + confidence.

  • edge_lists (nested list of ints) – List of keypoint indices for edges.

  • size (tuple of int) – Output size.

  • basic_points_only (bool) – Whether to use only the basic keypoints.

  • remove_face_labels (bool) – Whether to remove face labels.

  • random_drop_prob (float) – Probability to randomly drop keypoints.

Returns

Output label map.

Return type

(HxWxC numpy array)

imaginaire.utils.visualization.pose.define_edge_lists(basic_points_only)[source]

Define the list of keypoints that should be connected to form the edges.

Parameters

basic_points_only (bool) – Whether to use only the basic keypoints.

imaginaire.utils.visualization.pose.draw_edges(canvas, keypoints, edges_list, bw, use_one_hot, random_drop_prob=0, edge_len=2, colors=None, draw_end_points=False)[source]

Draw all the edges in the edge list on the canvas.

Parameters

  • canvas (HxWxK numpy array) – Canvas to draw.

  • keypoints (Px2 numpy array) – Keypoints.

  • edge_list (nested list of ints) – List of keypoint indices for edges.

  • bw (int) – Stroke width.

  • use_one_hot (bool) – Use one-hot encoding or not.

  • random_drop_prob (float) – Probability to randomly drop keypoints.

  • edge_len (int) – Number of keypoints in an edge.

  • colors (tuple of int) – Color to draw.

  • draw_end_points (bool) – Whether to draw end points for edges.

Returns

Output.

Return type

(HxWxK numpy array)

imaginaire.utils.visualization.pose.draw_openpose_npy(resize_h, resize_w, crop_h, crop_w, original_h, original_w, is_flipped, cfgdata, keypoints_npy)[source]

Connect the OpenPose keypoints to edges and draw the pose map.

Parameters

  • resize_h (int) – Height the input image was resized to.

  • resize_w (int) – Width the input image was resized to.

  • crop_h (int) – Height the input image was cropped.

  • crop_w (int) – Width the input image was cropped.

  • original_h (int) – Original height of the input image.

  • original_w (int) – Original width of the input image.

  • is_flipped (bool) – Is the input image flipped.

  • cfgdata (obj) – Data configuration.

  • keypoints_npy (dict) – OpenPose keypoint dict.

Returns

Drawn label map.

Return type

(list of HxWxC numpy array)

imaginaire.utils.visualization.pose.extract_valid_keypoints(pts, edge_lists)[source]

Use only the valid keypoints by looking at the detection confidences. If the confidences for all keypoints in an edge are above threshold, keep the keypoints. Otherwise, their coordinates will be set to zero.

Parameters

  • pts (Px3 numpy array) – Keypoint xy coordinates + confidence.

  • edge_lists (nested list of ints) – List of keypoint indices for edges.

Returns

Output keypoints.

Return type

(Px2 numpy array)

imaginaire.utils.visualization.pose.openpose_to_npy(inputs)[source]

Conver OpenPose dicts to numpy arrays of keypoints.

Parameters

inputs (list of dicts) – List of OpenPose dicts.

Returns

Keypoints.

Return type

(list of numpy arrays)

imaginaire.utils.visualization.pose.openpose_to_npy_largest_only(inputs)[source]

Convert OpenPose dicts to numpy arrays of keypoints. Only return the largest/tallest person in each dict.

Parameters

inputs (list of dicts) – List of OpenPose dicts.

Returns

Keypoints.

Return type

(list of numpy arrays)

imaginaire.utils.visualization.pose.tensor2pose(cfg, label_tensor)[source]

Convert output tensor to a numpy pose map.

Parameters

label_tensor (3D/4D/5D tensor) – Label tensor.

Returns

Pose map.

Return type

(HxWx3 numpy array or list of numpy arrays)

Module contents

imaginaire.utils.visualization.tensor2im(image_tensor, imtype=<class 'numpy.uint8'>, normalize=True, three_channel_output=True)[source]

Convert tensor to image.

Parameters

  • image_tensor (torch.tensor or list of torch.tensor) – If tensor then (NxCxHxW) or (NxTxCxHxW) or (CxHxW).

  • imtype (np.dtype) – Type of output image.

  • normalize (bool) – Is the input image normalized or not? three_channel_output (bool): Should single channel images be made 3 channel in output?

Returns

(numpy.ndarray, list if case 1, 2 above).

imaginaire.utils.visualization.tensor2flow(tensor, imtype=<class 'numpy.uint8'>)[source]

Convert flow tensor to color image.

Parameters

  • tensor (tensor) –

  • tensor then (If) –

  • imtype (np.dtype) – Type of output image.

Returns

(numpy.ndarray or normalized torch image).

imaginaire.utils.visualization.tensor2label(segmap, n_label=None, imtype=<class 'numpy.uint8'>, colorize=True, output_normalized_tensor=False)[source]

Convert segmentation mask tensor to color image. :param segmap: :type segmap: tensor :param If tensor then: :type If tensor then: NxCxHxW) or (NxTxCxHxW) or (CxHxW :param n_label: If None, then segmap.size(0). :type n_label: int :param imtype: Type of output image. :type imtype: np.dtype :param colorize: Put colors in. :type colorize: bool

Returns

(numpy.ndarray or normalized torch image).

imaginaire.utils.visualization.tensor2pilimage(image, width=None, height=None, minus1to1_normalized=False)[source]

Convert a 3 dimensional torch tensor to a PIL image with the desired width and height.

Parameters

  • image (3 x W1 x H1 tensor) – Image tensor

  • width (int) – Desired width for the result PIL image.

  • height (int) – Desired height for the result PIL image.

  • minus1to1_normalized (bool) – True if the tensor values are in [-1,

  • Otherwise, we assume the values are in [0, 1] (1]) –

Returns

The resulting PIL image.

Return type

(PIL image)

imaginaire.utils.visualization.save_tensor_image(filename, image, minus1to1_normalized=False)[source]

Convert a 3 dimensional torch tensor to a PIL image with the desired width and height.

Parameters

  • filename (str) – Image filename to be saved to.

  • image (3 x W1 x H1 tensor) – Image tensor

  • minus1to1_normalized (bool) – True if the tensor values are in [-1,

  • Otherwise, we assume the values are in [0, 1] (1]) –

Returns

The resulting PIL image.

Return type

(PIL image)