|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import base64 |
|
|
import os |
|
|
|
|
|
import mmcv |
|
|
import torch |
|
|
|
|
|
from mmpose.apis import (inference_bottom_up_pose_model, |
|
|
inference_top_down_pose_model, init_pose_model) |
|
|
from mmpose.models.detectors import AssociativeEmbedding, TopDown |
|
|
|
|
|
try: |
|
|
from ts.torch_handler.base_handler import BaseHandler |
|
|
except ImportError: |
|
|
raise ImportError('Please install torchserve.') |
|
|
|
|
|
|
|
|
class MMPoseHandler(BaseHandler): |
|
|
|
|
|
def initialize(self, context): |
|
|
properties = context.system_properties |
|
|
self.map_location = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
self.device = torch.device(self.map_location + ':' + |
|
|
str(properties.get('gpu_id')) if torch.cuda. |
|
|
is_available() else self.map_location) |
|
|
self.manifest = context.manifest |
|
|
|
|
|
model_dir = properties.get('model_dir') |
|
|
serialized_file = self.manifest['model']['serializedFile'] |
|
|
checkpoint = os.path.join(model_dir, serialized_file) |
|
|
self.config_file = os.path.join(model_dir, 'config.py') |
|
|
|
|
|
self.model = init_pose_model(self.config_file, checkpoint, self.device) |
|
|
self.initialized = True |
|
|
|
|
|
def preprocess(self, data): |
|
|
images = [] |
|
|
|
|
|
for row in data: |
|
|
image = row.get('data') or row.get('body') |
|
|
if isinstance(image, str): |
|
|
image = base64.b64decode(image) |
|
|
image = mmcv.imfrombytes(image) |
|
|
images.append(image) |
|
|
|
|
|
return images |
|
|
|
|
|
def inference(self, data, *args, **kwargs): |
|
|
if isinstance(self.model, TopDown): |
|
|
results = self._inference_top_down_pose_model(data) |
|
|
elif isinstance(self.model, (AssociativeEmbedding, )): |
|
|
results = self._inference_bottom_up_pose_model(data) |
|
|
else: |
|
|
raise NotImplementedError( |
|
|
f'Model type {type(self.model)} is not supported.') |
|
|
|
|
|
return results |
|
|
|
|
|
def _inference_top_down_pose_model(self, data): |
|
|
results = [] |
|
|
for image in data: |
|
|
|
|
|
preds, _ = inference_top_down_pose_model( |
|
|
self.model, image, person_results=None) |
|
|
results.append(preds) |
|
|
return results |
|
|
|
|
|
def _inference_bottom_up_pose_model(self, data): |
|
|
results = [] |
|
|
for image in data: |
|
|
preds, _ = inference_bottom_up_pose_model(self.model, image) |
|
|
results.append(preds) |
|
|
return results |
|
|
|
|
|
def postprocess(self, data): |
|
|
output = [[{ |
|
|
'keypoints': pred['keypoints'].tolist() |
|
|
} for pred in preds] for preds in data] |
|
|
|
|
|
return output |
|
|
|
