Пример #1
0
def compute_frames_psnr(frames,isize):
    skip_fields = ["clean"]
    psnrs = edict()
    for field in frames.keys():
        if field in skip_fields: continue
        psnrs[field] = compute_aligned_psnr(frames[field],frames.clean,isize)
    return psnrs
Пример #2
0
def compute_frames_psnr(frames, isize):
    psnrs = edict()
    psnrs.of = compute_aligned_psnr(frames.of, frames.clean, isize)
    psnrs.nnf = compute_aligned_psnr(frames.nnf, frames.clean, isize)
    psnrs.split = compute_aligned_psnr(frames.split, frames.clean, isize)
    psnrs.ave_simp = compute_aligned_psnr(frames.ave_simp, frames.clean, isize)
    psnrs.ave = compute_aligned_psnr(frames.ave, frames.clean, isize)
    psnrs.est = compute_aligned_psnr(frames.est, frames.clean, isize)
    return psnrs
Пример #3
0
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)
Пример #4
0
def test_global_dynamics():

    # -- run exp --
    cfg = get_cfg_defaults()
    cfg.random_seed = 123
    nbatches = 20

    # -- set random seed --
    set_seed(cfg.random_seed)

    # -- load dataset --
    print("load image dataset.")
    data, loaders = load_image_dataset(cfg)
    image_iter = iter(loaders.tr)

    # -- save path for viz --
    save_dir = Path(
        f"{settings.ROOT_PATH}/output/tests/datasets/test_global_dynamics/")
    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)
        # 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 = sample['flow_gt']
        index = sample['image_index'][0][0].item()
        nframes, nimages, c, h, w = noisy.shape
        mid_pix = h * w // 2 + 2 * cfg.nblocks
        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 = repeat(flow, 'i fm1 two -> i s fm1 two', s=h * w)
        pix_gt = flow_to_pix(flow_gt.clone(), nframes, isize=isize)

        def cc(image):
            return tvF.center_crop(image, (psize.h, psize.w))

        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())
        # aligned_gt = warp_burst_flow(clean, flow_global)
        aligned_gt = align_from_pix(clean, pix_gt, cfg.nblocks)
        # isize = edict({'h':h,'w':w})
        # aligned_gt = align_from_flow(clean,flow_global,cfg.nblocks,isize=isize)
        # psnr = compute_aligned_psnr(sclean[[nframes//2]],clean[[nframes//2]],psize)
        psnr = compute_aligned_psnr(sclean, aligned_gt, psize)
        print(f"[GT Alignment] PSNR: {psnr}")

        # -- compute with nvidia's opencv optical flow --
        nd_clean = rearrange(clean.numpy(), 't 1 c h w -> t h w c')
        ref_t = nframes // 2
        frames, flows = [], []
        for t in range(nframes):
            if t == ref_t:
                frames.append(nd_clean[t][None, :])
                flows.append(torch.zeros(flows[-1].shape))
                continue
            from_frame = 255. * cv2.cvtColor(nd_clean[ref_t],
                                             cv2.COLOR_RGB2GRAY)
            to_frame = 255. * cv2.cvtColor(nd_clean[t], cv2.COLOR_RGB2GRAY)
            _flow = cv2.calcOpticalFlowFarneback(to_frame, from_frame, None,
                                                 0.5, 3, 3, 10, 5, 1.2, 0)
            _flow = np.round(_flow).astype(np.float32)  # not good for later
            w_frame = warp_flow(nd_clean[t], -_flow)
            _flow[..., 0] = -_flow[..., 0]  # my OF is probably weird.
            # print("w_frame.shape ",w_frame.shape)
            flows.append(torch.FloatTensor(_flow))
            frames.append(torch.FloatTensor(w_frame[None, :]))
        flows = torch.stack(flows)
        flows = rearrange(flows, 't h w two -> 1 (h w) t two')
        frames = torch.FloatTensor(np.stack(frames))
        frames = rearrange(frames, 't i h w c -> t i c h w')
        # print("flows.shape ",flows.shape)
        # print("frames.shape ",frames.shape)
        # print("sclean.shape ",sclean.shape)
        psnr = compute_aligned_psnr(sclean, frames, psize)
        print(f"[NVOF Alignment] PSNR: {psnr}")

        pix_nvof = flow_to_pix(flows.clone(), nframes, isize=isize)
        aligned_nvof = align_from_pix(clean, pix_nvof, cfg.nblocks)
        psnr = compute_aligned_psnr(sclean, aligned_nvof, psize)
        print(f"[NVOF Alignment v2] PSNR: {psnr}")

        psnr = compute_aligned_psnr(frames, aligned_nvof, psize)
        print(f"[NVOF Alignment Methods] PSNR: {psnr}")

        print(pix_global[0, mid_pix])
        print(pix_gt[0, mid_pix])
        print(flow_global[0, mid_pix])
        print(flow_gt[0, mid_pix])
        print(flows[0, mid_pix])

        #
        # -- 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)

        # return

        # -- NNF Global --
        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())
        # aligned_global = align_from_flow(clean,flow_gt,cfg.nblocks)
        aligned_global = align_from_pix(clean, pix_gt, cfg.nblocks)
        psnr = compute_aligned_psnr(sclean, aligned_global, psize)
        print(f"[NNF Global] PSNR: {psnr}")

        # -- NNF Local (old) --
        iterations, K, subsizes = 0, 1, []
        optim = AlignOptimizer("v3")
        score_fxn_ave = get_score_function("ave")
        eval_ave = EvalBlockScores(score_fxn_ave, "ave", cfg.patchsize, 256,
                                   None)
        flow_local = optim.run(clean, cfg.patchsize, eval_ave, cfg.nblocks,
                               iterations, subsizes, K)
        pix_local = flow_to_pix(flow_local.clone(), nframes, isize=isize)
        # aligned_local = align_from_flow(clean,flow_gt,cfg.nblocks)
        aligned_local = align_from_pix(clean, pix_local, cfg.nblocks)
        psnr = compute_aligned_psnr(sclean, aligned_local, psize)
        print(f"[NNF Local (Old)] PSNR: {psnr}")

        # -- NNF Local (new) --
        _, pix_local = nnf_utils.runNnfBurst(clean,
                                             cfg.patchsize,
                                             cfg.nblocks,
                                             1,
                                             valMean=0.,
                                             blockLabels=None)
        pix_local = rearrange(pix_local, 't i h w 1 two -> i (h w) t two')
        flow_local = pix_to_flow(pix_local.clone())
        # aligned_local = align_from_flow(clean,flow_gt,cfg.nblocks)
        aligned_local = align_from_pix(clean, pix_local, cfg.nblocks)
        psnr = compute_aligned_psnr(sclean, aligned_local, psize)
        print(f"[NNF Local (New)] PSNR: {psnr}")

        # -- remove boundary from pix --
        pixes = {'gt': pix_gt, 'global': pix_global, 'local': pix_local}
        for field, pix in pixes.items():
            pix_img = rearrange(pix, 'i (h w) t two -> (i t) two h w', h=h)
            pix_cc = cc(pix_img)
            pixes[field] = pix_cc

        # -- pairwise diffs --
        field2 = "gt"
        for field1 in pixes.keys():
            if field1 == field2: continue
            delta = pixes[field1] - pixes[field2]
            delta = delta.type(torch.float)
            delta_fn = image_dir / f"delta_{field1}_{field2}.png"
            save_image(delta, delta_fn, normalize=True, vrange=None)
        print(pix_gt[0, mid_pix])
        print(pix_global[0, mid_pix])
        print(pix_local[0, mid_pix])

        print(flow_gt[0, mid_pix])
        print(flow_global[0, mid_pix])
        print(flow_local[0, mid_pix])

        #
        # -- 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)

        fn = image_dir / "aligned_global.png"
        save_image(cc(aligned_global), fn, normalize=True, vrange=None)

        fn = image_dir / "aligned_local.png"
        save_image(cc(aligned_local), fn, normalize=True, vrange=None)
Пример #5
0
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())
Пример #6
0
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
Пример #7
0
def test_set8_dataset():

    # -- PID --
    pid = os.getpid()
    print(f"PID: {pid}")

    # -- run exp --
    cfg = get_cfg_defaults()
    cfg.random_seed = 123
    nbatches = 3

    # -- set random seed --
    set_seed(cfg.random_seed)

    # -- load dataset --
    cfg.nframes = 0
    cfg.frame_size = [256, 256]
    # cfg.frame_size = [128,128]
    # cfg.frame_size = [32,32]
    cfg.dataset.name = "set8"
    data, loaders = load_dataset(cfg, "dynamic")
    image_iter = iter(loaders.tr)
    sample = next(image_iter)
    print(len(loaders.tr))
    fn = "./set8_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)
        # 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 = sample['ref_flow']
        index = sample['index'][0][0].item()
        nframes, nimages, c, h, w = noisy.shape
        mid_pix = h * w // 2 + 2 * cfg.nblocks

        # -- 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_to_pix(flow_gt.clone(), nframes, isize=isize)

        def cc(image):
            return tvF.center_crop(image, (psize.h, psize.w))

        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())
        # aligned_gt = warp_burst_flow(clean, flow_global)
        aligned_gt = align_from_pix(clean, pix_gt, cfg.nblocks)
        # isize = edict({'h':h,'w':w})
        # aligned_gt = align_from_flow(clean,flow_global,cfg.nblocks,isize=isize)
        # psnr = compute_aligned_psnr(sclean[[nframes//2]],clean[[nframes//2]],psize)
        psnr = compute_aligned_psnr(sclean, aligned_gt, psize)
        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)