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

# Read image
frame = cv2.imread("051-kukkutasana-resize.jpg")
# 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 file extension
basename = os.path.basename("051-kukkutasana-resize.jpg")
filename = os.path.splitext(basename)[0]
newfilename = filename + "-opencv.jpg"


cv2.imwrite(newfilename, frame)
cv2.destroyAllWindows()
#==============================================

#cv2.dnn_DetectionModel
# Convert the image to grayscale
#gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# Detect keypoints using a feature detector, such as SIFT
#detector = cv2.human
#keypoints, descriptors = detector.detectAndCompute(gray, None)

# Draw the keypoints on the image
#img_keypoints = cv2.drawKeypoints(img, keypoints, None, color=(0,255,0), flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)

# Display the image with keypoints
#cv2.imshow("Keypoints", img_keypoints)
#cv2.waitKey(0)
#cv2.destroyAllWindows()