210 lines
6.2 KiB
Python
210 lines
6.2 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
摄像头实时 YOLOv5 简易测试脚本。
|
||
|
||
特点：
|
||
- 完全独立脚本，直接 python test/test_yolo_camera_simple.py 运行，不需要传参。
|
||
- 不 import config，不依赖项目模块。
|
||
- 直接调用 maix.nn.YOLOv5(model=..., dual_buff=False)。
|
||
- camera.read() 得到的 Maix image 直接送 det.detect()。
|
||
- 在画面上画检测框、类别、置信度，并显示到屏幕。
|
||
|
||
运行环境：MaixCAM / MaixPy。
|
||
"""
|
||
|
||
import os
|
||
|
||
|
||
CAMERA_WIDTH = 640
|
||
CAMERA_HEIGHT = 480
|
||
# 默认与主项目 config.TRIANGLE_YOLO_MODEL_PATH 一致（勿用 /root/yolo26_int8.mud，那是占位路径）
|
||
_MODEL_DEFAULT = "/maixapp/apps/t11/model_270139.mud"
|
||
try:
|
||
import config as _cfg
|
||
|
||
MODEL_PATH = getattr(_cfg, "TRIANGLE_YOLO_MODEL_PATH", _MODEL_DEFAULT) or _MODEL_DEFAULT
|
||
except Exception:
|
||
MODEL_PATH = _MODEL_DEFAULT
|
||
CONF_TH = 0.7
|
||
IOU_TH = 0.45
|
||
# native: Maix detect 返回框已映射到 camera.read() 图像坐标；letterbox: 需要从网络输入坐标反算
|
||
COORD_MODE = "native"
|
||
# 只用于 DRAW_ONLY_CLASS_IDS=True 时过滤显示；默认画所有框
|
||
CLASS_IDS = (0,)
|
||
DRAW_ONLY_CLASS_IDS = False # True=只画 CLASS_IDS 里的类别；False=画所有 YOLO 返回框
|
||
|
||
|
||
def _det_obj_class_id(o):
|
||
for key in ("class_id", "cls", "label", "category", "cat_id", "id"):
|
||
if hasattr(o, key):
|
||
v = getattr(o, key)
|
||
if v is None:
|
||
continue
|
||
try:
|
||
return int(float(v))
|
||
except (TypeError, ValueError):
|
||
continue
|
||
return None
|
||
|
||
|
||
def _det_obj_from_seq(t):
|
||
if not isinstance(t, (list, tuple)) or len(t) < 6:
|
||
return None
|
||
|
||
class Box:
|
||
pass
|
||
|
||
b = Box()
|
||
b.x = float(t[0])
|
||
b.y = float(t[1])
|
||
b.w = float(t[2])
|
||
b.h = float(t[3])
|
||
b.score = float(t[4])
|
||
b.class_id = int(float(t[5]))
|
||
return b
|
||
|
||
|
||
def _normalize_objs(objs):
|
||
out = []
|
||
for o in objs or []:
|
||
if isinstance(o, (list, tuple)):
|
||
m = _det_obj_from_seq(o)
|
||
if m is not None:
|
||
out.append(m)
|
||
else:
|
||
out.append(o)
|
||
return out
|
||
|
||
|
||
def _letterbox_net_to_src_xyxy(x, y, w, h, src_w, src_h, net_w, net_h):
|
||
scale = min(net_w / float(src_w), net_h / float(src_h))
|
||
new_w = src_w * scale
|
||
new_h = src_h * scale
|
||
pad_x = (net_w - new_w) * 0.5
|
||
pad_y = (net_h - new_h) * 0.5
|
||
x0 = (x - pad_x) / scale
|
||
y0 = (y - pad_y) / scale
|
||
x1 = (x + w - pad_x) / scale
|
||
y1 = (y + h - pad_y) / scale
|
||
return x0, y0, x1, y1
|
||
|
||
|
||
def _det_to_src_xyxy(o, coord_mode, src_w, src_h, net_w, net_h):
|
||
x = float(getattr(o, "x", 0.0))
|
||
y = float(getattr(o, "y", 0.0))
|
||
w = float(getattr(o, "w", 0.0))
|
||
h = float(getattr(o, "h", 0.0))
|
||
if coord_mode in ("native", "source", "camera", "full"):
|
||
return x, y, x + w, y + h
|
||
return _letterbox_net_to_src_xyxy(x, y, w, h, src_w, src_h, net_w, net_h)
|
||
|
||
|
||
def _clip_xywh(x0, y0, x1, y1, src_w, src_h):
|
||
x0 = max(0, min(int(round(x0)), src_w - 1))
|
||
y0 = max(0, min(int(round(y0)), src_h - 1))
|
||
x1 = max(x0 + 1, min(int(round(x1)), src_w))
|
||
y1 = max(y0 + 1, min(int(round(y1)), src_h))
|
||
return x0, y0, x1 - x0, y1 - y0
|
||
|
||
|
||
def _label(det, cid):
|
||
labels = getattr(det, "labels", None)
|
||
if labels is None:
|
||
return str(cid)
|
||
try:
|
||
return str(labels[int(cid)])
|
||
except Exception:
|
||
return str(cid)
|
||
|
||
|
||
def main():
|
||
from maix import camera, display, nn, time, image
|
||
|
||
if not MODEL_PATH or not os.path.isfile(MODEL_PATH):
|
||
print("[ERR] 模型文件不存在:", MODEL_PATH)
|
||
return
|
||
|
||
print("[INFO] 初始化 YOLO 模型:", MODEL_PATH)
|
||
det = nn.YOLOv26(model=MODEL_PATH, dual_buff=False)
|
||
net_w = int(det.input_width())
|
||
net_h = int(det.input_height())
|
||
print(
|
||
"[INFO] net_in=%dx%d conf=%.2f iou=%.2f coord=%s class_ids=%s"
|
||
% (net_w, net_h, CONF_TH, IOU_TH, COORD_MODE, str(CLASS_IDS))
|
||
)
|
||
|
||
print("[INFO] 初始化摄像头: %dx%d" % (CAMERA_WIDTH, CAMERA_HEIGHT))
|
||
cam = camera.Camera(CAMERA_WIDTH, CAMERA_HEIGHT)
|
||
disp = display.Display()
|
||
|
||
color_cycle = []
|
||
for name in ("RED", "GREEN", "BLUE", "ORANGE", "YELLOW", "CYAN", "MAGENTA"):
|
||
c = getattr(image, "COLOR_" + name, None)
|
||
if c is not None:
|
||
color_cycle.append(c)
|
||
if not color_cycle:
|
||
color_cycle = [getattr(image, "COLOR_RED", 0)]
|
||
|
||
frame_idx = 0
|
||
last_log_ms = time.ticks_ms()
|
||
fps_count = 0
|
||
|
||
while True:
|
||
frame = cam.read()
|
||
src_w = frame.width()
|
||
src_h = frame.height()
|
||
|
||
t0 = time.ticks_ms()
|
||
raw = det.detect(frame, conf_th=CONF_TH, iou_th=IOU_TH)
|
||
detect_ms = time.ticks_ms() - t0
|
||
objs = _normalize_objs(raw if raw is not None else [])
|
||
|
||
draw_count = 0
|
||
for i, o in enumerate(objs):
|
||
cid = _det_obj_class_id(o)
|
||
if cid is None:
|
||
cid = -1
|
||
if DRAW_ONLY_CLASS_IDS and cid not in CLASS_IDS:
|
||
continue
|
||
try:
|
||
score = float(getattr(o, "score", 0.0))
|
||
except Exception:
|
||
score = 0.0
|
||
|
||
x0, y0, x1, y1 = _det_to_src_xyxy(o, COORD_MODE, src_w, src_h, net_w, net_h)
|
||
ix, iy, iw, ih = _clip_xywh(x0, y0, x1, y1, src_w, src_h)
|
||
col = color_cycle[cid % len(color_cycle)] if cid >= 0 else color_cycle[0]
|
||
frame.draw_rect(ix, iy, iw, ih, color=col)
|
||
frame.draw_string(ix, max(0, iy - 16), "%s %.2f" % (_label(det, cid), score), color=col)
|
||
draw_count += 1
|
||
|
||
frame.draw_string(4, 4, "YOLO boxes:%d draw:%d %dms" % (len(objs), draw_count, detect_ms), color=color_cycle[0])
|
||
disp.show(frame)
|
||
|
||
frame_idx += 1
|
||
fps_count += 1
|
||
now = time.ticks_ms()
|
||
if now - last_log_ms >= 1000:
|
||
print(
|
||
"[INFO] frame=%d fps=%d raw_boxes=%d draw_boxes=%d detect_ms=%d"
|
||
% (frame_idx, fps_count, len(objs), draw_count, detect_ms)
|
||
)
|
||
fps_count = 0
|
||
last_log_ms = now
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except KeyboardInterrupt:
|
||
print("[INFO] exit")
|
||
except Exception as e:
|
||
print("[ERR]", e)
|
||
try:
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
except Exception:
|
||
pass
|