data_juicer.ops.common.hawor_func module#

data_juicer.ops.common.hawor_func.default_config() CfgNode[source]#

Get a yacs CfgNode object with the default config values.

data_juicer.ops.common.hawor_func.get_config(config_file: str, merge: bool = True, update_cachedir: bool = False) CfgNode[source]#

Read a config file and optionally merge it with the default config file. :param config_file: Path to config file. :type config_file: str :param merge: Whether to merge with the default config or not. :type merge: bool

Returns:

Config as a yacs CfgNode object.

Return type:

CfgNode

data_juicer.ops.common.hawor_func.create_backbone(cfg)[source]#
data_juicer.ops.common.hawor_func.normalization_layer(norm: str | None, dim: int, norm_cond_dim: int = -1)[source]#
class data_juicer.ops.common.hawor_func.AdaptiveLayerNorm1D(data_dim: int, norm_cond_dim: int)[source]#

Bases: Module

__init__(data_dim: int, norm_cond_dim: int)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, t: Tensor) Tensor[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.FrequencyEmbedder(num_frequencies, max_freq_log2)[source]#

Bases: Module

__init__(num_frequencies, max_freq_log2)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

data_juicer.ops.common.hawor_func.get_transform(center, scale, res, rot=0)[source]#

Generate transformation matrix.

data_juicer.ops.common.hawor_func.transform(pt, center, scale, res, invert=0, rot=0, asint=True)[source]#

Transform pixel location to different reference.

data_juicer.ops.common.hawor_func.crop(img, center, scale, res, rot=0)[source]#

Crop image according to the supplied bounding box.

data_juicer.ops.common.hawor_func.boxes_2_cs(boxes)[source]#
data_juicer.ops.common.hawor_func.exists(val)[source]#
data_juicer.ops.common.hawor_func.default(val, d)[source]#
data_juicer.ops.common.hawor_func.perspective_projection(points, rotation, translation, focal_length, camera_center, distortion=None)[source]#

This function computes the perspective projection of a set of points. Input:

points (bs, N, 3): 3D points rotation (bs, 3, 3): Camera rotation translation (bs, 3): Camera translation focal_length (bs,) or scalar: Focal length camera_center (bs, 2): Camera center

data_juicer.ops.common.hawor_func.avg_rot(rot)[source]#
data_juicer.ops.common.hawor_func.rot9d_to_rotmat(x)[source]#

Convert 9D rotation representation to 3x3 rotation matrix. Based on Levinson et al., “An Analysis of SVD for Deep Rotation Estimation” Input:

(B,9) or (B,J*9) Batch of 9D rotation (interpreted as 3x3 est rotmat)

Output:

(B,3,3) or (B*J,3,3) Batch of corresponding rotation matrices

data_juicer.ops.common.hawor_func.quat_to_rotmat(quat)[source]#

Convert quaternion coefficients to rotation matrix. :param quat: size = [B, 4] 4 <===>(w, x, y, z)

Returns:

Rotation matrix corresponding to the quaternion – size = [B, 3, 3]

data_juicer.ops.common.hawor_func.rot6d_to_rotmat(x: Tensor) Tensor[source]#

Convert 6D rotation representation to 3x3 rotation matrix. Based on Zhou et al., “On the Continuity of Rotation Representations in Neural Networks”, CVPR 2019 :param x: (B,6) Batch of 6-D rotation representations. :type x: torch.Tensor

Returns:

Batch of corresponding rotation matrices with shape (B,3,3).

Return type:

torch.Tensor

data_juicer.ops.common.hawor_func.rotmat_to_rot6d(rotmat)[source]#

Inverse function of the above. Input:

(B,3,3) Batch of corresponding rotation matrices

Output:

(B,6) Batch of 6-D rotation representations

data_juicer.ops.common.hawor_func.rotation_matrix_to_angle_axis(rotation_matrix)[source]#

This function is borrowed from kornia/kornia

Convert 3x4 rotation matrix to Rodrigues vector

Parameters:

rotation_matrix (Tensor) – rotation matrix.

Returns:

Rodrigues vector transformation.

Return type:

Tensor

Shape:
  • Input: \((N, 3, 4)\)

  • Output: \((N, 3)\)

Example

>>> input = torch.rand(2, 3, 4)  # Nx4x4
>>> output = tgm.rotation_matrix_to_angle_axis(input)  # Nx3
data_juicer.ops.common.hawor_func.quaternion_to_angle_axis(quaternion: Tensor) Tensor[source]#

This function is borrowed from kornia/kornia

Convert quaternion vector to angle axis of rotation.

Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

Parameters:

quaternion (torch.Tensor) – tensor with quaternions.

Returns:

tensor with angle axis of rotation.

Return type:

torch.Tensor

Shape:
  • Input: \((*, 4)\) where * means, any number of dimensions

  • Output: \((*, 3)\)

Example

>>> quaternion = torch.rand(2, 4)  # Nx4
>>> angle_axis = tgm.quaternion_to_angle_axis(quaternion)  # Nx3
data_juicer.ops.common.hawor_func.rotation_matrix_to_quaternion(rotation_matrix, eps=1e-06)[source]#

This function is borrowed from kornia/kornia Convert rotation matrix to 4d quaternion vector This algorithm is based on algorithm described in KieranWynn/pyquaternion

:param rotation_matrix (N, 3, 3)

data_juicer.ops.common.hawor_func.estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000.0, img_size=224.0)[source]#

This function is borrowed from nkolot/SPIN

Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. Input:

S: (25, 3) 3D joint locations joints: (25, 3) 2D joint locations and confidence

Returns:

(3,) camera translation vector

data_juicer.ops.common.hawor_func.estimate_translation(S, joints_2d, focal_length=5000.0, img_size=224.0)[source]#

Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. Input:

S: (B, 49, 3) 3D joint locations joints: (B, 49, 3) 2D joint locations and confidence

Returns:

(B, 3) camera translation vectors

data_juicer.ops.common.hawor_func.get_keypoints_rectangle(keypoints: array, threshold: float) Tuple[float, float, float][source]#

Compute rectangle enclosing keypoints above the threshold. :param keypoints: Keypoint array of shape (N, 3). :type keypoints: np.array :param threshold: Confidence visualization threshold. :type threshold: float

Returns:

Rectangle width, height and area.

Return type:

Tuple[float, float, float]

data_juicer.ops.common.hawor_func.render_keypoints(img: array, keypoints: array, pairs: List, colors: List, thickness_circle_ratio: float, thickness_line_ratio_wrt_circle: float, pose_scales: List, threshold: float = 0.1, alpha: float = 1.0) array[source]#

Render keypoints on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param keypoints: Keypoint array of shape (N, 3). :type keypoints: np.array :param pairs: List of keypoint pairs per limb. :type pairs: List :param colors: (List): List of colors per keypoint. :param thickness_circle_ratio: Circle thickness ratio. :type thickness_circle_ratio: float :param thickness_line_ratio_wrt_circle: Line thickness ratio wrt the circle. :type thickness_line_ratio_wrt_circle: float :param pose_scales: List of pose scales. :type pose_scales: List :param threshold: Only visualize keypoints with confidence above the threshold. :type threshold: float

Returns:

Image of shape (H, W, 3) with keypoints drawn on top of the original image.

Return type:

(np.array)

data_juicer.ops.common.hawor_func.render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=<function <lambda>>, alpha=1.0)[source]#
data_juicer.ops.common.hawor_func.render_hand_landmarks(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=<function <lambda>>, alpha=1.0)[source]#
data_juicer.ops.common.hawor_func.render_body_keypoints(img: array, body_keypoints: array) array[source]#

Render OpenPose body keypoints on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param body_keypoints: Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). :type body_keypoints: np.array

Returns:

Image of shape (H, W, 3) with keypoints drawn on top of the original image.

Return type:

(np.array)

data_juicer.ops.common.hawor_func.render_openpose(img: array, hand_keypoints: array) array[source]#

Render keypoints in the OpenPose format on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param body_keypoints: Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). :type body_keypoints: np.array

Returns:

Image of shape (H, W, 3) with keypoints drawn on top of the original image.

Return type:

(np.array)

data_juicer.ops.common.hawor_func.render_openpose_landmarks(img: array, hand_keypoints: array) array[source]#

Render keypoints in the OpenPose format on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param body_keypoints: Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). :type body_keypoints: np.array

Returns:

Image of shape (H, W, 3) with keypoints drawn on top of the original image.

Return type:

(np.array)

data_juicer.ops.common.hawor_func.get_pylogger(name='data_juicer.ops.common.hawor_func') Logger[source]#

Initializes multi-GPU-friendly python command line logger.

class data_juicer.ops.common.hawor_func.PreNorm(dim: int, fn: Callable, norm: str = 'layer', norm_cond_dim: int = -1)[source]#

Bases: Module

__init__(dim: int, fn: Callable, norm: str = 'layer', norm_cond_dim: int = -1)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, *args, **kwargs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.FeedForward(dim, hidden_dim, dropout=0.0)[source]#

Bases: Module

__init__(dim, hidden_dim, dropout=0.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.Attention(dim, heads=8, dim_head=64, dropout=0.0)[source]#

Bases: Module

__init__(dim, heads=8, dim_head=64, dropout=0.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.CrossAttention(dim, context_dim=None, heads=8, dim_head=64, dropout=0.0)[source]#

Bases: Module

__init__(dim, context_dim=None, heads=8, dim_head=64, dropout=0.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, context=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.Transformer(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1)[source]#

Bases: Module

__init__(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, *args)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.TransformerCrossAttn(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None)[source]#

Bases: Module

__init__(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, *args, context=None, context_list=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.DropTokenDropout(p: float = 0.1)[source]#

Bases: Module

__init__(p: float = 0.1)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.ZeroTokenDropout(p: float = 0.1)[source]#

Bases: Module

__init__(p: float = 0.1)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.TransformerEncoder(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', emb_dropout_loc: str = 'token', norm: str = 'layer', norm_cond_dim: int = -1, token_pe_numfreq: int = -1)[source]#

Bases: Module

__init__(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', emb_dropout_loc: str = 'token', norm: str = 'layer', norm_cond_dim: int = -1, token_pe_numfreq: int = -1)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inp: Tensor, *args, **kwargs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.TransformerDecoder(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None, skip_token_embedding: bool = False)[source]#

Bases: Module

__init__(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None, skip_token_embedding: bool = False)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inp: Tensor, *args, context=None, context_list=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.MANOTransformerDecoderHead(cfg)[source]#

Bases: Module

HMR2 Cross-attention based SMPL Transformer decoder

__init__(cfg)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, **kwargs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.temporal_attention(in_dim=1280, out_dim=1280, hdim=512, nlayer=6, nhead=4, residual=False)[source]#

Bases: Module

__init__(in_dim=1280, out_dim=1280, hdim=512, nlayer=6, nhead=4, residual=False)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.PositionalEncoding(d_model, dropout=0.1, max_len=100)[source]#

Bases: Module

__init__(d_model, dropout=0.1, max_len=100)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class data_juicer.ops.common.hawor_func.TrackDatasetEval(imgfiles, boxes, crop_size=256, dilate=1.0, img_focal=None, img_center=None, normalization=True, item_idx=0, do_flip=False)[source]#

Bases: Dataset

Track Dataset Class - Load images/crops of the tracked boxes.

__init__(imgfiles, boxes, crop_size=256, dilate=1.0, img_focal=None, img_center=None, normalization=True, item_idx=0, do_flip=False)[source]#
data_juicer.ops.common.hawor_func.dataset_config() CfgNode[source]#

Get dataset config file :returns: Dataset config as a yacs CfgNode object. :rtype: CfgNode

data_juicer.ops.common.hawor_func.parse_chunks(frame, boxes, min_len=16)[source]#

If a track disappear in the middle, we separate it to different segments to estimate the HPS independently. If a segment is less than 16 frames, we get rid of it for now.

data_juicer.ops.common.hawor_func.interpolate_bboxes(bboxes)[source]#
data_juicer.ops.common.hawor_func.aa_to_rotmat(theta: Tensor)[source]#

Convert axis-angle representation to rotation matrix. Works by first converting it to a quaternion. :param theta: Tensor of shape (B, 3) containing axis-angle representations. :type theta: torch.Tensor

Returns:

Corresponding rotation matrices with shape (B, 3, 3).

Return type:

torch.Tensor

class data_juicer.ops.common.hawor_func.MANO(*args, joint_regressor_extra: str | None = None, **kwargs)[source]#

Bases: MANOLayer

__init__(*args, joint_regressor_extra: str | None = None, **kwargs)[source]#

Extension of the official MANO implementation to support more joints. :param Same as MANOLayer.: :param joint_regressor_extra: Path to extra joint regressor. :type joint_regressor_extra: str

forward(*args, **kwargs) MANOOutput[source]#

Run forward pass. Same as MANO and also append an extra set of joints if joint_regressor_extra is specified.

query(hmr_output)[source]#
data_juicer.ops.common.hawor_func.block_print()[source]#
data_juicer.ops.common.hawor_func.enable_print()[source]#
data_juicer.ops.common.hawor_func.run_mano(trans, root_orient, hand_pose, is_right=None, betas=None, use_cuda=True)[source]#

Forward pass of the SMPL model and populates pred_data accordingly with joints3d, verts3d, points3d.

trans : B x T x 3 root_orient : B x T x 3 body_pose : B x T x J*3 betas : (optional) B x D

data_juicer.ops.common.hawor_func.run_mano_left(trans, root_orient, hand_pose, is_right=None, betas=None, use_cuda=True, fix_shapedirs=True)[source]#

Forward pass of the SMPL model and populates pred_data accordingly with joints3d, verts3d, points3d.

trans : B x T x 3 root_orient : B x T x 3 body_pose : B x T x J*3 betas : (optional) B x D

data_juicer.ops.common.hawor_func.run_mano_twohands(init_trans, init_rot, init_hand_pose, is_right, init_betas, use_cuda=True, fix_shapedirs=True)[source]#
data_juicer.ops.common.hawor_func.quaternion_mul(q0, q1)[source]#

EXPECTS WXYZ :param q0 (, 4) :param q1 (, 4)

data_juicer.ops.common.hawor_func.quaternion_inverse(q, eps=1e-08)[source]#

EXPECTS WXYZ :param q (*, 4)

data_juicer.ops.common.hawor_func.quaternion_slerp(t, q0, q1, eps=1e-08)[source]#

:param t (, 1) must be between 0 and 1 :param q0 (, 4) :param q1 (*, 4)

data_juicer.ops.common.hawor_func.angle_axis_to_rotation_matrix(angle_axis)[source]#

:param angle_axis (, 3) return (, 3, 3)

data_juicer.ops.common.hawor_func.quaternion_to_rotation_matrix(quaternion)[source]#

Convert a quaternion to a rotation matrix. Taken from kornia/kornia, based on matthew-brett/transforms3d tensorflow/graphics :param quaternion (N, 4) expects WXYZ order returns rotation matrix (N, 3, 3)

data_juicer.ops.common.hawor_func.angle_axis_to_quaternion(angle_axis)[source]#

This function is borrowed from kornia/kornia Convert angle axis to quaternion in WXYZ order :param angle_axis (, 3) :returns quaternion (, 4) WXYZ order

class data_juicer.ops.common.hawor_func.HAWOR(cfg: CfgNode)[source]#

Bases: LightningModule

__init__(cfg: CfgNode)[source]#

Setup HAWOR model :param cfg: Config file as a yacs CfgNode :type cfg: CfgNode

get_parameters()[source]#
configure_optimizers() Optimizer[source]#

Setup model and distriminator Optimizers :returns: Model and discriminator optimizers :rtype: Tuple[torch.optim.Optimizer, torch.optim.Optimizer]

forward_step(batch: Dict, train: bool = False) Dict[source]#

Run a forward step of the network :param batch: Dictionary containing batch data :type batch: Dict :param train: Flag indicating whether it is training or validation mode :type train: bool

Returns:

Dictionary containing the regression output

Return type:

Dict

compute_loss(batch: Dict, output: Dict, train: bool = True) Tensor[source]#

Compute losses given the input batch and the regression output :param batch: Dictionary containing batch data :type batch: Dict :param output: Dictionary containing the regression output :type output: Dict :param train: Flag indicating whether it is training or validation mode :type train: bool

Returns:

Total loss for current batch

Return type:

torch.Tensor

tensorboard_logging(batch: Dict, output: Dict, step_count: int, train: bool = True, write_to_summary_writer: bool = True, render_log: bool = True) None[source]#

Log results to Tensorboard :param batch: Dictionary containing batch data :type batch: Dict :param output: Dictionary containing the regression output :type output: Dict :param step_count: Global training step count :type step_count: int :param train: Flag indicating whether it is training or validation mode :type train: bool

forward(batch: Dict) Dict[source]#

Run a forward step of the network in val mode :param batch: Dictionary containing batch data :type batch: Dict

Returns:

Dictionary containing the regression output

Return type:

Dict

training_step(joint_batch: Dict, batch_idx: int) Dict[source]#

Run a full training step :param joint_batch: Dictionary containing image and mocap batch data :type joint_batch: Dict :param batch_idx: Unused. :type batch_idx: int :param batch_idx: Unused. :type batch_idx: torch.Tensor

Returns:

Dictionary containing regression output.

Return type:

Dict

inference(imgfiles, boxes, img_focal, img_center, device='cuda', do_flip=False)[source]#
validation_step(batch: Dict, batch_idx: int, dataloader_idx=0) Dict[source]#

Run a validation step and log to Tensorboard :param batch: Dictionary containing batch data :type batch: Dict :param batch_idx: Unused. :type batch_idx: int

Returns:

Dictionary containing regression output.

Return type:

Dict

visualize_tensorboard(images, pred_keypoints, gt_project_j2d, gt_keypoints)[source]#
project(points, pred_cam, center, scale, img_focal, img_center, return_full=False)[source]#
get_trans(pred_cam, center, scale, img_focal, img_center)[source]#
bbox_est(center, scale, img_focal, img_center)[source]#