def align_fxn(inputs, db=None, gt_info=None):
if db is None: db = inputs
burst = inputs
flow = gt_info['flow']
isize = gt_info['isize']
nimages, npix, nframes_m1, two = flow.shape
aligned = align_from_flow(burst, flow, 0, isize=isize)
aligned = aligned.to(burst.device, non_blocking=True)
return aligned, flow
def get_train_log_info(cfg, model, denoised, loss, dyn_noisy, dyn_clean, sims,
masks, aligned, flow, flow_gt):
# -- init info --
info = {}
nframes, nimages, ncolor, h, w = dyn_clean.shape
ref_t = nframes // 2
# -- image psnrs --
image_psnrs = images_to_psnrs(denoised, dyn_clean[ref_t])
info['image_psnrs'] = image_psnrs
# -- sim images psnrs
nimages = sims.shape[1]
nsims = sims.shape[0] - 1
clean = repeat(dyn_clean[ref_t], 'b c h w -> s b c h w', s=nsims)
ref_clean = rearrange(clean, 's b c h w -> (s b) c h w')
sims = rearrange(sims[1:], 's b c h w -> (s b) c h w')
sim_psnrs = images_to_psnrs(ref_clean, sims)
sim_psnrs = rearrange(sim_psnrs, '(t b) -> t b', b=nimages)
info['sim_psnrs'] = sim_psnrs
# -- aligned image psnrs --
T, B, C, H, W = dyn_noisy.shape
isize = edict({'h': H, 'w': W})
clean = repeat(dyn_clean[ref_t], 'b c h w -> t b c h w', t=nframes)
ref_clean = rearrange(clean, 't b c h w -> (t b) c h w')
if not (flow is None):
aligned_clean = align_from_flow(dyn_clean,
flow,
cfg.nblocks,
isize=isize)
aligned_clean = aligned_clean.to(dyn_clean.device, non_blocking=True)
aligned_rs = rearrange(aligned_clean, 't b c h w -> (t b) c h w')
aligned_psnrs = images_to_psnrs(ref_clean, aligned_rs)
aligned_psnrs = rearrange(aligned_psnrs, '(t b) -> t b', t=nframes)
info['aligned_psnrs'] = aligned_psnrs
else:
info['aligned_psnrs'] = np.zeros(1)
# -- epe errors --
if not (flow is None):
info['epe'] = compute_epe(flow, flow_gt)
else:
info['epe'] = np.zeros(1)
# -- nnf acc --
if not (flow is None):
info['nnf_acc'] = compute_pair_flow_acc(flow, flow_gt)
else:
info['nnf_acc'] = np.zeros(1)
return info
def run_check(ex):
blocks_aligned = align_from_blocks(ex.burst,
ex.blocks,
ex.nblocks,
ex.patchsize,
isize=ex.isize)
verify_aligned(blocks_aligned, ex.burst, ex.patchsize)
flow_aligned = align_from_flow(ex.burst,
ex.flow,
ex.patchsize,
isize=ex.isize)
verify_aligned(flow_aligned, ex.burst, ex.patchsize)
pix_aligned = align_from_pix(ex.burst, ex.pix, ex.nblocks)
verify_aligned(pix_aligned, ex.burst, ex.patchsize)
def align_fxn(burst, db=None, gt_info=None):
if db is None: db = burst
T, B, C, H, W = burst.shape
isize = edict({'h': H, 'w': W})
# if cfg.noise_params.ntype == "pn" or cfg.use_anscombe:
# search = anscombe.forward(burst)
# search -= search.min()
# search /= search.max()
burst = burst.to(gpuid, non_blocking=True)
clean = return_optional(gt_info, 'clean', None)
search = search_img_select(burst, burst)
flow = flownet.burst2flow(search).cpu()
flow = rearrange(flow, 't i h w two -> i (h w) t two')
aligned = None
if comp_align:
aligned = align_from_flow(burst, flow, cfg.pads, isize=isize)
aligned = aligned.to(burst.device, non_blocking=True)
return aligned, flow
def align_fxn(burst, db=None, gt_info=None):
if db is None: db = burst
T, B, C, H, W = burst.shape
isize = edict({'h': H, 'w': W})
# if cfg.noise_params.ntype == "pn" or cfg.use_anscombe:
# search = anscombe.forward(burst)
# search -= search.min()
# search /= search.max()
clean = return_optional(gt_info, 'clean', None)
search = align_noise(burst, clean)
_, flow, _ = runJointSearch(search,
cfg.patchsize,
cfg.nblocks,
k=1,
valMean=valMode)
flow = rearrange(flow, 't i h w 1 two -> i (h w) t two')
aligned = align_from_flow(burst, flow, cfg.nblocks, isize=isize)
aligned = aligned.to(burst.device, non_blocking=True)
return aligned, flow
def align_fxn(burst, db=None, gt_info=None):
if db is None: db = burst
T, B, C, H, W = burst.shape
isize = edict({'h': H, 'w': W})
# burst = inputs['burst']
# isize = inputs['isize']
# if cfg.noise_params.ntype == "pn" or cfg.use_anscombe:
# search = anscombe.forward(burst)
# search -= search.min()
# search /= search.max()
clean = return_optional(gt_info, 'clean', None)
search = search_img_select(burst, clean)
nnf_vals, nnf_pix = nnf.compute_burst_nnf(search,
ref_t,
cfg.patchsize,
K=1)
shape_str = 't b h w two -> b (h w) t two'
nnf_pix_best = torch.LongTensor(
rearrange(nnf_pix[..., 0, :], shape_str))
flow = pix_to_flow(nnf_pix_best)
aligned = align_from_flow(burst, flow, 0, isize=isize)
aligned = aligned.to(burst.device, non_blocking=True)
return aligned, flow
def test_nnf():
# -- get config --
cfg = config()
# -- set random seed --
set_seed(cfg.random_seed)
# -- load dataset --
data, loaders = load_image_dataset(cfg)
image_iter = iter(loaders.tr)
# -- get score function --
score_fxn = get_score_function(cfg.score_fxn_name)
# -- some constants --
NUM_BATCHES = 2
nframes, nblocks = cfg.nframes, cfg.nblocks
patchsize = cfg.patchsize
check_parameters(nblocks, patchsize)
# -- create evaluator
iterations, K = 10, 2
subsizes = [2, 2, 2, 2, 2]
evaluator = combo.eval_scores.EvalBlockScores(score_fxn, patchsize, 100,
None)
# -- iterate over images --
for image_bindex in range(NUM_BATCHES):
# -- sample & unpack batch --
sample = next(image_iter)
sample_to_cuda(sample)
dyn_noisy = sample['noisy'] # dynamics and noise
dyn_clean = sample['burst'] # dynamics and no noise
static_noisy = sample['snoisy'] # no dynamics and noise
static_clean = sample['sburst'] # no dynamics and no noise
flow_gt = sample['flow']
# -- shape info --
T, B, C, H, W = dyn_noisy.shape
isize = edict({'h': H, 'w': W})
ref_t = nframes // 2
npix = H * W
# -- groundtruth flow --
flow_gt = repeat(flow_gt, 'i tm1 two -> i p tm1 two', p=npix)
print("sample['flow']: ", flow_gt.shape)
aligned_of = align_from_flow(dyn_clean,
flow_gt,
patchsize,
isize=isize)
# -- compute nearest neighbor fields --
shape_str = 't b h w two -> b (h w) t two'
nnf_vals, nnf_pix = nnf.compute_burst_nnf(dyn_clean, ref_t, patchsize)
nnf_pix_best = torch.LongTensor(
rearrange(nnf_pix[..., 0, :], shape_str))
nnf_pix_best = torch.LongTensor(nnf_pix_best)
flow_nnf = pix_to_flow(nnf_pix_best)
aligned_nnf = align_from_pix(dyn_clean, nnf_pix_best, patchsize)
# -- compute proposed search of nnf --
flow_split = optim.v1.run_image_burst(dyn_clean, patchsize, evaluator,
nblocks, iterations, subsizes, K)
isize = edict({'h': H, 'w': W})
aligned_split = align_from_flow(dyn_clean,
flow_split,
patchsize,
isize=isize)
# -- compute proposed search of nnf --
flow_est = optim.v3.run_image_burst(dyn_clean, patchsize, evaluator,
nblocks, iterations, subsizes, K)
aligned_est = align_from_flow(dyn_clean,
flow_est,
patchsize,
isize=isize)
# -- banner --
print("-" * 25 + " Results " + "-" * 25)
# -- compare gt v.s. nnf computations --
nnf_of = compute_epe(flow_nnf, flow_gt)
split_of = compute_epe(flow_split, flow_gt)
est_of = compute_epe(flow_est, flow_gt)
split_nnf = compute_epe(flow_split, flow_nnf)
est_nnf = compute_epe(flow_est, flow_nnf)
print("-" * 50)
print("EPE Errors")
print("-" * 50)
print("NNF v.s. Optical Flow.")
print(nnf_of)
print("Split v.s. Optical Flow.")
print(split_of)
print("Proposed v.s. Optical Flow.")
print(est_of)
print("Split v.s. NNF")
print(split_nnf)
print("Proposed v.s. NNF")
print(est_nnf)
# -- psnr eval --
pad = 2 * patchsize
isize = edict({'h': H - pad, 'w': W - pad})
psnr_of = compute_aligned_psnr(aligned_of, static_clean, isize)
psnr_nnf = compute_aligned_psnr(aligned_nnf, static_clean, isize)
psnr_split = compute_aligned_psnr(aligned_split, static_clean, isize)
psnr_est = compute_aligned_psnr(aligned_est, static_clean, isize)
print("-" * 50)
print("PSNR Values")
print("-" * 50)
print("Optical Flow [groundtruth v1]")
print(psnr_of)
print("NNF [groundtruth v2]")
print(psnr_nnf)
print("Split [old method]")
print(psnr_split)
print("Proposed [new method]")
print(psnr_est)
def run_multiscale_nnf(cfg, noisy, clean, nlevels=3, verbose=False):
T, C, H, W = noisy.shape
nframes = T
noisy = noisy[:, None]
clean = clean[:, None]
isize = edict({'h': H, 'w': W})
isize_l = [H, W]
pad3 = cfg.nblocks // 2 + 3 // 2
psize3 = edict({'h': H - pad3, 'w': W - pad3})
pad = cfg.nblocks // 2 + cfg.patchsize // 2
psize = edict({'h': H - pad, 'w': W - pad})
cfg_patchsize = cfg.patchsize
factor = 2
noisy = tvF.resize(noisy[:, 0], [H // factor, W // factor],
interpolation=InterpMode.BILINEAR)[:, None]
clean = tvF.resize(clean[:, 0], [H // factor, W // factor],
interpolation=InterpMode.BILINEAR)[:, None]
T, _, C, H, W = noisy.shape
isize = edict({'h': H, 'w': W})
isize_l = [H, W]
pad3 = cfg.nblocks // 2 + 3 // 2
psize3 = edict({'h': H - pad3, 'w': W - pad3})
pad = cfg.nblocks // 2 + cfg.patchsize // 2
psize = edict({'h': H - pad, 'w': W - pad})
cfg_patchsize = cfg.patchsize
# -- looks good --
cfg.patchsize = 3
align_fxn = get_align_method(cfg, "l2_global")
_, flow = align_fxn(clean.to(0))
aclean = align_from_flow(clean, flow, cfg.nblocks, isize=isize)
save_image("aclean.png", aclean)
apsnr = align_psnr(aclean, psize3)
print("[global] clean: ", apsnr)
# -- looks not good --
_, flow = align_fxn(noisy.to(0))
isize = edict({'h': H, 'w': W})
aclean = align_from_flow(clean, flow, cfg.nblocks, isize=isize)
save_image("aclean_rs1.png", aclean)
apsnr = align_psnr(aclean, psize3)
print("noisy: ", apsnr)
# -- fix it --
cfg.nblocks = 5
align_fxn = get_align_method(cfg, "pair_l2_local")
_, flow = align_fxn(aclean.to(0))
isize = edict({'h': H, 'w': W})
aclean = align_from_flow(aclean, flow, cfg.nblocks, isize=isize)
save_image("aclean_rs1.png", aclean)
apsnr = align_psnr(aclean, psize3)
print("[fixed] noisy: ", apsnr)
#
# -- [Tiled] try it again and to fix it --
#
img_ps = 3
cfg.patchsize = img_ps
cfg.nblocks = 50
tnoisy = padAndTileBatch(noisy, cfg.patchsize, cfg.nblocks)
tclean = padAndTileBatch(clean, cfg.patchsize, cfg.nblocks)
t2i_clean = tvF.center_crop(tiled_to_img(tclean, img_ps), isize_l)
print(t2i_clean.shape, clean.shape)
save_image("atiled_to_img.png", t2i_clean)
delta = torch.sum(torch.abs(clean - t2i_clean)).item()
assert delta < 1e-8, "tiled to image must work!"
cfg.patchsize = 3
align_fxn = get_align_method(cfg, "pair_l2_local")
_, flow = align_fxn(tnoisy.to(0))
print(flow.shape, tclean.shape, clean.shape, np.sqrt(flow.shape[1]))
nbHalf = cfg.nblocks // 2
pisize = edict({'h': H + 2 * nbHalf, 'w': W + 2 * nbHalf})
aclean = align_from_flow(tclean, flow, cfg.nblocks, isize=pisize)
aclean_img = tvF.center_crop(tiled_to_img(aclean, img_ps), isize_l)
save_image("aclean_rs1_tiled.png", aclean_img)
apsnr = align_psnr(aclean_img, psize3)
print("[tiled] noisy: ", apsnr)
# i want to use a different block size but I need to correct the image padding..?
# def shrink_search_space(tclean,flow,nblocks_prev,nblocks_curr):
# print("tclean.shape: ",tclean.shape)
# print("flow.shape: ",flow.shape)
# T,_,C,H,W = tclean.shape
# flow = rearrange(flow,'i (h w) t two -> t i two h w',h=H)
# tclean = tvF.center_crop(tclean,new_size)
# = tvF.center_crop(tclean,new_size)
nblocks_prev = cfg.nblocks
cfg.nblocks = 5
# tclean,flow = shrink_search_space(tclean,flow,nblocks_prev,cfg.nblocks)
align_fxn = get_align_method(cfg, "pair_l2_local")
at_clean = align_from_flow(tclean, flow, cfg.nblocks, isize=pisize)
_, flow_at = align_fxn(at_clean.to(0))
aaclean = align_from_flow(at_clean, flow_at, cfg.nblocks, isize=pisize)
aaclean_img = tvF.center_crop(tiled_to_img(aaclean, img_ps), isize_l)
save_image("aclean_rs1_fixed.png", aaclean_img)
apsnr = align_psnr(aaclean_img, psize3)
print("[fixed] noisy: ", apsnr)
exit()
cfg.patchsize = 1 #cfg_patchsize
align_fxn = get_align_method(cfg, "pair_l2_local")
# clusters = cluster_flow(flow,H,nclusters=4)
cflow = flow #replace_flow_median(flow,clusters,H,cfg.nblocks)
# save_image("clusters.png",clusters.type(torch.float))
cflow_img = flow2img(cflow, H, T)
save_image("cflow.png", cflow_img)
aclean = align_from_flow(clean, cflow, cfg.nblocks, isize=isize)
save_image("aclean_rs1_cf.png", aclean)
print(cflow[:, 64 * 64 + 64])
apsnr = align_psnr(aclean, psize)
print("noisy_cf: ", apsnr)
print(flow.shape)
# flow = rearrange(flow,'i (h w) t two -> t i two h w',h=H)
# print_stats(flow)
flow_img = flow2img(flow, H, T)
save_image("flow.png", flow_img)
print(torch.histc(flow.type(torch.float)))
factor = 2
cfg.nblocks = max(cfg.nblocks // 2, 3)
cfg.patchsize = 1
# cfg.patchsize = max(cfg.patchsize//2,3)
noisy_rs = tvF.resize(noisy[:, 0], [H // factor, W // factor],
interpolation=InterpMode.BILINEAR)[:, None]
_, flow_rs = align_fxn(noisy_rs.to(0))
clean_rs = tvF.resize(clean[:, 0], [H // factor, W // factor],
interpolation=InterpMode.BILINEAR)[:, None]
isize = edict({'h': H // factor, 'w': W // factor})
aclean = align_from_flow(clean_rs, flow_rs, cfg.nblocks, isize=isize)
save_image("aclean_rs2.png", aclean)
apsnr = align_psnr(aclean, psize)
print("rs2", apsnr, cfg.nblocks, cfg.patchsize)
clusters = cluster_flow(flow_rs, H // factor, nclusters=3)
save_image("clusters_rs.png", clusters.type(torch.float))
# cflow_rs = cluster_flow(flow_rs,H//factor,nclusters=5)
# print(cflow_rs)
aclean = align_from_flow(clean_rs, cflow_rs, cfg.nblocks, isize=isize)
save_image("aclean_rs2_cl.png", aclean)
apsnr = align_psnr(aclean, psize)
print("rs2_cl", apsnr, cfg.nblocks, cfg.patchsize)
exit()
print(flow_rs.shape)
# flow_rs = rearrange(flow_rs,'i (h w) t two -> t i two h w',h=H//factor)
print(flow_rs.shape)
flow_img = flow2img(flow_rs, H // factor, T)
save_image("flow_rs2.png", flow_img)
fmin, fmax, fmean = print_stats(flow_rs)
print(torch.histc(flow_rs.type(torch.float), max=50, min=-50))
factor = 4
cfg.nblocks = max(cfg.nblocks // 2, 3)
# cfg.patchsize = max(cfg.patchsize//2,3)
noisy_rs = tvF.resize(noisy[:, 0], [H // factor, W // factor],
interpolation=InterpMode.BILINEAR)[:, None]
_, flow_rs = align_fxn(noisy_rs.to(0))
clean_rs = tvF.resize(clean[:, 0], [H // factor, W // factor],
interpolation=InterpMode.BILINEAR)[:, None]
isize = edict({'h': H // factor, 'w': W // factor})
aclean = align_from_flow(clean_rs, flow_rs, cfg.nblocks, isize=isize)
save_image("aclean_rs4.png", aclean)
apsnr = align_psnr(aclean, psize)
print(apsnr, cfg.nblocks, cfg.patchsize)
print(flow_rs.shape)
# flow_rs = rearrange(flow_rs,'i (h w) t two -> t i two h w',h=H//factor)
print(flow_rs.shape)
flow_img = flow2img(flow_rs, H // factor, T)
save_image("flow_rs4.png", flow_img)
fmin, fmax, fmean = print_stats(flow_rs)
print(torch.histc(flow_rs.type(torch.float), max=50, min=-50))
def execute_experiment(cfg):
# -- init exp! --
print("RUNNING EXP.")
print(cfg)
# -- create results record to save --
dims = {
'batch_results': None,
'batch_to_record': None,
'record_results': {
'default': 0
},
'stack': {
'default': 0
},
'cat': {
'default': 0
}
}
record = cache_io.ExpRecord(dims)
# -- set random seed --
set_seed(cfg.random_seed)
# -- load dataset --
data, loaders = load_image_dataset(cfg)
image_iter = iter(loaders.tr)
# -- get score function --
score_fxn_ave = get_score_function("ave")
score_fxn_mse = get_score_function("mse")
score_fxn_bs = get_score_function("bootstrapping")
score_fxn_bs_cf = get_score_function("bootstrapping_cf")
# score_fxn_bs = get_score_function("bootstrapping_mod2")
# score_fxn_bs = get_score_function(cfg.score_fxn_name)
score_fxn_bsl = get_score_function("bootstrapping_limitB")
# -- some constants --
NUM_BATCHES = 3
nframes, nblocks = cfg.nframes, cfg.nblocks
patchsize = cfg.patchsize
ppf = cfg.dynamic_info.ppf
check_parameters(nblocks, patchsize)
# -- create evaluator for ave; simple --
iterations, K = 1, 1
subsizes = []
block_batchsize = 256
eval_ave_simp = EvalBlockScores(score_fxn_ave, "ave", patchsize,
block_batchsize, None)
# -- create evaluator for ave --
iterations, K = 1, 1
subsizes = []
eval_ave = EvalBlockScores(score_fxn_ave, "ave", patchsize,
block_batchsize, None)
eval_mse = EvalBlockScores(score_fxn_mse, "mse", patchsize,
block_batchsize, None)
# -- create evaluator for bootstrapping --
block_batchsize = 81
eval_plimb = EvalBootBlockScores(score_fxn_bsl, score_fxn_bs, "bsl",
patchsize, block_batchsize, None)
eval_prop = EvalBlockScores(score_fxn_bs, "bs", patchsize, block_batchsize,
None)
eval_prop_cf = EvalBlockScores(score_fxn_bs_cf, "bs_cf", patchsize,
block_batchsize, None)
# -- iterate over images --
for image_bindex in range(NUM_BATCHES):
print("-=" * 30 + "-")
print(f"Running image batch index: {image_bindex}")
print("-=" * 30 + "-")
torch.cuda.empty_cache()
# -- sample & unpack batch --
sample = next(image_iter)
sample_to_cuda(sample)
dyn_noisy = sample['noisy'] # dynamics and noise
dyn_clean = sample['burst'] # dynamics and no noise
static_noisy = sample['snoisy'] # no dynamics and noise
static_clean = sample['sburst'] # no dynamics and no noise
flow_gt = sample['flow']
image_index = sample['index']
tl_index = sample['tl_index']
rng_state = sample['rng_state']
if cfg.noise_params.ntype == "pn":
dyn_noisy = anscombe.forward(dyn_noisy)
# dyn_nosiy = dyn_clean
dyn_noisy = static_clean
# -- shape info --
T, B, C, H, W = dyn_noisy.shape
isize = edict({'h': H, 'w': W})
ref_t = nframes // 2
nimages, npix, nframes = B, H * W, T
# -- create results dict --
flows = edict()
aligned = edict()
runtimes = edict()
optimal_scores = edict() # score function at optimal
# -- groundtruth flow --
flow_gt_rs = rearrange(flow_gt, 'i tm1 two -> i 1 tm1 two')
blocks_gt = flow_to_blocks(flow_gt_rs, nblocks)
flows.of = repeat(flow_gt, 'i tm1 two -> i p tm1 two', p=npix)
aligned.of = align_from_flow(dyn_clean, flows.of, nblocks, isize=isize)
runtimes.of = 0. # given
optimal_scores.of = np.zeros(
(nimages, npix, nframes)) # clean target is zero
aligned.clean = static_clean
# -- compute nearest neighbor fields [global] --
start_time = time.perf_counter()
shape_str = 't b h w two -> b (h w) t two'
nnf_vals, nnf_pix = nnf.compute_burst_nnf(dyn_clean, ref_t, patchsize)
nnf_pix_best = torch.LongTensor(
rearrange(nnf_pix[..., 0, :], shape_str))
nnf_pix_best = torch.LongTensor(nnf_pix_best)
flows.nnf = pix_to_flow(nnf_pix_best)
aligned.nnf = align_from_pix(dyn_clean, nnf_pix_best, nblocks)
runtimes.nnf = time.perf_counter() - start_time
optimal_scores.nnf = np.zeros(
(nimages, npix, nframes)) # clean target is zero
# -- compute proposed search of nnf --
print("[Bootstrap] loss function")
iterations = 1
K, subsizes = get_boot_hyperparams(cfg.nframes, cfg.nblocks)
start_time = time.perf_counter()
optim = AlignOptimizer("v3")
# flows.est = optim.run(dyn_noisy,patchsize,eval_prop,
# nblocks,iterations,subsizes,K)
flows.est = flows.of.clone()
aligned.est = align_from_flow(dyn_clean,
flows.est,
patchsize,
isize=isize)
runtimes.est = time.perf_counter() - start_time
# -- load adjacent blocks --
nframes, nimages, ncolor, h, w = dyn_noisy.shape
nsegs = h * w
brange = exh_block_range(nimages, nsegs, nframes, nblocks)
ref_block = get_ref_block(nblocks)
curr_blocks = init_optim_block(nimages, nsegs, nframes, nblocks)
curr_blocks = curr_blocks[:, :, :,
None] # nimages, nsegs, nframes, naligns
frames = np.r_[np.arange(nframes // 2),
np.arange(nframes // 2 + 1, nframes)]
frames = repeat(frames, 't -> i s t', i=nimages, s=nsegs)
search_blocks = get_search_blocks(frames, brange, curr_blocks,
f'cuda:{cfg.gpuid}')
print("search_blocks.shape ", search_blocks.shape)
init_blocks = rearrange(curr_blocks,
'i s t a -> i s a t').to(dyn_noisy.device)
print("init_blocks.shape ", init_blocks.shape)
search_blocks = search_blocks[0, 0]
init_blocks_ = init_blocks[0, 0]
search_blocks = eval_plimb.filter_blocks_to_1skip_neighbors(
search_blocks, init_blocks_)
search_blocks = repeat(search_blocks,
'a t -> i s a t',
i=nimages,
s=nsegs)
print("search_blocks.shape ", search_blocks.shape)
# -- compute MSE for the batch --
est = edict()
bscf = edict()
plimb = edict()
print("curr_blocks.shape ", curr_blocks.shape)
eval_prop.score_fxn_name = ""
scores, scores_t, blocks = eval_mse.score_burst_from_blocks(
dyn_noisy, search_blocks, patchsize, nblocks)
est.scores = scores
est.scores_t = scores_t
est.blocks = blocks
print("Done with est.")
# -- compute bootrapping in closed form --
scores, scores_t, blocks = eval_prop_cf.score_burst_from_blocks(
dyn_noisy, search_blocks, patchsize, nblocks)
bscf.scores = scores
bscf.scores_t = scores_t
bscf.blocks = blocks
# -- compute bootstrapping for the batch --
print("Get init block from original bootstrap.")
print(init_blocks[0, 0])
scores, scores_t, blocks = eval_prop.score_burst_from_blocks(
dyn_noisy, init_blocks, patchsize, nblocks)
print("Starting prop.")
eval_prop.score_fxn_name = "bs"
eval_prop.score_cfg.bs_type = ""
state = edict({'scores': scores, 'blocks': blocks})
scores, scores_t, blocks = eval_plimb.score_burst_from_blocks(
dyn_noisy, search_blocks, state, patchsize, nblocks)
plimb.scores = scores
plimb.scores_t = scores_t
plimb.blocks = blocks
print(est.scores.shape)
print(bscf.scores.shape)
print(state.scores.shape)
print(plimb.scores.shape)
diff_plimb = plimb.scores[0] - est.scores[0]
perc_delta = torch.abs(diff_plimb) / est.scores[0]
diff_bscf = bscf.scores[0] - est.scores[0]
perc_delta_cf = torch.abs(diff_bscf) / est.scores[0]
pix_idx_list = [0, 20, 30] #np.arange(h*w)
for p in pix_idx_list:
print("-" * 10 + f" @ {p}")
print("est", est.scores[0, p].cpu().numpy())
print("state", state.scores[0, p].cpu().numpy())
print("bscf", bscf.scores[0, p].cpu().numpy())
print("plimb", plimb.scores[0, p].cpu().numpy())
print("plimb/est", plimb.scores[0, p] / est.scores[0, p])
print("plimb - est", plimb.scores[0, p] - est.scores[0, p])
print("plimb - bscf", plimb.scores[0, p] - bscf.scores[0, p])
print("%Delta [plimb]", perc_delta[p])
print("L2-Norm [plimb]", torch.sum(diff_plimb[p]**2))
print("Nmlz L2-Norm [plimb]", torch.mean(diff_plimb[p]**2))
print("%Delta [bscf]", perc_delta_cf[p])
print("L2-Norm [bscf]", torch.sum(diff_bscf[p]**2))
print("Nmlz L2-Norm [bscf]", torch.mean(diff_bscf[p]**2))
print("[Overall: plimb] %Delta: ", torch.mean(perc_delta).item())
print("[Overall: plimb] L2-Norm: ", torch.sum(diff_plimb**2).item())
print("[Overall: plimb] Nmlz L2-Norm: ",
torch.mean(diff_plimb**2).item())
print("[Overall: bscf] %Delta: ", torch.mean(perc_delta_cf).item())
print("[Overall: bscf] L2-Norm: ", torch.sum(diff_bscf**2).item())
print("[Overall: bscf] Nmlz L2-Norm: ",
torch.mean(diff_bscf**2).item())
# -- format results --
pad = 3 #2*(nframes-1)*ppf+4
isize = edict({'h': H - pad, 'w': W - pad})
# -- flows to numpy --
is_even = cfg.frame_size % 2 == 0
mid_pix = cfg.frame_size * cfg.frame_size // 2 + (cfg.frame_size //
2) * is_even
mid_pix = 32 * 10 + 23
flows_np = edict_torch_to_numpy(flows)
# -- End-Point-Errors --
epes_of = compute_flows_epe_wrt_ref(flows, "of")
epes_nnf = compute_flows_epe_wrt_ref(flows, "nnf")
epes_nnf_local = compute_flows_epe_wrt_ref(flows, "nnf_local")
nnf_acc = compute_acc_wrt_ref(flows, "nnf")
nnf_local_acc = compute_acc_wrt_ref(flows, "nnf_local")
# -- PSNRs --
aligned = remove_frame_centers(aligned)
psnrs = compute_frames_psnr(aligned, isize)
# -- print report ---
print("\n" * 3) # banner
print("-" * 25 + " Results " + "-" * 25)
print_dict_ndarray_0_midpix(flows_np, mid_pix)
print_runtimes(runtimes)
print_verbose_psnrs(psnrs)
print_delta_summary_psnrs(psnrs)
print_verbose_epes(epes_of, epes_nnf)
print_nnf_acc(nnf_acc)
print_nnf_local_acc(nnf_local_acc)
print_summary_epes(epes_of, epes_nnf)
print_summary_psnrs(psnrs)
# -- prepare results to be appended --
psnrs = edict_torch_to_numpy(psnrs)
epes_of = edict_torch_to_numpy(epes_of)
epes_nnf = edict_torch_to_numpy(epes_nnf)
epes_nnf_local = edict_torch_to_numpy(epes_nnf_local)
nnf_acc = edict_torch_to_numpy(nnf_acc)
nnf_local_acc = edict_torch_to_numpy(nnf_local_acc)
image_index = torch_to_numpy(image_index)
batch_results = {
'runtimes': runtimes,
'optimal_scores': optimal_scores,
'psnrs': psnrs,
'epes_of': epes_of,
'epes_nnf': epes_nnf,
'epes_nnf_local': epes_nnf_local,
'nnf_acc': nnf_acc,
'nnf_local_acc': nnf_local_acc
}
# -- format results --
batch_results = flatten_internal_dict(batch_results)
format_fields(batch_results, image_index, rng_state)
print("shape check.")
for key, value in batch_results.items():
print(key, value.shape)
record.append(batch_results)
# print("\n"*3)
# print("-"*20)
# print(record.record)
# print("-"*20)
# print("\n"*3)
# record.stack_record()
record.cat_record()
# print("\n"*3)
# print("-"*20)
# print(record.record)
# print("-"*20)
print("\n" * 3)
print("\n" * 3)
print("-" * 20)
# df = pd.DataFrame().append(record.record,ignore_index=True)
for key, val in record.record.items():
print(key, val.shape)
# print(df)
print("-" * 20)
print("\n" * 3)
return record.record
def test_nnf():
# -- get config --
cfg = config()
print("Config for Testing.")
print(cfg)
# -- set random seed --
set_seed(cfg.random_seed)
# -- load dataset --
data, loaders = load_image_dataset(cfg)
image_iter = iter(loaders.tr)
nskips = 2 + 4 + 2 + 4 + 1
for skip in range(nskips):
next(image_iter)
# -- get score function --
score_fxn_ave = get_score_function("ave")
score_fxn_bs = get_score_function(cfg.score_fxn_name)
# -- some constants --
NUM_BATCHES = 10
nframes, nblocks = cfg.nframes, cfg.nblocks
patchsize = cfg.patchsize
ppf = cfg.dynamic_info.ppf
check_parameters(nblocks, patchsize)
# -- create evaluator for ave; simple --
iterations, K = 1, 1
subsizes = []
block_batchsize = 256
eval_ave_simp = EvalBlockScores(score_fxn_ave, "ave", patchsize,
block_batchsize, None)
# -- create evaluator for ave --
iterations, K = 1, 1
subsizes = []
eval_ave = EvalBlockScores(score_fxn_ave, "ave", patchsize,
block_batchsize, None)
# -- create evaluator for bootstrapping --
block_batchsize = 64
eval_prop = EvalBlockScores(score_fxn_bs, "bs", patchsize, block_batchsize,
None)
# -- iterate over images --
for image_bindex in range(NUM_BATCHES):
print("-=" * 30 + "-")
print(f"Running image batch index: {image_bindex}")
print("-=" * 30 + "-")
# -- sample & unpack batch --
sample = next(image_iter)
sample_to_cuda(sample)
dyn_noisy = sample['noisy'] # dynamics and noise
dyn_clean = sample['burst'] # dynamics and no noise
static_noisy = sample['snoisy'] # no dynamics and noise
static_clean = sample['sburst'] # no dynamics and no noise
flow_gt = sample['ref_flow']
# flow_gt = sample['seq_flow']
if cfg.noise_params.ntype == "pn":
dyn_noisy = anscombe.forward(dyn_noisy)
# -- shape info --
T, B, C, H, W = dyn_noisy.shape
isize = edict({'h': H, 'w': W})
ref_t = nframes // 2
npix = H * W
# -- groundtruth flow --
# print("flow_gt",flow_gt)
flow_gt_rs = rearrange(flow_gt, 'i tm1 two -> i 1 tm1 two')
blocks_gt = flow_to_blocks(flow_gt_rs, nblocks)
# print("\n\n")
# print("flow_gt[0,0] ",flow_gt)
# print("blocks_gt[0,0] ",blocks_gt[0,0])
flow_gt = repeat(flow_gt, 'i tm1 two -> i p tm1 two', p=npix)
aligned_of = align_from_flow(dyn_clean, flow_gt, nblocks, isize=isize)
pix_gt = flow_to_pix(flow_gt.clone(), isize=isize)
# -- compute nearest neighbor fields --
start_time = time.perf_counter()
shape_str = 't b h w two -> b (h w) t two'
nnf_vals, nnf_pix = nnf.compute_burst_nnf(dyn_clean, ref_t, patchsize)
nnf_pix_best = torch.LongTensor(
rearrange(nnf_pix[..., 0, :], shape_str))
nnf_pix_best = torch.LongTensor(nnf_pix_best)
pix_nnf = nnf_pix_best.clone()
flow_nnf = pix_to_flow(nnf_pix_best)
aligned_nnf = align_from_pix(dyn_clean, nnf_pix_best, nblocks)
time_nnf = time.perf_counter() - start_time
# -- compute proposed search of nnf --
start_time = time.perf_counter()
print(dyn_noisy.shape)
# split_vals,split_pix = nnf.compute_burst_nnf(dyn_noisy,ref_t,patchsize)
split_pix = np.copy(nnf_pix)
split_pix_best = torch.LongTensor(
rearrange(split_pix[..., 0, :], shape_str))
split_pix_best = torch.LongTensor(split_pix_best)
pix_split = split_pix_best.clone()
flow_split = pix_to_flow(split_pix_best)
aligned_split = align_from_pix(dyn_clean, split_pix_best, nblocks)
time_split = time.perf_counter() - start_time
# -- compute simple ave --
iterations, K = 0, 1
subsizes = []
print("[simple] Ave loss function")
start_time = time.perf_counter()
optim = AlignOptimizer("v3")
# flow_ave_simp = optim.run(dyn_noisy,patchsize,eval_ave_simp,
# nblocks,iterations,subsizes,K)
flow_ave_simp = flow_gt.clone().cpu()
aligned_ave_simp = align_from_flow(dyn_clean,
flow_ave_simp,
nblocks,
isize=isize)
time_ave_simp = time.perf_counter() - start_time
print(flow_ave_simp.shape)
# -- compute complex ave --
iterations, K = 0, 1
subsizes = []
print("[complex] Ave loss function")
start_time = time.perf_counter()
optim = AlignOptimizer("v3")
flow_ave = optim.run(dyn_noisy, patchsize, eval_ave, nblocks,
iterations, subsizes, K)
# flow_ave = flow_gt.clone()
pix_ave = flow_to_pix(flow_ave.clone(), isize=isize)
aligned_ave = align_from_flow(dyn_clean,
flow_ave,
nblocks,
isize=isize)
time_ave = time.perf_counter() - start_time
# -- compute proposed search of nnf --
# iterations,K = 50,3
# subsizes = [2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2]
#iterations,K = 1,nblocks**2
# K is a function of noise level.
# iterations,K = 1,nblocks**2
iterations, K = 1, 2 * nblocks #**2
# subsizes = [3]#,3,3,3,3,3,3,3,3,3]
# subsizes = [3,3,3,3,3,3,3,]
subsizes = [3, 3, 3, 3, 3, 3, 3, 3]
# subsizes = [nframes]
# subsizes = [nframes]
print("[Bootstrap] loss function")
start_time = time.perf_counter()
optim = AlignOptimizer("v3")
flow_est = optim.run(dyn_noisy, patchsize, eval_prop, nblocks,
iterations, subsizes, K)
pix_est = flow_to_pix(flow_est.clone(), isize=isize)
aligned_est = align_from_flow(dyn_clean,
flow_est,
patchsize,
isize=isize)
time_est = time.perf_counter() - start_time
# flow_est = flow_gt.clone()
# aligned_est = aligned_of.clone()
# time_est = 0.
# -- banner --
print("\n" * 3)
print("-" * 25 + " Results " + "-" * 25)
# -- examples of flow --
print("-" * 50)
is_even = cfg.frame_size % 2 == 0
mid_pix = cfg.frame_size * cfg.frame_size // 2 + (cfg.frame_size //
2) * is_even
mid_pix = 32 * 10 + 23
# mid_pix = 32*23+10
flow_gt_np = torch_to_numpy(flow_gt)
flow_nnf_np = torch_to_numpy(flow_nnf)
flow_split_np = torch_to_numpy(flow_split)
flow_ave_simp_np = torch_to_numpy(flow_ave_simp)
flow_ave_np = torch_to_numpy(flow_ave)
flow_est_np = torch_to_numpy(flow_est)
print(flow_gt_np[0, mid_pix])
print(flow_nnf_np[0, mid_pix])
print(flow_split_np[0, mid_pix])
print(flow_ave_simp_np[0, mid_pix])
print(flow_ave_np[0, mid_pix])
print(flow_est_np[0, mid_pix])
print("-" * 50)
pix_gt_np = torch_to_numpy(pix_gt)
pix_nnf_np = torch_to_numpy(pix_nnf)
pix_ave_np = torch_to_numpy(pix_ave)
pix_est_np = torch_to_numpy(pix_est)
print(pix_gt_np[0, mid_pix])
print(pix_nnf_np[0, mid_pix])
print(pix_ave_np[0, mid_pix])
print(pix_est_np[0, mid_pix])
# print(aligned_of[0,0,:,10,23].cpu() - static_clean[0,0,:,10,23].cpu())
# print(aligned_ave[0,0,:,10,23].cpu() - static_clean[0,0,:,10,23].cpu())
# print(aligned_of[0,0,:,23,10].cpu() - static_clean[0,0,:,23,10].cpu())
# print(aligned_ave[0,0,:,23,10].cpu() - static_clean[0,0,:,23,10].cpu())
print("-" * 50)
# -- compare compute time --
print("-" * 50)
print("Compute Time [smaller is better]")
print("-" * 50)
print("[NNF]: %2.3e" % time_nnf)
print("[Split]: %2.3e" % time_split)
print("[Ave [Simple]]: %2.3e" % time_ave_simp)
print("[Ave]: %2.3e" % time_ave)
print("[Proposed]: %2.3e" % time_est)
# -- compare gt v.s. nnf computations --
nnf_of = compute_epe(flow_nnf, flow_gt)
split_of = compute_epe(flow_split, flow_gt)
ave_simp_of = compute_epe(flow_ave_simp, flow_gt)
ave_of = compute_epe(flow_ave, flow_gt)
est_of = compute_epe(flow_est, flow_gt)
split_nnf = compute_epe(flow_split, flow_nnf)
ave_simp_nnf = compute_epe(flow_ave_simp, flow_nnf)
ave_nnf = compute_epe(flow_ave, flow_nnf)
est_nnf = compute_epe(flow_est, flow_nnf)
# -- End-Point-Errors --
print("-" * 50)
print("EPE Errors [smaller is better]")
print("-" * 50)
print("NNF v.s. Optical Flow.")
print(nnf_of)
print("Split v.s. Optical Flow.")
print(split_of)
print("Ave [Simple] v.s. Optical Flow.")
print(ave_simp_of)
print("Ave v.s. Optical Flow.")
print(ave_of)
print("Proposed v.s. Optical Flow.")
print(est_of)
print("Split v.s. NNF")
print(split_nnf)
print("Ave [Simple] v.s. NNF")
print(ave_simp_nnf)
print("Ave v.s. NNF")
print(ave_nnf)
print("Proposed v.s. NNF")
print(est_nnf)
# -- compare accuracy of method nnf v.s. actual nnf --
def compute_flow_acc(guess, gt):
both = torch.all(guess.type(torch.long) == gt.type(torch.long),
dim=-1)
ncorrect = torch.sum(both)
acc = 100 * float(ncorrect) / both.numel()
return acc
split_nnf_acc = compute_flow_acc(flow_split, flow_nnf)
ave_simp_nnf_acc = compute_flow_acc(flow_ave_simp, flow_nnf)
ave_nnf_acc = compute_flow_acc(flow_ave, flow_nnf)
est_nnf_acc = compute_flow_acc(flow_est, flow_nnf)
# -- PSNR to Reference Image --
pad = 2 * (nframes - 1) * ppf + 4
isize = edict({'h': H - pad, 'w': W - pad})
# print("isize: ",isize)
aligned_of = remove_center_frame(aligned_of)
aligned_nnf = remove_center_frame(aligned_nnf)
aligned_split = remove_center_frame(aligned_split)
aligned_ave_simp = remove_center_frame(aligned_ave_simp)
aligned_ave = remove_center_frame(aligned_ave)
aligned_est = remove_center_frame(aligned_est)
static_clean = remove_center_frame(static_clean)
psnr_of = compute_aligned_psnr(aligned_of, static_clean, isize)
psnr_nnf = compute_aligned_psnr(aligned_nnf, static_clean, isize)
psnr_split = compute_aligned_psnr(aligned_split, static_clean, isize)
psnr_ave_simp = compute_aligned_psnr(aligned_ave_simp, static_clean,
isize)
psnr_ave = compute_aligned_psnr(aligned_ave, static_clean, isize)
psnr_est = compute_aligned_psnr(aligned_est, static_clean, isize)
print("-" * 50)
print("PSNR Values [bigger is better]")
print("-" * 50)
print("Optical Flow [groundtruth v1]")
print(psnr_of)
print("NNF [groundtruth v2]")
print(psnr_nnf)
print("Split [old method]")
print(psnr_split)
print("Ave [simple; old method]")
print(psnr_ave_simp)
print("Ave [old method]")
print(psnr_ave)
print("Proposed [new method]")
print(psnr_est)
# -- print nnf accuracy here --
print("-" * 50)
print("NNF Accuracy [bigger is better]")
print("-" * 50)
print("Split v.s. NNF")
print(split_nnf_acc)
print("Ave [Simple] v.s. NNF")
print(ave_simp_nnf_acc)
print("Ave v.s. NNF")
print(ave_nnf_acc)
print("Proposed v.s. NNF")
print(est_nnf_acc)
# -- location of PSNR errors --
csize = 30
# aligned_of = torch_to_numpy(tvF.center_crop(aligned_of,(csize,csize)))
# aligned_ave = torch_to_numpy(tvF.center_crop(aligned_ave,(csize,csize)))
# static_clean = torch_to_numpy(tvF.center_crop(static_clean,(csize,csize)))
flow_gt = torch_to_numpy(flow_gt)
flow_ave = torch_to_numpy(flow_ave)
aligned_of = torch_to_numpy(aligned_of)
aligned_ave = torch_to_numpy(aligned_ave)
static_clean = torch_to_numpy(static_clean)
# print("WHERE?")
# print("OF")
# print(aligned_of.shape)
# for row in range(30):
# print(np.abs(aligned_of[0,0,0,row]- static_clean[0,0,0,row]))
# print(np.where(~np.isclose(aligned_of,aligned_of)))
# print(np.where(~np.isclose(flow_gt,flow_ave)))
# print(np.where(~np.isclose(aligned_of,aligned_of)))
# print(np.where(~np.isclose(aligned_of,static_clean)))
# print("Ave")
# indices = np.where(~np.isclose(aligned_ave,static_clean))
# row,col = indices[-2:]
# for elem in range(len(row)):
# print(np.c_[row,col][elem])
# print(np.where(~np.isclose(aligned_ave,static_clean)))
# -- Summary of End-Point-Errors --
print("-" * 50)
print("Summary of EPE Errors [smaller is better]")
print("-" * 50)
print("[NNF v.s. Optical Flow]: %2.3f" % nnf_of.mean().item())
print("[Split v.s. Optical Flow]: %2.3f" % split_of.mean().item())
print("[Ave [Simple] v.s. Optical Flow]: %2.3f" %
ave_simp_of.mean().item())
print("[Ave v.s. Optical Flow]: %2.3f" % ave_of.mean().item())
print("[Proposed v.s. Optical Flow]: %2.3f" % est_of.mean().item())
print("[Split v.s. NNF]: %2.3f" % split_nnf.mean().item())
print("[Ave [Simple] v.s. NNF]: %2.3f" % ave_simp_nnf.mean().item())
print("[Ave v.s. NNF]: %2.3f" % ave_nnf.mean().item())
print("[Proposed v.s. NNF]: %2.3f" % est_nnf.mean().item())
# -- Summary of PSNR to Reference Image --
print("-" * 50)
print("Summary PSNR Values [bigger is better]")
print("-" * 50)
print("[Optical Flow]: %2.3f" % psnr_of.mean().item())
print("[NNF]: %2.3f" % psnr_nnf.mean().item())
print("[Split]: %2.3f" % psnr_split.mean().item())
print("[Ave [Simple]]: %2.3f" % psnr_ave_simp.mean().item())
print("[Ave]: %2.3f" % psnr_ave.mean().item())
print("[Proposed]: %2.3f" % psnr_est.mean().item())
print("-" * 50)
print("PSNR Comparisons [smaller is better]")
print("-" * 50)
delta_split = psnr_nnf - psnr_split
delta_ave_simp = psnr_nnf - psnr_ave_simp
delta_ave = psnr_nnf - psnr_ave
delta_est = psnr_nnf - psnr_est
print("ave([NNF] - [Split]): %2.3f" % delta_split.mean().item())
print("ave([NNF] - [Ave [Simple]]): %2.3f" %
delta_ave_simp.mean().item())
print("ave([NNF] - [Ave]): %2.3f" % delta_ave.mean().item())
print("ave([NNF] - [Proposed]): %2.3f" % delta_est.mean().item())
def execute_experiment(cfg):
# -- init exp! --
print("RUNNING EXP.")
print(cfg)
# -- create results record to save --
dims = {
'batch_results': None,
'batch_to_record': None,
'record_results': {
'default': 0
},
'stack': {
'default': 0
},
'cat': {
'default': 0
}
}
record = cache_io.ExpRecord(dims)
# -- set random seed --
set_seed(cfg.random_seed)
# -- load dataset --
# data,loaders = load_image_dataset(cfg)
data, loaders = load_dataset(cfg, cfg.dataset.mode)
image_iter = iter(loaders.tr)
# -- get score function --
score_fxn_ave = get_score_function("ave")
score_fxn_bs = get_score_function(cfg.score_fxn_name)
# -- some constants --
NUM_BATCHES = 10
nframes, nblocks = cfg.nframes, cfg.nblocks
patchsize = cfg.patchsize
ps = patchsize
ppf = cfg.dynamic_info.ppf
check_parameters(nblocks, patchsize)
# -- theory constants --
std = cfg.noise_params.g.std / 255.
p = cfg.patchsize**2 * 3
t = cfg.nframes
theory = edict()
theory.c2 = ((t - 1) / t)**2 * std**2 + (t - 1) / t**2 * std**2
theory.mean = theory.c2
theory.mode = (1 - 2 / p) * theory.c2
theory.var = 2 / p * theory.c2**2
theory.std = np.sqrt(theory.var)
pp.pprint(theory)
# npn = no patch normalization
theory_npn = edict()
theory_npn.c2 = ((t - 1) / t)**2 * std**2 + (t - 1) / t**2 * std**2
theory_npn.mean = theory_npn.c2 * p
theory_npn.mode = (1 - 2 / p) * theory_npn.c2 * p
theory_npn.var = 2 * theory_npn.c2**2 * p
theory_npn.std = np.sqrt(theory_npn.var)
pp.pprint(theory_npn)
# oracle = clean reference frame
theory_oracle = edict()
theory_oracle.c2 = std**2
theory_oracle.mean = theory_oracle.c2 * p
theory_oracle.mode = (1 - 2 / p) * theory_oracle.c2 * p
theory_oracle.var = 2 * theory_oracle.c2**2 * p
theory_oracle.std = np.sqrt(theory_oracle.var)
pp.pprint(theory_oracle)
# -- create evaluator for ave; simple --
iterations, K = 1, 1
subsizes = []
block_batchsize = 32
eval_ave_simp = EvalBlockScores(score_fxn_ave, "ave", patchsize,
block_batchsize, None)
# -- create evaluator for ave --
iterations, K = 1, 1
subsizes = []
eval_ave = EvalBlockScores(score_fxn_ave, "ave", patchsize,
block_batchsize, None)
# -- create evaluator for bootstrapping --
block_batchsize = 32
eval_prop = EvalBlockScores(score_fxn_bs, "bs", patchsize, block_batchsize,
None)
# -- init flownet model --
cfg.gpuid = 1 - cfg.gpuid # flip. flop.
flownet_align = get_align_method(cfg, "flownet_v2", comp_align=False)
cfg.gpuid = 1 - cfg.gpuid # flippity flop.
# -- get an image transform --
image_xform = get_image_xform(cfg.image_xform, cfg.gpuid, cfg.frame_size)
blockLabels, _ = nnf_utils.getBlockLabels(None, nblocks, np.int32,
cfg.device, True)
# -- iterate over images --
NUM_BATCHES = min(NUM_BATCHES, len(image_iter))
for image_bindex in range(NUM_BATCHES):
print("-=" * 30 + "-")
print(f"Running image batch index: {image_bindex}")
print("-=" * 30 + "-")
torch.cuda.empty_cache()
# -- sample & unpack batch --
nwaste = 0
for w in range(nwaste):
sample = next(image_iter) # waste one
sample = next(image_iter)
sample_to_cuda(sample)
convert_keys(sample)
torch.cuda.synchronize()
# for key,val in sample.items():
# print(key,type(val))
# if torch.is_tensor(val):
# print(key,val.device)
dyn_noisy = sample['dyn_noisy'] # dynamics and noise
dyn_clean = sample['dyn_clean'] # dynamics and no noise
static_noisy = sample['static_noisy'] # no dynamics and noise
static_clean = sample['static_clean'] # no dynamics and no noise
nnf_gt = sample['nnf']
flow_gt = sample['flow']
if nnf_gt.ndim == 6:
nnf_gt = nnf_gt[:, :, 0] # pick top 1 out of K
image_index = sample['image_index']
rng_state = sample['rng_state']
# TODO: anscombe is a type of image transform
if not (image_xform is None):
dyn_clean_ftrs = image_xform(dyn_clean)
dyn_noisy_ftrs = image_xform(dyn_noisy)
else:
dyn_clean_ftrs = dyn_clean
dyn_noisy_ftrs = dyn_noisy
if "resize" in cfg.image_xform:
vprint("Images, Flows, and NNF Modified.")
dyn_clean = image_xform(dyn_clean)
dyn_noisy = image_xform(dyn_noisy)
T, B, C, H, W = dyn_noisy.shape
flow_gt = torch.zeros((B, 1, T, H, W, 2))
nnf_gt = torch.zeros((1, T, B, H, W, 2))
save_image(dyn_clean, "dyn_clean.png")
# print("SHAPES")
# print(dyn_noisy.shape)
# print(dyn_clean.shape)
# print(nnf_gt.shape)
# -- shape info --
pad = cfg.nblocks // 2 + cfg.patchsize // 2
T, B, C, H, W = dyn_noisy.shape
isize = edict({'h': H, 'w': W})
psize = edict({'h': H - 2 * pad, 'w': W - 2 * pad})
ref_t = nframes // 2
nimages, npix, nframes = B, H * W, T
frame_size = [H, W]
ifsize = [H - 2 * pad, W - 2 * pad]
print("flow_gt.shape: ", flow_gt.shape)
print("flow_gt: ", flow_gt[0, 0, :, H // 2, W // 2, :])
# -- create results dict --
pixs = edict()
flows = edict()
anoisy = edict()
aligned = edict()
runtimes = edict()
optimal_scores = edict() # score function at optimal
# -- compute proposed search of nnf --
# ave = torch.mean(dyn_noisy_ftrs[:,0,:,4:4+ps,4:4+ps],dim=0)
# frames = dyn_noisy_ftrs[:,0,:,4:4+ps,4:4+ps]
# gt_offset = torch.sum((frames - ave)**2/nframes).item()
# print("Optimal: ",gt_offset)
# gt_offset = -1.
# -- FIND MODE of BURST --
vprint("Our Method")
flow_fmt = rearrange(flow_gt, 'i 1 t h w two -> t i h w 1 two')
locs_fmt = flow2locs(flow_fmt)
print("locs_fmt.shape: ", locs_fmt.shape)
print(dyn_noisy_ftrs.min(), dyn_noisy_ftrs.max())
vals, _ = evalAtLocs(dyn_noisy_ftrs,
locs_fmt,
patchsize,
nblocks,
return_mode=False)
vals = torch.zeros_like(vals)
# flow_fmt = rearrange(flow_gt,'i t h w two -> i (h w) t two')
# vals,_ = bnnf_utils.evalAtFlow(dyn_noisy_ftrs, flow_fmt, patchsize,
# nblocks, return_mode=False)
mode = mode_vals(vals, ifsize)
cc_vals = vals[0, 5:29, 5:29, 0].ravel()
vstd = torch.std(cc_vals).item()
print("[SubBurst] Computed Mode: ", mode)
print("[SubBurst] Computed Std: ", vstd)
# -- compute proposed search of nnf --
vprint("dyn_noisy_ftrs.shape ", dyn_noisy_ftrs.shape)
valMean = theory_npn.mode
vprint("valMean: ", valMean)
start_time = time.perf_counter()
if cfg.nframes < 5:
_, flows.est = bnnf_utils.runBurstNnf(dyn_noisy_ftrs,
patchsize,
nblocks,
k=1,
valMean=valMean,
blockLabels=None,
fmt=True,
to_flow=True)
else:
flows.est = rearrange(flow_gt,
'i 1 t h w two -> 1 i (h w) t two').clone()
flows.est = flows.est[0]
runtimes.est = time.perf_counter() - start_time
pixs.est = flow_to_pix(flows.est.clone(), nframes, isize=isize)
aligned.est = align_from_flow(dyn_clean,
flows.est,
patchsize,
isize=isize)
if cfg.nframes > 7: aligned.est = torch.zeros_like(aligned.est)
anoisy.est = align_from_flow(dyn_noisy,
flows.est,
patchsize,
isize=isize)
optimal_scores.est = np.zeros((nimages, npix, 1, nframes))
# -- the proposed method --
std = cfg.noise_params.g.std
start_time = time.perf_counter()
_flow = flow_gt.clone()
# _,_flow = runKmSearch(dyn_noisy_ftrs, patchsize, nblocks, k = 1,
# std = std/255.,mode="cuda")
runtimes.kmb = time.perf_counter() - start_time
flows.kmb = rearrange(_flow, 'i 1 t h w two -> i (h w) t two')
pixs.kmb = flow_to_pix(flows.kmb.clone(), nframes, isize=isize)
aligned.kmb = align_from_flow(dyn_clean, flows.kmb, 0, isize=isize)
optimal_scores.kmb = torch_to_numpy(optimal_scores.est)
# -- compute proposed search of nnf --
vprint("Our BpSearch Method")
# print(flow_gt)
# std = cfg.noise_params.g.std/255.
valMean = theory_npn.mode
start_time = time.perf_counter()
bp_nblocks = 3
# _,bp_est,a_noisy = runBpSearch(dyn_noisy_ftrs, dyn_noisy_ftrs,
# patchsize, bp_nblocks, k = 1,
# valMean = valMean, std=std,
# blockLabels=None,
# l2_nblocks=nblocks,
# fmt = True, to_flow=True,
# search_type=cfg.bp_type,
# gt_info={'flow':flow_gt})
bp_est = flows.est[None, :].clone()
flows.bp_est = bp_est[0]
# flows.bp_est = rearrange(flow_gt,'i t h w two -> i (h w) t two')
runtimes.bp_est = time.perf_counter() - start_time
pixs.bp_est = flow_to_pix(flows.bp_est.clone(), nframes, isize=isize)
# aligned.bp_est = a_clean
aligned.bp_est = align_from_flow(dyn_clean,
flows.bp_est,
patchsize,
isize=isize)
anoisy.bp_est = align_from_flow(dyn_noisy,
flows.bp_est,
patchsize,
isize=isize)
optimal_scores.bp_est = np.zeros((nimages, npix, 1, nframes))
# -- compute proposed search of nnf [with tiling ]--
vprint("Our Burst Method (Tiled)")
valMean = 0.
start_time = time.perf_counter()
if cfg.nframes < 5:
_, flows.est_tile = bnnf_utils.runBurstNnf(dyn_noisy_ftrs,
patchsize,
nblocks,
k=1,
valMean=valMean,
blockLabels=None,
fmt=True,
to_flow=True,
tile_burst=True)
else:
flows.est_tile = rearrange(
flow_gt, 'i 1 t h w two -> 1 i (h w) t two').clone()
flows.est_tile = flows.est_tile[0]
# flows.est_tile = rearrange(flow_gt,'i t h w two -> i (h w) t two')
runtimes.est_tile = time.perf_counter() - start_time
pixs.est_tile = flow_to_pix(flows.est_tile.clone(),
nframes,
isize=isize)
aligned.est_tile = align_from_flow(dyn_clean,
flows.est_tile,
patchsize,
isize=isize)
if cfg.nframes > 7:
aligned.est_tile = torch.zeros_like(aligned.est_tile)
anoisy.est_tile = align_from_flow(dyn_noisy,
flows.est_tile,
patchsize,
isize=isize)
optimal_scores.est_tile = np.zeros((nimages, npix, 1, nframes))
# -- compute new est method --
vprint("[Burst-LK] loss function")
vprint(flow_gt.shape)
# print(flow_gt[0,:3,32,32,:])
vprint(flow_gt.shape)
start_time = time.perf_counter()
if frame_size[0] <= 64 and cfg.nblocks < 10 and True:
flows.blk = burstNnf.run(dyn_noisy_ftrs, patchsize, nblocks)
else:
flows.blk = rearrange(flow_gt, 'i 1 t h w two -> i (h w) t two')
runtimes.blk = time.perf_counter() - start_time
pixs.blk = flow_to_pix(flows.blk.clone(), nframes, isize=isize)
aligned.blk = align_from_flow(dyn_clean,
flows.blk,
patchsize,
isize=isize)
optimal_scores.blk = np.zeros((nimages, npix, 1, nframes))
# optimal_scores.blk = eval_prop.score_burst_from_flow(dyn_noisy,flows.nnf_local,
# patchsize,nblocks)[1]
optimal_scores.blk = torch_to_numpy(optimal_scores.blk)
# -- compute new est method --
vprint("Oracle")
vprint(flow_gt.shape)
# print(flow_gt[0,:3,32,32,:])
vprint(flow_gt.shape)
valMean = theory_oracle.mode
oracle_burst = dyn_noisy_ftrs.clone()
oracle_burst[nframes // 2] = dyn_clean_ftrs[nframes // 2]
start_time = time.perf_counter()
vals_oracle, pix_oracle = nnf_utils.runNnfBurst(
oracle_burst,
patchsize,
nblocks,
1,
valMean=valMean,
blockLabels=blockLabels)
runtimes.oracle = time.perf_counter() - start_time
pixs.oracle = rearrange(pix_oracle, 't i h w 1 two -> i (h w) t two')
flows.oracle = pix_to_flow(pixs.oracle.clone())
aligned.oracle = align_from_flow(dyn_clean,
flows.oracle,
patchsize,
isize=isize)
optimal_scores.oracle = np.zeros((nimages, npix, 1, nframes))
optimal_scores.oracle = torch_to_numpy(optimal_scores.blk)
# -- compute optical flow --
vprint("[C Flow]")
vprint(dyn_noisy_ftrs.shape)
start_time = time.perf_counter()
# flows.cflow = cflow.runBurst(dyn_clean_ftrs)
# flows.cflow[...,1] = -flows.cflow[...,1]
flows.cflow = torch.LongTensor(flows.blk.clone().cpu().numpy())
# flows.cflow = flows.blk.clone()
# flows.cflow = torch.round(flows.cflow)
runtimes.cflow = time.perf_counter() - start_time
pixs.cflow = flow_to_pix(flows.cflow.clone(), nframes, isize=isize)
aligned.cflow = align_from_flow(dyn_clean,
flows.cflow,
patchsize,
isize=isize)
optimal_scores.cflow = np.zeros((nimages, npix, 1, nframes))
# optimal_scores.blk = eval_prop.score_burst_from_flow(dyn_noisy,flows.nnf_local,
# patchsize,nblocks)[1]
optimal_scores.blk = torch_to_numpy(optimal_scores.blk)
# -- compute groundtruth flow --
dsname = cfg.dataset.name
if "kitti" in dsname or 'bsd_burst' == dsname:
pix_gt = nnf_gt.type(torch.float)
if pix_gt.ndim == 3:
pix_gt_rs = rearrange(pix_gt, 'i tm1 two -> i 1 tm1 two')
pix_gt = repeat(pix_gt, 'i tm1 two -> i p tm1 two', p=npix)
if pix_gt.ndim == 5:
pix_gt = rearrange(pix_gt, 't i h w two -> i (h w) t two')
pix_gt = torch.LongTensor(pix_gt.cpu().numpy().copy())
# flows.of = torch.zeros_like(pix_gt)#pix_to_flow(pix_gt.clone())
flows.of = pix_to_flow(pix_gt.clone())
else:
flows.of = flow_gt
flows.of = rearrange(flow_gt, 'i 1 t h w two -> i (h w) t two')
# -- align groundtruth flow --
aligned.of = align_from_flow(dyn_clean, flows.of, nblocks, isize=isize)
pixs.of = flow_to_pix(flows.of.clone(), nframes, isize=isize)
runtimes.of = 0. # given
optimal_scores.of = np.zeros(
(nimages, npix, 1, nframes)) # clean target is zero
aligned.clean = static_clean
anoisy.clean = static_clean
# optimal_scores.of = eval_ave.score_burst_from_flow(dyn_noisy,
# flows.of,
# patchsize,nblocks)[0]
# -- compute nearest neighbor fields [global] --
vprint("NNF Global.")
start_time = time.perf_counter()
shape_str = 't b h w two -> b (h w) t two'
nnf_vals, nnf_pix = nnf.compute_burst_nnf(dyn_clean_ftrs, ref_t,
patchsize)
runtimes.nnf = time.perf_counter() - start_time
pixs.nnf = torch.LongTensor(rearrange(nnf_pix[..., 0, :], shape_str))
flows.nnf = pix_to_flow(pixs.nnf.clone())
vprint(dyn_clean.shape, pixs.nnf.shape, nblocks)
aligned.nnf = align_from_pix(dyn_clean, pixs.nnf, nblocks)
anoisy.nnf = align_from_pix(dyn_noisy, pixs.nnf, nblocks)
# aligned.nnf = align_from_flow(dyn_clean,flows.nnf,nblocks,isize=isize)
optimal_scores.nnf = np.zeros(
(nimages, npix, 1, nframes)) # clean target is zero
# -- compute nearest neighbor fields [local] --
vprint("NNF Local.")
start_time = time.perf_counter()
valMean = 0.
vals_local, pix_local = nnf_utils.runNnfBurst(dyn_clean_ftrs,
patchsize,
nblocks,
1,
valMean=valMean,
blockLabels=blockLabels)
runtimes.nnf_local = time.perf_counter() - start_time
torch.cuda.synchronize()
vprint("pix_local.shape ", pix_local.shape)
pixs.nnf_local = rearrange(pix_local, 't i h w 1 two -> i (h w) t two')
flows.nnf_local = pix_to_flow(pixs.nnf_local.clone())
# aligned_local = align_from_flow(clean,flow_gt,cfg.nblocks)
# aligned_local = align_from_pix(dyn_clean,pix_local,cfg.nblocks)
vprint(flows.nnf_local.min(), flows.nnf_local.max())
aligned.nnf_local = align_from_pix(dyn_clean, pixs.nnf_local, nblocks)
anoisy.nnf_local = align_from_pix(dyn_noisy, pixs.nnf_local, nblocks)
optimal_scores.nnf_local = optimal_scores.nnf
# optimal_scores.nnf_local = eval_ave.score_burst_from_flow(dyn_noisy,
# flows.nnf,
# patchsize,nblocks)[1]
optimal_scores.nnf_local = torch_to_numpy(optimal_scores.nnf_local)
# -----------------------------------
#
# -- old way to compute NNF local --
#
# -----------------------------------
# pixs.nnf = torch.LongTensor(rearrange(nnf_pix[...,0,:],shape_str))
# flows.nnf = pix_to_flow(pixs.nnf.clone())
# aligned.nnf = align_from_pix(dyn_clean,pixs.nnf,nblocks)
# aligned.nnf = align_from_flow(dyn_clean,flows.nnf,nblocks,isize=isize)
# flows.nnf_local = optim.run(dyn_clean_ftrs,patchsize,eval_ave,
# nblocks,iterations,subsizes,K)
# -----------------------------------
# -----------------------------------
# -- compute proposed search of nnf --
vprint("Global NNF Noisy")
start_time = time.perf_counter()
split_vals, split_pix = nnf.compute_burst_nnf(dyn_noisy_ftrs, ref_t,
patchsize)
runtimes.split = time.perf_counter() - start_time
# split_pix = np.copy(nnf_pix)
split_pix_best = torch.LongTensor(
rearrange(split_pix[..., 0, :], shape_str))
split_pix_best = torch.LongTensor(split_pix_best)
pixs.split = split_pix_best.clone()
flows.split = pix_to_flow(split_pix_best)
aligned.split = align_from_pix(dyn_clean, split_pix_best, nblocks)
anoisy.split = align_from_pix(dyn_noisy, split_pix_best, nblocks)
optimal_scores.split = optimal_scores.nnf_local
# optimal_scores.split = eval_ave.score_burst_from_flow(dyn_noisy,flows.nnf_local,
# patchsize,nblocks)[1]
optimal_scores.split = torch_to_numpy(optimal_scores.split)
# -- compute complex ave --
iterations, K = 0, 1
subsizes = []
vprint("[Ours] Ave loss function")
start_time = time.perf_counter()
estVar = torch.std(dyn_noisy_ftrs.reshape(-1)).item()**2
valMean = 0. #2 * estVar# * patchsize**2# / patchsize**2
vals_local, pix_local = nnf_utils.runNnfBurst(dyn_noisy_ftrs,
patchsize,
nblocks,
1,
valMean=valMean,
blockLabels=blockLabels)
runtimes.ave = time.perf_counter() - start_time
pixs.ave = rearrange(pix_local, 't i h w 1 two -> i (h w) t two')
flows.ave = pix_to_flow(pixs.ave.clone())
optimal_scores.ave = optimal_scores.split # same "ave" function
aligned.ave = align_from_flow(dyn_clean,
flows.ave,
nblocks,
isize=isize)
anoisy.ave = align_from_flow(dyn_noisy,
flows.ave,
nblocks,
isize=isize)
optimal_scores.ave = optimal_scores.split # same "ave" function
# -- compute ave with smoothing --
iterations, K = 0, 1
subsizes = []
vprint("[Ours] Ave loss function")
start_time = time.perf_counter()
pix_local = smooth_locs(pix_local, nclusters=1)
runtimes.ave_smooth = time.perf_counter() - start_time + runtimes.ave
pixs.ave_smooth = rearrange(pix_local,
't i h w 1 two -> i (h w) t two')
flows.ave_smooth = pix_to_flow(pixs.ave_smooth.clone())
optimal_scores.ave_smooth = optimal_scores.split # same "ave" function
aligned.ave_smooth = align_from_flow(dyn_clean,
flows.ave_smooth,
nblocks,
isize=isize)
anoisy.ave_smooth = align_from_flow(dyn_noisy,
flows.ave_smooth,
nblocks,
isize=isize)
optimal_scores.ave_smooth = optimal_scores.split # same "ave_smooth" function
# -- compute flow --
vprint("L2-Local Recursive")
start_time = time.perf_counter()
vals_local, pix_local, wburst = nnf_utils.runNnfBurstRecursive(
dyn_noisy_ftrs,
dyn_clean,
patchsize,
nblocks,
isize,
1,
valMean=valMean,
blockLabels=blockLabels)
runtimes.l2r = time.perf_counter() - start_time
pixs.l2r = rearrange(pix_local, 't i h w 1 two -> i (h w) t two')
flows.l2r = pix_to_flow(pixs.l2r.clone())
aligned.l2r = wburst #align_from_flow(dyn_clean,flows.l2r,nblocks,isize=isize)
optimal_scores.l2r = optimal_scores.split # same "ave" function
# -- compute nvof flow --
vprint("NVOF")
start_time = time.perf_counter()
# flows.nvof = nvof.nvof_burst(dyn_noisy_ftrs)
flows.nvof = flows.ave.clone()
runtimes.nvof = time.perf_counter() - start_time
pixs.nvof = flow_to_pix(flows.nvof.clone(), nframes, isize=isize)
aligned.nvof = align_from_flow(dyn_clean,
flows.nvof,
nblocks,
isize=isize)
anoisy.nvof = align_from_flow(dyn_noisy,
flows.nvof,
nblocks,
isize=isize)
optimal_scores.nvof = optimal_scores.split # same "ave" function
# -- compute flownet --
vprint("FlowNetv2")
start_time = time.perf_counter()
_, flows.flownet = flownet_align(dyn_noisy_ftrs)
# flows.flownet = flows.ave.clone().cpu()
runtimes.flownet = time.perf_counter() - start_time
pixs.flownet = flow_to_pix(flows.flownet.clone(), nframes, isize=isize)
aligned.flownet = align_from_flow(dyn_clean,
flows.flownet,
nblocks,
isize=isize)
anoisy.flownet = align_from_flow(dyn_noisy,
flows.flownet,
nblocks,
isize=isize)
optimal_scores.flownet = optimal_scores.split
# -- compute simple ave --
iterations, K = 0, 1
subsizes = []
vprint("[simple] Ave loss function")
start_time = time.perf_counter()
optim = AlignOptimizer("v3")
if cfg.patchsize < 11 and cfg.frame_size[0] <= 64 and False:
flows.ave_simp = optim.run(dyn_noisy, patchsize, eval_ave_simp,
nblocks, iterations, subsizes, K)
else:
flows.ave_simp = flows.ave.clone().cpu()
runtimes.ave_simp = time.perf_counter() - start_time
pixs.ave_simp = flow_to_pix(flows.ave_simp.clone(),
nframes,
isize=isize)
aligned.ave_simp = align_from_flow(dyn_clean,
flows.ave_simp,
nblocks,
isize=isize)
anoisy.ave_simp = align_from_flow(dyn_noisy,
flows.ave_simp,
nblocks,
isize=isize)
optimal_scores.ave_simp = optimal_scores.split # same "ave" function
# -- format results --
#pad = 2*(nframes-1)*ppf+4
# pad = 2*(cfg.nblocks//2)#2*(nframes-1)*ppf+4
# isize = edict({'h':H-pad,'w':W-pad})
# -- flows to numpy --
frame_size = cfg.frame_size[0]
is_even = frame_size % 2 == 0
mid_pix = frame_size * frame_size // 2 + (frame_size // 2) * is_even
mid_pix = 32 * 10 + 23
flows_np = edict_torch_to_numpy(flows)
pixs_np = edict_torch_to_numpy(pixs)
# -- End-Point-Errors --
epes_of = compute_flows_epe_wrt_ref(flows, "of")
epes_nnf = compute_flows_epe_wrt_ref(flows, "nnf")
epes_nnf_local = compute_flows_epe_wrt_ref(flows, "nnf_local")
nnf_acc = compute_acc_wrt_ref(flows, "nnf")
nnf_local_acc = compute_acc_wrt_ref(flows, "nnf_local")
# -- PSNRs --
aligned = remove_center_frames(aligned)
psnrs = compute_frames_psnr(aligned, psize)
# -- denoised PSNRS --
def burst_mean(in_burst):
return torch.mean(in_burst, dim=0)[None, :]
anoisy = remove_center_frames(anoisy)
anoisy = apply_across_dict(anoisy, burst_mean)
dn_psnrs = compute_frames_psnr(anoisy, psize)
vprint(dn_psnrs)
# -- print report ---
print("\n" * 3) # banner
print("-" * 25 + " Results " + "-" * 25)
# print_dict_ndarray_0_midpix(flows_np,mid_pix)
# print_dict_ndarray_0_midpix(pixs_np,mid_pix)
# print_verbose_psnrs(psnrs)
# print_delta_summary_psnrs(psnrs)
# print_verbose_epes(epes_of,epes_nnf)
# print_nnf_acc(nnf_acc)
# print_nnf_local_acc(nnf_local_acc)
# print_summary_epes(epes_of,epes_nnf)
# print_summary_denoised_psnrs(dn_psnrs)
print_summary_psnrs(psnrs)
print_runtimes(runtimes)
# -- prepare results to be appended --
psnrs = edict_torch_to_numpy(psnrs)
epes_of = edict_torch_to_numpy(epes_of)
epes_nnf = edict_torch_to_numpy(epes_nnf)
epes_nnf_local = edict_torch_to_numpy(epes_nnf_local)
nnf_acc = edict_torch_to_numpy(nnf_acc)
nnf_local_acc = edict_torch_to_numpy(nnf_local_acc)
image_index = torch_to_numpy(image_index)
batch_results = {
'runtimes': runtimes,
'optimal_scores': optimal_scores,
'psnrs': psnrs,
'epes_of': epes_of,
'epes_nnf': epes_nnf,
'epes_nnf_local': epes_nnf_local,
'nnf_acc': nnf_acc,
'nnf_local_acc': nnf_local_acc
}
# -- format results --
batch_results = flatten_internal_dict(batch_results)
format_fields(batch_results, image_index, rng_state)
# print("shape check.")
# for key,value in batch_results.items():
# print(key,value.shape)
record.append(batch_results)
# print("\n"*3)
# print("-"*20)
# print(record.record)
# print("-"*20)
# print("\n"*3)
# record.stack_record()
record.cat_record()
# print("\n"*3)
# print("-"*20)
# print(record.record)
# print("-"*20)
print("\n" * 3)
print("-" * 20)
# df = pd.DataFrame().append(record.record,ignore_index=True)
for key, val in record.record.items():
vprint(key, val.shape)
# vprint(df)
vprint("-" * 20)
vprint("\n" * 3)
return record.record
def run_with_seed(seed):
# -- settings --
cfg = get_cfg_defaults()
cfg.use_anscombe = False
cfg.noise_params.ntype = 'g'
cfg.noise_params.g.std = 10.
cfg.nframes = 5
cfg.patchsize = 11
# -- seeds --
cfg.seed = seed
# cfg.seed = 123 # sky of a forest
# cfg.seed = 345 # handrail and stairs
# cfg.seed = 567 # cloudy blue sky
# cfg.seed = 567 # cloudy blue sky
# -- set seed --
set_seed(cfg.seed)
# -- load dataset --
data, loaders = load_image_dataset(cfg)
train_iter = iter(loaders.tr)
# -- fetch sample --
sample = next(train_iter)
sample_to_cuda(sample)
# -- unpack data --
noisy, clean = sample['noisy'], sample['burst']
nframes, nimages, ncolors, H, W = noisy.shape
isize = edict({'h': H, 'w': W})
# -- boxes for plotting --
boxes = edict()
aligned = edict()
# -- compute clean nnf --
vprint("[start] clean nnf.")
align_fxn = get_align_method(cfg, "l2_global")
aligned.gt, flow = align_fxn(clean, None, None)
boxes.gt = boxes_from_flow(flow, H, W)
vprint("[done] clean nnf.")
# -- compute nnf --
vprint("[start] global nnf.")
align_fxn = get_align_method(cfg, "l2_global")
_, flow = align_fxn(noisy, None, None)
aligned.global_l2 = align_from_flow(clean, flow, cfg.nblocks, isize=isize)
boxes.global_l2 = boxes_from_flow(flow, H, W)
vprint("[done] global nnf.")
# -- compute local nnf --
vprint("[start] local nnf.")
align_fxn = get_align_method(cfg, "l2_local")
_, flow = align_fxn(noisy, None, None)
aligned.local_l2 = align_from_flow(clean, flow, cfg.nblocks, isize=isize)
boxes.local_l2 = boxes_from_flow(flow, H, W)
vprint("[done] local nnf.")
# -- compute proposed score --
vprint("[start] bootstrapping.")
align_fxn = get_align_method(cfg, "bs_local_v2")
_, flow = align_fxn(noisy, None, None)
aligned.local_bs = align_from_flow(clean, flow, cfg.nblocks, isize=isize)
boxes.local_bs = boxes_from_flow(flow, H, W)
vprint("[done] bootstrapping.")
# -- reshape to image --
noisy = rearrange(noisy, 't b c h w -> t b h w c')
clean = rearrange(clean, 't b c h w -> t b h w c')
# -- normalize to [0,1] --
noisy -= noisy.min()
clean -= clean.min()
noisy /= noisy.max()
clean /= clean.max()
# -- clamp to [0,1] --
# noisy = noisy.clamp(0,1)
# clean = clean.clamp(0,1)
# print_tensor_stats("noisy",noisy)
# print_tensor_stats("clean",clean)
# -- cuda to cpu --
noisy = noisy.cpu()
clean = clean.cpu()
for field in boxes.keys():
boxes[field] = boxes[field].cpu().numpy()
# -- plot boxes for middle pix --
ref_pix = edict({'x': H // 2, 'y': W // 2})
field = 'global_l2'
plot_boxes(noisy, clean, aligned, field, boxes, ref_pix, cfg.patchsize,
seed)
def test_davis_dataset():
# -- run exp --
cfg = get_cfg_defaults()
cfg.random_seed = 123
nbatches = 20
# -- set random seed --
set_seed(cfg.random_seed)
# -- load dataset --
cfg.nframes = 5
# cfg.frame_size = None
# cfg.frame_size = [256,256]
cfg.frame_size = [96, 96]
cfg.dataset.name = "davis"
data, loaders = load_dataset(cfg, "dynamic")
image_iter = iter(loaders.tr)
sample = next(image_iter)
print(len(loaders.tr))
fn = "./davis_example.png"
save_image(sample['dyn_noisy'], fn, normalize=True, vrange=(0., 1.))
# -- save path for viz --
save_dir = SAVE_PATH
if not save_dir.exists(): save_dir.mkdir(parents=True)
# -- sample data --
for image_index in range(nbatches):
# -- sample image --
index = -1
# while index != 3233:
# sample = next(image_iter)
# convert_keys(sample)
# index = sample['image_index'][0][0].item()
sample = next(image_iter)
sample = next(image_iter)
# batch_dim0(sample)
# convert_keys(sample)
# -- extract info --
noisy = sample['dyn_noisy']
clean = sample['dyn_clean']
snoisy = sample['static_noisy']
sclean = sample['static_clean']
flow_est = sample['ref_flow']
pix_gt = sample['ref_pix']
index = sample['index'][0][0].item()
nframes, nimages, c, h, w = noisy.shape
mid_pix = h * w // 2 + 2 * cfg.nblocks
shape_str = 'b k t h w two -> k b (h w) t two'
pix_gt = rearrange(pix_gt, shape_str)[0]
flow_est = rearrange(flow_est, shape_str)[0]
# -- print shapes --
print("-" * 50)
for key, val in sample.items():
if isinstance(val, list): continue
print("{}: {}".format(key, val.shape))
print("-" * 50)
print(f"Image Index {index}")
# -- io info --
image_dir = save_dir / f"index{index}/"
if not image_dir.exists(): image_dir.mkdir()
#
# -- Compute NNF to Ensure things are OKAY --
#
isize = edict({'h': h, 'w': w})
# pad = cfg.patchsize//2 if cfg.patchsize > 1 else 1
pad = cfg.nblocks // 2 + 1
psize = edict({'h': h - 2 * pad, 'w': w - 2 * pad})
# flow_gt = rearrange(flow,'i 1 fm1 h w two -> i (h w) fm1 two')
# pix_gt = flow_gt.clone()
# pix_gt = flow_to_pix(flow_gt.clone(),nframes,isize=isize)
def cc(image):
return tvF.center_crop(image, (psize.h, psize.w))
# -- align from [pix or flow] --
pix_gt = pix_gt.type(torch.long)
aligned_gt = align_from_pix(clean, pix_gt.clone(), 0) #cfg.nblocks)
psnr = compute_aligned_psnr(sclean, aligned_gt, psize).T[0]
print(f"[Dataset Pix Alignment] PSNR: {psnr}")
flow_gt = pix_to_flow(pix_gt)
aligned_gt = align_from_flow(clean, flow_gt.clone(), 0, isize=isize)
psnr = compute_aligned_psnr(sclean, aligned_gt, psize).T[0]
print(f"[Dataset Flow Alignment] PSNR: {psnr}")
aligned_gt = align_from_flow(clean, flow_est.clone(), 0, isize=isize)
psnr = compute_aligned_psnr(sclean, aligned_gt, psize).T[0]
print(f"[(est) Dataset Flow Alignment] PSNR: {psnr}")
# aligned_gt = warp_burst_flow(clean, flow_global)
# isize = edict({'h':h,'w':w})
# flow_est = pix_to_flow_est(pix_gt)
print(torch.stack([flow_gt, flow_est], -1))
assert torch.sum((flow_est - flow_gt)**2).item() < 1e-8
# -- compute the nnf again [for checking] --
nnf_vals, nnf_pix = compute_burst_nnf(clean, nframes // 2,
cfg.patchsize)
shape_str = 't b h w two -> b (h w) t two'
pix_global = torch.LongTensor(rearrange(nnf_pix[..., 0, :], shape_str))
flow_global = pix_to_flow(pix_global.clone())
# print(flow_gt.shape,flow_global.shape)
# -- explore --
# print("NFrames: ",nframes)
print(pix_global.shape, pix_gt.shape)
for t in range(nframes):
delta = pix_global[:, :, t] != pix_gt[:, :, t]
delta = delta.type(torch.float)
delta = 100 * torch.mean(delta)
print("[%d]: %2.3f" % (t, delta))
# print(torch.sum(torch.abs(pix_global - pix_gt)))
print(pix_global[:, :, 41])
print(pix_gt[:, :, 41])
# print(torch.stack([pix_global,pix_gt],-1))
# print(torch.where(pix_global!=pix_gt))
# print(torch.sum((pix_global-pix_gt)**2))
# print(torch.sum((pix_global!=pix_gt).type(torch.float)))
agg = torch.stack([pix_global.ravel(), pix_gt.ravel()], -1)
# print(agg)
# print(agg.shape)
# print(torch.sum((agg[:,0]!=agg[:,1]).type(torch.float)))
# print(torch.where(agg[:,0]!=agg[:,1]))
# -- create report --
print(pix_global.shape, pix_gt.shape)
print(pix_global.min(), pix_gt.min())
print(pix_global.max(), pix_gt.max())
print(type(pix_global), type(pix_gt))
aligned_gt = align_from_pix(clean, pix_global, cfg.nblocks)
psnr = compute_aligned_psnr(sclean, aligned_gt, psize).T[0]
print(f"[GT Alignment] PSNR: {psnr}")
#
# -- Save Images to Qualitative Inspect --
#
fn = image_dir / "noisy.png"
save_image(cc(noisy), fn, normalize=True, vrange=None)
fn = image_dir / "clean.png"
save_image(cc(clean), fn, normalize=True, vrange=None)
print(cc(sclean).shape)
fn = image_dir / "diff.png"
save_image(cc(sclean) - cc(aligned_gt),
fn,
normalize=True,
vrange=None)
fn = image_dir / "aligned_gt.png"
save_image(cc(aligned_gt), fn, normalize=True, vrange=None)
print(image_dir)
return