import os
import sys
import numpy as np
import cv2
sys.path.append(os.path.dirname(__file__) + "/../")

from scipy.misc import imread, imsave

from config import load_config
from dataset.factory import create as create_dataset
from nnet import predict
from util import visualize
from dataset.pose_dataset import data_to_input

from multiperson.detections import extract_detections
from multiperson.predict import SpatialModel, eval_graph, get_person_conf_multicut
from multiperson.visualize import PersonDraw, visualize_detections

import matplotlib.pyplot as plt


cfg = load_config("demo/pose_cfg_multi.yaml")

dataset = create_dataset(cfg)

sm = SpatialModel(cfg)
sm.load()

draw_multi = PersonDraw()
i=1
# Load and setup CNN part detector
sess, inputs, outputs = predict.setup_pose_prediction(cfg)
cap = cv2.VideoCapture(0)

    # Read image from file
while(True):
    if i%4==0:        
        ret, frame=cap.read()
        file_name = frame
        #image = cv2.imread(file_name)
        image=frame
        frame=cv2.resize(frame,(100,100))
        image_batch = data_to_input(image)

        # Compute prediction with the CNN
        outputs_np = sess.run(outputs, feed_dict={inputs: image_batch})
        scmap, locref, pairwise_diff = predict.extract_cnn_output(outputs_np, cfg, dataset.pairwise_stats)

        detections = extract_detections(cfg, scmap, locref, pairwise_diff)
        unLab, pos_array, unary_array, pwidx_array, pw_array = eval_graph(sm, detections)
        person_conf_multi = get_person_conf_multicut(sm, unLab, unary_array, pos_array)

        img = np.copy(image)

        visim_multi = img.copy()

        fig = cv2.imshow('frame', visim_multi)
        draw_multi.draw(visim_multi, dataset, person_conf_multi)
        #fig.axes.get_xaxis().set_visible(False)
        #fig.axes.get_yaxis().set_visible(False)

        plt.show()
        visualize.waitforbuttonpress()
        if cv2.waitKey(0):
            break
    i+=1
cap.release()
cv2.destroyAllWindows()