Update app.py

app.py

@@ -1,7 +1,5 @@
 import cv2
 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
@@ -27,73 +25,101 @@ 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)
 
+# ----------------------------
+# Load Model
+# ----------------------------
 model = GestureNet()
-model.load_state_dict(torch.load("gesture_model1.pth", map_location="cpu"))
+
+model.load_state_dict(
+    torch.load("gesture_model1.pth", map_location=torch.device("cpu"))
+)
+
 model.eval()
 
 # ----------------------------
-# MediaPipe HandLandmarker
+# MediaPipe
 # ----------------------------
-# 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)
+mp_hands = mp.solutions.hands
+
+hands = mp_hands.Hands(
+    static_image_mode=True,
+    max_num_hands=2,
+    min_detection_confidence=0.5
+)
+
+# ----------------------------
+# Prediction Function
+# ----------------------------
+def predict(image):
+
+    if image is None:
+        return "No image uploaded"
+
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    results = hands.process(image_rgb)
+
     coords = []
 
-    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
+    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
             hand_coords -= hand_coords[0]
+
             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 or truncate to fixed length 126 (21*3*2)
+    # Padding
     if len(coords) < 126:
         coords.extend([0.0] * (126 - len(coords)))
+
     elif len(coords) > 126:
         coords = coords[:126]
-    return coords
 
-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}%)"
+    input_tensor = torch.tensor(
+        coords,
+        dtype=torch.float32
+    ).unsqueeze(0)
+
+    with torch.no_grad():
+        outputs = model(input_tensor)
+
+        probs = torch.softmax(outputs, dim=1).cpu().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"
 
 # ----------------------------
@@ -107,4 +133,7 @@ app = gr.Interface(
     description="Upload an image of a hand gesture"
 )
 
-app.launch()
+# ----------------------------
+# Launch
+# ----------------------------
+app.launch(server_name="0.0.0.0", server_port=7860)