| import tensorflow as tf
|
| import numpy as np
|
|
|
| def random_crop(img_list, crop_h, crop_w):
|
| img_size = tf.shape(img_list[0])
|
|
|
| rand_offset_h = tf.random_uniform([], 0, img_size[0]-crop_h+1, dtype=tf.int32)
|
| rand_offset_w = tf.random_uniform([], 0, img_size[1]-crop_w+1, dtype=tf.int32)
|
|
|
| for i, img in enumerate(img_list):
|
| img_list[i] = tf.image.crop_to_bounding_box(img, rand_offset_h, rand_offset_w, crop_h, crop_w)
|
|
|
| return img_list
|
|
|
| def flow_vertical_flip(flow):
|
| flow = tf.image.flip_up_down(flow)
|
| flow_u, flow_v = tf.unstack(flow, axis=-1)
|
| flow_v = flow_v * -1
|
| flow = tf.stack([flow_u, flow_v], axis=-1)
|
| return flow
|
|
|
| def flow_horizontal_flip(flow):
|
| flow = tf.image.flip_left_right(flow)
|
| flow_u, flow_v = tf.unstack(flow, axis=-1)
|
| flow_u = flow_u * -1
|
| flow = tf.stack([flow_u, flow_v], axis=-1)
|
| return flow
|
|
|
| def random_flip(img_list):
|
| is_flip = tf.random_uniform([2], minval=0, maxval=2, dtype=tf.int32)
|
|
|
| for i in range(len(img_list)):
|
| img_list[i] = tf.where(is_flip[0] > 0, tf.image.flip_left_right(img_list[i]), img_list[i])
|
| img_list[i] = tf.where(is_flip[1] > 0, tf.image.flip_up_down(img_list[i]), img_list[i])
|
| return img_list
|
|
|
| def random_flip_with_flow(img_list, flow_list):
|
| is_flip = tf.random_uniform([2], minval=0, maxval=2, dtype=tf.int32)
|
| for i in range(len(img_list)):
|
| img_list[i] = tf.where(is_flip[0] > 0, tf.image.flip_left_right(img_list[i]), img_list[i])
|
| img_list[i] = tf.where(is_flip[1] > 0, tf.image.flip_up_down(img_list[i]), img_list[i])
|
| for i in range(len(flow_list)):
|
| flow_list[i] = tf.where(is_flip[0] > 0, flow_horizontal_flip(flow_list[i]), flow_list[i])
|
| flow_list[i] = tf.where(is_flip[1] > 0, flow_vertical_flip(flow_list[i]), flow_list[i])
|
| return img_list, flow_list
|
|
|
|
|
| def random_channel_swap(img_list):
|
| channel_permutation = tf.constant([[0, 1, 2],
|
| [0, 2, 1],
|
| [1, 0, 2],
|
| [1, 2, 0],
|
| [2, 0, 1],
|
| [2, 1, 0]])
|
| rand_i = tf.random_uniform([], minval=0, maxval=6, dtype=tf.int32)
|
| perm = channel_permutation[rand_i]
|
| for i, img in enumerate(img_list):
|
| channel_1 = img[:, :, perm[0]]
|
| channel_2 = img[:, :, perm[1]]
|
| channel_3 = img[:, :, perm[2]]
|
| img_list[i] = tf.stack([channel_1, channel_2, channel_3], axis=-1)
|
| return img_list
|
|
|
| def flow_resize(flow, out_size, is_scale=True, method=0):
|
| '''
|
| method: 0 mean bilinear, 1 means nearest
|
| '''
|
| flow_size = tf.to_float(tf.shape(flow)[-3:-1])
|
| flow = tf.image.resize_images(flow, out_size, method=method, align_corners=True)
|
| if is_scale:
|
| scale = tf.to_float(out_size) / flow_size
|
| scale = tf.stack([scale[1], scale[0]])
|
| flow = tf.multiply(flow, scale)
|
| return flow
|
|
|
|
|
