import gradio as gr
import PIL.Image
from poser import draw_bones, movenet
import tensorflow as tf
import numpy as np

def predict(image: PIL.Image.Image):
    input_size = 256
    image = image.resize((1280, 1280))
    image_tf = tf.keras.preprocessing.image.img_to_array(image)

    # Prepare input for the model
    input_image = tf.expand_dims(image_tf, axis=0)
    input_image = tf.image.resize_with_pad(input_image, input_size, input_size)

    # Run MoveNet pose estimation
    keypoints = movenet(input_image)

    # Draw bones on the image
    joints = draw_bones(image, keypoints)

    # Format points as text
    points = [f"{label} → ({int(float(x))}, {int(float(y))})" for label, x, y in joints]

    return image, joints, points

with gr.Blocks(title="MoveNet Pose Estimation") as demo:
    gr.Markdown("# 🧍‍♀️ Human Pose Estimation with MoveNet")
    gr.Markdown("Upload an image to detect body keypoints and view the skeleton overlay.")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            run_button = gr.Button("Detect Pose", variant="primary")
        with gr.Column():
            output_image = gr.Image(type="numpy", label="Skeleton Output")
            joints_table = gr.Dataframe(headers=["Label", "X", "Y"], row_count=17, col_count=(3, "fixed"))
            point_text = gr.Textbox(label="Formatted Keypoints", lines=8)

    run_button.click(
        fn=predict,
        inputs=[input_image],
        outputs=[output_image, joints_table, point_text]
    )

demo.launch()