commit abf715360b3bd891dba17b6441a0c9da4700d1b5 Author: newbieQQ Date: Mon Jul 6 11:35:27 2026 +0800 Init: 制动器角度读取工具 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4845697 --- /dev/null +++ b/.gitignore @@ -0,0 +1,14 @@ +__pycache__/ +*.pyc +*.pyo +.venv/ +venv/ +*.egg-info/ +dist/ +build/ +.DS_Store +*.swp +*.swo +*~ +*.csv +*_frames/ diff --git a/actuator/__init__.py b/actuator/__init__.py new file mode 100644 index 0000000..e34fe6e --- /dev/null +++ b/actuator/__init__.py @@ -0,0 +1,8 @@ +from .video_reader import VideoReader +from .angle_calc import signed_angle, curvature_from_three_points +from .automation import ( + PointDetector, + ManualDetector, + ColorBlobDetector, + ActuatorAnalyzer, +) diff --git a/actuator/angle_calc.py b/actuator/angle_calc.py new file mode 100644 index 0000000..912ddce --- /dev/null +++ b/actuator/angle_calc.py @@ -0,0 +1,81 @@ +"""角度计算模块 + +角度定义(∠M2-O-M): + - O: 圆心 / 固定旋转中心(第一帧选定,全程不变) + - M: 初始尾端参考位置(第一帧选定,全程不变) + - M2: 当前帧尾端位置(每帧重新选定) + - 角度 = 向量 O→M 与 O→M2 之间的夹角(0-180°) +""" + +import math +from typing import Tuple + + +def angle_between_vectors( + v1: Tuple[float, float], v2: Tuple[float, float] +) -> float: + """计算两个向量之间的夹角(度)""" + dot = v1[0] * v2[0] + v1[1] * v2[1] + mag1 = math.hypot(v1[0], v1[1]) + mag2 = math.hypot(v2[0], v2[1]) + if mag1 < 1e-9 or mag2 < 1e-9: + return 0.0 + cos_theta = max(-1.0, min(1.0, dot / (mag1 * mag2))) + return math.degrees(math.acos(cos_theta)) + + +def signed_angle( + o: Tuple[float, float], + m: Tuple[float, float], + m2: Tuple[float, float], +) -> float: + """计算带符号的 ∠M2-O-M 角度 + + Args: + o: 圆心(顶点) + m: 初始尾端参考位置 + m2: 当前尾端位置 + + Returns: + 带符号的角度(度),正=顺时针,负=逆时针 + """ + v_ref = (m[0] - o[0], m[1] - o[1]) # O → M + v_cur = (m2[0] - o[0], m2[1] - o[1]) # O → M2 + + angle = angle_between_vectors(v_ref, v_cur) + + # 叉积 Z 分量判断方向 + cross_z = v_ref[0] * v_cur[1] - v_ref[1] * v_cur[0] + if cross_z < 0: + angle = -angle + + return angle + + +def curvature_from_three_points( + o: Tuple[float, float], + m: Tuple[float, float], + m2: Tuple[float, float], + scale: float = 1.0, +) -> float: + """由三点计算曲率(Menger curvature) + + Args: + o, m, m2: 三个点 + scale: 像素到物理单位的比例(mm/pixel) + + Returns: + 曲率 (1/mm) + """ + a = math.hypot(o[0] - m[0], o[1] - m[1]) * scale + b = math.hypot(m[0] - m2[0], m[1] - m2[1]) * scale + c = math.hypot(m2[0] - o[0], m2[1] - o[1]) * scale + + s = (a + b + c) / 2.0 + area = math.sqrt(max(0, s * (s - a) * (s - b) * (s - c))) + + if area < 1e-12: + return 0.0 + + radius = (a * b * c) / (4.0 * area) + return 1.0 / radius if radius > 1e-12 else 0.0 diff --git a/actuator/automation.py b/actuator/automation.py new file mode 100644 index 0000000..451cfec --- /dev/null +++ b/actuator/automation.py @@ -0,0 +1,134 @@ +"""自动化接口模块 + +工作流程: + 1. 第一帧:选 O(圆心)+ M(初始尾端参考) + 2. 后续帧:选 M2(当前尾端),自动计算 ∠M2-O-M +""" + +from abc import ABC, abstractmethod +from typing import List, Optional, Tuple + +import numpy as np + + +class PointDetector(ABC): + """标记点检测器抽象基类 — 自动化接口""" + + @abstractmethod + def detect(self, frame: np.ndarray) -> List[Tuple[float, float]]: + """从一帧图像中检测所有候选标记点 + + Returns: + [(x, y), ...] 检测到的标记点列表 + """ + ... + + +class ManualDetector(PointDetector): + """手动选点(当前版本)""" + + def __init__(self): + self._point: Optional[Tuple[float, float]] = None + + def set_point(self, pt: Tuple[float, float]) -> None: + self._point = pt + + def detect(self, frame: np.ndarray) -> List[Tuple[float, float]]: + return [self._point] if self._point else [] + + +class ColorBlobDetector(PointDetector): + """基于颜色的自动检测器(预留) + + 通过 HSV 颜色范围 + 轮廓检测来定位标记点。 + """ + + def __init__( + self, + lower_hsv: Tuple[int, int, int] = (20, 100, 100), + upper_hsv: Tuple[int, int, int] = (35, 255, 255), + ): + self.lower = np.array(lower_hsv) + self.upper = np.array(upper_hsv) + + def detect(self, frame: np.ndarray) -> List[Tuple[float, float]]: + import cv2 + + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + mask = cv2.inRange(hsv, self.lower, self.upper) + contours, _ = cv2.findContours( + mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE + ) + + points = [] + for c in contours: + M = cv2.moments(c) + if M["m00"] > 100: + cx = M["m10"] / M["m00"] + cy = M["m01"] / M["m00"] + points.append((cx, cy)) + return points + + +class ActuatorAnalyzer: + """制动器分析器 — 封装完整分析流程 + + 使用方式: + analyzer = ActuatorAnalyzer(video_path) + analyzer.set_reference(o, m) # 圆心 + 初始尾端 + analyzer.set_actuator_length(120) # mm + for frame_idx in range(n): + result = analyzer.analyze_frame(frame_idx) + """ + + def __init__(self, video_path: str): + from .video_reader import VideoReader + + self.video = VideoReader(video_path) + self.point_o: Optional[Tuple[float, float]] = None # 圆心 + self.point_m: Optional[Tuple[float, float]] = None # 初始尾端 + self.actuator_length_mm: float = 100.0 + self.detector: PointDetector = ManualDetector() + + def set_reference( + self, + o: Tuple[float, float], + m: Tuple[float, float], + ) -> None: + """设置圆心 O 和初始尾端参考位置 M""" + self.point_o = o + self.point_m = m + + def set_detector(self, detector: PointDetector) -> None: + self.detector = detector + + def analyze_frame(self, frame_idx: int) -> Optional[dict]: + """分析单帧,返回 ∠M2-O-M 等数据""" + from .angle_calc import signed_angle + + frame = self.video.read_frame(frame_idx) + if frame is None or self.point_o is None or self.point_m is None: + return None + + points = self.detector.detect(frame) + if len(points) < 1: + return None + + m2 = points[0] # 当前尾端 + angle = signed_angle(self.point_o, self.point_m, m2) + + return { + "frame_idx": frame_idx, + "point_m2": m2, + "angle_deg": angle, + "actuator_length_mm": self.actuator_length_mm, + } + + def analyze_all_frames(self) -> List[dict]: + """分析所有帧""" + results = [] + for i in range(self.video.frame_count): + r = self.analyze_frame(i) + if r is not None: + results.append(r) + return results diff --git a/actuator/video_reader.py b/actuator/video_reader.py new file mode 100644 index 0000000..01db8bb --- /dev/null +++ b/actuator/video_reader.py @@ -0,0 +1,97 @@ +"""视频读取模块""" + +import cv2 +from pathlib import Path +from typing import Iterator, Optional +import numpy as np + + +class VideoReader: + """读取视频文件,按帧提取""" + + def __init__(self, path: str | Path): + self.path = Path(path) + self._cap: Optional[cv2.VideoCapture] = None + self._frame_count: Optional[int] = None + self._fps: Optional[float] = None + self._width: Optional[int] = None + self._height: Optional[int] = None + + @property + def frame_count(self) -> int: + if self._frame_count is None: + self._open() + return self._frame_count + + @property + def fps(self) -> float: + if self._fps is None: + self._open() + return self._fps + + @property + def width(self) -> int: + if self._width is None: + self._open() + return self._width + + @property + def height(self) -> int: + if self._height is None: + self._open() + return self._height + + def _open(self) -> None: + self._cap = cv2.VideoCapture(str(self.path)) + if not self._cap.isOpened(): + raise FileNotFoundError(f"Cannot open video: {self.path}") + self._frame_count = int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT)) + self._fps = self._cap.get(cv2.CAP_PROP_FPS) + self._width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + self._height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + def read_frame(self, index: int) -> Optional[np.ndarray]: + """读取指定索引的一帧 (BGR) — 随机访问,可能慢""" + if self._cap is None: + self._open() + self._cap.set(cv2.CAP_PROP_POS_FRAMES, index) + ret, frame = self._cap.read() + if not ret: + return None + return frame + + def reset(self) -> None: + """重开视频,回到第 0 帧,用于顺序读取""" + self.close() + self._open() + + def read_next(self) -> Optional[np.ndarray]: + """顺序读取下一帧(不 seek,比 read_frame 可靠)""" + if self._cap is None: + self._open() + ret, frame = self._cap.read() + if not ret: + return None + return frame + + def iter_frames(self) -> Iterator[tuple[int, np.ndarray]]: + """逐帧迭代,返回 (index, frame)""" + self.reset() + idx = 0 + while True: + frame = self.read_next() + if frame is None: + break + yield idx, frame + idx += 1 + + def close(self) -> None: + if self._cap is not None: + self._cap.release() + self._cap = None + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/main.py b/main.py new file mode 100644 index 0000000..0123dff --- /dev/null +++ b/main.py @@ -0,0 +1,766 @@ +#!/usr/bin/env python3 +"""制动器角度读取工具 — PyQt6 GUI + +工作流程: + 1. 打开视频 + 2. 第一帧:点击 O(圆心)→ 点击 M(初始尾端参考) + 3. 点击「开始测量」→ 逐帧点击 M2,每点完一帧自动进下一帧 + 4. 可随时用「方向矫正」翻转角度符号 +""" + +import sys +import csv +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(): + 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() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..71f4815 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +opencv-python>=4.8 +numpy>=1.24 +PyQt6>=6.5 diff --git a/static/test.mp4 b/static/test.mp4 new file mode 100755 index 0000000..cec79da Binary files /dev/null and b/static/test.mp4 differ