text.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
if os.environ.get("XDG_SESSION_TYPE", "").lower() == "wayland" and "QT_QPA_PLATFORM" not in os.environ:
    os.environ["QT_QPA_PLATFORM"] = "xcb"

import sys, json, time, argparse, datetime, subprocess
from collections import deque, Counter
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import mediapipe as mp

mp_hands = mp.solutions.hands

# =====================================================
# 🔹 GROQ API Setup (optional but recommended)
# =====================================================
try:
    from dotenv import load_dotenv
    load_dotenv()
except Exception:
    pass

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
_groq_client = None
if GROQ_API_KEY:
    try:
        from groq import Groq
        _groq_client = Groq(api_key=GROQ_API_KEY)
        print("[Groq] ✅ API client ready")
    except Exception as e:
        print(f"[Groq] ❌ Init error: {e}")
        _groq_client = None
else:
    print("⚠️  GROQ_API_KEY not found in .env (correction/translation will be skipped)")

GROQ_MODELS = [
    "llama-3.1-8b-instant",
    "llama-3.1-70b-versatile",
    "mixtral-8x7b-32768",
    "gemma2-9b-it",
    "whisper-large-v3",
]

SYSTEM_CORRECTOR = (
    "You are an expert linguistic corrector for Indian Sign Language (ISL) gesture output. "
    "Infer the intended English meaning even if the text is phonetically distorted or grammatically broken. "
    "Fix spelling, spacing, grammar, capitalization, and punctuation. "
    "Return ONLY the corrected English sentence."
)

def _groq_try(messages):
    if _groq_client is None:
        return None
    for m in GROQ_MODELS:
        try:
            resp = _groq_client.chat.completions.create(
                model=m, temperature=0.2, max_tokens=400, messages=messages
            )
            out = (resp.choices[0].message.content or "").strip()
            if out:
                return out
        except Exception as e:
            print(f"[Groq] ❌ {m} failed: {e}")
            continue
    return None

def groq_correct_text(raw_text: str) -> str:
    if not raw_text.strip() or _groq_client is None:
        return raw_text
    print(f"[Groq] → Correcting: {raw_text!r}")
    out = _groq_try(
        [{"role": "system", "content": SYSTEM_CORRECTOR},
         {"role": "user", "content": f"Input: {raw_text}\nOutput:"}]
    )
    if out:
        print(f"[Groq] ✅ Corrected: {out}")
        return out
    return raw_text

# =====================================================
# 🔊 TTS + Translation Utilities
# =====================================================

# Canonical names for menu display
LANG_MAP = {
    "gu": "Gujarati",
    "ne": "Nepali",
    "mr": "Marathi",
    "as": "Assamese",     # gTTS support varies
    "hi": "Hindi",
    "pa": "Punjabi",
    "bn": "Bengali",
    "or": "Odia",         # gTTS code 'or'
    "ur": "Urdu",
    "si": "Sinhala",
    "mag": "Magahi",      # may not be supported by gTTS
    "bho": "Bhojpuri",    # may not be supported by gTTS
    "mai": "Maithili",    # may not be supported by gTTS
    "awa": "Awadhi",      # may not be supported by gTTS
    "sd": "Sindhi",       # may not be supported by gTTS
    "ta": "Tamil",
    "te": "Telugu",
    "kn": "Kannada",
    "ml": "Malayalam",
    "en": "English",
}

# Conservatively known/common gTTS language codes
_GTTs_COMMON = {
    "en","hi","bn","gu","kn","ml","mr","ne","pa","ta","te","ur","si","or"
}

def is_lang_supported_by_gtts(code: str) -> bool:
    return code in _GTTs_COMMON

def groq_translate(text: str, target_lang_code: str) -> str:
    if not text.strip():
        return text
    if _groq_client is None:
        print("[Groq] ⚠️ No API key; skipping translation.")
        return text
    lang_name = LANG_MAP.get(target_lang_code, target_lang_code)
    print(f"[Groq] → Translate to {lang_name} ({target_lang_code})")
    prompt = (
        f"Translate the following sentence into {lang_name}. "
        f"Use natural, fluent {lang_name}. Return ONLY the translated sentence.\n\nText: {text}"
    )
    out = _groq_try([
        {"role": "system", "content": "You are an expert translator for Indian languages."},
        {"role": "user", "content": prompt}
    ])
    if out:
        print(f"[Groq] ✅ Translation: {out}")
        return out
    return text

def tts_gtts_save(text: str, lang_code: str, path: str) -> bool:
    try:
        from gtts import gTTS
    except Exception as e:
        print(f"[TTS] ❌ gTTS not installed: {e}. Install with: pip install gTTS")
        return False
    try:
        tts = gTTS(text=text, lang=lang_code)
        tts.save(path)
        print(f"[TTS] 💾 Saved: {path}")
        return True
    except Exception as e:
        print(f"[TTS] ❌ Failed to synthesize with gTTS ({lang_code}): {e}")
        return False

def open_audio_file(path: str):
    try:
        if sys.platform.startswith("win"):
            os.startfile(path)  # type: ignore
        elif sys.platform == "darwin":
            subprocess.Popen(["open", path])
        else:
            subprocess.Popen(["xdg-open", path])
        print(f"[TTS] ▶️  Playing: {path}")
    except Exception as e:
        print(f"[TTS] ⚠️ Could not auto-open audio: {e}")

def ensure_lang_or_fallback(requested: str) -> str:
    """Prefer requested if supported by gTTS; else fallback to hi -> en."""
    if is_lang_supported_by_gtts(requested):
        return requested
    print(f"[TTS] ⚠️ '{requested}' not in common gTTS list; falling back.")
    return "hi" if requested != "hi" and is_lang_supported_by_gtts("hi") else "en"

def timestamp_str(): 
    return datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

def top1(probs):
    i = int(np.argmax(probs)); 
    return i, float(probs[i])

def majority(q): 
    return Counter(q).most_common(1)[0][0] if q else None

# ----- Interactive prompts -----
def _menu_lines_for_langs():
    ordered = [
        ("gu", "Gujarati"), ("ne", "Nepali"), ("mr", "Marathi"), ("as", "Assamese"),
        ("hi", "Hindi"), ("pa", "Punjabi"), ("bn", "Bengali"), ("or", "Odia"),
        ("ur", "Urdu"), ("si", "Sinhala"), ("mag", "Magahi"), ("bho", "Bhojpuri"),
        ("mai", "Maithili"), ("awa", "Awadhi"), ("sd", "Sindhi"),
        ("ta", "Tamil"), ("te", "Telugu"), ("kn", "Kannada"), ("ml", "Malayalam"),
        ("en", "English"),
    ]
    lines = []
    for code, name in ordered:
        sup = "✅" if is_lang_supported_by_gtts(code) else "⚠️"
        lines.append((code, f"{name} [{code}] {sup}"))
    return lines

def print_language_menu():
    print("\n=== Choose TTS Languages ===")
    print("Enter comma-separated codes or names (e.g., hi,gu,English).")
    print("We’ll translate from corrected English before TTS when needed.")
    print("Legend: ✅ commonly supported by gTTS, ⚠️ may fall back to Hindi/English.\n")
    for _code, line in _menu_lines_for_langs():
        print(f"  {line}")
    print("\nExamples:")
    print("  hi               -> Hindi only")
    print("  gu, mr, pa       -> Gujarati + Marathi + Punjabi")
    print("  Urdu, Bengali    -> using names is fine")

def _normalize_to_code(token: str) -> str | None:
    t = token.strip().lower()
    if not t: return None
    if t in LANG_MAP: 
        return t
    for code, name in LANG_MAP.items():
        if t == name.lower(): 
            return code
    return None

def prompt_for_languages() -> list[str]:
    print_language_menu()
    raw = input("\nEnter languages: ").strip()
    chosen = []
    for part in raw.split(","):
        code = _normalize_to_code(part)
        if code and code not in chosen:
            chosen.append(code)
    if not chosen:
        print("[TTS] No valid entries; defaulting to Hindi (hi).")
        chosen = ["hi"]
    resolved = []
    for c in chosen:
        c2 = ensure_lang_or_fallback(c)
        if c2 != c:
            print(f"[TTS] '{c}' not well-supported; using '{c2}' for synthesis.")
        resolved.append(c2)
    return resolved

def prompt_for_what_to_speak() -> str:
    print("\nWhat should be spoken?")
    print("  1) corrected (default)")
    print("  2) raw")
    print("  3) both")
    ans = input("Choose 1/2/3: ").strip()
    return {"2":"raw", "3":"both"}.get(ans, "corrected")

def prompt_auto_play() -> bool:
    ans = input("\nAuto-play audio after saving? [Y/n]: ").strip().lower()
    return ans in ("", "y", "yes")

# =====================================================
# 🖐️ MediaPipe Hand Extractor
# =====================================================
class HandExtractor:
    def __init__(self, max_hands=2, det_conf=0.5, track_conf=0.5):
        self.h = mp_hands.Hands(False, max_hands, 1, det_conf, track_conf)
    def from_bgr(self, img):
        rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        res = self.h.process(rgb)
        all_lm = []
        if res.multi_hand_landmarks:
            for hnd in res.multi_hand_landmarks[:2]:
                all_lm.append(np.array([[lm.x, lm.y, lm.z] for lm in hnd.landmark], np.float32))
        while len(all_lm) < 2:
            all_lm.append(np.zeros((21, 3), np.float32))
        lm = np.concatenate(all_lm)
        lm[:, 2] = np.clip(lm[:, 2], -0.2, 0.2) / 0.2
        return lm.flatten()

# =====================================================
# 🧠 MLP Model
# =====================================================
class MLP(nn.Module):
    def __init__(self, in_dim, num_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, 512), nn.LayerNorm(512), nn.GELU(),
            nn.Dropout(0.2), nn.Linear(512, 256), nn.GELU(),
            nn.Dropout(0.2), nn.Linear(256, num_classes)
        )
    def forward(self, x): 
        return self.net(x)

def load_ckpt(path, max_hands):
    ckpt = torch.load(path, map_location="cpu")
    classes = ckpt.get("classes")
    if not classes:
        with open(os.path.join(os.path.dirname(path), "label_map.json")) as f:
            idx2name = json.load(f)
        classes = [idx2name[str(i)] for i in sorted(map(int, idx2name))]
    model = MLP(21*max_hands*3, len(classes))
    model.load_state_dict(ckpt["model"] if "model" in ckpt else ckpt)
    model.eval()
    return model, classes

@torch.no_grad()
def predict_one(model, feat):
    x = torch.from_numpy(feat[None, :]).float()
    return F.softmax(model(x), dim=1).cpu().numpy()[0]

# =====================================================
# 💾 Final Save and Exit (interactive TTS)
# =====================================================
def finalize_and_exit(cap, show_gui, final_text, args):
    base = f"isl_output_{timestamp_str()}"
    save_path = f"{base}.txt"
    try:
        with open(save_path, "w", encoding="utf-8") as f:
            f.write(final_text + "\n")
        print(f"\n📝 RAW Detected: {final_text}\n[SAVED] {save_path}")
    except Exception as e:
        print(f"[ERROR] Saving raw failed: {e}")

    corrected = groq_correct_text(final_text)
    corrected_path = f"{base}_corrected.txt"
    try:
        with open(corrected_path, "w", encoding="utf-8") as f:
            f.write(corrected + "\n")
        print(f"✅ Corrected: {corrected}\n[SAVED] {corrected_path}")
    except Exception as e:
        print(f"[ERROR] Saving corrected failed: {e}")

    # === INTERACTIVE TTS ===
    try:
        langs = prompt_for_languages()             # e.g., ["hi","gu","mr"]
        speak_choice = prompt_for_what_to_speak()  # "raw" | "corrected" | "both"
        do_play = prompt_auto_play()

        to_do = []
        if speak_choice in ("raw","both") and final_text.strip():
            to_do.append(("raw", final_text.strip()))
        if speak_choice in ("corrected","both") and corrected.strip():
            to_do.append(("corrected", corrected.strip()))

        for tgt in langs:
            for tag, text0 in to_do:
                speak_text = text0
                if tgt != "en":
                    speak_text = groq_translate(text0, tgt)
                mp3_path = f"{base}_{tag}_{tgt}.mp3"
                ok = tts_gtts_save(speak_text, tgt, mp3_path)
                if ok and do_play:
                    open_audio_file(mp3_path)
    except KeyboardInterrupt:
        print("\n[TTS] Skipped by user.")
    except Exception as e:
        print(f"[TTS] ⚠️ TTS step skipped due to error: {e}")

    if cap: cap.release()
    if show_gui: cv2.destroyAllWindows()
    sys.exit(0)

# =====================================================
# 🔤 Spelling Mode with Multiple Words & Digit Display
# =====================================================
def run_spell(args):
    print(f"[OK] Using checkpoint: {args.ckpt}")
    model, classes = load_ckpt(args.ckpt, args.max_hands)
    extractor = HandExtractor(args.max_hands)

    cap = cv2.VideoCapture(args.camera)
    if not cap.isOpened():
        print("❌ Camera not found")
        return

    win = "ISL Speller (Multi-word)"
    cv2.namedWindow(win, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(win, 960, 540)

    label_q, t_dwell = deque(maxlen=args.smooth), None
    state, target_len, word, sentence = "NUM", None, "", ""

    print("[INFO] SPACE=confirm word | ENTER=save+Groq(+TTS) | q=quit | b=backspace | r=reset")

    while True:
        ok, frame = cap.read()
        if not ok: 
            continue
        frame = cv2.flip(frame, 1)
        feat = extractor.from_bgr(frame)
        probs = predict_one(model, feat)
        idx, p = top1(probs)
        cls = classes[idx]
        label_q.append(cls)
        mv = majority(label_q)

        # HUD
        cv2.putText(frame, f"STATE: {state}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)
        cv2.putText(frame, f"Sentence: {sentence}{word}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)
        if state == "NUM":
            cv2.putText(frame, f"Detected Digit: {mv}", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (100,255,255), 2)
        elif state == "LETTERS":
            cv2.putText(frame, f"Word: {word} (target {target_len})", (10, 90),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,200), 2)

        # Logic
        if state == "NUM":
            if mv and isinstance(mv, str) and mv.isdigit() and p > args.commit_conf:
                if t_dwell is None: 
                    t_dwell = time.time()
                elif (time.time()-t_dwell)*1000 > args.dwell_ms:
                    target_len = int(mv)
                    word = ""
                    state = "LETTERS"
                    label_q.clear()
                    t_dwell = None
            else:
                t_dwell = None

        elif state == "LETTERS":
            if mv and isinstance(mv, str) and mv.isalpha() and p > args.commit_conf and (len(word) < (target_len or 9)):
                if t_dwell is None: 
                    t_dwell = time.time()
                elif (time.time()-t_dwell)*1000 > args.dwell_ms:
                    word += mv
                    label_q.clear()
                    t_dwell = None
            else:
                t_dwell = None

        # Display + key handling
        cv2.imshow(win, frame)
        k = cv2.waitKey(1) & 0xFF
        if k == ord('q'):
            break
        elif k == ord('r'):
            state, target_len, word, sentence = "NUM", None, "", ""
            label_q.clear(); t_dwell = None
        elif k == ord('b'):
            if word: 
                word = word[:-1]
            elif sentence: 
                sentence = sentence[:-1]
        elif k == ord(' '):  # confirm word
            if word:
                sentence += word + " "
                word = ""
                state, target_len = "NUM", None
        elif k in (10, 13):  # ENTER
            final_text = (sentence + word).strip()
            finalize_and_exit(cap, True, final_text, args)

# =====================================================
# 🏁 CLI
# =====================================================
def main():
    ap = argparse.ArgumentParser("ISL + Groq Auto-Corrector (Multi-word) + Interactive Multilingual TTS")
    ap.add_argument("mode", choices=["spell"])
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--camera", type=int, default=0)
    ap.add_argument("--smooth", type=int, default=7)
    ap.add_argument("--commit_conf", type=float, default=0.6)
    ap.add_argument("--dwell_ms", type=int, default=550)
    ap.add_argument("--max_hands", type=int, default=2)
    # Flags kept for future non-interactive mode (currently not used)
    ap.add_argument("--tts_lang", type=str, default="hi")
    ap.add_argument("--speak", type=str, default="corrected", choices=["none","raw","corrected","both"])
    ap.add_argument("--no_play", action="store_true")
    args = ap.parse_args()

    if args.mode == "spell":
        run_spell(args)

if __name__ == "__main__":
    main()