import gzip
import math
import numpy as np
import os
from PIL import Image
import random
import torch
import torch.utils.data as data
def load_mnist(root):
# Load MNIST dataset for generating training data.
path = os.path.join(root, 'train-images-idx3-ubyte.gz')
with gzip.open(path, 'rb') as f:
mnist = np.frombuffer(f.read(), np.uint8, offset=16)
mnist = mnist.reshape(-1, 28, 28)
return mnist
def load_fixed_set(root, is_train):
# Load the fixed dataset
filename = 'mnist_test_seq.npy'
path = os.path.join(root, filename)
dataset = np.load(path)
dataset = dataset[..., np.newaxis]
return dataset
class MovingMNIST(data.Dataset):
def __init__(self, root, is_train, n_frames_input, n_frames_output, num_objects,
transform=None):
'''
param num_objects: a list of number of possible objects.
'''
super(MovingMNIST, self).__init__()
self.dataset = None
if is_train:
self.mnist = load_mnist(root)
else:
if num_objects[0] != 2:
self.mnist = load_mnist(root)
else:
self.dataset = load_fixed_set(root, False)
self.length = int(1e4) if self.dataset is None else self.dataset.shape[1]
self.is_train = is_train
self.num_objects = num_objects
self.n_frames_input = n_frames_input
self.n_frames_output = n_frames_output
self.n_frames_total = self.n_frames_input + self.n_frames_output
self.transform = transform
# For generating data
self.image_size_ = 64
self.digit_size_ = 28
self.step_length_ = 0.1
def get_random_trajectory(self, seq_length):
''' Generate a random sequence of a MNIST digit '''
canvas_size = self.image_size_ - self.digit_size_
x = random.random()
y = random.random()
theta = random.random() * 2 * np.pi
v_y = np.sin(theta)
v_x = np.cos(theta)
start_y = np.zeros(seq_length)
start_x = np.zeros(seq_length)
for i in range(seq_length):
# Take a step along velocity.
y += v_y * self.step_length_
x += v_x * self.step_length_
# Bounce off edges.
if x <= 0:
x = 0
v_x = -v_x
if x >= 1.0:
x = 1.0
v_x = -v_x
if y <= 0:
y = 0
v_y = -v_y
if y >= 1.0:
y = 1.0
v_y = -v_y
start_y[i] = y
start_x[i] = x
# Scale to the size of the canvas.
start_y = (canvas_size * start_y).astype(np.int32)
start_x = (canvas_size * start_x).astype(np.int32)
return start_y, start_x
def generate_moving_mnist(self, num_digits=2):
'''
Get random trajectories for the digits and generate a video.
'''
data = np.zeros((self.n_frames_total, self.image_size_, self.image_size_), dtype=np.float32)
for n in range(num_digits):
# Trajectory
start_y, start_x = self.get_random_trajectory(self.n_frames_total)
ind = random.randint(0, self.mnist.shape[0] - 1)
digit_image = self.mnist[ind]
for i in range(self.n_frames_total):
top = start_y[i]
left = start_x[i]
bottom = top + self.digit_size_
right = left + self.digit_size_
# Draw digit
data[i, top:bottom, left:right] = np.maximum(data[i, top:bottom, left:right], digit_image)
data = data[..., np.newaxis]
return data
def __getitem__(self, idx):
length = self.n_frames_input + self.n_frames_output
if self.is_train or self.num_objects[0] != 2:
# Sample number of objects
num_digits = random.choice(self.num_objects)
# Generate data on the fly
images = self.generate_moving_mnist(num_digits)
else:
images = self.dataset[:, idx, ...]
if self.transform is not None:
images = self.transform(images)
input = images[:self.n_frames_input]
if self.n_frames_output > 0:
output = images[self.n_frames_input:length]
else:
output = []
return input, output
def __len__(self):
return self.length
