run_pose.py

import argparse
# import logging
import sys
import time
import glob
# import os
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
# from tf_pose import common
import cv2
from tf_pose.estimator import TfPoseEstimator
from tf_pose.networks import get_graph_path, model_wh



def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")

@profile
def load_model(args):
    w, h = model_wh(args.resize)
    if w == 0 or h == 0:
        e = TfPoseEstimator(get_graph_path(args.model), target_size=(432, 368), trt_bool=str2bool(args.tensorrt))
    else:
        e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h), trt_bool=str2bool(args.tensorrt))    
    return e

@profile
def first_inference(e,imagepath,args):
    w, h = model_wh(args.resize)
    frame = cv2.imread(imagepath)
    img = cv2.resize(frame, (w, h))
    humans = e.inference(img, resize_to_default=(w > 0 and h > 0), upsample_size=args.resize_out_ratio)
    return humans

@profile
def remaining_inference(e,imagepath,args):
    w, h = model_wh(args.resize)
    frame = cv2.imread(imagepath)
    img = cv2.resize(frame, (w, h))
    humans = e.inference(img, resize_to_default=(w > 0 and h > 0), upsample_size=args.resize_out_ratio)
    return humans

@profile
def main(args):
    # estimate human poses from images !
    e = load_model(args)
    image_paths = []
    if len(args.images) > 0:
        for image_set in args.images:
            image_paths.extend(glob.glob(image_set))

    image_paths.sort()
    hn = []
    a = 1
    for imagepath in image_paths:
        if (a != 1):
            humans = remaining_inference(e, imagepath, args)
        else:
            humans = first_inference(e, imagepath, args)
            a+=1
        hn.append(humans)
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='tf-pose-estimation run')
    parser.add_argument(
        "-i", "--images", nargs="*", default=[], help="Images to process"
    )
    parser.add_argument('--image', type=str, default='./images/p1.jpg')
    parser.add_argument('--model', type=str, default='mobilenet_thin',
                        help='cmu / mobilenet_thin / mobilenet_v2_large / mobilenet_v2_small')
    parser.add_argument('--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')
    parser.add_argument('--tensorrt', type=str, default="False",
                        help='for tensorrt process.')
    args = parser.parse_args()
    main(args)