def train(train_loader, model, optimizer, epoch, train_writer):
global n_iter, args
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
flow2_EPEs = AverageMeter()
epoch_size = len(train_loader) if args.epoch_size == 0 else min(
len(train_loader), args.epoch_size)
# switch to train mode
model.train()
end = time.time()
#import ipdb; ipdb.set_trace()
for i, (input, target) in enumerate(train_loader):
# 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)
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[0], target, sparse=args.sparse)
# 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 train(train_loader, model, optimizer, epoch, train_writer):
global n_iter, args
batch_time = AverageMeter()#个人习惯??
data_time = AverageMeter()
losses = AverageMeter()
flow2_EPEs = AverageMeter()
epoch_size = len(train_loader) if args.epoch_size == 0 else min(len(train_loader), args.epoch_size)
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):#len(train_loader)= 22 (170 / batch_size = 22 )
# 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)# while train n*[b,2,h,w]
#-------------------
if args.sparse:#数据集gt为稀疏型(比如kitti这种来自真实场景而不是虚拟数据集)
# Since Target pooling is not very precise when sparse,
# take the highest resolution prediction and upsample it instead of downsampling target
#output[0] is flow2, 1 is flow3, so on
h, w = target.size()[-2:]
output = [F.interpolate(output[0], (h,w)), *output[1:]]
#in trainning model, outputs is a list
loss = multiscaleEPE(output, target, weights=args.multiscale_weights, sparse=args.sparse)
#default as 20
flow2_EPE = args.div_flow * realEPE(output[0], target, sparse=args.sparse)
# 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}]\nTrain:[{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:#epoch_size default as 1000
break
return losses.avg, flow2_EPEs.avg# every epoch has a losses.avg, EPE
def partsEPE(output, gtflow, seg_mask):
parts_epe_dict = {}
for bk in BODY_MAP.keys():
mask = seg_mask == BODY_MAP[bk]
gt_partflow = mask.type_as(gtflow)[:, :2, :, :] * gtflow
epe_part = realEPE(output, gt_partflow, sparse=True)
parts_epe_dict[bk] = epe_part
return parts_epe_dict
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 train(train_loader, model, optimizer, epoch, train_writer):
global n_iter, args
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
flow2_EPEs = AverageMeter()
epoch_size = len(train_loader) if args.epoch_size == 0 else min(
len(train_loader), args.epoch_size)
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.cuda(async=True)
input = [j.cuda() for j in input]
input_var = torch.autograd.Variable(torch.cat(input, 1))
target_var = torch.autograd.Variable(target)
# compute output
output = model(input_var)
loss = multiscaleEPE(output,
target_var,
weights=args.multiscale_weights)
flow2_EPE = args.div_flow * realEPE(output[0], target_var)
# record loss and EPE
losses.update(loss.data[0], target.size(0))
train_writer.add_scalar('train_loss', loss.data[0], n_iter)
flow2_EPEs.update(flow2_EPE.data[0], 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(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.cuda(async=True)
input_var = torch.autograd.Variable(torch.cat(input, 1).cuda(),
volatile=True)
target_var = torch.autograd.Variable(target, volatile=True)
# compute output
output = model(input_var)
flow2_EPE = args.div_flow * realEPE(output, target_var)
# record EPE
flow2_EPEs.update(flow2_EPE.data[0], 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:
output_writers[i].add_image(
'GroundTruth',
flow2rgb(args.div_flow * target[0].cpu().numpy(),
max_value=10), 0)
output_writers[i].add_image(
'Inputs', input[0][0].numpy().transpose(1, 2, 0), 0)
output_writers[i].add_image(
'Inputs', input[1][0].numpy().transpose(1, 2, 0), 1)
output_writers[i].add_image(
'FlowNet Outputs',
flow2rgb(args.div_flow * output.data[0].cpu().numpy(),
max_value=10), epoch)
if i % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'EPE {flow2_EPE.val:.3f} ({flow2_EPE.avg:.3f})'.format(
i,
len(val_loader),
batch_time=batch_time,
flow2_EPE=flow2_EPEs))
print(' * EPE {flow2_EPE.avg:.3f}'.format(flow2_EPE=flow2_EPEs))
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 main():
global args
args = parser.parse_args()
test_list = make_dataset(args.data)
flow_epe = AverageMeter()
avg_parts_epe = {}
for bk in BODY_MAP.keys():
avg_parts_epe[bk] = AverageMeter()
for i, (pred_path, flow_path, seg_path) in enumerate(tqdm(test_list)):
predflow = flow_transforms.ArrayToTensor()(load_flo(pred_path))
gtflow = flow_transforms.ArrayToTensor()(load_flo(flow_path))
segmask = flow_transforms.ArrayToTensor()(cv2.imread(seg_path))
predflow_var = predflow.unsqueeze(0)
gtflow_var = gtflow.unsqueeze(0)
segmask_var = segmask.unsqueeze(0)
predflow_var = predflow_var.to(device)
gtflow_var = gtflow_var.to(device)
segmask_var = segmask_var.to(device)
# compute output
epe = realEPE(predflow_var, gtflow_var)
epe_parts = partsEPE(predflow_var, gtflow_var, segmask_var)
#epe_parts.update((x, args.div_flow*y) for x, y in epe_parts.items() )
# record EPE
flow_epe.update(epe.item(), gtflow_var.size(0))
for bk in avg_parts_epe:
if epe_parts[bk].item() > 0:
avg_parts_epe[bk].update(epe_parts[bk].item(),
gtflow_var.size(0))
epe_dict = {}
for bk in BODY_MAP.keys():
epe_dict[bk] = avg_parts_epe[bk].avg
epe_dict['full_epe'] = flow_epe.avg
np.save(os.path.join('results', args.save_name), epe_dict)
print("Averge EPE", flow_epe.avg)
def validate(val_loader, model, epoch):
global args
batch_time = AverageMeter()
flow2_EPEs = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for i, batch in enumerate(val_loader):
target_x = batch['Dispx'].to(device)
target_y = batch['Dispy'].to(device)
target = torch.cat([target_x, target_y], 1).to(device)
in_ref = batch['Ref'].float().to(device)
in_ref = torch.cat([in_ref, in_ref, in_ref], 1).to(device)
in_def = batch['Def'].float().to(device)
in_def = torch.cat([in_def, in_def, in_def], 1).to(device)
input = torch.cat([in_ref, in_def], 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 % args.print_freq == 0:
print('Test: [{0}/{1}]\t Time {2}\t EPE {3}'.format(
i, len(val_loader), batch_time, flow2_EPEs))
#break
print(' * EPE {:.3f}'.format(flow2_EPEs.avg))
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, 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(val_loader, model, epoch, output_writers):
global args
batch_time = AverageMeter()
flow2_EPEs = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
total_photons = []
EPEs = []
for i, (input, target) in enumerate(val_loader):
if evaluate:
#print(np.shape(input))
#print(np.shape(input)[0])
#print(input.size()[0])
#print(np.shape(input.numpy()))
#print("####################")
photons = []
for q in range(input.size()[0]):
array = input.numpy()[q,0]
array = array*dataset_std[0] + dataset_mean[0]
sum = np.sum(array)
photons.append(sum)
total_photons.extend(photons)
target = target.to(device)
#################
#?????????????????
#input = torch.cat(input,1).to(device)
input = input.to(device)
# compute output
output = model(input)
flow2_EPE = args.div_flow*realEPE(output, target, sparse=args.sparse)
if evaluate:
#perImgEPE(output, target, sparse=args.sparse).tolist()
EPEs.extend(perImgEPE(output, target, sparse=args.sparse))
#print("flow to epe: ", flow2_EPE.item())
# record EPE
flow2_EPEs.update(flow2_EPE.item(), target.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if False:
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, :2].cpu() + mean_values).clamp(0, 1), 0)
output_writers[i].add_image('Inputs', (input[0, 2:].cpu() + mean_values).clamp(0, 1), 1)
#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 epoch % 1 == 0 and i % 5 ==0:
#try:
a = flow2rgb(args.div_flow * output[0], max_value=10)
#print(np.shape(a))
a = np.transpose(a, (1, 2, 0))
#print(np.shape(a))
#print(a[15:19, 27:40])
print(a.dtype)
print("#######################")
flow_map = output[0].detach().cpu().numpy()
flow_map = np.transpose(flow_map, (1, 2, 0))
print(np.shape(flow_map))
IMG = ((a / np.max(a))*255).astype(np.uint8)
viewImg(IMG, to_8bit = False, filename = "flow_" + str(i) + '.jpg', show = False)
A = args.div_flow * flow_map
#viewFlow(flow_array, dimx, dimy, to_8bit=False, filename=None, show=True):
viewFlow(A, 64, 48, to_8bit = False, filename = 'gen_flow_' + str(i) + '.jpg', show = False)
flow_map_gt = target[0].detach().cpu().numpy()
flow_map_gt = np.transpose(flow_map_gt, (1, 2, 0))
print(np.shape(flow_map_gt))
viewFlow(flow_map_gt, 256, 192, to_8bit=False, filename = 'gt_flow_' + str(i) + '.jpg', show = False)
#except:
# continue
#viewImg(flow2rgb(args.div_flow * output[0], max_value=10))
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))
if evaluate:
print("length of total_photons: ", len(total_photons))
print("length of EPEs: ", len(EPEs))
fig = plt.figure()
plt.plot(total_photons,EPEs,'.')
plt.title("End Point Error vs Collected Photons")
plt.xlabel("Photons in Example")
plt.ylabel("EPE of Example")
else:
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()
epoch_size = len(train_loader) if args.epoch_size == 0 else min(len(train_loader), args.epoch_size)
# switch to train mode
model.train()
end = time.time()
#rot = flow_transforms.RandomRotate(30)
#trans = flow_transforms.RandomTranslate(5)
for i, (input, target) in enumerate(train_loader):
######Data Augmentation#############
#input, target = rot(input, target)
#input, target = trans(input, target)
# measure data loading time
data_time.update(time.time() - end)
target = target.to(device)
#?????
#input = torch.cat(input,1).to(device)
input = input.to(device)
#print("input is: ", input)
#print("input[0] is: ", input[0])
#print("input[0] size: ",input[0].size())
if False:
print("time input:")
print(input[0,1,75:85,150:160])
print("boolean input")
print(input[0, 0, 75:85, 150:160])
# compute output
output = model(input)
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[0], target, sparse=args.sparse)
if False:
print("type: ", type(output))
print("length: ",len(output))
print("it is: ",output[0])
print("nextddddddddddddddddddddddd")
print("it is: ", output[1])
print("nextddddddddddddddddddddddd")
print("it is: ", output[2])
print("EPEEEEEEEEEEEEEEEEEEEEEEE")
print(flow2_EPE)
# 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
#if i > 5:
# break
return losses.avg, flow2_EPEs.avg
def main():
global args
args = parser.parse_args()
test_list = make_dataset(args.data)
# test_list = make_real_dataset(args.data)
if args.arch == 'pwc':
model = models.pwc_dc_net('models/pwc_net_ft.pth.tar').cuda()
elif args.arch == 'spynet':
model = models.spynet(nlevels=5, strmodel='F').cuda()
elif args.arch == 'flownet2':
model = models.FlowNet2().cuda()
print("=> using pre-trained weights for FlowNet2")
weights = torch.load('models/FlowNet2_checkpoint.pth.tar')
model.load_state_dict(weights['state_dict'])
if args.pretrained is not None:
network_data = torch.load(args.pretrained)
args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](data=network_data).cuda()
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
model.eval()
flow_epe = AverageMeter()
avg_mot_err = AverageMeter()
avg_parts_epe = {}
for bk in BODY_MAP.keys():
avg_parts_epe[bk] = AverageMeter()
if args.no_norm:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255])
])
else:
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])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
for i, (img_paths, flow_path, seg_path) in enumerate(tqdm(test_list)):
raw_im1 = flow_transforms.ArrayToTensor()(imread(img_paths[0],
mode='RGB')[:, :, :3])
raw_im2 = flow_transforms.ArrayToTensor()(imread(img_paths[1],
mode='RGB')[:, :, :3])
img1 = input_transform(imread(img_paths[0], mode='RGB')[:, :, :3])
img2 = input_transform(imread(img_paths[1], mode='RGB')[:, :, :3])
if flow_path is None:
_, h, w = img1.size()
new_h = int(np.floor(h / 256) * 256)
new_w = int(np.floor(w / 448) * 448)
# if i>744:
# import ipdb; ipdb.set_trace()
img1 = F.upsample(img1.unsqueeze(0), (new_h, new_w),
mode='bilinear').squeeze()
img2 = F.upsample(img2.unsqueeze(0), (new_h, new_w),
mode='bilinear').squeeze()
if flow_path is not None:
gtflow = target_transform(load_flo(flow_path))
segmask = flow_transforms.ArrayToTensor()(cv2.imread(seg_path))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if flow_path is not None:
gtflow_var = gtflow.unsqueeze(0)
segmask_var = segmask.unsqueeze(0)
input_var = input_var.to(device)
if flow_path is not None:
gtflow_var = gtflow_var.to(device)
segmask_var = segmask_var.to(device)
# compute output
output = model(input_var)
if flow_path is not None:
epe = args.div_flow * realEPE(
output,
gtflow_var,
sparse=True if 'KITTI' in args.dataset else False)
epe_parts = partsEPE(output, gtflow_var, segmask_var)
epe_parts.update(
(x, args.div_flow * y) for x, y in epe_parts.items())
# record EPE
flow_epe.update(epe.item(), gtflow_var.size(0))
for bk in avg_parts_epe:
if epe_parts[bk].item() > 0:
avg_parts_epe[bk].update(epe_parts[bk].item(),
gtflow_var.size(0))
# record motion warping error
raw_im1 = raw_im1.cuda().unsqueeze(0)
raw_im2 = raw_im2.cuda().unsqueeze(0)
mot_err = motion_warping_error(raw_im1, raw_im2,
args.div_flow * output)
avg_mot_err.update(mot_err.item(), raw_im1.size(0))
if args.output_dir is not None:
if flow_path is not None:
_, h, w = gtflow.size()
output_path = flow_path.replace(args.data, args.output_dir)
output_path = output_path.replace('/test/', '/')
os.system('mkdir -p ' + output_path[:-15])
else:
output_path = img_paths[0].replace(args.data, args.output_dir)
os.system('mkdir -p ' + output_path[:-10])
output_path = output_path.replace('.png', '.flo')
output_path = output_path.replace('/flow/', '/')
upsampled_output = F.interpolate(output, (h, w),
mode='bilinear',
align_corners=False)
flow_write(output_path,
upsampled_output.cpu()[0].data.numpy()[0],
upsampled_output.cpu()[0].data.numpy()[1])
if args.save_name is not None:
epe_dict = {}
for bk in BODY_MAP.keys():
epe_dict[bk] = avg_parts_epe[bk].avg
epe_dict['full_epe'] = flow_epe.avg
np.save(os.path.join('results', args.save_name), epe_dict)
print("Averge EPE", flow_epe.avg)
print("Motion warping error", avg_mot_err.avg)
def train(train_loader, model, optimizer, epoch, train_writer, config):
global n_iter
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
flow2_EPEs = AverageMeter()
epoch_size = len(train_loader) if args.epoch_size == 0 else min(len(train_loader), args.epoch_size)
# switch to train mode
model.train()
end = time.time()
if not args.self_supervised_loss:
# use old loss
for i, (input, target) in enumerate(train_loader):
# 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)
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[0], target, sparse=args.sparse)
# 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
elif args.unflow:
weights = [0.005, 0.01, 0.02, 0.08, 0.32]
for it, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.to(device)
im1 = input[0].to(device)
im2 = input[1].to(device)
input_fw = torch.cat(input, 1).to(device)
pred_fw = model(input_fw)
input_bw = torch.cat((im2, im1), 1).to(device)
pred_bw = model(input_bw)
census_loss = 0
census_loss_list = []
if config['census']:
#weights = [1, 0.34, 0.31, 0.27, 0.09]
#max_dist = [3, 2, 2, 1, 1]
for i in range(len(pred_fw)):
flow_fw = pred_fw[i] * args.div_flow
flow_bw = pred_bw[i] * args.div_flow
loss = ternary_loss(im2, im1, flow_fw, max_distance=1) +\
ternary_loss(im1, im2, flow_bw,max_distance=1)
census_loss += loss
census_loss_list.append(loss.item())
if not config['multiscale_census_loss']:
break
train_writer.add_scalar('train_loss_census', census_loss.item(), n_iter)
sl_loss = 0
sl_loss_list = []
if config['sl']:
for i in range(len(pred_fw)):
flow_fw = pred_fw[i] * args.div_flow
flow_bw = pred_bw[i] * args.div_flow
loss = smoothness_loss(flow_fw,config) + smoothness_loss(flow_bw,config)
#loss = smoothness_loss(flow_bw, config)
sl_loss += loss
sl_loss_list.append(loss.item())
if not config['multiscale_sl_loss']:
break
train_writer.add_scalar('train_loss_sl', sl_loss.item(), n_iter)
ssim_loss = 0
ssim_loss_list = []
if config['ssim']:
for i in range(len(pred_bw)):
flow_bw = pred_bw[i] * args.div_flow
loss = ssim(im1,im2,flow_bw)
ssim_loss += loss
ssim_loss_list.append(loss.item())
if not config['multiscale_ssim_loss']:
break
train_writer.add_scalar('train_loss_ssim', ssim_loss.item(), n_iter)
fb_loss = 0
fb_loss_list = []
if config['fb']:
for i in range(len(pred_bw)):
flow_fw = pred_fw[i] * args.div_flow
flow_bw = pred_bw[i] * args.div_flow
loss = forward_backward_loss(im1=im1, im2=im2, flow_fw=flow_fw, flow_bw=flow_bw, config=config)
fb_loss += loss
fb_loss_list.append(loss.item())
if not config['multiscale_fb_loss']:
break
train_writer.add_scalar('train_loss_fb', fb_loss.item(), n_iter)
# to check the magnitude of both losses
if it % 500 == 0:
print("[DEBUG] census_loss:", str(census_loss_list))
print("[DEBUG] sl_loss:", str(sl_loss_list))
print("[DEBUG] ssim_loss:", str(ssim_loss_list))
print("[DEBUG] fb_loss:", str(fb_loss_list))
loss = census_loss + sl_loss + ssim_loss + 0.001*fb_loss
# record loss and EPE
flow = pred_bw[0]
losses.update(loss.item(), target.size(0))
flow2_EPE = args.div_flow * realEPE(flow, target, sparse=args.sparse)
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 it % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t Time {3}\t Data {4}\t Loss {5}\t EPE {6}'
.format(epoch, it, epoch_size, batch_time,
data_time, losses, flow2_EPEs))
n_iter += 1
if it >= epoch_size:
break
return losses.avg, flow2_EPEs.avg
else:
# use self-supervised loss
for it, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.to(device)
im1 = input[0].to(device)
im2 = input[1].to(device)
input = torch.cat(input,1).to(device)
pred = model(input)
pl_loss = 0
pl_loss_list = []
for i in range(len(pred)):
flow = pred[i] * args.div_flow
loss = photometric_loss(im1, im2, flow, config)
pl_loss += loss
pl_loss_list.append(loss.item())
if not config['multiscale_pl_loss']:
break
sl_loss = 0
sl_loss_list = []
if config['weighted_sl_loss']:
for i in range(len(pred)):
flow = pred[i] * args.div_flow
loss = weighted_smoothness_loss(im1, im2, flow, config)
sl_loss += loss
sl_loss_list.append(loss.item())
if not config['multiscale_sl_loss']:
break
else:
# smoothness loss for multi resolution flow pyramid
for i in range(len(pred)):
flow = pred[i] * args.div_flow
loss = smoothness_loss(flow, config)
sl_loss += loss
sl_loss_list.append(loss.item())
if not config['multiscale_sl_loss']:
break
# to check the magnitude of both losses
if it % 500 == 0:
print("[DEBUG] pl_loss:", str(pl_loss_list))
print("[DEBUG] sl_loss:", str(sl_loss_list))
loss = pl_loss + sl_loss
# record loss and EPE
flow = pred[0]
losses.update(loss.item(), target.size(0))
flow2_EPE = args.div_flow * realEPE(flow, target, sparse=args.sparse)
train_writer.add_scalar('train_loss', loss.item(), n_iter)
train_writer.add_scalar('train_loss_pl', pl_loss.item(), n_iter)
train_writer.add_scalar('train_loss_sl', sl_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 it % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t Time {3}\t Data {4}\t Loss {5}\t EPE {6}'
.format(epoch, it, epoch_size, batch_time,
data_time, losses, flow2_EPEs))
n_iter += 1
if it >= epoch_size:
break
return losses.avg, flow2_EPEs.avg
