import gradio as gr
from matplotlib.pyplot import draw
import mediapipe as mp
import numpy as np
import tempfile
import mediapy as media

mp_hands = mp.solutions.hands
mp_hands_connections = mp.solutions.hands_connections
hands = mp_hands.Hands()
mp_draw = mp.solutions.drawing_utils 

connections = {
    'HAND_CONNECTIONS': mp_hands_connections.HAND_CONNECTIONS,
    'HAND_PALM_CONNECTIONS': mp_hands_connections.HAND_PALM_CONNECTIONS,
    'HAND_THUMB_CONNECTIONS': mp_hands_connections.HAND_THUMB_CONNECTIONS,
    'HAND_INDEX_FINGER_CONNECTIONS': mp_hands_connections.HAND_INDEX_FINGER_CONNECTIONS,
    'HAND_MIDDLE_FINGER_CONNECTIONS': mp_hands_connections.HAND_MIDDLE_FINGER_CONNECTIONS,
    'HAND_RING_FINGER_CONNECTIONS': mp_hands_connections.HAND_RING_FINGER_CONNECTIONS,
    'HAND_PINKY_FINGER_CONNECTIONS': mp_hands_connections.HAND_PINKY_FINGER_CONNECTIONS,
}

def process_hands(img, selected_connection, draw_background):
  results = hands.process(img)
  output_img = img if draw_background else np.zeros_like(img)  
  if results.multi_hand_landmarks:
    for hand_landmarks in results.multi_hand_landmarks:
      mp_draw.draw_landmarks(output_img, hand_landmarks, connections[selected_connection])    
  return output_img


def process_video(video_path, selected_connection, draw_background):
  with tempfile.NamedTemporaryFile() as f:
    out_path = f"{f.name}.{video_path.split('.')[-1]}"
    with media.VideoReader(video_path) as r:
      with media.VideoWriter(
          out_path, shape=r.shape, fps=r.fps, bps=r.bps) as w:
        for image in r:
          w.add_image(process_hands(image, selected_connection, draw_background))
    return out_path


demo = gr.Blocks()

with demo:
  gr.Markdown(
    """
    # Hand & Finger Tracking
    This is a demo of hand and finger tracking using [Google's MediaPipe](https://google.github.io/mediapipe/solutions/hands.html).
    """)

  with gr.Column():  
    draw_background = gr.Checkbox(value=True, label="Draw background?")
    connection_keys = list(connections.keys())
    selected_connection = gr.Dropdown(
      label="Select connections to draw", 
      choices=connection_keys,
      value=connection_keys[0],
      )
    with gr.Tabs():
      with gr.TabItem(label="Record a video"):
        recorded_video = gr.Video(source="webcam", format="mp4")
        submit_recorded_video = gr.Button(value="Process Video")
      with gr.TabItem(label="Upload a video"):
        uploaded_video = gr.Video(format="mp4")
        submit_uploaded_video = gr.Button(value="Process Video")  

  with gr.Column():
    processed_video = gr.Video()

  gr.Markdown('<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=kristyc.mediapipe-hands" />')
  submit_recorded_video.click(fn=process_video, inputs=[recorded_video, selected_connection, draw_background], outputs=[processed_video])
  submit_uploaded_video.click(fn=process_video, inputs=[recorded_video, selected_connection, draw_background], outputs=[processed_video])

demo.launch()