243 lines
9.7 KiB
Python
243 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
单张图片快速测试：三角形四角标记识别 + 单应性落点 + PnP 估距
|
||
|
||
用法（在板子上）：
|
||
python3 test/test_triangle_one_image.py --image /root/phot/xxx.jpg --out /root/phot/tri_out.jpg
|
||
|
||
调参对比（不改代码，临时覆盖 config.TRIANGLE_*）：
|
||
python3 test/test_triangle_one_image.py --image /root/phot/xxx.jpg --preset shadow
|
||
python3 test/test_triangle_one_image.py --image /root/phot/xxx.jpg --max-interior-gray 160 --min-dark-ratio 0.20
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import time
|
||
from typing import Any, Dict, Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
import config
|
||
import triangle_target as tri_mod
|
||
from triangle_target import (
|
||
detect_triangle_markers,
|
||
load_camera_from_xml,
|
||
load_triangle_positions,
|
||
try_triangle_scoring,
|
||
)
|
||
|
||
|
||
def _apply_overrides(args) -> None:
|
||
# 预设：阴影/低对比度场景更宽松（尽量保持速度：不启 CLAHE）
|
||
if args.preset == "shadow":
|
||
setattr(config, "TRIANGLE_ENABLE_CLAHE_FALLBACK", False)
|
||
setattr(config, "TRIANGLE_MIN_CONTRAST_DIFF", 0)
|
||
setattr(config, "TRIANGLE_MAX_INTERIOR_GRAY", 160)
|
||
setattr(config, "TRIANGLE_DARK_PIXEL_GRAY", 160)
|
||
setattr(config, "TRIANGLE_MIN_DARK_RATIO", 0.20)
|
||
# adaptive 只在 Otsu 失败时尝试，保持尝试次数很少
|
||
setattr(config, "TRIANGLE_ADAPTIVE_BLOCK_SIZES", (21,))
|
||
|
||
# 手动覆盖（优先级高于 preset）
|
||
if args.max_interior_gray is not None:
|
||
setattr(config, "TRIANGLE_MAX_INTERIOR_GRAY", int(args.max_interior_gray))
|
||
if args.dark_pixel_gray is not None:
|
||
setattr(config, "TRIANGLE_DARK_PIXEL_GRAY", int(args.dark_pixel_gray))
|
||
if args.min_dark_ratio is not None:
|
||
setattr(config, "TRIANGLE_MIN_DARK_RATIO", float(args.min_dark_ratio))
|
||
if args.min_contrast_diff is not None:
|
||
setattr(config, "TRIANGLE_MIN_CONTRAST_DIFF", int(args.min_contrast_diff))
|
||
if args.detect_scale is not None:
|
||
setattr(config, "TRIANGLE_DETECT_SCALE", float(args.detect_scale))
|
||
if args.adaptive_blocks is not None:
|
||
bs = tuple(int(x) for x in args.adaptive_blocks.split(",") if x.strip())
|
||
setattr(config, "TRIANGLE_ADAPTIVE_BLOCK_SIZES", bs)
|
||
|
||
|
||
def _dump_config() -> Dict[str, Any]:
|
||
keys = [
|
||
"TRIANGLE_DETECT_SCALE",
|
||
"TRIANGLE_SIZE_RANGE",
|
||
"TRIANGLE_MAX_INTERIOR_GRAY",
|
||
"TRIANGLE_DARK_PIXEL_GRAY",
|
||
"TRIANGLE_MIN_DARK_RATIO",
|
||
"TRIANGLE_MIN_CONTRAST_DIFF",
|
||
"TRIANGLE_ADAPTIVE_BLOCK_SIZES",
|
||
"TRIANGLE_MAX_FILTERED_FOR_COMBO",
|
||
"TRIANGLE_EARLY_EXIT_CANDIDATES",
|
||
"TRIANGLE_ENABLE_CLAHE_FALLBACK",
|
||
]
|
||
out = {}
|
||
for k in keys:
|
||
out[k] = getattr(config, k, None)
|
||
return out
|
||
|
||
|
||
def _draw_tri_debug(img_bgr: np.ndarray, tri: Dict[str, Any]) -> np.ndarray:
|
||
out = img_bgr.copy()
|
||
markers = tri.get("markers") or []
|
||
|
||
# 画三角形轮廓 + center + id
|
||
for m in markers:
|
||
corners = np.array(m.get("corners", []), dtype=np.int32)
|
||
if corners.size == 0:
|
||
continue
|
||
cv2.polylines(out, [corners], True, (0, 255, 0), 2)
|
||
c = m.get("center") or (corners[:, 0].mean(), corners[:, 1].mean())
|
||
cx, cy = int(c[0]), int(c[1])
|
||
cv2.circle(out, (cx, cy), 4, (0, 0, 255), -1)
|
||
mid = m.get("id", "?")
|
||
cv2.putText(out, f"T{mid}", (cx - 18, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0), 1)
|
||
|
||
# 若有 homography，画靶心（把 (0,0) 反投影到图像）
|
||
H = tri.get("homography")
|
||
if H is not None:
|
||
try:
|
||
H = np.array(H, dtype=np.float64)
|
||
H_inv = np.linalg.inv(H)
|
||
c_img = cv2.perspectiveTransform(np.array([[[0.0, 0.0]]], dtype=np.float32), H_inv)[0][0]
|
||
ocx, ocy = int(c_img[0]), int(c_img[1])
|
||
cv2.circle(out, (ocx, ocy), 5, (0, 0, 255), -1)
|
||
cv2.circle(out, (ocx, ocy), 10, (0, 0, 255), 1)
|
||
except Exception:
|
||
pass
|
||
|
||
# 叠加结果信息
|
||
lines = []
|
||
if tri.get("ok"):
|
||
lines.append("tri_ok=True")
|
||
if tri.get("dx_cm") is not None and tri.get("dy_cm") is not None:
|
||
lines.append(f"dx,dy=({tri['dx_cm']:.2f},{tri['dy_cm']:.2f})cm")
|
||
if tri.get("distance_m") is not None:
|
||
lines.append(f"dist={float(tri['distance_m']):.2f}m")
|
||
else:
|
||
lines.append("tri_ok=False")
|
||
|
||
y0 = 22
|
||
for i, t in enumerate(lines):
|
||
cv2.putText(out, t, (10, y0 + i * 18), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0), 1)
|
||
return out
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--image", required=True, help="输入图片路径（jpg/png）")
|
||
ap.add_argument("--out", default="", help="输出标注图片路径（可选）")
|
||
ap.add_argument("--laser-x", type=int, default=-1, help="激光点 x（像素），默认用图像中心")
|
||
ap.add_argument("--laser-y", type=int, default=-1, help="激光点 y（像素），默认用图像中心")
|
||
ap.add_argument("--preset", choices=["", "shadow"], default="", help="调参预设（shadow=阴影更鲁棒，不启 CLAHE）")
|
||
ap.add_argument("--max-interior-gray", type=int, default=None)
|
||
ap.add_argument("--dark-pixel-gray", type=int, default=None)
|
||
ap.add_argument("--min-dark-ratio", type=float, default=None)
|
||
ap.add_argument("--min-contrast-diff", type=int, default=None)
|
||
ap.add_argument("--detect-scale", type=float, default=None)
|
||
ap.add_argument("--adaptive-blocks", default=None, help="例如: 11,21 ；为空表示不改")
|
||
ap.add_argument("--verbose", action="store_true", help="输出更多检测阶段信息")
|
||
args = ap.parse_args()
|
||
|
||
_apply_overrides(args)
|
||
# triangle_target.py 的日志默认写到 logger_manager；在离线脚本里 logger 可能未初始化。
|
||
# verbose 模式下把 _log 重定向为 print，方便直接看到诊断信息。
|
||
if args.verbose:
|
||
try:
|
||
tri_mod._log = lambda msg: print(str(msg))
|
||
except Exception:
|
||
pass
|
||
|
||
img_bgr = cv2.imread(args.image, cv2.IMREAD_COLOR)
|
||
if img_bgr is None:
|
||
raise SystemExit(f"读图失败：{args.image}")
|
||
# triangle_target.try_triangle_scoring 约定输入为 RGB；OpenCV imread 返回 BGR
|
||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||
|
||
h, w = img_bgr.shape[:2]
|
||
if args.laser_x >= 0 and args.laser_y >= 0:
|
||
laser_point = (int(args.laser_x), int(args.laser_y))
|
||
else:
|
||
laser_point = (w // 2, h // 2)
|
||
|
||
K, dist = load_camera_from_xml(getattr(config, "CAMERA_CALIB_XML", ""))
|
||
pos = load_triangle_positions(getattr(config, "TRIANGLE_POSITIONS_JSON", ""))
|
||
|
||
print("[tri-test] image:", args.image, "shape:", (h, w))
|
||
print("[tri-test] laser_point:", laser_point)
|
||
print("[tri-test] calib_ok:", bool(K is not None and dist is not None), "pos_ok:", bool(pos))
|
||
print("[tri-test] config:", json.dumps(_dump_config(), ensure_ascii=False))
|
||
|
||
# 先单独跑一次三角形候选检测，便于区分“没找到候选” vs “找到候选但评分/单应性失败”
|
||
scale = float(getattr(config, "TRIANGLE_DETECT_SCALE", 0.5) or 0.5)
|
||
if not (0.05 <= scale <= 1.0):
|
||
scale = 0.5
|
||
long_side = max(h, w)
|
||
max_dim = max(64, int(long_side * scale))
|
||
if long_side > max_dim:
|
||
det_scale = max_dim / long_side
|
||
det_w = int(w * det_scale)
|
||
det_h = int(h * det_scale)
|
||
img_det = cv2.resize(img_bgr, (det_w, det_h), interpolation=cv2.INTER_LINEAR)
|
||
inv_scale = 1.0 / det_scale
|
||
size_range_det = (
|
||
max(4, int(getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500))[0] * det_scale)),
|
||
max(8, int(getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500))[1] * det_scale)),
|
||
)
|
||
else:
|
||
img_det = img_bgr
|
||
inv_scale = 1.0
|
||
size_range_det = getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500))
|
||
|
||
gray = cv2.cvtColor(img_det, cv2.COLOR_BGR2GRAY)
|
||
markers_det = detect_triangle_markers(
|
||
gray,
|
||
orig_gray=gray,
|
||
size_range=size_range_det,
|
||
verbose=bool(args.verbose),
|
||
)
|
||
if inv_scale != 1.0 and markers_det:
|
||
for m in markers_det:
|
||
m["center"] = [m["center"][0] * inv_scale, m["center"][1] * inv_scale]
|
||
m["corners"] = [[c[0] * inv_scale, c[1] * inv_scale] for c in m["corners"]]
|
||
|
||
print("[tri-test] markers_found:", len(markers_det), "ids:", [m.get("id") for m in markers_det])
|
||
|
||
t0 = time.time()
|
||
tri = try_triangle_scoring(
|
||
img_rgb, # try_triangle_scoring 期望 RGB
|
||
laser_point,
|
||
pos,
|
||
K,
|
||
dist,
|
||
size_range=getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500)),
|
||
)
|
||
dt_ms = int(round((time.time() - t0) * 1000))
|
||
|
||
print("[tri-test] elapsed_ms:", dt_ms)
|
||
print(json.dumps(tri, ensure_ascii=False, indent=2))
|
||
|
||
if args.out:
|
||
out_path = args.out
|
||
# 允许传目录（如 ./），自动生成文件名；未带扩展名时默认 .jpg
|
||
if out_path.endswith("/") or out_path.endswith("\\") or os.path.isdir(out_path):
|
||
out_path = os.path.join(out_path, "tri_out.jpg")
|
||
root, ext = os.path.splitext(out_path)
|
||
if not ext:
|
||
out_path = root + ".jpg"
|
||
|
||
# 若 try_triangle_scoring 失败且没带回 markers，至少把候选 markers 画出来，方便肉眼判断
|
||
tri_for_draw = tri if isinstance(tri, dict) else {"ok": False}
|
||
if not tri_for_draw.get("markers") and markers_det:
|
||
tri_for_draw = dict(tri_for_draw)
|
||
tri_for_draw["markers"] = markers_det
|
||
out_img = _draw_tri_debug(img_bgr, tri_for_draw)
|
||
ok = cv2.imwrite(out_path, out_img)
|
||
if not ok:
|
||
raise SystemExit(f"写图失败（可能是不支持的扩展名）：{out_path}")
|
||
print("[tri-test] wrote:", out_path)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|