Source code for QVideo.overlays.yolo
'''Real-time object detection overlay using YOLO.
References
----------
Jocher, G., Chaurasia, A., & Qiu, J. (2023). Ultralytics YOLO.
https://github.com/ultralytics/ultralytics
Redmon, J., Divvala, S., Girshick, R., & Farhadi, A. (2016).
You only look once: Unified, real-time object detection.
Proceedings of the IEEE Conference on Computer Vision and Pattern
Recognition, 779-788. https://doi.org/10.1109/CVPR.2016.91
'''
from qtpy import QtCore, QtGui, QtWidgets
import pyqtgraph as pg
from QVideo.lib.videotypes import Image
import numpy as np
import logging
try:
import pandas as pd
except ImportError:
pd = None
logger = logging.getLogger(__name__)
__all__ = ['QYoloOverlay', 'QYoloWidget']
try:
from ultralytics import YOLO
except Exception:
YOLO = None
class _YoloWorker(QtCore.QObject):
'''Runs YOLO object detection in a background thread.
Parameters
----------
model_name : str
Name of the YOLO model weights file. Default: ``'yolo11n.pt'``.
confidence : float
Minimum confidence threshold for reported detections.
Default: ``0.25``.
Signals
-------
newData(object)
Emitted after each detection with a :class:`pandas.DataFrame`
containing columns ``x1``, ``y1``, ``x2``, ``y2``,
``confidence``, ``class``, and ``label``.
``None`` on error or when no objects are detected.
'''
newData = QtCore.Signal(object)
def __init__(self,
model_name: str = 'yolo11n.pt',
confidence: float = 0.25) -> None:
super().__init__()
if pd is None:
raise ImportError(
'pandas is required for QYoloWidget.'
'\n\tInstall it with: pip install pandas'
'\n\tor: pip install "QVideo[overlays]"')
if YOLO is None:
raise ImportError(
'ultralytics is required for QYoloWidget.'
'\n\tInstall it with: pip install ultralytics'
'\n\tSee https://docs.ultralytics.com/ for more information.')
try:
self.model = YOLO(model_name)
except FileNotFoundError:
raise FileNotFoundError(
f'YOLO model "{model_name}" not found.'
'\n\tProvide the name of a pretrained ultralytics model'
'\n\tor the full path to a custom YOLO weights file.'
'\n\tSee https://docs.ultralytics.com/models/ '
'for available pretrained models.')
self.confidence = confidence
@QtCore.Slot(np.ndarray)
def detect(self, image: Image) -> None:
'''Run YOLO detection on *image* and emit :attr:`newData`.
Parameters
----------
image : Image
Video frame to analyze.
'''
try:
results = self.model(image, verbose=False, conf=self.confidence)
boxes = results[0].boxes
if len(boxes) == 0:
features = None
else:
xyxy = boxes.xyxy.cpu().numpy()
conf = boxes.conf.cpu().numpy()
cls = boxes.cls.cpu().numpy().astype(int)
features = pd.DataFrame({
'x1': xyxy[:, 0],
'y1': xyxy[:, 1],
'x2': xyxy[:, 2],
'y2': xyxy[:, 3],
'confidence': conf,
'class': cls,
'label': [results[0].names[c] for c in cls],
})
except Exception as exc:
logger.warning(f'YOLO detection failed: {exc}')
features = None
self.newData.emit(features)
[docs]
class QYoloOverlay(pg.GraphicsObject):
'''Bounding-box overlay that marks YOLO detected objects.
A :class:`pyqtgraph.GraphicsObject` that draws axis-aligned
bounding boxes over detected objects. Add it to a
:class:`~QVideo.lib.QVideoScreen.QVideoScreen` via
``screen.addOverlay(widget.overlay)``.
'''
def __init__(self) -> None:
super().__init__()
self._features = None
self._pen = pg.mkPen('g', width=2)
[docs]
def paint(self, painter: QtGui.QPainter, option, widget=None) -> None:
'''Draw bounding boxes for all current detections.'''
if self._features is None or len(self._features) == 0:
return
painter.setPen(self._pen)
for _, row in self._features.iterrows():
painter.drawRect(QtCore.QRectF(
row['x1'], row['y1'],
row['x2'] - row['x1'],
row['y2'] - row['y1']))
[docs]
@QtCore.Slot(object)
def setFeatures(self, features) -> None:
'''Update bounding boxes from a YOLO detections DataFrame.
Parameters
----------
features : pandas.DataFrame or None
DataFrame with ``x1``, ``y1``, ``x2``, ``y2`` columns
returned by :class:`_YoloWorker`. ``None`` or empty clears
the overlay.
'''
self._features = features
self.update()
[docs]
class QYoloWidget(QtWidgets.QGroupBox):
'''Control widget for the YOLO object-detection overlay.
Runs YOLO inference in a background thread and renders detected
object bounding boxes as a :class:`QYoloOverlay` on a
:class:`~QVideo.lib.QVideoScreen.QVideoScreen`.
Use ``screen.addOverlay(widget.overlay)`` to register the overlay
graphics item with a screen, and set :attr:`source` to supply video frames.
Parameters
----------
parent : QtWidgets.QWidget or None
Parent widget.
model_name : str
YOLO model weights file. Default: ``'yolo11n.pt'``.
confidence : float
Initial confidence threshold. Default: ``0.25``.
'''
#: Emitted for each processed frame with the detections
#: :class:`~pandas.DataFrame`, or ``None`` on error / no detections.
newData = QtCore.Signal(object)
_detect = QtCore.Signal(np.ndarray)
def __init__(self,
parent: QtWidgets.QWidget | None = None,
model_name: str = 'yolo11n.pt',
confidence: float = 0.25) -> None:
if pd is None:
raise ImportError(
'pandas is required for QYoloWidget.'
'\n\tInstall it with: pip install pandas'
'\n\tor: pip install "QVideo[overlays]"')
if YOLO is None:
raise ImportError(
'ultralytics is required for QYoloWidget.'
'\n\tInstall it with: pip install ultralytics'
'\n\tSee https://docs.ultralytics.com/ for more information.')
super().__init__('YOLO', parent)
self._source = None
self._ready = True
self._overlay = QYoloOverlay()
self._worker = _YoloWorker(model_name=model_name, confidence=confidence)
self._thread = QtCore.QThread(self)
self._worker.moveToThread(self._thread)
self._detect.connect(self._worker.detect)
self._worker.newData.connect(self._onNewData)
self._thread.start()
self._setupUi()
QtCore.QCoreApplication.instance().aboutToQuit.connect(self._cleanup)
def _setupUi(self) -> None:
self.setCheckable(True)
self.setChecked(False)
self.setFlat(True)
layout = QtWidgets.QFormLayout(self)
layout.setContentsMargins(2, 5, 2, 5)
self._confidenceSpinBox = QtWidgets.QDoubleSpinBox()
self._confidenceSpinBox.setRange(0., 1.)
self._confidenceSpinBox.setSingleStep(0.05)
self._confidenceSpinBox.setDecimals(2)
self._confidenceSpinBox.setValue(self._worker.confidence)
self._confidenceSpinBox.valueChanged.connect(self._setConfidence)
layout.addRow('Confidence', self._confidenceSpinBox)
self.toggled.connect(self._overlay.setVisible)
@property
def source(self):
'''The :class:`~QVideo.lib.QVideoSource.QVideoSource` being analyzed.'''
return self._source
@source.setter
def source(self, source) -> None:
if self._source is not None:
self._source.newFrame.disconnect(self._onNewFrame)
self._source = source
if source is not None:
source.newFrame.connect(self._onNewFrame)
@property
def overlay(self) -> QYoloOverlay:
'''The :class:`QYoloOverlay` graphics item for this widget.'''
return self._overlay
@QtCore.Slot(np.ndarray)
def _onNewFrame(self, image: Image) -> None:
if self._ready and self.isChecked():
self._ready = False
self._detect.emit(image)
@QtCore.Slot(object)
def _onNewData(self, features) -> None:
self._ready = True
self._overlay.setFeatures(features)
self.newData.emit(features)
@QtCore.Slot(float)
def _setConfidence(self, value: float) -> None:
self._worker.confidence = value
[docs]
def closeEvent(self, event: QtGui.QCloseEvent) -> None:
'''Stop the worker thread when the widget is closed.'''
self._cleanup()
super().closeEvent(event)
@QtCore.Slot()
def _cleanup(self) -> None:
self.source = None
self._thread.quit()
self._thread.wait()