def infer(image, model='cmu', resize='0x0', resize_out_ratio=4.0):
"""
:param image:
:param model:
:param resize:
:param resize_out_ratio:
:return: coco_style_keypoints array
"""
w, h = model_wh(resize)
e = get_estimator(model, resize)
# estimate human poses from a single image !
image = common.read_imgfile(image, None, None)
if image is None:
raise Exception('Image can not be read, path=%s' % image)
humans = e.inference(image,
resize_to_default=(w > 0 and h > 0),
upsample_size=resize_out_ratio)
image_h, image_w = image.shape[:2]
if "TERM_PROGRAM" in os.environ and 'iTerm' in os.environ["TERM_PROGRAM"]:
image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
image_str = cv2.imencode(".jpg", image)[1].tostring()
print("\033]1337;File=name=;inline=1:" +
base64.b64encode(image_str).decode("utf-8") + "\a")
return [(eval.write_coco_json(human, image_w, image_h), human.score)
for human in humans]
def get_estimator(model='cmu', resize='0x0'):
w, h = model_wh(resize)
if w == 0 or h == 0:
e = TfPoseEstimator(get_graph_path(model), target_size=(432, 368))
else:
e = TfPoseEstimator(get_graph_path(model), target_size=(w, h))
return e
def get_model(self):
# Returns the appropriate model used for pose detection. For realtime use "mobilenet_thin"
try:
self.w, self.h = model_wh(constants.POSE_RESIZE_OPTION)
if self.w > 0 and self.h > 0:
model = TfPoseEstimator(get_graph_path(
constants.POSE_MODEL_NAME),
target_size=(self.w, self.h),
trt_bool=False)
else:
model = TfPoseEstimator(get_graph_path(
constants.POSE_MODEL_NAME),
target_size=(432, 368),
trt_bool=False)
return model
except:
pose_logger.fatal("Could not find pose estimation model")
return None
fps_time = 0
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='tf-pose-estimation Video')
parser.add_argument('--video', type=str, default='')
parser.add_argument('--resolution', type=str, default='432x368', help='network input resolution. default=432x368')
parser.add_argument('--model', type=str, default='mobilenet_thin', help='cmu / mobilenet_thin / mobilenet_v2_large / mobilenet_v2_small')
parser.add_argument('--show-process', type=bool, default=False,
help='for debug purpose, if enabled, speed for inference is dropped.')
parser.add_argument('--showBG', type=bool, default=True, help='False to show skeleton only.')
args = parser.parse_args()
logger.debug('initialization %s : %s' % (args.model, get_graph_path(args.model)))
w, h = model_wh(args.resolution)
e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
cap = cv2.VideoCapture(args.video)
if cap.isOpened() is False:
print("Error opening video stream or file")
while cap.isOpened():
ret_val, image = cap.read()
humans = e.inference(image)
if not args.showBG:
image = np.zeros(image.shape)
image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
cv2.putText(image, "FPS: %f" % (1.0 / (time.time() - fps_time)), (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
cv2.imshow('tf-pose-estimation result', image)
'--resize',
type=str,
default='0x0',
help='if provided, resize images before they are processed. '
'default=0x0, Recommends : 432x368 or 656x368 or 1312x736 ')
parser.add_argument(
'--resize-out-ratio',
type=float,
default=4.0,
help=
'if provided, resize heatmaps before they are post-processed. default=1.0'
)
args = parser.parse_args()
w, h = model_wh(args.resize)
if w == 0 or h == 0:
e = TfPoseEstimator(get_graph_path(args.model), target_size=(432, 368))
else:
e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
# estimate human poses from a single image !
image = common.read_imgfile(args.image, None, None)
if image is None:
logger.error('Image can not be read, path=%s' % args.image)
sys.exit(-1)
t = time.time()
humans = e.inference(image,
resize_to_default=(w > 0 and h > 0),
upsample_size=args.resize_out_ratio)