Upload clashcr/models/evidence_model.py with huggingface_hub

clashcr/models/evidence_model.py

@@ -0,0 +1,298 @@
+"""Card evidence model: detect troops, buildings, spells, evolutions, heroes.
+
+Outputs evidence, not just card names:
+- detected units/effects,
+- bounding boxes/masks,
+- spawn time,
+- location,
+- side,
+- possible cards,
+- confidence,
+- ambiguity reason.
+
+For normal live view, we use a YOLO-based unit detector (inspired by KataCR)
+plus heuristic spell/effect detectors for transient effects.
+"""
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class UnitEvidence:
+    unit_name: str
+    bbox: Tuple[int, int, int, int]  # x, y, w, h
+    confidence: float
+    side: str
+    frame_idx: int
+    timestamp: float
+    is_evolution: bool = False
+    is_hero: bool = False
+    is_building: bool = False
+    is_spell_effect: bool = False
+
+
+@dataclass
+class SpellEvidence:
+    spell_name: str
+    effect_mask: np.ndarray = field(repr=False)
+    bbox: Tuple[int, int, int, int]
+    confidence: float
+    side: str
+    frame_idx: int
+    timestamp: float
+    ambiguity_reason: str = ""
+
+
+@dataclass
+class EvidenceBundle:
+    timestamp: float
+    frame_idx: int
+    side: str
+    units: List[UnitEvidence] = field(default_factory=list)
+    spells: List[SpellEvidence] = field(default_factory=list)
+    possible_cards: List[str] = field(default_factory=list)
+    confidence: float = 0.0
+    ambiguity_reason: str = ""
+
+
+class EvidenceModel:
+    """Wraps a YOLO unit detector and adds spell/heuristic effect detection.
+
+    Args:
+        model_path: Path to YOLO .pt file (e.g., KataCR-trained model).
+        img_size: Inference size.
+        conf_threshold: Minimum confidence for unit detections.
+        spell_enabled: Whether to run heuristic spell detectors.
+    """
+
+    # Heuristic spell effect signatures (BGR color ranges in HSV)
+    SPELL_SIGNATURES = {
+        "zap": {
+            "hsv_ranges": [(90, 50, 200), (130, 255, 255)],  # bright blue-white flash
+            "min_area": 300,
+            "max_duration_frames": 8,
+        },
+        "fireball": {
+            "hsv_ranges": [(0, 100, 200), (20, 255, 255)],  # orange-red explosion
+            "min_area": 500,
+            "max_duration_frames": 12,
+        },
+        "arrows": {
+            "hsv_ranges": [(0, 0, 180), (180, 30, 255)],  # white/grey streaks
+            "min_area": 200,
+            "max_duration_frames": 6,
+        },
+        "poison": {
+            "hsv_ranges": [(35, 50, 50), (85, 255, 200)],  # green cloud
+            "min_area": 400,
+            "max_duration_frames": 20,
+        },
+        "freeze": {
+            "hsv_ranges": [(80, 30, 200), (120, 100, 255)],  # icy blue-white
+            "min_area": 300,
+            "max_duration_frames": 15,
+        },
+        "rage": {
+            "hsv_ranges": [(150, 100, 150), (180, 255, 255)],  # purple-pink
+            "min_area": 400,
+            "max_duration_frames": 18,
+        },
+        "tornado": {
+            "hsv_ranges": [(0, 0, 100), (180, 50, 200)],  # grey swirl
+            "min_area": 500,
+            "max_duration_frames": 15,
+        },
+        "earthquake": {
+            "hsv_ranges": [(10, 50, 50), (30, 200, 150)],  # brown cracks
+            "min_area": 600,
+            "max_duration_frames": 15,
+        },
+        "log": {
+            "hsv_ranges": [(10, 50, 80), (30, 200, 180)],  # brown rolling log
+            "min_area": 400,
+            "max_duration_frames": 10,
+        },
+        "barbarian-barrel": {
+            "hsv_ranges": [(10, 50, 80), (30, 200, 180)],  # similar to log
+            "min_area": 400,
+            "max_duration_frames": 10,
+        },
+        "vines": {
+            "hsv_ranges": [(35, 80, 80), (75, 255, 200)],  # green tangling vines
+            "min_area": 300,
+            "max_duration_frames": 20,
+        },
+        "void": {
+            "hsv_ranges": [(120, 50, 20), (160, 255, 80)],  # dark purple/black hole
+            "min_area": 400,
+            "max_duration_frames": 15,
+        },
+    }
+
+    def __init__(self,
+                 model_path: Optional[str] = None,
+                 img_size: int = 640,
+                 conf_threshold: float = 0.5,
+                 spell_enabled: bool = True,
+                 device: str = "cpu"):
+        self.model_path = model_path
+        self.img_size = img_size
+        self.conf_threshold = conf_threshold
+        self.spell_enabled = spell_enabled
+        self.device = device
+        self._model = None
+        self._spell_history: Dict[str, List[int]] = {}  # spell_name -> list of frame indices seen
+        self._frame_idx = 0
+
+        if model_path and Path(model_path).exists():
+            self._load_yolo()
+        else:
+            logger.warning("YOLO model not found at %s; unit detection disabled.", model_path)
+
+    def _load_yolo(self) -> None:
+        try:
+            from ultralytics import YOLO
+            self._model = YOLO(self.model_path)
+            logger.info("Loaded YOLO model from %s", self.model_path)
+        except Exception as e:
+            logger.error("Failed to load YOLO model: %s", e)
+            self._model = None
+
+    def detect_units(self, frame: np.ndarray, side: str) -> List[UnitEvidence]:
+        """Run YOLO unit detection on the frame."""
+        if self._model is None:
+            return []
+
+        results = self._model(frame, imgsz=self.img_size, verbose=False, device=self.device)
+        evidence = []
+        for box in results[0].boxes:
+            conf = float(box.conf)
+            if conf < self.conf_threshold:
+                continue
+            cls_name = self._model.names[int(box.cls)]
+            x1, y1, x2, y2 = box.xyxy[0].tolist()
+            bbox = (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
+
+            # Infer flags from class name
+            is_evolution = "-evolution" in cls_name or "evolution" in cls_name
+            is_hero = cls_name in {
+                "skeleton-king", "golden-knight", "archer-queen",
+                "monk", "mighty-miner", "little-prince", "royal-guardian"
+            }
+            is_building = cls_name in {
+                "cannon", "tesla", "inferno-tower", "bomb-tower",
+                "mortar", "x-bow", "elixir-collector", "furnace",
+                "goblin-hut", "barbarian-hut", "tombstone"
+            }
+            is_spell = cls_name in {
+                "fireball", "zap", "arrows", "poison", "freeze",
+                "rage", "tornado", "earthquake", "the-log",
+                "barbarian-barrel", "clone", "mirror", "royal-delivery",
+                "giant-snowball", "lightning", "rocket", "graveyard"
+            }
+
+            ev = UnitEvidence(
+                unit_name=cls_name,
+                bbox=bbox,
+                confidence=conf,
+                side=side,
+                frame_idx=self._frame_idx,
+                timestamp=time.monotonic(),
+                is_evolution=is_evolution,
+                is_hero=is_hero,
+                is_building=is_building,
+                is_spell_effect=is_spell,
+            )
+            evidence.append(ev)
+        return evidence
+
+    def detect_spells(self, frame: np.ndarray, side: str) -> List[SpellEvidence]:
+        """Heuristic spell effect detection based on color/motion signatures."""
+        if not self.spell_enabled:
+            return []
+
+        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+        evidence = []
+        for spell_name, sig in self.SPELL_SIGNATURES.items():
+            lower, upper = sig["hsv_ranges"]
+            mask = cv2.inRange(hsv, np.array(lower), np.array(upper))
+            # Morphological cleanup
+            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
+            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+
+            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)
+            for i in range(1, num_labels):
+                area = stats[i, cv2.CC_STAT_AREA]
+                if area < sig["min_area"]:
+                    continue
+                x = stats[i, cv2.CC_STAT_LEFT]
+                y = stats[i, cv2.CC_STAT_TOP]
+                w = stats[i, cv2.CC_STAT_WIDTH]
+                h = stats[i, cv2.CC_STAT_HEIGHT]
+
+                # Track duration to avoid re-detecting same spell
+                history = self._spell_history.setdefault(spell_name, [])
+                history.append(self._frame_idx)
+                history[:] = [f for f in history if self._frame_idx - f <= sig["max_duration_frames"]]
+                if len(history) > 1:
+                    # Already detected recently
+                    continue
+
+                ev = SpellEvidence(
+                    spell_name=spell_name,
+                    effect_mask=mask[y:y+h, x:x+w].copy(),
+                    bbox=(x, y, w, h),
+                    confidence=min(area / (sig["min_area"] * 3), 1.0),
+                    side=side,
+                    frame_idx=self._frame_idx,
+                    timestamp=time.monotonic(),
+                    ambiguity_reason="heuristic_color",
+                )
+                evidence.append(ev)
+        return evidence
+
+    def process(self, frame: np.ndarray, side: str = "opponent") -> EvidenceBundle:
+        """Run full evidence extraction on a frame crop."""
+        units = self.detect_units(frame, side)
+        spells = self.detect_spells(frame, side)
+
+        # Build possible cards from evidence
+        possible = set()
+        for u in units:
+            possible.add(u.unit_name)
+        for s in spells:
+            possible.add(s.spell_name)
+
+        ambiguity = ""
+        if not units and not spells:
+            ambiguity = "no_evidence"
+        elif len(possible) > 3:
+            ambiguity = f"too_many_candidates:{len(possible)}"
+
+        bundle = EvidenceBundle(
+            timestamp=time.monotonic(),
+            frame_idx=self._frame_idx,
+            side=side,
+            units=units,
+            spells=spells,
+            possible_cards=sorted(possible),
+            confidence=0.0 if ambiguity else 0.7,
+            ambiguity_reason=ambiguity,
+        )
+        self._frame_idx += 1
+        return bundle
+
+    def reset(self) -> None:
+        self._spell_history.clear()
+        self._frame_idx = 0