308 lines
10 KiB
Python
308 lines
10 KiB
Python
# utils.py
|
||
# -*- coding: utf-8 -*-
|
||
import cv2
|
||
import numpy as np
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 物品名称映射
|
||
ITEM_TO_CLASS_MAP = {
|
||
"红牛": "Red_Bull",
|
||
"AD钙奶": "AD_milk",
|
||
"ad钙奶": "AD_milk",
|
||
"钙奶": "AD_milk",
|
||
}
|
||
|
||
# 英文类别名到中文的映射
|
||
_OBSTACLE_NAME_CN = {
|
||
'person': '人',
|
||
'bicycle': '自行车',
|
||
'car': '车',
|
||
'motorcycle': '摩托车',
|
||
'bus': '公交车',
|
||
'truck': '卡车',
|
||
'animal': '动物',
|
||
'scooter': '电瓶车',
|
||
'stroller': '婴儿车',
|
||
'dog': '狗',
|
||
}
|
||
|
||
# 动态类别名称列表
|
||
DYNAMIC_CLASS_NAMES = {'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'animal', 'dog'}
|
||
|
||
def extract_english_label(item_cn: str) -> tuple:
|
||
"""
|
||
提取中文物品名称对应的英文标签
|
||
:param item_cn: 中文物品名称
|
||
:return: (英文标签, 来源)
|
||
"""
|
||
# 先查找本地映射
|
||
if item_cn in ITEM_TO_CLASS_MAP:
|
||
return ITEM_TO_CLASS_MAP[item_cn], "local"
|
||
|
||
# 如果没有找到，返回原始名称
|
||
return item_cn, "direct"
|
||
|
||
def _to_cn_obstacle(name: str) -> str:
|
||
"""
|
||
将英文障碍物名称转换为中文
|
||
:param name: 英文名称
|
||
:return: 中文名称
|
||
"""
|
||
try:
|
||
key = (name or '').strip().lower()
|
||
return _OBSTACLE_NAME_CN.get(key, '障碍物')
|
||
except Exception:
|
||
return '障碍物'
|
||
|
||
def estimate_global_affine(prev_gray, curr_gray, mask=None):
|
||
"""
|
||
估计两帧之间的全局仿射变换
|
||
:param prev_gray: 前一帧灰度图
|
||
:param curr_gray: 当前帧灰度图
|
||
:param mask: 可选的掩码，只在掩码区域内计算
|
||
:return: (仿射矩阵, 内点数)
|
||
"""
|
||
try:
|
||
# 提取特征点
|
||
detector = cv2.ORB_create(nfeatures=500)
|
||
kp1, des1 = detector.detectAndCompute(prev_gray, mask)
|
||
kp2, des2 = detector.detectAndCompute(curr_gray, mask)
|
||
|
||
if des1 is None or des2 is None or len(kp1) < 10 or len(kp2) < 10:
|
||
return np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32), 0
|
||
|
||
# 匹配特征点
|
||
matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
|
||
matches = matcher.match(des1, des2)
|
||
|
||
if len(matches) < 4:
|
||
return np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32), 0
|
||
|
||
# 提取匹配的点对
|
||
src_pts = np.float32([kp1[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2)
|
||
dst_pts = np.float32([kp2[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2)
|
||
|
||
# 使用RANSAC估计仿射变换
|
||
M, inliers = cv2.estimateAffinePartial2D(src_pts, dst_pts, method=cv2.RANSAC,
|
||
ransacReprojThreshold=3.0)
|
||
|
||
if M is None:
|
||
return np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32), 0
|
||
|
||
inlier_count = np.sum(inliers) if inliers is not None else 0
|
||
return M, inlier_count
|
||
|
||
except Exception as e:
|
||
logger.warning(f"estimate_global_affine failed: {e}")
|
||
return np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32), 0
|
||
|
||
def warp_mask(mask, M, output_shape):
|
||
"""
|
||
使用仿射变换对掩码进行变换
|
||
:param mask: 输入掩码
|
||
:param M: 2x3的仿射变换矩阵
|
||
:param output_shape: 输出形状 (width, height)
|
||
:return: 变换后的掩码
|
||
"""
|
||
try:
|
||
if mask is None or M is None:
|
||
return None
|
||
|
||
W, H = output_shape
|
||
warped = cv2.warpAffine(mask, M, (W, H),
|
||
flags=cv2.INTER_NEAREST,
|
||
borderMode=cv2.BORDER_CONSTANT,
|
||
borderValue=0)
|
||
return warped
|
||
|
||
except Exception as e:
|
||
logger.warning(f"warp_mask failed: {e}")
|
||
return None
|
||
|
||
def estimate_translation_flow(prev_gray, curr_gray, mask=None):
|
||
"""
|
||
估计两帧之间的平移光流
|
||
:param prev_gray: 前一帧灰度图
|
||
:param curr_gray: 当前帧灰度图
|
||
:param mask: 可选的掩码
|
||
:return: (中位光流幅度, 平移矩阵)
|
||
"""
|
||
try:
|
||
# 计算稀疏光流
|
||
corners = cv2.goodFeaturesToTrack(prev_gray, maxCorners=100,
|
||
qualityLevel=0.3, minDistance=7,
|
||
mask=mask)
|
||
|
||
if corners is None or len(corners) < 10:
|
||
return 0.0, np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32)
|
||
|
||
# 计算光流
|
||
next_pts, status, _ = cv2.calcOpticalFlowPyrLK(prev_gray, curr_gray,
|
||
corners, None)
|
||
|
||
# 筛选有效点
|
||
valid_old = corners[status == 1]
|
||
valid_new = next_pts[status == 1]
|
||
|
||
if len(valid_old) < 5:
|
||
return 0.0, np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32)
|
||
|
||
# 计算位移
|
||
flow_vectors = valid_new - valid_old
|
||
flow_magnitudes = np.linalg.norm(flow_vectors, axis=1)
|
||
median_flow = np.median(flow_magnitudes)
|
||
|
||
# 估计平均平移
|
||
mean_translation = np.mean(flow_vectors, axis=0)
|
||
M = np.array([[1, 0, mean_translation[0]],
|
||
[0, 1, mean_translation[1]]], dtype=np.float32)
|
||
|
||
return median_flow, M
|
||
|
||
except Exception as e:
|
||
logger.warning(f"estimate_translation_flow failed: {e}")
|
||
return 0.0, np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32)
|
||
|
||
def is_stationary_frame(prev_gray, curr_gray, mask=None, threshold=0.35):
|
||
"""
|
||
判断用户是否静止
|
||
:param prev_gray: 前一帧灰度图
|
||
:param curr_gray: 当前帧灰度图
|
||
:param mask: 可选的掩码
|
||
:param threshold: 静止判定阈值
|
||
:return: True表示静止，False表示运动
|
||
"""
|
||
try:
|
||
median_flow, _ = estimate_translation_flow(prev_gray, curr_gray, mask)
|
||
return median_flow < threshold
|
||
except:
|
||
return False
|
||
|
||
def compute_approach_metrics(prev_obstacles, curr_obstacles, M, H, W):
|
||
"""
|
||
计算障碍物的接近度量
|
||
:param prev_obstacles: 前一帧障碍物列表
|
||
:param curr_obstacles: 当前帧障碍物列表
|
||
:param M: 仿射变换矩阵
|
||
:param H: 图像高度
|
||
:param W: 图像宽度
|
||
:return: 接近度量列表
|
||
"""
|
||
metrics = []
|
||
|
||
for curr_obs in curr_obstacles:
|
||
# 寻找最佳匹配的前一帧障碍物
|
||
best_match = None
|
||
best_iou = 0.0
|
||
|
||
curr_mask = curr_obs.get('mask')
|
||
if curr_mask is None:
|
||
metrics.append(None)
|
||
continue
|
||
|
||
for prev_obs in prev_obstacles:
|
||
prev_mask = prev_obs.get('mask')
|
||
if prev_mask is None:
|
||
continue
|
||
|
||
# 将前一帧掩码变换到当前帧
|
||
warped_prev = warp_mask(prev_mask, M, (W, H))
|
||
if warped_prev is None:
|
||
continue
|
||
|
||
# 计算IoU
|
||
intersection = np.logical_and(curr_mask > 0, warped_prev > 0).sum()
|
||
union = np.logical_or(curr_mask > 0, warped_prev > 0).sum()
|
||
iou = intersection / union if union > 0 else 0.0
|
||
|
||
if iou > best_iou:
|
||
best_iou = iou
|
||
best_match = prev_obs
|
||
|
||
if best_match is None:
|
||
metrics.append(None)
|
||
continue
|
||
|
||
# 计算度量
|
||
curr_area = curr_obs.get('area', 0)
|
||
prev_area = best_match.get('area', 0)
|
||
area_growth = (curr_area - prev_area) / prev_area if prev_area > 0 else 0.0
|
||
|
||
curr_bottom_y = curr_obs.get('bottom_y_ratio', 0)
|
||
prev_bottom_y = best_match.get('bottom_y_ratio', 0)
|
||
v_forward = curr_bottom_y - prev_bottom_y
|
||
|
||
metrics.append({
|
||
'area_growth': area_growth,
|
||
'v_forward': v_forward,
|
||
'iou': best_iou
|
||
})
|
||
|
||
return metrics
|
||
|
||
def compute_risk_scores(obstacles, prev_obstacles, M, path_mask, image_shape,
|
||
stop_th=0.6, avoid_th=0.56):
|
||
"""
|
||
计算障碍物的风险评分
|
||
:param obstacles: 当前障碍物列表
|
||
:param prev_obstacles: 前一帧障碍物列表
|
||
:param M: 仿射变换矩阵
|
||
:param path_mask: 路径掩码
|
||
:param image_shape: 图像形状
|
||
:param stop_th: 停止阈值
|
||
:param avoid_th: 避让阈值
|
||
:return: (评分后的障碍物列表, 是否需要停止, 是否需要避让, 可视化元素)
|
||
"""
|
||
H, W = image_shape[:2]
|
||
has_stop = False
|
||
has_avoid = False
|
||
risk_vis = []
|
||
|
||
# 计算接近度量
|
||
metrics = compute_approach_metrics(prev_obstacles, obstacles, M, H, W)
|
||
|
||
for obs, met in zip(obstacles, metrics):
|
||
risk_score = 0.0
|
||
|
||
if met is not None:
|
||
# 基于接近速度和面积增长计算风险
|
||
if met['v_forward'] > 0.004: # 向下移动
|
||
risk_score += 0.3
|
||
if met['area_growth'] > 0.01: # 面积增长
|
||
risk_score += 0.3
|
||
|
||
# 基于距离的风险
|
||
bottom_y = obs.get('bottom_y_ratio', 0)
|
||
area_ratio = obs.get('area_ratio', 0)
|
||
|
||
if bottom_y > 0.8 or area_ratio > 0.15:
|
||
risk_score += 0.3
|
||
|
||
# 动态物体额外风险
|
||
name_lower = str(obs.get('name', '')).lower()
|
||
if name_lower in DYNAMIC_CLASS_NAMES:
|
||
risk_score *= 1.2
|
||
|
||
obs['risk_score'] = risk_score
|
||
|
||
# 更新标志
|
||
if risk_score >= stop_th:
|
||
has_stop = True
|
||
elif risk_score >= avoid_th:
|
||
has_avoid = True
|
||
|
||
# 添加风险可视化
|
||
if risk_score > 0.3:
|
||
risk_color = "rgba(255, 0, 0, 0.3)" if risk_score >= stop_th else "rgba(255, 165, 0, 0.3)"
|
||
risk_vis.append({
|
||
"type": "risk_indicator",
|
||
"score": risk_score,
|
||
"color": risk_color,
|
||
"position": [int(obs.get('center_x', W/2)), int(obs.get('center_y', H/2))]
|
||
})
|
||
|
||
return obstacles, has_stop, has_avoid, risk_vis
|
||
|