pocketpose.models.wholebody#

Submodules#

Package Contents#

Classes#

TFLiteModel

Interface for all TensorFlow Lite models.

BlazePose

Base class for the BlazePose models.

BlazePoseLite

BlazePose-Lite model.

BlazePoseFull

BlazePose-Full model.

BlazePoseHeavy

BlazePose-Heavy model.

Attributes#

class pocketpose.models.wholebody.TFLiteModel(model_path: str, model_url: str, **kwargs)#

Bases: pocketpose.models.interfaces.imodel.IModel

Interface for all TensorFlow Lite models.

We assume that the model has a single input, but it can have multiple outputs.

process_image(image)#

Default implementation of process_image() for models that don’t need preprocessing.

This method can be overridden by subclasses to implement model-specific preprocessing.

Args:
image (np.ndarray): The image to prepare for prediction. The image is a numpy

array with shape (1, height, width, channels) and dtype uint8 (range [0, 255]).

get_output(output_idx: int) numpy.ndarray#

Returns the output tensor of the model.

Args:

output_idx (int): The index of the output tensor to return.

Returns:

The output tensor as a numpy array.

predict(image: numpy.ndarray) Any#

Predicts the pose of the image.

Args:
image (np.ndarray): The image to predict the pose of. The image has

the shape and dtype expected by the model.

Returns:

The prediction returned by the model. This can be a single tensor or a tuple of tensors, depending on the model.

pocketpose.models.wholebody.model_registry#
class pocketpose.models.wholebody.BlazePose(model_path: str, model_url: str, input_size: tuple)#

Bases: pocketpose.models.interfaces.TFLiteModel

Base class for the BlazePose models.

NUM_KEYPOINTS = 33#
NUM_LANDMARKS = 39#
LANDMARKS_DIM = 5#
HEATMAPS_DIM = 39#
process_image(image)#

Default implementation of process_image() for models that don’t need preprocessing.

This method can be overridden by subclasses to implement model-specific preprocessing.

Args:
image (np.ndarray): The image to prepare for prediction. The image is a numpy

array with shape (1, height, width, channels) and dtype uint8 (range [0, 255]).

_calculate_keypoints(landmark_points, heatmap, index, original_size)#
postprocess_prediction(prediction, original_size)#

Postprocess the prediction.

Args:

prediction (list): List of outputs from the model. original_size (tuple): Original size of the image as (height, width).

class pocketpose.models.wholebody.BlazePoseLite#

Bases: BlazePose

BlazePose-Lite model.

class pocketpose.models.wholebody.BlazePoseFull#

Bases: BlazePose

BlazePose-Full model.

class pocketpose.models.wholebody.BlazePoseHeavy#

Bases: BlazePose

BlazePose-Heavy model.