新增双目yolo实现python代码
This commit is contained in:
keita
@@ -0,0 +1,286 @@
|
||||
"""demo_stereo_yolo.py —— 双目 YOLO 检测前端演示（规格 docs/双目YOLO目标识别需求说明.md）。
|
||||
|
||||
用真实模型（models/yolo26n-drone.pt）与一张占位素材（utils/assets/stereo_yolo/
|
||||
drone_sky.png，AI 合成图、非真实飞行数据）驱动，覆盖：
|
||||
|
||||
① 单图检测冒烟：真实推理 → 原生坐标框 + 出口自检（§2.2）
|
||||
② 合成双目对（已知视差）：同一无人机 patch 按已知视差贴到左右图 →
|
||||
真检测框喂 StereoRanger.update() → 与解析真值 Z=F·B/D 对比（§6 下游对接）
|
||||
③ 空场景：无目标 → 两目空帧 → 测距端返回空列表（正常空帧，不告警）
|
||||
④ 形状不一致：左右图分辨率不同 → 整帧拒收输出空（§2.2 第 2 条）
|
||||
⑤ 帧周期：双流满负荷计时（concurrency=parallel，§5 并发要求）
|
||||
|
||||
用法（从仓库根运行）：
|
||||
python utils/demo_stereo_yolo.py
|
||||
python utils/demo_stereo_yolo.py --frames 20
|
||||
|
||||
⚠ 占位参数 + 合成素材：仅限台架/回放验证，不得接入实飞控制回路（§7 红线）。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from stereo_ranger import (
|
||||
Detection,
|
||||
DetectionFrame,
|
||||
RangeRecord,
|
||||
StereoConfig,
|
||||
StereoRanger,
|
||||
load_from_json as load_ranger_config,
|
||||
)
|
||||
from stereo_yolo import StereoYoloConfig, StereoYoloFrontend, load_from_json
|
||||
|
||||
logger = logging.getLogger("demo_stereo_yolo")
|
||||
|
||||
_DEMO_DIR = Path(__file__).resolve().parent # utils/
|
||||
_REPO_ROOT = _DEMO_DIR.parent
|
||||
_DEFAULT_CONFIG = _DEMO_DIR / "config" / "stereo_yolo.json"
|
||||
_RANGER_CONFIG = _DEMO_DIR / "config" / "stereo_ranger.json"
|
||||
_ASSET = _DEMO_DIR / "assets" / "stereo_yolo" / "drone_sky.png"
|
||||
|
||||
# 合成双目对参数：粘贴同一 patch 的整数水平偏移 = 精确视差
|
||||
_DISPARITY_PX = 12 # ⇔ Z ≈ 4.8 m（F·B=57.6 px·m，域内）
|
||||
_LEFT_CENTER_X = 1200
|
||||
_PATCH_Y_CENTER = 560
|
||||
_CANVAS = (1920, 1080)
|
||||
_FEATHER_PX = 32 # 羽化混合宽度：硬边接缝会被检测器误当目标边缘（实测数 px 框偏移）
|
||||
_RANGER_NOW_MS = 1000
|
||||
|
||||
|
||||
def log_scene(title: str) -> None:
|
||||
logger.info("")
|
||||
logger.info("── %s", title)
|
||||
|
||||
|
||||
def describe(record: RangeRecord) -> str:
|
||||
z_text = f"{record.forward_distance_m:7.3f}" if record.forward_distance_m is not None else " — "
|
||||
d_text = f"{record.disparity_px:7.2f}" if record.disparity_px is not None else " — "
|
||||
return (
|
||||
f"[F{record.frame_id:02d}] {record.degradation.name:<12} valid={int(record.valid)} "
|
||||
f"Z={z_text} m D={d_text} px "
|
||||
f"c_mid=({record.center_mid_px[0]:7.1f},{record.center_mid_px[1]:6.1f}) "
|
||||
f"conf={record.confidence:.2f}"
|
||||
)
|
||||
|
||||
|
||||
def log_frame(eye: str, frame: DetectionFrame) -> None:
|
||||
if not frame.detections:
|
||||
logger.info(" %s：无检出", eye)
|
||||
return
|
||||
for detection in frame.detections:
|
||||
logger.info(
|
||||
" %s：框=(%7.1f,%7.1f,%7.1f,%7.1f) conf=%.3f class=%d",
|
||||
eye,
|
||||
detection.x1,
|
||||
detection.y1,
|
||||
detection.x2,
|
||||
detection.y2,
|
||||
detection.confidence,
|
||||
detection.class_id,
|
||||
)
|
||||
|
||||
|
||||
def _best_box(frame: DetectionFrame) -> Optional[Detection]:
|
||||
if not frame.detections:
|
||||
return None
|
||||
return max(frame.detections, key=lambda detection: detection.confidence)
|
||||
|
||||
|
||||
def _crop_patch(asset: np.ndarray, box: Detection, *, margin_x: float = 0.3, margin_y: float = 0.5) -> np.ndarray:
|
||||
width, height = box.width, box.height
|
||||
x1 = max(0, int(box.x1 - margin_x * width))
|
||||
x2 = min(asset.shape[1], int(box.x2 + margin_x * width))
|
||||
y1 = max(0, int(box.y1 - margin_y * height))
|
||||
y2 = min(asset.shape[0], int(box.y2 + margin_y * height))
|
||||
return asset[y1:y2, x1:x2].copy()
|
||||
|
||||
|
||||
def _sky_canvas(asset: np.ndarray) -> np.ndarray:
|
||||
"""用素材顶部天空条拉伸成画布，避免 patch 与背景出现色差硬边。"""
|
||||
strip = asset[0 : max(8, asset.shape[0] // 12), :, :]
|
||||
return cv2.resize(strip, (_CANVAS[0], _CANVAS[1]))
|
||||
|
||||
|
||||
def _paste_blended(canvas: np.ndarray, patch: np.ndarray, center_x: int, center_y: int) -> None:
|
||||
"""羽化混合粘贴：硬边接缝会被检测器当作目标边缘，造成数 px 的框偏移。"""
|
||||
height, width = patch.shape[:2]
|
||||
alpha = np.zeros((height, width), dtype=np.float32)
|
||||
alpha[_FEATHER_PX : height - _FEATHER_PX, _FEATHER_PX : width - _FEATHER_PX] = 1.0
|
||||
alpha = cv2.GaussianBlur(alpha, (0, 0), _FEATHER_PX * 0.5)[..., None]
|
||||
x0, y0 = int(center_x - width / 2), int(center_y - height / 2)
|
||||
region = canvas[y0 : y0 + height, x0 : x0 + width]
|
||||
canvas[y0 : y0 + height, x0 : x0 + width] = (patch * alpha + region * (1.0 - alpha)).astype(np.uint8)
|
||||
|
||||
|
||||
def build_synthetic_pair(frontend: StereoYoloFrontend, asset: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""素材上检出无人机 → 取 patch → 同一 patch 按已知整数视差贴到左右天空画布。"""
|
||||
frame, _ = frontend.detect(asset, asset, frame_id=0)
|
||||
box = _best_box(frame)
|
||||
if box is None:
|
||||
raise RuntimeError("占位素材未检出无人机，请更换 utils/assets/stereo_yolo/drone_sky.png")
|
||||
patch = _crop_patch(asset, box)
|
||||
left_img, right_img = _sky_canvas(asset), _sky_canvas(asset)
|
||||
_paste_blended(left_img, patch, _LEFT_CENTER_X, _PATCH_Y_CENTER)
|
||||
_paste_blended(right_img, patch, _LEFT_CENTER_X - _DISPARITY_PX, _PATCH_Y_CENTER)
|
||||
return left_img, right_img
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- 场景
|
||||
|
||||
|
||||
def scene_single_image(frontend: StereoYoloFrontend, asset: np.ndarray) -> None:
|
||||
log_scene("① 单图检测冒烟：真实模型 + 占位素材（同一图喂左右目，仅看输出契约）")
|
||||
left_frame, right_frame = frontend.detect(asset, asset, frame_id=1)
|
||||
log_frame("左目", left_frame)
|
||||
log_frame("右目", right_frame)
|
||||
logger.info(
|
||||
" → 框位于输入图原生坐标系（%d×%d）；出口自检/白名单已在识别端完成（§2.2 第 6 条）",
|
||||
asset.shape[1],
|
||||
asset.shape[0],
|
||||
)
|
||||
|
||||
|
||||
def scene_synthetic_stereo(
|
||||
frontend: StereoYoloFrontend,
|
||||
ranger: StereoRanger,
|
||||
ranger_cfg: StereoConfig,
|
||||
asset: np.ndarray,
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
log_scene(f"② 合成双目对（粘贴真值视差 {_DISPARITY_PX} px）→ StereoRanger.update()")
|
||||
left_img, right_img = build_synthetic_pair(frontend, asset)
|
||||
left_frame, right_frame = frontend.detect(left_img, right_img, frame_id=2)
|
||||
log_frame("左目", left_frame)
|
||||
log_frame("右目", right_frame)
|
||||
|
||||
left_box, right_box = _best_box(left_frame), _best_box(right_frame)
|
||||
if left_box is not None and right_box is not None:
|
||||
disparity = left_box.x1 - right_box.x1
|
||||
vertical_residual = abs(left_box.y1 - right_box.y1)
|
||||
logger.info(
|
||||
" 实测边缘视差 D=x1_L−x1_R=%.2f px（真值 %d px）；|Δy|=%.2f px",
|
||||
disparity,
|
||||
_DISPARITY_PX,
|
||||
vertical_residual,
|
||||
)
|
||||
logger.info(
|
||||
" 解析对照 Z=F·B/D=%.3f m（F·B=%.2f px·m）",
|
||||
ranger_cfg.focal_baseline_px_m / disparity,
|
||||
ranger_cfg.focal_baseline_px_m,
|
||||
)
|
||||
|
||||
records = ranger.update(left_frame, right_frame, now_ms=_RANGER_NOW_MS)
|
||||
for record in records:
|
||||
logger.info(" %s", describe(record))
|
||||
if left_box is not None and right_box is not None:
|
||||
logger.info(
|
||||
" → 误差来源：检测框边缘定位抖动（本合成对 |ΔD|=%.2f px），小视差下被 Z=F·B/D 放大（§7.1）",
|
||||
abs(disparity - _DISPARITY_PX),
|
||||
)
|
||||
logger.info(" → 与真值投影法（demo_stereo_ranger）同量级；真值验证待真实双目素材回填（§6）")
|
||||
return left_img, right_img
|
||||
|
||||
|
||||
def scene_empty(frontend: StereoYoloFrontend, ranger: StereoRanger, sky: np.ndarray) -> None:
|
||||
log_scene("③ 空场景：无目标 → 两目空帧 → 测距端空列表（正常空帧，不告警）")
|
||||
left_frame, right_frame = frontend.detect(sky, sky, frame_id=3)
|
||||
records = ranger.update(left_frame, right_frame, now_ms=2000)
|
||||
logger.info(
|
||||
" 两目检出数：%d / %d；测距记录数：%d（应为 0）",
|
||||
len(left_frame.detections),
|
||||
len(right_frame.detections),
|
||||
len(records),
|
||||
)
|
||||
logger.info(" → 无检测即无输出；\"目标丢失就该没有距离\"（§2.5 第 3 条）")
|
||||
|
||||
|
||||
def scene_shape_mismatch(frontend: StereoYoloFrontend, sky: np.ndarray) -> None:
|
||||
log_scene("④ 形状不一致拒收：左 1920×1080 vs 右 1280×720")
|
||||
resized = cv2.resize(sky, (1280, 720))
|
||||
left_frame, right_frame = frontend.detect(sky, resized, frame_id=4)
|
||||
logger.info(" 左右检出数：%d / %d（应为 0 / 0，WARN 节流打印）", len(left_frame.detections), len(right_frame.detections))
|
||||
logger.info(" → 禁止左右各自缩放/填充，否则破坏校正后的极线对齐（§2.2 第 2 条）")
|
||||
|
||||
|
||||
def scene_timing(frontend: StereoYoloFrontend, left_img: np.ndarray, right_img: np.ndarray, frames: int) -> None:
|
||||
log_scene(f"⑤ 帧周期：双流满负荷 {frames} 帧（concurrency=parallel）")
|
||||
for index in range(2): # 预热
|
||||
frontend.detect(left_img, right_img, frame_id=100 + index)
|
||||
started = time.perf_counter()
|
||||
for index in range(frames):
|
||||
frontend.detect(left_img, right_img, frame_id=200 + index)
|
||||
per_frame_ms = (time.perf_counter() - started) / frames * 1000.0
|
||||
logger.info(" 平均 %.1f ms/帧（%.1f fps）；30 fps 预算 33.3 ms/帧（§5）", per_frame_ms, 1000.0 / per_frame_ms)
|
||||
logger.info(" → 本值为台架开发机实测（device 配置 None = 自动选择）；平台算力预算另见平台工程文档")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- 入口
|
||||
|
||||
|
||||
def _rebase_model_path(cfg: StereoYoloConfig) -> StereoYoloConfig:
|
||||
"""demo 便利：相对 model_path 按仓库根解析（库内规则是启动工作目录）。"""
|
||||
path = Path(cfg.model_path).expanduser()
|
||||
if not path.is_absolute() and not path.is_file():
|
||||
candidate = _REPO_ROOT / path
|
||||
if candidate.is_file():
|
||||
return replace(cfg, model_path=str(candidate))
|
||||
return cfg
|
||||
|
||||
|
||||
def run_demo(config_path: str, frames: int) -> None:
|
||||
cfg = _rebase_model_path(load_from_json(config_path))
|
||||
ranger_cfg = load_ranger_config(str(_RANGER_CONFIG))
|
||||
logger.info("配置来源：%s", config_path)
|
||||
logger.info(
|
||||
"检测前端：model=%s imgsz=%d conf=%.2f iou=%.2f whitelist=%s device=%s concurrency=%s",
|
||||
Path(cfg.model_path).name,
|
||||
cfg.imgsz,
|
||||
cfg.conf_threshold,
|
||||
cfg.iou_threshold,
|
||||
cfg.class_whitelist,
|
||||
cfg.device,
|
||||
cfg.concurrency,
|
||||
)
|
||||
logger.info(
|
||||
"测距几何（唯一来源 utils/config/stereo_ranger.json）：F=%.1f px B=%.3f m F·B=%.2f px·m",
|
||||
ranger_cfg.fx_px,
|
||||
ranger_cfg.baseline_m,
|
||||
ranger_cfg.focal_baseline_px_m,
|
||||
)
|
||||
logger.warning("⚠ 占位参数与合成素材：仅限台架/回放验证，不得接入实飞控制回路（§7 红线同源）")
|
||||
|
||||
asset = cv2.imread(str(_ASSET))
|
||||
if asset is None:
|
||||
raise FileNotFoundError(f"占位素材缺失: {_ASSET}")
|
||||
|
||||
frontend = StereoYoloFrontend(cfg)
|
||||
ranger = StereoRanger(ranger_cfg)
|
||||
try:
|
||||
scene_single_image(frontend, asset)
|
||||
left_img, right_img = scene_synthetic_stereo(frontend, ranger, ranger_cfg, asset)
|
||||
scene_empty(frontend, ranger, _sky_canvas(asset))
|
||||
scene_shape_mismatch(frontend, _sky_canvas(asset))
|
||||
scene_timing(frontend, left_img, right_img, frames)
|
||||
finally:
|
||||
frontend.close()
|
||||
logger.info("")
|
||||
logger.info("演示结束：素材为 AI 生成合成图（替换 utils/assets/stereo_yolo/drone_sky.png 可换真实素材）")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="双目 YOLO 检测前端演示（真实模型 + 占位素材）")
|
||||
parser.add_argument("--config", default=str(_DEFAULT_CONFIG), help="检测前端配置 JSON 路径")
|
||||
parser.add_argument("--frames", type=int, default=10, help="帧周期场景的测量帧数")
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
||||
run_demo(args.config, args.frames)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user