# streamlit_app.py

import streamlit as st
import numpy as np
import pandas as pd
from keras.models import load_model
from PIL import Image, ImageDraw
import io

# Load the trained model
model = load_model('keypoint_model.h5')

def load_image(image):
    image = Image.open(image).convert('L')  # Convert to grayscale
    image = image.resize((96, 96))  # Resize to match model input
    image_array = np.array(image)
    image_array = image_array / 255.0  # Normalize
    return image_array.reshape(-1, 96, 96, 1)  # Reshape for model input

def draw_keypoints(image, keypoints):
    # Draw keypoints on the image
    draw = ImageDraw.Draw(image)
    for (x, y) in keypoints:
        draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill='red')  # Draw a circle for each keypoint
    return image

# Title of the app
st.title("Keypoint Prediction App")

# Upload an image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Load and preprocess the image
    image = load_image(uploaded_file)
    
    # Display the uploaded image
    original_image = Image.open(uploaded_file).convert('L').resize((96, 96))  # Convert and resize for displaying
    st.image(original_image, caption='Uploaded Image.', use_column_width=True)
    
    # Make predictions
    if st.button("Predict"):
        predictions = model.predict(image)
        # Reshape predictions to (15, 2) for x and y coordinates
        keypoints = predictions.reshape(-1, 2)
        
        # Draw keypoints on the original image
        keypoint_image = draw_keypoints(original_image.copy(), keypoints)
        
        # Display the image with keypoints
        st.image(keypoint_image, caption='Image with Predicted Keypoints', use_column_width=True)

        # Display the keypoints
        st.write("Predicted Keypoints:")
        for i, (x, y) in enumerate(keypoints):
            st.write(f"Keypoint {i+1}: (X: {x:.2f}, Y: {y:.2f})")