135 lines
3.9 KiB
Python
135 lines
3.9 KiB
Python
"""自动化接口模块
|
||
|
||
工作流程：
|
||
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
|