ml_video / app.py
Mathematicaljuice's picture
Update app.py
c072565 verified
Raw
History Blame Contribute Delete
6.57 kB
import cv2
import mediapipe as mp
import numpy as np
import torch
import torch.nn as nn
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"
}
CONF_THRESHOLD = 0.6
BLOCK_SIZE = 10 # frames per block
# ----------------------------
# Model
# ----------------------------
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)
model = GestureNet()
model.load_state_dict(torch.load("gesture_model1.pth", map_location="cpu"))
model.eval()
# ----------------------------
# MediaPipe
# ----------------------------
mp_hands = mp.solutions.hands
# ----------------------------
# Landmark extraction + prediction
# ----------------------------
def extract_coords(results):
coords = []
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])
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())
return coords
def run_model(coords):
"""Return label string (either gesture or 'Unknown')"""
if len(coords) < 126:
coords = coords + [0.0] * (126 - len(coords))
elif len(coords) > 126:
coords = coords[:126]
input_tensor = torch.tensor(coords, dtype=torch.float32).unsqueeze(0)
with torch.no_grad():
outputs = model(input_tensor)
probs = torch.softmax(outputs, dim=1)
pred_class = torch.argmax(probs, dim=1).item()
confidence = probs[0][pred_class].item()
if confidence >= CONF_THRESHOLD:
return GESTURE_LABELS[pred_class]
return "Unknown"
# ----------------------------
# Image prediction (unchanged format)
# ----------------------------
def predict_image_from_rgb(image_rgb):
with mp_hands.Hands(static_image_mode=True, max_num_hands=2) as hands:
results = hands.process(image_rgb)
coords = extract_coords(results)
if not coords:
return "No hand detected"
label = run_model(coords)
# run_model returns label or "Unknown" – we can keep the simple label
return label
# ----------------------------
# 🆕 Video prediction → build word from 10-frame blocks
# ----------------------------
def predict_video_word(video_path):
"""
Divides video into non-overlapping blocks of BLOCK_SIZE frames.
If every frame in a block yields the *same* valid gesture,
that gesture is added to the output word.
Returns the concatenated word, or 'Unknown' if no block is valid.
"""
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return "Cannot open video file"
frame_labels = [] # list of strings (gesture, "Unknown", or "No hand detected")
with mp_hands.Hands(
static_image_mode=False,
max_num_hands=2,
min_detection_confidence=0.5,
min_tracking_confidence=0.5,
) as hands:
while True:
ret, frame = cap.read()
if not ret:
break
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = hands.process(frame_rgb)
coords = extract_coords(results)
if coords:
label = run_model(coords) # gesture name or "Unknown"
frame_labels.append(label)
else:
frame_labels.append("No hand detected")
cap.release()
# Build word from non-overlapping blocks
word = []
for i in range(0, len(frame_labels) - BLOCK_SIZE + 1, BLOCK_SIZE):
block = frame_labels[i : i + BLOCK_SIZE]
# Check: all labels identical AND the label is not invalid
if (len(set(block)) == 1 and
block[0] not in ("Unknown", "No hand detected")):
word.append(block[0])
if word:
return "".join(word)
return "Unknown"
# ----------------------------
# Universal predictor
# ----------------------------
def predict_file(file_path, is_video=False):
if file_path is None:
return "No file uploaded"
# Direct numpy array (image)
if isinstance(file_path, np.ndarray):
if len(file_path.shape) == 3 and file_path.shape[2] == 3:
image_rgb = cv2.cvtColor(file_path, cv2.COLOR_BGR2RGB)
return predict_image_from_rgb(image_rgb)
else:
return "Invalid image format"
file_ext = str(file_path).lower()
is_video_file = is_video or file_ext.endswith(('.mp4', '.avi', '.mov', '.mkv', '.webm'))
if is_video_file:
return predict_video_word(file_path) # returns the word
else:
image = cv2.imread(file_path)
if image is None:
return "Could not read image file"
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return predict_image_from_rgb(image_rgb)
# ----------------------------
# Gradio UI
# ----------------------------
with gr.Blocks(title="Hand Gesture Recognition") as app:
gr.Markdown("# ✋ Hand Gesture Recognition (Word‑Builder)")
gr.Markdown(
"Upload an image or video. For videos, non‑overlapping blocks of **10 frames** "
"are checked. If all 10 frames in a block show the same gesture, that letter is "
"added to the output word. Blocks with mixed or invalid gestures are ignored."
)
with gr.Row():
input_file = gr.File(
label="Upload Image or Video",
file_types=["image", "video"],
type="filepath"
)
output_text = gr.Textbox(label="Prediction", lines=3)
predict_btn = gr.Button("Predict", variant="primary")
predict_btn.click(
predict_file,
inputs=input_file,
outputs=output_text
)
app.launch()