这是我的wuliao_detection库，你学习一下 from maix import camera, display, image, nn, app, uart from collections import defaultdict, deque import random import struct device = "/dev/ttyS0" serial0 = uart.UART(device, 115200) class YOLODetector: def __init__(self, model_path, conf_th=0.9, iou_th=0.45): self.model = nn.YOLOv5(model=model_path, dual_buff=True) self.conf_th = conf_th self.iou_th = iou_th self.labels = self.model.labels self.color_map = self._generate_color_map() def _generate_color_map(self): # 生成 image.Color 对象 return { idx: image.Color.from_rgb( random.randint(0, 255), random.randint(0, 255), random.randint(0, 255) ) for idx in range(len(self.labels)) } def detect(self, img): return self.model.detect(img, conf_th=self.conf_th, iou_th=self.iou_th) class DetectionVisualizer: def __init__(self, trail_length=15): self.trail_history = defaultdict(lambda: deque(maxlen=trail_length)) def draw_detection(self, img, obj, label, color): # 确保所有绘图函数使用 image.Color 类型 img.draw_rect(obj.x, obj.y, obj.w, obj.h, color) img.draw_string(obj.x, obj.y, f"{label}: {obj.score:.2f}", color) cx = obj.x + obj.w // 2 cy = obj.y + obj.h // 2 self._draw_center_marker(img, cx, cy, color) img.draw_string(cx-20, cy+15, f"({cx}, {cy})", color) return cx, cy def update_trail(self, class_id, cx, cy): self.trail_history[class_id].append((cx, cy)) def draw_trails(self, img, class_id, color): history = self.trail_history[class_id] for i in range(1, len(history)): x1, y1 = history[i-1] x2, y2 = history[i] img.draw_line(x1, y1, x2, y2, color, thickness=2) def _draw_center_marker(self, img, x, y, color): img.draw_circle(x, y, 3, color) img.draw_line(x-10, y, x+10, y, color) img.draw_line(x, y-10, x, y+10, color) def main(): detector = YOLODetector("/root/my_wuliao/model_183042.mud", conf_th=0.7) cam = camera.Camera(detector.model.input_width(), detector.model.input_height(), detector.model.input_format()) visualizer = DetectionVisualizer(trail_length=20) disp = display.Display() while not app.need_exit():
最新发布
扫一扫
590
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
