为多个对象创建 TFrecord 文件

Question

我正在关注来自 GitHub (https://github.com/datitran/raccoon_dataset) 的浣熊检测教程，以使用 google 对象检测 API 检测动物。为此，我需要生成 tfrecord 文件，该文件已在第 29 到 34 行 (https://github.com/datitran/raccoon_dataset/blob/master/generate_tfrecord.py) 上生成。

但他只为一种动物（浣熊）第 29 行到第 34 行编写了代码。我有多种动物，如浣熊、螳螂、寄居蟹等，我如何为多种动物修改此 tfrecord 文件。我发现的一种方法是更改​​ generatetfrecord 文件中的第 29 到 34 行，如下所示

def class_text_to_int(row_label):
    if row_label == 'raccoon':
        return 1
    if row_label == 'prayingmantis':
        return 2
    else:
        None

这种方法是否可以在同一个文件中包含多个 if 或者我需要生成多个 tfrecord 文件来训练多个对象

Answer 1

def class_text_to_int(row_label):
    if row_label == 'raccoon':
        return 1
    elif row_label == 'prayingmantis':
        return 2
    else:
        None

在 generatetfrecord 中进行上述更改是同时训练多个自定义图像的正确方法。此外，在 object detection.pbtxt 文件中，进行以下更改：

item{
       id:1
       name:'racoon'
      }
item{
       id:2
       name: 'prayingmantis'
      }

并从第一步开始重新训练模型。