如何使用pycoral API重写edgetpu旧版智能喂鸟器项目代码

前置操作 #

• 先删除你之前手动拷贝到/usr/lib/python3/dist-packages的旧版edgetpu目录，避免多版本库路径冲突
• 确保系统安装了官方最新的libedgetpu14和pycoral包，直接通过apt安装即可，无需手动处理版本
• 将原项目使用的tflite模型用最新版edgetpu_compiler重新编译，适配v14运行时，原模型训练逻辑不需要修改

迁移后完整代码 #

#!/usr/bin/python3
"""
Coral智能喂鸟器
使用PyCoral API调用EdgeTPU分析摄像头画面中的物体，检测到松鼠时播放驱离音效
支持自定义模型、标签文件、存储路径、驱离音效，可切换训练模式收集自定义模型训练用图片
"""
import argparse
import time
import logging
import gstreamer
from PIL import Image
from playsound import playsound

from pycoral.adapters import classify
from pycoral.adapters import common
from pycoral.utils.dataset import read_label_file
from pycoral.utils.edgetpu import make_interpreter

def save_data(image, results, path, ext='png'):
    """保存摄像头帧和模型推理结果到指定目录"""
    tag = '%010d' % int(time.monotonic()*1000)
    name = '%s/img-%s.%s' %(path, tag, ext)
    image.save(name)
    print('帧已保存为: %s' % name)
    logging.info('图片: %s 结果: %s', tag, results)

def print_results(start_time, last_time, end_time, results):
    """打印推理结果到终端用于调试"""
    inference_rate = ((end_time - start_time) * 1000)
    fps = (1.0/(end_time - last_time))
    print('\n推理耗时: %.2f ms, 帧率: %.2f fps' % (inference_rate, fps))
    for label, score in results:
        print(' %s, 置信度=%.2f' %(label, score))

def do_training(results, last_results, top_k):
    """对比当前和上一帧的推理结果，检测到标签变化时返回True，用于收集训练图片"""
    new_labels = [label[0] for label in results]
    old_labels = [label[0] for label in last_results]
    shared_labels  = set(new_labels).intersection(old_labels)
    if len(shared_labels) < top_k:
        print('检测到内容变化')
        return True
    return False

def user_selections():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', required=True,
                        help='.tflite模型路径')
    parser.add_argument('--labels', required=True,
                        help='标签文件路径')
    parser.add_argument('--top_k', type=int, default=3,
                        help='显示置信度最高的N个类别')
    parser.add_argument('--threshold', type=float, default=0.1,
                        help='类别置信度阈值')
    parser.add_argument('--storage', required=True,
                        help='图片和结果存储路径')
    parser.add_argument('--sound', required=True,
                        help='驱离音效文件路径')
    parser.add_argument('--print', default=False, required=False,
                        help='是否打印推理结果到终端')
    parser.add_argument('--training', default=False, required=False,
                        help='是否开启训练模式（收集图片）')
    args = parser.parse_args()
    return args

def main():
    args = user_selections()
    print("加载模型 %s，标签文件 %s" % (args.model, args.labels))
    # 替换旧版ClassificationEngine为PyCoral推理器
    interpreter = make_interpreter(args.model)
    interpreter.allocate_tensors()
    # 直接使用PyCoral内置的标签读取工具
    labels = read_label_file(args.labels)
    storage_dir = args.storage

    # 初始化日志文件
    logging.basicConfig(filename='%s/results.log' % storage_dir,
                        format='%(asctime)s-%(message)s',
                        level=logging.DEBUG)

    last_time = time.monotonic()
    last_results = [('label', 0)]
    def user_callback(image, svg_canvas):
        nonlocal last_time
        nonlocal last_results
        start_time = time.monotonic()
        # PyCoral推理流程
        size = common.input_size(interpreter)
        resized_image = image.convert('RGB').resize(size, Image.Resampling.LANCZOS)
        common.set_input(interpreter, resized_image)
        interpreter.invoke()
        results = classify.get_classes(interpreter, top_k=args.top_k, score_threshold=args.threshold)
        end_time = time.monotonic()
        results = [(labels.get(i.id, f'类别_{i.id}'), i.score) for i in results]

        if args.print:
            print_results(start_time, last_time, end_time, results)

        if args.training:
            if do_training(results, last_results, args.top_k):
                save_data(image, results, storage_dir)
        else:
            # 自定义模型模式：可以修改下面的标签判断逻辑检测/驱离指定物体
            if results and results[0][0] != 'background':
                save_data(image, results, storage_dir)
            label_list = [item[0] for item in results]
            if 'fox squirrel, eastern fox squirrel, Sciurus niger' in label_list:
                playsound(args.sound)
                logging.info('已触发驱离音效')

        last_results = results
        last_time = end_time
    result = gstreamer.run_pipeline(user_callback)

if __name__ == '__main__':
    main()

核心修改说明 #

• 完全移除了旧版edgetpu库的导入和路径追加逻辑，所有推理功能基于官方PyCoral API实现，适配v14运行时
• 替换废弃的ClassificationEngine为PyCoral的Interpreter推理流程，兼容最新版EdgeTPU运行时
• 复用PyCoral内置的read_label_file读取标签，避免自定义解析逻辑出错
• 修复了原代码非训练模式下save_data函数参数顺序错误的问题
• 增加了标签不存在时的兜底处理，避免标签缺失导致程序崩溃
• 新增推理前的图片尺寸适配和格式转换，兼容不同输入尺寸的模型

内容的提问来源于stack exchange，提问作者taboopython