inference.py

import cv2
import torch
import numpy as np
import mediapipe as mp
from collections import Counter

# -----------------------
# Device
# -----------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -----------------------
# Alphabet mapping (FIX)
# -----------------------
alphabet = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")

# -----------------------
# Model
# -----------------------
class ASLModel(torch.nn.Module):
    def __init__(self, input_size, num_classes):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(input_size, 256),
            torch.nn.ReLU(),
            torch.nn.BatchNorm1d(256),
            torch.nn.Dropout(0.4),
            torch.nn.Linear(256, 128),
            torch.nn.ReLU(),
            torch.nn.BatchNorm1d(128),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(128, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, num_classes)
        )

    def forward(self, x):
        return self.net(x)


# -----------------------
# Feature Engineering
# -----------------------
def normalize_landmarks(landmarks):
    landmarks = landmarks.reshape(1, 21, 3)
    wrist = landmarks[:, 0:1, :]
    landmarks = landmarks - wrist
    scale = np.linalg.norm(landmarks)
    return landmarks / (scale + 1e-6)

def compute_angle(a, b, c):
    ba = a - b
    bc = c - b
    cos_angle = np.dot(ba, bc) / (np.linalg.norm(ba)*np.linalg.norm(bc)+1e-6)
    return np.arccos(np.clip(cos_angle, -1.0, 1.0))

def extract_features(landmarks):
    landmarks = normalize_landmarks(landmarks)
    flat = landmarks.reshape(-1)

    d1 = np.linalg.norm(landmarks[0,4] - landmarks[0,8])
    d2 = np.linalg.norm(landmarks[0,8] - landmarks[0,12])
    d3 = np.linalg.norm(landmarks[0,12] - landmarks[0,16])
    d4 = np.linalg.norm(landmarks[0,16] - landmarks[0,20])
    distances = np.array([d1,d2,d3,d4])

    angles = []
    triplets = [
        (0,1,2),(1,2,3),(2,3,4),
        (0,5,6),(5,6,7),(6,7,8),
        (0,9,10),(9,10,11),(10,11,12),
        (0,13,14),(13,14,15),(14,15,16),
        (0,17,18),(17,18,19),(18,19,20)
    ]

    for a,b,c in triplets:
        angles.append(compute_angle(landmarks[0,a], landmarks[0,b], landmarks[0,c]))

    return np.concatenate([flat, distances, angles])


# -----------------------
# Load Model
# -----------------------
num_classes = 28  # adjust if needed
input_size = 82

model = ASLModel(input_size, num_classes).to(device)
model.load_state_dict(torch.load("best_model.pth", map_location=device))
model.eval()


# -----------------------
# MediaPipe
# -----------------------
mp_hands = mp.solutions.hands
hands = mp_hands.Hands()
mp_draw = mp.solutions.drawing_utils


# -----------------------
# TEXT SYSTEM
# -----------------------
buffer = []
buffer_size = 8
last_letter = ""
text = ""
pending_letter = "" # 🔥 New: Holds the letter waiting for confirmation
no_hand_frames = 0

def get_stable_letter(buffer):
    if len(buffer) == 0:
        return None
    return Counter(buffer).most_common(1)[0][0]


# -----------------------
# Webcam
# -----------------------
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    frame = cv2.flip(frame, 1)

    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    result = hands.process(rgb)

    if result.multi_hand_landmarks:
        no_hand_frames = 0

        for hand_landmarks in result.multi_hand_landmarks:
            lm = np.array([[p.x, p.y, p.z] for p in hand_landmarks.landmark])

            features = extract_features(lm)
            features = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)

            with torch.no_grad():
                output = model(features)
                pred = torch.argmax(output, dim=1).item()

            # 🔥 Convert to alphabet
            if pred < 26:
                letter = alphabet[pred]
            else:
                letter = ""  # ignore extra classes

            buffer.append(letter)

            if len(buffer) >= buffer_size:
                stable_letter = get_stable_letter(buffer)
                buffer.clear()

                if stable_letter and stable_letter != last_letter and not pending_letter:
                    pending_letter = stable_letter
                    last_letter = stable_letter

            cv2.putText(frame, f"{letter}", (50,100),
                        cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,0), 3)

            mp_draw.draw_landmarks(frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)

    else:
        no_hand_frames += 1

        # Add space if no hand for some time
        if no_hand_frames > 15:
            if text and text[-1] != " ":
                text += " "
            last_letter = ""

    # 🔥 Display Proposed Letter (Accept/Reject)
    if pending_letter:
        cv2.rectangle(frame, (30, 380), (350, 450), (220, 255, 220), -1)
        cv2.rectangle(frame, (30, 380), (350, 450), (0, 150, 0), 3)
        cv2.putText(frame, f"'{pending_letter}' detected!", (50, 420),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 100, 0), 2)
        cv2.putText(frame, "ENTER=Accept  R=Reject", (50, 455),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (100, 100, 100), 1)

    # 🔥 Display full text
    cv2.putText(frame, f"Text: {text[-30:]}", (50, 480),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    
    # Instructions
    cv2.putText(frame, "ENTER: Accept | R: Reject | D: Delete | ESC: Exit", (50, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (50, 50, 50), 1)

    cv2.imshow("ASL Text System", frame)

    key = cv2.waitKey(1) & 0xFF
    if key == 27: # ESC
        break
    elif key == 13: # ENTER (Accept)
        if pending_letter:
            print(f"Accepted: {pending_letter}")
            text += pending_letter
            pending_letter = ""
    elif key == ord('r'): # R (Reject)
        if pending_letter:
            print(f"Rejected: {pending_letter}")
            pending_letter = ""
    elif key == ord('d') or key == 8: # D or Backspace
        if len(text) > 0:
            print(f"Deleting: {text[-1]}")
            text = text[:-1]
        buffer.clear()
        pending_letter = ""

cap.release()
cv2.destroyAllWindows()