def save_flow(input_path, ref_path, flow_path):
inputs = glob.glob(os.path.join(input_path, '*.png'))
for i, input in enumerate(tqdm(inputs)):
img_i = cv2.imread(input)
img_i = cv2.cvtColor(img_i, cv2.COLOR_BGR2GRAY)
img_r = cv2.imread(os.path.join(ref_path, input.split('/')[-1]))
img_r = cv2.cvtColor(img_r, cv2.COLOR_BGR2GRAY)
tmp_flow = compute_flow(img_i, img_r)
# tmp_flow = ToImg(tmp_flow)
if not os.path.exists(os.path.join(flow_path, 'color')):
# os.makedirs os.mkdir
os.makedirs(os.path.join(flow_path, 'color'))
tmp_flow_color = flow_vis.flow_to_color(tmp_flow, convert_to_bgr=False)
cv2.imwrite(os.path.join(flow_path, 'color', input.split('/')[-1]), tmp_flow_color)
if not os.path.exists(os.path.join(flow_path, 'u')):
# os.makedirs os.mkdir
os.makedirs(os.path.join(flow_path, 'u'))
cv2.imwrite(os.path.join(flow_path, 'u', input.split('/')[-1]), tmp_flow[:, :, 0])
if not os.path.exists(os.path.join(flow_path, 'v')):
# os.makedirs os.mkdir
os.makedirs(os.path.join(flow_path, 'v'))
cv2.imwrite(os.path.join(flow_path, 'v', input.split('/')[-1]), tmp_flow[:, :, 1])
print('complete:' + flow_path)
return
def process_all_images(self):
# get a list of png imgs
curr_img_list = glob.glob(os.path.join(self.input_dir,
"image_2/*.png"))
curr_img_list = sorted(curr_img_list)
log_fn = os.path.join(self.output_dir, 'get_optical_flow.log')
for curr_img_path in curr_img_list:
im2_fn = curr_img_path
seq = curr_img_path.split("/")[-1].split(".")[
0] # get serial number of the img
im1_fn = os.path.join(self.prev_dir, seq +
'_01.png') # get path of current-1 image
im0_fn = os.path.join(self.prev_dir, seq +
'_02.png') # get path of current-2 image
print("processing img #" + seq)
import time
start = time.time()
of = self.pwc_fusion(
im0_fn, im1_fn, im2_fn
) # select either self.pwc_net(im1_fn, im2_fn) or self.pwc_fusion(im0_fn, im1_fn, im2_fn)
of_fn = os.path.join(self.output_dir, seq + '_pwc_fusion.npy')
np.save(of_fn, of)
flow_color = flow_vis.flow_to_color(
of, convert_to_bgr=True
) # Apply the coloring (for OpenCV, set convert_to_bgr=True)
flow_color_fn = os.path.join(self.output_dir,
seq + '_pwc_fusion.png')
cv2.imwrite(flow_color_fn, flow_color)
end = time.time()
print("completed in " + str(end - start) + "s")
with open(log_fn, "a") as f:
f.write("completed img #" + seq + " in " + str(end - start) +
"s")
def main():
dataset_dir = cfg.DATA_ROOT_DIR
alignments_path = os.path.join(
dataset_dir, "{0}_selfsupervised.json".format(cfg.DATA_TYPE))
flow_normalization = 100.0
# in pixels
with open(alignments_path, "r") as f:
pairs = json.load(f)
for pair in pairs:
source_color = cv2.imread(
os.path.join(dataset_dir, pair["source_color"]))
target_color = cv2.imread(
os.path.join(dataset_dir, pair["target_color"]))
optical_flow = utils.load_flow(
os.path.join(dataset_dir, pair["optical_flow"]))
optical_flow = np.moveaxis(optical_flow, 0, -1) # (h, w, 2)
invalid_flow = optical_flow == -np.Inf
optical_flow[invalid_flow] = 0.0
flow_color = flow_vis.flow_to_color(optical_flow, convert_to_bgr=False)
cv2.imshow("Source", source_color)
cv2.imshow("Target", target_color)
cv2.imshow("Flow", flow_color)
cv2.waitKey(0)
def save_flow_video(flow_array, save_path):
def get_img_size(image):
height, width, _ = image.shape
return width, height
save_dir = os.path.dirname(save_path)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
flow_imgs = []
for flow in flow_array:
flow_img = flow_vis.flow_to_color(flow, convert_to_bgr=False)
flow_imgs.append(flow_img)
fps = 10
size = get_img_size(flow_imgs[0])
out = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'DIVX'), fps,
size)
for i in range(len(flow_imgs)):
out.write(flow_imgs[i])
if i % 2500 == 0:
logging.info('Building video: {0}'.format(i))
out.release()
print("Video saved at " + save_path)
def visualize_flows(frames, flows, num_imgs, video_save_path=None):
combined = list(zip(frames, flows))
enumerated = list(enumerate(combined))
random_samples = random.choices(enumerated[1:len(enumerated)-1], k=num_imgs)
f, axes = plt.subplots(math.ceil(num_imgs/4), 4, figsize = (60, 40))
print("Raw Images")
for num, axs in enumerate(axes):
for i, ax in enumerate(axs):
index = (num * 4 + i)
if index < num_imgs:
img = random_samples[index][1][0]
ax.imshow(img, aspect='auto')
plt.show()
print("FlowNet2 Images")
f2, axes2 = plt.subplots(math.ceil(num_imgs/4), 4, figsize = (60, 40))
for num, axs in enumerate(axes2):
for i, ax in enumerate(axs):
index = (num * 4 + i)
if index < num_imgs:
flow_color = flow_vis.flow_to_color(random_samples[index][1][1], convert_to_bgr=False)
ax.imshow(flow_color, aspect='auto')
plt.show()
print("OpenCV Images")
f3, axes3 = plt.subplots(math.ceil(num_imgs/4), 4, figsize = (60, 40))
for num, axs in enumerate(axes3):
for i, ax in enumerate(axs):
index = (num * 4 + i)
if index < num_imgs:
img_index = random_samples[index][0]
frame1 = frames[img_index]
frame2 = frames[img_index + 1]
flow = FlowUtils.calc_opt_flow(frame1, frame2)
flow_color = flow_vis.flow_to_color(flow, convert_to_bgr=False)
ax.imshow(flow_color, aspect='auto')
plt.show()
if video_save_path:
VideoUtils.save_flow_video(flows, video_save_path)
def magnitudeOP(self,flow_img, path):
img = plt.imread(path)
flow_color = flow_vis.flow_to_color(flow_img[:, :, :2], convert_to_bgr=False)
plt.imshow(flow_color)
plt.imshow(img, alpha=0.2, cmap='gray')
plt.title('Magnitude OF')
plt.xticks([])
plt.yticks([])
plt.show()
def magnitudeOP_save(self, flow_img, path, alg, type):
img = plt.imread(path)
flow_color = flow_vis.flow_to_color(flow_img[:, :, :2],
convert_to_bgr=False)
plt.imshow(flow_color)
plt.imshow(img, alpha=0.2, cmap='gray')
plt.title('Magnitude OF - ' + alg + ' - ' + type)
plt.xticks([])
plt.yticks([])
plt.show()
plt.savefig('magnitudeOP - ' + alg + ' - ' + type)
def show_df_from_tensor(tensor, path):
"""
tensor = torch.sqrt(torch.abs(tensor.cpu()))
temp = torch.zeros(1, H, H)
tensor_ = torch.cat([tensor, temp], 0).permute(1, 2, 0)
tensor_ = tensor_.numpy() * 255
cv2.imwrite(path, tensor_)
"""
tensor = tensor.cpu().permute(1, 2, 0).numpy()
#tensor = np.rollaxis(tensor, 2, 0)
flow_color = flow_vis.flow_to_color(tensor, convert_to_bgr=True)
cv2.imwrite(path, flow_color)
def get_optical_flow(pre, last):
#use flownet to get flow
with torch.no_grad():
first = torch.FloatTensor(np.ascontiguousarray(pre))
second = torch.FloatTensor(np.ascontiguousarray(last))
flow = estimate(first.permute(2, 0, 1), second.permute(2, 0, 1),
G.flow_network).detach()
#flow to rgb
flow = flow_vis.flow_to_color(flow.numpy().transpose(1, 2, 0),
convert_to_bgr=True)
return flow / 255.0
def run(imagefile1, imagefile2, save_file):
tensorFirst = torch.FloatTensor(
numpy.array(PIL.Image.open(imagefile1))[:, :, ::-1].transpose(
2, 0, 1).astype(numpy.float32) * (1.0 / 255.0))
tensorSecond = torch.FloatTensor(
numpy.array(PIL.Image.open(imagefile2))[:, :, ::-1].transpose(
2, 0, 1).astype(numpy.float32) * (1.0 / 255.0))
tensorOutput = estimate(tensorFirst, tensorSecond)
flow_color = flow_vis.flow_to_color(tensorOutput.numpy().transpose(
1, 2, 0),
convert_to_bgr=True)
cv2.imwrite(save_file, flow_color)
def get_optical_flow(pre, last):
'''
use liteflownet to get flow
input:pre is the last rgb picture;last is the current rgb picture
output:dense flow picture;need flow_vis to convert flow array to rgb picture
'''
with torch.no_grad():
first = torch.FloatTensor(np.ascontiguousarray(pre))
second = torch.FloatTensor(np.ascontiguousarray(last))
flow = estimate(first.permute(2, 0, 1), second.permute(2, 0, 1),
G.flow_network).detach()
#flow to rgb
flow = flow_vis.flow_to_color(flow.numpy().transpose(1, 2, 0),
convert_to_bgr=True)
return flow / 255.0
def gen_optflow(flow_filename, target_folder):
with open(str(flow_filename), 'r') as f:
header = np.fromfile(f, dtype=np.uint8, count=4)
size = np.fromfile(f, dtype=np.int32, count=2)
flow_uv = np.fromfile(f, dtype=np.float32) \
.reshape(config.cropped_height, config.cropped_width, 2)# .transpose(2,0,1)
# visuzlize using color model
flow_color = flow_vis.flow_to_color(flow_uv, convert_to_bgr=False)
cv2.imwrite(str(target_folder / (flow_filename.stem + '_color.png')),
flow_color)
# visualize with arrows
h, w = flow_uv.shape[:2]
x, y = np.meshgrid(np.arange(w), np.arange(h))
new_x = np.rint(x + flow_uv[:, :, 0]).astype(dtype=np.int64)
new_y = np.rint(y + flow_uv[:, :, 1]).astype(dtype=np.int64)
# new_x = new_x * ((new_x >= 0) & (new_x < w))
new_x = np.clip(new_x, 0, w)
# new_y = new_y * ((new_y >= 0) & (new_y < h))
new_y = np.clip(new_y, 0, h)
# blank image
coords_origin = np.array([x.flatten(), y.flatten()]).T
coords_new = np.array([new_x.flatten(), new_y.flatten()]).T
img = np.ones(flow_color.shape, np.uint8) * 255
for i in range(0, len(coords_origin), 1000):
cv2.arrowedLine(img, tuple(coords_origin[i]), tuple(coords_new[i]),
(255, 0, 0), 2)
cv2.imwrite(str(target_folder / (flow_filename.stem + '_pos_arr.png')),
img)
img_inv = np.ones(flow_color.shape, np.uint8) * 255
for i in range(0, len(coords_origin), 1000):
cv2.arrowedLine(img_inv, tuple(coords_new[i]),
tuple(coords_origin[i]), (255, 0, 0), 2)
cv2.imwrite(
str(target_folder / (flow_filename.stem + '_neg_arr.png')),
img_inv)
def infer_sequence(self, lr_data, device):
"""
Parameters:
:param lr_data: torch.FloatTensor in shape tchw
:param device: torch.device
:return hr_seq: uint8 np.ndarray in shape tchw
"""
# setup params
tot_frm, c, h, w = lr_data.size()
s = self.scale
# forward
hr_seq = []
lr_prev = torch.zeros(1, c, h, w, dtype=torch.float32).to(device)
hr_prev = torch.zeros(1, c, s * h, s * w,
dtype=torch.float32).to(device)
for i in range(tot_frm):
with torch.no_grad():
self.eval()
lr_curr = lr_data[i:i + 1, ...].to(device)
hr_curr, hr_flow, hr_warp = self.forward(
lr_curr, lr_prev, hr_prev)
lr_prev, hr_prev = lr_curr, hr_curr
hr_warp = hr_warp.squeeze(0).cpu().numpy() # chw|rgb|uint8
hr_frm = hr_warp.transpose(1, 2, 0) # hwc
flow_frm = hr_flow.squeeze(0).cpu().numpy() # chw|rgb|uint8
flow_uv = flow_frm.transpose(1, 2, 0) # hwc
flow_color = flow_vis.flow_to_color(flow_uv,
convert_to_bgr=False)
hr_seq.append(float32_to_uint8(hr_frm))
# hr_seq.append(float32_to_uint8(flow_color))
return np.stack(hr_seq)
def run_training(args):
print('---------- Perform Training ----------')
savedir = args.savepath
if not os.path.exists(savedir):
os.mkdir(savedir)
head_tail = os.path.split(args.dataset)
savedir = os.path.join(savedir, head_tail[1])
if not os.path.exists(savedir):
os.mkdir(savedir)
if not os.path.exists(os.path.join(savedir, "trained model")):
os.mkdir(os.path.join(savedir, "trained model"))
print('creating directory %s' %
(os.path.join(savedir, "trained model")))
if not os.path.exists(os.path.join(savedir, "saved training")):
os.mkdir(os.path.join(savedir, "saved training"))
print('creating directory %s' %
(os.path.join(savedir, "saved training")))
print('XField type: %s' % (args.type))
print('Dimension of input xfield: %s' % (args.dim))
#loading images
images, coordinates, all_pairs, h_res, w_res = load_imgs(args)
dims = args.dim
num_n = args.num_n # number of neighbors
min_ = np.min(coordinates)
max_ = np.max(coordinates)
print('\n ------- Creating the model -------')
# batch size is num_n + 1 (number of neighbors + target)
inputs = tf.placeholder(tf.float32, shape=[num_n + 1, 1, 1, len(dims)])
# Jacobian network
num_output = len(args.type) * 2
with tf.variable_scope("gen_flows"):
flows = Flow(inputs, h_res, w_res, num_output, args.nfg, min_, max_)
nparams_decoder = np.sum([
np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()
if v.name.startswith("gen_flows")
])
print('Number of learnable parameters (decoder): %d' % (nparams_decoder))
# learnt albedo
# The albedos are initialized with constant 1.0
if args.type == ['light', 'view', 'time']:
with tf.variable_scope("gen_flows"):
# For light-view-time interpolation, we consider num_views*num_times albedos
albedos = tf.Variable(tf.constant(
1.0, shape=[dims[1] * dims[2], h_res, w_res, 3]),
name='albedo')
index_albedo = tf.placeholder(tf.int32, shape=(1, ))
albedo = tf.gather(albedos, index_albedo, 0)
nparams = np.sum([
np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()
if v.name.startswith("gen_flows")
])
print(
'Number of learnable parameters (%d albedos with res %d x %d ): %d'
% (dims[1] * dims[2], h_res, w_res, nparams - nparams_decoder))
elif args.type == ['light']:
with tf.variable_scope("gen_flows"):
# For light interpolation, we consider just one albedo
albedo = tf.Variable(tf.constant(1.0, shape=[1, h_res, w_res, 3]),
name='albedo')
nparams = np.sum([
np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()
if v.name.startswith("gen_flows")
])
print(
'Number of learnable parameters (%d albedos with res %d x %d ): %d'
% (1, h_res, w_res, nparams - nparams_decoder))
else:
# For view and time interpolation, we do not train for albedo, we consider it as a constant non-learnable parameter
albedo = tf.constant(1.0, shape=[1, h_res, w_res, 3])
Neighbors = tf.placeholder(tf.float32, shape=[num_n, h_res, w_res, 3])
# soft blending
interpolated = Blending_train(inputs, Neighbors, flows, albedo, h_res,
w_res, args)
Reference = tf.placeholder(tf.float32, shape=[1, h_res, w_res, 3])
# L1 loss
loss = tf.reduce_mean((tf.abs(interpolated - Reference)))
gen_tvars = [
var for var in tf.trainable_variables()
if var.name.startswith("gen_flows")
]
learning_rate = tf.placeholder(tf.float32, shape=())
gen_optim = tf.train.AdamOptimizer(learning_rate)
gen_grads = gen_optim.compute_gradients(loss, var_list=gen_tvars)
gen_train = gen_optim.apply_gradients(gen_grads)
saver = tf.train.Saver(max_to_keep=1000)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
if args.load_pretrained:
ckpt = tf.train.get_checkpoint_state("%s\\trained model" % (savedir))
if ckpt:
print('\n loading pretrained model ' + ckpt.model_checkpoint_path)
saver.restore(sess, ckpt.model_checkpoint_path)
print('------------ Start Training ------------')
lr = args.lr
print('Starting learning rate with %0.4f' % (lr))
stop_l1_thr = 0.01
iter_end = 100000 # total number of iterations
indices = np.array([i for i in range(len(all_pairs))])
if len(indices
) < 500: # we considered around 500 iterations per each epoch
indices = np.repeat(indices, 500 // len(indices))
epoch_size = len(indices)
epoch_end = iter_end // epoch_size # total number of epochs
if args.type == ['light', 'view', 'time']:
st = time.time()
min_loss = 1000
l1_loss_t = 1
epoch = 0
while l1_loss_t > stop_l1_thr and epoch <= epoch_end:
l1_loss_t = 0
np.random.shuffle(indices)
for id in range(epoch_size):
pair = all_pairs[indices[id], ::]
input_coords = coordinates[pair[:num_n + 1], ::]
reference_img = images[pair[:1], ::]
Neighbors_img = images[pair[1:num_n + 1], ::]
_index = [pair[-1]]
_, l1loss = sess.run(
[gen_train, loss],
feed_dict={
inputs: input_coords,
Reference: reference_img,
Neighbors: Neighbors_img,
learning_rate: lr,
index_albedo: _index
})
l1_loss_t = l1_loss_t + l1loss
print(
'\r Epoch %3.0d Iteration %3.0d of %3.0d Cumulative L1 loss = %3.3f'
% (epoch, id + 1, epoch_size, l1_loss_t),
end=" ")
l1_loss_t = l1_loss_t / epoch_size
print(" elapsed time %3.1f m Averaged L1 loss = %3.5f " %
((time.time() - st) / 60, l1_loss_t))
if l1_loss_t < min_loss:
saver.save(sess, "%s\\trained model\\model.ckpt" % (savedir))
min_loss = l1_loss_t
center = np.prod(dims) // 2
cv2.imwrite("%s/saved training/reference.png" % (savedir),
np.uint8(images[center, ::] * 255))
pair = all_pairs[3 * center + 0, ::]
out_img, flows_out = sess.run(
[interpolated, flows],
feed_dict={
inputs: coordinates[pair[:num_n + 1], ::],
Neighbors: images[pair[1:num_n + 1], ::],
index_albedo: [pair[-1]]
})
out_img = np.minimum(np.maximum(out_img, 0.0), 1.0)
cv2.imwrite("%s/saved training/recons_light.png" % (savedir),
np.uint8(out_img[0, ::] * 255))
flow_color = flow_vis.flow_to_color(flows_out[0, :, :, 0:2],
convert_to_bgr=False)
cv2.imwrite("%s/saved training/flow_light.png" % (savedir),
np.uint8(flow_color))
flow_color = flow_vis.flow_to_color(flows_out[0, :, :, 2:4],
convert_to_bgr=False)
cv2.imwrite("%s/saved training/flow_view.png" % (savedir),
np.uint8(flow_color))
flow_color = flow_vis.flow_to_color(flows_out[0, :, :, 4:6],
convert_to_bgr=False)
cv2.imwrite("%s/saved training/flow_time.png" % (savedir),
np.uint8(flow_color))
pair = all_pairs[3 * center + 1, ::]
out_img = sess.run(interpolated,
feed_dict={
inputs: coordinates[pair[:num_n + 1], ::],
Neighbors: images[pair[1:num_n + 1], ::],
index_albedo: [pair[-1]]
})
out_img = np.minimum(np.maximum(out_img, 0.0), 1.0)
cv2.imwrite("%s/saved training/recons_view.png" % (savedir),
np.uint8(out_img[0, ::] * 255))
pair = all_pairs[3 * center + 2, ::]
out_img = sess.run(interpolated,
feed_dict={
inputs: coordinates[pair[:num_n + 1], ::],
Neighbors: images[pair[1:num_n + 1], ::],
index_albedo: [pair[-1]]
})
out_img = np.minimum(np.maximum(out_img, 0.0), 1.0)
cv2.imwrite("%s/saved training/recons_time.png" % (savedir),
np.uint8(out_img[0, ::] * 255))
epoch = epoch + 1
if epoch == epoch_end // 2:
lr = 0.00005
if args.type == ['view'] or args.type == ['time'
] or args.type == ['light']:
st = time.time()
img_mov = cv2.VideoWriter(
'%s/saved training/epoch_recons.mp4' % (savedir),
cv2.VideoWriter_fourcc(*'mp4v'), 10, (w_res, h_res))
flow_mov = cv2.VideoWriter(
'%s/saved training/epoch_flows.mp4' % (savedir),
cv2.VideoWriter_fourcc(*'mp4v'), 10, (w_res, h_res))
min_loss = 1000
l1_loss_t = 1
epoch = 0
while l1_loss_t > stop_l1_thr and epoch <= epoch_end:
l1_loss_t = 0
np.random.shuffle(indices)
for id in range(epoch_size):
pair = all_pairs[indices[id], ::]
input_coords = coordinates[pair[:num_n + 1], ::]
reference_img = images[pair[:1], ::]
Neighbors_img = images[pair[1:num_n + 1], ::]
_, l1loss = sess.run(
[gen_train, loss],
feed_dict={
inputs: input_coords,
Reference: reference_img,
Neighbors: Neighbors_img,
learning_rate: lr,
})
l1_loss_t = l1_loss_t + l1loss
print(
'\r Epoch %3.0d Iteration %3.0d of %3.0d Cumulative L1 loss = %3.3f'
% (epoch, id + 1, epoch_size, l1_loss_t),
end=" ")
l1_loss_t = l1_loss_t / epoch_size
print(" elapsed time %3.1f m Averaged L1 loss = %3.5f" %
((time.time() - st) / 60, l1_loss_t))
if l1_loss_t < min_loss:
saver.save(sess, "%s\\trained model\\model.ckpt" % (savedir))
min_loss = l1_loss_t
if args.type == ['light']:
albedo_out = np.minimum(np.maximum(sess.run(albedo), 0.0), 1.0)
cv2.imwrite("%s/saved training/albedo.png" % (savedir),
np.uint8(albedo_out[0, :, :, :] * 255))
center = np.prod(dims) // 2
cv2.imwrite("%s/saved training/reference.png" % (savedir),
np.uint8(images[center, ::] * 255))
pair = all_pairs[(len(all_pairs) // len(images)) * center, ::]
out_img, flows_out = sess.run(
[interpolated, flows],
feed_dict={
inputs: coordinates[pair[:num_n + 1], ::],
Neighbors: images[pair[1:num_n + 1], ::]
})
out_img = np.minimum(np.maximum(out_img, 0.0), 1.0)
cv2.imwrite("%s/saved training/recons.png" % (savedir),
np.uint8(out_img[0, ::] * 255))
flow_color = flow_vis.flow_to_color(flows_out[0, :, :, 0:2],
convert_to_bgr=False)
cv2.imwrite("%s/saved training/flow.png" % (savedir),
np.uint8(flow_color))
img_mov.write(np.uint8(out_img[0, ::] * 255))
flow_mov.write(np.uint8(flow_color))
epoch = epoch + 1
if epoch == epoch_end // 2:
lr = 0.00005
img_mov.release()
flow_mov.release()
from skimage.io import imread
from scipy import signal, ndimage
import numpy as np
import time
import scipy.io as sio
from matplotlib.pyplot import imshow, show, figure
import skimage.transform as tf
import IPython
import flow_vis
image1 = imread('frame1.png', as_gray=True)
image2 = imread('frame2.png', as_gray=True)
flow_gt = sio.loadmat('flow_gt.mat')['groundTruth']
flow_image_gt = flow_vis.flow_to_color(flow_gt)
#PARTA:Lukas_kanade
def lukas_kanade(I1, I2, window_size=5):
w = window_size // 2 # window_size is odd, all the pixels with offset in between [-w, w] are inside the window
I1 = I1 / 255. # normalize pixels
I2 = I2 / 255. # normalize pixels
# Define convolution kernels.
kernel_x = [[-1 / 4, 1 / 4], [-1 / 4, 1 / 4]]
kernel_y = [[-1 / 4, -1 / 4], [1 / 4, 1 / 4]]
kernel_t = [[1 / 4, 1 / 4], [1 / 4, 1 / 4]]
# Compute partial derivatives.
Ix = signal.convolve2d(I1 + I2, kernel_x, mode='same')
def main():
transform = torchvision.transforms.ColorJitter(brightness=(0.6, 1.0),
contrast=(0.6, 1.0))
train_loader = torch.utils.data.DataLoader(KITTI(mode="training",
transform=transform),
batch_size=4,
shuffle=True)
test_loader = torch.utils.data.DataLoader(KITTI(mode="testing"),
batch_size=4,
shuffle=False)
model = PWCDCNet()
device = torch.device("cuda")
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
wandb.watch(model)
for epoch in range(1, 300):
training_loss = 0.0
for i, data in enumerate(train_loader, 0):
optimizer.zero_grad()
left_image, right_image, target_flows, valids = data[0].to(
device), data[1].to(device), data[2].to(device), data[3].to(
device)
# Shapes:
## Left Image: (batch_size, 3, 320, 896)
## Right Image: (batch_size, 3, 320, 896)
## Target Flows: (batch_size, 2, 320, 896)
## Valids: (batch_size, 320, 896)
flow2, flow3, flow4, flow5, flow6 = model(
torch.cat((left_image, right_image), 1))
loss = ProbabilisticNLLLoss([flow2, flow3, flow4, flow5, flow6],
target_flows, valids)
loss.backward()
optimizer.step()
training_loss += loss.item()
print("[Epoch {}] Training loss: {:.5f}".format(epoch, training_loss))
wandb.log({"Epoch Number": epoch, "Training Loss": training_loss})
running_loss = 0.0
with torch.no_grad():
model.eval()
validation_loss = 0.0
rgb_images = []
predicted_flow_images = []
target_flow_images = []
for i, data in enumerate(test_loader, 0):
left_image, right_image, target_flows, valids = data[0].to(
device), data[1].to(device), data[2].to(
device), data[3].to(device)
flow2 = model(torch.cat((left_image, right_image), 1))
upsample = torch.nn.Upsample(size=(320, 896),
mode='bilinear',
align_corners=False)
upsampled_flow2 = upsample(flow2)
upsampled_flow2 = upsampled_flow2 * (320 // flow2.shape[2])
validation_loss += total_EPE_loss([upsampled_flow2],
target_flows, valids)
if i == 0:
for idx in range(4):
img_show = left_image[idx].permute(1, 2,
0).cpu().numpy()
img_show = img_show.clip(0, 1)
rgb_images.append(wandb.Image(img_show, grouping=3))
flowcolored_pred = flow_vis.flow_to_color(
upsampled_flow2[idx, :2, :, :].permute(
1, 2, 0).cpu().numpy())
predicted_flow_images.append(
wandb.Image(flowcolored_pred))
flowcolored_gt = flow_vis.flow_to_color(
(target_flows[idx].permute(
1, 2, 0)[:, :, :2]).cpu().numpy())
target_flow_images.append(wandb.Image(flowcolored_gt))
images_for_wandb = list(
chain.from_iterable(
zip(rgb_images, predicted_flow_images,
target_flow_images)))
wandb.log({'examples': images_for_wandb}, commit=False)
print("Validation loss: {:.5f}".format(validation_loss))
wandb.log({"Validation Loss": validation_loss})
model.train()
torch.save(model.state_dict(), './feb22/last_aug_300_06.pth')
else:
embedding[..., 0] -= np.arange(embedding.shape[0])[:, None]
embedding[..., 1] -= np.arange(embedding.shape[1])[None, :]
print(embedding.shape)
outpath = f"/groups/funke/home/wolfs2/local/data/lsl/evaluation/{expname_plus_setup}"
base_name = f"{i:02}_{bw}_{sf}_{idx}_{postfix}"
os.makedirs(outpath, exist_ok=True)
# imsave(
# f"{outpath}/{base_name}_0_embedding.png", embedding[..., :3])
# print(embedding.shape)
# exit()
if SAVEIMG:
rgb = flow_vis.flow_to_color(embedding[..., :2], convert_to_bgr=False)
imsave(
f"{outpath}/{base_name}_0_embedding.png", rgb)
seg = pred_zarr[f"inference/{expname_plus_setup}/pn_embedding/{idx}/inst_embedding"][:].transpose(1,2,0)
if postfix=="_full":
predicted_segmentation += 1
predicted_segmentation[gt_segmentation==0] = 0
# size filter
if sf > 0:
values, counts = np.unique(
predicted_segmentation, return_counts=True)
# could be vectorized!
def train(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
kwargs = {'num_workers': 0, 'pin_memory': False}
else:
kwargs = {}
if args.model_type == "rnn":
transformer = transformer_net.TransformerRNN(args.pad_type)
seq_size = 4
else:
transformer = transformer_net.TransformerNet(args.pad_type)
seq_size = 2
train_dataset = dataset.DAVISDataset(args.dataset,
seq_size=seq_size,
use_flow=args.flow)
train_loader = DataLoader(train_dataset, batch_size=1, **kwargs)
if args.model_type == "rnn":
transformer = transformer_net.TransformerRNN(args.pad_type)
else:
transformer = transformer_net.TransformerNet(args.pad_type)
model_path = args.init_model
print("=> Load from model file %s" % model_path)
transformer.load_state_dict(torch.load(model_path))
transformer.train()
if args.model_type == "rnn":
transformer.conv1 = transformer_net.ConvLayer(6,
32,
kernel_size=9,
stride=1,
pad_type=args.pad_type)
optimizer = torch.optim.Adam(transformer.parameters(), args.lr)
mse_loss = torch.nn.MSELoss()
l1_loss = torch.nn.SmoothL1Loss()
vgg = Vgg16()
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model, "vgg16.weight")))
vgg.eval()
if args.cuda:
transformer.cuda()
vgg.cuda()
mse_loss.cuda()
l1_loss.cuda()
style = utils.tensor_load_resize(args.style_image, args.style_size)
style = style.unsqueeze(0)
print("=> Style image size: " + str(style.size()))
print("=> Pixel OFB loss weight: %f" % args.time_strength)
style = utils.preprocess_batch(style)
if args.cuda: style = style.cuda()
utils.tensor_save_bgrimage(
style[0].detach(), os.path.join(args.save_model_dir,
'train_style.jpg'), args.cuda)
style = utils.subtract_imagenet_mean_batch(style)
features_style = vgg(style)
gram_style = [utils.gram_matrix(y).detach() for y in features_style]
for e in range(args.epochs):
train_loader.dataset.reset()
transformer.train()
transformer.cuda()
agg_content_loss = agg_style_loss = agg_pixelofb_loss = 0.
iters = 0
anormaly = False
for batch_id, (x, flow, conf) in enumerate(train_loader):
x, flow, conf = x[0], flow[0], conf[0]
iters += 1
optimizer.zero_grad()
x = utils.preprocess_batch(x) # (N, 3, 256, 256)
if args.cuda:
x = x.cuda()
flow = flow.cuda()
conf = conf.cuda()
y = transformer(x) # (N, 3, 256, 256)
xc = center_crop(x.detach(), y.size(2), y.size(3))
vgg_y = utils.subtract_imagenet_mean_batch(y)
vgg_x = utils.subtract_imagenet_mean_batch(xc)
features_y = vgg(vgg_y)
features_xc = vgg(vgg_x)
#content target
f_xc_c = features_xc[2].detach()
# content
f_c = features_y[2]
#content_feature_target = center_crop(f_xc_c, f_c.size(2), f_c.size(3))
content_loss = args.content_weight * mse_loss(f_c, f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_s = gram_style[m]
gram_y = utils.gram_matrix(features_y[m])
batch_style_loss = 0
for n in range(gram_y.shape[0]):
batch_style_loss += args.style_weight * mse_loss(
gram_y[n], gram_s[0])
style_loss += batch_style_loss / gram_y.shape[0]
warped_y, warped_y_mask = warp(y[1:], flow)
warped_y = warped_y.detach()
warped_y_mask *= conf
pixel_ofb_loss = args.time_strength * weighted_mse(
y[:-1], warped_y, warped_y_mask)
total_loss = content_loss + style_loss + pixel_ofb_loss
total_loss.backward()
optimizer.step()
if (batch_id + 1) % 100 == 0:
prefix = args.save_model_dir + "/"
idx = (batch_id + 1) // 100
flow_image = flow_to_color(
flow[0].detach().cpu().numpy().transpose(1, 2, 0))
utils.save_image(prefix + "forward_flow_%d.png" % idx,
flow_image)
warped_x, warped_x_mask = warp(x[1:], flow)
warped_x = warped_x.detach()
warped_x_mask *= conf
for i in range(2):
utils.tensor_save_bgrimage(
y.data[i], prefix + "out_%d-%d.png" % (idx, i),
args.cuda)
utils.tensor_save_bgrimage(
x.data[i], prefix + "in_%d-%d.png" % (idx, i),
args.cuda)
if i < warped_y.shape[0]:
utils.tensor_save_bgrimage(
warped_y.data[i],
prefix + "wout_%d-%d.png" % (idx, i), args.cuda)
utils.tensor_save_bgrimage(
warped_x.data[i],
prefix + "win_%d-%d.png" % (idx, i), args.cuda)
utils.tensor_save_image(
prefix + "conf_%d-%d.png" % (idx, i),
warped_x_mask.data[i])
agg_content_loss += content_loss.data
agg_style_loss += style_loss.data
agg_pixelofb_loss += pixel_ofb_loss.data
agg_total = agg_content_loss + agg_style_loss + agg_pixelofb_loss
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.6f}\tstyle: {:.6f}\tpixel ofb: {:.6f}\ttotal: {:.6f}".format(
time.ctime(), e + 1, batch_id + 1, len(train_loader),
agg_content_loss / iters, agg_style_loss / iters,
agg_pixelofb_loss / iters, agg_total / iters)
print(mesg)
agg_content_loss = agg_style_loss = agg_pixelofb_loss = 0.0
iters = 0
# save model
transformer.eval()
transformer.cpu()
save_model_filename = "epoch_" + str(e) + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir,
save_model_filename)
torch.save(transformer.state_dict(), save_model_path)
print("\nDone, trained model saved at", save_model_path)
_, filename, _ = fileparts(image_path)
frame2 = load_image(image_path)
next = rgb2gray(frame2)
# save_image(next, '/home/xinshuo/aa.jpg')
# zxc
# cv2.imshow('asd', next)
# k = cv2.waitKey(30)
# zxc
# ret, frame2 = cap.read()
# next = cv2.cvtColor(frame2,cv2.COLOR_BGR2GRAY)
flow_uv = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
rgb = flow_to_color(flow_uv, convert_to_bgr=False)
# 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)
# cv2.imshow('frame2', rgb)
# k = cv2.waitKey(30) & 0xff
# if k == 27:
# break
# elif k == ord('s'):
save_path_tmp = os.path.join(save_dir, filename+'.jpg')
save_image(rgb, save_path=save_path_tmp)
# cv2.imwrite('opticalfb.png', frame2)
# cv2.imwrite('opticalhsv.png', rgb)
def visualiseFlow(img_pair, flow_gt, flow_pred, flow_prob, ds, idx=0):
dim = flow_pred[0].size(1)
H, W = img_pair[0].size()[2:]
with torch.no_grad():
full_vis_list = []
for idx in range(img_pair[0].size(0)):
raw_img0 = recoverImage(img_pair[0][idx].data)
raw_img1 = recoverImage(img_pair[1][idx].data)
raw_img0 = raw_img0 * 255.0
raw_img1 = raw_img1 * 255.0
for l in range(len(flow_pred)):
# Image
# for i in range(3):
# raw_img0[i, :, :] = raw_img0[i, :, :] * 255.0 / torch.max(raw_img0[i, :, :])
# raw_img1[i, :, :] = raw_img1[i, :, :] * 255.0 / torch.max(raw_img1[i, :, :])
vis_list = [raw_img0, raw_img1]
# Ground truth flow
gt_flow, valid_mask = downsample_flow(flow_gt, 1. / 2 ** (ds - l))
gt_flow = F.interpolate(gt_flow, (H, W), mode="nearest", recompute_scale_factor=True)[idx]
valid_mask = F.interpolate(valid_mask, (H, W), mode="nearest", recompute_scale_factor=True)[idx]
max_mag1 = torch.max(torch.norm(gt_flow, 2, 0))
# predicted flow
pred_flow = flow_pred[l]
pred_flow = F.interpolate(pred_flow, (H, W), mode='nearest', recompute_scale_factor=True)[idx]
max_mag2 = torch.max(torch.norm(pred_flow, 2, 0))
max_mag = max(float(max_mag1), float(max_mag2))
# print("GT Flow", gt_flow.size())
# print("GT Flow", gt_flow.shape)
gt_flow_np = gt_flow.detach().cpu().permute(1, 2, 0).numpy()
pred_flow_np = pred_flow.cpu().permute(1, 2, 0).numpy()
gt_flow_vis = flow_vis.flow_to_color(gt_flow_np, convert_to_bgr=False)
pred_flow_vis = flow_vis.flow_to_color(pred_flow_np, convert_to_bgr=False)
vis_list.append(torch.from_numpy(gt_flow_vis).permute(2, 0, 1).cuda())
vis_list.append(torch.from_numpy(pred_flow_vis).permute(2, 0, 1).cuda())
# epe error visualization
epe_error = torch.norm(pred_flow - gt_flow, 2, 0, keepdim=False) * valid_mask[0, :, :]
normalizer = max(torch.max(epe_error), 1)
epe_error = 1 - epe_error / normalizer
vis_list.append(visualiseHeatmap(epe_error))
# confidence map visualization
prob = flow_prob[l].data
prob = probGather(prob, normalize=True)
if prob.size(2) != H or prob.size(3) != W:
prob = F.interpolate(prob, (H, W), mode='nearest', recompute_scale_factor=True)
vis_list.append(visualiseHeatmap(prob[idx].squeeze(), cv2.COLORMAP_BONE))
vis = torch.cat(vis_list, dim=2)
if l == 0:
ms_vis = vis
else:
ms_vis = torch.cat([ms_vis, vis], dim=1)
full_vis_list.append(ms_vis.unsqueeze(0))
return full_vis_list
from skimage.io import imread
from scipy import signal, ndimage
import numpy as np
import time
import scipy.io as sio
from matplotlib.pyplot import imshow, show, figure
import skimage.transform as tf
import IPython
import flow_vis
image1 = imread('frame1.png', as_gray=True)
image2 = imread('frame2.png', as_gray=True)
flow_gt = sio.loadmat('flow_gt.mat')['groundTruth']
flow_image_gt = flow_vis.flow_to_color(flow_gt)
#PARTA:Lukas_kanade
def lukas_kanade(I1, I2, window_size=5):
w = window_size // 2 # window_size is odd, all the pixels with offset in between [-w, w] are inside the window
I1 = I1 / 255. # normalize pixels
I2 = I2 / 255. # normalize pixels
# Define convolution kernels.
kernel_x = [[-1 / 4, 1 / 4], [-1 / 4, 1 / 4]]
kernel_y = [[-1 / 4, -1 / 4], [1 / 4, 1 / 4]]
kernel_t = [[1 / 4, 1 / 4], [1 / 4, 1 / 4]]
# Compute partial derivatives.
Ix = signal.convolve2d(I1 + I2, kernel_x, mode='same')
align_corners=False)
img2 = F.interpolate(img2,
size=input_size,
mode='bilinear',
align_corners=False)
else:
input_size = orig_size
input_t = torch.cat([img1, img2], 1).cuda()
output = model(input_t).data
flow = div_flow * F.interpolate(
output, size=input_size, mode='bilinear', align_corners=False)
if input_size != orig_size:
scale_h = orig_size[0] / input_size[0]
scale_w = orig_size[1] / input_size[1]
flow = F.interpolate(flow,
size=orig_size,
mode='bilinear',
align_corners=False)
flow[:, 0, :, :] *= scale_w
flow[:, 1, :, :] *= scale_h
flow = flow[0].cpu().permute(1, 2, 0).numpy()
flow_color = flow_to_color(flow, convert_to_bgr=True)
cv2.imwrite('./data/flow.png', flow_color)
# Sets image saturation to maximum
mask[..., 1] = 255
while(cap.isOpened()):
# ret = a boolean return value from getting the frame, frame = the current frame being projected in the video
ret, frame = cap.read()
# Opens a new window and displays the input frame
cv.imshow("input", frame)
# Converts each frame to grayscale - we previously only converted the first frame to grayscale
gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
# Calculates dense optical flow by Farneback method
# https://docs.opencv.org/3.0-beta/modules/video/doc/motion_analysis_and_object_tracking.html#calcopticalflowfarneback
flow = cv.calcOpticalFlowFarneback(prev_gray, gray, None, 0.5, 3, 15, 3, 5, 1.2, 0)
# Apply the coloring (for OpenCV, set convert_to_bgr=True)
flow_color = flow_vis.flow_to_color(flow, convert_to_bgr=False)
# Computes the magnitude and angle of the 2D vectors
magnitude, angle = cv.cartToPolar(flow[..., 0], flow[..., 1])
# Sets image hue according to the optical flow direction
mask[..., 0] = angle * 180 / np.pi / 2
# Sets image value according to the optical flow magnitude (normalized)
mask[..., 2] = cv.normalize(magnitude, None, 0, 255, cv.NORM_MINMAX)
# Converts HSV to RGB (BGR) color representation
rgb = cv.cvtColor(mask, cv.COLOR_HSV2BGR)
# Open a new window and displays the output frame
dense_flow = cv.addWeighted(frame, 1,rgb, 2, 0)
cv.imshow('flow', draw_flow(gray, flow))
def main():
data = get_data()
# ==========================================================================
# Sidebar
# ==========================================================================
select_flux_scaling = st.sidebar.selectbox(
"Flux scaling",
["Linear", "Log"]
)
select_band = st.sidebar.selectbox(
"Band",
[i for i in range(data["flux"].shape[-1])]
)
select_n = st.sidebar.select_slider(
"nth Closest Source",
[str(i) for i in range(1, data["cm"].shape[-1]+1)]
)
select_src = st.sidebar.select_slider(
"Select Deblended Source",
[str(i) for i in range(1, len(data["srcs"])+1)]
)
select_blend_alpha = st.sidebar.slider(
"Select Source Mask Alpha",
min_value=0.0,
max_value=1.0,
step=0.1,
value=0.5,
)
# ==========================================================================
st.write("""
# Morpheus-Deblend View
""")
# ==========================================================================
# Input Flux Display
# ==========================================================================
flux = data["flux"][...,select_band]
if select_flux_scaling=="Linear":
flux_img = scale_flux(flux)
elif select_flux_scaling=="Log":
flux_img = rescale(
np.log10(
flux - flux.min() + 1e-3
)
)
st.write("### Input Image")
st.image(flux_img,use_column_width=True)
# ==========================================================================
# ==========================================================================
# Center of Mass
# ==========================================================================
com = data["com"]
suppressed = data["suppressed"]
st.write("### Identified Sources")
col_com_raw, col_com_suppressed = st.beta_columns(2)
col_com_raw.write("Center of Mass Raw")
col_com_raw.image(colorize_array(com[...,0]), use_column_width=True)
col_com_suppressed.write("Center of Mass Suppressed")
col_com_suppressed.image(suppressed, use_column_width=True)
# ==========================================================================
# ==========================================================================
# Claim Vectors/Claim Maps
# ==========================================================================
n = int(select_n) - 1
cm = data["cm"][...,select_band, n]
cv = flow_vis.flow_to_color(
data["cv"][:, :, n, [1,0]]
)
nth_string = select_n + resolve_n_post(select_n)
col_cv, col_cm = st.beta_columns(2)
col_cv.write(f"Claim Vector for {nth_string} Source")
col_cv.image(cv, use_column_width=True)
col_cm.write(f"Claim Map for {nth_string} Source")
col_cm.image(colorize_array(cm), use_column_width=True)
# ==========================================================================
# ==========================================================================
# Deblended Sources
# ==========================================================================
n_src = int(select_src) - 1
src = data["srcs"][n_src][:, :, select_band]
if select_flux_scaling=="Linear":
src_img = scale_flux(src)
elif select_flux_scaling=="Log":
src_img = rescale(
np.log10(
src - src.min() + 1e-3
)
)
mask = src != 0
mask_img = np.zeros(list(mask.shape) + [4], dtype=np.float32)
mask_img[mask, :] = np.array([1,0,0,1.0])
#mask_img[~mask, :] = np.array([0,0,0,0])
rgba_flux = (np.dstack(
[flux_img, flux_img, flux_img, np.ones_like(flux_img)]
) * 255).astype(np.uint8)
img1 = Image.fromarray((mask_img * 255).astype(np.uint8))
img2 = Image.fromarray(rgba_flux)
blended = Image.blend(img1, img2, select_blend_alpha)
st.write("Source {select_src}")
col_mask, col_src_flux = st.beta_columns(2)
col_mask.image(np.array(blended), use_column_width=True)
col_src_flux.image(src_img, use_column_width=True)
def vis_flow(self, flow):
return flow_vis.flow_to_color(flow, convert_to_bgr=True)
def read_flow(filename):
TAG_FLOAT = 202021.25
with open(filename, 'rb') as f:
flo_number = np.fromfile(f, np.float32, count=1)[0]
assert flo_number == TAG_FLOAT, 'Flow number %r incorrect. Invalid .flo file' % flo_number
w = np.fromfile(f, np.int32, count=1)[0]
h = np.fromfile(f, np.int32, count=1)[0]
data = np.fromfile(f, np.float32, count=2 * w * h)
return np.resize(data, (int(h), int(w), 2))
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('input', type=str, help='input .flo files')
parser.add_argument('output', type=str, help='output .png files')
args = parser.parse_args()
input = args.input
output = args.output
images = glob.glob(input)
images.sort()
for i, img_path in enumerate(images):
name = os.path.basename(img_path)
img = read_flow(img_path)
img = flow_vis.flow_to_color(img, convert_to_bgr=True)
cv2.imwrite(output + f'/{i:06d}.png', img)
