|
|
|
import sys |
|
import os |
|
from tqdm import tqdm |
|
dir_path = os.path.dirname(os.path.realpath(__file__)) |
|
sys.path.append(os.path.join(dir_path, 'core')) |
|
from pathlib import Path |
|
from data_test_flow import * |
|
from models.network_test_flow import NeuralNRT |
|
from options_test_flow import TestOptions |
|
import torch |
|
import numpy as np |
|
|
|
|
|
|
|
def save_flow_numpy(filename, flow_input): |
|
np.save(filename, flow_input) |
|
|
|
|
|
def predict(data): |
|
with torch.no_grad(): |
|
model.eval() |
|
path_flow = data["path_flow"] |
|
src_crop_im = data["src_crop_color"].cuda() |
|
tar_crop_im = data["tar_crop_color"].cuda() |
|
src_im = data["src_color"].cuda() |
|
tar_im = data["tar_color"].cuda() |
|
src_mask = data["src_mask"].cuda() |
|
crop_param = data["Crop_param"].cuda() |
|
B = src_mask.shape[0] |
|
flow = model(src_crop_im, tar_crop_im, src_im, tar_im, crop_param) |
|
for i in range(B): |
|
flow_tmp = flow[i].cpu().numpy() * src_mask[i].cpu().numpy() |
|
save_flow_numpy(os.path.join(save_path, os.path.basename( |
|
path_flow[i])[:-6]+".npy"), flow_tmp) |
|
|
|
|
|
if __name__ == "__main__": |
|
width = 272 |
|
height = 480 |
|
|
|
test_opts = TestOptions().parse() |
|
test_opts.pretrain_model_path = os.path.join( |
|
dir_path, 'pretrain_model/raft-small.pth') |
|
data_loader = CreateDataLoader(test_opts) |
|
testloader = data_loader.load_data() |
|
model_path = os.path.join(dir_path, 'sgd_NNRT_model_epoch19008_50000.pth') |
|
model = NeuralNRT(test_opts, os.path.join( |
|
dir_path, 'pretrain_model/raft-small.pth')) |
|
state_dict = torch.load(model_path) |
|
|
|
model.CorresPred.load_state_dict(state_dict["net_C"]) |
|
model.ImportanceW.load_state_dict(state_dict["net_W"]) |
|
|
|
model = model.cuda() |
|
|
|
save_path = test_opts.savepath |
|
Path(save_path).mkdir(parents=True, exist_ok=True) |
|
total_length = len(testloader) |
|
|
|
for batch_idx, data in tqdm(enumerate(testloader), total=total_length): |
|
predict(data) |
|
|
