TF-Pose_NAOdir.py

import argparse
import logging
import sys
import time
import glob
import ast
import os

from tf_pose import common
import cv2
import numpy as np
from tf_pose.estimator import TfPoseEstimator
from tf_pose.networks import get_graph_path, model_wh

logger = logging.getLogger('TfPoseEstimatorRun')
logger.handlers.clear()
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='tf-pose-estimation run')
    parser.add_argument('--folder', type=str, default='./images/')
    parser.add_argument('--model', type=str, default='cmu',
                        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')

    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 !
    files_grabbed = glob.glob(os.path.join(args.folder, '*.jpg'))
    all_humans = dict()
    for i,file in enumerate(files_grabbed):
        print("+++++++++++++++++++++++++++++++++++++++++")
        print(file)
        image = common.read_imgfile(file, None, None)
        image2 = np.power(image,[1.0, 0.7, 1.0])

        if image is None:
            logger.error('Image can not be read, path=%s' % args.folder)
            sys.exit(-1)

        t = time.time()
        humans = e.inference(image2, resize_to_default=(w > 0 and h > 0), upsample_size=args.resize_out_ratio)

        # for human in humans: 
        #     print("*************human total score :***************** "+ str(human.score))
        #     print("*************human max score :***************** "+ str(human.get_max_score()))
        #     for k in human.body_parts:
        #         print("--------- "+ str( human.body_parts[k].get_part_name() ) +" = "+ str(human.body_parts[k].score))
        #         print("----keypoint----- "+str( human.body_parts[k].get_part_name() )+" : "+ str( human.body_parts[k].x ) +" , "+ str( human.body_parts[k].y ))

        elapsed = time.time() - t

        logger.info('inference image: %s in %.4f seconds.' % (args.folder, elapsed))

        image = TfPoseEstimator.draw_humans(image, humans, imgcopy=False)
        try:
            import matplotlib.pyplot as plt

            fig = plt.figure()
            a = fig.add_subplot(2, 2, 1)
            a.set_title('Result')
            plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

            bgimg = cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_BGR2RGB)
            bgimg = cv2.resize(bgimg, (e.heatMat.shape[1], e.heatMat.shape[0]), interpolation=cv2.INTER_AREA)

            # show network output
            a = fig.add_subplot(2, 2, 2)
            plt.imshow(bgimg, alpha=0.5)
            tmp = np.amax(e.heatMat[:, :, :-1], axis=2)
            plt.imshow(tmp, cmap=plt.cm.gray, alpha=0.5)
            plt.colorbar()

            tmp2 = e.pafMat.transpose((2, 0, 1))
            tmp2_odd = np.amax(np.absolute(tmp2[::2, :, :]), axis=0)
            tmp2_even = np.amax(np.absolute(tmp2[1::2, :, :]), axis=0)

            a = fig.add_subplot(2, 2, 3)
            a.set_title('Vectormap-x')
            # plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
            plt.imshow(tmp2_odd, cmap=plt.cm.gray, alpha=0.5)
            plt.colorbar()

            a = fig.add_subplot(2, 2, 4)
            a.set_title('Vectormap-y')
            # plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
            plt.imshow(tmp2_even, cmap=plt.cm.gray, alpha=0.5)
            plt.colorbar()
            plt.show()
        except Exception as e:
            logger.warning('matplitlib error, %s' % e)
            cv2.imshow('result', image)
            cv2.waitKey()