import streamlit as st
from PIL import Image
from ultralytics import YOLO
import torch
import utils

import utils
from drawing import draw_keypoints


@st.cache_resource()
def load_model():
    print('Loading model...')
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_pose = YOLO('yolov8l-pose.pt')
    model_pose.to(device)
    return model_pose


def draw_output(image_pil: Image.Image, keypoints: dict):    
    output_image = draw_keypoints(image_pil, keypoints).convert("RGB")
    return output_image


st.title('Pose Estimation App')

device = 'cuda' if torch.cuda.is_available() else 'cpu'
st.caption(f'Using device: {device}')

mode = st.radio('Select mode:', ['Upload an Image', 'Webcam Capture'])

if mode == 'Upload an Image':
    img_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
elif mode == 'Webcam Capture':
    img_file = st.camera_input("Take a picture")

img = None
if img_file is not None:
    img = Image.open(img_file)

st.divider()
if img is not None:
    # predict
    with st.spinner('Predicting...'):
        model = load_model()
        result = model(img)[0]
    st.markdown('**Results:**')
    keypoints = utils.get_keypoints(result)
    if keypoints is not None:
        img = draw_output(img, keypoints)
        st.image(img, caption='Predicted image', use_column_width=True)
        # calculate angles
        lea, rea = utils.get_eye_angles(keypoints)
        lba, rba = utils.get_elbow_angles(keypoints)
        angles  = {'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba}
        st.write('Angles:')
        st.json(angles)
        st.write('Raw keypoints:')
        st.json(keypoints)
    else:
        st.error('No keypoints detected!')
        st.image(img, caption='Original image', use_column_width=True)