771 lines
28 KiB
Python
771 lines
28 KiB
Python
#!/usr/bin/env python3
|
||
"""制动器角度读取工具 — PyQt6 GUI
|
||
|
||
工作流程：
|
||
1. 打开视频
|
||
2. 第一帧：点击 O（圆心）→ 点击 M（初始尾端参考）
|
||
3. 点击「开始测量」→ 逐帧点击 M2，每点完一帧自动进下一帧
|
||
4. 可随时用「方向矫正」翻转角度符号
|
||
"""
|
||
|
||
import sys
|
||
import csv
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from PyQt6.QtCore import Qt, QTimer
|
||
from PyQt6.QtGui import (
|
||
QImage, QPixmap, QPainter, QPen, QColor, QFont, QAction,
|
||
)
|
||
from PyQt6.QtWidgets import (
|
||
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
|
||
QLabel, QPushButton, QSlider, QFileDialog, QGroupBox,
|
||
QDoubleSpinBox, QFormLayout, QMessageBox, QSplitter, QCheckBox,
|
||
QSizePolicy,
|
||
)
|
||
|
||
from actuator import VideoReader, signed_angle, curvature_from_three_points
|
||
|
||
|
||
# ── 颜色 ──────────────────────────────────────────────────
|
||
COLOR_O = QColor(0, 255, 0)
|
||
COLOR_M = QColor(255, 200, 0)
|
||
COLOR_M2 = QColor(255, 50, 50)
|
||
COLOR_M2_PREV = QColor(120, 120, 120) # 灰 — 上一帧 M2
|
||
COLOR_REF_LINE = QColor(100, 200, 100)
|
||
COLOR_CUR_LINE = QColor(255, 120, 120)
|
||
COLOR_PREV_LINE = QColor(100, 100, 100) # 深灰 — O→M2' 上一帧线
|
||
|
||
RADIUS = 8
|
||
|
||
|
||
class FrameView(QLabel):
|
||
"""视频帧显示，鼠标点击选点：O → M → M2 → M2 ..."""
|
||
|
||
def __init__(self, parent=None):
|
||
super().__init__(parent)
|
||
self.setMinimumSize(320, 180)
|
||
self.setSizePolicy(
|
||
QSizePolicy.Policy.Expanding,
|
||
QSizePolicy.Policy.Expanding,
|
||
)
|
||
self.setAlignment(Qt.AlignmentFlag.AlignCenter)
|
||
self.setStyleSheet("background: #1a1a1a; border: 1px solid #333;")
|
||
self.setMouseTracking(True)
|
||
|
||
self._frame: Optional[np.ndarray] = None
|
||
self._pixmap: Optional[QPixmap] = None
|
||
self._scale: float = 1.0
|
||
self._offset_x: float = 0.0
|
||
self._offset_y: float = 0.0
|
||
|
||
self.point_o: Optional[tuple] = None
|
||
self.point_m: Optional[tuple] = None
|
||
self.point_m2: Optional[tuple] = None
|
||
self.point_m2_prev: Optional[tuple] = None # 上一帧 M2
|
||
|
||
self._select_mode: str = "o" # o → m → m2
|
||
self._block_clicks: bool = False
|
||
|
||
self.on_point_selected = None
|
||
|
||
@property
|
||
def reference_ready(self) -> bool:
|
||
return self.point_o is not None and self.point_m is not None
|
||
|
||
def set_frame(self, frame: np.ndarray) -> None:
|
||
self._frame = frame.copy()
|
||
self._update_pixmap()
|
||
|
||
def resizeEvent(self, event) -> None:
|
||
super().resizeEvent(event)
|
||
self._update_pixmap()
|
||
|
||
def _update_pixmap(self) -> None:
|
||
if self._frame is None:
|
||
return
|
||
h, w = self._frame.shape[:2]
|
||
rgb = cv2.cvtColor(self._frame, cv2.COLOR_BGR2RGB)
|
||
qimg = QImage(rgb.data, w, h, 3 * w, QImage.Format.Format_RGB888)
|
||
self._pixmap = QPixmap.fromImage(qimg).scaled(
|
||
self.width(), self.height(),
|
||
Qt.AspectRatioMode.KeepAspectRatio,
|
||
Qt.TransformationMode.SmoothTransformation,
|
||
)
|
||
self._scale = self._pixmap.width() / w
|
||
self._offset_x = (self.width() - self._pixmap.width()) // 2
|
||
self._offset_y = (self.height() - self._pixmap.height()) // 2
|
||
self.update()
|
||
|
||
def widget_to_image(self, wx: int, wy: int) -> tuple:
|
||
return ((wx - self._offset_x) / self._scale,
|
||
(wy - self._offset_y) / self._scale)
|
||
|
||
def image_to_widget(self, ix: float, iy: float) -> tuple:
|
||
return (ix * self._scale + self._offset_x,
|
||
iy * self._scale + self._offset_y)
|
||
|
||
def mousePressEvent(self, event) -> None:
|
||
if event.button() != Qt.MouseButton.LeftButton:
|
||
return
|
||
if self._frame is None or self._block_clicks:
|
||
return
|
||
|
||
pt = self.widget_to_image(
|
||
int(event.position().x()), int(event.position().y()))
|
||
|
||
if self._select_mode == "o":
|
||
self.point_o = pt
|
||
self._select_mode = "m"
|
||
elif self._select_mode == "m":
|
||
self.point_m = pt
|
||
self._select_mode = "m2"
|
||
else:
|
||
self.point_m2 = pt
|
||
|
||
self.update()
|
||
if self.on_point_selected:
|
||
self.on_point_selected(self._select_mode, pt)
|
||
|
||
def paintEvent(self, event) -> None:
|
||
super().paintEvent(event)
|
||
painter = QPainter(self)
|
||
painter.setRenderHint(QPainter.RenderHint.Antialiasing)
|
||
|
||
if self._pixmap:
|
||
painter.drawPixmap(
|
||
int(self._offset_x), int(self._offset_y), self._pixmap)
|
||
|
||
if self._frame is None:
|
||
painter.setPen(QColor(150, 150, 150))
|
||
painter.setFont(QFont("sans-serif", 14))
|
||
painter.drawText(
|
||
self.rect(), Qt.AlignmentFlag.AlignCenter, "文件 → 打开视频")
|
||
painter.end()
|
||
return
|
||
|
||
# 尺寸跟随窗口缩放
|
||
r = max(2, min(6, int(3 * self._scale)))
|
||
lw = max(1, int(2 * self._scale))
|
||
fs = max(8, min(24, int(14 * self._scale)))
|
||
|
||
# O → M
|
||
if self.point_o and self.point_m:
|
||
p1 = self.image_to_widget(*self.point_o)
|
||
p2 = self.image_to_widget(*self.point_m)
|
||
painter.setPen(QPen(COLOR_REF_LINE, lw, Qt.PenStyle.DashLine))
|
||
painter.drawLine(int(p1[0]), int(p1[1]), int(p2[0]), int(p2[1]))
|
||
|
||
# O → M2' (上一帧)
|
||
if self.point_o and self.point_m2_prev:
|
||
p1 = self.image_to_widget(*self.point_o)
|
||
p2 = self.image_to_widget(*self.point_m2_prev)
|
||
painter.setPen(QPen(COLOR_PREV_LINE, max(1, lw // 2), Qt.PenStyle.DotLine))
|
||
painter.drawLine(int(p1[0]), int(p1[1]), int(p2[0]), int(p2[1]))
|
||
|
||
# O → M2
|
||
if self.point_o and self.point_m2:
|
||
p1 = self.image_to_widget(*self.point_o)
|
||
p2 = self.image_to_widget(*self.point_m2)
|
||
painter.setPen(QPen(COLOR_CUR_LINE, lw, Qt.PenStyle.DashLine))
|
||
painter.drawLine(int(p1[0]), int(p1[1]), int(p2[0]), int(p2[1]))
|
||
|
||
self._draw_point(painter, self.point_o, COLOR_O, "O", r, fs)
|
||
self._draw_point(painter, self.point_m, COLOR_M, "M", r, fs)
|
||
self._draw_point(painter, self.point_m2_prev, COLOR_M2_PREV, "M2'", r, fs)
|
||
self._draw_point(painter, self.point_m2, COLOR_M2, "M2", r, fs)
|
||
painter.end()
|
||
|
||
def _draw_point(self, painter, pt, color, label, r, fs):
|
||
if pt is None:
|
||
return
|
||
wx, wy = self.image_to_widget(*pt)
|
||
x, y = int(wx), int(wy)
|
||
painter.setBrush(color)
|
||
painter.setPen(QPen(Qt.GlobalColor.white, max(1, r // 4)))
|
||
painter.drawEllipse(x - r, y - r, r * 2, r * 2)
|
||
painter.setPen(Qt.GlobalColor.white)
|
||
painter.setFont(QFont("monospace", fs, QFont.Weight.Bold))
|
||
painter.drawText(x + r + 2, y - r - 2, label)
|
||
|
||
def reset_all(self) -> None:
|
||
self.point_o = None
|
||
self.point_m = None
|
||
self.point_m2 = None
|
||
self.point_m2_prev = None
|
||
self._select_mode = "o"
|
||
self.update()
|
||
|
||
def reset_m2(self) -> None:
|
||
"""清除 M2，当前 M2 变为上一帧"""
|
||
if self.point_m2 is not None:
|
||
self.point_m2_prev = self.point_m2
|
||
self.point_m2 = None
|
||
self.update()
|
||
|
||
|
||
class MainWindow(QMainWindow):
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.setWindowTitle("制动器角度读取 — ∠M2-O-M")
|
||
self.setMinimumSize(1100, 650)
|
||
|
||
self._video: Optional[VideoReader] = None
|
||
self._current_frame_idx: int = 0
|
||
|
||
# 测量模式
|
||
self._measuring: bool = False
|
||
self._flip_sign: bool = False
|
||
self._prev_angle: Optional[float] = None # 上帧角度
|
||
self._measurements: list = []
|
||
|
||
self._setup_ui()
|
||
self._setup_menu()
|
||
|
||
# ══════════════ UI ═══════════════════════════════════
|
||
|
||
def _setup_ui(self) -> None:
|
||
central = QWidget()
|
||
self.setCentralWidget(central)
|
||
root = QHBoxLayout(central)
|
||
|
||
self.frame_view = FrameView()
|
||
self.frame_view.on_point_selected = self._on_point_changed
|
||
|
||
# 右侧面板（弹性宽度）
|
||
panel = QWidget()
|
||
panel.setMinimumWidth(240)
|
||
panel.setMaximumWidth(400)
|
||
pl = QVBoxLayout(panel)
|
||
pl.setContentsMargins(4, 4, 4, 4)
|
||
pl.setSpacing(4)
|
||
|
||
# ── 选点状态 ──
|
||
grp = QGroupBox("选点状态")
|
||
gf = QFormLayout(grp)
|
||
gf.setHorizontalSpacing(4)
|
||
|
||
lbl_style = "font-size: 11px; font-weight: bold;"
|
||
coord_style = "font-size: 10px;"
|
||
|
||
self.lbl_o = QLabel("未选")
|
||
self.lbl_o.setStyleSheet(f"{lbl_style} color: #0f0;")
|
||
self.lbl_o.setWordWrap(True)
|
||
self.lbl_m = QLabel("未选")
|
||
self.lbl_m.setStyleSheet(f"{lbl_style} color: #fa0;")
|
||
self.lbl_m.setWordWrap(True)
|
||
self.lbl_m2 = QLabel("未选")
|
||
self.lbl_m2.setStyleSheet(f"{lbl_style} color: #f33;")
|
||
self.lbl_m2.setWordWrap(True)
|
||
gf.addRow("O:", self.lbl_o)
|
||
gf.addRow("M:", self.lbl_m)
|
||
gf.addRow("M2:", self.lbl_m2)
|
||
|
||
btn_row = QHBoxLayout()
|
||
self.btn_reset_all = QPushButton("重置全部")
|
||
self.btn_reset_all.clicked.connect(self._reset_all)
|
||
self.btn_reset_m2 = QPushButton("清除 M2")
|
||
self.btn_reset_m2.clicked.connect(self._reset_m2)
|
||
btn_row.addWidget(self.btn_reset_all)
|
||
btn_row.addWidget(self.btn_reset_m2)
|
||
gf.addRow(btn_row)
|
||
|
||
self.lbl_hint = QLabel("点击画面选 O → M → M2")
|
||
self.lbl_hint.setStyleSheet("color: #888; font-size: 11px;")
|
||
gf.addRow(self.lbl_hint)
|
||
pl.addWidget(grp)
|
||
|
||
# ── 方向矫正 ──
|
||
grp_corr = QGroupBox("方向矫正")
|
||
corr_layout = QVBoxLayout(grp_corr)
|
||
self.chk_flip = QCheckBox("翻转角度符号 (± 取反)")
|
||
self.chk_flip.toggled.connect(self._on_flip_toggled)
|
||
corr_layout.addWidget(self.chk_flip)
|
||
pl.addWidget(grp_corr)
|
||
|
||
# ── 参数 ──
|
||
grp_p = QGroupBox("参数")
|
||
pf = QFormLayout(grp_p)
|
||
self.spin_length = QDoubleSpinBox()
|
||
self.spin_length.setRange(10, 5000)
|
||
self.spin_length.setValue(100.0)
|
||
self.spin_length.setSuffix(" mm")
|
||
self.spin_length.valueChanged.connect(self._update_display)
|
||
pf.addRow("致动器长度:", self.spin_length)
|
||
|
||
self.spin_scale = QDoubleSpinBox()
|
||
self.spin_scale.setRange(0.001, 100.0)
|
||
self.spin_scale.setValue(1.0)
|
||
self.spin_scale.setDecimals(4)
|
||
self.spin_scale.setSuffix(" mm/px")
|
||
self.spin_scale.valueChanged.connect(self._update_display)
|
||
pf.addRow("像素比例:", self.spin_scale)
|
||
pl.addWidget(grp_p)
|
||
|
||
# ── 读数 ──
|
||
grp_r = QGroupBox("∠M2-O-M")
|
||
rf = QFormLayout(grp_r)
|
||
self.lbl_angle = QLabel("—")
|
||
self.lbl_angle.setStyleSheet(
|
||
"font-size: 32px; font-weight: bold; color: #0af;")
|
||
self.lbl_angle_prev = QLabel("")
|
||
self.lbl_angle_prev.setStyleSheet(
|
||
"font-size: 14px; color: #888;")
|
||
self.lbl_curvature = QLabel("—")
|
||
self.lbl_curvature.setStyleSheet("font-size: 14px; color: #aaa;")
|
||
self.lbl_frame = QLabel("—")
|
||
self.lbl_progress = QLabel("")
|
||
self.lbl_progress.setStyleSheet("color: #888; font-size: 11px;")
|
||
rf.addRow("角度:", self.lbl_angle)
|
||
rf.addRow("上帧:", self.lbl_angle_prev)
|
||
rf.addRow("曲率:", self.lbl_curvature)
|
||
rf.addRow("当前帧:", self.lbl_frame)
|
||
rf.addRow(self.lbl_progress)
|
||
pl.addWidget(grp_r)
|
||
|
||
# ── 测量控制 ──
|
||
grp_ctrl = QGroupBox("测量控制")
|
||
ctrl_layout = QVBoxLayout(grp_ctrl)
|
||
self.btn_start = QPushButton("▶ 开始测量")
|
||
self.btn_start.setStyleSheet(
|
||
"QPushButton { background: #1a6b1a; color: #fff; font-weight: bold; "
|
||
"padding: 8px; border-radius: 4px; }"
|
||
"QPushButton:hover { background: #228b22; }")
|
||
self.btn_start.clicked.connect(self._start_measuring)
|
||
ctrl_layout.addWidget(self.btn_start)
|
||
|
||
self.btn_stop = QPushButton("⏹ 停止测量")
|
||
self.btn_stop.setEnabled(False)
|
||
self.btn_stop.setStyleSheet(
|
||
"QPushButton { background: #8b1a1a; color: #fff; font-weight: bold; "
|
||
"padding: 8px; border-radius: 4px; }"
|
||
"QPushButton:hover { background: #a52a2a; }")
|
||
self.btn_stop.clicked.connect(self._stop_measuring)
|
||
ctrl_layout.addWidget(self.btn_stop)
|
||
pl.addWidget(grp_ctrl)
|
||
|
||
# ── 帧导航 ──
|
||
grp_n = QGroupBox("手动导航")
|
||
nl = QVBoxLayout(grp_n)
|
||
self.slider = QSlider(Qt.Orientation.Horizontal)
|
||
self.slider.setMinimum(0)
|
||
self.slider.setMaximum(0)
|
||
self.slider.valueChanged.connect(self._on_slider)
|
||
nl.addWidget(self.slider)
|
||
bn = QHBoxLayout()
|
||
self.btn_prev = QPushButton("◀")
|
||
self.btn_prev.clicked.connect(self._prev_frame)
|
||
self.btn_next = QPushButton("▶")
|
||
self.btn_next.clicked.connect(self._next_frame)
|
||
bn.addWidget(self.btn_prev)
|
||
bn.addWidget(self.btn_next)
|
||
nl.addLayout(bn)
|
||
pl.addWidget(grp_n)
|
||
|
||
# ── 导出 ──
|
||
self.btn_export = QPushButton("导出测量数据 (CSV)")
|
||
self.btn_export.clicked.connect(self._export_csv)
|
||
self.btn_export.setEnabled(False)
|
||
pl.addWidget(self.btn_export)
|
||
pl.addStretch()
|
||
|
||
splitter = QSplitter(Qt.Orientation.Horizontal)
|
||
splitter.addWidget(self.frame_view)
|
||
splitter.addWidget(panel)
|
||
splitter.setStretchFactor(0, 1)
|
||
root.addWidget(splitter)
|
||
|
||
self.statusBar().showMessage("请打开视频文件")
|
||
|
||
def _setup_menu(self) -> None:
|
||
m = self.menuBar()
|
||
fm = m.addMenu("文件")
|
||
fm.addAction(QAction("打开视频...", self, shortcut="Ctrl+O",
|
||
triggered=self._open_video))
|
||
fm.addSeparator()
|
||
fm.addAction(QAction("退出", self, triggered=self.close))
|
||
|
||
# ══════════════ 视频 ═════════════════════════════════
|
||
|
||
def _open_video(self) -> None:
|
||
path, _ = QFileDialog.getOpenFileName(
|
||
self, "打开视频", "",
|
||
"视频文件 (*.mp4 *.avi *.mov *.mkv);;所有文件 (*)")
|
||
if path:
|
||
self.load_video(path)
|
||
|
||
def load_video(self, path: str) -> None:
|
||
try:
|
||
self._video = VideoReader(path)
|
||
except Exception as e:
|
||
QMessageBox.critical(self, "错误", str(e))
|
||
return
|
||
self._current_frame_idx = 0
|
||
self.slider.setMaximum(self._video.frame_count - 1)
|
||
self.slider.setValue(0)
|
||
self.frame_view.reset_all()
|
||
self._measurements.clear()
|
||
self._prev_angle = None
|
||
self.btn_export.setEnabled(True)
|
||
self._show_frame(0)
|
||
self.statusBar().showMessage(
|
||
f"已加载: {Path(path).name} | "
|
||
f"{self._video.frame_count} 帧 | "
|
||
f"{self._video.fps:.1f} fps")
|
||
|
||
def _show_frame(self, idx: int) -> None:
|
||
"""随机访问显示帧（手动导航用）"""
|
||
if self._video is None:
|
||
return
|
||
frame = self._video.read_frame(idx)
|
||
if frame is None:
|
||
return
|
||
self._current_frame_idx = idx
|
||
self.frame_view.set_frame(frame)
|
||
self.slider.blockSignals(True)
|
||
self.slider.setValue(idx)
|
||
self.slider.blockSignals(False)
|
||
self._update_display()
|
||
|
||
def _show_frame_seq(self, idx: int) -> None:
|
||
"""顺序读取显示帧（测量模式用，不 seek）"""
|
||
if self._video is None:
|
||
return
|
||
frame = self._video.read_next()
|
||
if frame is None:
|
||
return
|
||
self._current_frame_idx = idx
|
||
self.frame_view.set_frame(frame)
|
||
self.slider.blockSignals(True)
|
||
self.slider.setValue(idx)
|
||
self.slider.blockSignals(False)
|
||
self._update_display()
|
||
|
||
def _update_display(self) -> None:
|
||
self.lbl_frame.setText(str(self._current_frame_idx))
|
||
o, m, m2 = self.frame_view.point_o, self.frame_view.point_m, self.frame_view.point_m2
|
||
|
||
if o and m and m2:
|
||
angle = signed_angle(o, m, m2)
|
||
if self._flip_sign:
|
||
angle = -angle
|
||
self.lbl_angle.setText(f"{angle:+.2f}°")
|
||
scale = self.spin_scale.value()
|
||
curv = curvature_from_three_points(o, m, m2, scale)
|
||
self.lbl_curvature.setText(f"{curv:.6f} mm⁻¹")
|
||
else:
|
||
self.lbl_angle.setText("—")
|
||
self.lbl_curvature.setText("—")
|
||
|
||
# 上帧角度
|
||
if self._prev_angle is not None:
|
||
self.lbl_angle_prev.setText(f"{self._prev_angle:+.2f}°")
|
||
else:
|
||
self.lbl_angle_prev.setText("")
|
||
|
||
if self._measuring:
|
||
total = self._video.frame_count if self._video else 0
|
||
self.lbl_progress.setText(
|
||
f"测量中: {len(self._measurements)}/{total} 帧")
|
||
else:
|
||
self.lbl_progress.setText("")
|
||
|
||
# ══════════════ 导航 ═════════════════════════════════
|
||
|
||
def _on_slider(self, value: int) -> None:
|
||
if not self._measuring and value != self._current_frame_idx:
|
||
self._show_frame(value)
|
||
|
||
def _prev_frame(self) -> None:
|
||
if self._video and not self._measuring:
|
||
self._show_frame(max(0, self._current_frame_idx - 1))
|
||
|
||
def _next_frame(self) -> None:
|
||
if self._video and not self._measuring:
|
||
self._show_frame(min(self._video.frame_count - 1,
|
||
self._current_frame_idx + 1))
|
||
|
||
# ══════════════ 测量模式 ═════════════════════════════
|
||
|
||
def _start_measuring(self) -> None:
|
||
if self._video is None:
|
||
return
|
||
if not self.frame_view.reference_ready:
|
||
QMessageBox.warning(self, "警告", "请先在首帧选定 O 和 M")
|
||
return
|
||
|
||
# 重置测量
|
||
self._measurements.clear()
|
||
self._prev_angle = None
|
||
self._measuring = True
|
||
self.frame_view._select_mode = "m2"
|
||
self.frame_view._block_clicks = False
|
||
|
||
# 禁用无关控件
|
||
self.btn_start.setEnabled(False)
|
||
self.btn_stop.setEnabled(True)
|
||
self.slider.setEnabled(False)
|
||
self.btn_prev.setEnabled(False)
|
||
self.btn_next.setEnabled(False)
|
||
self.btn_reset_all.setEnabled(False)
|
||
self.btn_reset_m2.setEnabled(False)
|
||
|
||
# 重开视频，从第 0 帧开始顺序读取
|
||
self._video.reset()
|
||
self.frame_view.reset_m2()
|
||
self._show_frame_seq(0)
|
||
self.lbl_hint.setText("⏳ 点击 M2 → 自动进入下一帧")
|
||
self.lbl_hint.setStyleSheet(
|
||
"color: #0af; font-size: 12px; font-weight: bold;")
|
||
self._update_display()
|
||
|
||
def _stop_measuring(self) -> None:
|
||
self._measuring = False
|
||
self.frame_view._block_clicks = False
|
||
self.btn_start.setEnabled(True)
|
||
self.btn_stop.setEnabled(False)
|
||
self.slider.setEnabled(True)
|
||
self.btn_prev.setEnabled(True)
|
||
self.btn_next.setEnabled(True)
|
||
self.btn_reset_all.setEnabled(True)
|
||
self.btn_reset_m2.setEnabled(True)
|
||
self.lbl_hint.setText(
|
||
f"已停止 · 共 {len(self._measurements)} 帧数据")
|
||
self.lbl_hint.setStyleSheet("color: #888; font-size: 11px;")
|
||
self._update_display()
|
||
|
||
def _go_to_frame(self, idx: int) -> None:
|
||
"""测量模式中顺序前进并清空 M2"""
|
||
self.frame_view.reset_m2()
|
||
self._show_frame_seq(idx)
|
||
|
||
# ══════════════ 选点回调 ═════════════════════════════
|
||
|
||
def _on_point_changed(self, mode: str, pt: tuple) -> None:
|
||
x, y = int(round(pt[0])), int(round(pt[1]))
|
||
if mode == "o":
|
||
self.lbl_o.setText(f"({x}, {y})")
|
||
self.lbl_hint.setText("点击画面选 M")
|
||
elif mode == "m":
|
||
self.lbl_m.setText(f"({x}, {y})")
|
||
if not self._measuring:
|
||
self.lbl_hint.setText("点击画面选 M2，或点「开始测量」")
|
||
elif mode == "m2":
|
||
self.lbl_m2.setText(f"({x}, {y})")
|
||
self._update_display()
|
||
if self._measuring:
|
||
self._record_and_advance()
|
||
|
||
def _record_and_advance(self) -> None:
|
||
o = self.frame_view.point_o
|
||
m = self.frame_view.point_m
|
||
m2 = self.frame_view.point_m2
|
||
if o is None or m is None or m2 is None:
|
||
return
|
||
|
||
angle = signed_angle(o, m, m2)
|
||
if self._flip_sign:
|
||
angle = -angle
|
||
|
||
self._measurements.append((
|
||
self._current_frame_idx, angle, m2[0], m2[1],
|
||
))
|
||
|
||
# 先保存当前角度为上帧，再进下一帧
|
||
self._prev_angle = angle
|
||
QApplication.processEvents()
|
||
QTimer.singleShot(150, self._advance_frame)
|
||
|
||
def _advance_frame(self) -> None:
|
||
if not self._measuring:
|
||
return
|
||
next_idx = self._current_frame_idx + 1
|
||
if next_idx >= self._video.frame_count:
|
||
self._stop_measuring()
|
||
self.lbl_hint.setText(f"✓ 测量完成 · 共 {len(self._measurements)} 帧")
|
||
self.lbl_hint.setStyleSheet(
|
||
"color: #0f0; font-size: 12px; font-weight: bold;")
|
||
return
|
||
self._go_to_frame(next_idx)
|
||
|
||
# ══════════════ 方向矫正 ═════════════════════════════
|
||
|
||
def _on_flip_toggled(self, checked: bool) -> None:
|
||
self._flip_sign = checked
|
||
self._update_display()
|
||
|
||
# ══════════════ 重置 ═════════════════════════════════
|
||
|
||
def _reset_all(self) -> None:
|
||
if self._measuring:
|
||
self._stop_measuring()
|
||
self.frame_view.reset_all()
|
||
self._measurements.clear()
|
||
self._prev_angle = None
|
||
self.lbl_o.setText("未选")
|
||
self.lbl_m.setText("未选")
|
||
self.lbl_m2.setText("未选")
|
||
self.lbl_hint.setText("点击画面选 O → M → M2")
|
||
self.lbl_hint.setStyleSheet("color: #888; font-size: 11px;")
|
||
self._update_display()
|
||
|
||
def _reset_m2(self) -> None:
|
||
if self._measuring:
|
||
return
|
||
self.frame_view.reset_m2()
|
||
self.lbl_m2.setText("未选")
|
||
self.lbl_hint.setText("点击画面选 M2（当前尾端）")
|
||
self._update_display()
|
||
|
||
# ══════════════ 导出 ═════════════════════════════════
|
||
|
||
def _export_csv(self) -> None:
|
||
if not self._measurements:
|
||
QMessageBox.warning(self, "警告", "没有测量数据，请先「开始测量」")
|
||
return
|
||
|
||
csv_path = QFileDialog.getSaveFileName(
|
||
self, "导出 CSV", "angle_data.csv", "CSV (*.csv)")[0]
|
||
if not csv_path:
|
||
return
|
||
path = Path(csv_path)
|
||
|
||
o = self.frame_view.point_o
|
||
m = self.frame_view.point_m
|
||
scale = self.spin_scale.value()
|
||
|
||
# 截图目录
|
||
frames_dir = path.parent / f"{path.stem}_frames"
|
||
frames_dir.mkdir(exist_ok=True)
|
||
|
||
# 用顺序读取导出（避免 seek 卡死）
|
||
self._video.reset()
|
||
frame_buf = self._video.read_next()
|
||
buf_idx = 0
|
||
|
||
with open(path, "w", newline="", encoding="utf-8") as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow([
|
||
"frame", "angle_deg",
|
||
"O_x", "O_y", "M_x", "M_y",
|
||
"M2_x", "M2_y", "curvature_mm-1",
|
||
"screenshot",
|
||
])
|
||
for frame_idx, angle, m2x, m2y in self._measurements:
|
||
# 跳过不匹配的帧（理论不会发生）
|
||
while buf_idx < frame_idx and frame_buf is not None:
|
||
frame_buf = self._video.read_next()
|
||
buf_idx += 1
|
||
|
||
if frame_buf is not None and buf_idx == frame_idx:
|
||
annotated = self._draw_annotations(
|
||
frame_buf, o, m, (m2x, m2y))
|
||
img_name = f"frame_{frame_idx:04d}.png"
|
||
cv2.imwrite(str(frames_dir / img_name), annotated)
|
||
frame_buf = self._video.read_next()
|
||
buf_idx += 1
|
||
else:
|
||
img_name = ""
|
||
|
||
curv = curvature_from_three_points(o, m, (m2x, m2y), scale)
|
||
writer.writerow([
|
||
frame_idx, f"{angle:.4f}",
|
||
f"{o[0]:.1f}", f"{o[1]:.1f}",
|
||
f"{m[0]:.1f}", f"{m[1]:.1f}",
|
||
f"{m2x:.1f}", f"{m2y:.1f}",
|
||
f"{curv:.8f}",
|
||
f"{frames_dir.name}/{img_name}" if img_name else "",
|
||
])
|
||
|
||
QMessageBox.information(
|
||
self, "导出成功",
|
||
f"已导出 {len(self._measurements)} 条数据 + 截图到:\n"
|
||
f"{path}\n{frames_dir}/")
|
||
|
||
def _draw_annotations(self, frame, o, m, m2):
|
||
"""在帧上绘制 O/M/M2 标记和连线"""
|
||
annotated = frame.copy()
|
||
|
||
def pt(p):
|
||
return (int(round(p[0])), int(round(p[1])))
|
||
|
||
cv2.line(annotated, pt(o), pt(m), (100, 200, 100), 2, cv2.LINE_AA)
|
||
cv2.line(annotated, pt(o), pt(m2), (100, 100, 255), 2, cv2.LINE_AA)
|
||
|
||
cv2.circle(annotated, pt(o), 8, (0, 255, 0), -1, cv2.LINE_AA)
|
||
cv2.circle(annotated, pt(o), 9, (255, 255, 255), 2, cv2.LINE_AA)
|
||
cv2.circle(annotated, pt(m), 8, (0, 200, 255), -1, cv2.LINE_AA)
|
||
cv2.circle(annotated, pt(m), 9, (255, 255, 255), 2, cv2.LINE_AA)
|
||
cv2.circle(annotated, pt(m2), 8, (50, 50, 255), -1, cv2.LINE_AA)
|
||
cv2.circle(annotated, pt(m2), 9, (255, 255, 255), 2, cv2.LINE_AA)
|
||
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
cv2.putText(annotated, "O", (pt(o)[0] + 12, pt(o)[1] - 8),
|
||
font, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
|
||
cv2.putText(annotated, "M", (pt(m)[0] + 12, pt(m)[1] - 8),
|
||
font, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
|
||
cv2.putText(annotated, "M2", (pt(m2)[0] + 12, pt(m2)[1] - 8),
|
||
font, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
|
||
|
||
return annotated
|
||
|
||
# ══════════════ 生命周期 ═════════════════════════════
|
||
|
||
def closeEvent(self, event) -> None:
|
||
if self._video:
|
||
self._video.close()
|
||
super().closeEvent(event)
|
||
|
||
|
||
def main():
|
||
# 禁用自动高分屏缩放，修复光标变小问题
|
||
os.environ["QT_AUTO_SCREEN_SCALE_FACTOR"] = "0"
|
||
|
||
app = QApplication(sys.argv)
|
||
app.setStyle("Fusion")
|
||
app.setStyleSheet("""
|
||
QMainWindow { background: #1e1e1e; }
|
||
QGroupBox {
|
||
color: #ccc; border: 1px solid #444; border-radius: 6px;
|
||
margin-top: 12px; padding-top: 10px; font-weight: bold;
|
||
}
|
||
QGroupBox::title {
|
||
subcontrol-origin: margin; left: 10px; padding: 0 6px;
|
||
}
|
||
QLabel { color: #ddd; }
|
||
QPushButton {
|
||
background: #3a3a3a; color: #ddd; border: 1px solid #555;
|
||
border-radius: 4px; padding: 5px 12px;
|
||
}
|
||
QPushButton:hover { background: #4a4a4a; }
|
||
QPushButton:pressed { background: #2a2a2a; }
|
||
QPushButton:disabled { color: #666; }
|
||
QSlider::groove:horizontal {
|
||
height: 6px; background: #444; border-radius: 3px;
|
||
}
|
||
QSlider::handle:horizontal {
|
||
width: 14px; height: 14px; margin: -5px 0;
|
||
background: #0af; border-radius: 7px;
|
||
}
|
||
QSpinBox, QDoubleSpinBox {
|
||
background: #2a2a2a; color: #ddd; border: 1px solid #555;
|
||
border-radius: 3px; padding: 3px;
|
||
}
|
||
QCheckBox { color: #ddd; }
|
||
QStatusBar { color: #888; }
|
||
""")
|
||
|
||
window = MainWindow()
|
||
if len(sys.argv) > 1:
|
||
window.load_video(sys.argv[1])
|
||
window.show()
|
||
sys.exit(app.exec())
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|