def validate(val_loader, model, epoch, output_writers):
global args
batch_time = AverageMeter()
flow2_EPEs = AverageMeter()
# switch to evaluate mode
model.eval()
# end = time.time();
runtime_count = 0;
sum_runtime = 0;
end = time.time()
for i, (input, target) in enumerate(val_loader):
target = target.to(device)
# concatnate the tensor
input = torch.cat(input,1).to(device)
# compute output
start_time = time.time()
if args.graymodel:
output = model(torch.cat([input[:,0:1,:,:], input[:,3:4,:,:]],1))
else:
output = model(input)
runtime = (time.time()-start_time)
#sum_runtime += runtime
# if runtime_count == 0:
# sum_runtime = 0
# else: print ('AvgRunTime: ', sum_runtime/runtime_count)
runtime_count += 1
sum_runtime += runtime
# print (runtime)
flow2_EPE = args.div_flow*realEPE(output, target, sparse=args.sparse)
# record EPE
flow2_EPEs.update(flow2_EPE.item(), target.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i < len(output_writers): # log first output of first batches
if epoch == 0:
mean_values = torch.tensor([0.45,0.432,0.411], dtype=input.dtype).view(3,1,1)
output_writers[i].add_image('GroundTruth', flow2rgb(args.div_flow * target[0], max_value=10), 0)
output_writers[i].add_image('Inputs', (input[0,:3].cpu() + mean_values).clamp(0,1), 0)
output_writers[i].add_image('Inputs', (input[0,3:].cpu() + mean_values).clamp(0,1), 1)
output_writers[i].add_image('FlowNet Outputs', flow2rgb(args.div_flow * output[0], max_value=10), epoch)
if i % args.print_freq == 0:
print('Test: [{0}/{1}]\t Time {2}\t EPE {3}'
.format(i, len(val_loader), batch_time, flow2_EPEs))
print(' * EPE {:.3f}'.format(flow2_EPEs.avg))
print ('sum_runtime', sum_runtime/runtime_count, runtime_count)
return flow2_EPEs.avg
def validate(val_loader, model, epoch, output_writers):
global args
batch_time = AverageMeter()
flow2_EPEs = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for i, (input, target) in enumerate(val_loader):
target = target.to(device)
input = torch.cat(input, 1).to(device)
# compute output
output = model(input)
flow2_EPE = args.div_flow * realEPE(output, target, sparse=args.sparse)
# record EPE
flow2_EPEs.update(flow2_EPE.item(), target.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i < len(output_writers): # log first output of first batches
if epoch == 0:
mean_values = torch.tensor([0.45, 0.432, 0.411],
dtype=input.dtype).view(3, 1, 1)
output_writers[i].add_image(
'GroundTruth',
flow2rgb(args.div_flow * target[0], max_value=10), 0)
output_writers[i].add_image(
'Inputs', (input[0, :3].cpu() + mean_values).clamp(0,
1), 0)
output_writers[i].add_image(
'Inputs', (input[0, 3:].cpu() + mean_values).clamp(0,
1), 1)
output_writers[i].add_image(
'FlowNet Outputs',
flow2rgb(args.div_flow * output[0], max_value=10), epoch)
if i % args.print_freq == 0:
print('Test: [{0}/{1}]\t Time {2}\t EPE {3}'.format(
i, len(val_loader), batch_time, flow2_EPEs))
print(' * EPE {:.3f}'.format(flow2_EPEs.avg))
return flow2_EPEs.avg
def run(img1, img2, bidirectional=False, upsampling='bilinear', div_flow=20, max_flow=None, to_rgb=False):
if model is None:
setup()
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1])
])
img1 = input_transform(np.array(img1))
img2 = input_transform(np.array(img2))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if bidirectional:
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
flow = model(input_var)
if upsampling is not None:
assert upsampling in ['nearest', 'bilinear'], \
'Upsampling mode {} not recognized'.format(upsampling)
flow = torch.nn.functional.interpolate(
flow,
size=img1.size()[-2:],
mode=upsampling,
align_corners=False)
if to_rgb:
rgb_flow = flow2rgb(div_flow * flow[0], max_value=max_flow)
rgb_flow = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)
return rgb_flow
else:
flow = (div_flow * flow[0]).detach().cpu().numpy().transpose(1,2,0)
return flow
def validate(val_loader, model, epoch, output_writers):
global args
batch_time = AverageMeter()
flow2_EPEs = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for i, data in enumerate(val_loader):
if args.data_loader == "torch":
(input, target) = data
if args.data_loader == "dali":
input = [
data[0]["images"][:, 0:3, :, :],
data[0]["images"][:, 3:6, :, :],
]
target = data[0]["flow"]
if args.show_val_images:
for k in range(len(input[0].cpu())):
f, axarr = plt.subplots(2, 2)
axarr[0,
0].imshow(np.moveaxis(np.array(input[0].cpu()[k]), 0, 2))
axarr[0,
1].imshow(np.moveaxis(np.array(input[1].cpu()[k]), 0, 2))
axarr[1, 0].imshow(
np.moveaxis(
flow2rgb(args.div_flow * np.squeeze(target.cpu()[k]),
max_value=10),
0,
2,
))
plt.show()
target = target.to(device)
input = torch.cat(input, 1).to(device)
# compute output
output = model(input) # [0], input[1])
flow2_EPE = args.div_flow * realEPE(output, target, sparse=args.sparse)
# record EPE
flow2_EPEs.update(flow2_EPE.item(), target.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i < len(output_writers): # log first output of first batches
if epoch == 0:
# input = torch.cat(input, 1).to(device)
mean_values = torch.tensor([0.45, 0.432, 0.411],
dtype=input.dtype).view(3, 1, 1)
output_writers[i].add_image(
"GroundTruth",
flow2rgb(args.div_flow * target[0], max_value=10), 0)
output_writers[i].add_image(
"Inputs", (input[0, :3].cpu() + mean_values).clamp(0,
1), 0)
output_writers[i].add_image(
"Inputs", (input[0, 3:].cpu() + mean_values).clamp(0,
1), 1)
output_writers[i].add_image(
"FlowNet Outputs",
flow2rgb(args.div_flow * output[0], max_value=10),
epoch,
)
if i % args.print_freq == 0:
print("Test: [{0}/{1}]\t Time {2}\t EPE {3}".format(
i, len(val_loader), batch_time, flow2_EPEs))
print(" * EPE {:.3f}".format(flow2_EPEs.avg))
return flow2_EPEs.avg
def train(train_loader, model, optimizer, epoch, train_writer):
global n_iter, args
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
flow2_EPEs = AverageMeter()
if args.data_loader == "torch":
epoch_size = (len(train_loader) if args.epoch_size == 0 else min(
len(train_loader), args.epoch_size))
if args.data_loader == "dali":
epoch_size = (9999 if args.epoch_size == 0 else min(
len(train_loader), args.epoch_size))
# switch to train mode
model.train()
end = time.time()
for i, data in enumerate(train_loader):
if args.data_loader == "torch":
(input, target) = data
if args.data_loader == "dali":
input = [
data[0]["images"][:, 0:3, :, :],
data[0]["images"][:, 3:6, :, :],
]
target = data[0]["flow"]
if args.show_train_images:
for k in range(len(input[0].cpu())):
f, axarr = plt.subplots(2, 2)
axarr[0,
0].imshow(np.moveaxis(np.array(input[0].cpu()[k]), 0, 2))
axarr[0,
1].imshow(np.moveaxis(np.array(input[1].cpu()[k]), 0, 2))
axarr[1, 0].imshow(
np.moveaxis(
flow2rgb(args.div_flow * np.squeeze(target.cpu()[k]),
max_value=10),
0,
2,
))
plt.show()
# measure data loading time
data_time.update(time.time() - end)
target = target.to(device)
input = torch.cat(input, 1).to(device)
# compute output
output = model(input) # [0], input[1])
if args.sparse:
# Since Target pooling is not very precise when sparse,
# take the highest resolution prediction and upsample it instead of downsampling target
h, w = target.size()[-2:]
output = [F.interpolate(output[0], (h, w)), *output[1:]]
loss = multiscaleEPE(output,
target,
weights=args.multiscale_weights,
sparse=args.sparse)
flow2_EPE = args.div_flow * realEPE(
output, target,
sparse=args.sparse) # output[0] if using multi-scale loss
# record loss and EPE
losses.update(loss.item(), target.size(0))
train_writer.add_scalar("train_loss", loss.item(), n_iter)
flow2_EPEs.update(flow2_EPE.item(), target.size(0))
# compute gradient and do optimization step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print(
"Epoch: [{0}][{1}/{2}]\t Time {3}\t Data {4}\t Loss {5}\t EPE {6}"
.format(epoch, i, epoch_size, batch_time, data_time, losses,
flow2_EPEs))
n_iter += 1
if i >= epoch_size:
break
return losses.avg, flow2_EPEs.avg
def validate(test_loader, model, epoch, output_writers):
global args, image_resize, sp_threshold
d_label = h5py.File(gt_file, 'r')
gt_temp = np.float32(d_label['davis']['left']['flow_dist'])
gt_ts_temp = np.float64(d_label['davis']['left']['flow_dist_ts'])
d_label = None
d_set = h5py.File(testfile, 'r')
gray_image = d_set['davis']['left']['image_raw']
batch_time = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
batch_size_v = 4
sp_threshold = 0.5
AEE_sum = 0.
AEE_sum_sum = 0.
AEE_sum_gt = 0.
AEE_sum_sum_gt = 0.
percent_AEE_sum = 0.
iters = 0.
scale = 1
for i, data in enumerate(test_loader, 0):
former_inputs_on, former_inputs_off, latter_inputs_on, latter_inputs_off, st_time, ed_time = data
if torch.sum(former_inputs_on + former_inputs_off) > 0:
input_representation = torch.zeros(
former_inputs_on.size(0), batch_size_v, image_resize,
image_resize, former_inputs_on.size(3)).float()
for b in range(batch_size_v):
if b == 0:
input_representation[:, 0, :, :, :] = former_inputs_on
elif b == 1:
input_representation[:, 1, :, :, :] = former_inputs_off
elif b == 2:
input_representation[:, 2, :, :, :] = latter_inputs_on
elif b == 3:
input_representation[:, 3, :, :, :] = latter_inputs_off
# compute output
input_representation = input_representation.to(device)
output = model(input_representation.type(torch.cuda.FloatTensor),
image_resize, sp_threshold)
# pred_flow = output
pred_flow = np.zeros((image_resize, image_resize, 2))
output_temp = output.cpu()
pred_flow[:, :, 0] = cv2.resize(np.array(output_temp[0, 0, :, :]),
(image_resize, image_resize),
interpolation=cv2.INTER_LINEAR)
pred_flow[:, :, 1] = cv2.resize(np.array(output_temp[0, 1, :, :]),
(image_resize, image_resize),
interpolation=cv2.INTER_LINEAR)
U_gt_all = np.array(gt_temp[:, 0, :, :])
V_gt_all = np.array(gt_temp[:, 1, :, :])
U_gt, V_gt = estimate_corresponding_gt_flow(
U_gt_all, V_gt_all, gt_ts_temp, np.array(st_time),
np.array(ed_time))
gt_flow = np.stack((U_gt, V_gt), axis=2)
# ----------- Visualization
if epoch < 0:
mask_temp = former_inputs_on + former_inputs_off + latter_inputs_on + latter_inputs_off
mask_temp = torch.sum(torch.sum(mask_temp, 0), 2)
mask_temp_np = np.squeeze(np.array(mask_temp)) > 0
spike_image = mask_temp
spike_image[spike_image > 0] = 255
if args.render:
cv2.imshow('Spike Image',
np.array(spike_image, dtype=np.uint8))
gray = cv2.resize(gray_image[i],
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
if args.render:
cv2.imshow('Gray Image',
cv2.cvtColor(gray, cv2.COLOR_BGR2RGB))
out_temp = np.array(output_temp.cpu().detach())
x_flow = cv2.resize(
np.array(out_temp[0, 0, :, :]),
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
y_flow = cv2.resize(
np.array(out_temp[0, 1, :, :]),
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
flow_rgb = flow_viz_np(x_flow, y_flow)
if args.render:
cv2.imshow('Predicted Flow Output',
cv2.cvtColor(flow_rgb, cv2.COLOR_BGR2RGB))
gt_flow_x = cv2.resize(
gt_flow[:, :,
0], (scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
gt_flow_y = cv2.resize(
gt_flow[:, :,
1], (scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
gt_flow_large = flow_viz_np(gt_flow_x, gt_flow_y)
if args.render:
cv2.imshow('GT Flow',
cv2.cvtColor(gt_flow_large, cv2.COLOR_BGR2RGB))
masked_x_flow = cv2.resize(
np.array(out_temp[0, 0, :, :] * mask_temp_np),
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
masked_y_flow = cv2.resize(
np.array(out_temp[0, 1, :, :] * mask_temp_np),
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
flow_rgb_masked = flow_viz_np(masked_x_flow, masked_y_flow)
if args.render:
cv2.imshow(
'Masked Predicted Flow',
cv2.cvtColor(flow_rgb_masked, cv2.COLOR_BGR2RGB))
gt_flow_cropped = gt_flow[2:-2, 45:-45]
gt_flow_masked_x = cv2.resize(
gt_flow_cropped[:, :, 0] * mask_temp_np,
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
gt_flow_masked_y = cv2.resize(
gt_flow_cropped[:, :, 1] * mask_temp_np,
(scale * image_resize, scale * image_resize),
interpolation=cv2.INTER_LINEAR)
gt_masked_flow = flow_viz_np(gt_flow_masked_x,
gt_flow_masked_y)
if args.render:
cv2.imshow('GT Masked Flow',
cv2.cvtColor(gt_masked_flow, cv2.COLOR_BGR2RGB))
cv2.waitKey(1)
image_size = pred_flow.shape
full_size = gt_flow.shape
xsize = full_size[1]
ysize = full_size[0]
xcrop = image_size[1]
ycrop = image_size[0]
xoff = (xsize - xcrop) // 2
yoff = (ysize - ycrop) // 2
gt_flow = gt_flow[yoff:-yoff, xoff:-xoff, :]
AEE, percent_AEE, n_points, AEE_sum_temp, AEE_gt, AEE_sum_temp_gt = flow_error_dense(
gt_flow,
pred_flow,
(torch.sum(torch.sum(torch.sum(input_representation, dim=0),
dim=0),
dim=2)).cpu(),
is_car=False)
AEE_sum = AEE_sum + args.div_flow * AEE
AEE_sum_sum = AEE_sum_sum + AEE_sum_temp
AEE_sum_gt = AEE_sum_gt + args.div_flow * AEE_gt
AEE_sum_sum_gt = AEE_sum_sum_gt + AEE_sum_temp_gt
percent_AEE_sum += percent_AEE
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i < len(output_writers): # log first output of first batches
# if epoch == 0:
# mean_values = torch.tensor([0.411,0.432,0.45], dtype=input_representation.dtype).view(3,1,1)
output_writers[i].add_image(
'FlowNet Outputs',
flow2rgb(args.div_flow * output[0], max_value=10), epoch)
iters += 1
print('-------------------------------------------------------')
print(
'Mean AEE: {:.2f}, sum AEE: {:.2f}, Mean AEE_gt: {:.2f}, sum AEE_gt: {:.2f}, mean %AEE: {:.2f}, # pts: {:.2f}'
.format(AEE_sum / iters, AEE_sum_sum / iters, AEE_sum_gt / iters,
AEE_sum_sum_gt / iters, percent_AEE_sum / iters, n_points))
print('-------------------------------------------------------')
gt_temp = None
return AEE_sum / iters
def main():
global args, save_path
args = parser.parse_args()
if args.output_value == 'both':
output_string = "raw output and RGB visualization"
elif args.output_value == 'raw':
output_string = "raw output"
elif args.output_value == 'vis':
output_string = "RGB visualization"
print("=> will save " + output_string)
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir / 'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*0.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '1.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
# create model
network_data = torch.load(Path(args.pretrained),
map_location=torch.device('cpu'))
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
print("load model success!")
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
input_var = torch.cat([img1, img2]).unsqueeze(0)
print("read image success!")
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
print("prepare to compute flow!")
# compute output
output = model(input_var)
print("flow computing success!")
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = save_path / '{}{}'.format(img1_file.namebase[:-1],
suffix)
if args.output_value in ['vis', 'both']:
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
if args.output_value in ['raw', 'both']:
# Make the flow map a HxWx3 array as in .flo files
to_save = (args.div_flow *
flow_output).cpu().numpy().transpose(1, 2, 0)
np.save(filename + '.npy', to_save)
input_var = input_var.cuda()
# compute output
output = model(input_var)
outs.append(np.array(output.cpu().detach()))
# Upsample
output = F.interpolate(output,
size=img1.size()[-2:],
mode='bilinear',
align_corners=False)
# Convert to an RGBim1.
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
rgb_flow = flow2rgb(div_flow * flow_output, max_value=max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
combined = np.zeros([nY, nX * 2, nC], dtype=np.uint8)
combined[:, :nX, :] = mp4_ims[aa][:, :, [2, 1,
0]] # Invert back for BGR
combined[:, nX:, :] = to_save
export_ims.append(combined) # Save for list
# Write out the MP4
fn_out = os.path.join('/'.join(mp4_fn.split('/')[:-1]),
'flowflownet_' + mp4_fn.split('/')[-1])
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(fn_out, fourcc, 30.0, (int(2 * nX), int(nY)))
def generate_flow(self):
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
# create model
network_data = torch.load(self.model_dir, map_location='cpu')
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).cuda()
model.eval()
cudnn.benchmark = True
# if 'div_flow' in network_data.keys():
# args.div_flow = network_data['div_flow']
mp4_list = glob.glob(os.path.join(self.args.input_dir, "*.mp4"))
print(self.args.input_dir)
print(mp4_list)
for mp4_fn in mp4_list:
print(mp4_fn)
mp4_ims = mp4_load(mp4_fn)
n_im = len(mp4_ims) # Number of images
export_ims = []
[nY, nX, nC] = mp4_ims[0].shape
outs = []
#frame_ints = list(range(n_im))
for ii in tqdm(range(n_im - args.window)):
aa = ii
bb = ii + args.window
img1 = input_transform(mp4_ims[aa])
img2 = input_transform(mp4_ims[bb])
input_var = torch.cat([img1, img2]).unsqueeze(0)
# if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.cuda()
# compute output
output = model(input_var)
outs.append(np.array(output.cpu().detach()))
# Upsample
output = F.interpolate(output,
size=img1.size()[-2:],
mode='bilinear',
align_corners=False)
# Convert to an RGBim1.
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(
1, 2, 0)
combined = np.zeros([nY, args.join * nX, nC], dtype=np.uint8)
if args.join == 2:
combined[:, :nX, :] = mp4_ims[aa][:, :, [
2, 1, 0
]] # Invert back for BGR
combined[:, nX:, :] = to_save
else:
combined[:, :, :] = to_save
export_ims.append(combined) # Save for list
self.save_flow(mp4_fn, export_ims, mp4_ims[0].shape)
def main():
#0. assert args
global args, save_path
args = parser.parse_args()
if args.output_value == 'both':
output_string = "raw output and RGB visualization"
elif args.output_value == 'raw':
output_string = "raw output"
elif args.output_value == 'vis':
output_string = "RGB visualization"
print("=> will save " + output_string)
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir / 'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
#1. Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*1.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
#2. create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(
device
) #models = {module}<module 'models' from '/home/roit/ws/flownet_pt/models/__init__.py'>
model.eval() #flownetc or flownets
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = imread(img1_file)
img2 = imread(img2_file)
img1 = input_transform(img1)[:3, :, :]
img2 = input_transform(img2)[:3, :, :]
input_var = torch.cat([img1, img2
]).unsqueeze(0) #input_var={Tensor}, [1,6,h,w]
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(
0) # 如果用flownetc 到时候再分开
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var) #outTensor [1,2,h/4,w/4]
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
out1 = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = save_path / '{}{}'.format(img1_file.namebase[:-1],
suffix)
if args.output_value in ['vis', 'both']:
tmp = args.div_flow * flow_output
rgb_flow = flow2rgb(tmp, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
if args.output_value in ['raw', 'both']:
# Make the flow map a HxWx2 array as in .flo files
to_save = (args.div_flow *
flow_output).cpu().numpy().transpose(1, 2, 0)
np.save(filename + '.npy', to_save)
def main():
global args, save_path
args = parser.parse_args()
data_dir = args.data
print("=> fetching img pairs in '{}'".format(args.data))
save_path = args.output
print('=> will save everything to {}'.format(save_path))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
names = []
# for ext in args.img_exts:
# test_files = data_dir.files('*1.{}'.format(ext))
# for file in test_files:
# img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
# if img_pair.isfile():
# img_pairs.append([file, img_pair])
# print(img_pairs)
canon_im1num = [
1, 1, 2, 10, 10, 11, 21, 21, 21, 21, 22, 22, 22, 24, 25, 31, 31, 31,
32, 32, 33, 41, 43, 51, 54
] # 25
canon_im2num = [
2, 3, 3, 11, 12, 12, 22, 24, 25, 26, 23, 24, 27, 26, 27, 32, 33, 34,
33, 34, 34, 42, 45, 52, 55
]
# im1 = '%s/1/(1).JPG'%data_dir
# im2 = '%s/2/(1).JPG'%data_dir
# img_pairs.append([im1,im2])
s = 0
for i in range(0, 25):
foldernum1 = canon_im1num[i]
foldernum2 = canon_im2num[i]
rawnum = len(glob.glob('../../data/canon/%d/*.CR2' % (foldernum1)))
for j in range(rawnum):
image_path1 = '../../data/canon/%d/(%d).CR2' % (foldernum1, j + 1)
image_path2 = '../../data/canon/%d/(%d).CR2' % (foldernum2, j + 1)
img_pairs.append([image_path1, image_path2])
n = '%d-%d_%d' % (foldernum1, foldernum2, j + 1)
names.append(n)
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
count = -1
for (img1_file, img2_file) in tqdm(img_pairs):
raw1 = rawpy.imread(img1_file)
raw2 = rawpy.imread(img2_file)
im1 = raw1.raw_image_visible.astype(np.float32)
im1 = (im1 - 2047) / (16383 - 2047)
ratio = 0.3 / np.mean(im1)
im1 = np.minimum(np.maximum(im1 * ratio, 0.0), 1.0)
im2 = raw2.raw_image_visible.astype(np.float32)
im2 = (im2 - 2047) / (16383 - 2047)
ratio = 0.3 / np.mean(im2)
im2 = np.minimum(np.maximum(im2 * ratio, 0.0), 1.0)
# print(ratio)
im1 = np.expand_dims(im1, axis=2)
H = im1.shape[0]
W = im1.shape[1]
image1 = np.concatenate(
(
im1[0:H:2, 0:W:2, :], #r
(im1[0:H:2, 1:W:2, :] + im1[1:H:2, 0:W:2, :]) / 2.0, #g
im1[1:H:2, 1:W:2, :]),
axis=2) #b
im2 = np.expand_dims(im2, axis=2)
H = im2.shape[0]
W = im2.shape[1]
image2 = np.concatenate(
(
im2[0:H:2, 0:W:2, :], #r
(im2[0:H:2, 1:W:2, :] + im2[1:H:2, 0:W:2, :]) / 2.0, #g
im2[1:H:2, 1:W:2, :]),
axis=2) #b
i1 = crop(image1, 1920, 2944, 2)
i2 = crop(image2, 1920, 2944, 2)
i1 = (i1 * 255).astype('uint8')
i2 = (i2 * 255).astype('uint8')
i1 = cv.fastNlMeansDenoisingColored(i1, None, 10, 10, 7, 21)
i2 = cv.fastNlMeansDenoisingColored(i2, None, 10, 10, 7, 21)
img1 = input_transform(i1)
img2 = input_transform(i2)
# print(i1.shape)
# print(img1.shape)
input_var = torch.cat([img1, img2]).unsqueeze(0)
# if args.bidirectional:
# # feed inverted pair along with normal pair
# inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
# input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
count += 1
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = '%s/%s' % (save_path, names[count])
print(filename)
# if args.output_value in['vis', 'both']:
rgb_flow = flow2rgb(flow_output, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
# if args.output_value in ['raw', 'both']:
# # Make the flow map a HxWx2 array as in .flo files
to_save = (flow_output).cpu().numpy().transpose(1, 2, 0)
# print(to_save.shape)
flow_write(to_save, filename + '.flo')
def main():
global args, save_path
args = parser.parse_args()
save_path = args.output
print('=> will save everything to {}'.format(save_path))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1])
])
img_pairs = []
names = []
# for ext in args.img_exts:
# test_files = data_dir.files('*1.{}'.format(ext))
# for file in test_files:
# img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
# if img_pair.isfile():
# img_pairs.append([file, img_pair])
# print(img_pairs)
canon_im1num = [1,1,2,10,10,11,21,21,21,21,22,22,22,24,25,31,31,31,32,32,33,41,43,51,54] # 25
canon_im2num = [2,3,3,11,12,12,22,24,25,26,23,24,27,26,27,32,33,34,33,34,34,42,45,52,55]
sony_im1num = [1,1,1,2,11,11,12,15] # 8
sony_im2num = [2,3,4,4,12,13,13,16]
# im1 = '%s/1/(1).JPG'%data_dir
# im2 = '%s/2/(1).JPG'%data_dir
# img_pairs.append([im1,im2])
s = 0
for i in range(0,8):
foldernum1 = sony_im1num[i]
foldernum2 = sony_im2num[i]
rawnum = len(glob.glob('../../data/sid_Sony/%d/*.png'%(foldernum1)))
for j in range(rawnum):
image_path1 = '../../data/nlm/%d_%d.jpg'%(foldernum1,j+1)
image_path2 = '../../data/nlm/%d_%d.jpg'%(foldernum2,j+1)
img_pairs.append([image_path1, image_path2])
n = '%d-%d_%d'%(foldernum1,foldernum2,j+1)
names.append(n)
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
count = -1
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
# print(i1.shape)
# print(img1.shape)
input_var = torch.cat([img1, img2]).unsqueeze(0)
# if args.bidirectional:
# # feed inverted pair along with normal pair
# inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
# input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False)
count += 1
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = '%s/%s'%(save_path,names[count])
print(filename)
# if args.output_value in['vis', 'both']:
rgb_flow = flow2rgb(flow_output, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)
imwrite(filename + '.png', to_save)
# if args.output_value in ['raw', 'both']:
# # Make the flow map a HxWx2 array as in .flo files
to_save = (flow_output).cpu().numpy().transpose(1,2,0)
# print(to_save.shape)
flow_write(to_save, filename+'.flo')
def main():
global args, save_path
args = parser.parse_args()
if args.output_value == 'both':
output_string = "raw output and RGB visualization"
elif args.output_value == 'raw':
output_string = "raw output"
elif args.output_value == 'vis':
output_string = "RGB visualization"
print("=> will save " + output_string)
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir / 'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
# Make sure the folder has extracted images in the format of "name + num + extension", where num start from 0, 1, 2, ...
# Currently only test on MPI-Sintel, which comprises .png images
img_pairs = []
img_list = os.listdir(data_dir)
img_list = sorted(img_list)
#img_list.sort()
img_list = img_list[1:]
print(img_list)
#for ext in args.img_exts:
for ind in range(len(img_list) - 1):
img1 = img_list[ind]
print("img1: ", img1)
#test_files = data_dir.files('*1.{}'.format(ext))
img1_name, ext = img1.split('.')
basename, num = img1_name.split('_')
img2 = '.'.join(
['_'.join([basename, str(int(num) + 1).zfill(4)]), ext])
print("img 2: ", img2)
img_pairs.append([data_dir + '/' + img1, data_dir + '/' + img2])
#print("Test files: ")
#print(test_files)
#for file in test_files:
#img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
#print("file: ")
#print(file)
#print("img_pair: ")
#print(img_pair)
#if img_pair.isfile():
# img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
#print(img_pairs)
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
output = data_dir + '/flow/' + 'output.mp4'
print(output)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output, fourcc, 10.0, (640, 480))
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
print(img1_file.namebase)
filename = save_path / '{}{}'.format(img1_file.namebase, suffix)
print(filename)
if args.output_value in ['vis', 'both']:
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
out.write(to_save)
cv2.imshow('video', to_save)
if (cv2.waitKey(1) & 0xFF) == ord('q'): break
if args.output_value in ['raw', 'both']:
# Make the flow map a HxWx2 array as in .flo files
to_save = (args.div_flow *
flow_output).cpu().numpy().transpose(1, 2, 0)
np.save(filename + '.npy', to_save)
out.release()
cv2.destroyAllWindows()
def main():
global args, save_path
args = parser.parse_args()
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.233], std=[0.255]),
])
ref_list = glob.glob(os.path.join(args.data, '*_img1.tif'))
mov_list = glob.glob(os.path.join(args.data, '*_img2.tif'))
# create model
args = parser.parse_args()
args.lambda_list = [0.15, 0.85, 1, 0.1]
args.linear_resolution = 128
args.nonlinear_resolution = 512
model = Baseline(18, 18, args)
state_dict = torch.load(args.pretrained, map_location='cpu')
model.load_state_dict(state_dict['state_dict'], strict=True)
model = model.cuda()
model.eval()
cudnn.benchmark = True
i = 0
for ref_path, mov_path in zip(ref_list, mov_list):
ref = input_transform(Image.open(ref_path))
mov = input_transform(Image.open(mov_path))
input = torch.cat([ref, mov], dim=0).unsqueeze(0).cuda()
_, warp1, warp2, flow = model(input)
timg, simg = input.chunk(2, dim=1)
timg = (timg * 0.255 + 0.233) * 255
simg = (simg * 0.255 + 0.233) * 255
warp1 = (warp1 * 0.255 + 0.233) * 255
warp2 = (warp2 * 0.255 + 0.233) * 255
img = warp1.squeeze().cpu().numpy().astype(np.uint8)
path = '../eval2/' + str(i) + '_warp1.tif'
imsave(path, img)
img = warp2.squeeze().cpu().numpy().astype(np.uint8)
path = '../eval2/' + str(i) + '_warp2.tif'
imsave(path, img)
img = simg.squeeze().cpu().numpy().astype(np.uint8)
path = '../eval2/' + str(i) + '_src.tif'
imsave(path, img)
img = timg.squeeze().cpu().numpy().astype(np.uint8)
path = '../eval2/' + str(i) + '_tgt.tif'
imsave(path, img)
flow = flow.squeeze().cpu().numpy()
rgb_flow = flow2rgb(10 * flow, max_value=None)
image = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
path = '../eval2/' + str(i) + '_flow.png'
plt.imsave(path, image)
image = image[50:461, 50:461]
path = '../eval2/' + str(i) + '_flowcrop.png'
plt.imsave(path, image)
i = i + 1
