hf_inference.py

"""
Public VoxelMind inference — no Model.py / Affectors.py required.

Load the TorchScript artifact exported with scripts/export_scripted.py:
    python scripts/export_scripted.py
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Sequence

import cv2
import numpy as np
import torch
import torch.nn.functional as F

TARGET_FPS = 22
TARGET_SIZE = (64, 64)

CLASS_NAMES: list[str] = [
    "s",
    "d",
    "idle",
    "mouse_up",
    "jump",
    "mouse_right",
    "mouse_down",
    "w",
    "a",
    "mouse_left",
    "drop",
]


def frame_to_gray(frame_bgr: np.ndarray, target_size: tuple[int, int] = TARGET_SIZE) -> np.ndarray:
    resized = cv2.resize(frame_bgr, target_size, interpolation=cv2.INTER_AREA)
    return cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)


def read_video_gray(path: str | Path, target_size: tuple[int, int] = TARGET_SIZE) -> list[np.ndarray]:
    cap = cv2.VideoCapture(str(path))
    frames: list[np.ndarray] = []
    while cap.isOpened():
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(frame_to_gray(frame, target_size))
    cap.release()
    if not frames:
        raise ValueError(f"no frames in video: {path}")
    return frames


def crop_temporal(frames: Sequence[np.ndarray], target_fps: int = TARGET_FPS) -> list[np.ndarray]:
    n = len(frames)
    if n >= target_fps:
        start = (n - target_fps) // 2
        return list(frames[start : start + target_fps])
    padded = list(frames)
    while len(padded) < target_fps:
        padded.append(frames[-1])
    return padded[:target_fps]


def frames_to_tensor(
    frames_gray: Sequence[np.ndarray],
    target_fps: int = TARGET_FPS,
    target_size: tuple[int, int] = TARGET_SIZE,
) -> torch.Tensor:
    video = torch.stack([torch.from_numpy(f) for f in frames_gray]).float() / 255.0
    video = video.unsqueeze(0).unsqueeze(0)
    video = F.interpolate(
        video,
        size=(len(frames_gray), target_size[0], target_size[1]),
        mode="trilinear",
        align_corners=False,
    )
    return video.squeeze(0)


def load_model(path: str | Path, device: str | torch.device = "cpu") -> torch.jit.ScriptModule:
    model = torch.jit.load(str(path), map_location=device)
    model.eval()
    return model


def predict_tensor(
    model: torch.jit.ScriptModule,
    clip: torch.Tensor,
    device: str | torch.device = "cpu",
    topk: int = 3,
) -> list[tuple[str, float]]:
    """clip: [1, 22, 64, 64] or [1, 1, 22, 64, 64]"""
    if clip.dim() == 4:
        clip = clip.unsqueeze(0)
    clip = clip.to(device)
    with torch.inference_mode():
        logits = model(clip)
        probs = torch.softmax(logits, dim=-1).squeeze(0)
    k = min(topk, probs.numel())
    values, indices = probs.topk(k)
    return [(CLASS_NAMES[i], float(v)) for v, i in zip(values, indices)]


def predict_video(
    model: torch.jit.ScriptModule,
    video_path: str | Path,
    device: str | torch.device = "cpu",
    topk: int = 3,
) -> list[tuple[str, float]]:
    frames = crop_temporal(read_video_gray(video_path))
    clip = frames_to_tensor(frames)
    return predict_tensor(model, clip, device=device, topk=topk)


def main() -> None:
    p = argparse.ArgumentParser(description="VoxelMind inference (TorchScript, no source model code)")
    p.add_argument("--model", default="models/voxel_scripted.pt", help="TorchScript artifact")
    p.add_argument("--video", required=True, help="Input .mp4 clip")
    p.add_argument("--device", default="cpu", help="cpu | cuda")
    p.add_argument("--topk", type=int, default=3)
    args = p.parse_args()

    model = load_model(args.model, args.device)
    preds = predict_video(model, args.video, device=args.device, topk=args.topk)
    for name, prob in preds:
        print(f"{name:12s} {prob * 100:5.1f}%")


if __name__ == "__main__":
    main()