Spaces:
Runtime error
Runtime error
File size: 2,649 Bytes
04fbff5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import os
import sys
import torch
import argparse
import numpy as np
import os.path as osp
import torch.nn.functional as F
sys.path.append('.')
from utils.utils import read, write
from flow_generation.liteflownet.run import estimate
parser = argparse.ArgumentParser(
prog = 'AMT',
description = 'Flow generation',
)
parser.add_argument('-r', '--root', default='data/vimeo_triplet')
args = parser.parse_args()
vimeo90k_dir = args.root
vimeo90k_sequences_dir = osp.join(vimeo90k_dir, 'sequences')
vimeo90k_flow_dir = osp.join(vimeo90k_dir, 'flow')
def pred_flow(img1, img2):
img1 = torch.from_numpy(img1).float().permute(2, 0, 1) / 255.0
img2 = torch.from_numpy(img2).float().permute(2, 0, 1) / 255.0
flow = estimate(img1, img2)
flow = flow.permute(1, 2, 0).cpu().numpy()
return flow
print('Built Flow Path')
if not osp.exists(vimeo90k_flow_dir):
os.makedirs(vimeo90k_flow_dir)
for sequences_path in sorted(os.listdir(vimeo90k_sequences_dir)):
vimeo90k_sequences_path_dir = osp.join(vimeo90k_sequences_dir, sequences_path)
vimeo90k_flow_path_dir = osp.join(vimeo90k_flow_dir, sequences_path)
if not osp.exists(vimeo90k_flow_path_dir):
os.mkdir(vimeo90k_flow_path_dir)
for sequences_id in sorted(os.listdir(vimeo90k_sequences_path_dir)):
vimeo90k_flow_id_dir = osp.join(vimeo90k_flow_path_dir, sequences_id)
if not osp.exists(vimeo90k_flow_id_dir):
os.mkdir(vimeo90k_flow_id_dir)
for sequences_path in sorted(os.listdir(vimeo90k_sequences_dir)):
vimeo90k_sequences_path_dir = os.path.join(vimeo90k_sequences_dir, sequences_path)
vimeo90k_flow_path_dir = os.path.join(vimeo90k_flow_dir, sequences_path)
for sequences_id in sorted(os.listdir(vimeo90k_sequences_path_dir)):
vimeo90k_sequences_id_dir = os.path.join(vimeo90k_sequences_path_dir, sequences_id)
vimeo90k_flow_id_dir = os.path.join(vimeo90k_flow_path_dir, sequences_id)
img0_path = vimeo90k_sequences_id_dir + '/im1.png'
imgt_path = vimeo90k_sequences_id_dir + '/im2.png'
img1_path = vimeo90k_sequences_id_dir + '/im3.png'
flow_t0_path = vimeo90k_flow_id_dir + '/flow_t0.flo'
flow_t1_path = vimeo90k_flow_id_dir + '/flow_t1.flo'
img0 = read(img0_path)
imgt = read(imgt_path)
img1 = read(img1_path)
flow_t0 = pred_flow(imgt, img0)
flow_t1 = pred_flow(imgt, img1)
write(flow_t0_path, flow_t0)
write(flow_t1_path, flow_t1)
print('Written Sequences {}'.format(sequences_path)) |