Update app.py

app.py

@@ -1,5 +1,7 @@
 import cv2
-import mediapipe.python.solutions.hands as mp_hands
+import mediapipe as mp
+from mediapipe.tasks.python import vision
+from mediapipe.tasks.python import BaseOptions
 import numpy as np
 import torch
 import torch.nn as nn
@@ -9,31 +11,12 @@ import gradio as gr
 # Labels
 # ----------------------------
 GESTURE_LABELS = {
-    0: "A",
-    1: "B",
-    2: "L",
-    3: "U",
-    4: "V",
-    5: "W",
-    6: "Z",
-    7: "F",
-    8: "five",
-    9: "one",
-    10: "three",
-    11: "two",
-    12: "six",
-    13: "seven",
-    14: "eight",
-    15: "nine",
-    16: "ten",
-    17: "E",
-    18: "four",
-    19: "i",
-    20: "k",
-    21: "r",
-    22: "zero",
-    23: "m",
-    24: "s"
+    0: "A", 1: "B", 2: "L", 3: "U", 4: "V", 5: "W",
+    6: "Z", 7: "F", 8: "five", 9: "one", 10: "three",
+    11: "two", 12: "six", 13: "seven", 14: "eight",
+    15: "nine", 16: "ten", 17: "E", 18: "four",
+    19: "i", 20: "k", 21: "r", 22: "zero",
+    23: "m", 24: "s"
 }
 
 CONF_THRESHOLD = 0.6
@@ -44,128 +27,84 @@ CONF_THRESHOLD = 0.6
 class GestureNet(nn.Module):
     def __init__(self, input_size=126, num_classes=len(GESTURE_LABELS)):
         super().__init__()
-
         self.fc1 = nn.Linear(input_size, 256)
         self.fc2 = nn.Linear(256, 128)
         self.fc3 = nn.Linear(128, num_classes)
-
         self.relu = nn.ReLU()
         self.dropout = nn.Dropout(0.3)
 
     def forward(self, x):
         x = self.relu(self.fc1(x))
         x = self.dropout(x)
-
         x = self.relu(self.fc2(x))
         x = self.dropout(x)
+        return self.fc3(x)
 
-        x = self.fc3(x)
-
-        return x
-
-# ----------------------------
-# Load Model
-# ----------------------------
 model = GestureNet()
-
-model.load_state_dict(
-    torch.load(
-        "gesture_model1.pth",
-        map_location=torch.device("cpu")
-    )
-)
-
+model.load_state_dict(torch.load("gesture_model1.pth", map_location="cpu"))
 model.eval()
 
 # ----------------------------
-# Prediction Function
+# MediaPipe HandLandmarker
 # ----------------------------
-def predict(image):
-
-    if image is None:
-        return "No image uploaded"
-
-    # Convert RGB
-    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-    # MediaPipe Hands
-    with mp_hands.Hands(
-        static_image_mode=True,
-        max_num_hands=2,
-        min_detection_confidence=0.5
-    ) as hands:
-
-        results = hands.process(image_rgb)
-
+# Initialize the HandLandmarker once for efficiency
+model_path = "hand_landmarker.task"   # You may need to download this file
+# If you don't have the model file, uncomment the next line to download automatically:
+# import urllib.request; urllib.request.urlretrieve("https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/1/hand_landmarker.task", "hand_landmarker.task")
+
+base_options = BaseOptions(model_asset_path=model_path)
+options = vision.HandLandmarkerOptions(base_options=base_options, num_hands=2)
+landmarker = vision.HandLandmarker.create_from_options(options)
+
+def extract_hand_features(image_rgb):
+    """Extract normalized 63 landmarks (x,y,z) from up to 2 hands, returns flattened list of length 126."""
+    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image_rgb)
+    detection_result = landmarker.detect(mp_image)
     coords = []
 
-    # Extract landmarks
-    if results.multi_hand_landmarks:
-
-        for hand_landmarks in results.multi_hand_landmarks:
-
-            hand_coords = np.array([
-                [lm.x, lm.y, lm.z]
-                for lm in hand_landmarks.landmark
-            ])
-
-            # Normalize
+    if detection_result.hand_landmarks:
+        for hand_landmarks in detection_result.hand_landmarks:
+            # Convert to numpy array of shape (21, 3)
+            hand_coords = np.array([[lm.x, lm.y, lm.z] for lm in hand_landmarks])
+            # Normalize: subtract wrist (index 0) and scale by max distance
             hand_coords -= hand_coords[0]
-
-            max_val = np.max(
-                np.linalg.norm(hand_coords, axis=1)
-            )
-
+            max_val = np.max(np.linalg.norm(hand_coords, axis=1))
             if max_val > 0:
                 hand_coords /= max_val
-
             coords.extend(hand_coords.flatten())
 
-    # Pad / truncate
+    # Pad or truncate to fixed length 126 (21*3*2)
     if len(coords) < 126:
         coords.extend([0.0] * (126 - len(coords)))
-
     elif len(coords) > 126:
         coords = coords[:126]
+    return coords
 
-    # Predict
-    input_tensor = torch.tensor(
-        coords,
-        dtype=torch.float32
-    ).unsqueeze(0)
-
-    with torch.no_grad():
-
-        outputs = model(input_tensor)
-
-        probs = torch.softmax(outputs, dim=1)
-
-        probs = probs.numpy()[0]
-
-    pred_class = np.argmax(probs)
-
-    confidence = probs[pred_class]
-
-    if confidence >= CONF_THRESHOLD:
-
-        label = GESTURE_LABELS[pred_class]
-
-        return f"{label} ({confidence * 100:.2f}%)"
-
+def predict(image):
+    # Convert BGR (Gradio default) to RGB
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    features = extract_hand_features(image_rgb)
+
+    if len(features) == 126:
+        input_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0)
+        with torch.no_grad():
+            outputs = model(input_tensor)
+            probs = torch.softmax(outputs, dim=1).numpy()[0]
+        pred_class = np.argmax(probs)
+        confidence = probs[pred_class]
+        if confidence >= CONF_THRESHOLD:
+            return f"{GESTURE_LABELS[pred_class]} ({confidence*100:.2f}%)"
     return "Unknown"
 
 # ----------------------------
-# Gradio Interface
+# Gradio UI
 # ----------------------------
 app = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy"),
     outputs="text",
     title="Hand Gesture Recognition",
-    description="Upload a hand gesture image"
+    description="Upload an image of a hand gesture"
 )
 
-# ----------------------------
-# Launch
-# ----------------------------
-app.launch(server_name="0.0.0.0", server_port=7860)
+app.launch()