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

Question

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

到目前为止，我已经尝试过了，我使用了一个名为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示例集合(View的Dataset )，并将所有对象实例写入由类标签分隔的文件夹中的磁盘：

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数据集的一个子集，并将所有"person“实例写入磁盘：

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可能会很有用。