#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os, sys, argparse
import cv2
from PIL import Image
import numpy as np
import time
from timeit import default_timer as timer
from tensorflow.keras.models import Model, load_model
import tensorflow.keras.backend as K

from hourglass.model import get_hourglass_model
from hourglass.postprocess import post_process_heatmap, post_process_heatmap_single
from common.data_utils import preprocess_image
from common.utils import get_classes, get_skeleton, render_skeleton, optimize_tf_gpu

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

default_config = {
        "num_stacks": 2,
        "mobile" : False,
        "tiny" : False,
        "model_image_size": (256, 256),
        "num_channels": 256,
        "conf_threshold": 0.1,
        "classes_path": 'configs/mpii_classes.txt',
        "skeleton_path": None,
        "weights_path": 'weights/hourglass_mobile.h5',
        "gpu_num" : 1,
    }


class Hourglass(object):
    _defaults = default_config

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, **kwargs):
        super(Hourglass, self).__init__()
        self.__dict__.update(self._defaults) # set up default values
        self.__dict__.update(kwargs) # and update with user overrides
        if self.skeleton_path:
            self.skeleton_lines = get_skeleton(self.skeleton_path)
        else:
            self.skeleton_lines = None
        self.class_names = get_classes(self.classes_path)
        self.hourglass_model = self._generate_model()
        K.set_learning_phase(0)

    def _generate_model(self):
        '''to generate the bounding boxes'''
        weights_path = os.path.expanduser(self.weights_path)
        assert weights_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        num_classes = len(self.class_names)

        # update param for tiny model
        if self.tiny is True:
            self.num_channels = 128

        # construct model and load weights.
        hourglass_model = get_hourglass_model(num_classes, self.num_stacks, self.num_channels, input_size=self.model_image_size, mobile=self.mobile)
        hourglass_model.load_weights(weights_path, by_name=True)#, skip_mismatch=True)
        hourglass_model.summary()
        return hourglass_model


    def detect_image(self, image):
        image_data = preprocess_image(image, self.model_image_size)

        image_size = image.size
        scale = (image_size[0] * 1.0 / self.model_image_size[0], image_size[1] * 1.0 / self.model_image_size[1])

        start = time.time()
        keypoints = self.predict(image_data)
        end = time.time()
        print("Inference time: {:.8f}s".format(end - start))

        # rescale keypoints back to origin image size
        keypoints_dict = dict()
        for i, keypoint in enumerate(keypoints):
            keypoints_dict[self.class_names[i]] = (keypoint[0] * scale[0] * 4, keypoint[1] * scale[1] * 4, keypoint[2])

        # draw the keypoint skeleton on image
        image_array = np.array(image, dtype='uint8')
        image_array = render_skeleton(image_array, keypoints_dict, self.skeleton_lines, self.conf_threshold)

        return Image.fromarray(image_array)

    def predict(self, image_data):
        # get final predict heatmap
        heatmap = self.hourglass_model.predict(image_data)[-1]
        heatmap = heatmap[0]
        # parse out predicted keypoint from heatmap
        keypoints = post_process_heatmap(heatmap)

        return keypoints

    def dump_model_file(self, output_model_file):
        # Dump out the final heatmap output model as inference model,
        # since we don't need the intermediate heatmap in inference stage
        model = Model(inputs=self.hourglass_model.input, outputs=self.hourglass_model.outputs[-1])
        model.save(output_model_file)


def detect_video(hourglass, video_path, output_path=""):
    import cv2
    vid = cv2.VideoCapture(0 if video_path == '0' else video_path)
    if not vid.isOpened():
        raise IOError("Couldn't open webcam or video")
    video_FourCC    = cv2.VideoWriter_fourcc(*'XVID') if video_path == '0' else int(vid.get(cv2.CAP_PROP_FOURCC))
    video_fps       = vid.get(cv2.CAP_PROP_FPS)
    video_size      = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
                        int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    isOutput = True if output_path != "" else False
    if isOutput:
        print("!!! TYPE:", type(output_path), type(video_FourCC), type(video_fps), type(video_size))
        out = cv2.VideoWriter(output_path, video_FourCC, (5. if video_path == '0' else video_fps), video_size)
    accum_time = 0
    curr_fps = 0
    fps = "FPS: ??"
    prev_time = timer()
    while True:
        return_value, frame = vid.read()
        image = Image.fromarray(frame)
        image = hourglass.detect_image(image)
        result = np.asarray(image)
        curr_time = timer()
        exec_time = curr_time - prev_time
        prev_time = curr_time
        accum_time = accum_time + exec_time
        curr_fps = curr_fps + 1
        if accum_time > 1:
            accum_time = accum_time - 1
            fps = "FPS: " + str(curr_fps)
            curr_fps = 0
        cv2.putText(result, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.50, color=(255, 0, 0), thickness=2)
        cv2.namedWindow("result", cv2.WINDOW_NORMAL)
        cv2.imshow("result", result)
        if isOutput:
            out.write(result)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    # Release everything if job is finished
    vid.release()
    if isOutput:
        out.release()
    cv2.destroyAllWindows()



def detect_img(hourglass):
    while True:
        img = input('Input image filename:')
        try:
            image = Image.open(img)
        except:
            print('Open Error! Try again!')
            continue
        else:
            r_image = hourglass.detect_image(image)
            r_image.show()



if __name__ == "__main__":
    # class Hourglass defines the default value, so suppress any default here
    parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description='demo or dump out Hourglass h5 model')
    '''
    Command line options
    '''
    parser.add_argument(
        '--num_stacks', type=int,
        help='num of stacks, default ' + str(Hourglass.get_defaults("num_stacks"))
    )
    parser.add_argument(
        '--mobile', default=False, action="store_true",
        help='use depthwise conv in hourglass, default ' + str(Hourglass.get_defaults("mobile"))
    )
    parser.add_argument(
        '--tiny', default=False, action="store_true",
        help='tiny network for speed, feature channel=128, default ' + str(Hourglass.get_defaults("tiny"))
    )
    parser.add_argument(
        '--model_image_size', type=str,
        help='model image input size as <num>x<num>, default ' +
        str(Hourglass.get_defaults("model_image_size")[0])+'x'+str(Hourglass.get_defaults("model_image_size")[1]),
        default=str(Hourglass.get_defaults("model_image_size")[0])+'x'+str(Hourglass.get_defaults("model_image_size")[1])
    )
    parser.add_argument(
        '--weights_path', type=str,
        help='path to model weight file, default ' + Hourglass.get_defaults("weights_path")
    )
    parser.add_argument(
        '--classes_path', type=str, required=False,
        help='path to keypoint class definitions, default ' + Hourglass.get_defaults("classes_path")
    )
    parser.add_argument(
        '--skeleton_path', type=str, required=False,
        help='path to keypoint skeleton definitions, default ' + str(Hourglass.get_defaults("skeleton_path"))
    )
    parser.add_argument(
        '--conf_threshold', type=float,
        help='confidence threshold, default ' + str(Hourglass.get_defaults("conf_threshold"))
    )

    parser.add_argument(
        '--image', default=False, action="store_true",
        help='Image detection mode, will ignore all positional arguments'
    )
    '''
    Command line positional arguments -- for video detection mode
    '''
    parser.add_argument(
        "--input", nargs='?', type=str,required=False,default='./path2your_video',
        help = "Video input path"
    )
    parser.add_argument(
        "--output", nargs='?', type=str, default="",
        help = "[Optional] Video output path"
    )
    '''
    Command line positional arguments -- for model dump
    '''
    parser.add_argument(
        '--dump_model', default=False, action="store_true",
        help='Dump out training model to inference model'
    )

    parser.add_argument(
        '--output_model_file', type=str,
        help='output inference model file'
    )

    args = parser.parse_args()
    # param parse
    if args.model_image_size:
        height, width = args.model_image_size.split('x')
        args.model_image_size = (int(height), int(width))

    # get wrapped inference object
    hourglass = Hourglass(**vars(args))

    if args.dump_model:
        """
        Dump out training model to inference model
        """
        if not args.output_model_file:
            raise ValueError('output model file is not specified')

        print('Dumping out training model to inference model')
        hourglass.dump_model_file(args.output_model_file)
        sys.exit()

    if args.image:
        """
        Image detection mode, disregard any remaining command line arguments
        """
        print("Image detection mode")
        if "input" in args:
            print(" Ignoring remaining command line arguments: " + args.input + "," + args.output)
        detect_img(hourglass)
    elif "input" in args:
        detect_video(hourglass, args.input, args.output)
    else:
        print("Must specify at least video_input_path.  See usage with --help.")


numpy.asarray tensorflow.keras.backend.set_learning_phase numpy.array tensorflow.keras.models.Model
[(113, 'cv2.VideoCapture', 'cv2.VideoCapture', (["(0 if video_path == '0' else video_path)"], {}), False, 'import cv2\n'), (127, 'timeit.default_timer', 'timer', ([], {}), True, 'from timeit import default_timer as timer\n'), (154, 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ([], {}), False, 'import cv2\n'), (174, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'argument_default': 'argparse.SUPPRESS', 'description': '"""demo or dump out Hourglass h5 model"""'}), False, 'import os, sys, argparse\n'), (51, 'common.utils.get_classes', 'get_classes', (['self.classes_path'], {}), False, 'from common.utils import get_classes, get_skeleton, render_skeleton, optimize_tf_gpu\n'), (53, 'tensorflow.keras.backend.set_learning_phase', 'K.set_learning_phase', (['(0)'], {}), True, 'import tensorflow.keras.backend as K\n'), (57, 'os.path.expanduser', 'os.path.expanduser', (['self.weights_path'], {}), False, 'import os, sys, argparse\n'), (67, 'hourglass.model.get_hourglass_model', 'get_hourglass_model', (['num_classes', 'self.num_stacks', 'self.num_channels'], {'input_size': 'self.model_image_size', 'mobile': 'self.mobile'}), False, 'from hourglass.model import get_hourglass_model\n'), (74, 'common.data_utils.preprocess_image', 'preprocess_image', (['image', 'self.model_image_size'], {}), False, 'from common.data_utils import preprocess_image\n'), (79, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (81, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (90, 'numpy.array', 'np.array', (['image'], {'dtype': '"""uint8"""'}), True, 'import numpy as np\n'), (91, 'common.utils.render_skeleton', 'render_skeleton', (['image_array', 'keypoints_dict', 'self.skeleton_lines', 'self.conf_threshold'], {}), False, 'from common.utils import get_classes, get_skeleton, render_skeleton, optimize_tf_gpu\n'), (93, 'PIL.Image.fromarray', 'Image.fromarray', (['image_array'], {}), False, 'from PIL import Image\n'), (100, 'hourglass.postprocess.post_process_heatmap', 'post_process_heatmap', (['heatmap'], {}), False, 'from hourglass.postprocess import post_process_heatmap, post_process_heatmap_single\n'), (107, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': 'self.hourglass_model.input', 'outputs': 'self.hourglass_model.outputs[-1]'}), False, 'from tensorflow.keras.models import Model, load_model\n'), (116, 'cv2.VideoWriter_fourcc', 'cv2.VideoWriter_fourcc', (["*'XVID'"], {}), False, 'import cv2\n'), (123, 'cv2.VideoWriter', 'cv2.VideoWriter', (['output_path', 'video_FourCC', "(5.0 if video_path == '0' else video_fps)", 'video_size'], {}), False, 'import cv2\n'), (130, 'PIL.Image.fromarray', 'Image.fromarray', (['frame'], {}), False, 'from PIL import Image\n'), (132, 'numpy.asarray', 'np.asarray', (['image'], {}), True, 'import numpy as np\n'), (133, 'timeit.default_timer', 'timer', ([], {}), True, 'from timeit import default_timer as timer\n'), (142, 'cv2.putText', 'cv2.putText', (['result'], {'text': 'fps', 'org': '(3, 15)', 'fontFace': 'cv2.FONT_HERSHEY_SIMPLEX', 'fontScale': '(0.5)', 'color': '(255, 0, 0)', 'thickness': '(2)'}), False, 'import cv2\n'), (144, 'cv2.namedWindow', 'cv2.namedWindow', (['"""result"""', 'cv2.WINDOW_NORMAL'], {}), False, 'import cv2\n'), (145, 'cv2.imshow', 'cv2.imshow', (['"""result"""', 'result'], {}), False, 'import cv2\n'), (259, 'sys.exit', 'sys.exit', ([], {}), False, 'import os, sys, argparse\n'), (48, 'common.utils.get_skeleton', 'get_skeleton', (['self.skeleton_path'], {}), False, 'from common.utils import get_classes, get_skeleton, render_skeleton, optimize_tf_gpu\n'), (162, 'PIL.Image.open', 'Image.open', (['img'], {}), False, 'from PIL import Image\n'), (148, 'cv2.waitKey', 'cv2.waitKey', (['(1)'], {}), False, 'import cv2\n')]