425 lines
15 KiB
Python
425 lines
15 KiB
Python
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
import os
|
|||
|
|
import json
|
|||
|
|
import argparse
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Target2D:
|
|||
|
|
"""二维坐标下的叶尖跟踪器"""
|
|||
|
|
|
|||
|
|
def __init__(self, positions=None, hub=None):
|
|||
|
|
if positions is None:
|
|||
|
|
positions = np.array([[-100, -100], [-200, -200], [-300, -300]])
|
|||
|
|
if hub is None:
|
|||
|
|
hub = np.array([400, 300])
|
|||
|
|
|
|||
|
|
self.positions = np.array(positions, dtype=float) # 叶尖位置
|
|||
|
|
self.hub = np.array(hub, dtype=float) # 轮毂位置
|
|||
|
|
|
|||
|
|
def update(self, new_positions, new_hub=None):
|
|||
|
|
"""更新叶尖位置，使用最近点匹配算法"""
|
|||
|
|
if new_positions is None or len(new_positions) == 0:
|
|||
|
|
return self.positions, []
|
|||
|
|
|
|||
|
|
# 确保新位置是二维数组
|
|||
|
|
new_positions = np.array(new_positions, dtype=float)
|
|||
|
|
if new_positions.ndim == 1:
|
|||
|
|
new_positions = new_positions.reshape(1, -1)
|
|||
|
|
|
|||
|
|
# 更新轮毂位置
|
|||
|
|
if new_hub is not None:
|
|||
|
|
self.hub = np.array(new_hub, dtype=float)
|
|||
|
|
|
|||
|
|
# 如果还没有初始化位置，使用新位置
|
|||
|
|
if self.positions.size == 0:
|
|||
|
|
self.positions = new_positions
|
|||
|
|
return self.positions
|
|||
|
|
# 创建位置列表副本
|
|||
|
|
new_pos_list = new_positions.tolist()
|
|||
|
|
update_positons = self.positions.copy()
|
|||
|
|
used = set()
|
|||
|
|
# 对于每个新位置，找到最近的点
|
|||
|
|
for i, new_pos in enumerate(new_pos_list):
|
|||
|
|
# 计算所有距离
|
|||
|
|
distances = np.linalg.norm(self.positions - new_pos, axis=1)
|
|||
|
|
available_indices = [j for j in range(len(distances)) if j not in used]
|
|||
|
|
# 找到最近点的索引
|
|||
|
|
closest_idx = min(available_indices, key=lambda j: distances[j])
|
|||
|
|
|
|||
|
|
# 记录ID并更新位置
|
|||
|
|
update_positons[closest_idx] = new_pos # new_pos
|
|||
|
|
|
|||
|
|
# 更新位置
|
|||
|
|
self.positions = np.array(update_positons)
|
|||
|
|
|
|||
|
|
return self.positions
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ClearanceCalculator:
|
|||
|
|
def __init__(self, model_path, scale_factor=0.02,width=544, height=960):
|
|||
|
|
self.size = (width, height)
|
|||
|
|
# 初始化模型配置
|
|||
|
|
self.model = YOLO(model_path)
|
|||
|
|
self.cap = None
|
|||
|
|
self.video_path = ""
|
|||
|
|
self.is_playing = False
|
|||
|
|
self.scale_factor = scale_factor
|
|||
|
|
self.reference_pixel = 0.0
|
|||
|
|
self.reference_actual = 0.0
|
|||
|
|
self.current_clearance = 0.0
|
|||
|
|
self.use_camera = False
|
|||
|
|
self.hub_x_range = 100 # 轮毂x坐标的左右区间（像素）
|
|||
|
|
self.hub_y_range = 50
|
|||
|
|
self.inside_region = False # 轨迹是否在轮毂区域内
|
|||
|
|
self.entered_from_left = False # 轨迹是否从左侧进入
|
|||
|
|
self.last_pixel_clearance = 0.0 # 保存上一次的像素净空
|
|||
|
|
self.last_actual_clearance = 0.0 # 保存上一次的实际净空
|
|||
|
|
# 轨迹相关变量
|
|||
|
|
self.trajectories = []
|
|||
|
|
# self.lowest_points = []
|
|||
|
|
self.calculate_points = []
|
|||
|
|
self.hub_position = None
|
|||
|
|
self.last_tip_positions = []
|
|||
|
|
self.frame_count = 0
|
|||
|
|
self.nearest_tip_trajectory = [] # 只存储与轮毂x坐标最近的叶尖轨迹
|
|||
|
|
self.traj_cnt = 0
|
|||
|
|
self.tip_clearances = {
|
|||
|
|
0: None, # tip1
|
|||
|
|
1: None, # tip2
|
|||
|
|
2: None # tip3
|
|||
|
|
}
|
|||
|
|
# 净空历史记录
|
|||
|
|
self.clearance_history = {0: [], 1: [], 2: []}
|
|||
|
|
|
|||
|
|
# 初始化轨迹颜色
|
|||
|
|
self.tip_colors = [
|
|||
|
|
(255, 0, 0), # 红色
|
|||
|
|
(0, 255, 0), # 绿色
|
|||
|
|
(0, 0, 255) # 蓝色
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 初始化2D跟踪器
|
|||
|
|
self.tracker = Target2D()
|
|||
|
|
|
|||
|
|
def set_hub_x_range(self, x_value, y_value=None):
|
|||
|
|
"""设置轮毂x和y坐标的区间"""
|
|||
|
|
self.hub_x_range = x_value
|
|||
|
|
if y_value is not None:
|
|||
|
|
self.hub_y_range = y_value
|
|||
|
|
|
|||
|
|
def load_video(self, path):
|
|||
|
|
"""加载视频文件"""
|
|||
|
|
self.use_camera = False
|
|||
|
|
self.video_path = path
|
|||
|
|
self.cap = cv2.VideoCapture(path)
|
|||
|
|
return self.cap.isOpened()
|
|||
|
|
|
|||
|
|
def process_frame(self, frame):
|
|||
|
|
"""处理视频帧并检测叶尖和轮毂"""
|
|||
|
|
frame = cv2.resize(frame, self.size)
|
|||
|
|
# 使用模型进行检测
|
|||
|
|
results = self.model(frame, verbose=False)
|
|||
|
|
|
|||
|
|
# 初始化检测结果字典
|
|||
|
|
detected = {'tip': [], 'hub': []}
|
|||
|
|
|
|||
|
|
# 解析检测结果
|
|||
|
|
for result in results:
|
|||
|
|
boxes = result.boxes
|
|||
|
|
classes = boxes.cls
|
|||
|
|
class_names = result.names
|
|||
|
|
|
|||
|
|
for i in range(len(boxes)):
|
|||
|
|
class_id = int(classes[i])
|
|||
|
|
box = boxes.xyxy[i]
|
|||
|
|
x1, y1, x2, y2 = box.tolist()
|
|||
|
|
class_name = class_names[class_id]
|
|||
|
|
|
|||
|
|
# 只关注叶尖(tip)和轮毂(hub)
|
|||
|
|
if class_name in detected:
|
|||
|
|
# 使用底部中心点作为位置
|
|||
|
|
if class_name == 'tip':
|
|||
|
|
tip_x = (x1 + x2) / 2
|
|||
|
|
tip_y = y2 # 使用底部作为叶尖位置
|
|||
|
|
detected[class_name].append([tip_x, tip_y])
|
|||
|
|
else: # hub
|
|||
|
|
hub_x = (x1 + x2) / 2
|
|||
|
|
hub_y = (y1 + y2) / 2
|
|||
|
|
detected[class_name].append([hub_x, hub_y])
|
|||
|
|
|
|||
|
|
# 限制检测数量
|
|||
|
|
if len(detected['tip']) > 3:
|
|||
|
|
detected['tip'] = detected['tip'][:3]
|
|||
|
|
if len(detected['hub']) > 1:
|
|||
|
|
detected['hub'] = detected['hub'][:1]
|
|||
|
|
|
|||
|
|
return frame, detected['tip'], detected['hub']
|
|||
|
|
|
|||
|
|
def check_trajectory_state(self, tip, hub, thresh=20):
|
|||
|
|
"""检查轨迹是否进入或离开轮毂区域"""
|
|||
|
|
hub_x = hub[0] if hub else None
|
|||
|
|
if hub_x is None or tip is None:
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
# 检查轨迹是否在轮毂区域内
|
|||
|
|
inside = (hub_x - self.hub_x_range) < tip[0] < (hub_x + self.hub_x_range)
|
|||
|
|
# 如果轨迹进入区域，并且是从左侧进入
|
|||
|
|
if inside and not self.inside_region:
|
|||
|
|
self.inside_region = True
|
|||
|
|
self.entered_from_left = tip[0] >= (hub_x - self.hub_x_range)
|
|||
|
|
|
|||
|
|
if inside and self.inside_region and self.entered_from_left and (hub_x - thresh) < tip[0] < (hub_x + thresh):
|
|||
|
|
distance, actual_distance = self.calculate_clearance()
|
|||
|
|
return distance, actual_distance
|
|||
|
|
# 如果轨迹离开区域，并且是从左侧进入且从右侧离开
|
|||
|
|
if not inside and self.inside_region and self.entered_from_left and tip[0] > (hub_x + self.hub_x_range):
|
|||
|
|
self.inside_region = False
|
|||
|
|
self.traj_cnt += 1
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
def update_trajectories(self, tips, hub):
|
|||
|
|
"""更新叶尖轨迹，只关注轮毂x和y坐标附近的叶尖"""
|
|||
|
|
# 使用2D跟踪器更新位置
|
|||
|
|
self.tracker.update(tips, hub)
|
|||
|
|
sorted_tips = self.tracker.positions
|
|||
|
|
|
|||
|
|
# 如果没有轮毂位置，无法判断附近
|
|||
|
|
if hub is None or len(hub) == 0:
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
hub_x = hub[0] # 轮毂的x坐标
|
|||
|
|
hub_y = hub[1] # 轮毂的y坐标
|
|||
|
|
|
|||
|
|
# 初始化轨迹
|
|||
|
|
if not self.trajectories:
|
|||
|
|
for _ in range(len(sorted_tips)):
|
|||
|
|
self.trajectories.append([])
|
|||
|
|
|
|||
|
|
# 找到与轮毂x和y坐标都最近的叶尖
|
|||
|
|
nearest_tip = None
|
|||
|
|
min_distance = float('inf')
|
|||
|
|
|
|||
|
|
for tip in sorted_tips:
|
|||
|
|
tip_x, tip_y = tip
|
|||
|
|
# 检查是否在轮毂x坐标左右区间内
|
|||
|
|
if abs(tip_x - hub_x) > self.hub_x_range + 20:
|
|||
|
|
continue
|
|||
|
|
# 检查是否大于轮毂y坐标加y区间
|
|||
|
|
if tip_y < hub_y + self.hub_y_range:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 计算与轮毂的综合距离（考虑x）
|
|||
|
|
distance = abs(tip_x - hub_x) # + (tip_y - hub_y) ** 2
|
|||
|
|
if distance < min_distance:
|
|||
|
|
min_distance = distance
|
|||
|
|
nearest_tip = tip
|
|||
|
|
|
|||
|
|
# 更新最近的叶尖轨迹
|
|||
|
|
if nearest_tip is not None:
|
|||
|
|
if not self.nearest_tip_trajectory:
|
|||
|
|
self.nearest_tip_trajectory = [nearest_tip]
|
|||
|
|
else:
|
|||
|
|
self.nearest_tip_trajectory.append(nearest_tip)
|
|||
|
|
|
|||
|
|
# 限制轨迹长度
|
|||
|
|
if len(self.nearest_tip_trajectory) > 200:
|
|||
|
|
self.nearest_tip_trajectory.pop(0)
|
|||
|
|
|
|||
|
|
return self.check_trajectory_state(nearest_tip, hub)
|
|||
|
|
|
|||
|
|
def find_closest_tip_x(self):
|
|||
|
|
"""找到与轮毂x坐标最接近的叶尖x坐标的叶尖"""
|
|||
|
|
if not self.nearest_tip_trajectory:
|
|||
|
|
return None
|
|||
|
|
hub_x = self.hub_position[0] if self.hub_position else None
|
|||
|
|
if hub_x is None:
|
|||
|
|
return None
|
|||
|
|
trajectories = np.array(self.nearest_tip_trajectory[:][0])
|
|||
|
|
distances = np.abs(trajectories - hub_x)
|
|||
|
|
idx = np.argmin(distances)
|
|||
|
|
tip = self.nearest_tip_trajectory[idx]
|
|||
|
|
self.calculate_points = [tip] # 只保留一个最低点
|
|||
|
|
return tip
|
|||
|
|
|
|||
|
|
def calculate_clearance(self):
|
|||
|
|
"""计算叶尖到轮毂的净空"""
|
|||
|
|
if not self.nearest_tip_trajectory or self.hub_position is None:
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
hub_position = self.hub_position
|
|||
|
|
tip_x, tip_y = self.find_closest_tip_x()
|
|||
|
|
self.calculate_points = [[tip_x, tip_y]] # 只保留一个最低点
|
|||
|
|
hub_x, hub_y = hub_position
|
|||
|
|
distance = np.sqrt((tip_x - hub_x) ** 2 + (tip_y - hub_y) ** 2)
|
|||
|
|
|
|||
|
|
# 更新UI显示净空
|
|||
|
|
if self.scale_factor:
|
|||
|
|
actual_distance = distance * self.scale_factor
|
|||
|
|
tip_idx = self.traj_cnt % 3 # 根据实际匹配逻辑
|
|||
|
|
self.tip_clearances[tip_idx] = actual_distance
|
|||
|
|
# 记录到历史
|
|||
|
|
self.clearance_history[tip_idx].append(actual_distance)
|
|||
|
|
return distance, actual_distance
|
|||
|
|
|
|||
|
|
def draw_annotations(self, frame):
|
|||
|
|
"""在帧上绘制轨迹和标记"""
|
|||
|
|
# 绘制轮毂位置
|
|||
|
|
if self.hub_position:
|
|||
|
|
hub_x, hub_y = self.hub_position
|
|||
|
|
# 绘制轮毂区域标记
|
|||
|
|
cv2.rectangle(frame,
|
|||
|
|
(int(hub_x - self.hub_x_range), int(hub_y - 10)),
|
|||
|
|
(int(hub_x + self.hub_x_range), int(hub_y + 10)),
|
|||
|
|
(0, 255, 255), 2)
|
|||
|
|
cv2.circle(frame, (int(hub_x), int(hub_y)), 8, (0, 255, 255), -1)
|
|||
|
|
cv2.putText(frame, "Hub", (int(hub_x) + 15, int(hub_y)),
|
|||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
|
|||
|
|
|
|||
|
|
# 绘制最近的叶尖轨迹
|
|||
|
|
if len(self.nearest_tip_trajectory) >= 2:
|
|||
|
|
# 绘制轨迹线（黄色）
|
|||
|
|
for j in range(1, len(self.nearest_tip_trajectory)):
|
|||
|
|
start_point = (int(self.nearest_tip_trajectory[j - 1][0]),
|
|||
|
|
int(self.nearest_tip_trajectory[j - 1][1]))
|
|||
|
|
end_point = (int(self.nearest_tip_trajectory[j][0]),
|
|||
|
|
int(self.nearest_tip_trajectory[j][1]))
|
|||
|
|
cv2.line(frame, start_point, end_point, (0, 255, 255), 1)
|
|||
|
|
# 绘制三个叶尖的净空值
|
|||
|
|
y_offset = 30
|
|||
|
|
for i in range(3):
|
|||
|
|
clearance = self.tip_clearances.get(i, None)
|
|||
|
|
if clearance is not None:
|
|||
|
|
text = f"tip{i + 1}: {clearance:.2f} m"
|
|||
|
|
else:
|
|||
|
|
text = f"tip{i + 1}: -- m"
|
|||
|
|
cv2.putText(frame, text, (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
|||
|
|
y_offset += 25
|
|||
|
|
|
|||
|
|
# 绘制其他叶尖（不绘制轨迹）
|
|||
|
|
if self.tracker.positions.size > 0:
|
|||
|
|
for i, tip in enumerate(self.tracker.positions):
|
|||
|
|
if not any(np.array_equal(tip, self.nearest_tip_trajectory[-1]) for p in self.nearest_tip_trajectory if
|
|||
|
|
len(self.nearest_tip_trajectory) > 0):
|
|||
|
|
color = self.tip_colors[i % len(self.tip_colors)]
|
|||
|
|
cv2.circle(frame, (int(tip[0]), int(tip[1])), 6, color, -1)
|
|||
|
|
cv2.putText(frame, f"Tip {i + 1}", (int(tip[0]) + 15, int(tip[1])),
|
|||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
|
|||
|
|
|
|||
|
|
return frame
|
|||
|
|
|
|||
|
|
def process_video(self, output_dir):
|
|||
|
|
"""处理整个视频并保存结果"""
|
|||
|
|
if not os.path.exists(output_dir):
|
|||
|
|
os.makedirs(output_dir)
|
|||
|
|
|
|||
|
|
last_frame = None
|
|||
|
|
frame_count = 0
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
ret, frame = self.cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
# 处理帧
|
|||
|
|
processed_frame, tips, hubs = self.process_frame(frame)
|
|||
|
|
|
|||
|
|
# 更新轮毂位置
|
|||
|
|
if hubs:
|
|||
|
|
self.hub_position = hubs[0]
|
|||
|
|
|
|||
|
|
# 更新轨迹
|
|||
|
|
if tips and self.hub_position is not None:
|
|||
|
|
pixel_clearance, actual_clearance = self.update_trajectories(tips, self.hub_position)
|
|||
|
|
|
|||
|
|
# 绘制注释
|
|||
|
|
annotated_frame = self.draw_annotations(processed_frame)
|
|||
|
|
last_frame = annotated_frame.copy()
|
|||
|
|
frame_count += 1
|
|||
|
|
|
|||
|
|
# 显示处理进度
|
|||
|
|
if frame_count % 10 == 0:
|
|||
|
|
print(f"已处理 {frame_count} 帧...")
|
|||
|
|
|
|||
|
|
# 保存最后一帧
|
|||
|
|
last_frame_path = os.path.join(output_dir, "last_frame.jpg")
|
|||
|
|
cv2.imwrite(last_frame_path, last_frame)
|
|||
|
|
print(f"最后一帧已保存至: {last_frame_path}")
|
|||
|
|
|
|||
|
|
# 准备结果数据
|
|||
|
|
result = {
|
|||
|
|
"clearance_history": {
|
|||
|
|
"tip0": self.clearance_history[0],
|
|||
|
|
"tip1": self.clearance_history[1],
|
|||
|
|
"tip2": self.clearance_history[2]
|
|||
|
|
},
|
|||
|
|
"last_frame_clearances": {
|
|||
|
|
"tip0": self.tip_clearances[0],
|
|||
|
|
"tip1": self.tip_clearances[1],
|
|||
|
|
"tip2": self.tip_clearances[2]
|
|||
|
|
},
|
|||
|
|
"last_frame_image": last_frame_path
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 保存JSON结果
|
|||
|
|
json_path = os.path.join(output_dir, "results.json")
|
|||
|
|
with open(json_path, 'w') as f:
|
|||
|
|
json.dump(result, f, indent=4)
|
|||
|
|
print(f"结果已保存至: {json_path}")
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='叶片净空计算系统')
|
|||
|
|
parser.add_argument('--video_path', type=str, required=True, help='输入视频路径')
|
|||
|
|
parser.add_argument('--model_path', type=str, required=True, help='模型文件路径')
|
|||
|
|
parser.add_argument('--output_dir', type=str, required=True, help='结果保存目录')
|
|||
|
|
parser.add_argument('--scale_factor', type=float, default=0.02, help='比例尺因子 (实际距离/像素距离)')
|
|||
|
|
parser.add_argument('--width', type=int, default=544, help='图像宽度（像素）')
|
|||
|
|
parser.add_argument('--height', type=int, default=960, help='图像高度（像素）')
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
calculator = ClearanceCalculator(
|
|||
|
|
model_path=args.model_path,
|
|||
|
|
scale_factor=args.scale_factor,
|
|||
|
|
width=args.width,
|
|||
|
|
height=args.height
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not calculator.load_video(args.video_path):
|
|||
|
|
print(f"无法加载视频: {args.video_path}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
print(f"开始处理视频: {args.video_path}")
|
|||
|
|
print(f"图像尺寸设置为: ({args.width}, {args.height})")
|
|||
|
|
|
|||
|
|
result = calculator.process_video(args.output_dir)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|
|||
|
|
"""
|
|||
|
|
python cli_clearance.py \
|
|||
|
|
--video test.mp4 \
|
|||
|
|
--model best.pt \
|
|||
|
|
--output ./results \
|
|||
|
|
--scale_factor 0.02 \
|
|||
|
|
--width 640 \
|
|||
|
|
--height 480
|
|||
|
|
"""
|
|||
|
|
"""
|
|||
|
|
{
|
|||
|
|
"clearances": {
|
|||
|
|
"0": [5.12, 5.09, 5.11],
|
|||
|
|
"1": [4.98, 5.01],
|
|||
|
|
"2": []
|
|||
|
|
},
|
|||
|
|
"last_frame_clearances": {
|
|||
|
|
"tip1": 5.11,
|
|||
|
|
"tip2": 5.01,
|
|||
|
|
"tip3": null
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
"""
|