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

Question

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

到目前为止，我使用了一个名为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样本的任何集合(视图的数据集(，并将所有对象实例写入磁盘中由类标签分隔的文件夹中：

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)

如果这种功能将被定期使用，那么为这种格式编写一个自定义数据集导出器可能会很有用。