if args.loadmodel is not None:
pretrained_dict = torch.load(args.loadmodel)
mean_L = pretrained_dict['mean_L']
mean_R = pretrained_dict['mean_R']
pretrained_dict['state_dict'] = {
k: v
for k, v in pretrained_dict['state_dict'].items()
}
model.load_state_dict(pretrained_dict['state_dict'], strict=False)
else:
print('dry run')
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
mkdir_p('%s/%s/' % (args.outdir, args.dataset))
def main():
model.eval()
ttime_all = []
for inx in range(len(test_left_img)):
print(test_left_img[inx])
imgL_o = cv2.imread(test_left_img[inx])[:, :, ::-1]
imgR_o = cv2.imread(test_right_img[inx])[:, :, ::-1]
# for gray input images
if len(imgL_o.shape) == 2:
imgL_o = np.tile(imgL_o[:, :, np.newaxis], (1, 1, 3))
imgR_o = np.tile(imgR_o[:, :, np.newaxis], (1, 1, 3))
test_right_img = [
test_right_img[i] for i, flag in enumerate(split) if flag == 2
]
flow_paths = [flow_paths[i] for i, flag in enumerate(split) if flag == 2]
for i, gtflow_path in enumerate(flow_paths):
num = gtflow_path.split('/')[-1].strip().replace('flow.flo', 'img1.png')
if not 'test' in args.dataset and not 'clip' in args.dataset:
gtflow = read_flow(gtflow_path)
num = num.replace('jpg', 'png')
flow = read_flow('%s/%s/flo-%s' %
(args.path, args.dataset, num.replace('.png', '.pfm')))
if args.vis == 'yes':
flowimg = flow_to_image(flow) * np.linalg.norm(
flow[:, :, :2], 2, 2)[:, :, np.newaxis] / 100. / 255.
mkdir_p('%s/%s/flowimg' % (args.path, args.dataset))
plt.imsave('%s/%s/flowimg/%s' % (args.path, args.dataset, num),
flowimg)
if 'test' in args.dataset or 'clip' in args.dataset:
continue
gtflowimg = flow_to_image(gtflow)
mkdir_p('%s/%s/gtimg' % (args.path, args.dataset))
plt.imsave('%s/%s/gtimg/%s' % (args.path, args.dataset, num),
gtflowimg)
mask = gtflow[:, :, 2] == 1
gtflow = gtflow[:, :, :2]
flow = flow[:, :, :2]
epe = np.sqrt(np.power(gtflow - flow, 2).sum(-1))[mask]
def main():
model.eval()
ttime_all = []
for inx in range(len(test_left_img)):
print(test_left_img[inx])
imgL_o = skimage.io.imread(test_left_img[inx])
imgR_o = skimage.io.imread(test_right_img[inx])
# for gray input images
if len(imgL_o.shape) == 2:
imgL_o = np.tile(imgL_o[:, :, np.newaxis], (1, 1, 3))
imgR_o = np.tile(imgR_o[:, :, np.newaxis], (1, 1, 3))
# resize
maxh = imgL_o.shape[0] * args.testres
maxw = imgL_o.shape[1] * args.testres
max_h = int(maxh // 64 * 64)
max_w = int(maxw // 64 * 64)
if max_h < maxh: max_h += 64
if max_w < maxw: max_w += 64
input_size = imgL_o.shape
imgL = cv2.resize(imgL_o, (max_w, max_h))
imgR = cv2.resize(imgR_o, (max_w, max_h))
# flip channel, subtract mean
imgL = imgL[:, :, ::-1].copy() / 255. - np.asarray(mean_L).mean(0)[
np.newaxis, np.newaxis, :]
imgR = imgR[:, :, ::-1].copy() / 255. - np.asarray(mean_R).mean(0)[
np.newaxis, np.newaxis, :]
imgL = np.transpose(imgL, [2, 0, 1])[np.newaxis]
imgR = np.transpose(imgR, [2, 0, 1])[np.newaxis]
# support for any resolution inputs
from models.VCN import WarpModule, flow_reg
if hasattr(model.module, 'flow_reg64'):
model.module.flow_reg64 = flow_reg(
[1, max_w // 64, max_h // 64],
ent=model.module.flow_reg64.ent,
maxdisp=model.module.flow_reg64.md,
fac=model.module.flow_reg64.fac).cuda()
if hasattr(model.module, 'flow_reg32'):
model.module.flow_reg32 = flow_reg(
[1, max_w // 64 * 2, max_h // 64 * 2],
ent=model.module.flow_reg32.ent,
maxdisp=model.module.flow_reg32.md,
fac=model.module.flow_reg32.fac).cuda()
if hasattr(model.module, 'flow_reg16'):
model.module.flow_reg16 = flow_reg(
[1, max_w // 64 * 4, max_h // 64 * 4],
ent=model.module.flow_reg16.ent,
maxdisp=model.module.flow_reg16.md,
fac=model.module.flow_reg16.fac).cuda()
if hasattr(model.module, 'flow_reg8'):
model.module.flow_reg8 = flow_reg(
[1, max_w // 64 * 8, max_h // 64 * 8],
ent=model.module.flow_reg8.ent,
maxdisp=model.module.flow_reg8.md,
fac=model.module.flow_reg8.fac).cuda()
if hasattr(model.module, 'flow_reg4'):
model.module.flow_reg4 = flow_reg(
[1, max_w // 64 * 16, max_h // 64 * 16],
ent=model.module.flow_reg4.ent,
maxdisp=model.module.flow_reg4.md,
fac=model.module.flow_reg4.fac).cuda()
model.module.warp5 = WarpModule([1, max_w // 32, max_h // 32]).cuda()
model.module.warp4 = WarpModule([1, max_w // 16, max_h // 16]).cuda()
model.module.warp3 = WarpModule([1, max_w // 8, max_h // 8]).cuda()
model.module.warp2 = WarpModule([1, max_w // 4, max_h // 4]).cuda()
model.module.warpx = WarpModule([1, max_w, max_h]).cuda()
# forward
imgL = Variable(torch.FloatTensor(imgL).cuda())
imgR = Variable(torch.FloatTensor(imgR).cuda())
with torch.no_grad():
imgLR = torch.cat([imgL, imgR], 0)
model.eval()
torch.cuda.synchronize()
start_time = time.time()
rts = model(imgLR)
torch.cuda.synchronize()
ttime = (time.time() - start_time)
print('time = %.2f' % (ttime * 1000))
ttime_all.append(ttime)
pred_disp, entropy = rts
# upsampling
pred_disp = torch.squeeze(pred_disp).data.cpu().numpy()
pred_disp = cv2.resize(np.transpose(pred_disp, (1, 2, 0)),
(input_size[1], input_size[0]))
pred_disp[:, :, 0] *= input_size[1] / max_w
pred_disp[:, :, 1] *= input_size[0] / max_h
flow = np.ones([pred_disp.shape[0], pred_disp.shape[1], 3])
flow[:, :, :2] = pred_disp
entropy = torch.squeeze(entropy).data.cpu().numpy()
entropy = cv2.resize(entropy, (input_size[1], input_size[0]))
# save predictions
if args.dataset == 'mbstereo':
dirname = '%s/%s/%s' % (args.outdir, args.dataset,
test_left_img[inx].split('/')[-2])
mkdir_p(dirname)
idxname = ('%s/%s') % (dirname.rsplit(
'/', 1)[-1], test_left_img[inx].split('/')[-1])
else:
idxname = test_left_img[inx].split('/')[-1]
if args.dataset == 'mbstereo':
with open(test_left_img[inx].replace('im0.png', 'calib.txt')) as f:
lines = f.readlines()
#max_disp = int(int(lines[9].split('=')[-1]))
max_disp = int(int(lines[6].split('=')[-1]))
with open(
'%s/%s/%s' % (args.outdir, args.dataset,
idxname.replace('im0.png', 'disp0IO.pfm')),
'w') as f:
save_pfm(
f,
np.clip(-flow[::-1, :, 0].astype(np.float32), 0, max_disp))
with open(
'%s/%s/%s/timeIO.txt' %
(args.outdir, args.dataset, idxname.split('/')[0]), 'w') as f:
f.write(str(ttime))
elif args.dataset == 'k15stereo' or args.dataset == 'k12stereo':
skimage.io.imsave(
'%s/%s/%s.png' %
(args.outdir, args.dataset, idxname.split('.')[0]),
(-flow[:, :, 0].astype(np.float32) * 256).astype('uint16'))
else:
write_flow(
'%s/%s/%s.png' %
(args.outdir, args.dataset, idxname.rsplit('.', 1)[0]),
flow.copy())
cv2.imwrite(
'%s/%s/ent-%s.png' %
(args.outdir, args.dataset, idxname.rsplit('.', 1)[0]),
entropy * 200)
torch.cuda.empty_cache()
print(np.mean(ttime_all))
def main():
model.eval()
ttime_all = []
rmses = 0
nrmses = 0
inx = test_left_img.index(args.image)
print(test_left_img[inx])
flo = read_flow(test_flow[inx])
imgL_o = np.asarray(Image.open(test_left_img[inx]))
imgR_o = np.asarray(Image.open(test_right_img[inx]))
# resize
maxh = imgL_o.shape[0]*args.testres
maxw = imgL_o.shape[1]*args.testres
max_h = int(maxh // 64 * 64)
max_w = int(maxw // 64 * 64)
if max_h < maxh: max_h += 64
if max_w < maxw: max_w += 64
input_size = imgL_o.shape
imgL = cv2.resize(imgL_o,(max_w, max_h))
imgR = cv2.resize(imgR_o,(max_w, max_h))
# flip channel, subtract mean
imgL = imgL[:, :, None].copy() / 255. - np.asarray(mean_L).mean(0)[np.newaxis,np.newaxis,:]
imgR = imgR[:, :, None].copy() / 255. - np.asarray(mean_R).mean(0)[np.newaxis,np.newaxis,:]
print(imgL.shape)
imgL = np.transpose(imgL, [2,0,1])[np.newaxis]
imgR = np.transpose(imgR, [2,0,1])[np.newaxis]
# forward
imgL = Variable(torch.FloatTensor(imgL).cuda())
imgR = Variable(torch.FloatTensor(imgR).cuda())
with torch.no_grad():
imgLR = torch.cat([imgL,imgR],0)
model.eval()
torch.cuda.synchronize()
start_time = time.time()
rts = model(imgLR)
torch.cuda.synchronize()
ttime = (time.time() - start_time); print('time = %.2f' % (ttime*1000) )
ttime_all.append(ttime)
pred_disp, entropy = rts
# upsampling
pred_disp = torch.squeeze(pred_disp).data.cpu().numpy()
pred_disp = cv2.resize(np.transpose(pred_disp,(1,2,0)), (input_size[1], input_size[0]))
pred_disp[:,:,0] *= input_size[1] / max_w
pred_disp[:,:,1] *= input_size[0] / max_h
flow = np.ones([pred_disp.shape[0],pred_disp.shape[1],3])
flow[:,:,:2] = pred_disp
rmse = np.sqrt((np.linalg.norm(flow[:,:,:2] - flo[:,:,:2], ord=2, axis=-1) ** 2).mean())
rmses += rmse
nrmses += rmse / np.sqrt((np.linalg.norm(flo[:,:,:2], ord=2, axis=-1) ** 2).mean())
error = np.linalg.norm(flow[:,:,:2] - flo[:,:,:2], ord=2, axis=-1) ** 2
error = 255 - 255 * error / error.max()
entropy = torch.squeeze(entropy).data.cpu().numpy()
entropy = cv2.resize(entropy, (input_size[1], input_size[0]))
# save predictions
if args.dataset == 'mbstereo':
dirname = '%s/%s/%s'%(args.outdir, args.dataset, test_left_img[inx].split('/')[-2])
mkdir_p(dirname)
idxname = ('%s/%s')%(dirname.rsplit('/',1)[-1],test_left_img[inx].split('/')[-1])
else:
idxname = test_left_img[inx].split('/')[-1]
if args.dataset == 'mbstereo':
with open(test_left_img[inx].replace('im0.png','calib.txt')) as f:
lines = f.readlines()
#max_disp = int(int(lines[9].split('=')[-1]))
max_disp = int(int(lines[6].split('=')[-1]))
with open('%s/%s/%s'% (args.outdir, args.dataset,idxname.replace('im0.png','disp0IO.pfm')),'w') as f:
save_pfm(f,np.clip(-flow[::-1,:,0].astype(np.float32),0,max_disp) )
with open('%s/%s/%s/timeIO.txt'%(args.outdir, args.dataset,idxname.split('/')[0]),'w') as f:
f.write(str(ttime))
elif args.dataset == 'k15stereo' or args.dataset == 'k12stereo':
skimage.io.imsave('%s/%s/%s.png'% (args.outdir, args.dataset,idxname.split('.')[0]),(-flow[:,:,0].astype(np.float32)*256).astype('uint16'))
else:
# write_flow('%s/%s/%s.png'% (args.outdir, args.dataset,idxname.rsplit('.',1)[0]), flow.copy())
cv2.imwrite('%s/%s/%s.png' % (args.outdir, args.dataset,idxname.rsplit('.',1)[0]), flow_to_image(flow)[:, :, ::-1])
cv2.imwrite('%s/%s/%s-gt.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0]), flow_to_image(flo)[:, :, ::-1])
arrow_pic(flo, '%s/%s/%s-vec-gt.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0]))
arrow_pic(flow, '%s/%s/%s-vec.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0]))
test_compressibility(
flo, flow,
'%s/%s/%s-compr.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0])
)
test_energy_spectrum(
flo, flow,
'%s/%s/%s-energy.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0])
)
test_intermittency_r(
flo, flow,
'%s/%s/%s-interm-r.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0])
)
test_intermittency_n(
flo, flow,
'%s/%s/%s-interm-n.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0])
)
cv2.imwrite('%s/%s/%s-err.png' % (args.outdir, args.dataset, idxname.rsplit('.', 1)[0]), error)
# cv2.imwrite('%s/%s/ent-%s.png'% (args.outdir, args.dataset,idxname.rsplit('.',1)[0]), entropy*200)
torch.cuda.empty_cache()
rmses /= len(test_left_img)
nrmses /= len(test_left_img)
print(np.mean(ttime_all), rmses, nrmses)
uz_mesh[i, j] = ipu(x_mesh[i, j], y_mesh[i, j])
ipv = RectBivariateSpline(np.arange(X), np.arange(Y), field[..., 1])
vz_mesh = np.zeros_like(x_mesh)
for i in range(nx):
for j in range(ny):
vz_mesh[i, j] = ipv(x_mesh[i, j], y_mesh[i, j])
fig, ax = plt.subplots(figsize=(10, 10))
ax.imshow(flow_to_image(field))
ax.quiver(xs, ys, uz_mesh, vz_mesh, angles='xy')
ax.axis('off')
fig.savefig(fname)
mkdir_p('%s/%s/' % (args.outdir, "generated"))
def main():
model.eval()
flowl0 = "/gpfs/gpfs0/y.maximov/kolya/piv/SQG/SQG_00001_flow.flo"
iml0 = "/gpfs/gpfs0/y.maximov/kolya/piv/SQG/SQG_00001_img1.tif"
iml1 = "/gpfs/gpfs0/y.maximov/kolya/piv/SQG/SQG_00001_img2.tif"
iml0 = default_loader(iml0)
iml1 = default_loader(iml1)
iml1 = np.asarray(iml1) / 255.
iml0 = np.asarray(iml0) / 255.
iml0 = iml0[:, :, None].copy() # iml0[:,:,::-1].copy()
iml1 = iml1[:, :, None].copy() # iml1[:,:,::-1].copy()
