text: from glob import glob
  import os
  import numpy as np
  import random
  import torch.utils.data as data
  import json
  import cv2
  
  def make_dataset(root, is_train):
    if is_train:
      folder = 'balls_n4_t60_ex50000'
    else:
      folder = 'balls_n4_t60_ex2000'
  
    dataset = np.load(os.path.join(root, folder, 'dataset_info.npy'))
    return dataset
  
  class BouncingBalls(data.Dataset):
    '''
    Bouncing balls dataset.
    '''
    def __init__(self, root, is_train, n_frames_input, n_frames_output, image_size,
                 transform=None, return_positions=False):
      super(BouncingBalls, self).__init__()
      self.n_frames = n_frames_input + n_frames_output
      self.dataset = make_dataset(root, is_train)
  
      self.size = image_size
      self.scale = self.size / 800
      self.radius = int(60 * self.scale)
  
      self.root = root
      self.is_train = is_train
      self.n_frames_input = n_frames_input
      self.n_frames_output = n_frames_output
      self.transform = transform
      self.return_positions = return_positions
  
    def __getitem__(self, idx):
      # traj sizeL (n_frames, n_balls, 4)
      traj = self.dataset[idx]
      vid_len, n_balls = traj.shape[:2]
      if self.is_train:
        start = random.randint(0, vid_len - self.n_frames)
      else:
        start = 0
  
      n_channels = 1
      images = np.zeros([self.n_frames, self.size, self.size, n_channels], np.uint8)
      positions = []
      for fid in range(self.n_frames):
        xy = []
        for bid in range(n_balls):
          # each ball:
          ball = traj[start + fid, bid]
          x, y = int(round(self.scale * ball[0])), int(round(self.scale * ball[1]))
          images[fid] = cv2.circle(images[fid], (x, y), int(self.radius * ball[3]),
                                   255, -1)
          xy.append([x / self.size, y / self.size])
        positions.append(xy)
  
      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:]
      else:
        output = []
  
      if not self.return_positions:
        return input, output
      else:
        positions = np.array(positions)
        return input, output, positions
  
    def __len__(self):
      return len(self.dataset)
  