import cv2
import numpy as np
import os


#create a dictionary of the body parts that we want to detect, index, bodypart name
#list using mode body 25 from openpose
#initialize a blank dictionary
body_parts = {}
#add the body parts to the dictionary one by one
body_parts[0] = "Nose"
body_parts[1] = "Neck"
body_parts[2] = "RShoulder"
body_parts[3] = "RElbow"
body_parts[4] = "RWrist"
body_parts[5] = "LShoulder"
body_parts[6] = "LElbow"
body_parts[7] = "LWrist"
body_parts[8] = "MidHip"
body_parts[9] = "RHip"
body_parts[10] = "RKnee"
body_parts[11] = "RAnkle"
body_parts[12] = "LHip"
body_parts[13] = "LKnee"
body_parts[14] = "LAnkle"
body_parts[15] = "REye"
body_parts[16] = "LEye"
#body_parts[17] = "REar"
#body_parts[18] = "LEar"
#body_parts[19] = "LBigToe"
#body_parts[20] = "LSmallToe"
body_parts[21] = "LHeel"
body_parts[22] = "RBigToe"
body_parts[23] = "RSmallToe"
body_parts[24] = "RHeel"
#body_parts[25] = "Background"
# Specify the paths for the 2 files
protoFile = "/home/ok/Desktop/opencv/openpose-clone/openpose/models/pose/body_25/pose_deploy.prototxt"
weightsFile = "/home/ok/Desktop/opencv/openpose-clone/openpose/models/pose/body_25/pose_iter_584000.caffemodel"

# Read the network into Memory
net = cv2.dnn.readNetFromCaffe(protoFile, weightsFile)
# Specify the input image dimensions
inWidth = 368
inHeight = 368

#folder of images
folder = "./yoga_asanas/"
write_folder = "./yoga_asanas_keypoints_pose_estimation/"
#loop through the folder of images making sure it is a jpg using glob
for filename in os.listdir(folder):
    if filename.endswith(".jpg"):
        # Read image
        frame = cv2.imread(os.path.join(folder, filename))
        # resize image to 368x368
        frame = cv2.resize(frame, (inWidth, inHeight))

        # Prepare the frame to be fed to the network
        inpBlob = cv2.dnn.blobFromImage(frame, 1.0 / 255, (inWidth, inHeight), (0, 0, 0), swapRB=False, crop=False)

        # Set the prepared object as the input blob of the network
        net.setInput(inpBlob)

        out = net.forward()


        H = out.shape[2]
        W = out.shape[3]
        # Empty list to store the detected keypoints
        points = []
        for i in range(len(out[0])):
            # confidence map of corresponding body's part.
            probMap = out[0, i, :, :]

            # Find global maxima of the probMap.
            minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)

            # Scale the point to fit on the original image
            x = (inWidth * point[0]) / W
            y = (inHeight * point[1]) / H

            threshold = 0
            #make sure i is in the dictionary
            print(i)

            if prob > threshold  and i in body_parts:
                cv2.circle(frame, (int(x), int(y)), 3, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
                #lookup the name of the body part
                cv2.putText(frame, "{}".format(body_parts[i]), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1, lineType=cv2.LINE_AA)
                #cv2.putText(frame, "{}".format(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1, lineType=cv2.LINE_AA)

                # Add the point to the list if the probability is greater than the threshold
                points.append((int(x), int(y)))
            else :
                points.append(None)

        #save the image using the base name plus dash then opencv then the extension
        cv2.imwrite(write_folder + os.path.splitext(filename)[0] + "-opencv" + os.path.splitext(filename)[1], frame)