109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""Run from MaixVision on PC to inspect the box's live 20/40 YOLO output."""
|
|
|
|
import os
|
|
|
|
from maix import app, camera, display, image, nn, time
|
|
|
|
|
|
# This file is sent to /tmp/maixpy_run by MaixVision. Keep the model path
|
|
# absolute so the script uses the model already installed on the box.
|
|
MODEL_PATH = "/maixapp/apps/t11/model_317181.mud"
|
|
CAMERA_WIDTH = 640
|
|
CAMERA_HEIGHT = 480
|
|
CONF_TH = 0.65
|
|
IOU_TH = 0.45
|
|
|
|
|
|
def _flatten_objects(raw):
|
|
if raw is None:
|
|
return []
|
|
if isinstance(raw, (list, tuple)):
|
|
result = []
|
|
for item in raw:
|
|
if isinstance(item, (list, tuple)):
|
|
result.extend(_flatten_objects(item))
|
|
else:
|
|
result.append(item)
|
|
return result
|
|
return [raw]
|
|
|
|
|
|
def main():
|
|
if not os.path.isfile(MODEL_PATH):
|
|
raise FileNotFoundError("model not found on box: " + MODEL_PATH)
|
|
|
|
detector = nn.YOLOv5(model=MODEL_PATH, dual_buff=False)
|
|
cam = camera.Camera(CAMERA_WIDTH, CAMERA_HEIGHT)
|
|
disp = display.Display()
|
|
|
|
labels = tuple(str(label) for label in detector.labels)
|
|
print("[YOLO] model:", MODEL_PATH)
|
|
print("[YOLO] labels:", labels)
|
|
print("[YOLO] conf=%.2f iou=%.2f" % (CONF_TH, IOU_TH))
|
|
|
|
fps = 0.0
|
|
frame_count = 0
|
|
last_log_ms = time.ticks_ms()
|
|
|
|
while not app.need_exit():
|
|
loop_start_ms = time.ticks_ms()
|
|
img = cam.read()
|
|
|
|
detect_start_ms = time.ticks_ms()
|
|
raw = detector.detect(img, conf_th=CONF_TH, iou_th=IOU_TH)
|
|
detect_ms = max(0, time.ticks_diff(time.ticks_ms(), detect_start_ms))
|
|
objects = _flatten_objects(raw)
|
|
|
|
candidates = []
|
|
for obj in objects:
|
|
class_id = int(obj.class_id)
|
|
score = float(obj.score)
|
|
label = labels[class_id] if 0 <= class_id < len(labels) else "unknown"
|
|
color = image.COLOR_GREEN if label in ("20", "40") else image.COLOR_RED
|
|
|
|
img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=color)
|
|
img.draw_string(
|
|
obj.x,
|
|
max(0, obj.y - 16),
|
|
"%scm %.2f" % (label, score),
|
|
color=color,
|
|
)
|
|
if label in ("20", "40"):
|
|
candidates.append((score, label))
|
|
|
|
loop_ms = max(1, time.ticks_diff(time.ticks_ms(), loop_start_ms))
|
|
instant_fps = 1000.0 / float(loop_ms)
|
|
fps = instant_fps if frame_count == 0 else fps * 0.9 + instant_fps * 0.1
|
|
|
|
if candidates:
|
|
best_score, best_label = max(candidates, key=lambda item: item[0])
|
|
status = "TARGET %scm %.2f" % (best_label, best_score)
|
|
status_color = image.COLOR_GREEN
|
|
else:
|
|
status = "TARGET UNKNOWN"
|
|
status_color = image.COLOR_RED
|
|
|
|
img.draw_string(5, 5, status, color=status_color)
|
|
img.draw_string(
|
|
5,
|
|
25,
|
|
"infer=%dms fps=%.1f boxes=%d" % (detect_ms, fps, len(objects)),
|
|
color=image.COLOR_YELLOW,
|
|
)
|
|
disp.show(img)
|
|
|
|
frame_count += 1
|
|
now_ms = time.ticks_ms()
|
|
if time.ticks_diff(now_ms, last_log_ms) >= 1000:
|
|
print(
|
|
"[YOLO] %s infer=%dms fps=%.1f boxes=%d"
|
|
% (status, detect_ms, fps, len(objects))
|
|
)
|
|
last_log_ms = now_ms
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|