# Path: util/config.py
# def load_config(filename = "pose_cfg.yaml"):
#     if 'POSE_PARAM_PATH' in os.environ:
#         filename = os.environ['POSE_PARAM_PATH'] + '/' + filename
#     return cfg_from_file(filename)
# 
# Path: dataset/factory.py
# def create(cfg):
#     dataset_type = cfg.dataset_type
#     if dataset_type == "mpii":
#         from dataset.mpii import MPII
#         data = MPII(cfg)
#     elif dataset_type == "coco":
#         from dataset.mscoco import MSCOCO
#         data = MSCOCO(cfg)
#     elif dataset_type == "penn_action":
#         from dataset.penn_action import PennAction
#         data = PennAction(cfg)
#     elif dataset_type == "default":
#         data = PoseDataset(cfg)
#     else:
#         raise Exception("Unsupported dataset_type: \"{}\"".format(dataset_type))
#     return data
# 
# Path: dataset/pose_dataset.py
# class Batch(Enum):
#     inputs = 0
#     part_score_targets = 1
#     part_score_weights = 2
#     locref_targets = 3
#     locref_mask = 4
#     pairwise_targets = 5
#     pairwise_mask = 6
#     data_item = 7
# 
# Path: nnet/predict.py
# def setup_pose_prediction(cfg):
#     inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
# 
#     outputs = pose_net(cfg).test(inputs)
# 
#     restorer = tf.train.Saver()
# 
#     sess = tf.Session()
# 
#     sess.run(tf.global_variables_initializer())
#     sess.run(tf.local_variables_initializer())
# 
#     # Restore variables from disk.
#     restorer.restore(sess, cfg.init_weights)
# 
#     return sess, inputs, outputs
# 
# def extract_cnn_output(outputs_np, cfg, pairwise_stats = None):
#     scmap = outputs_np['part_prob']
#     scmap = np.squeeze(scmap)
#     locref = None
#     pairwise_diff = None
#     if cfg.location_refinement:
#         locref = np.squeeze(outputs_np['locref'])
#         shape = locref.shape
#         locref = np.reshape(locref, (shape[0], shape[1], -1, 2))
#         locref *= cfg.locref_stdev
#     if cfg.pairwise_predict:
#         pairwise_diff = np.squeeze(outputs_np['pairwise_pred'])
#         shape = pairwise_diff.shape
#         pairwise_diff = np.reshape(pairwise_diff, (shape[0], shape[1], -1, 2))
#         num_joints = cfg.num_joints
#         for pair in pairwise_stats:
#             pair_id = (num_joints - 1) * pair[0] + pair[1] - int(pair[0] < pair[1])
#             pairwise_diff[:, :, pair_id, 0] *= pairwise_stats[pair]["std"][0]
#             pairwise_diff[:, :, pair_id, 0] += pairwise_stats[pair]["mean"][0]
#             pairwise_diff[:, :, pair_id, 1] *= pairwise_stats[pair]["std"][1]
#             pairwise_diff[:, :, pair_id, 1] += pairwise_stats[pair]["mean"][1]
#     return scmap, locref, pairwise_diff
# 
# def argmax_pose_predict(scmap, offmat, stride):
#     """Combine scoremat and offsets to the final pose."""
#     num_joints = scmap.shape[2]
#     pose = []
#     for joint_idx in range(num_joints):
#         maxloc = np.unravel_index(np.argmax(scmap[:, :, joint_idx]),
#                                   scmap[:, :, joint_idx].shape)
#         offset = np.array(offmat[maxloc][joint_idx])[::-1] if offmat is not None else 0
#         pos_f8 = (np.array(maxloc).astype('float') * stride + 0.5 * stride +
#                   offset)
#         pose.append(np.hstack((pos_f8[::-1],
#                                [scmap[maxloc][joint_idx]])))
#     return np.array(pose)
# 
# Path: util/visualize.py
# def _npcircle(image, cx, cy, radius, color, transparency=0.0):
# def check_point(cur_x, cur_y, minx, miny, maxx, maxy):
# def visualize_joints(image, pose):
# def show_heatmaps(cfg, img, scmap, pose, cmap="jet"):
# def show_arrows(cfg, img, pose, arrows):
# def waitforbuttonpress():


import argparse
import logging
import os
import numpy as np
import scipy.io
import scipy.ndimage
from util.config import load_config
from dataset.factory import create as create_dataset
from dataset.pose_dataset import Batch
from nnet.predict import setup_pose_prediction, extract_cnn_output, argmax_pose_predict
from util import visualize



def test_net(visualise, cache_scoremaps):
    logging.basicConfig(level=logging.INFO)

    cfg = load_config()
    dataset = create_dataset(cfg)
    dataset.set_shuffle(False)
    dataset.set_test_mode(True)

    sess, inputs, outputs = setup_pose_prediction(cfg)

    if cache_scoremaps:
        out_dir = cfg.scoremap_dir
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

    num_images = dataset.num_images
    predictions = np.zeros((num_images,), dtype=np.object)

    for k in range(num_images):
        print('processing image {}/{}'.format(k, num_images-1))

        batch = dataset.next_batch()

        outputs_np = sess.run(outputs, feed_dict={inputs: batch[Batch.inputs]})

        scmap, locref, pairwise_diff = extract_cnn_output(outputs_np, cfg)


# Path: util/config.py
# def load_config(filename = "pose_cfg.yaml"):
#     if 'POSE_PARAM_PATH' in os.environ:
#         filename = os.environ['POSE_PARAM_PATH'] + '/' + filename
#     return cfg_from_file(filename)
# 
# Path: dataset/factory.py
# def create(cfg):
#     dataset_type = cfg.dataset_type
#     if dataset_type == "mpii":
#         from dataset.mpii import MPII
#         data = MPII(cfg)
#     elif dataset_type == "coco":
#         from dataset.mscoco import MSCOCO
#         data = MSCOCO(cfg)
#     elif dataset_type == "penn_action":
#         from dataset.penn_action import PennAction
#         data = PennAction(cfg)
#     elif dataset_type == "default":
#         data = PoseDataset(cfg)
#     else:
#         raise Exception("Unsupported dataset_type: \"{}\"".format(dataset_type))
#     return data
# 
# Path: dataset/pose_dataset.py
# class Batch(Enum):
#     inputs = 0
#     part_score_targets = 1
#     part_score_weights = 2
#     locref_targets = 3
#     locref_mask = 4
#     pairwise_targets = 5
#     pairwise_mask = 6
#     data_item = 7
# 
# Path: nnet/predict.py
# def setup_pose_prediction(cfg):
#     inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
# 
#     outputs = pose_net(cfg).test(inputs)
# 
#     restorer = tf.train.Saver()
# 
#     sess = tf.Session()
# 
#     sess.run(tf.global_variables_initializer())
#     sess.run(tf.local_variables_initializer())
# 
#     # Restore variables from disk.
#     restorer.restore(sess, cfg.init_weights)
# 
#     return sess, inputs, outputs
# 
# def extract_cnn_output(outputs_np, cfg, pairwise_stats = None):
#     scmap = outputs_np['part_prob']
#     scmap = np.squeeze(scmap)
#     locref = None
#     pairwise_diff = None
#     if cfg.location_refinement:
#         locref = np.squeeze(outputs_np['locref'])
#         shape = locref.shape
#         locref = np.reshape(locref, (shape[0], shape[1], -1, 2))
#         locref *= cfg.locref_stdev
#     if cfg.pairwise_predict:
#         pairwise_diff = np.squeeze(outputs_np['pairwise_pred'])
#         shape = pairwise_diff.shape
#         pairwise_diff = np.reshape(pairwise_diff, (shape[0], shape[1], -1, 2))
#         num_joints = cfg.num_joints
#         for pair in pairwise_stats:
#             pair_id = (num_joints - 1) * pair[0] + pair[1] - int(pair[0] < pair[1])
#             pairwise_diff[:, :, pair_id, 0] *= pairwise_stats[pair]["std"][0]
#             pairwise_diff[:, :, pair_id, 0] += pairwise_stats[pair]["mean"][0]
#             pairwise_diff[:, :, pair_id, 1] *= pairwise_stats[pair]["std"][1]
#             pairwise_diff[:, :, pair_id, 1] += pairwise_stats[pair]["mean"][1]
#     return scmap, locref, pairwise_diff
# 
# def argmax_pose_predict(scmap, offmat, stride):
#     """Combine scoremat and offsets to the final pose."""
#     num_joints = scmap.shape[2]
#     pose = []
#     for joint_idx in range(num_joints):
#         maxloc = np.unravel_index(np.argmax(scmap[:, :, joint_idx]),
#                                   scmap[:, :, joint_idx].shape)
#         offset = np.array(offmat[maxloc][joint_idx])[::-1] if offmat is not None else 0
#         pos_f8 = (np.array(maxloc).astype('float') * stride + 0.5 * stride +
#                   offset)
#         pose.append(np.hstack((pos_f8[::-1],
#                                [scmap[maxloc][joint_idx]])))
#     return np.array(pose)
# 
# Path: util/visualize.py
# def _npcircle(image, cx, cy, radius, color, transparency=0.0):
# def check_point(cur_x, cur_y, minx, miny, maxx, maxy):
# def visualize_joints(image, pose):
# def show_heatmaps(cfg, img, scmap, pose, cmap="jet"):
# def show_arrows(cfg, img, pose, arrows):
# def waitforbuttonpress():


# Path: test.py
import argparse
import logging
import os
import numpy as np
import scipy.io
import scipy.ndimage
from util.config import load_config
from dataset.factory import create as create_dataset
from dataset.pose_dataset import Batch
from nnet.predict import setup_pose_prediction, extract_cnn_output, argmax_pose_predict
from util import visualize




def test_net(visualise, cache_scoremaps):
    logging.basicConfig(level=logging.INFO)

    cfg = load_config()
    dataset = create_dataset(cfg)
    dataset.set_shuffle(False)
    dataset.set_test_mode(True)

    sess, inputs, outputs = setup_pose_prediction(cfg)

    if cache_scoremaps:
        out_dir = cfg.scoremap_dir
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

    num_images = dataset.num_images
    predictions = np.zeros((num_images,), dtype=np.object)

    for k in range(num_images):
        print('processing image {}/{}'.format(k, num_images-1))

        batch = dataset.next_batch()

        outputs_np = sess.run(outputs, feed_dict={inputs: batch[Batch.inputs]})

        scmap, locref, pairwise_diff = extract_cnn_output(outputs_np, cfg)


        pose = argmax_pose_predict(scmap, locref, cfg.stride)