mmdeploy/demo/python/pose_detection.py
lvhan028 2c18fbd2c8
[Enhancement] support kwargs in SDK python bindings (#794)
* support-kwargs

* make '__call__' as single image inference and add 'batch' API to deal with batch images inference

* fix linting error and typo

* fix lint
2022-07-29 12:32:42 +08:00

55 lines
1.4 KiB
Python

# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import cv2
import numpy as np
from mmdeploy_python import PoseDetector
def parse_args():
parser = argparse.ArgumentParser(
description='show how to use sdk python api')
parser.add_argument('device_name', help='name of device, cuda or cpu')
parser.add_argument(
'model_path',
help='path of mmdeploy SDK model dumped by model converter')
parser.add_argument('image_path', help='path of an image')
parser.add_argument(
'--bbox',
default=None,
nargs='+',
type=int,
help='bounding box of an object in format (x, y, w, h)')
args = parser.parse_args()
return args
def main():
args = parse_args()
img = cv2.imread(args.image_path)
detector = PoseDetector(
model_path=args.model_path, device_name=args.device_name, device_id=0)
if args.bbox is None:
result = detector(img)
else:
# converter (x, y, w, h) -> (left, top, right, bottom)
print(args.bbox)
bbox = np.array(args.bbox, dtype=int)
bbox[2:] += bbox[:2]
result = detector(img, bbox)
print(result)
_, point_num, _ = result.shape
points = result[:, :, :2].reshape(point_num, 2)
for [x, y] in points.astype(int):
cv2.circle(img, (x, y), 1, (0, 255, 0), 2)
cv2.imwrite('output_pose.png', img)
if __name__ == '__main__':
main()