File size: 6,330 Bytes
73c5503 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | 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
|