|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Reads data that is produced by dataset/gen_data.py."""
|
|
|
| from __future__ import absolute_import
|
| from __future__ import division
|
| from __future__ import print_function
|
|
|
| import os
|
| import random
|
| from absl import logging
|
| import tensorflow as tf
|
| import util
|
|
|
| gfile = tf.gfile
|
|
|
| QUEUE_SIZE = 2000
|
| QUEUE_BUFFER = 3
|
|
|
|
|
| class DataReader(object):
|
| """Reads stored sequences which are produced by dataset/gen_data.py."""
|
|
|
| def __init__(self, data_dir, batch_size, img_height, img_width, seq_length,
|
| num_scales):
|
| self.data_dir = data_dir
|
| self.batch_size = batch_size
|
| self.img_height = img_height
|
| self.img_width = img_width
|
| self.seq_length = seq_length
|
| self.num_scales = num_scales
|
|
|
| def read_data(self):
|
| """Provides images and camera intrinsics."""
|
| with tf.name_scope('data_loading'):
|
| with tf.name_scope('enqueue_paths'):
|
| seed = random.randint(0, 2**31 - 1)
|
| self.file_lists = self.compile_file_list(self.data_dir, 'train')
|
| image_paths_queue = tf.train.string_input_producer(
|
| self.file_lists['image_file_list'], seed=seed, shuffle=True)
|
| cam_paths_queue = tf.train.string_input_producer(
|
| self.file_lists['cam_file_list'], seed=seed, shuffle=True)
|
| img_reader = tf.WholeFileReader()
|
| _, image_contents = img_reader.read(image_paths_queue)
|
| image_seq = tf.image.decode_jpeg(image_contents)
|
|
|
| with tf.name_scope('load_intrinsics'):
|
| cam_reader = tf.TextLineReader()
|
| _, raw_cam_contents = cam_reader.read(cam_paths_queue)
|
| rec_def = []
|
| for _ in range(9):
|
| rec_def.append([1.0])
|
| raw_cam_vec = tf.decode_csv(raw_cam_contents, record_defaults=rec_def)
|
| raw_cam_vec = tf.stack(raw_cam_vec)
|
| intrinsics = tf.reshape(raw_cam_vec, [3, 3])
|
|
|
| with tf.name_scope('convert_image'):
|
| image_seq = self.preprocess_image(image_seq)
|
|
|
| with tf.name_scope('image_augmentation'):
|
| image_seq = self.augment_image_colorspace(image_seq)
|
|
|
| image_stack = self.unpack_images(image_seq)
|
|
|
| with tf.name_scope('image_augmentation_scale_crop'):
|
| image_stack, intrinsics = self.augment_images_scale_crop(
|
| image_stack, intrinsics, self.img_height, self.img_width)
|
|
|
| with tf.name_scope('multi_scale_intrinsics'):
|
| intrinsic_mat = self.get_multi_scale_intrinsics(intrinsics,
|
| self.num_scales)
|
| intrinsic_mat.set_shape([self.num_scales, 3, 3])
|
| intrinsic_mat_inv = tf.matrix_inverse(intrinsic_mat)
|
| intrinsic_mat_inv.set_shape([self.num_scales, 3, 3])
|
|
|
| with tf.name_scope('batching'):
|
| image_stack, intrinsic_mat, intrinsic_mat_inv = (
|
| tf.train.shuffle_batch(
|
| [image_stack, intrinsic_mat, intrinsic_mat_inv],
|
| batch_size=self.batch_size,
|
| capacity=QUEUE_SIZE + QUEUE_BUFFER * self.batch_size,
|
| min_after_dequeue=QUEUE_SIZE))
|
| logging.info('image_stack: %s', util.info(image_stack))
|
| return image_stack, intrinsic_mat, intrinsic_mat_inv
|
|
|
| def unpack_images(self, image_seq):
|
| """[h, w * seq_length, 3] -> [h, w, 3 * seq_length]."""
|
| with tf.name_scope('unpack_images'):
|
| image_list = [
|
| image_seq[:, i * self.img_width:(i + 1) * self.img_width, :]
|
| for i in range(self.seq_length)
|
| ]
|
| image_stack = tf.concat(image_list, axis=2)
|
| image_stack.set_shape(
|
| [self.img_height, self.img_width, self.seq_length * 3])
|
| return image_stack
|
|
|
| @classmethod
|
| def preprocess_image(cls, image):
|
|
|
| return tf.image.convert_image_dtype(image, dtype=tf.float32)
|
|
|
|
|
| @classmethod
|
| def augment_image_colorspace(cls, image_seq):
|
| """Apply data augmentation to inputs."""
|
|
|
| random_gamma = tf.random_uniform([], 0.8, 1.2)
|
| image_seq_aug = image_seq**random_gamma
|
|
|
| random_brightness = tf.random_uniform([], 0.5, 2.0)
|
| image_seq_aug *= random_brightness
|
|
|
| random_colors = tf.random_uniform([3], 0.8, 1.2)
|
| white = tf.ones([tf.shape(image_seq)[0], tf.shape(image_seq)[1]])
|
| color_image = tf.stack([white * random_colors[i] for i in range(3)], axis=2)
|
| image_seq_aug *= color_image
|
|
|
| image_seq_aug = tf.clip_by_value(image_seq_aug, 0, 1)
|
| return image_seq_aug
|
|
|
| @classmethod
|
| def augment_images_scale_crop(cls, im, intrinsics, out_h, out_w):
|
| """Randomly scales and crops image."""
|
|
|
| def scale_randomly(im, intrinsics):
|
| """Scales image and adjust intrinsics accordingly."""
|
| in_h, in_w, _ = im.get_shape().as_list()
|
| scaling = tf.random_uniform([2], 1, 1.15)
|
| x_scaling = scaling[0]
|
| y_scaling = scaling[1]
|
| out_h = tf.cast(in_h * y_scaling, dtype=tf.int32)
|
| out_w = tf.cast(in_w * x_scaling, dtype=tf.int32)
|
|
|
| im = tf.expand_dims(im, 0)
|
| im = tf.image.resize_area(im, [out_h, out_w])
|
| im = im[0]
|
| fx = intrinsics[0, 0] * x_scaling
|
| fy = intrinsics[1, 1] * y_scaling
|
| cx = intrinsics[0, 2] * x_scaling
|
| cy = intrinsics[1, 2] * y_scaling
|
| intrinsics = cls.make_intrinsics_matrix(fx, fy, cx, cy)
|
| return im, intrinsics
|
|
|
|
|
| def crop_randomly(im, intrinsics, out_h, out_w):
|
| """Crops image and adjust intrinsics accordingly."""
|
|
|
| in_h, in_w, _ = tf.unstack(tf.shape(im))
|
| offset_y = tf.random_uniform([1], 0, in_h - out_h + 1, dtype=tf.int32)[0]
|
| offset_x = tf.random_uniform([1], 0, in_w - out_w + 1, dtype=tf.int32)[0]
|
| im = tf.image.crop_to_bounding_box(im, offset_y, offset_x, out_h, out_w)
|
| fx = intrinsics[0, 0]
|
| fy = intrinsics[1, 1]
|
| cx = intrinsics[0, 2] - tf.cast(offset_x, dtype=tf.float32)
|
| cy = intrinsics[1, 2] - tf.cast(offset_y, dtype=tf.float32)
|
| intrinsics = cls.make_intrinsics_matrix(fx, fy, cx, cy)
|
| return im, intrinsics
|
|
|
| im, intrinsics = scale_randomly(im, intrinsics)
|
| im, intrinsics = crop_randomly(im, intrinsics, out_h, out_w)
|
| return im, intrinsics
|
|
|
| def compile_file_list(self, data_dir, split, load_pose=False):
|
| """Creates a list of input files."""
|
| logging.info('data_dir: %s', data_dir)
|
| with gfile.Open(os.path.join(data_dir, '%s.txt' % split), 'r') as f:
|
| frames = f.readlines()
|
| subfolders = [x.split(' ')[0] for x in frames]
|
| frame_ids = [x.split(' ')[1][:-1] for x in frames]
|
| image_file_list = [
|
| os.path.join(data_dir, subfolders[i], frame_ids[i] + '.jpg')
|
| for i in range(len(frames))
|
| ]
|
| cam_file_list = [
|
| os.path.join(data_dir, subfolders[i], frame_ids[i] + '_cam.txt')
|
| for i in range(len(frames))
|
| ]
|
| file_lists = {}
|
| file_lists['image_file_list'] = image_file_list
|
| file_lists['cam_file_list'] = cam_file_list
|
| if load_pose:
|
| pose_file_list = [
|
| os.path.join(data_dir, subfolders[i], frame_ids[i] + '_pose.txt')
|
| for i in range(len(frames))
|
| ]
|
| file_lists['pose_file_list'] = pose_file_list
|
| self.steps_per_epoch = len(image_file_list) // self.batch_size
|
| return file_lists
|
|
|
| @classmethod
|
| def make_intrinsics_matrix(cls, fx, fy, cx, cy):
|
| r1 = tf.stack([fx, 0, cx])
|
| r2 = tf.stack([0, fy, cy])
|
| r3 = tf.constant([0., 0., 1.])
|
| intrinsics = tf.stack([r1, r2, r3])
|
| return intrinsics
|
|
|
| @classmethod
|
| def get_multi_scale_intrinsics(cls, intrinsics, num_scales):
|
| """Returns multiple intrinsic matrices for different scales."""
|
| intrinsics_multi_scale = []
|
|
|
| for s in range(num_scales):
|
| fx = intrinsics[0, 0] / (2**s)
|
| fy = intrinsics[1, 1] / (2**s)
|
| cx = intrinsics[0, 2] / (2**s)
|
| cy = intrinsics[1, 2] / (2**s)
|
| intrinsics_multi_scale.append(cls.make_intrinsics_matrix(fx, fy, cx, cy))
|
| intrinsics_multi_scale = tf.stack(intrinsics_multi_scale)
|
| return intrinsics_multi_scale
|
|
|
