如何移除.pb格式训练模型中的无关物体？Tello无人机模型优化咨询

1. 移除.pb格式已训练模型中的指定物体 #

直接修改已冻结的.pb模型（TensorFlow格式）难度极大，因为模型的输出层结构、权重都是和原始训练的类别绑定的，更高效的方案是以下两种：

• 方案1：基于原始训练流程重新训练

  • 找到项目中的类别配置文件（如labelmap.pbtxt），删除自行车等不需要的类别条目，仅保留车辆相关类别（如car、bus、truck等）
  • 调整训练脚本中的num_classes参数，使其与修改后的类别数量一致
  • 使用仅包含车辆类标注的数据集重新训练，生成只识别目标类别的新.pb模型

• 方案2：推理阶段做后处理过滤

  • 在模型推理代码中，获取检测结果后添加过滤逻辑，仅保留车辆类别的检测结果，丢弃无关类数据
  • 示例代码片段：

    # 假设车辆类别对应的ID为[2, 3, 7]（需根据项目实际label映射调整）
    target_class_ids = [2, 3, 7]
    # 过滤检测结果
    filtered_detections = [det for det in raw_detections if det['class_id'] in target_class_ids]
    

2. 针对该Tello无人机检测模型删除无关识别类别 #

结合你使用的Tello项目，从数据、配置、训练、推理四个环节调整：

• 清理训练数据集
  • 检查项目中的数据集文件，删除所有无关类别的图片及对应的标注文件（如VOC格式的xml标注中，删除为bicycle等的条目）
• 修改类别配置文件
  • 编辑项目中的labelmap.pbtxt，移除自行车等无关类的定义，确保剩余类别ID连续
  • 更新训练配置中的num_classes参数，设置为调整后的类别总数
• 重新训练模型
  • 运行项目中的训练脚本，使用清理后的数据集和修改后的配置，重新训练得到仅识别车辆的模型
• 调整推理代码（快速临时方案）
  • 若暂时无需重新训练，直接修改Tello实时检测的代码逻辑：在检测结果输出前，判断类别ID，仅绘制车辆类别的检测框，忽略无关类结果

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