import cv2
import mediapipe as mp
import numpy as np
import os
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
mp_pose = mp.solutions.pose


#pip install mediapipe
# For static images:
IMAGE_FILES = [os.path.join('/home/ok/Desktop/tensorflow/yoga_asanas_keypoints_thunder/yoga_asanas/', filename) for filename in os.listdir('/home/ok/Desktop/tensorflow/yoga_asanas_keypoints_thunder/yoga_asanas/') if filename.endswith(".jpg")]
BG_COLOR = (192,192,192) # gray
LIST_OF_POSE_LANDMARKS = [ mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.LEFT_EYE_INNER,
mp_pose.PoseLandmark.LEFT_EYE, mp_pose.PoseLandmark.LEFT_EYE_OUTER,
mp_pose.PoseLandmark.RIGHT_EYE_INNER, mp_pose.PoseLandmark.RIGHT_EYE,
mp_pose.PoseLandmark.RIGHT_EYE_OUTER, mp_pose.PoseLandmark.LEFT_EAR,
mp_pose.PoseLandmark.RIGHT_EAR, mp_pose.PoseLandmark.MOUTH_LEFT,
mp_pose.PoseLandmark.MOUTH_RIGHT, mp_pose.PoseLandmark.LEFT_SHOULDER,
mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.LEFT_ELBOW,
mp_pose.PoseLandmark.RIGHT_ELBOW, mp_pose.PoseLandmark.LEFT_WRIST,
mp_pose.PoseLandmark.RIGHT_WRIST, mp_pose.PoseLandmark.LEFT_PINKY,
mp_pose.PoseLandmark.RIGHT_PINKY, mp_pose.PoseLandmark.LEFT_INDEX,
mp_pose.PoseLandmark.RIGHT_INDEX, mp_pose.PoseLandmark.LEFT_THUMB,
mp_pose.PoseLandmark.RIGHT_THUMB, mp_pose.PoseLandmark.LEFT_HIP,
mp_pose.PoseLandmark.RIGHT_HIP, mp_pose.PoseLandmark.LEFT_KNEE,
mp_pose.PoseLandmark.RIGHT_KNEE, mp_pose.PoseLandmark.LEFT_ANKLE,
mp_pose.PoseLandmark.RIGHT_ANKLE, mp_pose.PoseLandmark.LEFT_HEEL,
mp_pose.PoseLandmark.RIGHT_HEEL, mp_pose.PoseLandmark.LEFT_FOOT_INDEX, mp_pose.PoseLandmark.RIGHT_FOOT_INDEX ]
csv_list_body_parts_dict = {}
with mp_pose.Pose(
    static_image_mode=True,
    model_complexity=2,
    enable_segmentation=True,
    min_detection_confidence=0.1) as pose:
  for idx, file in enumerate(IMAGE_FILES):
    image = cv2.imread(file)
    basename = os.path.basename(file).replace('.', '_');
    image_height, image_width, _ = image.shape
    # Convert the BGR image to RGB before processing.
    results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    if not results.pose_landmarks:
      continue
    else:
      # store all the body parts in a dictionary csv_list_body_parts_dict, with the file basename and body part name as strings
      #get the x y and z coordinates of the body parts also, must multiple by image width and height to get the actual coordinates
      #use pythagoras to get the distance between the body parts in 3d space

      for body_part in LIST_OF_POSE_LANDMARKS:
        csv_list_body_parts_dict[basename + '____' + str(body_part)] = (results.pose_landmarks.landmark[body_part].x , results.pose_landmarks.landmark[body_part].y , results.pose_landmarks.landmark[body_part].z , results.pose_landmarks.landmark[body_part].visibility, image_width, image_height )



    annotated_image = image.copy()
    # Draw segmentation on the image.
    # To improve segmentation around boundaries, consider applying a joint
    # bilateral filter to "results.segmentation_mask" with "image".
    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
    bg_image = np.zeros(image.shape, dtype=np.uint8)
    bg_image[:] = BG_COLOR
    annotated_image = np.where(condition, annotated_image, bg_image)
    # Draw pose landmarks on the image.
    mp_drawing.draw_landmarks(
        annotated_image,
        results.pose_landmarks,
        mp_pose.POSE_CONNECTIONS,
        landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style())
    #cv2.imwrite('./tmp/mediapipe_pose_' + str(idx) + '.png', annotated_image)
    cv2.imwrite('./tmp128/'  + basename + '.png', annotated_image)
    # Plot pose world landmarks.
    #mp_drawing.plot_landmarks(
    #    results.pose_world_landmarks, mp_pose.POSE_CONNECTIONS)

#save the dictionary to a csv file
import csv
with open('mediapipe_pose_landmarks.csv', 'w') as csv_file:
    writer = csv.writer(csv_file)
    for key, value in csv_list_body_parts_dict.items():
       writer.writerow([key, value])