456 lines
17 KiB
Python
456 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PySide6.QtCore import Qt, QTimer, Slot
|
|
from PySide6.QtGui import QAction, QImage, QPixmap
|
|
from PySide6.QtWidgets import (
|
|
QApplication,
|
|
QHBoxLayout,
|
|
QLabel,
|
|
QFileDialog,
|
|
QMainWindow,
|
|
QMessageBox,
|
|
QPushButton,
|
|
QTextEdit,
|
|
QToolButton,
|
|
QVBoxLayout,
|
|
QWidget,
|
|
QStyle,
|
|
)
|
|
|
|
from app.camera import CameraWorker, DetectionWorker
|
|
from app.config import AppConfig
|
|
from app.detection import DetectionResult
|
|
from app.media import MediaStore, VideoRecorder
|
|
from app.settings_dialog import SettingsDialog
|
|
|
|
|
|
class MainWindow(QMainWindow):
|
|
def __init__(self, app_config: AppConfig) -> None:
|
|
super().__init__()
|
|
self.app_config = app_config
|
|
self.config = app_config.data
|
|
self.last_frame: np.ndarray | None = None
|
|
self.overlay_result: DetectionResult | None = None
|
|
self.last_detection: DetectionResult | None = None
|
|
self.detecting = False
|
|
self.detection_frame_count = 0
|
|
self.fps_frame_count = 0
|
|
self.fps_last_time = time.monotonic()
|
|
self.display_fps = 0.0
|
|
self.video_capture: cv2.VideoCapture | None = None
|
|
self.video_timer = QTimer(self)
|
|
self.video_timer.timeout.connect(self._read_video_frame)
|
|
self.video_playing = False
|
|
self.media_store = MediaStore(self.config, self.app_config)
|
|
self.video_recorder = VideoRecorder(self.config, self.app_config)
|
|
|
|
self.setWindowTitle("Duck Stain YOLO")
|
|
self.resize(1280, 720)
|
|
self._build_ui()
|
|
|
|
self.worker = CameraWorker(self.config, self.app_config)
|
|
self.worker.frame_ready.connect(self.on_frame_ready)
|
|
self.worker.camera_error.connect(self.on_camera_error)
|
|
self.worker.start()
|
|
|
|
self.detection_worker = DetectionWorker(self.config, self.app_config)
|
|
self.detection_worker.detection_ready.connect(self.on_detection_ready)
|
|
self.detection_worker.start()
|
|
|
|
def _build_ui(self) -> None:
|
|
self.stage = QWidget()
|
|
self.setCentralWidget(self.stage)
|
|
|
|
self.video_label = QLabel(self.stage)
|
|
self.video_label.setAlignment(Qt.AlignCenter)
|
|
self.video_label.setStyleSheet("background: #111; color: #ddd;")
|
|
self.video_label.setText("Kamera")
|
|
|
|
self.result_panel = QWidget(self.stage)
|
|
self.result_panel.setObjectName("resultPanel")
|
|
self.result_panel.setStyleSheet(
|
|
"""
|
|
QWidget#resultPanel {
|
|
background: rgba(20, 20, 20, 170);
|
|
border-radius: 8px;
|
|
}
|
|
QTextEdit {
|
|
background: transparent;
|
|
color: white;
|
|
border: 0;
|
|
font-size: 13px;
|
|
}
|
|
QPushButton {
|
|
min-height: 28px;
|
|
padding: 4px 12px;
|
|
}
|
|
"""
|
|
)
|
|
panel_layout = QVBoxLayout(self.result_panel)
|
|
self.result_text = QTextEdit()
|
|
self.result_text.setReadOnly(True)
|
|
self.result_text.setFixedHeight(118)
|
|
panel_layout.addWidget(self.result_text)
|
|
panel_buttons = QHBoxLayout()
|
|
self.detect_button = QPushButton("wykryj")
|
|
self.ok_button = QPushButton("ok")
|
|
panel_buttons.addStretch(1)
|
|
panel_buttons.addWidget(self.detect_button)
|
|
panel_buttons.addWidget(self.ok_button)
|
|
panel_layout.addLayout(panel_buttons)
|
|
self.detect_button.clicked.connect(self.start_detection)
|
|
self.ok_button.clicked.connect(self.accept_detection)
|
|
|
|
self.toolbar = QWidget(self.stage)
|
|
self.toolbar.setObjectName("bottomToolbar")
|
|
self.toolbar.setStyleSheet(
|
|
"""
|
|
QWidget#bottomToolbar {
|
|
background: rgba(20, 20, 20, 175);
|
|
border-radius: 8px;
|
|
}
|
|
QToolButton {
|
|
min-width: 44px;
|
|
min-height: 38px;
|
|
padding: 4px;
|
|
}
|
|
"""
|
|
)
|
|
toolbar_layout = QHBoxLayout(self.toolbar)
|
|
toolbar_layout.setContentsMargins(8, 6, 8, 6)
|
|
self.load_video_button = self._tool_button(QStyle.SP_DirOpenIcon, "Wczytaj film")
|
|
self.video_play_button = self._tool_button(QStyle.SP_MediaPlay, "Play/pauza filmu")
|
|
self.photo_button = self._tool_button(QStyle.SP_DialogSaveButton, "Zrob zdjecie")
|
|
self.record_button = self._tool_button(QStyle.SP_MediaPlay, "Start/stop nagrywania")
|
|
self.settings_button = self._tool_button(QStyle.SP_FileDialogDetailedView, "Ustawienia obrazu")
|
|
toolbar_layout.addWidget(self.load_video_button)
|
|
toolbar_layout.addWidget(self.video_play_button)
|
|
toolbar_layout.addWidget(self.photo_button)
|
|
toolbar_layout.addWidget(self.record_button)
|
|
toolbar_layout.addWidget(self.settings_button)
|
|
self.video_play_button.setEnabled(False)
|
|
self.load_video_button.clicked.connect(self.load_video)
|
|
self.video_play_button.clicked.connect(self.toggle_video_playback)
|
|
self.photo_button.clicked.connect(self.take_photo)
|
|
self.record_button.clicked.connect(self.toggle_recording)
|
|
self.settings_button.clicked.connect(self.open_settings)
|
|
|
|
quit_action = QAction("Zamknij", self)
|
|
quit_action.triggered.connect(self.close)
|
|
self.addAction(quit_action)
|
|
|
|
def _tool_button(self, icon_id: QStyle.StandardPixmap, tooltip: str) -> QToolButton:
|
|
button = QToolButton()
|
|
button.setIcon(self.style().standardIcon(icon_id))
|
|
button.setToolTip(tooltip)
|
|
return button
|
|
|
|
def resizeEvent(self, event: Any) -> None:
|
|
super().resizeEvent(event)
|
|
self.video_label.setGeometry(self.stage.rect())
|
|
|
|
panel_width = min(420, max(280, self.stage.width() // 3))
|
|
self.result_panel.setGeometry(self.stage.width() - panel_width - 18, 18, panel_width, 190)
|
|
|
|
self.toolbar.adjustSize()
|
|
toolbar_size = self.toolbar.sizeHint()
|
|
self.toolbar.setGeometry(
|
|
(self.stage.width() - toolbar_size.width()) // 2,
|
|
self.stage.height() - toolbar_size.height() - 18,
|
|
toolbar_size.width(),
|
|
toolbar_size.height(),
|
|
)
|
|
|
|
def closeEvent(self, event: Any) -> None:
|
|
if self.video_recorder.is_recording:
|
|
self.video_recorder.stop(self.current_metadata("video"))
|
|
self.video_timer.stop()
|
|
self._close_video_capture()
|
|
self._stop_camera_worker()
|
|
self.detection_worker.stop()
|
|
self.detection_worker.wait(2000)
|
|
super().closeEvent(event)
|
|
|
|
@Slot(object)
|
|
def on_frame_ready(self, frame: np.ndarray) -> None:
|
|
self._update_fps()
|
|
self.last_frame = frame.copy()
|
|
if self.video_recorder.is_recording:
|
|
self.video_recorder.write(frame)
|
|
self._maybe_request_detection(frame)
|
|
self._show_frame(frame)
|
|
|
|
@Slot(object)
|
|
def on_detection_ready(self, result: DetectionResult) -> None:
|
|
if not self.detecting:
|
|
return
|
|
self.last_detection = result
|
|
self.overlay_result = result if result.xyxy else None
|
|
self._update_result_text(result)
|
|
|
|
@Slot(str)
|
|
def on_camera_error(self, message: str) -> None:
|
|
self.result_text.setPlainText(message)
|
|
|
|
def start_detection(self) -> None:
|
|
self.overlay_result = None
|
|
self.detecting = True
|
|
self.detection_frame_count = 0
|
|
self.result_text.setPlainText("Wykrywanie...")
|
|
|
|
def accept_detection(self) -> None:
|
|
self.detecting = False
|
|
self.overlay_result = None
|
|
if self.last_detection:
|
|
self._update_result_text(self.last_detection, accepted=True)
|
|
|
|
def take_photo(self) -> None:
|
|
if self.last_frame is None:
|
|
QMessageBox.warning(self, "Zdjecie", "Brak klatki z kamery")
|
|
return
|
|
path = self.media_store.save_photo(self.last_frame, self.current_metadata("photo"))
|
|
self.statusBar().showMessage(f"Zapisano zdjecie: {path}", 5000)
|
|
|
|
def toggle_recording(self) -> None:
|
|
if self.last_frame is None:
|
|
QMessageBox.warning(self, "Wideo", "Brak klatki z kamery")
|
|
return
|
|
|
|
if self.video_recorder.is_recording:
|
|
path = self.video_recorder.stop(self.current_metadata("video"))
|
|
self.record_button.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))
|
|
self.statusBar().showMessage(f"Zapisano film: {path}", 5000)
|
|
return
|
|
|
|
try:
|
|
path = self.video_recorder.start(self.last_frame)
|
|
except RuntimeError as exc:
|
|
QMessageBox.warning(self, "Wideo", str(exc))
|
|
return
|
|
self.record_button.setIcon(self.style().standardIcon(QStyle.SP_MediaStop))
|
|
self.statusBar().showMessage(f"Nagrywanie: {path}", 5000)
|
|
|
|
def load_video(self) -> None:
|
|
path, _ = QFileDialog.getOpenFileName(
|
|
self,
|
|
"Wczytaj film",
|
|
"",
|
|
"Filmy (*.mp4 *.avi *.mov *.mkv *.m4v);;Wszystkie pliki (*)",
|
|
)
|
|
if not path:
|
|
return
|
|
|
|
capture = cv2.VideoCapture(path)
|
|
if not capture.isOpened():
|
|
QMessageBox.warning(self, "Film", "Nie mozna otworzyc pliku wideo")
|
|
capture.release()
|
|
return
|
|
|
|
if self.video_recorder.is_recording:
|
|
self.video_recorder.stop(self.current_metadata("video"))
|
|
self.record_button.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))
|
|
|
|
self._stop_camera_worker()
|
|
self._close_video_capture()
|
|
self.video_capture = capture
|
|
self.video_play_button.setEnabled(True)
|
|
self._set_video_playing(False)
|
|
self.overlay_result = None
|
|
self.last_detection = None
|
|
self.result_text.setPlainText(f"Wczytano film: {path}")
|
|
self._read_video_frame()
|
|
|
|
def toggle_video_playback(self) -> None:
|
|
if self.video_capture is None:
|
|
return
|
|
self._set_video_playing(not self.video_playing)
|
|
|
|
def open_settings(self) -> None:
|
|
dialog = SettingsDialog(self.config, self)
|
|
dialog.settings_saved.connect(self.save_camera_settings)
|
|
dialog.exec()
|
|
|
|
@Slot(dict)
|
|
def save_camera_settings(self, camera_config: dict[str, Any]) -> None:
|
|
self.config["camera"] = camera_config
|
|
self.app_config.save(self.config)
|
|
if self.worker is not None:
|
|
self.worker.update_camera_config(camera_config)
|
|
|
|
def _read_video_frame(self) -> None:
|
|
if self.video_capture is None:
|
|
return
|
|
|
|
ok, frame = self.video_capture.read()
|
|
if not ok or frame is None:
|
|
self._set_video_playing(False)
|
|
self.video_capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
|
self.statusBar().showMessage("Koniec filmu", 3000)
|
|
return
|
|
|
|
self.on_frame_ready(frame)
|
|
|
|
def _set_video_playing(self, playing: bool) -> None:
|
|
self.video_playing = playing
|
|
if self.video_capture is None:
|
|
self.video_timer.stop()
|
|
self.video_play_button.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))
|
|
self.video_play_button.setEnabled(False)
|
|
return
|
|
|
|
if playing:
|
|
fps = self.video_capture.get(cv2.CAP_PROP_FPS)
|
|
if fps <= 0:
|
|
fps = float(self.config["camera"].get("fps", 30))
|
|
interval_ms = max(1, int(round(1000 / fps)))
|
|
self.video_timer.start(interval_ms)
|
|
self.video_play_button.setIcon(self.style().standardIcon(QStyle.SP_MediaPause))
|
|
else:
|
|
self.video_timer.stop()
|
|
self.video_play_button.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))
|
|
|
|
def _close_video_capture(self) -> None:
|
|
self._set_video_playing(False)
|
|
if self.video_capture is not None:
|
|
self.video_capture.release()
|
|
self.video_capture = None
|
|
self.video_play_button.setEnabled(False)
|
|
|
|
def _stop_camera_worker(self) -> None:
|
|
if self.worker is None:
|
|
return
|
|
self.worker.stop()
|
|
self.worker.wait(2000)
|
|
self.worker = None
|
|
|
|
def _maybe_request_detection(self, frame: np.ndarray) -> None:
|
|
if not self.detecting:
|
|
return
|
|
|
|
frame_stride = max(1, int(self.config["detection"].get("frame_stride", 5)))
|
|
self.detection_frame_count += 1
|
|
if self.detection_frame_count % frame_stride != 0:
|
|
return
|
|
|
|
self.detection_worker.request_detection(frame)
|
|
|
|
def _update_fps(self) -> None:
|
|
self.fps_frame_count += 1
|
|
now = time.monotonic()
|
|
elapsed = now - self.fps_last_time
|
|
if elapsed < 1.0:
|
|
return
|
|
|
|
self.display_fps = self.fps_frame_count / elapsed
|
|
self.fps_frame_count = 0
|
|
self.fps_last_time = now
|
|
|
|
def current_metadata(self, media_type: str) -> dict[str, Any]:
|
|
return {
|
|
"media_type": media_type,
|
|
"created_at": datetime.now().isoformat(timespec="seconds"),
|
|
"detection": self.last_detection.to_metadata() if self.last_detection else None,
|
|
"camera": {
|
|
"width": self.config["camera"].get("width"),
|
|
"height": self.config["camera"].get("height"),
|
|
"fps": self.config["camera"].get("fps"),
|
|
"properties": self.config["camera"].get("properties", {}),
|
|
},
|
|
"detection_config": self.config.get("detection", {}),
|
|
}
|
|
|
|
def _update_result_text(self, result: DetectionResult, accepted: bool = False) -> None:
|
|
status = "Zatwierdzono" if accepted else "Wynik"
|
|
lines = [status]
|
|
if result.error:
|
|
lines.append(f"Komunikat: {result.error}")
|
|
if result.confidence is not None:
|
|
lines.append(f"YOLO confidence: {result.confidence:.3f}")
|
|
if result.ocr_engine:
|
|
lines.append(f"OCR: {result.ocr_engine}")
|
|
if result.ocr_confidence is not None:
|
|
lines.append(f"OCR confidence: {result.ocr_confidence:.3f}")
|
|
if result.ocr_elapsed_ms is not None:
|
|
lines.append(f"OCR czas: {result.ocr_elapsed_ms:.0f} ms")
|
|
if result.parsed:
|
|
lines.append(f"Zamowienie: {result.parsed.order_number or '-'}")
|
|
color_score = _format_score(result.parsed.color_score)
|
|
model_score = _format_score(result.parsed.product_model_score)
|
|
lines.append(f"Kolor: {result.parsed.color_code or '-'}{color_score}")
|
|
lines.append(f"Model: {result.parsed.product_model or '-'}{model_score}")
|
|
if result.raw_text:
|
|
lines.append("")
|
|
lines.append(result.raw_text)
|
|
self.result_text.setPlainText("\n".join(lines))
|
|
|
|
def _show_frame(self, frame_bgr: np.ndarray) -> None:
|
|
display_frame = frame_bgr.copy()
|
|
if self.overlay_result is not None:
|
|
self._draw_detection(display_frame, self.overlay_result)
|
|
if self.config.get("display", {}).get("show_fps", True):
|
|
self._draw_fps(display_frame)
|
|
|
|
frame_rgb = cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB)
|
|
h, w, channels = frame_rgb.shape
|
|
image = QImage(frame_rgb.data, w, h, channels * w, QImage.Format_RGB888).copy()
|
|
pixmap = QPixmap.fromImage(image)
|
|
self.video_label.setPixmap(
|
|
pixmap.scaled(self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
|
|
)
|
|
|
|
def _draw_detection(self, frame_bgr: np.ndarray, result: DetectionResult) -> None:
|
|
mode = self.config["detection"].get("mode", "best")
|
|
boxes = result.all_boxes if mode == "all" else [result.to_metadata()]
|
|
for item in boxes:
|
|
xyxy = item.get("xyxy") or item.get("bbox_xyxy")
|
|
if not xyxy:
|
|
continue
|
|
x1, y1, x2, y2 = [int(value) for value in xyxy]
|
|
confidence = item.get("confidence")
|
|
class_name = item.get("class_name") or "label"
|
|
cv2.rectangle(frame_bgr, (x1, y1), (x2, y2), (0, 220, 0), 3)
|
|
caption = f"{class_name} {confidence:.2f}" if confidence is not None else class_name
|
|
cv2.putText(
|
|
frame_bgr,
|
|
caption,
|
|
(x1, max(24, y1 - 8)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.8,
|
|
(0, 220, 0),
|
|
2,
|
|
cv2.LINE_AA,
|
|
)
|
|
|
|
def _draw_fps(self, frame_bgr: np.ndarray) -> None:
|
|
label = f"FPS: {self.display_fps:.1f}"
|
|
cv2.rectangle(frame_bgr, (12, 12), (142, 46), (0, 0, 0), -1)
|
|
cv2.putText(
|
|
frame_bgr,
|
|
label,
|
|
(20, 36),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.7,
|
|
(255, 255, 255),
|
|
2,
|
|
cv2.LINE_AA,
|
|
)
|
|
|
|
|
|
def run_app(app_config: AppConfig) -> int:
|
|
app = QApplication([])
|
|
window = MainWindow(app_config)
|
|
window.show()
|
|
return app.exec()
|
|
|
|
|
|
def _format_score(score: float | None) -> str:
|
|
if score is None:
|
|
return ""
|
|
return f" ({score:.2f})"
|