import gradio as gr
from mediapipe import solutions
from mediapipe.framework.formats import landmark_pb2
import numpy as np
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision



def draw_landmarks_on_image(rgb_image, detection_result):
  pose_landmarks_list = detection_result.pose_landmarks
  annotated_image = np.copy(rgb_image)

  # Loop through the detected poses to visualize.
  for idx in range(len(pose_landmarks_list)):
    pose_landmarks = pose_landmarks_list[idx]

    # Draw the pose landmarks.
    pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
    pose_landmarks_proto.landmark.extend([
      landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in pose_landmarks
    ])
    solutions.drawing_utils.draw_landmarks(
      annotated_image,
      pose_landmarks_proto,
      solutions.pose.POSE_CONNECTIONS,
      solutions.drawing_styles.get_default_pose_landmarks_style())
  return annotated_image



def app(image):
  try:
    image_annotated = section(image)
    pass
  except Exception as e:
    return gr.Error(f"An error occurred: {str(e)}")

  return image_annotated

with gr.Blocks() as demo:
  with gr.Row():
    ## Input
    image = gr.Image(label="Input Image")
    ## Output
    image_annotated = gr.Image(label="Output Image")
  with gr.Row():
    detection_results = gr.Text("Pose Landmarks")

  section_btn = gr.Button("Identify keypoints")



  def section(image):
    base_options = python.BaseOptions(model_asset_path='pose_landmarker_heavy.task')
    options = vision.PoseLandmarkerOptions(
        base_options=base_options,
        output_segmentation_masks=True)
    detector = vision.PoseLandmarker.create_from_options(options)

    # STEP 3: Load the input image.
    # image = mp.Image.create_from_file(image)
    # image = mp.Image(ImageFormat.SRGB, image)
    # image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    # image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # Create a MediaPipe image using the BGR image
    # Using 1 as the image format for BGR (which is the integer representation for mediapipe::ImageFormat::SRGB)
    print(type(image))
    mediapipe_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)

    # STEP 4: Detect pose landmarks from the input image.
    detection_result = detector.detect(mediapipe_image)

    # STEP 5: Process the detection result. In this case, visualize it.
    annotated_image = draw_landmarks_on_image(image, detection_result)
    # cv2.imshow(cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR))
    return annotated_image, detection_result


  section_btn.click(section, image, [image_annotated, detection_results])

demo.launch(share=True)
# demo.queue(max_size=99).launch(debug=True)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