inference.py

from __future__ import annotations

import os
import warnings
from pathlib import Path
from typing import Union

# Load .env from the inference directory (or parent) before HuggingFace imports
try:
    from dotenv import load_dotenv
    _env = Path(__file__).parent / '.env'
    if not _env.exists():
        _env = Path(__file__).parent.parent / '.env'
    load_dotenv(_env)
except ImportError:
    pass  # python-dotenv not installed; rely on shell environment

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

from model import FERModel, ID_TO_EMOTION, VIT_MEAN, VIT_STD, IMG_SIZE
from detect_face import detect_and_crop_faces

_INFERENCE_TRANSFORM = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=VIT_MEAN, std=VIT_STD),
])


def _resolve_device(device: str) -> torch.device:
    if device == 'auto':
        if torch.cuda.is_available():
            chosen = torch.device('cuda')
            print(f"[INFO] Using GPU: {torch.cuda.get_device_name(0)}")
        else:
            chosen = torch.device('cpu')
            print("[INFO] GPU not available — using CPU.")
        return chosen
    return torch.device(device)


def _to_pil(image_input) -> Image.Image:
    """Accepts file path, PIL Image, numpy array (BGR), or tensor."""
    if isinstance(image_input, str) or isinstance(image_input, Path):
        path = str(image_input)
        if not os.path.exists(path):
            raise FileNotFoundError(f"Image not found: {path}")
        try:
            return Image.open(path).convert('RGB')
        except Exception as e:
            raise ValueError(f"Cannot open image '{path}': {e}") from e

    if isinstance(image_input, Image.Image):
        return image_input.convert('RGB')

    if isinstance(image_input, np.ndarray):
        import cv2
        if image_input.ndim == 3 and image_input.shape[2] == 3:
            rgb = cv2.cvtColor(image_input, cv2.COLOR_BGR2RGB)
        else:
            rgb = image_input
        return Image.fromarray(rgb.astype(np.uint8)).convert('RGB')

    if isinstance(image_input, torch.Tensor):
        # Assume already preprocessed [3,H,W] or [1,3,H,W]; return as-is sentinel
        return image_input

    raise TypeError(f"Unsupported image input type: {type(image_input)}")


def _remap_vit_keys(state: dict) -> dict:
    """
    Translate saved ViT backbone keys from the old HuggingFace transformers layout
    (encoder.layer.N.attention.attention.query) to the new layout
    (layers.N.attention.q_proj) so weights load regardless of transformers version.
    """
    import re

    if not any('backbone.encoder.layer.' in k for k in state):
        return state  # already new format or custom keys — nothing to do

    remapped = {}
    for key, val in state.items():
        k = key

        # 1. backbone.encoder.layer.N.* → backbone.layers.N.*
        k = re.sub(r'backbone\.encoder\.layer\.(\d+)\.', lambda m: f'backbone.layers.{m.group(1)}.', k)

        # 2. Attention projections (most-specific patterns first)
        k = k.replace('.attention.attention.query.', '.attention.q_proj.')
        k = k.replace('.attention.attention.key.',   '.attention.k_proj.')
        k = k.replace('.attention.attention.value.', '.attention.v_proj.')
        k = k.replace('.attention.output.dense.',    '.attention.o_proj.')

        # 3. FFN layers
        k = k.replace('.intermediate.dense.', '.mlp.fc1.')
        k = k.replace('.output.dense.',       '.mlp.fc2.')

        # 4. Drop pooler (we use add_pooling_layer=False)
        if 'backbone.pooler' in k:
            continue

        remapped[k] = val

    return remapped


def _build_result(logits: torch.Tensor) -> dict:
    probs = F.softmax(logits, dim=-1).squeeze()
    probs_np = probs.cpu().numpy()

    top1_idx = int(probs_np.argmax())
    emotion = ID_TO_EMOTION[top1_idx]
    confidence = float(probs_np[top1_idx])

    all_probs = {ID_TO_EMOTION[i]: float(probs_np[i]) for i in range(len(ID_TO_EMOTION))}

    sorted_probs = sorted(all_probs.items(), key=lambda x: x[1], reverse=True)
    top3 = sorted_probs[:3]

    return {
        'emotion': emotion,
        'confidence': confidence,
        'probabilities': all_probs,
        'top3': top3,
    }


class FERPredictor:
    def __init__(self, weights_path: str = '../models/model_weights.pth', device: str = 'auto'):
        self.device = _resolve_device(device)

        weights_path = str(weights_path)
        if not os.path.exists(weights_path):
            raise FileNotFoundError(
                f"Model weights not found at '{weights_path}'. "
                "Place model_weights.pth inside models/ at the project root, or pass --weights."
            )

        print(f"[INFO] Loading model weights from: {weights_path}")
        self.model = FERModel(num_classes=7, num_domains=2)
        state = torch.load(weights_path, map_location=self.device)

        # Support both raw state_dict and full checkpoint dicts
        if isinstance(state, dict) and 'model_state' in state:
            state = state['model_state']

        state = _remap_vit_keys(state)
        self.model.load_state_dict(state)
        self.model.to(self.device)
        self.model.eval()
        print("[INFO] Model loaded successfully.")

        self._transform = _INFERENCE_TRANSFORM
        self._face_detect_device = 'cuda' if self.device.type == 'cuda' else 'cpu'

    @torch.no_grad()
    def predict_image(self, image_input) -> dict:
        """
        Predict emotion for a single (pre-cropped) image.

        Accepts: file path str, PIL Image, numpy array (BGR), or preprocessed tensor.
        Returns dict with keys: emotion, confidence, probabilities, top3.
        """
        raw = _to_pil(image_input)

        if isinstance(raw, torch.Tensor):
            tensor = raw
            if tensor.ndim == 3:
                tensor = tensor.unsqueeze(0)
        else:
            tensor = self._transform(raw).unsqueeze(0)

        tensor = tensor.to(self.device)
        emotion_logits, _ = self.model(tensor)
        return _build_result(emotion_logits)

    @torch.no_grad()
    def predict_batch(self, image_list: list) -> list[dict]:
        """Batch inference over a list of image inputs."""
        tensors = []
        for img in image_list:
            raw = _to_pil(img)
            if isinstance(raw, torch.Tensor):
                t = raw if raw.ndim == 4 else raw.unsqueeze(0)
            else:
                t = self._transform(raw).unsqueeze(0)
            tensors.append(t)

        batch = torch.cat(tensors, dim=0).to(self.device)
        emotion_logits, _ = self.model(batch)

        results = []
        for i in range(emotion_logits.shape[0]):
            results.append(_build_result(emotion_logits[i].unsqueeze(0)))
        return results

    @torch.no_grad()
    def predict_with_face_detection(
        self,
        image_input,
        method: str = 'mtcnn',
        margin: int = 20
    ) -> list[dict]:
        """
        Detect all faces in image, predict emotion for each crop.

        Returns list of dicts with extra keys: bbox (x,y,w,h) and face_index.
        """
        raw = _to_pil(image_input)
        if isinstance(raw, torch.Tensor):
            warnings.warn("Face detection not supported for pre-processed tensors. Running predict_image instead.")
            result = self.predict_image(raw)
            result.update({'bbox': None, 'face_index': 0})
            return [result]

        faces = detect_and_crop_faces(raw, method=method, margin=margin, device=self._face_detect_device)

        results = []
        for idx, (crop, bbox) in enumerate(faces):
            pred = self.predict_image(crop)
            pred['bbox'] = bbox
            pred['face_index'] = idx
            results.append(pred)
        return results