def get_batch(BS, img1_list, img2_list, flow_list, HEIGHT, WIDTH, FLO_MAX):
N = len(img2_list)
I1 = np.zeros((BS, HEIGHT, WIDTH, 3), dtype=np.float32)
I2 = np.zeros((BS, HEIGHT, WIDTH, 3), dtype=np.float32)
F = np.zeros((BS, HEIGHT, WIDTH, 2), dtype=np.float32)
E = np.zeros((BS, HEIGHT, WIDTH, 1), dtype=np.float32)
good_ids = []
for i in range(BS):
try:
idx = randint(0, N - 1)
img1 = io.imread(img1_list[idx]).astype(np.float32) / 255
img2 = io.imread(img2_list[idx]).astype(np.float32) / 255
edge = io.imread(img2_list[idx]).astype(np.float32) / 255
flow = (readFlow(flow_list[idx]).astype(np.float32))
I1[i][:][:][:] = img1
I2[i][:][:][:] = img2
F[i][:][:][:] = flow
E[i][:][:][0] = edge
good_ids.append(True)
except:
print('Some images were not read')
good_ids.append(False)
I1 = I1[good_ids, :, :, :]
I2 = I2[good_ids, :, :, :]
F = F[good_ids, :, :, :]
E = E[good_ids, :, :, :]
return I1, I2, F, E
def get_batch_jitter(batch_size, img1_list, img2_list, edge_list, flow_list,
height, width):
n_files = len(img1_list)
imgs1 = np.zeros((batch_size, height, width, 3), dtype=np.float32)
imgs2 = np.zeros((batch_size, height, width, 3), dtype=np.float32)
flows = np.zeros((batch_size, height, width, 2), dtype=np.float32)
edges = np.zeros((batch_size, height, width, 1), dtype=np.float32)
good_ids = []
for i in range(batch_size):
try:
idx = randint(0, n_files - 1)
img1 = io.imread(img1_list[idx]).astype(np.float32) / 255
img2 = io.imread(img2_list[idx]).astype(np.float32) / 255
flow = (readFlow(flow_list[idx]).astype(np.float32))
edge = np.expand_dims(
io.imread(edge_list[idx])[:, :, 1].astype(np.float32) / 255,
axis=2)
imgs1[i][:][:][:] = img1
imgs2[i][:][:][:] = img2
flows[i][:][:][:] = flow
edges[i][:][:][:] = edge
good_ids.append(True)
except:
good_ids.append(False)
imgs1 = imgs1[good_ids][:][:][:]
imgs2 = imgs2[good_ids][:][:][:]
flows = flows[good_ids][:][:][:]
edges = edges[good_ids][:][:][:]
# print(good_ids)
# apply horizontal flipping
if toss():
imgs1 = np.flip(imgs1, axis=2)
imgs2 = np.flip(imgs2, axis=2)
edges = np.flip(edges, axis=2)
flows = np.flip(flows, axis=2)
flows[:, :, :, 0] = -flows[:, :, :, 0]
# apply vertical flipping
if toss():
imgs1 = np.flip(imgs1, axis=1)
imgs2 = np.flip(imgs2, axis=1)
edges = np.flip(edges, axis=1)
flows = np.flip(flows, axis=1)
flows[:, :, :, 1] = -flows[:, :, :, 1]
return imgs1, imgs2, edges, flows
def get_batch(batch_size, img1_list, img2_list, edge_list, miss_list,
flow_list, height, width):
n_files = len(img2_list)
imgs1 = np.zeros((batch_size, height, width, 3), dtype=np.float32)
imgs2 = np.zeros((batch_size, height, width, 3), dtype=np.float32)
flows = np.zeros((batch_size, height, width, 2), dtype=np.float32)
edges = np.zeros((batch_size, height, width, 1), dtype=np.float32)
misses = np.zeros((batch_size, height, width, 1), dtype=np.float32)
good_ids = []
for i in range(batch_size):
try:
idx = randint(0, n_files - 1)
img1 = io.imread(img1_list[idx]).astype(np.float32) / 255
img2 = io.imread(img2_list[idx]).astype(np.float32) / 255
flow = (readFlow(flow_list[idx]).astype(np.float32))
edge = np.expand_dims(
io.imread(edge_list[idx])[:, :, 1].astype(np.float32) / 255,
axis=2)
miss = np.expand_dims(
io.imread(miss_list[idx])[:, :, 1].astype(np.float32) / 255,
axis=2)
imgs1[i][:][:][:] = img1
imgs2[i][:][:][:] = img2
flows[i][:][:][:] = flow
misses[i][:][:][:] = miss
edges[i][:][:][:] = edge
good_ids.append(True)
except:
good_ids.append(False)
imgs1 = imgs1[good_ids][:][:][:]
imgs2 = imgs2[good_ids][:][:][:]
flows = flows[good_ids][:][:][:]
edges = edges[good_ids][:][:][:]
misses = misses[good_ids][:][:][:]
# print(good_ids)
return imgs1, imgs2, edges, misses, flows
def get_batch_jitter_scale(batch_size,
img1_list,
img2_list,
edge_list,
miss_list,
flow_list,
height,
width,
scale=1):
start_time = time.time()
n_files = len(img1_list)
imgs1 = np.zeros((batch_size, height, width, 3), dtype=np.float32)
imgs2 = np.zeros((batch_size, height, width, 3), dtype=np.float32)
flows = np.zeros((batch_size, height, width, 2), dtype=np.float32)
edges = np.zeros((batch_size, height, width, 1), dtype=np.float32)
misses = np.zeros((batch_size, height, width, 1), dtype=np.float32)
temp_f = np.zeros((height, width, 2), dtype=np.float32)
good_ids = []
for i in range(batch_size):
try:
idx = randint(0, n_files - 1)
img1 = io.imread(img1_list[idx]).astype(np.float32) / 255
img2 = io.imread(img2_list[idx]).astype(np.float32) / 255
flow = readFlow(flow_list[idx]).astype(np.float32)
edge = io.imread(edge_list[idx]).astype(np.float32) / 255
miss = io.imread(miss_list[idx]).astype(np.float32) / 255
jump = int(1 / scale)
if scale < 1:
img1 = img1[::jump, ::jump, :]
img2 = img2[::jump, ::jump, :]
edge = edge[::jump, ::jump]
miss = miss[::jump, ::jump]
temp_f = flow[::jump, ::jump, :] * scale
else:
temp_f = flow
edge = np.expand_dims(edge, axis=2)
miss = np.expand_dims(miss, axis=2)
imgs1[i][:][:][:] = img1
imgs2[i][:][:][:] = img2
flows[i][:][:][:] = temp_f
edges[i][:][:][:] = edge
misses[i][:][:][:] = miss
good_ids.append(True)
except:
good_ids.append(False)
imgs1 = imgs1[good_ids][:][:][:]
imgs2 = imgs2[good_ids][:][:][:]
flows = flows[good_ids][:][:][:]
edges = edges[good_ids][:][:][:]
misses = misses[good_ids][:][:][:]
# print(good_ids)
# apply horizontal flipping
if toss():
imgs1 = np.flip(imgs1, axis=2)
imgs2 = np.flip(imgs2, axis=2)
edges = np.flip(edges, axis=2)
misses = np.flip(misses, axis=2)
flows = np.flip(flows, axis=2)
flows[:, :, :, 0] = -flows[:, :, :, 0]
# apply vertical flipping
if toss():
imgs1 = np.flip(imgs1, axis=1)
imgs2 = np.flip(imgs2, axis=1)
edges = np.flip(edges, axis=1)
misses = np.flip(misses, axis=1)
flows = np.flip(flows, axis=1)
flows[:, :, :, 1] = -flows[:, :, :, 1]
elapsed_time = time.time() - start_time
# print('time elapse in reading batch is', elapsed_time)
return imgs1, imgs2, edges, misses, flows
learning_rate = tf.placeholder(dtype=tf.float32, shape=[])
model_path = './models/fn_sup'
net = FlowNetS(HEIGHT, WIDTH)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(net.loss)
init = tf.global_variables_initializer()
sess = tf.InteractiveSession()
saver = tf.train.Saver()
# sess.run(init)
# saver = tf.train.import_meta_graph(model_path + '.meta')
saver.restore(sess, model_path)
test_dir = './middlebury_data/Grove3'
img1 = read_img(test_dir + '/' + 'frame10.png')[0:HEIGHT, 0:WIDTH, :]
img2 = read_img(test_dir + '/' + 'frame11.png')[0:HEIGHT, 0:WIDTH, :]
flow_gt = readFlow(test_dir + '/' + 'flow10.flo')[0:HEIGHT, 0:WIDTH, :]
print('gt flow shape', flow_gt.shape)
flow_eval = np.squeeze(net.flow2.eval(feed_dict={net.inp1: np.expand_dims(img1, 0),
net.inp2: np.expand_dims(img2, 0)}))
print('eval flow shape', flow_eval.shape)
img_vis = np.hstack((img1, img2))
flow_vis = np.hstack((imresize(flow2hsv(flow_gt), 1./4), flow2hsv(-flow_eval)))
plt.imshow(img_vis)
plt.pause(2)
plt.imshow(flow_vis)
plt.pause(2)