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

Question

Currently, I am preparing a synthetic dataset for object detection task. There are annotated datasets available for this kind of tasks like COCO dataset and Open Images V6. I am trying to download the images from there but only the foreground objects for a specific class e.g. person, in other words images without transparent background. The reason I am doing this is that I want to insert those images after editing them into a new images e.g. a street scene.

What I have tried so far, I used a library called FiftyOne and I downloaded the dataset with their semantic label and I am stuck here and I don`t what else to do.

It is not necessary to use FiftyOne any other method would work.

Here is the code that I have used to download a sample of the dataset with their labels

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",
)

Thank you

Answer 1

The easiest way to do this is by using FiftyOne to iterate over your dataset in a simple Python loop, using OpenCV and Numpy to format and write the images of object instances to disk.

For example, this function will take in any collection of FiftyOne samples (either a Dataset for View) and write all object instances to disk in folders separated by class label:

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)

Here is a complete example that loads a subset of the COCO2017 dataset and writes all "person" instances to disk:

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)

If this capability is something that will be used regularly, it may be useful to write a custom dataset exporter for this format.