141 lines
5.7 KiB
Python
141 lines
5.7 KiB
Python
"""Widget that renders an image and draws detection overlays.
|
|
|
|
The confidence threshold is applied at paint time, so changing it is instant. One
|
|
detection can be *highlighted* (selected in the detail table) — it is drawn boldly
|
|
even if below the threshold, while the others dim, so the user can inspect exactly
|
|
what the detector found.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import zlib
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PySide6.QtCore import QPointF, QRectF, Qt
|
|
from PySide6.QtGui import QBrush, QColor, QFont, QImage, QPainter, QPen, QPolygonF
|
|
from PySide6.QtWidgets import QWidget
|
|
|
|
from ..config import OverlayConfig
|
|
from ..core.detection.types import Detection
|
|
|
|
|
|
class ImageView(QWidget):
|
|
def __init__(self, overlay_cfg: OverlayConfig, parent: QWidget | None = None) -> None:
|
|
super().__init__(parent)
|
|
self._cfg = overlay_cfg
|
|
self._qimage: QImage | None = None
|
|
self._dets: list[Detection] = []
|
|
self._threshold = 0.0
|
|
self._model_thresholds: dict[str, float] = {} # model stem -> override threshold
|
|
self._highlight: int | None = None
|
|
self.setMinimumSize(480, 360)
|
|
|
|
# ------------------------------------------------------------------ slots
|
|
def set_image(self, image_bgr: np.ndarray | None, dets: list[Detection]) -> None:
|
|
if image_bgr is None:
|
|
self._qimage = None
|
|
else:
|
|
rgb = np.ascontiguousarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
|
|
h, w, ch = rgb.shape
|
|
# .copy() so the QImage owns its pixels (the numpy buffer can be freed).
|
|
self._qimage = QImage(rgb.data, w, h, ch * w, QImage.Format_RGB888).copy()
|
|
self._dets = dets
|
|
self._highlight = None
|
|
self.update()
|
|
|
|
def set_threshold(self, threshold: float) -> None:
|
|
self._threshold = threshold
|
|
self.update()
|
|
|
|
def set_model_thresholds(self, thresholds: dict[str, float]) -> None:
|
|
"""Per-model overlay threshold overrides (keyed by model stem). Empty = none."""
|
|
self._model_thresholds = dict(thresholds)
|
|
self.update()
|
|
|
|
def _eff_threshold(self, d: Detection) -> float:
|
|
"""The threshold this detection must clear — its model's override, else global."""
|
|
return self._model_thresholds.get(d.model, self._threshold)
|
|
|
|
def set_highlight(self, index: int | None) -> None:
|
|
self._highlight = index
|
|
self.update()
|
|
|
|
# ------------------------------------------------------------------ paint
|
|
def _color(self, key: str) -> QColor:
|
|
"""Colour for a detection category — fixed if configured, else a stable palette pick."""
|
|
rgb = self._cfg.colors.get(key)
|
|
if rgb is None:
|
|
palette = self._cfg.palette
|
|
rgb = palette[zlib.crc32(key.encode("utf-8")) % len(palette)]
|
|
return QColor(*rgb)
|
|
|
|
def paintEvent(self, event) -> None:
|
|
painter = QPainter(self)
|
|
painter.fillRect(self.rect(), QColor(18, 18, 18))
|
|
|
|
if self._qimage is None:
|
|
painter.setPen(QColor(160, 160, 160))
|
|
painter.drawText(self.rect(), Qt.AlignCenter, "Откройте проект (Файл → Открыть проект…)")
|
|
painter.end()
|
|
return
|
|
|
|
iw, ih = self._qimage.width(), self._qimage.height()
|
|
scale = min(self.width() / iw, self.height() / ih)
|
|
dw, dh = iw * scale, ih * scale
|
|
ox, oy = (self.width() - dw) / 2, (self.height() - dh) / 2
|
|
|
|
painter.setRenderHint(QPainter.SmoothPixmapTransform, True)
|
|
painter.drawImage(QRectF(ox, oy, dw, dh), self._qimage)
|
|
|
|
if self._dets:
|
|
painter.setRenderHint(QPainter.Antialiasing, True)
|
|
for i, d in enumerate(self._dets):
|
|
highlighted = i == self._highlight
|
|
# A highlighted detection is always drawn; others respect the threshold
|
|
# (per-model override if any, else the global one).
|
|
if not highlighted and d.score < self._eff_threshold(d):
|
|
continue
|
|
dim = self._highlight is not None and not highlighted
|
|
self._draw_detection(painter, d, ox, oy, scale, highlighted, dim)
|
|
painter.end()
|
|
|
|
def _draw_detection(
|
|
self, painter: QPainter, d: Detection, ox: float, oy: float,
|
|
scale: float, highlighted: bool, dim: bool,
|
|
) -> None:
|
|
color = self._color(d.display)
|
|
width = self._cfg.line_width * (2 if highlighted else 1)
|
|
pen_color = QColor(color)
|
|
if dim:
|
|
pen_color.setAlpha(70)
|
|
painter.setPen(QPen(pen_color, width))
|
|
fill = QColor(color)
|
|
fill.setAlpha(0 if dim else (self._cfg.fill_alpha * 2 if highlighted else self._cfg.fill_alpha))
|
|
painter.setBrush(QBrush(fill))
|
|
|
|
points = d.polygon if len(d.polygon) >= 3 else self._bbox_points(d.bbox)
|
|
poly = QPolygonF([QPointF(ox + x * scale, oy + y * scale) for x, y in points])
|
|
painter.drawPolygon(poly)
|
|
|
|
if self._cfg.show_labels and not dim:
|
|
x, y, _w, _h = d.bbox
|
|
self._draw_label(painter, f"{d.display} {d.score:.2f}", ox + x * scale, oy + y * scale, color)
|
|
|
|
@staticmethod
|
|
def _bbox_points(bbox: tuple[int, int, int, int]) -> list[tuple[int, int]]:
|
|
x, y, w, h = bbox
|
|
return [(x, y), (x + w, y), (x + w, y + h), (x, y + h)]
|
|
|
|
def _draw_label(self, painter: QPainter, text: str, x: float, y: float, color: QColor) -> None:
|
|
font = QFont()
|
|
font.setPointSize(9)
|
|
painter.setFont(font)
|
|
metrics = painter.fontMetrics()
|
|
tw = metrics.horizontalAdvance(text) + 8
|
|
th = metrics.height() + 2
|
|
bg = QRectF(x, max(0.0, y - th), tw, th)
|
|
painter.fillRect(bg, color)
|
|
painter.setPen(QColor(0, 0, 0))
|
|
painter.drawText(bg, Qt.AlignCenter, text)
|