def dump_primitive_data(self, primitive, tar_path, config):
temp_dir = Path(os.environ['TMPDIR'], primitive)
tf.logging.info('Generating tarfile for primitive {}.'.format(primitive))
synthetic_dataset.set_random_state(np.random.RandomState(
config['generation']['random_seed']))
for split, size in self.config['generation']['split_sizes'].items():
im_dir, pts_dir = [Path(temp_dir, i, split) for i in ['images', 'points']]
im_dir.mkdir(parents=True, exist_ok=True)
pts_dir.mkdir(parents=True, exist_ok=True)
for i in tqdm(range(size), desc=split, leave=False):
image = synthetic_dataset.generate_background(
config['generation']['image_size'],
**config['generation']['params']['generate_background'])
points = np.array(getattr(synthetic_dataset, primitive)(
image, **config['generation']['params'].get(primitive, {})))
points = np.flip(points, 1) # reverse convention with opencv
b = config['preprocessing']['blur_size']
image = cv2.GaussianBlur(image, (b, b), 0)
points = (points * np.array(config['preprocessing']['resize'], np.float)
/ np.array(config['generation']['image_size'], np.float))
image = cv2.resize(image, tuple(config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
cv2.imwrite(str(Path(im_dir, '{}.png'.format(i))), image)
np.save(Path(pts_dir, '{}.npy'.format(i)), points)
# Pack into a tar file
tar = tarfile.open(tar_path, mode='w:gz')
tar.add(temp_dir, arcname=primitive)
tar.close()
shutil.rmtree(temp_dir)
tf.logging.info('Tarfile dumped to {}.'.format(tar_path))
def _gen_shape():
primitives = parse_primitives(config['primitives'], self.drawing_primitives)
while True:
primitive = np.random.choice(primitives)
image = synthetic_dataset.generate_background(
config['generation']['image_size'],
**config['generation']['params']['generate_background'])
points = np.array(getattr(synthetic_dataset, primitive)(
image, **config['generation']['params'].get(primitive, {})))
yield (np.expand_dims(image, axis=-1).astype(np.float32),
np.flip(points.astype(np.float32), 1))
