Update app.py

kontur/web/app.py

@@ -11,6 +11,8 @@ from PyQt5.QtCore import Qt, pyqtSlot
 from PyQt5.QtWebEngineWidgets import QWebEngineView
 from PyQt5.QtGui import QImage, QPixmap
 from qasync import QEventLoop, asyncSlot
+from ultralytics import YOLO
+import os
 
 # Настройка логирования
 logging.basicConfig(level=logging.DEBUG)
@@ -21,6 +23,10 @@ class MainWindow(QWidget):
         super().__init__()
         self.initUI()
 
+        # Инициализация модели YOLO
+        self.model = YOLO('models/yolov8n.pt')
+        self.model.conf = 0.5  # Установка порога уверенности для детекции
+
     def initUI(self):
         self.setWindowTitle('Multi-Input PyQt Application')
         self.setGeometry(100, 100, 1920, 1080)
@@ -116,6 +122,11 @@ class MainWindow(QWidget):
             if not ret:
                 logger.error("Failed to read frame from RTSP stream")
                 break
+
+            # Применение модели YOLO к кадру
+            results = self.model.predict(frame, conf=0.6)
+            frame = self.draw_detections(frame, results)
+
             # Convert the frame to QImage
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             height, width, channel = frame_rgb.shape
@@ -176,6 +187,16 @@ class MainWindow(QWidget):
         self.telemetry_label.setText(f"Telemetry Data: {telemetry_info}")
         logger.debug("Telemetry data updated")
 
+    def draw_detections(self, frame, results):
+        for r in results:
+            for i, c in enumerate(r.boxes.cls):
+                if r.boxes.conf[i] > 0.5:
+                    x1, y1, x2, y2 = r.boxes.xyxy[i]  # Координаты рамок
+                    cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)  # Рисование рамки
+                    label = f"{self.model.names[int(c)]} {r.boxes.conf[i]:.2f}"
+                    cv2.putText(frame, label, (int(x1), int(y1 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        return frame
+
     def dark_theme_stylesheet(self):
         return """
         QWidget {