slr-easz / app.py
Niharmahesh's picture
Update app.py
ee3ddaf verified
raw
history blame
4.52 kB
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
# 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()
ax.scatter(landmarks[:, 0], landmarks[:, 1], c='blue')
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-')
ax.invert_yaxis()
ax.set_title(title)
ax.axis('off')
return fig
# Streamlit app
st.title("ASL Recognition App")
# Create two columns
col1, col2 = st.columns(2)
with col1:
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:
image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
st.image(image, caption="Uploaded Image", use_column_width=True)
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}")
else:
st.write("No hand detected in the image.")
with col2:
st.header("Draw Hand Landmarks")
all_alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
excluded_alphabets = 'DMNPTUVXZ'
available_alphabets = ''.join(set(all_alphabets) - set(excluded_alphabets))
selected_alphabet = st.selectbox("Select an alphabet to draw landmarks:", list(available_alphabets))
if selected_alphabet:
image_path = f'asl test set/{selected_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 {selected_alphabet}")
st.pyplot(fig)
else:
st.write(f"No hand detected for {selected_alphabet}")
else:
st.write(f"Image not found for {selected_alphabet}")