Download background images for YOLO model training

bg_images_yolo.py

# Script to download background image from COCO in YOLO format
# Adapted from https://stackoverflow.com/a/62770484/8061030


from pycocotools.coco import COCO
import requests
import os

DETECTOR_CLASSES = [] # Specify the COCO classes that you are detecting
NUM_IMAGES = # Number of background images to download

# instantiate COCO specifying the annotations json path
coco = COCO('annotations/instances_train2017.json')

# Get list of images to exclude based on DETECTOR_CLASSES
# Background images will not contain these.
# These should be classes included in training.
exc_cat_ids = coco.getCatIds(catNms=DETECTOR_CLASSES)

# Get the corresponding image ids to exclude
exc_img_ids = []
for cat_id in exc_cat_ids:
    exc_img_ids += coco.getImgIds(catIds=cat_id)

# Remove duplicates
exc_img_ids = set(exc_img_ids)

# Get all image ids
all_img_ids = coco.getImgIds()

# Remove img ids of classes that are included in training
bg_img_ids = set(all_img_ids) - set(exc_img_ids)

# Get background image metadata
bg_images = coco.loadImgs(bg_img_ids)

# Create dirs
os.makedirs("images", exist_ok=True)
os.makedirs("labels", exist_ok=True)

# Save the images into a local folder
for im in bg_images[:NUM_IMAGES]:
    img_data = requests.get(im["coco_url"]).content

    # Save the image
    with open("images/" + im["file_name"], "wb") as handler:
        handler.write(img_data)

    # Save the corresponding blank label txt file
    with open("labels/" + im["file_name"][:-3] + "txt", "wb") as handler:
        pass