model.py

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

# -------- Wayland-safe UI: force Qt to X11 (xcb) if session is Wayland ------
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, glob, argparse, datetime
from collections import deque, Counter
from typing import List
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F

# ---------- MediaPipe ----------
import mediapipe as mp
mp_hands = mp.solutions.hands

SUPPORTED = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

# =================== Small utils ===================
def top1(probs: np.ndarray):
    i = int(np.argmax(probs))
    return i, float(probs[i])

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

def list_images(root: str) -> List[str]:
    out = []
    for p in glob.glob(os.path.join(root, "**", "*"), recursive=True):
        if os.path.splitext(p)[1].lower() in SUPPORTED:
            out.append(p)
    return out

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

# =================== Landmark extractor ===================
class HandExtractor:
    """
    Returns flattened (21*max_hands*3,) features; z clipped to [-0.2,0.2] then scaled to [-1,1].
    """
    def __init__(self, max_hands=2, det_conf=0.5, track_conf=0.5):
        self.max_hands = max_hands
        self.h = mp_hands.Hands(static_image_mode=False,
                                max_num_hands=max_hands,
                                model_complexity=1,
                                min_detection_confidence=det_conf,
                                min_tracking_confidence=track_conf)

    def from_bgr(self, img_bgr: np.ndarray) -> np.ndarray:
        rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        res = self.h.process(rgb)
        all_lm = []
        if res.multi_hand_landmarks:
            for hand_lms in res.multi_hand_landmarks[:self.max_hands]:
                coords = [[lm.x, lm.y, lm.z] for lm in hand_lms.landmark]
                all_lm.append(np.array(coords, dtype=np.float32))
        k = len(all_lm)
        for _ in range(self.max_hands - k):
            all_lm.append(np.zeros((21,3), dtype=np.float32))
        lm = np.concatenate(all_lm, axis=0)
        lm[:,2] = np.clip(lm[:,2], -0.2, 0.2) / 0.2
        return lm.reshape(-1)

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

def load_ckpt(ckpt_path: str, max_hands: int):
    if not os.path.exists(ckpt_path):
        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location="cpu")
    classes = ckpt.get("classes", None)
    if classes is None:
        lm_path = os.path.join(os.path.dirname(ckpt_path), "label_map.json")
        if os.path.exists(lm_path):
            with open(lm_path, "r") as f:
                idx2name = json.load(f)
            classes = [idx2name[str(i)] for i in sorted(map(int, idx2name.keys()))] \
                      if isinstance(idx2name, dict) else idx2name
        else:
            raise RuntimeError("No class list in ckpt and missing label_map.json.")
    num_classes = len(classes)
    in_dim = 21 * max_hands * 3
    model = MLP(in_dim=in_dim, num_classes=num_classes, hidden=512, dropout=0.2)
    state = ckpt["model"] if "model" in ckpt else ckpt
    model.load_state_dict(state, strict=True)
    model.eval()
    return model, classes

@torch.no_grad()
def predict_one(model: nn.Module, feat: np.ndarray) -> np.ndarray:
    x = torch.from_numpy(feat[None, :]).float()
    logits = model(x)
    return F.softmax(logits, dim=1).cpu().numpy()[0]

def label_of_path(p: str):
    return os.path.basename(os.path.dirname(p))

# =================== Class filters ===================
DIGITS_1_9 = {str(d) for d in range(1, 10)}
LETTERS_A_Z = {chr(c) for c in range(ord('A'), ord('Z')+1)}

def is_digit_label(lbl: str) -> bool:
    return lbl in DIGITS_1_9

def is_letter_label(lbl: str) -> bool:
    return lbl in LETTERS_A_Z

# =================== Window helpers ===================
def ensure_window(win_name: str, width: int = 960, height: int = 540):
    # Create a normal resizable window and size it; this avoids “toolbar-only” repaint issues on some Qt builds.
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)   # not GUI_EXPANDED
    cv2.resizeWindow(win_name, width, height)
    # keep ratio on resize if available (OpenCV >= 4.5.1)
    try:
        cv2.setWindowProperty(win_name, cv2.WND_PROP_KEEP_RATIO, 1)
    except Exception:
        pass

# =================== Camera opener (robust + warmup) ===================
def _set_common_camera_props(cap, w=1280, h=720, fps=30):
    # Many Linux webcams need MJPG for >10fps
    try:
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
    except Exception:
        pass
    cap.set(cv2.CAP_PROP_CONVERT_RGB, 1)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH,  w)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)
    cap.set(cv2.CAP_PROP_FPS, fps)

def _warmup_and_validate(cap, tries=10):
    ok_any, last_shape = False, None
    for _ in range(tries):
        ok, frame = cap.read()
        if ok and frame is not None and frame.size != 0:
            ok_any = True
            last_shape = frame.shape
            break
        time.sleep(0.02)
    return ok_any, last_shape

def open_camera(device_index: int):
    # 1) V4L2
    cap = cv2.VideoCapture(device_index, cv2.CAP_V4L2)
    if cap.isOpened():
        _set_common_camera_props(cap)
        ok, shp = _warmup_and_validate(cap)
        if ok: return cap
        cap.release()

    # 2) ANY
    cap = cv2.VideoCapture(device_index, cv2.CAP_ANY)
    if cap.isOpened():
        _set_common_camera_props(cap)
        ok, shp = _warmup_and_validate(cap)
        if ok: return cap
        cap.release()

    # 3) GStreamer fallback (if device exists)
    dev = f"/dev/video{device_index}"
    if os.path.exists(dev):
        pipeline = (
            f"v4l2src device={dev} ! "
            f"image/jpeg,framerate=30/1 ! jpegdec ! "
            f"videoconvert ! video/x-raw,format=BGR ! "
            f"appsink drop=1 sync=false"
        )
        cap = cv2.VideoCapture(pipeline, cv2.CAP_GSTREAMER)
        if cap.isOpened():
            ok, shp = _warmup_and_validate(cap)
            if ok: return cap
            cap.release()

    return None

# =================== Modes ===================
def run_live(args):
    print(f"[OK] Using checkpoint: {args.ckpt}")
    model, classes = load_ckpt(args.ckpt, args.max_hands)
    extractor = HandExtractor(args.max_hands, args.det_conf, args.track_conf)

    cap = open_camera(args.camera)
    if cap is None or not cap.isOpened():
        print("ERROR: camera not available. Try:")
        print("  - ls /dev/video*  (pick the right --camera index)")
        print("  - Close other apps using the camera (Zoom/Teams/Browser)")
        print("  - QT_QPA_PLATFORM=xcb python model.py ...")
        print("  - sudo apt-get install qtwayland5 libgtk-3-0 gstreamer1.0-libav")
        return

    show_gui = not getattr(args, "no_gui", False)
    if show_gui:
        ensure_window("ISL Live (MediaPipe + MLP)")

    smooth_q = deque(maxlen=args.smooth)
    t_last = time.time()

    while True:
        ok, frame = cap.read()
        if not ok or frame is None or frame.size == 0:
            continue
        frame = cv2.flip(frame, 1)

        feat = extractor.from_bgr(frame)
        probs = predict_one(model, feat)
        idx, p = top1(probs)
        cls = classes[idx]

        smooth_q.append(cls)
        mv = majority(smooth_q)

        if show_gui:
            cv2.putText(frame, f"Pred: {cls} ({p:.2f})  MV:{mv}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,255,255), 2)
            now = time.time()
            fps = 1.0 / max(1e-6, now - t_last); t_last = now
            cv2.putText(frame, f"FPS: {fps:.1f}", (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (200,255,200), 2)

            cv2.imshow("ISL Live (MediaPipe + MLP)", frame)
            if (cv2.waitKey(1) & 0xFF) == ord('q'):
                break

    cap.release()
    if show_gui:
        cv2.destroyAllWindows()

def run_spell(args):
    """
    Flow:
      - NUM: model must predict digit 1–9 → sets target_len
      - LETTERS: only A–Z; add letters until target_len (or press SPACE sooner)
      - SPACE confirms word, returns to NUM
      - ENTER saves full phrase to .txt and exits
    """
    print(f"[OK] Using checkpoint: {args.ckpt}")
    model, classes = load_ckpt(args.ckpt, args.max_hands)
    extractor = HandExtractor(args.max_hands, args.det_conf, args.track_conf)

    cap = open_camera(args.camera)
    if cap is None or not cap.isOpened():
        print("ERROR: camera not available. Try:")
        print("  - ls /dev/video*  (pick the right --camera index)")
        print("  - Close other apps using the camera (Zoom/Teams/Browser)")
        print("  - QT_QPA_PLATFORM=xcb python model.py ...")
        print("  - sudo apt-get install qtwayland5 libgtk-3-0 gstreamer1.0-libav")
        return

    save_path = args.save_path if args.save_path else f"isl_output_{timestamp_str()}.txt"
    show_gui = not getattr(args, "no_gui", False)
    win = "ISL NUM→LETTERS Speller (MediaPipe + MLP)"
    if show_gui:
        ensure_window(win)

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

    print("[INFO] Keys: q=quit without saving, SPACE=confirm word, ENTER=finish & save, b=backspace, r=reset all")

    while True:
        ok, frame = cap.read()
        if not ok or frame is None or frame.size == 0:
            continue
        frame = cv2.flip(frame, 1)
        H, W = frame.shape[:2]

        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)

        if show_gui:
            cv2.putText(frame, f"STATE: {state}", (10, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,255,255), 2)
            cv2.putText(frame, f"Sentence: {sentence}{word}", (10, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)
            cv2.putText(frame, f"Pred: {cls} ({p:.2f}) | MV: {mv}", (10, 88), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)
            if target_len is not None:
                cv2.putText(frame, f"Target letters: {target_len} | Current: {len(word)}", (10, 118),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)

        if state == "NUM":
            valid = (mv == cls) and (p >= args.commit_conf) and is_digit_label(cls)
            if valid:
                if t_dwell is None:
                    t_dwell = time.time()
                elif (time.time() - t_dwell)*1000 >= args.dwell_ms:
                    target_len = int(cls)
                    word = ""
                    label_q.clear()
                    t_dwell = None
                    state = "LETTERS"
            else:
                t_dwell = None

            if show_gui:
                cv2.putText(frame, "Show a digit (1–9) to set word length", (10, 148),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)

        elif state == "LETTERS":
            can_append = (target_len is None) or (len(word) < target_len)
            valid = (mv == cls) and (p >= args.commit_conf) and is_letter_label(cls) and can_append

            if valid:
                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

            if show_gui:
                guide = "Press [SPACE] to confirm word"
                if target_len is not None:
                    guide += f" (limit {target_len})"
                cv2.putText(frame, guide, (10, 148),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2)

        k = 0
        if show_gui:
            cv2.putText(frame, "Keys: [q] Quit  [r] Reset  [b] Backspace  [SPACE] Confirm  [ENTER] Save & Exit",
                        (10, H-14), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255,255,255), 2)
            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 and sentence[-1] != " ":
                sentence = sentence[:-1]
        elif k == ord(' '):
            if word:
                if sentence and not sentence.endswith(" "): sentence += " "
                sentence += word + " "
                word = ""
            state, target_len = "NUM", None
            label_q.clear(); t_dwell = None
        elif k in (10, 13):  # Enter
            final_text = (sentence + word).strip()
            try:
                with open(save_path, "w", encoding="utf-8") as f:
                    f.write(final_text + "\n")
                print("\nFINAL OUTPUT:")
                print(final_text)
                print(f"[SAVED] {save_path}")
            except Exception as e:
                print(f"[ERROR] Could not save to {save_path}: {e}")
            break

        if not show_gui:
            # headless mode: no key handling; Ctrl+C to stop
            pass

    cap.release()
    if show_gui:
        cv2.destroyAllWindows()

def run_image(args):
    print(f"[OK] Using checkpoint: {args.ckpt}")
    model, classes = load_ckpt(args.ckpt, args.max_hands)
    extractor = HandExtractor(args.max_hands, args.det_conf, args.track_conf)
    img = cv2.imread(args.image_path, cv2.IMREAD_COLOR)
    if img is None:
        print(f"Cannot read: {args.image_path}"); return
    feat = extractor.from_bgr(img)
    probs = predict_one(model, feat)
    idx, p = top1(probs)
    print(f"Top1: {classes[idx]} ({p:.4f})")
    for i, c in sorted(enumerate(classes), key=lambda z: probs[z[0]], reverse=True)[:5]:
        print(f"{c:>8s}: {probs[i]:.4f}")

def run_folder(args):
    print(f"[OK] Using checkpoint: {args.ckpt}")
    model, classes = load_ckpt(args.ckpt, args.max_hands)
    extractor = HandExtractor(args.max_hands, args.det_conf, args.track_conf)
    paths = list_images(args.folder)
    if not paths:
        print("No images found."); return
    correct = 0; total = 0
    for p in paths:
        img = cv2.imread(p, cv2.IMREAD_COLOR)
        if img is None: continue
        feat = extractor.from_bgr(img)
        probs = predict_one(model, feat)
        idx, _ = top1(probs)
        pred = classes[idx]
        gt = label_of_path(p) if args.use_path_label else None
        if gt is not None:
            total += 1; correct += int(pred == gt)
            print(f"{p} -> {pred}  (gt={gt})")
        else:
            print(f"{p} -> {pred}")
    if args.use_path_label and total > 0:
        print(f"\nAccuracy: {correct}/{total} = {correct/total:.4f}")

# =================== CLI ===================
def build_parser():
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--ckpt", required=True, help="Path to runs_simple/best.pt")
    common.add_argument("--max_hands", type=int, default=2)
    common.add_argument("--det_conf", type=float, default=0.5)
    common.add_argument("--track_conf", type=float, default=0.5)

    ap = argparse.ArgumentParser("Test your ISL (0-9, A-Z) MLP trained on MediaPipe landmarks")
    sub = ap.add_subparsers(dest="mode", required=True)

    sp_live = sub.add_parser("live", parents=[common], help="Live webcam classification", add_help=True)
    sp_live.add_argument("--camera", type=int, default=0)
    sp_live.add_argument("--smooth", type=int, default=7)
    sp_live.add_argument("--no_gui", action="store_true", help="Run without OpenCV windows (headless)")

    sp_spell = sub.add_parser("spell", parents=[common], help="Digit→Letters speller (1-9 then A–Z)", add_help=True)
    sp_spell.add_argument("--camera", type=int, default=0)
    sp_spell.add_argument("--smooth", type=int, default=7)
    sp_spell.add_argument("--commit_conf", type=float, default=0.75)
    sp_spell.add_argument("--dwell_ms", type=int, default=550)
    sp_spell.add_argument("--save_path", type=str, default=None,
                          help="Path to save final text when pressing ENTER (default: isl_output_<timestamp>.txt)")
    sp_spell.add_argument("--no_gui", action="store_true", help="Run without OpenCV windows (headless)")

    sp_img = sub.add_parser("image", parents=[common], help="Single image", add_help=True)
    sp_img.add_argument("--image_path", required=True)

    sp_dir = sub.add_parser("folder", parents=[common], help="Folder of images (recursive)", add_help=True)
    sp_dir.add_argument("--folder", required=True)
    sp_dir.add_argument("--use_path_label", action="store_true",
                        help="If images are in .../<CLASS>/<file>, compute accuracy")
    return ap

def main():
    ap = build_parser()
    args = ap.parse_args()
    if not os.path.exists(args.ckpt):
        raise FileNotFoundError(f"Checkpoint not found: {args.ckpt}")
    if args.mode == "live":   run_live(args)
    elif args.mode == "spell": run_spell(args)
    elif args.mode == "image": run_image(args)
    elif args.mode == "folder": run_folder(args)
    else:
        print("Unknown mode")

if __name__ == "__main__":
    main()