Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,135 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import os.path as osp
|
||||
import warnings
|
||||
from argparse import ArgumentParser, Namespace
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import mmcv
|
||||
import torch
|
||||
from mmcv.runner import CheckpointLoader
|
||||
|
||||
try:
|
||||
from model_archiver.model_packaging import package_model
|
||||
from model_archiver.model_packaging_utils import ModelExportUtils
|
||||
except ImportError:
|
||||
package_model = None
|
||||
|
||||
|
||||
def mmpose2torchserve(config_file: str,
|
||||
checkpoint_file: str,
|
||||
output_folder: str,
|
||||
model_name: str,
|
||||
model_version: str = '1.0',
|
||||
force: bool = False):
|
||||
"""Converts MMPose model (config + checkpoint) to TorchServe `.mar`.
|
||||
|
||||
Args:
|
||||
config_file:
|
||||
In MMPose config format.
|
||||
The contents vary for each task repository.
|
||||
checkpoint_file:
|
||||
In MMPose checkpoint format.
|
||||
The contents vary for each task repository.
|
||||
output_folder:
|
||||
Folder where `{model_name}.mar` will be created.
|
||||
The file created will be in TorchServe archive format.
|
||||
model_name:
|
||||
If not None, used for naming the `{model_name}.mar` file
|
||||
that will be created under `output_folder`.
|
||||
If None, `{Path(checkpoint_file).stem}` will be used.
|
||||
model_version:
|
||||
Model's version.
|
||||
force:
|
||||
If True, if there is an existing `{model_name}.mar`
|
||||
file under `output_folder` it will be overwritten.
|
||||
"""
|
||||
|
||||
mmcv.mkdir_or_exist(output_folder)
|
||||
|
||||
config = mmcv.Config.fromfile(config_file)
|
||||
|
||||
with TemporaryDirectory() as tmpdir:
|
||||
model_file = osp.join(tmpdir, 'config.py')
|
||||
config.dump(model_file)
|
||||
handler_path = osp.join(osp.dirname(__file__), 'mmpose_handler.py')
|
||||
model_name = model_name or osp.splitext(
|
||||
osp.basename(checkpoint_file))[0]
|
||||
|
||||
# use mmcv CheckpointLoader if checkpoint is not from a local file
|
||||
if not osp.isfile(checkpoint_file):
|
||||
ckpt = CheckpointLoader.load_checkpoint(checkpoint_file)
|
||||
checkpoint_file = osp.join(tmpdir, 'checkpoint.pth')
|
||||
with open(checkpoint_file, 'wb') as f:
|
||||
torch.save(ckpt, f)
|
||||
|
||||
args = Namespace(
|
||||
**{
|
||||
'model_file': model_file,
|
||||
'serialized_file': checkpoint_file,
|
||||
'handler': handler_path,
|
||||
'model_name': model_name,
|
||||
'version': model_version,
|
||||
'export_path': output_folder,
|
||||
'force': force,
|
||||
'requirements_file': None,
|
||||
'extra_files': None,
|
||||
'runtime': 'python',
|
||||
'archive_format': 'default'
|
||||
})
|
||||
manifest = ModelExportUtils.generate_manifest_json(args)
|
||||
package_model(args, manifest)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = ArgumentParser(
|
||||
description='Convert MMPose models to TorchServe `.mar` format.')
|
||||
parser.add_argument('config', type=str, help='config file path')
|
||||
parser.add_argument('checkpoint', type=str, help='checkpoint file path')
|
||||
parser.add_argument(
|
||||
'--output-folder',
|
||||
type=str,
|
||||
required=True,
|
||||
help='Folder where `{model_name}.mar` will be created.')
|
||||
parser.add_argument(
|
||||
'--model-name',
|
||||
type=str,
|
||||
default=None,
|
||||
help='If not None, used for naming the `{model_name}.mar`'
|
||||
'file that will be created under `output_folder`.'
|
||||
'If None, `{Path(checkpoint_file).stem}` will be used.')
|
||||
parser.add_argument(
|
||||
'--model-version',
|
||||
type=str,
|
||||
default='1.0',
|
||||
help='Number used for versioning.')
|
||||
parser.add_argument(
|
||||
'-f',
|
||||
'--force',
|
||||
action='store_true',
|
||||
help='overwrite the existing `{model_name}.mar`')
|
||||
args = parser.parse_args()
|
||||
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
|
||||
# Following strings of text style are from colorama package
|
||||
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
|
||||
red_text, blue_text = '\x1b[31m', '\x1b[34m'
|
||||
white_background = '\x1b[107m'
|
||||
|
||||
msg = white_background + bright_style + red_text
|
||||
msg += 'DeprecationWarning: This tool will be deprecated in future. '
|
||||
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
|
||||
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
|
||||
msg += reset_style
|
||||
warnings.warn(msg)
|
||||
|
||||
if package_model is None:
|
||||
raise ImportError('`torch-model-archiver` is required.'
|
||||
'Try: pip install torch-model-archiver')
|
||||
|
||||
mmpose2torchserve(args.config, args.checkpoint, args.output_folder,
|
||||
args.model_name, args.model_version, args.force)
|
||||
@@ -0,0 +1,80 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
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:
|
||||
# use dummy person bounding box
|
||||
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
|
||||
@@ -0,0 +1,165 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from mmpose.apis import init_pose_model
|
||||
|
||||
try:
|
||||
import onnx
|
||||
import onnxruntime as rt
|
||||
except ImportError as e:
|
||||
raise ImportError(f'Please install onnx and onnxruntime first. {e}')
|
||||
|
||||
try:
|
||||
from mmcv.onnx.symbolic import register_extra_symbolics
|
||||
except ModuleNotFoundError:
|
||||
raise NotImplementedError('please update mmcv to version>=1.0.4')
|
||||
|
||||
|
||||
def _convert_batchnorm(module):
|
||||
"""Convert the syncBNs into normal BN3ds."""
|
||||
module_output = module
|
||||
if isinstance(module, torch.nn.SyncBatchNorm):
|
||||
module_output = torch.nn.BatchNorm3d(module.num_features, module.eps,
|
||||
module.momentum, module.affine,
|
||||
module.track_running_stats)
|
||||
if module.affine:
|
||||
module_output.weight.data = module.weight.data.clone().detach()
|
||||
module_output.bias.data = module.bias.data.clone().detach()
|
||||
# keep requires_grad unchanged
|
||||
module_output.weight.requires_grad = module.weight.requires_grad
|
||||
module_output.bias.requires_grad = module.bias.requires_grad
|
||||
module_output.running_mean = module.running_mean
|
||||
module_output.running_var = module.running_var
|
||||
module_output.num_batches_tracked = module.num_batches_tracked
|
||||
for name, child in module.named_children():
|
||||
module_output.add_module(name, _convert_batchnorm(child))
|
||||
del module
|
||||
return module_output
|
||||
|
||||
|
||||
def pytorch2onnx(model,
|
||||
input_shape,
|
||||
opset_version=11,
|
||||
show=False,
|
||||
output_file='tmp.onnx',
|
||||
verify=False):
|
||||
"""Convert pytorch model to onnx model.
|
||||
|
||||
Args:
|
||||
model (:obj:`nn.Module`): The pytorch model to be exported.
|
||||
input_shape (tuple[int]): The input tensor shape of the model.
|
||||
opset_version (int): Opset version of onnx used. Default: 11.
|
||||
show (bool): Determines whether to print the onnx model architecture.
|
||||
Default: False.
|
||||
output_file (str): Output onnx model name. Default: 'tmp.onnx'.
|
||||
verify (bool): Determines whether to verify the onnx model.
|
||||
Default: False.
|
||||
"""
|
||||
model.cpu().eval()
|
||||
|
||||
one_img = torch.randn(input_shape)
|
||||
|
||||
register_extra_symbolics(opset_version)
|
||||
torch.onnx.export(
|
||||
model,
|
||||
one_img,
|
||||
output_file,
|
||||
export_params=True,
|
||||
keep_initializers_as_inputs=True,
|
||||
verbose=show,
|
||||
opset_version=opset_version)
|
||||
|
||||
print(f'Successfully exported ONNX model: {output_file}')
|
||||
if verify:
|
||||
# check by onnx
|
||||
onnx_model = onnx.load(output_file)
|
||||
onnx.checker.check_model(onnx_model)
|
||||
|
||||
# check the numerical value
|
||||
# get pytorch output
|
||||
pytorch_results = model(one_img)
|
||||
if not isinstance(pytorch_results, (list, tuple)):
|
||||
assert isinstance(pytorch_results, torch.Tensor)
|
||||
pytorch_results = [pytorch_results]
|
||||
|
||||
# get onnx output
|
||||
input_all = [node.name for node in onnx_model.graph.input]
|
||||
input_initializer = [
|
||||
node.name for node in onnx_model.graph.initializer
|
||||
]
|
||||
net_feed_input = list(set(input_all) - set(input_initializer))
|
||||
assert len(net_feed_input) == 1
|
||||
sess = rt.InferenceSession(output_file)
|
||||
onnx_results = sess.run(None,
|
||||
{net_feed_input[0]: one_img.detach().numpy()})
|
||||
|
||||
# compare results
|
||||
assert len(pytorch_results) == len(onnx_results)
|
||||
for pt_result, onnx_result in zip(pytorch_results, onnx_results):
|
||||
assert np.allclose(
|
||||
pt_result.detach().cpu(), onnx_result, atol=1.e-5
|
||||
), 'The outputs are different between Pytorch and ONNX'
|
||||
print('The numerical values are same between Pytorch and ONNX')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Convert MMPose models to ONNX')
|
||||
parser.add_argument('config', help='test config file path')
|
||||
parser.add_argument('checkpoint', help='checkpoint file')
|
||||
parser.add_argument('--show', action='store_true', help='show onnx graph')
|
||||
parser.add_argument('--output-file', type=str, default='tmp.onnx')
|
||||
parser.add_argument('--opset-version', type=int, default=11)
|
||||
parser.add_argument(
|
||||
'--verify',
|
||||
action='store_true',
|
||||
help='verify the onnx model output against pytorch output')
|
||||
parser.add_argument(
|
||||
'--shape',
|
||||
type=int,
|
||||
nargs='+',
|
||||
default=[1, 3, 256, 192],
|
||||
help='input size')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
|
||||
assert args.opset_version == 11, 'MMPose only supports opset 11 now'
|
||||
|
||||
# Following strings of text style are from colorama package
|
||||
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
|
||||
red_text, blue_text = '\x1b[31m', '\x1b[34m'
|
||||
white_background = '\x1b[107m'
|
||||
|
||||
msg = white_background + bright_style + red_text
|
||||
msg += 'DeprecationWarning: This tool will be deprecated in future. '
|
||||
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
|
||||
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
|
||||
msg += reset_style
|
||||
warnings.warn(msg)
|
||||
|
||||
model = init_pose_model(args.config, args.checkpoint, device='cpu')
|
||||
model = _convert_batchnorm(model)
|
||||
|
||||
# onnx.export does not support kwargs
|
||||
if hasattr(model, 'forward_dummy'):
|
||||
model.forward = model.forward_dummy
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
'Please implement the forward method for exporting.')
|
||||
|
||||
# convert model to onnx file
|
||||
pytorch2onnx(
|
||||
model,
|
||||
args.shape,
|
||||
opset_version=args.opset_version,
|
||||
show=args.show,
|
||||
output_file=args.output_file,
|
||||
verify=args.verify)
|
||||
@@ -0,0 +1,79 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import os
|
||||
import os.path as osp
|
||||
import warnings
|
||||
from argparse import ArgumentParser
|
||||
|
||||
import requests
|
||||
|
||||
from mmpose.apis import (inference_bottom_up_pose_model,
|
||||
inference_top_down_pose_model, init_pose_model,
|
||||
vis_pose_result)
|
||||
from mmpose.models import AssociativeEmbedding, TopDown
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument('img', help='Image file')
|
||||
parser.add_argument('config', help='Config file')
|
||||
parser.add_argument('checkpoint', help='Checkpoint file')
|
||||
parser.add_argument('model_name', help='The model name in the server')
|
||||
parser.add_argument(
|
||||
'--inference-addr',
|
||||
default='127.0.0.1:8080',
|
||||
help='Address and port of the inference server')
|
||||
parser.add_argument(
|
||||
'--device', default='cuda:0', help='Device used for inference')
|
||||
parser.add_argument(
|
||||
'--out-dir', default='vis_results', help='Visualization output path')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main(args):
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
# Inference single image by native apis.
|
||||
model = init_pose_model(args.config, args.checkpoint, device=args.device)
|
||||
if isinstance(model, TopDown):
|
||||
pytorch_result, _ = inference_top_down_pose_model(
|
||||
model, args.img, person_results=None)
|
||||
elif isinstance(model, (AssociativeEmbedding, )):
|
||||
pytorch_result, _ = inference_bottom_up_pose_model(model, args.img)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
vis_pose_result(
|
||||
model,
|
||||
args.img,
|
||||
pytorch_result,
|
||||
out_file=osp.join(args.out_dir, 'pytorch_result.png'))
|
||||
|
||||
# Inference single image by torchserve engine.
|
||||
url = 'http://' + args.inference_addr + '/predictions/' + args.model_name
|
||||
with open(args.img, 'rb') as image:
|
||||
response = requests.post(url, image)
|
||||
server_result = response.json()
|
||||
|
||||
vis_pose_result(
|
||||
model,
|
||||
args.img,
|
||||
server_result,
|
||||
out_file=osp.join(args.out_dir, 'torchserve_result.png'))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
main(args)
|
||||
|
||||
# Following strings of text style are from colorama package
|
||||
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
|
||||
red_text, blue_text = '\x1b[31m', '\x1b[34m'
|
||||
white_background = '\x1b[107m'
|
||||
|
||||
msg = white_background + bright_style + red_text
|
||||
msg += 'DeprecationWarning: This tool will be deprecated in future. '
|
||||
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
|
||||
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
|
||||
msg += reset_style
|
||||
warnings.warn(msg)
|
||||
Reference in New Issue
Block a user