如何从 COCO 数据集或 Open Images V6 数据集中提取前景对象？

Question

目前，我正在为目标检测任务准备一个合成数据集。有可用于此类任务的注释数据集，例如 COCO 数据集和 Open Images V6。我正在尝试从那里下载图像，但只下载特定 class 的前景对象，例如人，换句话说，没有透明背景的图像。我这样做的原因是我想在将这些图像编辑成新图像后插入这些图像，例如街景。

到目前为止我已经尝试过，我使用了一个名为 FiftyOne 的库，我下载了带有语义标签的数据集，但我被困在这里，我无能为力。

没有必要使用 FiftyOne，任何其他方法都可以。

这是我用来下载带有标签的数据集样本的代码

import fiftyone as fo
import fiftyone.zoo as foz

dataset = foz.load_zoo_dataset(
    "coco-2017",
    split="validation",
    dataset_dir = "path/fiftyone",
    label_types=["segmentations"],
    classes = ["person"],
    max_samples=10,
    label_field="instances",
    dataset_name="coco-images-person",
)

# Export the dataset
dataset.export(
    export_dir = "path/fiftyone/image-segmentation-dataset",
    dataset_type=fo.types.ImageSegmentationDirectory,
    label_field="instances",
)

谢谢

Answer 1

最简单的方法是使用 FiftyOne 在简单的 Python 循环中迭代您的数据集，使用 OpenCV 和 Numpy 格式化对象实例的图像并将其写入磁盘。

例如，此函数将获取 FiftyOne 样本的任何集合（Dataset for View）并将所有对象实例写入由 class 标签分隔的文件夹中的磁盘：

import os

import cv2
import numpy as np

def extract_classwise_instances(samples, output_dir, label_field, ext=".png"):
    print("Extracting object instances...")
    for sample in samples.iter_samples(progress=True):
        img = cv2.imread(sample.filepath)
        img_h,img_w,c = img.shape
        for det in sample[label_field].detections:
            mask = det.mask
            [x,y,w,h] = det.bounding_box
            x = int(x * img_w)
            y = int(y * img_h)
            h, w = mask.shape
            mask_img = img[y:y+h, x:x+w, :] 
            alpha = mask.astype(np.uint8)*255
            alpha = np.expand_dims(alpha, 2)
            mask_img = np.concatenate((mask_img, alpha), axis=2)
    
            label = det.label
            label_dir = os.path.join(output_dir, label)
            if not os.path.exists(label_dir):
                os.mkdir(label_dir)
            output_filepath = os.path.join(label_dir, det.id+ext)
            cv2.imwrite(output_filepath, mask_img)

这是一个加载 COCO2017 数据集子集并将所有“人”实例写入磁盘的完整示例：

import fiftyone as fo
import fiftyone.zoo as foz
from fiftyone import ViewField as F

dataset_name = "coco-image-example"
if dataset_name in fo.list_datasets():
    fo.delete_dataset(dataset_name)

label_field = "ground_truth"
classes = ["person"]

dataset = foz.load_zoo_dataset(
    "coco-2017",
    split="validation",
    label_types=["segmentations"],
    classes=classes,
    max_samples=20,
    label_field=label_field,
    dataset_name=dataset_name,
)

view = dataset.filter_labels(label_field, F("label").is_in(classes))

output_dir = "/path/to/output/segmentations/dir/"
os.makedirs(output_dir, exist_ok=True)

extract_classwise_instances(view, output_dir, label_field)

如果这个功能是经常使用的东西，写一个custom dataset exporter这个格式可能会有用。

如何从 COCO 数据集或 Open Images V6 数据集中提取前景对象？

How to extract foreground objects from COCO dataset or Open Images V6 Dataset?

python

computer-vision

dataset

object-detection

fiftyone