import pyb
import sensor, image, time, math
import os, tf

# 初始化摄像头
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_brightness(800)
sensor.skip_frames(time=2000)  # 让摄像头稳定下来
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(True, (0, 0x80, 0))

clock = time.clock()

# 模型路径和标签加载
net_path = "mobilenet_v2-2024-03-08-99.6(1).tflite"
labels = [line.rstrip() for line in open("/sd/labels.txt")]
net = tf.load(net_path, load_to_fb=True)

# 创建字符串映射到数字的字典
string_to_number = {label: index + 1 for index, label in enumerate(labels)}

# 主循环
while True:
    clock.tick()  # 更新时钟

    # 拍摄图像
    img = sensor.snapshot()

    # 在图像中搜索矩形
    for r in img.find_rects(threshold=50000):
        img.draw_rectangle(r.rect(), color=(255, 0, 0))  # 绘制矩形外框

        # 提取矩形框内的图像
        img1 = img.copy(roi=r.rect())

        # 在矩形框内进行目标分类
        for obj in tf.classify(net, img1, min_scale=1.0, scale_mul=0.5, x_overlap=0.0, y_overlap=0.0):
            print("**********\nTop 1 Detections at [x=%d,y=%d,w=%d,h=%d]" % obj.rect())

            # 将分类结果和置信度进行排序
            sorted_list = sorted(zip(labels, obj.output()), key=lambda x: x[1], reverse=True)

            # 打印置信度最高的结果
            if sorted_list and len(sorted_list) > 0:
                print("%s = %f" % (sorted_list[0][0], sorted_list[0][1]))

                # 直接输出数字，而不是列表
                print(string_to_number[sorted_list[0][0]])
            else:
                print("No objects detected.")
