def visualize_batch(batch):
if len(batch.shape) == 4:
if batch.shape[3] == 2:
batch = [flow_to_image(batch[i]) for i in range(batch.shape[0])]
cv2.imshow('Optical flow set', np.hstack(batch))
else:
batch = [batch[i] for i in range(batch.shape[0])]
cv2.imshow('Image sets', np.hstack(batch))
cv2.waitKey(0)
else:
if batch.shape[4] == 2:
batch = [
np.hstack([
flow_to_image(batch[j][i])
for i in range(batch[j].shape[0])
]) for j in range(batch.shape[0])
]
cv2.imshow('Optical flow set', np.vstack(batch))
else:
batch = [
np.hstack([batch[j][i] for i in range(batch[j].shape[0])])
for j in range(batch.shape[0])
]
cv2.imshow('Image sets', np.vstack(batch))
cv2.waitKey(0)
def flow_visualize(flow, mode='Y'):
if mode == 'Y':
# Ccbcr color wheel
img = fl.flow_to_image(flow)
plt.imshow(img)
plt.show()
elif mode == 'RGB':
(h, w) = flow.shape[0:2]
du = flow[:, :, 0]
dv = flow[:, :, 1]
valid = flow[:, :, 2]
max_flow = max(np.max(du), np.max(dv))
img = np.zeros((h, w, 3), dtype=np.float64)
# angle layer
img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi)
# magnitude layer, normalized to 1
img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow
# phase layer
img[:, :, 2] = 8 - img[:, :, 1]
# clip to [0,1]
small_idx = img[:, :, 0:3] < 0
large_idx = img[:, :, 0:3] > 1
img[small_idx] = 0
img[large_idx] = 1
# convert to rgb
img = cl.hsv_to_rgb(img)
# remove invalid point
img[:, :, 0] = img[:, :, 0] * valid
img[:, :, 1] = img[:, :, 1] * valid
img[:, :, 2] = img[:, :, 2] * valid
# show
plt.imshow(img)
plt.show()
return
def flow_visualize(flow, mode='Y'):
if mode == 'Y':
# Ccbcr color wheel
img = fl.flow_to_image(flow)
plt.imshow(img)
plt.show()
elif mode == 'RGB':
(h, w) = flow.shape[0:2]
du = flow[:, :, 0]
dv = flow[:, :, 1]
valid = flow[:, :, 2]
max_flow = max(np.max(du), np.max(dv))
img = np.zeros((h, w, 3), dtype=np.float64)
# angle layer
img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi)
# magnitude layer, normalized to 1
img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow
# phase layer
img[:, :, 2] = 8 - img[:, :, 1]
# clip to [0,1]
small_idx = img[:, :, 0:3] < 0
large_idx = img[:, :, 0:3] > 1
img[small_idx] = 0
img[large_idx] = 1
# convert to rgb
img = cl.hsv_to_rgb(img)
# remove invalid point
img[:, :, 0] = img[:, :, 0] * valid
img[:, :, 1] = img[:, :, 1] * valid
img[:, :, 2] = img[:, :, 2] * valid
# show
plt.imshow(img)
plt.show()
return
def visualize(flow_path):
flow = load_of(flow_path)
print(flow.shape)
flow_img = flowlib.flow_to_image(flow[0, :, :, :])
print(flow_img.shape)
cv2.imwrite('flow1.png', flow_img)
def run_test(img1,
img2,
test_out_path,
test_checkpoints,
save_image=False,
save_flo=True):
prediction = net_structure.net_structure(img1, img2)
pred_flow = prediction['flow']
saver = tf.train.Saver()
with tf.Session() as sess:
if test_checkpoints:
saver.restore(sess, test_checkpoints)
print 'restored checkpoints'
pred_flow = sess.run(pred_flow)[0, :, :, :]
unique_name = 'flow_predict_old' + str(uuid.uuid4())
if save_image:
flow_img = flow_to_image(pred_flow)
full_out_path = os.path.join(test_out_path, unique_name + '.png')
print unique_name
imsave(full_out_path, flow_img)
if save_flo:
full_out_path = os.path.join(test_out_path, unique_name + '.flo')
write_flow(pred_flow, full_out_path)
def inference(args, model):
model.eval()
if args.save_flow:
flow_folder = "{}".format(args.save)
if not os.path.exists(flow_folder):
os.makedirs(flow_folder)
input_image_list = glob(args.input_dir + '*.jpg')
input_image_list.sort()
print(args.input_dir, "len: ", len(input_image_list))
for i in range(0, len(input_image_list) - 1, 2):
print("img1: ", input_image_list[i])
print("img2: ", input_image_list[i + 1])
img1 = frame_utils.read_gen(input_image_list[i])
img2 = frame_utils.read_gen(input_image_list[i + 1])
# resize to 512
# inputs of the net are 256/512/1024...
img1_in = imresize(img1, (512, 512))
img2_in = imresize(img2, (512, 512))
images = [img1_in, img2_in]
images = np.array(images).transpose(3, 0, 1, 2)
images = torch.from_numpy(images.astype(np.float32))
images = torch.unsqueeze(images, 0)
images = [images]
if args.cuda:
data = [d.cuda() for d in images]
data = [Variable(d) for d in data]
with torch.no_grad():
output = model(data[0])
if args.save_flow:
_pflow = output[0].data.cpu().numpy().transpose(1, 2, 0)
frame_name = input_image_list[i].split('/')[-1]
flow_path = join(flow_folder, '{}.flo'.format(frame_name))
print("flow saved as: ", flow_path)
flow_utils.writeFlow(flow_path, _pflow)
if args.save_img:
# and saved as image
flow = flow_utils.readFlow(flow_path)
if not os.path.exists(flow_folder + '_img'):
os.makedirs(flow_folder + '_img')
img_path = join(flow_folder + '_img', frame_name)
print('img saved as: ', img_path)
img = flow_to_image(flow)
img = imresize(img, (img1.shape[0], img1.shape[1]))
imsave(img_path, img)
return
def model_out_to_flow_png(output):
out_np = output[0].data.cpu().numpy()
#https://gitorchub.com/DediGadot/PatchBatch/blob/master/flowlib.py
out_np = np.squeeze(out_np)
out_np = np.moveaxis(out_np, 0, -1)
im_arr = flowlib.flow_to_image(out_np)
im = Image.fromarray(im_arr)
im.save('test.png')
def test(checkpoint,
input_a_path,
input_b_path,
img_num,
out_path,
save_image=True,
save_flo=True):
input_a = cv2.imread(input_a_path)
input_b = cv2.imread(input_b_path)
# input_a = cv2.resize(input_a, (512, 384))
# input_b = cv2.resize(input_b, (512, 384))
# Convert from RGB -> BGR
input_a = input_a[..., [2, 1, 0]]
input_b = input_b[..., [2, 1, 0]]
# Scale from [0, 255] -> [0.0, 1.0] if needed
# if input_a.max() > 1.0:
# input_a = input_a / 255.0
# if input_b.max() > 1.0:
# input_b = input_b / 255.0
input_a = input_a / 255.0
input_b = input_b / 255.0
# TODO: This is a hack, we should get rid of this
# training_schedule = LONG_SCHEDULE
inputs = {
'input_a': tf.expand_dims(tf.constant(input_a, dtype=tf.float32), 0),
'input_b': tf.expand_dims(tf.constant(input_b, dtype=tf.float32), 0),
}
batch_img1 = create_batch([img1_path])
batch_img2 = create_batch([img2_path])
predictions = net_structure.net_structure(batch_img1, batch_img2)
pred_flow = predictions['flow']
init = tf.global_variables_initializer()
saver = tf.train.Saver()
sess = tf.Session()
sess.run(init)
saver.restore(sess, checkpoint)
pred_flow = sess.run(pred_flow)[0, :, :, :]
# unique_name = 'flow-' + str(uuid.uuid4())
unique_name = 'flow_predict_old_' + img_num
print unique_name
if save_image:
flow_img = flow_to_image(pred_flow)
full_out_path = os.path.join(out_path, unique_name + '.png')
imsave(full_out_path, flow_img)
if save_flo:
full_out_path = os.path.join(out_path, unique_name + '.flo')
write_flow(pred_flow, full_out_path)
def model_out_to_flow_png(output, name):
out_np = output
print(np.shape(out_np))
#https://gitorchub.com/DediGadot/PatchBatch/blob/master/flowlib.py
out_np = np.squeeze(out_np)
print(np.shape(out_np))
#out_np = np.moveaxis(out_np, 0, -1)
im_arr = flowlib.flow_to_image(out_np)
im = Image.fromarray(im_arr)
im.save("%s.png"%name)
def visualize_pair(batch_1, batch_2):
if len(batch_1.shape) == 4:
if batch_1.shape[3] == 2:
batch_1 = [flow_to_image(batch_1[i]) for i in range(batch_1.shape[0])]
else:
batch_1 = [batch_1[i] for i in range(batch_1.shape[0])]
if batch_2.shape[3] == 2:
batch_2 = [flow_to_image(batch_2[i]) for i in range(batch_2.shape[0])]
else:
batch_2 = [batch_2[i] for i in range(batch_2.shape[0])]
cv2.imshow('Pair comparison', np.vstack([np.hstack(batch_1), np.hstack(batch_2)]))
cv2.waitKey(0)
else:
if batch_1.shape[4] == 2:
batch_1 = [flow_to_image(batch_1[-1][i]) for i in range(batch_1[-1].shape[0])]
else:
batch_1 = [batch_1[-1][i] for i in range(batch_1[-1].shape[0])]
if batch_2.shape[4] == 2:
batch_2 = [flow_to_image(batch_2[-1][i]) for i in range(batch_2[-1].shape[0])]
else:
batch_2 = [batch_2[-1][i] for i in range(batch_2[-1].shape[0])]
cv2.imshow('Pair comparison', np.vstack([np.hstack(batch_1), np.hstack(batch_2)]))
cv2.waitKey(0)
def get_opt_flow_rgb(dir, timestamp, datapoints=None, correction=True):
opt_flow_n = os.path.join(dir, 'images', 'flow_' + str(timestamp) + '.npz')
opt_flow = np.load(opt_flow_n)['opt_flow']
opt_flow = cv2.resize(opt_flow, (256, 144))
# optical_flow_correction_stats needs to be rune before to callibrate flowlib rgb conversion
opt_flow = o_f_compensate_for_rotation(
opt_flow,
datapoints.loc[timestamp][['AV_x', 'AV_y', 'AV_z']],
focal=4.47189418e+00,
xy_scaling=4.78224949e-06)
flowlib.enable_sticky_max_flow(False)
opt_flow_rgb = flowlib.flow_to_image(opt_flow)
return opt_flow_rgb
def plot_middlebury_color_code(
xlim=(-10, 10), ylim=(-10, 10), resolution=100, **kwargs):
'''For some reason, colors are flipped on the y axis
with respect to real middlebury color code.
'''
x = np.linspace(xlim[0], xlim[1], resolution)
y = np.linspace(ylim[0], ylim[1], resolution)
X, Y = np.meshgrid(x, y)
X = np.expand_dims(X, -1)
Y = np.expand_dims(Y, -1)
C = np.concatenate([X, Y], axis=-1)
plt.title('Middlebury color code')
plt.imshow(flow_to_image(C), origin='lower', **kwargs)
plt.plot([0, resolution - 1], [resolution / 2, resolution / 2], c='black')
plt.plot([resolution / 2, resolution / 2], [0, resolution - 1], c='black')
plt.axis('off')
def view_flow(flow_path, flow_name):
flow = read_flo(flow_path)
flow_img = flowlib.flow_to_image(flow)
png_name = flow_name.split('.')[0] + '.png'
png_path = os.path.join(flow2img_path, png_name)
imsave(png_path, flow_img)
def flow_color(flow):
return flow_to_image(flow.transpose(1, 2, 0))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Convert to float and add batch dimension.
return np.expand_dims(img.astype(np.float32) / 255., axis=0)
if __name__ == "__main__":
# Graph construction.
image1_placeholder = tf.placeholder(tf.float32, [1, 384, 512, 3])
image2_placeholder = tf.placeholder(tf.float32, [1, 384, 512, 3])
flownet2 = models.FlowNet2()
inputs = {"input_a": image1_placeholder, "input_b": image2_placeholder}
predicted_flow = flownet2.run(inputs)["flow"]
# Load inputs.
image1 = load_image("example/0img0.ppm")
image2 = load_image("example/0img1.ppm")
ckpt_file = "checkpoints/FlowNet2/flownet-2.ckpt-0"
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, ckpt_file)
feed_dict = {image1_placeholder: image1, image2_placeholder: image2}
predicted_flow_np = sess.run(predicted_flow, feed_dict=feed_dict)
# Show visualization.
flow_image = flowlib.flow_to_image(predicted_flow_np[0])
plt.imshow(flow_image)
plt.show()
optical_flow = cv2.DualTVL1OpticalFlow_create() ret, frame = cap.read() prvs_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) hsv = np.zeros_like(frame) hsv[...,1] = 255 while(1): ret, frame = cap.read() if ret==True: next_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) flow = optical_flow.calc(prvs_img, next_img, None) rgb = flow_to_image(flow) # mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1]) # hsv[...,0] = ang*180/np.pi/2 # hsv[...,2] = cv2.normalize(mag,None,0,255,cv2.NORM_MINMAX) # rgb = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR) if int(cap.get(1))%5==0: img_name = 'cook2'+str(int(cap.get(1)))+'.jpg' cv2.imwrite(img_name, rgb) k = cv2.waitKey(30) & 0xff if k==27: break prvs_img = next_img else:
def plot_one_image(im, title='', vmin=None, vmax=None):
if im.shape[0] == 2:
im = flowlib.flow_to_image(im.transpose(1, 2, 0))
plt.axis('off')
plt.imshow(im.squeeze(), origin='lower', vmin=vmin, vmax=vmax)
def eval_flow_avg(gt_flows,
noc_masks,
pred_flows,
cfg,
moving_masks=None,
write_img=False):
error, error_noc, error_occ, error_move, error_static, error_rate = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
error_move_rate, error_static_rate = 0.0, 0.0
num = len(gt_flows)
for gt_flow, noc_mask, pred_flow, i in zip(gt_flows, noc_masks, pred_flows,
range(len(gt_flows))):
H, W = gt_flow.shape[0:2]
pred_flow = np.copy(pred_flow)
pred_flow[:, :, 0] = pred_flow[:, :, 0] / cfg.img_hw[1] * W
pred_flow[:, :, 1] = pred_flow[:, :, 1] / cfg.img_hw[0] * H
flo_pred = cv2.resize(
pred_flow, (W, H), interpolation=cv2.INTER_LINEAR)
if write_img:
if not os.path.exists(os.path.join(cfg.model_dir, "pred_flow")):
os.mkdir(os.path.join(cfg.model_dir, "pred_flow"))
cv2.imwrite(
os.path.join(cfg.model_dir, "pred_flow",
str(i).zfill(6) + "_10.png"),
flow_to_image(flo_pred))
cv2.imwrite(
os.path.join(cfg.model_dir, "pred_flow",
str(i).zfill(6) + "_10_gt.png"),
flow_to_image(gt_flow[:, :, 0:2]))
cv2.imwrite(
os.path.join(cfg.model_dir, "pred_flow",
str(i).zfill(6) + "_10_err.png"),
flow_to_image(
(flo_pred - gt_flow[:, :, 0:2]) * gt_flow[:, :, 2:3]))
epe_map = np.sqrt(
np.sum(np.square(flo_pred[:, :, 0:2] - gt_flow[:, :, 0:2]),
axis=2))
error += np.sum(epe_map * gt_flow[:, :, 2]) / np.sum(gt_flow[:, :, 2])
error_noc += np.sum(epe_map * noc_mask) / np.sum(noc_mask)
error_occ += np.sum(epe_map * (gt_flow[:, :, 2] - noc_mask)) / max(
np.sum(gt_flow[:, :, 2] - noc_mask), 1.0)
error_rate += calculate_error_rate(epe_map, gt_flow[:, :, 0:2],
gt_flow[:, :, 2])
if moving_masks:
move_mask = moving_masks[i]
error_move_rate += calculate_error_rate(
epe_map, gt_flow[:, :, 0:2], gt_flow[:, :, 2] * move_mask)
error_static_rate += calculate_error_rate(
epe_map, gt_flow[:, :, 0:2],
gt_flow[:, :, 2] * (1.0 - move_mask))
error_move += np.sum(epe_map * gt_flow[:, :, 2] *
move_mask) / np.sum(gt_flow[:, :, 2] *
move_mask)
error_static += np.sum(epe_map * gt_flow[:, :, 2] * (
1.0 - move_mask)) / np.sum(gt_flow[:, :, 2] *
(1.0 - move_mask))
if moving_masks:
result = "{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10} \n".format(
'epe', 'epe_noc', 'epe_occ', 'epe_move', 'epe_static',
'move_err_rate', 'static_err_rate', 'err_rate')
result += "{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f} \n".format(
error / num, error_noc / num, error_occ / num, error_move / num,
error_static / num, error_move_rate / num, error_static_rate / num,
error_rate / num)
return result
else:
result = "{:>10}, {:>10}, {:>10}, {:>10} \n".format(
'epe', 'epe_noc', 'epe_occ', 'err_rate')
result += "{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f} \n".format(
error / num, error_noc / num, error_occ / num, error_rate / num)
return result
def test(self,
checkpoint,
input_path,
out_path,
save_image=False,
save_flo=False):
img_file_list = os.listdir(input_path)
img_file_list.sort(key=lambda f: int(filter(str.isdigit, f)))
img_list = []
for img_file_name in img_file_list:
img_file_path = os.path.join(input_path, img_file_name)
img = imread(img_file_path)
img = cv2.resize(img, (512, 384), interpolation=cv2.INTER_AREA)
# Convert from RGB -> BGR
img = img[..., [2, 1, 0]]
img_list.append(img)
input_a = np.stack(img_list[:-1], axis=0)
input_b = np.stack(img_list[1:], axis=0)
# Scale from [0, 255] -> [0.0, 1.0] if needed
if input_a.max() > 1.0:
input_a = input_a / 255.0
if input_b.max() > 1.0:
input_b = input_b / 255.0
# TODO: This is a hack, we should get rid of this
training_schedule = LONG_SCHEDULE
inputs = {
'input_a': tf.constant(input_a, dtype=tf.float32),
'input_b': tf.constant(input_b, dtype=tf.float32),
}
predictions = self.model(inputs, training_schedule)
pred_flow = predictions['flow']
saver = tf.train.Saver()
num = int(filter(str.isdigit, input_path))
file = open(os.path.join(out_path, str(num) + '_flow.csv'), 'w')
writer = csv.writer(file, delimiter=',', quotechar='|')
with tf.Session() as sess:
saver.restore(sess, checkpoint)
pred_flow_seq = sess.run(pred_flow)
pred_flow_lists = np.split(pred_flow_seq,
len(pred_flow_seq),
axis=0)
for t, pred_flow in enumerate(pred_flow_lists):
# unique_name = input_a_path.split('.')[-2].split('/')[-1]+'and'+input_b_path.split('.')[-2].split('/')[-1]
unique_name = str(t)
pred_flow = np.squeeze(pred_flow)
shape = np.shape(pred_flow)
mean_flow = np.mean(
pred_flow[shape[0] / 4:shape[0] / 4 * 3,
shape[1] / 4:shape[1] / 4 * 3, :],
axis=(0, 1))
writer.writerow(mean_flow)
if save_image:
flow_img = flow_to_image(pred_flow)
full_out_path = os.path.join(out_path,
unique_name + '.png')
imsave(full_out_path, flow_img)
if save_flo:
full_out_path = os.path.join(out_path,
unique_name + '.flo')
write_flow(pred_flow, full_out_path)
file.close()
def visualize_img(img):
if img.shape[2] == 2:
cv2.imshow('Optical flow', flow_to_image(img))
else:
cv2.imshow('Image', img)
cv2.waitKey(0)
def optical_flow_correction_stats(data_dir, batch_size):
orig_image_size = (270, 480) # (h,w)
FOV = 90 # in degeres
cam_f = orig_image_size[1] / 2 / np.tan(FOV / 180 * np.pi / 2)
corr_params = {
'raw': None,
'optimized': (4.47189418e+00, 4.78224949e-06),
'calculated':
(orig_image_size[1] / 2 / np.tan(FOV / 180 * np.pi / 2), 1)
}
datapoints_csv_name = os.path.join(data_dir, 'airsim_rec.csv')
datapoints = pd.read_csv(datapoints_csv_name,
header=None,
sep=',',
names=csv_col_names,
index_col=0)
ix = np.random.choice(np.arange(len(datapoints) - 1), batch_size) + 1
flow_to_save = []
avg_mag = {'raw': 0, 'optimized': 0, 'calculated': 0}
global max_flow # Yes, I feel bad for using a global but whatever, time shortcuts
print('Generating optical flow examples...')
for i, idx in (enumerate(ix)):
opt_flow = get_opt_flow(data_dir, datapoints.index[idx])
# print(hash(str(opt_flow)))
for name, params in corr_params.items():
if params is not None:
# corrected optical flow
opt_flow_alt = o_f_compensate_for_rotation(
opt_flow,
datapoints.iloc[idx][['AV_x', 'AV_y', 'AV_z']],
focal=params[0],
xy_scaling=params[1])
else:
opt_flow_alt = opt_flow
savename = os.path.abspath(
os.path.join(plotdir, 'opt_fl_eg',
name + '_{}.png'.format(str(i))))
flow_to_save.append((savename, opt_flow_alt))
flow_norm = np.linalg.norm(opt_flow_alt, axis=-1)
avg_mag[name] += np.mean(flow_norm) / batch_size
print('{}_{}'.format(name, i))
print(np.max(flow_norm))
if name != 'calculated':
# Calculated corrector gives overinflated flow values
max_flow = np.max(
flow_norm) if np.max(flow_norm) > max_flow else max_flow
savename = os.path.abspath(
os.path.join(plotdir, 'opt_fl_eg',
'farneback_{}.png'.format(str(i))))
opt_flow_alt = get_cv2_flow(datapoints,
os.path.join(data_dir, 'images'), idx)
flow_to_save.append((savename, opt_flow_alt))
savename = os.path.abspath(
os.path.join(plotdir, 'opt_fl_eg', 'scene_{}.png'.format(str(i))))
opt_flow_alt = get_scene(datapoints, os.path.join(data_dir, 'images'),
idx)
flow_to_save.append((savename, opt_flow_alt))
print('Max flow mag encountered: {}'.format(max_flow))
flowlib.set_max_flow(max_flow)
for flow in flow_to_save:
if not os.path.isdir(os.path.dirname(flow[0])):
os.makedirs(os.path.dirname(flow[0]))
if flow[1].shape[-1] != 3:
opt_flow_rgb = flowlib.flow_to_image(flow[1])
else:
opt_flow_rgb = flow[1]
cv2.imwrite(flow[0], opt_flow_rgb)
print('Max flow mag: {}'.format(max_flow))
for name, params in corr_params.items():
print('Correction parameters: {}\nAvg flow: {}\n\n'.format(
params, avg_mag[name]))
def train():
if not (os.path.exists(FLAGS['checkpoint_dir'])
and FLAGS['mode'] == 'train_rigid'):
os.makedirs(FLAGS['checkpoint_dir'])
else:
if len(os.listdir(FLAGS['checkpoint_dir'])) > 0:
print("Notice : Please remove exist checkpoints")
exit()
summary_writer = tf.summary.create_file_writer(FLAGS['summary_dir'])
if FLAGS['mode'] == 'train_flow':
checkpoint_path = os.path.join(FLAGS['init_ckpt_file'])
geonet.load_weights(checkpoint_path)
with summary_writer.as_default():
if FLAGS['mode'] == 'train_rigid':
start_time = time.time()
for step in range(FLAGS['max_steps']):
src_image_stack, tgt_image, intrinsics = data_loader.get_next()
loss, tgt_image_pyramid, src_image_concat_pyramid, \
pred_disp, pred_depth, pred_poses, \
fwd_rigid_error_pyramid, bwd_rigid_error_pyramid, \
fwd_rigid_warp_pyramid, bwd_rigid_warp_pyramid, \
fwd_rigid_flow_pyramid, bwd_rigid_flow_pyramid = train_rigid_step(src_image_stack, tgt_image, intrinsics)
if (step % 100 == 0):
time_per_iter = (time.time() - start_time) / 100
start_time = time.time()
print('Iteration: [%7d] | Time: %4.4fs/iter | Loss: %.3f' %
(step, time_per_iter, loss))
if (step % FLAGS['save_summary_freq'] == 0):
# Summary
# loss
tf.summary.scalar('loss', loss, step=step)
# input images
tf.summary.image('tgt_image',
deprocess_image(tgt_image_pyramid[0]) /
255,
step=step) # [4, 128, 416, 3]
tf.summary.image(
'src_image',
deprocess_image(src_image_concat_pyramid[0]) / 255,
step=step) # [8, 128, 416, 3]
# pred_depth # [12, 128, 416, 1] x num_scales=3
norm_pred_depth = normalize_depth_for_display(
pred_depth[0].numpy())
tf.summary.image('norm_pred_depth',
norm_pred_depth,
step=step)
# rigid_warp_img
tf.summary.image(
'fwd_rigid_warp_img',
deprocess_image(fwd_rigid_warp_pyramid[0]) / 255,
step=step) # [8, 128, 416, 3]
tf.summary.image(
'bwd_rigid_warp_img',
deprocess_image(bwd_rigid_warp_pyramid[0]) / 255,
step=step) # [8, 128, 416, 3]
# pred rigid flow
color_fwd_flow_list = []
color_bwd_flow_list = []
for i in range(FLAGS['batch_size']):
color_fwd_flow = fl.flow_to_image(
fwd_rigid_flow_pyramid[0][
i, :, :, :].numpy()) # [8, 128, 416, 2]
color_fwd_flow = cv2.cvtColor(color_fwd_flow,
cv2.COLOR_RGB2BGR)
# color_fwd_flow = tf.expand_dims(color_fwd_flow, axis=0)
color_fwd_flow_list.append(color_fwd_flow)
color_bwd_flow = fl.flow_to_image(
bwd_rigid_flow_pyramid[0][
i, :, :, :].numpy()) # [8, 128, 416, 2]
color_bwd_flow = cv2.cvtColor(color_bwd_flow,
cv2.COLOR_RGB2BGR)
# color_bwd_flow = tf.expand_dims(color_bwd_flow, axis=0)
color_bwd_flow_list.append(color_bwd_flow)
color_fwd_flow = np.array(color_fwd_flow_list)
# color_fwd_flow = tf.squeeze(color_fwd_flow)
color_bwd_flow = np.array(color_bwd_flow_list)
# color_bwd_flow = tf.squeeze(color_bwd_flow)
tf.summary.image('color_fwd_flow',
color_fwd_flow,
step=step)
tf.summary.image('color_bwd_flow',
color_bwd_flow,
step=step)
if (step % FLAGS['save_ckpt_freq'] == 0):
save_path = FLAGS['checkpoint_dir'] + "iter-" + str(step)
geonet.save_weights(save_path)
print("Saved checkpoint for step {}: {}".format(
step, FLAGS['checkpoint_dir']))
elif FLAGS['mode'] == 'train_flow':
start_time = time.time()
for step in range(FLAGS['max_steps']):
src_image_stack, tgt_image, intrinsics = data_loader.get_next()
loss, tgt_image_pyramid, src_image_concat_pyramid, \
fwd_full_flow_pyramid, bwd_full_flow_pyramid, \
fwd_full_error_pyramid, bwd_full_error_pyramid, \
noc_masks_tgt, noc_masks_src, \
tgt_image_tile_pyramid, src_image_concat_pyramid, \
fwd_flow_diff_pyramid, bwd_flow_diff_pyramid = train_flow_step(src_image_stack, tgt_image,
intrinsics)
if (step % 100 == 0):
time_per_iter = (time.time() - start_time) / 100
start_time = time.time()
print('Iteration: [%7d] | Time: %4.4fs/iter | Loss: %.3f' %
(step, time_per_iter, loss))
if (step % FLAGS['save_summary_freq'] == 0):
# Summary
# loss
tf.summary.scalar('loss', loss, step=step)
# input images
tf.summary.image('tgt_image',
deprocess_image(tgt_image_pyramid[0]) /
255,
step=step) # [4, 128, 416, 3]
tf.summary.image(
'src_image',
deprocess_image(src_image_concat_pyramid[0]) / 255,
step=step) # [8, 128, 416, 3]
# pred full flow
color_fwd_flow_list = []
color_bwd_flow_list = []
for i in range(FLAGS['batch_size']):
color_fwd_flow = fl.flow_to_image(
fwd_full_flow_pyramid[0][
i, :, :, :].numpy()) # [8, 128, 416, 2]
color_fwd_flow = cv2.cvtColor(color_fwd_flow,
cv2.COLOR_RGB2BGR)
# color_fwd_flow = tf.expand_dims(color_fwd_flow, axis=0)
color_fwd_flow_list.append(color_fwd_flow)
color_bwd_flow = fl.flow_to_image(
bwd_full_flow_pyramid[0][
i, :, :, :].numpy()) # [8, 128, 416, 2]
color_bwd_flow = cv2.cvtColor(color_bwd_flow,
cv2.COLOR_RGB2BGR)
# color_bwd_flow = tf.expand_dims(color_bwd_flow, axis=0)
color_bwd_flow_list.append(color_bwd_flow)
color_fwd_flow = np.array(color_fwd_flow_list)
# color_fwd_flow = tf.squeeze(color_fwd_flow)
color_bwd_flow = np.array(color_bwd_flow_list)
# color_bwd_flow = tf.squeeze(color_bwd_flow)
tf.summary.image('color_fwd_flow',
color_fwd_flow,
step=step)
tf.summary.image('color_bwd_flow',
color_bwd_flow,
step=step)
if (step % FLAGS['save_ckpt_freq'] == 0):
save_path = FLAGS['checkpoint_dir'] + "iter-" + str(step)
geonet.save_weights(save_path)
print("Saved checkpoint for step {}: {}".format(
step, FLAGS['checkpoint_dir']))
def plot_vel(xd):
plt.axis('off')
plt.imshow(flowlib.flow_to_image(xd.reshape(64, 64, 3)[:, :, 1:]))
def flow_to_image(flow):
return flowlib.flow_to_image(flow.transpose(1, 2, 0))
def plot_vel(x):
plt.imshow(flowlib.flow_to_image(x.reshape(64, 64, 3)[:, :, 1:]))
def visualize_pair_map(batch_1, batch_2):
if len(batch_1.shape) == 4:
if batch_1.shape[3] == 2:
batch_show_1 = [
flow_to_image(batch_1[i]) for i in range(batch_1.shape[0])
]
else:
batch_show_1 = [batch_1[i] for i in range(batch_1.shape[0])]
if batch_2.shape[3] == 2:
batch_show_2 = [
flow_to_image(batch_2[i]) for i in range(batch_2.shape[0])
]
else:
batch_show_2 = [batch_2[i] for i in range(batch_2.shape[0])]
if batch_1.shape[3] == 3 or batch_1.shape[3] == 1: # RGB or GRAYSCALE
batchtmp_1 = [
cv2.normalize(batch_1[i], None, 0, 255, cv2.NORM_MINMAX,
cv2.CV_8U) for i in range(batch_1.shape[0])
]
batchtmp_2 = [
cv2.normalize(batch_2[i], None, 0, 255, cv2.NORM_MINMAX,
cv2.CV_8U) for i in range(batch_1.shape[0])
]
if batch_1.shape[3] == 1:
error_gray = [
cv2.absdiff(batchtmp_1[i], batchtmp_2[i])
for i in range(batch_1.shape[0])
]
elif batch_1.shape[3] == 3:
error_rgb = [
cv2.absdiff(batchtmp_1[i], batchtmp_2[i])
for i in range(batch_1.shape[0])
]
error_gray = [
cv2.cvtColor(error_rgb[i], cv2.COLOR_BGR2GRAY)
for i in range(batch_1.shape[0])
]
heatmap = [
cv2.applyColorMap(error_gray[i], cv2.COLORMAP_JET)
for i in range(batch_1.shape[0])
]
if batch_1.shape[3] == 3: # RGB
cv2.imshow(
'Pair comparison AP',
np.vstack([
np.hstack(batch_show_1),
np.hstack(batch_show_2),
np.hstack(heatmap)
]))
else: # GRAYSCALE
cv2.imshow('Pair comparison AP',
np.vstack([
np.hstack(batch_show_1),
np.hstack(batch_show_2)
])) # GRAYSCALE
cv2.imshow('Error AP',
np.vstack([np.hstack(heatmap)
])) # different color space: RGB
else:
cv2.imshow(
'Pair comparison OF',
np.vstack([np.hstack(batch_show_1),
np.hstack(batch_show_2)]))
cv2.waitKey(0)
else:
if batch_1.shape[4] == 2:
batch_1 = [
flow_to_image(batch_1[-1][i])
for i in range(batch_1[-1].shape[0])
]
else:
batch_1 = [batch_1[-1][i] for i in range(batch_1[-1].shape[0])]
if batch_2.shape[4] == 2:
batch_2 = [
flow_to_image(batch_2[-1][i])
for i in range(batch_2[-1].shape[0])
]
else:
batch_2 = [batch_2[-1][i] for i in range(batch_2[-1].shape[0])]
cv2.imshow('Pair comparison',
np.vstack([np.hstack(batch_1),
np.hstack(batch_2)]))
cv2.waitKey(0)
def plot_flow_color(flow, **kwargs):
plt.imshow(flow_to_image(flow.transpose(1, 2, 0)),
origin='lower',
**kwargs)
def test_flow(opt):
##### load testing list #####
with open(opt.dataset_dir + "test_flow.txt", 'r') as f:
test_files = f.readlines()
input_list = []
for seq in test_files:
seq = seq.split(' ')
input_list.append(opt.dataset_dir+seq[0]+'/'+seq[1][:-1])
if not os.path.exists(opt.output_dir):
os.makedirs(opt.output_dir)
##### init #####
# TODO: currently assuming batch_size = 1
assert opt.batch_size == 1
tgt_image_uint8 = tf.placeholder(tf.uint8, [opt.batch_size,
opt.img_height, opt.img_width, 3],
name='tgt_input')
src_image_stack_uint8 = tf.placeholder(tf.uint8, [opt.batch_size,
opt.img_height, opt.img_width, opt.num_source * 3],
name='src_stack_input')
intrinsics = tf.placeholder(tf.float32, [opt.batch_size, 3, 3],
name='intrinsics_input')
loader = DataLoader(opt)
intrinsics_ms = loader.get_multi_scale_intrinsics(intrinsics, opt.num_scales)
# currently assume a sequence is fed and the tgt->src_id flow is computed
src_id = int(opt.num_source // 2)
bs = opt.batch_size
model = GeoNetModel(opt, tgt_image_uint8, src_image_stack_uint8, intrinsics_ms)
fetches = {}
fetches["pred_flow"] = model.fwd_full_flow_pyramid[0][bs*src_id:bs*(src_id+1)]
saver = tf.train.Saver([var for var in tf.model_variables()])
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
##### Go! #####
output_file = opt.output_dir + '/' + os.path.basename(opt.init_ckpt_file)
if not os.path.exists(output_file):
os.makedirs(output_file)
binary_dir = os.path.join(output_file, 'binary')
color_dir = os.path.join(output_file, 'color')
png_dir = os.path.join(output_file, 'png')
if (not os.path.exists(binary_dir)):
os.makedirs(binary_dir)
if (not os.path.exists(color_dir)):
os.makedirs(color_dir)
if (not os.path.exists(png_dir)):
os.makedirs(png_dir)
with tf.Session(config=config) as sess:
saver.restore(sess, opt.init_ckpt_file)
pred_all = []
img_num = len(input_list)
for tgt_idx in range(img_num):
if (tgt_idx+1) % 100 == 0:
print('processing: %d/%d' % (tgt_idx+1, img_num))
image_seq = cv2.imread(input_list[tgt_idx]+'.jpg')
tgt_image, src_image_stack = unpack_image_sequence(image_seq,
opt.img_height, opt.img_width, opt.num_source)
with open(input_list[tgt_idx]+'_cam.txt', 'r') as cf:
cam_file = cf.readlines()
cam_file = cam_file[0].split(',')
cam_file = np.array([float(d) for d in cam_file])
cam_file = np.reshape(cam_file, (3,3))
pred = sess.run(fetches, feed_dict={tgt_image_uint8: tgt_image[None, :, :, :],
src_image_stack_uint8: src_image_stack[None, :, :, :],
intrinsics: cam_file[None,:,:]})
pred_flow=pred['pred_flow'][0]
# save flow
flow_fn = '%.6d.png' % tgt_idx
color_fn = os.path.join(color_dir, flow_fn)
color_flow = fl.flow_to_image(pred_flow)
color_flow = cv2.cvtColor(color_flow, cv2.COLOR_RGB2BGR)
color_flow = cv2.imwrite(color_fn, color_flow)
png_fn = os.path.join(png_dir, flow_fn)
mask_blob = np.ones((opt.img_height, opt.img_width), dtype = np.uint16)
fl.write_kitti_png_file(png_fn, pred_flow, mask_blob)
binary_fn = flow_fn.replace('.png', '.flo')
binary_fn = os.path.join(binary_dir, binary_fn)
fl.write_flow(pred_flow, binary_fn)
def visualize_prediction(datapoint, display=True, metadata=None):
inputs, ground_truth, rgbs, timestamps = datapoint
predictions = loaded_model.predict(inputs)
if type(predictions) is list:
predictions, layers_debug = predictions
inputs = inputs[0]
ground_truth = one_hot_to_img(ground_truth, bins, inputs.shape[-3:-1])
print(predictions[0, 777, :])
predictions = one_hot_to_img(predictions, bins, inputs.shape[-3:-1])
predictions = predictions.astype(np.uint8)
ground_truth = ground_truth.astype(np.uint8)
# make into rgb grayscale
if len(ground_truth.shape) == 3:
ground_truth = np.expand_dims(ground_truth, axis=-1)
ground_truth = np.tile(ground_truth, (1, 1, 1, 3))
if len(predictions.shape) == 3:
predictions = np.expand_dims(predictions, axis=-1)
predictions = np.tile(predictions, (1, 1, 1, 3))
rgb_scaled = np.empty_like(ground_truth)
optical_flows = np.empty_like(ground_truth)
for idx, inp in enumerate(inputs):
rgb_scaled[idx] = cv2.resize(rgbs[idx],
(inputs.shape[2], inputs.shape[1]))
optical_flows[idx] = flowlib.flow_to_image(inputs[idx])
sep_thickness = 5
sep_pat = 4
separator = np.tile(
((np.arange(inputs.shape[-2], dtype=np.uint8) % sep_pat) / sep_pat) *
255, (sep_thickness, 1))
separator = np.tile(np.expand_dims(separator, axis=-1), (1, 1, 3))
separator = np.expand_dims(separator, axis=0)
separator = separator.astype(np.uint8)
collated_img = np.concatenate(
(rgb_scaled, separator, optical_flows, separator, ground_truth,
separator, predictions),
axis=1)
# Present also some metadata information
if metadata is not None:
metadata_idx = metadata['timestamp_idx'][timestamps[0]]
keys_to_display = ['cam_lin_vel', 'cam_ang_vel']
for key in keys_to_display:
if key == 'cam_lin_vel':
data = generator.velocity_form(metadata[key][metadata_idx])
else:
data = metadata[key][metadata_idx]
print(key + ': ' + str(data))
if display:
if metadata is not None:
print("timestamp {}\tmetadata_idx {}".format(
timestamps[0], metadata_idx))
else:
print("timestamp {}".format(timestamps[0]))
cv2.imshow("preview", collated_img[0])
k = cv2.waitKey(0)
if k == 27:
return None
elif k == ord('s'):
# Save the image example
filename = get_f_name_no_duplicates(auto_checkpoint_name, 'png',
'saved_preview')
cv2.imwrite(filename, collated_img[0])
return collated_img
##### Display flow #####
flo_im = flo[0] * 20.0
flo_im = flo_im.cpu().data.numpy()
# scale the flow back to the input size
flo_im = np.swapaxes(np.swapaxes(flo_im, 0, 1), 1, 2) #
u_ = cv2.resize(flo_im[:, :, 0], (W, H))
v_ = cv2.resize(flo_im[:, :, 1], (W, H))
u_ *= W / float(W_)
v_ *= H / float(H_)
flo_im = np.dstack((u_, v_))
writeFlowFile(flow_fn, flo_im)
flo_img = flow_to_image(flo_im)
plt.imshow(flo_img)
plt.show()
##### Display warped image #####
flo_warp = 20.0 * torch.nn.functional.interpolate(
flo, scale_factor=4, mode='bilinear', align_corners=False)
#pdb.set_trace()
img_warped = warp(im2, flo_warp)
img_warped = img_warped.squeeze()
img_warped = img_warped.cpu().data.numpy()
img_warped = np.transpose(img_warped, (1, 2, 0))
plt.imshow(img_warped)
plt.show()
