|
import sys |
|
sys.path.append('core') |
|
|
|
import argparse |
|
import os |
|
import cv2 |
|
import glob |
|
import numpy as np |
|
import torch |
|
from PIL import Image |
|
|
|
from raft import RAFT |
|
from utils import flow_viz |
|
from utils.utils import InputPadder |
|
|
|
|
|
|
|
DEVICE = 'cuda' |
|
|
|
def load_image(imfile): |
|
img = np.array(Image.open(imfile)).astype(np.uint8) |
|
img = torch.from_numpy(img).permute(2, 0, 1).float() |
|
return img[None].to(DEVICE) |
|
|
|
|
|
def viz(img, flo): |
|
img = img[0].permute(1,2,0).cpu().numpy() |
|
flo = flo[0].permute(1,2,0).cpu().numpy() |
|
|
|
|
|
flo = flow_viz.flow_to_image(flo) |
|
img_flo = np.concatenate([img, flo], axis=0) |
|
|
|
|
|
|
|
|
|
|
|
cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0) |
|
cv2.waitKey() |
|
|
|
|
|
def demo(args): |
|
model = torch.nn.DataParallel(RAFT(args)) |
|
model.load_state_dict(torch.load(args.model)) |
|
|
|
model = model.module |
|
model.to(DEVICE) |
|
model.eval() |
|
|
|
with torch.no_grad(): |
|
images = glob.glob(os.path.join(args.path, '*.png')) + \ |
|
glob.glob(os.path.join(args.path, '*.jpg')) |
|
|
|
images = sorted(images) |
|
for imfile1, imfile2 in zip(images[:-1], images[1:]): |
|
image1 = load_image(imfile1) |
|
image2 = load_image(imfile2) |
|
|
|
padder = InputPadder(image1.shape) |
|
image1, image2 = padder.pad(image1, image2) |
|
|
|
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) |
|
viz(image1, flow_up) |
|
|
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--model', help="restore checkpoint") |
|
parser.add_argument('--path', help="dataset for evaluation") |
|
parser.add_argument('--small', action='store_true', help='use small model') |
|
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') |
|
parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation') |
|
args = parser.parse_args() |
|
|
|
demo(args) |
|
|
