如何从 COCO 数据集中过滤类？答案

【问题标题】：How to filter classes from COCO Dataset?如何从 COCO 数据集中过滤类？
【发布时间】：2019-03-25 11:03:20
【问题描述】：

我想用 COCO 数据集训练一个 yolo 模型。由于有 80 多个类，我该如何过滤呢？我只需要上课的人和车。

【问题讨论】：

标签： artificial-intelligence object-detection yolo

【解决方案1】：

为了方便和简单的方法，请按照下列步骤操作：

修改（或复制用于备份）darknet\data\coco.names 中的coco.names 文件
删除除 person 和 car 之外的所有其他类
修改你的cfg文件（例如yolov3.cfg），将第610、696、783行的3个类从80改为2
将 cfg 文件中第 603、689、776 行的 3 个过滤器从 255 更改为 (classes+5)x3 = 21
运行检测器./darknet detector test cfg/coco.data cfg/yolov3.cfg yolov3.weights data/person.jpg

如需更高级的方式，您可以使用此 repo 创建基于 voc、coco 或 open 图像的 yolo 数据集。 https://github.com/holger-prause/yolo_utils

另请参考：How can I download a specific part of Coco Dataset?

【讨论】：

【解决方案2】：

您可以使用PyCoco API 来处理COCO 数据集。有了这个库，从数据集中过滤类变得如此简单！

# Define the classes (out of the 81) which you want to see. Others will not be shown.
filterClasses = ['person', 'dog']

# Fetch class IDs only corresponding to the filterClasses
catIds = coco.getCatIds(catNms=filterClasses) 
# Get all images containing the above Category IDs
imgIds = coco.getImgIds(catIds=catIds)
print("Number of images containing all the  classes:", len(imgIds))

# load and display a random image
img = coco.loadImgs(imgIds[np.random.randint(0,len(imgIds))])[0]
I = io.imread('{}/images/{}/{}'.format(dataDir,dataType,img['file_name']))/255.0

我最近在exploring and manipulating the COCO dataset 上写了一篇完整的帖子。看看吧。

【讨论】：

如何导入pycocotools？
我刚试了filterClasses = ['person', 'dog'] catIds = coco.getCatIds(catNms=filterClasses)，得到了KeyError: 'categories'

【解决方案3】：

在 Coco 数据集上无需重新训练模型即可过滤类的唯一方法是检查检测输出以避免为无用的类画一个框，但模型将继续检测背景中的所有类。

【讨论】：

【解决方案4】：

现在最简单的方法是使用fiftyone，在COCO网站上推荐download, visualize, and evaluate the dataset including any subset of classes。

import fiftyone as fo
import fiftyone.zoo as foz

#
# Only the required images will be downloaded (if necessary).
# By default, only detections are loaded
#

dataset = foz.load_zoo_dataset(
    "coco-2017",
    splits=["validation","train"],
    classes=["person", "car"],
    # max_samples=50,
)

# Visualize the dataset in the FiftyOne App
session = fo.launch_app(dataset)

您也可以直接在数据集上将其用于convert the dataset to YOLO format 和train models。

# Export the dataset in YOLO format
export_dir = "/path/for/yolov5-dataset"
label_field = "ground_truth"

dataset.export(
    export_dir=export_dir,
    dataset_type=fo.types.YOLOv5Dataset,
    label_field=label_field,
)

安装：

pip install fiftyone

【讨论】：