|
|
|
|
|
|
|
|
|
from pdb import set_trace as bb |
|
import os, os.path as osp |
|
from tqdm import tqdm |
|
import numpy as np |
|
|
|
SEQUENCES = [ 'lakeside', 'sand_box', 'storage_room', 'storage_room_2', 'tunnel', |
|
'delivery_area', 'electro', 'forest', 'playground', 'terrains'] |
|
|
|
RATES = [3, 5, 7, 9, 11, 13, 15] |
|
|
|
def parse_args(): |
|
import argparse |
|
parser = argparse.ArgumentParser('PUMP evaluation script for the ETH3D dataset') |
|
|
|
parser.add_argument('--root', default='datasets/eth3d') |
|
parser.add_argument('--output', default='results/eth3d') |
|
|
|
parser.add_argument('--just-print', action='store_true', help='just print commands') |
|
return parser.parse_args() |
|
|
|
|
|
def main( args ): |
|
run_pump(args) and run_eval(args) |
|
|
|
|
|
def run_pump(args): |
|
done = True |
|
for img1, img2 in tqdm(list_eth3d_pairs()): |
|
output_path = osp.join(args.output, img1, img2+'.corres') |
|
if osp.isfile(output_path): continue |
|
|
|
done = False |
|
_exec(f'''python test_multiscale_recursive.py |
|
--img1 {osp.join(args.root,img1)} |
|
--img2 {osp.join(args.root,img2)} |
|
--max-scale 1.5 |
|
--desc PUMP |
|
--post-filter "densify=True,dense_side='right'" |
|
--output {output_path}''') |
|
|
|
return done |
|
|
|
|
|
def run_eval( args ): |
|
for rate in RATES: |
|
mean_aepe_per_rate = 0 |
|
|
|
for seq in SEQUENCES: |
|
pairs = np.load(osp.join(args.root, 'info_ETH3D_files', f'{seq}_every_5_rate_of_{rate}'), allow_pickle=True) |
|
|
|
mean_aepe_per_seq = 0 |
|
for pair in pairs: |
|
img1, img2 = pair['source_image'], pair['target_image'] |
|
Ys, Xs, Yt, Xt = [np.float32(pair[k]) for k in 'Ys Xs Yt Xt'.split()] |
|
|
|
corres_path = osp.join(args.output, img1, img2+'.corres') |
|
corres = np.load(corres_path, allow_pickle=True)['corres'] |
|
|
|
|
|
W, H = np.int32(corres[-1, 2:4] + 1) |
|
flow = (corres[:,0:2] - corres[:,2:4]).reshape(H, W, 2) |
|
iYt, iXt = np.int32(np.round(Yt)), np.int32(np.round(Xt)) |
|
if 'correct way': |
|
gt_targets = np.c_[Xs - Xt, Ys - Yt] |
|
est_targets = flow[iYt, iXt] |
|
elif 'GLU-Net way (somewhat inaccurate because of overlapping points in the mask)': |
|
mask = np.zeros((H,W), dtype=bool) |
|
mask[iYt, iXt] = True |
|
gt_flow = np.full((H,W,2), np.nan, dtype=np.float32) |
|
gt_flow[iYt, iXt, 0] = Xs - Xt |
|
gt_flow[iYt, iXt, 1] = Ys - Yt |
|
gt_targets = gt_flow[mask] |
|
est_targets = flow[mask] |
|
|
|
|
|
aepe = np.linalg.norm(est_targets - gt_targets, axis=-1).mean() |
|
mean_aepe_per_seq += aepe |
|
|
|
mean_aepe_per_seq /= len(pairs) |
|
mean_aepe_per_rate += mean_aepe_per_seq |
|
print(f'mean AEPE for {rate=} {seq=}:', mean_aepe_per_seq) |
|
|
|
print(f'>> mean AEPE for {rate=}:', mean_aepe_per_rate / len(SEQUENCES)) |
|
|
|
|
|
def list_eth3d_pairs(): |
|
path = osp.join(args.root, 'info_ETH3D_files', 'list_pairs.txt') |
|
try: |
|
lines = open(path).read().splitlines() |
|
except OSError: |
|
lines = [] |
|
for seq in SEQUENCES: |
|
for rate in RATES: |
|
pairs = np.load(osp.join(args.root, 'info_ETH3D_files', f'{seq}_every_5_rate_of_{rate}'), allow_pickle=True) |
|
for pair in pairs: |
|
lines.append(pair['source_image'] + ' ' + pair['target_image']) |
|
open(path, 'w').write('\n'.join(lines)) |
|
|
|
pairs = [line.split() for line in lines if line[0] != '#'] |
|
return pairs |
|
|
|
|
|
def _exec(cmd): |
|
|
|
cmd = ' '.join(cmd.split()) |
|
if args.just_print: |
|
print(cmd) |
|
else: |
|
os.system(cmd) |
|
|
|
|
|
if __name__ == '__main__': |
|
args = parse_args() |
|
main( args ) |
|
|
