from maix import image, nn, display

# 1. 加载模型
detector = nn.YOLOv8(model="/root/models/yolov8n.mud", dual_buff=True)

# 2. 加载指定图片（根据模型输入尺寸自动缩放宽高）
img = image.load("/root/test.jpg")
if img is None:
    raise FileNotFoundError("图片加载失败，请检查路径")

# 3. 调整图片尺寸到模型输入要求（可选，detect内部会处理，但提前缩放可提高速度）
# img = img.resize(detector.input_width(), detector.input_height())

# 4. 检测
objs = detector.detect(img, conf_th=0.5, iou_th=0.45)

# 5. 在图片上绘制结果
for obj in objs:
    img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=image.COLOR_RED)
    msg = f'{detector.labels[obj.class_id]}: {obj.score:.2f}'
    img.draw_string(obj.x, obj.y, msg, color=image.COLOR_RED)

# 6. 显示结果（如果设备有屏幕）
disp = display.Display()
disp.show(img)

# 7. 保存结果（可选）
img.save("/root/result.jpg")

print("识别完成，结果已显示并保存为 result.jpg")