import gradio as gr
import pandas as pd
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
from tf_bodypix.draw import draw_poses  # utility function using OpenCV
from tensorflow.keras import preprocessing
import cv2
import json
import numpy as np
from calculations import measure_body_sizes

# Load BodyPix model
bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16))

def process_images(front_img, side_img, height):
    # Convert images to image arrays
    front_image_array = preprocessing.image.img_to_array(front_img)
    side_image_array = preprocessing.image.img_to_array(side_img)

    # BodyPix prediction
    result = bodypix_model.predict_single(front_image_array)
    mask = result.get_mask(threshold=0.75)
    # colored_mask = result.get_colored_part_mask(mask)

    poses = result.get_poses()
    print(f'shape of poses: {np.shape(poses)}')
    print(poses)
    # image_with_poses = draw_poses(
    #     front_image_array.copy(),  # create a copy to ensure we are not modifying the source image
    #     poses,
    #     keypoints_color=(255, 100, 100),
    #     skeleton_color=(100, 100, 255)
    # )

    # Measure body sizes using poses and real height
    body_sizes = measure_body_sizes(poses, height)
    print(f'Body sizes: {body_sizes}')

    # Prepare the output images
    # front_image_with_poses = preprocessing.image.array_to_img(image_with_poses)

    # Convert measurements to DataFrame for display
    measurements_df = pd.DataFrame(body_sizes)

    return measurements_df

# Create the Gradio interface
interface = gr.Interface(
    fn=process_images,
    inputs=[
        gr.Image(label="Upload Front Pose"),
        gr.Image(label="Upload Side Pose"),
        gr.Number(label="Enter Height (cm)")
    ],
    outputs=[
        gr.DataFrame(label="Body Measurements")
    ],
    title="Body Sizing System Demo",
    description="Upload two images: Front View and Side View, and input the height in cm."
)

# Launch the app
interface.launch(share=False)