File size: 3,087 Bytes
872b0a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import tensorflow as tf
import numpy as np
def random_crop(img_list, crop_h, crop_w):
img_size = tf.shape(img_list[0])
# crop image and flow
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
|