unipose.datasets.utils

module documentation

Undocumented

Function	`apply_gaussian`	Undocumented
Function	`crop_image`	Crop the image to the bounding box and resize it to the given size.
Function	`crop_keypoints`	Crop the keypoints to the bounding box.
Function	`enlarge_bounding_box`	Enlarge the bounding box to the smallest square that contains it without changing its center.
Function	`gaussian`	Undocumented
Function	`get_keypoints_images`	Generate the images of the keypoints.
Function	`process_batch`	Process a batch of data.

def apply_gaussian(image, mean_x, mean_y, sigma): ¶

Undocumented

def crop_image(image, bbox_rounded, image_size): ¶

Crop the image to the bounding box and resize it to the given size.

Parameters
image:`torch.Tensor`	A tensor of shape (batchsize, channels, height, width) containing the image.
bbox_rounded:`torch.Tensor`	A tensor of shape (batchsize, 4) containing the bounding box in the format [x, y, w, h].
image_size:`int`	The size to which the image should be resized.
Returns
`torch.Tensor`	A tensor of shape (batchsize, channels, image_size, image_size) containing the cropped and resized image.

def crop_keypoints(keypoints, bbox_rounded, mask): ¶

Crop the keypoints to the bounding box.

Parameters
keypoints:`torch.Tensor`	A tensor of shape (batchsize, num_keypoints, 2) containing the keypoints.
bbox_rounded:`torch.Tensor`	A tensor of shape (batchsize, 4) containing the bounding box in the format [x, y, w, h].
mask:`torch.Tensor`	A tensor of shape (batchsize, num_keypoints) containing the mask.
Returns
`torch.Tensor`	A tensor of shape (batchsize, num_keypoints, 3) containing the cropped keypoints.

def enlarge_bounding_box(bbox): ¶

Enlarge the bounding box to the smallest square that contains it without changing its center.

Parameters
bbox:`torch.Tensor`	A tensor of shape (4) containing the bounding box in the format [x, y, w, h].
Returns
`torch.Tensor`	A tensor of shape (4) containing the enlarged bounding box in the format [x1, y1, w1, h1].

def gaussian(x, y, mean_x, mean_y, sigma): ¶

Undocumented

def get_keypoints_images(keypoints, image_size, mask, sigma): ¶

Generate the images of the keypoints.

Parameters
keypoints:`torch.Tensor`	A tensor of shape (batchsize, num_keypoints, 2) containing the keypoints.
image_size:`int`	The size of the image.
mask:`torch.Tensor`	A tensor of shape (batchsize, num_keypoints) containing the mask.
sigma:`float`	The sigma of the gaussian.
Returns
`torch.Tensor`	A tensor of shape (batchsize, num_keypoints, image_size, image_size) containing the images of the keypoints.

def process_batch(image, bbox, keypoints, label_mask, image_size, scale_factor): ¶

Process a batch of data.

Parameters
image:`torch.Tensor`	A tensor of shape (batchsize, channels, height, width) containing the image.
bbox:`torch.Tensor`	A tensor of shape (batchsize, 4) containing the bounding box in the format [x, y, w, h].
keypoints:`torch.Tensor`	A tensor of shape (batchsize, num_keypoints, 2) containing the keypoints.
label_mask:`torch.Tensor`	A tensor of shape (batchsize, num_keypoints) containing the mask.
image_size:`int`	The size to which the image should be resized.
scale_factor:`int`	Patch size of output keypoint heatmaps. Must be 4 the default Unipose model is used without modification.
Returns
`List[torch.Tensor]`	A list of tensors containing the cropped and resized image and the images of the keypoints.