Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import tensorflow as tf | |
| import tensorflow_hub as hub | |
| import os | |
| class PoseDetection: | |
| """Base class for pose detection in images using various algorithm such | |
| as Movenet | |
| Warning: This class should not be used directly. | |
| Use derived classes instead. | |
| Parameters: | |
| model_name : name of the pose detection method | |
| input_size : image size of input required for model | |
| model_path : path to pose detection model | |
| """ | |
| def __init__(self, model_name=None, input_size=None, model_path=None): | |
| self.model_name = model_name | |
| self.input_size = input_size | |
| self.model_path = model_path | |
| self.load_model() | |
| def load_model(self): | |
| """Absrtact method to loads pose detection model. | |
| """ | |
| raise NotImplementedError | |
| def preprocess_image(self, frame): | |
| """Absrtact method to preprocess a image before running pose | |
| detection inference on it. | |
| """ | |
| raise NotImplementedError | |
| def run_inference(self, frames): | |
| """Absrtact method to run pose detection inference on it. | |
| """ | |
| raise NotImplementedError | |
| class MovenetPoseDetection(PoseDetection): | |
| """Class for pose detection in images using Movenet model | |
| Parameters: | |
| model_name : name of the pose detection method | |
| input_size : image size of input required for model | |
| model_path : path to pose detection model | |
| """ | |
| def __init__(self, model_name=None, input_size=None, model_path=None): | |
| model_name = model_name or 'movenet' | |
| input_size = input_size or 256 | |
| model_path = model_path or 'models/movenet/movenet_tf/1' | |
| super().__init__(model_name, input_size, model_path) | |
| def load_model(self): | |
| """Loads the pose detection model. | |
| """ | |
| if self.model_name == "movenet": | |
| module = hub.load(self.model_path) | |
| self.model = module.signatures['serving_default'] | |
| def preprocess_image(self, frame): | |
| """Preprocesses an image to transform it into required format for | |
| pose detection | |
| Args: | |
| frame: A numpy array representing the input image | |
| Returns: | |
| A tensor of 'int32' data type and resized input image. | |
| """ | |
| input_image = tf.expand_dims(frame, axis=0) | |
| input_image = tf.image.resize_with_pad( | |
| input_image, self.input_size, self.input_size | |
| ) | |
| input_image = tf.cast(input_image, dtype=tf.int32) | |
| return input_image | |
| def run_inference(self, frames): | |
| """Appllies pose dection model on a frame | |
| Args: | |
| frames: A list of numpy arrays representing the input images | |
| Returns: | |
| A [n, 1, 17, 3] float numpy array representing the keypoint | |
| coordinates | |
| and scores predicted by Movenet Model of list of images | |
| """ | |
| keypoints = np.zeros((len(frames), 17, 2)) | |
| for i, frame in enumerate(frames): | |
| frame = self.preprocess_image(frame) | |
| keypoints[i, :, :] = self.run_model(frame) | |
| return keypoints | |
| def run_model(self, input_image): | |
| """Runs detection on an input image. | |
| Args: | |
| input_image: A [1, height, width, 3] tensor represents the input image | |
| pixels. Note that the height/width should already be resized and | |
| match the | |
| expected input resolution of the model before passing into this | |
| function. | |
| Returns: | |
| A [1, 1, 17, 3] float numpy array representing the keypoint | |
| coordinates | |
| and scores predicted by Movenet Model | |
| """ | |
| outputs = self.model(input_image) | |
| outputs = outputs['output_0'].numpy() | |
| outputs[:, :, :, [0, 1]] = outputs[:, :, :, [1, 0]] | |
| return outputs[0, 0, :, :2] | |