src/detector.py

import os
from dataclasses import dataclass
from typing import List

import numpy as np
from PIL import Image

from threat_kb import CLASSES
from dataset_generator import CLASS_COLORS  # reaproveita a paleta usada na geração sintética


@dataclass
class Detection:
    label: str
    confidence: float
    box: tuple  # (x0, y0, x1, y1)


class ComponentDetector:
    def __init__(self, model_path: str = "../models/best.pt", conf_threshold: float = 0.35):
        self.conf_threshold = conf_threshold
        self.model = None
        self.demo_mode = True

        if model_path and os.path.exists(model_path):
            try:
                from ultralytics import YOLO
                self.model = YOLO(model_path)
                self.demo_mode = False
            except Exception as e:
                print(f"[detector] Falha ao carregar modelo treinado ({e}). Usando modo demo.")

    def predict(self, image: Image.Image) -> List[Detection]:
        if self.model is not None:
            return self._predict_yolo_tiled(image)
        return self._predict_heuristic(image)

    # ---------------------------------------------------------------
    # Caminho principal: modelo supervisionado treinado (YOLOv8),
    # com inferência por "tiles" (janelas deslizantes) para diagramas
    # grandes/densos — evita perder ícones pequenos ao redimensionar a
    # imagem inteira para o tamanho de treino (ex.: diagramas reais de
    # nuvem, com dezenas de ícones pequenos numa imagem grande).
    # ---------------------------------------------------------------
    def _predict_yolo_tiled(self, image: Image.Image, tile_size: int = 640,
                             overlap: float = 0.2, tile_threshold: int = 800) -> List[Detection]:
        img_w, img_h = image.size

        # Imagens pequenas/simples: um único passe é suficiente e mais rápido.
        if img_w <= tile_threshold and img_h <= tile_threshold:
            return self._predict_yolo_single(image, offset=(0, 0))

        step = int(tile_size * (1 - overlap))
        all_detections: List[Detection] = []

        xs = list(range(0, max(1, img_w - tile_size), step)) or [0]
        ys = list(range(0, max(1, img_h - tile_size), step)) or [0]
        if xs[-1] + tile_size < img_w:
            xs.append(img_w - tile_size)
        if ys[-1] + tile_size < img_h:
            ys.append(img_h - tile_size)

        for y0 in ys:
            for x0 in xs:
                x1 = min(x0 + tile_size, img_w)
                y1 = min(y0 + tile_size, img_h)
                tile = image.crop((x0, y0, x1, y1))
                all_detections.extend(self._predict_yolo_single(tile, offset=(x0, y0)))

        # também roda um passe na imagem inteira (redimensionada), para
        # capturar componentes grandes que ficariam cortados entre tiles
        all_detections.extend(self._predict_yolo_single(image, offset=(0, 0)))

        return self._simple_nms(all_detections, iou_threshold=0.3)

    def _predict_yolo_single(self, image_or_tile: Image.Image, offset: tuple) -> List[Detection]:
        results = self.model.predict(image_or_tile, conf=self.conf_threshold, verbose=False)
        detections = []
        ox, oy = offset
        for r in results:
            for box in r.boxes:
                cls_id = int(box.cls[0])
                label = CLASSES[cls_id] if cls_id < len(CLASSES) else str(cls_id)
                conf = float(box.conf[0])
                x0, y0, x1, y1 = [float(v) for v in box.xyxy[0]]
                detections.append(Detection(label, conf, (x0 + ox, y0 + oy, x1 + ox, y1 + oy)))
        return detections

    # ---------------------------------------------------------------
    # Modo demo (fallback): detecção heurística por cor/contorno,
    # usada apenas quando ainda não há modelo treinado disponível.
    # ---------------------------------------------------------------
    def _predict_heuristic(self, image: Image.Image) -> List[Detection]:
        """Detecção por segmentação de cor: como os diagramas sintéticos usam uma
        paleta fixa por classe (ver dataset_generator.CLASS_COLORS), procuramos
        diretamente por regiões de cada cor, em vez de depender de contorno externo
        (que tende a fundir blocos conectados por linhas em uma única região)."""
        import cv2

        img = np.array(image.convert("RGB"))
        detections = []
        tolerance = 30  # tolerância de cor (em cada canal RGB)

        for label, color in CLASS_COLORS.items():
            lower = np.array([max(0, c - tolerance) for c in color], dtype=np.uint8)
            upper = np.array([min(255, c + tolerance) for c in color], dtype=np.uint8)
            mask = cv2.inRange(img, lower, upper)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))

            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for cnt in contours:
                x, y, w, h = cv2.boundingRect(cnt)
                if w < 15 or h < 15:
                    continue
                detections.append(Detection(label, 0.5, (x, y, x + w, y + h)))
        return self._simple_nms(detections)

    @staticmethod
    def _simple_nms(detections: List["Detection"], iou_threshold: float = 0.4) -> List["Detection"]:
        """Remove caixas fortemente sobrepostas (cores parecidas podem casar com
        a mesma região), mantendo apenas uma detecção por área sobreposta."""

        def iou(a, b):
            ax0, ay0, ax1, ay1 = a
            bx0, by0, bx1, by1 = b
            ix0, iy0 = max(ax0, bx0), max(ay0, by0)
            ix1, iy1 = min(ax1, bx1), min(ay1, by1)
            iw, ih = max(0, ix1 - ix0), max(0, iy1 - iy0)
            inter = iw * ih
            area_a = (ax1 - ax0) * (ay1 - ay0)
            area_b = (bx1 - bx0) * (by1 - by0)
            union = area_a + area_b - inter
            return inter / union if union > 0 else 0

        kept = []
        # ordena por área (maior primeiro): blocos maiores tendem a ser a detecção mais confiável
        ordered = sorted(detections, key=lambda d: (d.box[2] - d.box[0]) * (d.box[3] - d.box[1]), reverse=True)
        for det in ordered:
            if all(iou(det.box, k.box) < iou_threshold for k in kept):
                kept.append(det)
        return kept