research/vid2depth/reader.py
# Copyright 2017 The TensorFlow Authors All Rights Reserved.
#
# 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
#
# http://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.
# ==============================================================================
"""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) # Converts to float.
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):
# Convert from uint8 to float.
return tf.image.convert_image_dtype(image, dtype=tf.float32)
# Source: https://github.com/mrharicot/monodepth.
@classmethod
def augment_image_colorspace(cls, image_seq):
"""Apply data augmentation to inputs."""
# Randomly shift gamma.
random_gamma = tf.random_uniform([], 0.8, 1.2)
image_seq_aug = image_seq**random_gamma
# Randomly shift brightness.
random_brightness = tf.random_uniform([], 0.5, 2.0)
image_seq_aug *= random_brightness
# Randomly shift color.
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
# Saturate.
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)
# Add batch.
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
# Random cropping
def crop_randomly(im, intrinsics, out_h, out_w):
"""Crops image and adjust intrinsics accordingly."""
# batch_size, in_h, in_w, _ = im.get_shape().as_list()
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 = []
# Scale the intrinsics accordingly for each 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