Generate Toy Images for Object Detection

toy.py

#!/usr/bin/env python3

import argparse
import math
import numpy as np
import scipy.misc
import time


def generate_rectangle(x_0, y_0, width, height, color, max_dimension):
    # (x_0, y_0) is the top left corner
    w = min(np.random.randint(width - x_0), max_dimension)
    h = min(np.random.randint(height - y_0), max_dimension)
    mask = np.zeros((width, height, 3))
    mask[y_0:y_0+h, x_0:x_0+w] = color

    return mask


def generate_circle(x_0, y_0, width, height, color, max_dimension):
    # (x_0, y_0) is the center
    left = x_0
    right = width - x_0
    top = y_0
    bottom = height - y_0
    available_radius = min(left, right, top, bottom, max_dimension)
    if available_radius == 0:
        return None
    radius = np.random.randint(available_radius)
    mask = np.zeros((width, height, 3))
    for x in range(x_0 - radius, x_0 + radius + 1):
         y = int(np.sqrt(radius**2 - (x - x_0)**2))
         mask[y_0-y:y_0+y+1, x] = color

    return mask


def generate_triangle(x_0, y_0, width, height, color, max_dimension):
    # (x_0, y_0) is the bottom left corner
    # We're making an equilateral triangle. Pick the side as the min of the
    # available space to the right and available space above.
    side = min(np.random.randint(min(width - x_0, y_0 + 1)), max_dimension)
    slope = np.sqrt(3) # this pops up after some math
    y = y_0
    mid_point = x_0 + side / 2
    mask = np.zeros((width, height, 3))
    # Drawing a linear function with positive slope to the right up to the
    # mid-point, then drawing a linear function with negative slope.
    for x in range(x_0, x_0 + side):
        mask[int(y):y_0, x] = color
        if x < mid_point:
            # Subtracting y means going up ((0, 0) is top left)
            y -= slope
        else:
            y += slope

    return mask


def pick_color():
    return np.random.randint(255, size=3)


def generate_image(width, height, min_objects, max_objects, max_dimension):
    shapes = [generate_rectangle, generate_circle, generate_triangle]
    image = np.zeros((width, height, 3))
    for _ in range(min_objects, max_objects + 1):
        x = np.random.randint(width)
        y = np.random.randint(width)
        color = pick_color()
        shape = np.random.choice(shapes)
        mask = shape(x, y, width, height, color, max_dimension)
        if mask is not None:
            image += mask
    # Turn black (0) in to white (255), 0 made the math/logic easier
    image[image==0] = 255

    return image


def generate_dataset(number,
                     width,
                     height,
                     min_objects,
                     max_objects,
                     max_dimension=None):
    images = []
    for _ in range(number):
        image = generate_image(width,
                               height,
                               min_objects,
                               max_objects,
                               max_dimension)
        images.append(image)
    return images


def parse():
    parser = argparse.ArgumentParser(description='Generate R-CNN Toy Dataset')
    parser.add_argument(
        '-n',
        '--number',
        type=int,
        default=42,
        help='The number of images to generate (42)')
    parser.add_argument(
        '--width',
        type=int,
        default=512,
        help='The width of generated images (128)')
    parser.add_argument(
        '--height',
        type=int,
        default=512,
        help='The height of generated images (128)')
    parser.add_argument(
        '--max-objects',
        type=int,
        default=100,
        help='The maximum number of objects per image (100)')
    parser.add_argument(
        '--min-objects',
        type=int,
        default=1,
        help='The maximum number of objects per image (1)')
    parser.add_argument(
        '--max-dimension',
        type=int,
        default=math.inf,
        help='The maximum dimension of an object (None)')

    return parser.parse_args()


def main():
    args = parse()
    start = time.time()
    images = generate_dataset(args.number,
                              args.width,
                              args.height,
                              args.min_objects,
                              args.max_objects,
                              args.max_dimension)
    end = time.time() - start
    print('Generated {0} images in {1:.2f}s'.format(len(images), end))
    for image in images:
        scipy.misc.toimage(image, mode='RGB').show()


if __name__ == '__main__':
    main()