DirtyLeague_Bot/card_embedder.py

"""
DirtyLeague Card Vector Embedding and Recognition Module.

Uses a lightweight MobileNetV3-small ONNX feature extractor (3.5 MB) to convert
card crops into 576-dimensional normalized embedding vectors.
Enables instant cosine similarity matching (<1ms) against 600+ creatures
without retraining the neural network.
"""

import os
from typing import Dict, List, Optional, Tuple
import cv2
import numpy as np
import onnxruntime as ort


# Standard ImageNet normalization parameters
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)


class CardEmbedder:
    def __init__(
        self,
        model_path: str = "models/card_embedder.onnx",
        db_path: str = "models/creatures_db.npz",
    ):
        if not os.path.exists(model_path):
            alt_path = os.path.join(os.path.dirname(__file__), model_path)
            if os.path.exists(alt_path):
                model_path = alt_path
            else:
                raise FileNotFoundError(f"Card embedder model not found at: {model_path}")

        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        sess_options.intra_op_num_threads = 2
        self.session = ort.InferenceSession(model_path, sess_options, providers=["CPUExecutionProvider"])
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

        self.db_path = db_path
        self.names: List[str] = []
        self.embeddings: Optional[np.ndarray] = None  # Shape: [N, 576]
        self.load_db()

    def load_db(self, path: Optional[str] = None):
        """Loads creature vector database from .npz file."""
        target_path = path or self.db_path
        if not os.path.isabs(target_path):
            target_path = os.path.join(os.path.dirname(__file__), target_path)

        if os.path.exists(target_path):
            try:
                data = np.load(target_path)
                self.names = list(data["names"])
                self.embeddings = data["embeddings"]
                print(f"[CardEmbedder] Loaded {len(self.names)} creature vectors from {target_path}")
            except Exception as e:
                print(f"[CardEmbedder] Failed to load DB: {e}")
                self.names = []
                self.embeddings = None
        else:
            self.names = []
            self.embeddings = None

    def save_db(self, path: Optional[str] = None):
        """Saves creature vector database to .npz file."""
        target_path = path or self.db_path
        if not os.path.isabs(target_path):
            target_path = os.path.join(os.path.dirname(__file__), target_path)

        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        if self.embeddings is not None and len(self.names) > 0:
            np.savez_compressed(
                target_path,
                names=np.array(self.names),
                embeddings=self.embeddings,
            )
            print(f"[CardEmbedder] Saved {len(self.names)} creature vectors to {target_path}")

    def embed(self, card_crop_bgr: np.ndarray) -> np.ndarray:
        """
        Extracts a 576-dimensional L2-normalized feature vector from a card crop.
        Takes ~0.8ms on CPU.
        """
        if card_crop_bgr is None or card_crop_bgr.size == 0:
            return np.zeros(576, dtype=np.float32)

        # 1. Resize to 128x128
        resized = cv2.resize(card_crop_bgr, (128, 128), interpolation=cv2.INTER_LINEAR)

        # 2. BGR -> RGB, normalize to [0, 1]
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0

        # 3. Standard ImageNet normalization
        normalized = (rgb - IMAGENET_MEAN) / IMAGENET_STD

        # 4. HWC -> CHW -> NCHW
        blob = np.transpose(normalized, (2, 0, 1))
        blob = np.expand_dims(blob, axis=0)

        # 5. Run ONNX model
        out = self.session.run([self.output_name], {self.input_name: blob})[0]
        feat = out[0]  # Shape: (576,)

        # 6. L2 Normalize vector for cosine distance
        norm = np.linalg.norm(feat)
        if norm > 1e-6:
            feat = feat / norm
        return feat

    def identify_card(
        self, card_crop_bgr: np.ndarray, min_similarity: float = 0.65
    ) -> Tuple[Optional[str], float]:
        """
        Matches a card crop against the reference vector database.
        Returns: (creature_name, similarity_score) or (None, score) if below threshold.
        """
        if self.embeddings is None or len(self.names) == 0:
            return (None, 0.0)

        query_vec = self.embed(card_crop_bgr)  # Shape: (576,)

        # Dot product with L2-normalized vectors = cosine similarity
        similarities = np.dot(self.embeddings, query_vec)
        best_idx = int(np.argmax(similarities))
        best_score = float(similarities[best_idx])

        if best_score >= min_similarity:
            return (self.names[best_idx], best_score)
        return (None, best_score)

    def register_creature(self, name: str, card_crop_bgr: np.ndarray) -> np.ndarray:
        """
        Registers or updates a creature's vector embedding.
        """
        vec = self.embed(card_crop_bgr)

        if self.embeddings is None or len(self.names) == 0:
            self.names = [name]
            self.embeddings = np.expand_dims(vec, axis=0)
        else:
            if name in self.names:
                idx = self.names.index(name)
                # Running average or replace
                self.embeddings[idx] = vec
            else:
                self.names.append(name)
                self.embeddings = np.vstack([self.embeddings, vec])

        return vec