Licensed under the Apache License, Version 2.0 (the "License")


In [0]:
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

In [0]:
# These datasets were used in:
# "Are GANs Created Equal? A Large-Scale Study"
# Mario Lucic*, Karol Kurach*, Marcin Michalski, Sylvain Gelly, Olivier Bousquet
# NeurIPS 2018

# The data set in the paper contains 80000 instances. First 60K are training, next 10K are testing,
# and last 10K additional hold-out. You can use this code to create and visualize the datasets.

from matplotlib import path
import numpy as np
from scipy.ndimage import filters

def GenerateConvexPolygon(n_vertices, min_segment_angle, scale, raster_dim,
                          subpixel_res, shift_to_mean=False):
  """Returns a random convex polygon with constraints on the minimum angle size.

  An ellipse is split into 'n_vertices' segments and a vertex is picked in
  each segment. Vertices are then selected within each segment such that the
  minimum segment angle (vertex A, zero, vertex B) is respected.

  Args:
    n_vertices: Integer > 2, number of vertices of the polygon.
    min_segment_angle: Float, minimum segment angle (in degrees).
    scale: Float in (0, 1], downscaling factor for the generated polygon.
    raster_dim: Integer > 1, size of the square grid of pixels to return.
    subpixel_res: Integer, each cell of the raster is split into subpixels.
    shift_to_mean: Boolean, whether to move the center of mass to the center
                   of the raster.
  """
  segment_angle = 360.0 / n_vertices
  max_angle = segment_angle - min_segment_angle/2.0
  grid_length = raster_dim * scale
  angles = []
  for i in range(n_vertices):
    offset_angle = np.random.rand() * max_angle
    angles.append(i * segment_angle + min_segment_angle/2.0 + offset_angle)
  angles_in_radians = np.array(angles) * np.pi / 180.0
  # Corresponding vertices are mapped to [0, 1]^2.
  x = np.column_stack((np.cos(angles_in_radians), np.sin(angles_in_radians)))
  x = (x + 1.0) / 2.0

  # Randomly rotate the vertices.
  theta = np.radians(np.random.rand() * 2 * np.pi)
  cos, sin = np.cos(theta), np.sin(theta)
  rot = np.matrix([[cos, -sin], [sin, cos]])
  x = np.dot(x, rot)

  # Corresponding vertices are mapped to [0, raster_dim]^2
  vertices = x * grid_length

  # If required, move the centre of mass to [raster_dim/2, raster_dim/2]^2.
  if shift_to_mean:
    vertices += raster_dim / 2.0 - vertices.mean(axis=0)

  # Create the bitmask by dividing each cell into subpixels.
  r, dim = subpixel_res, raster_dim
  poly_as_path = path.Path(vertices * r)
  grid_x, grid_y = np.mgrid[0:dim * r, 0:dim * r]
  flattened_grid = np.column_stack((grid_x.ravel(), grid_y.ravel()))

  mask = poly_as_path.contains_points(flattened_grid).reshape(dim * r, dim * r)
  mask = np.array(~mask, dtype=np.float32)
  mask = filters.convolve(mask, np.ones((r, r)), mode="constant", cval=1.0)
  return mask[::r, ::r] / (r * r)


def GenerateDataset(n_instances,
                    n_vertices,
                    min_segment_angle,
                    scale,
                    raster_dim,
                    subpixel_res,
                    shift_to_mean,
                    seed=0):
  """Returns a data set of random convex polygons with angle constraints.

  Args:
    n_instances: Integer, number of polygons to generate.
    n_vertices: Integer > 2, number of vertices of the polygon.
    min_segment_angle: Float, minimum angle between two vertices (in degrees).
    scale: Float in (0, 1], downscaling factor for the generated polygon.
    raster_dim: Integer > 1, size of the square grid of pixels to return.
    subpixel_res: Each pixel is split into a grid of size [subpixel_res,
                  subpixel_res] and each subpixel is checked when computing
                  whether the pixel belongs to the polygon.
    shift_to_mean: Boolean, whether to move the center of mass to the center
                   of the raster.
    seed: Random seed to ensure reproducibility.
  Raises:
    ValueError: If the number of vertices is too small, the min angle is too
                large, or scale is not within (0, 1].
  """
  if n_vertices < 3:
    raise ValueError("Need more than 2 vertices.")
  if min_segment_angle > 360.0/n_vertices:
    raise ValueError("The minimum segment angle is infeasible.")
  if scale < 0 or scale > 1:
    raise ValueError("Scale must be within (0, 1]")
  if raster_dim <= 1:
    raise ValueError("Raster sidelength has to be greater than 1.")
  np.random.seed(seed)
  x = np.zeros((n_instances, raster_dim, raster_dim))
  y = np.ones(n_instances, dtype=np.int8) * n_vertices
  for i in range(n_instances):
    x[i] = GenerateConvexPolygon(
        n_vertices=n_vertices,
        min_segment_angle=min_segment_angle,
        scale=scale,
        raster_dim=raster_dim,
        subpixel_res=subpixel_res,
        shift_to_mean=shift_to_mean)
  ids = np.random.permutation(x.shape[0])
  return x[ids], y[ids]

In [0]:
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid
 
def visualize(dataset, n_rows):
  fig = plt.figure(1, (n_rows, n_rows))
  n = dataset.shape[0]
  n_cols = n//n_rows + (n % n_rows > 0)
  print(n_rows, n_cols)
  grid = ImageGrid(fig, 111, nrows_ncols=(n_rows, n_cols), axes_pad=(0, 0))
  for i in range(n):
    grid[i].set_xlim([0, dataset[0].shape[0]])
    grid[i].set_ylim([0, dataset[0].shape[0]])
    grid[i].get_xaxis().set_ticks([])
    grid[i].get_yaxis().set_ticks([])
  for row in range(n_rows):
    for col in range(n_cols):
      target = dataset[row * n_cols + col]
      grid[row * n_cols + col].imshow(target, cmap='gray')

In [0]:
triangles, labels = GenerateDataset(
    n_instances=100, 
    n_vertices=3, 
    min_segment_angle=20, 
    scale=0.75, 
    raster_dim=28, 
    subpixel_res=8, 
    shift_to_mean=True, 
    seed=0) 
visualize(triangles, n_rows=10)