app.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import cv2
import mediapipe as mp
import gradio as gr
import os

# ==========================================
# 1. Core Translation Dictionary (Full 90 words)
# ==========================================
# This dictionary is used for looking up translations after prediction.
translation_dict = {
    'payung': {'ms': 'Payung', 'en': 'Umbrella', 'zh': '雨伞', 'ta': 'குடை'},
    'baik_2': {'ms': 'Baik-baik', 'en': 'Very good', 'zh': '非常好', 'ta': 'மிகவும் நல்லது'},
    'baca': {'ms': 'Baca', 'en': 'Read', 'zh': '读', 'ta': 'படி'},
    'ayah': {'ms': 'Ayah', 'en': 'Father', 'zh': '父亲', 'ta': 'தந்தை'},
    'mari': {'ms': 'Mari', 'en': 'Come', 'zh': '来', 'ta': 'வா'},
    'minum': {'ms': 'Minum', 'en': 'Drink', 'zh': '喝', 'ta': 'குடி'},
    'tolong': {'ms': 'Tolong', 'en': 'Help', 'zh': '帮助', 'ta': 'உதவி'},
    'pukul': {'ms': 'Pukul', 'en': 'Hit', 'zh': '打', 'ta': 'அடி'},
    'perlahan_2': {'ms': 'Perlahan-lahan', 'en': 'Very slow', 'zh': '非常慢', 'ta': 'மிகவும் மெதுவாக'},
    'hujan': {'ms': 'Hujan', 'en': 'Rain', 'zh': '雨', 'ta': 'மழை'},
    'perlahan': {'ms': 'Perlahan', 'en': 'Slow', 'zh': '慢', 'ta': 'மெதுவாக'},
    'bahasa_isyarat': {'ms': 'Bahasa Isyarat', 'en': 'Sign Language', 'zh': '手语', 'ta': 'சைகை மொழி'},
    'kacau': {'ms': 'Kacau', 'en': 'Disturb', 'zh': '打扰', 'ta': 'தொந்தரவு'},
    'marah': {'ms': 'Marah', 'en': 'Angry', 'zh': '生气', 'ta': 'கோபம்'},
    'kesakitan': {'ms': 'Kesakitan', 'en': 'Pain', 'zh': '疼痛', 'ta': 'வலி'},
    'saudara': {'ms': 'Saudara', 'en': 'Relative', 'zh': '亲戚', 'ta': 'உறவினர்'},
    'bapa_saudara': {'ms': 'Bapa Saudara', 'en': 'Uncle', 'zh': '叔伯', 'ta': 'மாமா'},
    'pinjam': {'ms': 'Pinjam', 'en': 'Borrow', 'zh': '借', 'ta': 'கடன் வாங்கு'},
    'bapa': {'ms': 'Bapa', 'en': 'Father', 'zh': '父亲', 'ta': 'அப்பா'},
    'bas': {'ms': 'Bas', 'en': 'Bus', 'zh': '巴士', 'ta': 'பேருந்து'},
    'pergi': {'ms': 'Pergi', 'en': 'Go', 'zh': '去', 'ta': 'செல்'},
    'tidur': {'ms': 'Tidur', 'en': 'Sleep', 'zh': '睡觉', 'ta': 'தூங்கு'},
    'pergi_2': {'ms': 'Pergi-pergi', 'en': 'Go away', 'zh': '走开', 'ta': 'வெளியே போ'},
    'beli_2': {'ms': 'Beli-beli', 'en': 'Shopping', 'zh': '购物', 'ta': 'பொருள் வாங்குதல்'},
    'tanya': {'ms': 'Tanya', 'en': 'Ask', 'zh': '问', 'ta': 'கேள்'},
    'suka': {'ms': 'Suka', 'en': 'Like', 'zh': '喜欢', 'ta': 'பிடிக்கும்'},
    'pandai_2': {'ms': 'Pandai-pandai', 'en': 'Very clever', 'zh': '非常聪明', 'ta': 'மிகவும் புத்திசாலி'},
    'sampai': {'ms': 'Sampai', 'en': 'Arrived', 'zh': '到达', 'ta': 'வந்து சேர்ந்தார்'},
    'sekolah': {'ms': 'Sekolah', 'en': 'School', 'zh': '学校', 'ta': 'பள்ளி'},
    'bomba': {'ms': 'Bomba', 'en': 'Firefighter', 'zh': '消防员', 'ta': 'தீயணைப்பு வீரர்'},
    'baik': {'ms': 'Baik', 'en': 'Good', 'zh': '好', 'ta': 'நல்லது'},
    'bila': {'ms': 'Bila', 'en': 'When', 'zh': '什么时候', 'ta': 'எப்பொழுது'},
    'berlari': {'ms': 'Berlari', 'en': 'Running', 'zh': '跑', 'ta': 'ஓடுதல்'},
    'pandai': {'ms': 'Pandai', 'en': 'Clever', 'zh': '聪明', 'ta': 'புத்திசாலி'},
    'emak': {'ms': 'Emak', 'en': 'Mother', 'zh': '母亲', 'ta': 'அம்மா'},
    'arah': {'ms': 'Arah', 'en': 'Direction', 'zh': '方向', 'ta': 'திசை'},
    'abang': {'ms': 'Abang', 'en': 'Elder Brother', 'zh': '哥哥', 'ta': 'அண்ணன்'},
    'mana': {'ms': 'Mana', 'en': 'Where', 'zh': '哪里', 'ta': 'எங்கே'},
    'teksi': {'ms': 'Teksi', 'en': 'Taxi', 'zh': '德士', 'ta': 'டாக்ஸி'},
    'boleh': {'ms': 'Boleh', 'en': 'Can', 'zh': '可以', 'ta': 'முடியும்'},
    'dari': {'ms': 'Dari', 'en': 'From', 'zh': '从', 'ta': 'இருந்து'},
    'lupa': {'ms': 'Lupa', 'en': 'Forget', 'zh': '忘记', 'ta': 'மறந்து விடு'},
    'dapat': {'ms': 'Dapat', 'en': 'Get', 'zh': '得到', 'ta': 'கிடைக்கும்'},
    'pensil': {'ms': 'Pensil', 'en': 'Pencil', 'zh': '铅笔', 'ta': 'பென்சில்'},
    'main': {'ms': 'Main', 'en': 'Play', 'zh': '玩', 'ta': 'விளையாடு'},
    'bagaimana': {'ms': 'Bagaimana', 'en': 'How', 'zh': '怎样', 'ta': 'எப்படி'},
    'keluarga': {'ms': 'Keluarga', 'en': 'Family', 'zh': '家族', 'ta': 'குடும்பம்'},
    'emak_saudara': {'ms': 'Emak Saudara', 'en': 'Aunt', 'zh': '姑姨', 'ta': 'அத்தை'},
    'mohon': {'ms': 'Mohon', 'en': 'Apply', 'zh': '申请', 'ta': 'விண்ணப்பிக்கவும்'},
    'nasi_lemak': {'ms': 'Nasi Lemak', 'en': 'Nasi Lemak', 'zh': '椰浆饭', 'ta': 'நாசி லெமாக்'},
    'ada': {'ms': 'Ada', 'en': 'Have', 'zh': '有', 'ta': 'உள்ளது'},
    'assalamualaikum': {'ms': 'Assalamualaikum', 'en': 'Peace be upon you', 'zh': '祝你平安', 'ta': 'அஸ்ஸலாமு அலைக்கும்'},
    'beli': {'ms': 'Beli', 'en': 'Buy', 'zh': '买', 'ta': 'வாங்கு'},
    'kakak': {'ms': 'Kakak', 'en': 'Elder Sister', 'zh': '姐姐', 'ta': 'அக்கா'},
    'anak_perempuan': {'ms': 'Anak Perempuan', 'en': 'Daughter', 'zh': '女儿', 'ta': 'மகள்'},
    'masa': {'ms': 'Masa', 'en': 'Time', 'zh': '时间', 'ta': 'நேரம்'},
    'panas': {'ms': 'Panas', 'en': 'Hot', 'zh': '热', 'ta': 'சூடு'},
    'sudah': {'ms': 'Sudah', 'en': 'Already', 'zh': '已经', 'ta': 'ஏற்கனவே'},
    'apa_khabar': {'ms': 'Apa khabar', 'en': 'How are you', 'zh': '你好吗', 'ta': 'எப்படி இருக்கிறீர்கள்'},
    'buat': {'ms': 'Buat', 'en': 'Do', 'zh': '做', 'ta': 'செய்'},
    'lelaki': {'ms': 'Lelaki', 'en': 'Male', 'zh': '男', 'ta': 'ஆண்'},
    'siapa': {'ms': 'Siapa', 'en': 'Who', 'zh': '谁', 'ta': 'யார்'},
    'jahat': {'ms': 'Jahat', 'en': 'Bad', 'zh': '坏', 'ta': 'கெட்ட'},
    'tandas': {'ms': 'Tandas', 'en': 'Toilet', 'zh': '厕所', 'ta': 'கழிப்பறை'},
    'bawa': {'ms': 'Bawa', 'en': 'Bring', 'zh': '带', 'ta': 'கொண்டு வா'},
    'kereta': {'ms': 'Kereta', 'en': 'Car', 'zh': '汽车', 'ta': 'கார்'},
    'apa': {'ms': 'Apa', 'en': 'What', 'zh': '什么', 'ta': 'என்ன'},
    'jangan': {'ms': 'Jangan', 'en': 'Don\'t', 'zh': '不要', 'ta': 'வேண்டாம்'},
    'ribut': {'ms': 'Ribut', 'en': 'Storm', 'zh': '暴风雨', 'ta': 'புயல்'},
    'berapa': {'ms': 'Berapa', 'en': 'How much', 'zh': '多少', 'ta': 'எவ்வளவு'},
    'berjalan': {'ms': 'Berjalan', 'en': 'Walking', 'zh': '走', 'ta': 'நடத்தல்'},
    'hari': {'ms': 'Hari', 'en': 'Day', 'zh': '天', 'ta': 'நாள்'},
    'teh_tarik': {'ms': 'Teh Tarik', 'en': 'Pulled Tea', 'zh': '拉茶', 'ta': 'தே தாரிக்'},
    'masalah': {'ms': 'Masalah', 'en': 'Problem', 'zh': '问题', 'ta': 'பிரச்சனை'},
    'makan': {'ms': 'Makan', 'en': 'Eat', 'zh': '吃', 'ta': 'சாப்பிடு'},
    'polis': {'ms': 'Polis', 'en': 'Police', 'zh': '警察', 'ta': 'போலீஸ்'},
    'panas_2': {'ms': 'Panas-panas', 'en': 'Very hot', 'zh': '非常热', 'ta': 'மிகவும் சூடாக'},
    'sejuk': {'ms': 'Sejuk', 'en': 'Cold', 'zh': '冷', 'ta': 'குளிர்'},
    'curi': {'ms': 'Curi', 'en': 'Steal', 'zh': '偷', 'ta': 'திருடு'},
    'lemak': {'ms': 'Lemak', 'en': 'Fat', 'zh': '脂肪', 'ta': 'கொழுப்பு'},
    'buang': {'ms': 'Buang', 'en': 'Throw', 'zh': '丢', 'ta': 'எறி'},
    'jam': {'ms': 'Jam', 'en': 'Clock', 'zh': '钟', 'ta': 'கடிகாரம்'},
    'ambil': {'ms': 'Ambil', 'en': 'Take', 'zh': '拿', 'ta': 'எடு'},
    'hi': {'ms': 'Hi', 'en': 'Hi', 'zh': '嗨', 'ta': 'வணக்கம்'},
    'anak_lelaki': {'ms': 'Anak Lelaki', 'en': 'Son', 'zh': '儿子', 'ta': 'மகன்'},
    'jumpa': {'ms': 'Jumpa', 'en': 'Meet', 'zh': '见面', 'ta': 'சந்திப்பு'},
    'nasi': {'ms': 'Nasi', 'en': 'Rice', 'zh': '饭', 'ta': 'அரிசி சாதம்'},
    'pen': {'ms': 'Pen', 'en': 'Pen', 'zh': '笔', 'ta': 'பேனா'},
    'bola': {'ms': 'Bola', 'en': 'Ball', 'zh': '球', 'ta': 'பந்து'},
    'perempuan': {'ms': 'Perempuan', 'en': 'Female', 'zh': '女', 'ta': 'பெண்'}
}

# ==========================================
# 2. Critical Fix: Hardcoded Gesture Order
# ==========================================
# [IMPORTANT] This list MUST match the exact order used during training (in mode_testing.py).
# Using sorted() on the dictionary keys here would break the model because
# model output indices (0, 1, 2...) are tied to the specific order below.
gestures = ['tandas', 'assalamualaikum', 'bawa', 'abang', 'sejuk', 'berlari', 'pandai_2',
 'jumpa', 'perempuan', 'teh_tarik', 'pensil', 'anak_perempuan', 'lelaki', 'main',
 'hari', 'bas', 'bola', 'minum', 'jam', 'beli_2', 'jahat', 'baik_2', 'ada', 'emak',
 'curi', 'kakak', 'emak_saudara', 'bila', 'bomba', 'mana', 'bapa_saudara',
 'saudara', 'anak_lelaki', 'siapa', 'kereta', 'pen', 'buang', 'bagaimana', 'hi',
 'kesakitan', 'pergi', 'masa', 'baca', 'mohon', 'nasi_lemak', 'sudah', 'dapat',
 'teksi', 'baik', 'marah', 'pandai', 'berjalan', 'sampai', 'bapa', 'makan',
 'kacau', 'bahasa_isyarat', 'berapa', 'perlahan_2', 'polis', 'boleh', 'dari',
 'suka', 'nasi', 'beli', 'apa_khabar', 'tidur', 'lupa', 'perlahan', 'masalah',
 'ribut', 'ambil', 'pergi_2', 'mari', 'keluarga', 'hujan', 'panas_2', 'payung',
 'lemak', 'panas', 'buat', 'sekolah', 'ayah', 'arah', 'pukul', 'apa', 'tolong',
 'jangan', 'tanya', 'pinjam']

num_classes = len(gestures) # Should be 10
print(f"🔥 Model using HARDCODED order ({num_classes} classes): {gestures}")

# ==========================================
# 3. Model Definition
# ==========================================
class CustomLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(CustomLSTM, self).__init__()
        # Bidirectional LSTM: Input 258, Hidden 64, 2 Layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers=2, batch_first=True, dropout=0.2, bidirectional=True)
        self.bn = nn.BatchNorm1d(hidden_size * 2)
        # Fully Connected Layer
        self.fc = nn.Sequential(
            nn.Linear(hidden_size * 2, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_classes)
        )
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        # Take the output of the last time step
        out = lstm_out[:, -1, :]
        out = self.bn(out)
        out = self.fc(out)
        return out

# ==========================================
# 4. Load Weights
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CustomLSTM(input_size=258, hidden_size=64, num_classes=num_classes).to(device)

MODEL_PATH = 'best_lstm_model_3.pth'

if os.path.exists(MODEL_PATH):
    # Load weights directly. No filtering needed since shape matches (10 classes).
    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
    model.eval()
    print(f"✅ Model Loaded Successfully! (Full Match)")
else:
    print(f"❌ Model not found at {MODEL_PATH}")

# ==========================================
# 5. Prediction Logic
# ==========================================
mp_holistic = mp.solutions.holistic

def extract_keypoints(results):
    """
    Extracts 258 landmarks: Pose (132) + Left Hand (63) + Right Hand (63)
    """
    pose = np.array([[res.x, res.y, res.z, res.visibility] for res in results.pose_landmarks.landmark]).flatten() if results.pose_landmarks else np.zeros(33*4)
    lh = np.array([[res.x, res.y, res.z] for res in results.left_hand_landmarks.landmark]).flatten() if results.left_hand_landmarks else np.zeros(21*3)
    rh = np.array([[res.x, res.y, res.z] for res in results.right_hand_landmarks.landmark]).flatten() if results.right_hand_landmarks else np.zeros(21*3)
    return np.concatenate([pose, lh, rh])

def translate_video(input_video):
    if not input_video:
        return "No video uploaded", "", "", ""

    sequence = []
    cap = cv2.VideoCapture(input_video)

    # Process video frames
    with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = holistic.process(frame_rgb)

            # Only collect frames where hands are detected
            if results.left_hand_landmarks or results.right_hand_landmarks:
                keypoints = extract_keypoints(results)
                sequence.append(keypoints)

    cap.release()

    curr_len = len(sequence)
    if curr_len < 5:
        return "No gesture detected", "", "", ""

    # Padding / Truncating logic to ensure sequence length is 30
    TARGET_LEN = 30
    if curr_len < TARGET_LEN:
        pad_len = TARGET_LEN - curr_len
        zero_frame = np.zeros(258)
        sequence = sequence + [zero_frame] * pad_len
    else:
        sequence = sequence[:TARGET_LEN]

    # Convert to tensor and predict
    inp = torch.tensor(np.array([sequence]), dtype=torch.float32).to(device)

    with torch.no_grad():
        res = model(inp)
        probs = F.softmax(res, dim=1)
        max_prob, max_idx = torch.max(probs, dim=1)

        # Map index to label using the HARDCODED list
        predicted_gesture = gestures[max_idx.item()]
        confidence = max_prob.item()

    print(f"Prediction Index: {max_idx.item()} -> Label: {predicted_gesture} (Conf: {confidence:.2f})")

    # Fetch translation from the large dictionary
    if confidence > 0.4:
        data = translation_dict.get(predicted_gesture, {})
        en = data.get('en', predicted_gesture)
        ms = data.get('ms', predicted_gesture)
        zh = data.get('zh', predicted_gesture)
        ta = data.get('ta', predicted_gesture)
        return en, ms, zh, ta
    else:
        return f"Uncertain ({confidence:.2f})", "Tidak Pasti", "不确定", "நிச்சயமற்ற"

# ==========================================
# 6. Launch Gradio Interface
# ==========================================
with gr.Blocks(title="BIM Sign Language Translator") as demo:
    gr.Markdown("# 🇲🇾 BIM Sign Language Translator (Trilingual: EN/MS/ZH)")
    gr.Markdown("Upload a video file to translate Malaysian Sign Language (BIM) into English, Malay, and Chinese. This app uses MediaPipe and a Bidirectional LSTM model.")
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Upload Video", format="mp4")
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Submit Translation", variant="primary")

        with gr.Column():
            en_output = gr.Textbox(label="English")
            ms_output = gr.Textbox(label="Malay")
            zh_output = gr.Textbox(label="Chinese")
            ta_output = gr.Textbox(label="Tamil")

    submit_btn.click(
        fn=translate_video,
        inputs=video_input,
        outputs=[en_output, ms_output, zh_output, ta_output]
    )
    clear_btn.click(lambda: [None]*5, outputs=[video_input, en_output, ms_output, zh_output, ta_output])


demo.launch(share=True, debug=True)