Update app.py

app.py

@@ -4,10 +4,11 @@ import mediapipe as mp
 import numpy as np
 import tensorflow as tf
 import tempfile
-from transformers import AutoModelForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-# Load gesture classification model from Hugging Face Hub
-model = AutoModelForImageClassification.from_pretrained("nateraw/gesture-classification")
+# Load gesture classification model from Hugging Face Hub (public model)
+processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
 
 # Mediapipe initialization
 mp_hands = mp.solutions.hands
@@ -15,10 +16,11 @@ hands = mp_hands.Hands()
 mp_draw = mp.solutions.drawing_utils
 
 # Function for gesture classification
-def classify_gesture(landmarks):
-    landmarks = np.array(landmarks).reshape(1, -1)
-    prediction = model(landmarks)
-    return np.argmax(prediction.logits.detach().numpy())
+def classify_gesture(image):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    prediction = outputs.logits.argmax(-1).item()
+    return prediction
 
 # Streamlit UI
 def main():
@@ -48,10 +50,8 @@ def main():
                 for hand_landmarks in results.multi_hand_landmarks:
                     mp_draw.draw_landmarks(frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)
 
-                    # Extract landmarks
-                    landmarks = [landmark.x for landmark in hand_landmarks.landmark]
-                    landmarks += [landmark.y for landmark in hand_landmarks.landmark]
-                    gesture = classify_gesture(landmarks)
+                    # Gesture classification
+                    gesture = classify_gesture(frame)
                     st.write(f"Gesture: {gesture}")
 
             frame_placeholder.image(frame, channels="RGB")