import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
import joblib
import pandas as pd
from numpy.linalg import norm
import matplotlib.pyplot as plt
import os
st.set_page_config(layout="wide")
# Function to load the Random Forest model
@st.cache_resource
def load_model():
    try:
        return joblib.load('best_random_forest_model.pkl')
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

# Load the model using the cached function
model = load_model()

# Ensure the model is loaded before proceeding
if model is None:
    st.stop()

# Initialize MediaPipe Hands
@st.cache_resource
def load_mediapipe_model():
    mp_hands = mp.solutions.hands
    return mp_hands.Hands(static_image_mode=True, max_num_hands=1, min_detection_confidence=0.5)

hands = load_mediapipe_model()
mp_drawing = mp.solutions.drawing_utils

# Function to normalize landmarks
def normalize_landmarks(landmarks):
    center = np.mean(landmarks, axis=0)
    landmarks_centered = landmarks - center
    std_dev = np.std(landmarks_centered, axis=0)
    landmarks_normalized = landmarks_centered / std_dev
    return np.nan_to_num(landmarks_normalized)

# Function to calculate angles between landmarks
def calculate_angles(landmarks):
    angles = []
    for i in range(20):
        for j in range(i + 1, 21):
            vector = landmarks[j] - landmarks[i]
            angle_x = np.arccos(np.clip(vector[0] / norm(vector), -1.0, 1.0))
            angle_y = np.arccos(np.clip(vector[1] / norm(vector), -1.0, 1.0))
            angles.extend([angle_x, angle_y])
    return angles

# Function to process image and predict alphabet
def process_and_predict(image):
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = hands.process(image_rgb)
    
    if results.multi_hand_landmarks:
        landmarks = np.array([[lm.x, lm.y] for lm in results.multi_hand_landmarks[0].landmark])
        landmarks_normalized = normalize_landmarks(landmarks)
        angles = calculate_angles(landmarks_normalized)
        
        angle_columns = [f'angle_{i}' for i in range(len(angles))]
        angles_df = pd.DataFrame([angles], columns=angle_columns)
        
        probabilities = model.predict_proba(angles_df)[0]
        return probabilities, landmarks
    
    return None, None

# Function to plot hand landmarks
def plot_hand_landmarks(landmarks, title):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.scatter(landmarks[:, 0], landmarks[:, 1], c='blue', s=50)
    for connection in mp.solutions.hands.HAND_CONNECTIONS:
        start_idx = connection[0]
        end_idx = connection[1]
        ax.plot([landmarks[start_idx, 0], landmarks[end_idx, 0]],
                [landmarks[start_idx, 1], landmarks[end_idx, 1]], 'r-', linewidth=2)
    ax.invert_yaxis()
    ax.set_title(title, fontsize=16)
    ax.axis('off')
    return fig

# Streamlit app
st.title("ASL Recognition App")

# Create tabs for different functionalities
tab1, tab2 = st.tabs(["Predict ASL Sign", "View Hand Landmarks"])

with tab1:
    st.header("Predict ASL Sign")
    uploaded_file = st.file_uploader("Upload an image of an ASL sign", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        col1, col2 = st.columns(2)
        with col1:
            image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
            st.image(image, caption="Uploaded Image", use_column_width=True)

        with col2:
            probabilities, landmarks = process_and_predict(image)
            
            if probabilities is not None:
                st.subheader("Top 5 Predictions:")
                top_indices = np.argsort(probabilities)[::-1][:5]
                for i in top_indices:
                    st.write(f"{model.classes_[i]}: {probabilities[i]:.2f}")
                
                fig = plot_hand_landmarks(landmarks, "Detected Hand Landmarks")
                st.pyplot(fig)
            else:
                st.write("No hand detected in the image.")

with tab2:
    st.header("View Hand Landmarks")
    all_alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    excluded_alphabets = 'DMNPTUVXZ'
    available_alphabets = ''.join(set(all_alphabets) - set(excluded_alphabets))
    
    selected_alphabets = st.multiselect("Select alphabets to view landmarks:", list(available_alphabets))

    if selected_alphabets:
        cols = st.columns(min(3, len(selected_alphabets)))
        for idx, alphabet in enumerate(selected_alphabets):
            with cols[idx % 3]:
                image_path = f'asl test set/{alphabet.lower()}.jpeg'
                if os.path.exists(image_path):
                    image = cv2.imread(image_path)
                    _, landmarks = process_and_predict(image)
                    if landmarks is not None:
                        fig = plot_hand_landmarks(landmarks, f"Hand Landmarks for {alphabet}")
                        st.pyplot(fig)
                    else:
                        st.write(f"No hand detected for {alphabet}")
                else:
                    st.write(f"Image not found for {alphabet}")

# Release MediaPipe resources
hands.close()