|
|
|
|
|
|
|
|
|
from pdb import set_trace as bb |
|
import os, os.path as osp |
|
|
|
from PIL import Image |
|
import numpy as np |
|
from tools.viz import pl, noticks |
|
|
|
""" This script will warp (deform) img2 so that it fits img1 |
|
|
|
>> In case of memory failure (not enough GPU memory): |
|
try adding '--resize 400 300' (or larger values if possible) to the _exec(...) command below. |
|
""" |
|
|
|
def parse_args(): |
|
import argparse |
|
parser = argparse.ArgumentParser('PUMP demo script for the image warping demo') |
|
|
|
parser.add_argument('--img1', default='datasets/demo_warp/mountains_src.jpg') |
|
parser.add_argument('--img2', default='datasets/demo_warp/mountains_tgt.jpg') |
|
parser.add_argument('--output', default='results/demo_warp') |
|
|
|
parser.add_argument('--just-print', action='store_true', help='just print commands') |
|
return parser.parse_args() |
|
|
|
|
|
def main( args ): |
|
run_pump(args) and run_demo_warp(args) |
|
|
|
|
|
def run_pump(args): |
|
output_path = osp.join(args.output, args.img1, args.img2+'.corres') |
|
if osp.isfile(output_path): return True |
|
|
|
return _exec(f'''python test_singlescale_recursive.py |
|
--img1 {args.img1} |
|
--img2 {args.img2} |
|
--post-filter densify=True |
|
--output {output_path}''') |
|
|
|
|
|
def run_demo_warp(args): |
|
corres_path = osp.join(args.output, args.img1, args.img2+'.corres') |
|
corres = np.load(corres_path)['corres'] |
|
|
|
img1 = Image.open(args.img1).convert('RGB') |
|
img2 = Image.open(args.img2).convert('RGB') |
|
|
|
W, H = img1.size |
|
warped_img2 = warp_img(np.asarray(img2), corres[:,2:4].reshape(H,W,2)) |
|
|
|
pl.figure('Warping demo') |
|
|
|
noticks(pl.subplot(211)) |
|
pl.imshow( img2 ) |
|
pl.title('Source image') |
|
|
|
noticks(pl.subplot(223)) |
|
pl.imshow( img1 ) |
|
pl.title('Target image') |
|
|
|
noticks(pl.subplot(224)) |
|
pl.imshow( warped_img2 ) |
|
pl.title('Source image warped to match target') |
|
|
|
pl.tight_layout() |
|
pl.show(block=True) |
|
|
|
|
|
def warp_img( img, absolute_flow ): |
|
H1, W1, TWO = absolute_flow.shape |
|
H2, W2, THREE = img.shape |
|
assert TWO == 2 and THREE == 3 |
|
|
|
warp = absolute_flow.round().astype(int) |
|
invalid = (warp[:,:,0]<0) | (warp[:,:,0]>=W2) | (warp[:,:,1]<0) | (warp[:,:,1]>=H2) |
|
|
|
warp[:,:,0] = warp[:,:,0].clip(min=0, max=W2-1) |
|
warp[:,:,1] = warp[:,:,1].clip(min=0, max=H2-1) |
|
warp = warp[:,:,0] + W2*warp[:,:,1] |
|
|
|
warped_img = np.asarray(img).reshape(-1,3)[warp].reshape(H1,W1,3) |
|
return warped_img |
|
|
|
|
|
def _exec(cmd): |
|
|
|
cmd = ' '.join(cmd.split()) |
|
|
|
if args.just_print: |
|
print(cmd) |
|
return False |
|
else: |
|
return os.WEXITSTATUS(os.system(cmd)) == 0 |
|
|
|
|
|
if __name__ == '__main__': |
|
args = parse_args() |
|
main( args ) |
|
|
