def inference(args):
# get the RAFT model
model = RAFT(args)
# load pretrained weights
pretrained_weights = torch.load(args.model)
save = args.save
if save:
if not os.path.exists("demo_frames"):
os.mkdir("demo_frames")
if torch.cuda.is_available():
device = "cuda"
# parallel between available GPUs
model = torch.nn.DataParallel(model)
# load the pretrained weights into model
model.load_state_dict(pretrained_weights)
model.to(device)
else:
device = "cpu"
# change key names for CPU runtime
pretrained_weights = get_cpu_model(pretrained_weights)
# load the pretrained weights into model
model.load_state_dict(pretrained_weights)
# change model's mode to evaluation
model.eval()
video_path = args.video
# capture the video and get the first frame
cap = cv2.VideoCapture(video_path)
ret, frame_1 = cap.read()
# frame preprocessing
frame_1 = frame_preprocess(frame_1, device)
counter = 0
with torch.no_grad():
while True:
# read the next frame
ret, frame_2 = cap.read()
if not ret:
break
# preprocessing
frame_2 = frame_preprocess(frame_2, device)
# predict the flow
flow_low, flow_up = model(frame_1,
frame_2,
iters=20,
test_mode=True)
# transpose the flow output and convert it into numpy array
ret = vizualize_flow(frame_1, flow_up, save, counter)
if not ret:
break
frame_1 = frame_2
counter += 1
def demo(args):
model = RAFT(args)
model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load(args.model))
model.to(DEVICE)
model.eval()
with torch.no_grad():
# sintel images
image1 = load_image('images/sintel_0.png')
image2 = load_image('images/sintel_1.png')
flow_predictions = model(image1,
image2,
iters=args.iters,
upsample=False)
display(image1[0], image2[0], flow_predictions[-1][0])
# kitti images
image1 = load_image('images/kitti_0.png')
image2 = load_image('images/kitti_1.png')
flow_predictions = model(image1, image2, iters=16)
display(image1[0], image2[0], flow_predictions[-1][0])
# davis images
image1 = load_image('images/davis_0.jpg')
image2 = load_image('images/davis_1.jpg')
flow_predictions = model(image1, image2, iters=16)
display(image1[0], image2[0], flow_predictions[-1][0])
def demo(args):
model = RAFT(args)
model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load(args.model))
model.to(DEVICE)
model.eval()
with torch.no_grad():
cap = cv2.VideoCapture('video.mp4')
_, left_frame = cap.read()
h, w, _ = left_frame.shape
left_tensor = preprocess(left_frame)
while (1):
_, right_frame = cap.read()
right_tensor = preprocess(right_frame)
start1 = time.time()
flow_predictions = model(left_tensor,
right_tensor,
iters=args.iters,
upsample=True)
print(time.time() - start1)
flow_image = postprocess(flow_predictions, w * 2, h * 2)
cv2.imshow('frame', flow_image)
k = cv2.waitKey(25)
if (k == 27):
break
left_tensor = right_tensor.clone()
def batch_calOpt(args):
model = RAFT(args)
model = torch.nn.DataParallel(model)
model.load_state_dict(torch.load(args.model))
model.to(DEVICE)
model.eval()
# Transform
transform = transforms.Compose([transforms.ToTensor()])
# ImageFolder and Loader
image_dataset = OpticalFlowFolder(IMAGE_FOLDER_PATH, transform=transform)
image_loader = torch.utils.data.DataLoader(image_dataset,
batch_size=BATCH_SIZE)
start = time.time()
with torch.no_grad():
for left, right in image_loader:
# Most of the time, this preprocessing is not needed
# Especially if the video dimensions are multiple of 8s
_, _, h, w = left.shape
if ((h % 8 != 0) or (w % 8 != 0)):
left = pad8(left)
right = pad8(right)
# Forward
flow_predictions = model(left,
right,
iters=args.iters,
upsample=False)
print("Time Elapsed: ", time.time() - start)
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)
flows = sorted(glob.glob(args.flow_dir + '*_up.flo'))
idx = 0
for imfile1, imfile2, flow_fn in zip(images[:-1], images[1:], flows):
idx = idx + 1
print(idx)
image1 = load_image(imfile1)
image2 = load_image(imfile2)
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
flow_up = readFlow(flow_fn)
# flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
# flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
out_fname = flow_fn.split('/')[-1].replace('.flo', '.jpg')
viz_opt_file(image1, flow_up, os.path.join(args.outdir, out_fname))
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 = natsorted(images)
for imfile1, imfile2 in tqdm(zip(images[:-1], images[1:]), total=len(images)):
try :
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) # Flow Up is the upsampled version
if args.save :
path = Path(args.path_save)
path.mkdir(parents=True, exist_ok=True)
flow = padder.unpad(flow_up[0]).permute(1, 2, 0).cpu().numpy()
frame_utils.writeFlow(imfile1.replace(args.path,args.path_save).replace('.png','.flo'), flow)
else :
viz(image1, flow_up)
except Exception as e :
print(f'Error with {imfile1} : {e}')
def demo(args):
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
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 index, (imfile1,
imfile2) in tqdm(enumerate(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, index, args.out_dir)
print(f'done processing {args.path}')
def demo(args):
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model)) # , map_location=torch.device('cpu')))
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)
flow = padder.unpad(flow_up[0]).permute(1, 2, 0).cpu().numpy()
image1 = padder.unpad(image1[0]).permute(1, 2, 0).cpu().numpy()
image2 = padder.unpad(image2[0]).permute(1, 2, 0).cpu().numpy()
subname = imfile1.split("/")
savename = os.path.join(args.result, subname[-1])
vizproject(savename, image1, image2, flow)
def run(args):
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load('models/raft-things.pth'))
model = model.module
model.to(DEVICE)
model.eval()
output_dir = Path(args.output_dir)
images_dir = Path(args.images_dir)
images = list(images_dir.glob('*.png')) + list(images_dir.glob('*.jpg'))
with torch.no_grad():
images = sorted(images)
for i in range(len(images) - 1):
im_f1 = str(images[i])
im_f2 = str(images[i + 1])
image1 = load_image(im_f1)
image2 = load_image(im_f2)
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
# 2.2 MB
of_f_name = output_dir / f'{i}.npy'
np.save(of_f_name, flow_up.cpu())
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 = load_image_list(images)
for i in range(images.shape[0] - 1):
image1 = images[i, None]
image2 = images[i + 1, None]
print(image1.size())
print(image2.size())
image1 = F.interpolate(image1, (480, 720))
image2 = F.interpolate(image2, (480, 720))
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
to_save = flow_up.squeeze(0).cpu().numpy()
np.save('tosave.npy', to_save)
viz(image1, flow_up)
def load_model(args) :
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model))
model = model.module
model.to(DEVICE)
model.eval()
return model
def prepare(args):
global model
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model))
model = model.module
model.to(DEVICE)
model.eval()
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():
for i, folder in enumerate(sorted(os.listdir(args.path))[:100]):
print(f"video {i+1}")
images = glob.glob(
os.path.join(
args.path + folder + "/images/", "*.png")) + glob.glob(
os.path.join(args.path + folder + "/images/", "*.jpg"))
masks = glob.glob(
os.path.join(
args.path + folder + "/masks/", "*.png")) + glob.glob(
os.path.join(args.path + folder + "/masks/", "*.jpg"))
fact_people = FactorsPeople(DEVICE)
ret = []
is_ret = False
images = sorted(images)[:100]
masks = sorted(masks)[:100]
count = 0
for i, imfiles in enumerate(zip(images[:-1], images[1:])):
if count >= 100:
print("count >= 100")
break
imfile1 = imfiles[0]
imfile2 = imfiles[1]
image1, _ = fact_people.get_image(imfile1, masks[i])
image2, _ = fact_people.get_image(imfile2, masks[i + 1])
image1 = image1.to(DEVICE)
image2 = image2.to(DEVICE)
image1 *= 255.0
image2 *= 255.0
_, flow_up = model(image1, image2, iters=20, test_mode=True)
if not is_ret:
is_ret = True
ret = flow_up
else:
ret = torch.cat((ret, flow_up), 0)
count += 1
np.save(f"/phoenix/S3/ab2383/data/flows/{i+1}.npy",
ret.detach().cpu().numpy())
if i == 99:
break
def save_results(args, save_path):
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(
torch.load(args.model, map_location=torch.device('cpu')))
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'))
save_list(path=f'{save_path}/image.list', list_obj=images)
length = len(images)
# images = load_image_list(images)
image1 = load_image(images[0])[None]
# Get sample points
points1 = []
for i in range(15):
for j in range(20):
# pt1 = (160, 160)
# x = pt1[0]
# y = pt1[1]
x = j * 30
y = i * 30
pt1 = (x, y)
points1.append(pt1)
for i in range(length - 1):
# image1 = images[i,None]
# image2 = images[i+1,None]
image2 = load_image(images[i + 1])[None]
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
filename = f'flow_{i}_{i+1}'
np.savez(f'{save_path}/{filename}.npz',
image1=images[i],
image2=images[i + 1],
flow=flow_up)
viz(image1,
image2,
flow_up,
uv1=points1,
filename=f'{save_path}/{filename}.png')
image1 = image2
print(f'Finish saving {i}/{length}')
def initModel(self):
from argparse import Namespace
modelname = "D:/git/RAFT/models/raft-sintel.pth"
#self.model = torch.nn.DataParallel(RAFT(Namespace(small=False,mixed_precision=False)))
self.model = torch.nn.DataParallel(
RAFT(Namespace(small=False, mixed_precision=False)))
self.model.load_state_dict(torch.load(modelname))
self.model = self.model.module
self.model.to('cuda')
self.model.eval()
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()
fps = None
with torch.no_grad():
fps = FPS().start()
try:
image1 = None
image2 = None
while True:
# Wait for a coherent pair of frames: depth and color
frames = pipeline.wait_for_frames()
color_frame = frames.get_color_frame()
if not color_frame:
continue
im = np.asanyarray(color_frame.get_data())[:, :, ::-1].astype(
np.uint8)
img = torch.from_numpy(im).permute(
2, 0, 1).float().unsqueeze(0).to(DEVICE)
if image2 is None:
image1 = img
else:
image1 = image2
image2 = img
flow_low, flow_up = model(image1,
image2,
iters=args.iters,
test_mode=True)
fps.update()
fps.stop()
print(f"FPS: {fps.fps():.2f}")
viz(image1, flow_up)
# check for escape key
key = cv2.waitKey(1)
if key == 27 or key == 1048603:
break
finally:
pipeline.stop()
cv2.destroyAllWindows()
def predict(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'))
folder = os.path.basename(args.path)
floout = os.path.join(args.outroot, folder)
rawfloout = os.path.join(args.raw_outroot, folder)
os.makedirs(floout, exist_ok=True)
os.makedirs(rawfloout, exist_ok=True)
gap = args.gap
images = sorted(images)
images_ = images[:-gap]
for index, imfile1 in enumerate(images_):
if args.reverse:
image1 = load_image(images[index+gap])
image2 = load_image(imfile1)
svfile = images[index+gap]
else:
image1 = load_image(imfile1)
image2 = load_image(images[index + gap])
svfile = imfile1
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
flopath = os.path.join(floout, os.path.basename(svfile))
rawflopath = os.path.join(rawfloout, os.path.basename(svfile))
flo = flow_up[0].permute(1, 2, 0).cpu().numpy()
# save raw flow
writeFlowFile(rawflopath[:-4]+'.flo', flo)
# save image.
flo = flow_viz.flow_to_image(flo)
cv2.imwrite(flopath[:-4]+'.png', flo[:, :, [2, 1, 0]])
def demo(args):
model = torch.nn.DataParallel(RAFT(args),
device_ids=[0, 1],
output_device=1).cuda()
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'))
with open('seqheads.txt') as fs:
image_seqheads = fs.read().splitlines()
fileidx = 0
for img1 in tqdm(image_seqheads):
city, n1, n2, suf = img1.split('/', maxsplit=1)[-1].split("_")
n2 = int(n2)
for idx in range(1, 11):
img2 = f"leftImg8bit_sequence/{city}_{n1}_{n2+idx:06d}_{suf}"
savepath = f"flow_sequence/{city}_{n1}_{n2+idx:06d}"
print(f"{img1}\n{img2}")
image1 = load_image(os.path.join(args.path, img1))
image2 = load_image(os.path.join(args.path, img2))
savepath = os.path.join(args.path, savepath)
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
image1_b = torch.cat([image1, image2], 0)
image2_b = torch.cat([image2, image1], 0)
flow_low, flow_up = model(image1_b,
image2_b,
iters=20,
test_mode=True)
flow_up_fwd = flow_up[0].unsqueeze(0)
flow_up_bwd = flow_up[1].unsqueeze(0)
# viz(image1, image2, flow_up_fwd, flow_up_bwd, fileidx)
save_flow(flow_up_fwd, f"{savepath}_flow_fwd.png")
save_flow(flow_up_bwd, f"{savepath}_flow_bwd.png")
fileidx += 1
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():
final_dict = {}
for i, folder in enumerate(sorted(os.listdir(args.path))[:100]):
images = glob.glob(
os.path.join(
args.path + folder + "/images/", "*.png")) + glob.glob(
os.path.join(args.path + folder + "/images/", "*.jpg"))
fact_people = FactorsPeople(DEVICE)
ret = []
is_ret = False
images = sorted(images)[:100]
mask_tmp_pth = "/phoenix/S3/ab2383/data/TikTok_dataset/00267/masks/0001.png"
for imfile1, imfile2 in zip(images[:-1], images[1:]):
image1, _ = fact_people.get_image(imfile1, mask_tmp_pth)
image2, _ = fact_people.get_image(imfile2, mask_tmp_pth)
image1 *= 255.0
image2 *= 255.0
_, flow_up = model(image1, image2, iters=20, test_mode=True)
if not is_ret:
is_ret = True
ret = flow_up
else:
ret = torch.cat((ret, flow_up), 0)
final_dict[i + 1] = ret.detach().cpu().numpy()
with open("/phoenix/S3/ab2383/data/hundred_hundreds.pickle",
"wb") as f:
pickle.dump(
final_dict,
f,
protocol=pickle.HIGHEST_PROTOCOL,
)
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 = load_image_list(images)
for i in range(images.shape[0] - 1):
image1 = images[i, None]
image2 = images[i + 1, None]
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
viz(image1, flow_up)
def demo(args):
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(
torch.load(args.model, map_location=torch.device('cpu')))
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 = load_image_list(images)
# Get sample points
points1 = []
for i in range(15):
for j in range(20):
# pt1 = (160, 160)
# x = pt1[0]
# y = pt1[1]
x = j * 30
y = i * 30
pt1 = (x, y)
points1.append(pt1)
for i in range(images.shape[0] - 1):
image1 = images[i, None]
image2 = images[i + 1, None]
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
flow_up_numpy = flow_up[0].permute(1, 2, 0).cpu().numpy()
viz(image1, image2, flow_up, uv1=points1)
k = cv2.waitKey(200) * 0xFF
if k == 27:
break
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)
def __init__(self, type):
self.type = type
if type == 'farneback':
self.type = 'farneback'
if type == 'raft':
parser = argparse.ArgumentParser()
parser.add_argument('--model',
nargs='?',
const='raft-models/raft-things.pth',
type=str,
help="restore checkpoint")
parser.add_argument('--path',
nargs='?',
const='frames',
type=int,
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()
args.model = 'raft-models/raft-things.pth'
self.flow_model = torch.nn.DataParallel(RAFT(args))
self.flow_model.load_state_dict(torch.load(args.model))
self.flow_model = self.flow_model.module
self.flow_model.to(DEVICE)
self.flow_model.eval()
if type == 'flownet':
print('Flownet')
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)
img = images[0]
for i in range(len(images)-1):
image1 = load_image(images[i])
image2 = load_image(images[i+1])
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, i)
grid = grid.cuda()
vgrid = Variable(grid) + flo
# scale grid to [-1,1]
vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:].clone() / max(W-1,1)-1.0
vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:].clone() / max(H-1,1)-1.0
vgrid = vgrid.permute(0,2,3,1)
output = nn.functional.grid_sample(x, vgrid)
mask = torch.autograd.Variable(torch.ones(x.size()))#.cuda()
if x.is_cuda:
mask = mask.cuda()
mask = nn.functional.grid_sample(mask, vgrid)
mask[mask<0.9999] = 0
mask[mask>0] = 1
return output*mask
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model , map_location=torch.device(DEVICE)))
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[100:101], images[101:102]):#zip(images[:-1], images[1:]):
print(imfile1 + ' ' + imfile2)
for scaling in range(8,9,1):
image1, img_orig1 = load_image(imfile1,scaling=scaling)
image2, img_orig2 = load_image(imfile2,scaling=scaling)
def run_optical_flows(args):
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model))
model = model.module
model.to(DEVICE)
model.eval()
basedir = "%s" % args.data_path
# print(basedir)
img_dir = glob.glob(basedir + '/images_*')[0] #basedir + '/images_*288'
img_path_train = os.path.join(glob.glob(img_dir)[0], '%05d.png' % 0)
img_train = cv2.imread(img_path_train)
interval = 1
of_dir = os.path.join(basedir, 'flow_i%d' % interval)
if not os.path.exists(of_dir):
os.makedirs(of_dir)
with torch.no_grad():
images = glob.glob(os.path.join(basedir, 'images/', '*.png')) + \
glob.glob(os.path.join(basedir, 'images/', '*.jpg'))
images = load_image_list(images, args.input_flow_w)
for i in range(images.shape[0] - 1):
print(i)
image1 = images[i, None]
image2 = images[i + 1, None]
_, flow_up_fwd = model(image1, image2, iters=20, test_mode=True)
_, flow_up_bwd = model(image2, image1, iters=20, test_mode=True)
flow_up_fwd = flow_up_fwd[0].cpu().numpy().transpose(1, 2, 0)
flow_up_bwd = flow_up_bwd[0].cpu().numpy().transpose(1, 2, 0)
img1 = cv2.resize(
np.uint8(
np.clip(image1[0].cpu().numpy().transpose(1, 2, 0), 0,
255)), (img_train.shape[1], img_train.shape[0]),
cv2.INTER_LINEAR)
img2 = cv2.resize(
np.uint8(
np.clip(image2[0].cpu().numpy().transpose(1, 2, 0), 0,
255)), (img_train.shape[1], img_train.shape[0]),
cv2.INTER_LINEAR)
fwd_flow = resize_flow(flow_up_fwd, img_train.shape[0],
img_train.shape[1])
# fwd_flow = refinement_flow(fwd_flow, img1, img2)
bwd_flow = resize_flow(flow_up_bwd, img_train.shape[0],
img_train.shape[1])
# bwd_flow = refinement_flow(bwd_flow, img1, img2)
fwd_mask, bwd_mask = compute_fwdbwd_mask(fwd_flow, bwd_flow)
if VIZ:
if not os.path.exists('./viz_flow'):
os.makedirs('./viz_flow')
if not os.path.exists('./viz_warp_imgs'):
os.makedirs('./viz_warp_imgs')
plt.figure(figsize=(12, 6))
plt.subplot(2, 3, 1)
plt.imshow(img1)
plt.subplot(2, 3, 4)
plt.imshow(img2)
plt.subplot(2, 3, 2)
plt.imshow(flow_viz.flow_to_image(fwd_flow) / 255.)
plt.subplot(2, 3, 3)
plt.imshow(flow_viz.flow_to_image(bwd_flow) / 255.)
plt.subplot(2, 3, 5)
plt.imshow(
flow_viz.flow_to_image(fwd_flow) / 255. *
np.float32(fwd_mask[..., np.newaxis]))
plt.subplot(2, 3, 6)
plt.imshow(
flow_viz.flow_to_image(bwd_flow) / 255. *
np.float32(bwd_mask[..., np.newaxis]))
plt.savefig('./viz_flow/%02d.jpg' % i)
plt.close()
warped_im2 = warp_flow(img2, fwd_flow)
warped_im0 = warp_flow(img1, bwd_flow)
cv2.imwrite('./viz_warp_imgs/im_%05d.jpg' % (i),
img1[..., ::-1])
cv2.imwrite('./viz_warp_imgs/im_%05d_fwd.jpg' % (i),
warped_im2[..., ::-1])
cv2.imwrite('./viz_warp_imgs/im_%05d_bwd.jpg' % (i + 1),
warped_im0[..., ::-1])
np.savez(os.path.join(of_dir, '%05d_fwd.npz' % i),
flow=fwd_flow,
mask=fwd_mask)
np.savez(os.path.join(of_dir, '%05d_bwd.npz' % (i + 1)),
flow=bwd_flow,
mask=bwd_mask)
def train(args):
model = nn.DataParallel(RAFT(args), device_ids=args.gpus)
print("Parameter Count: %d" % count_parameters(model))
if args.restore_ckpt is not None:
model.load_state_dict(torch.load(args.restore_ckpt), strict=False)
model.cuda()
model.train()
if args.stage != 'chairs':
model.module.freeze_bn()
train_loader = datasets.fetch_dataloader(args)
optimizer, scheduler = fetch_optimizer(args, model)
total_steps = 0
scaler = GradScaler(enabled=args.mixed_precision)
logger = Logger(model, scheduler)
VAL_FREQ = 5000
add_noise = True
should_keep_training = True
while should_keep_training:
for i_batch, data_blob in enumerate(train_loader):
optimizer.zero_grad()
image1, image2, flow, valid = [x.cuda() for x in data_blob]
# show_image(image1[0])
# show_image(image2[0])
if args.add_noise:
stdv = np.random.uniform(0.0, 5.0)
image1 = (image1 +
stdv * torch.randn(*image1.shape).cuda()).clamp(
0.0, 255.0)
image2 = (image2 +
stdv * torch.randn(*image2.shape).cuda()).clamp(
0.0, 255.0)
flow_predictions = model(image1, image2, iters=args.iters)
loss, metrics = sequence_loss(flow_predictions, flow, valid)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
scaler.step(optimizer)
scheduler.step()
scaler.update()
logger.push(metrics)
if total_steps % VAL_FREQ == VAL_FREQ - 1:
PATH = 'checkpoints/%d_%s.pth' % (total_steps + 1, args.name)
torch.save(model.state_dict(), PATH)
results = {}
for val_dataset in args.validation:
if val_dataset == 'chairs':
results.update(evaluate.validate_chairs(model.module))
elif val_dataset == 'sintel':
results.update(evaluate.validate_sintel(model.module))
elif val_dataset == 'kitti':
results.update(evaluate.validate_kitti(model.module))
logger.write_dict(results)
model.train()
if args.stage != 'chairs':
model.module.freeze_bn()
total_steps += 1
if total_steps > args.num_steps:
should_keep_training = False
break
logger.close()
PATH = 'checkpoints/%s.pth' % args.name
torch.save(model.state_dict(), PATH)
return PATH
def opt_flow_estimation(args):
"""
args.path: path to the directory of the dataset that contains the images.
- base_dir/
- ArgoVerse/
- video1/
- frame1.jpg
- frame2.jpg
- video2/
- BDD/
- Charades/
- LaSOT/
- YFCC100M/
"""
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():
base_dir = args.path
data_srcs = [
fn.split('/')[-1]
for fn in sorted(glob.glob(os.path.join(base_dir, '*')))
]
if args.datasrc:
data_srcs = [args.datasrc]
for data_src in data_srcs:
print("Processing", data_src)
videos = [
fn.split('/')[-1] for fn in sorted(
glob.glob(os.path.join(base_dir, data_src, '*')))
]
for idx, video in enumerate(tqdm.tqdm(videos)):
fpath = os.path.join(base_dir, data_src, video)
images = glob.glob(os.path.join(fpath, '*.png')) + \
glob.glob(os.path.join(fpath, '*.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)
# Store the flow vector
flow_fname = imfile1.split("/")[-1].replace(".jpg", ".flo")
flow_up_fname = flow_fname.split(".")
flow_up_fname[0] = flow_up_fname[0] + "_up"
flow_up_fname = ".".join(flow_up_fname)
flow_low_fname = flow_fname.split(".")
flow_low_fname[0] = flow_low_fname[0] + "_low"
flow_low_fname = ".".join(flow_low_fname)
up_fname = os.path.join(args.outdir, data_src, video,
flow_up_fname)
# low_fname = os.path.join(args.outdir, data_src, video, flow_low_fname)
if not os.path.exists(
os.path.join(args.outdir, data_src, video)):
os.makedirs(os.path.join(args.outdir, data_src, video))
flow_up = flow_up[0].permute(1, 2, 0).cpu().numpy()
# flow_low = flow_low[0].permute(1, 2, 0).cpu().numpy()
writeFlow(up_fname, flow_up)
def run(args):
in_path_left = os.path.join(args.path, "image_left")
in_path_right = os.path.join(args.path, "image_right")
out_path_foward = os.path.join(args.path, "flow_forward")
out_path_backward = os.path.join(args.path, "flow_backward")
print("Computing forward and backwrd optical flow between:")
print(f" {in_path_left}")
print(f" {in_path_right}\n")
create_dir(out_path_foward)
create_dir(out_path_backward)
num_gpus = torch.cuda.device_count()
batch_size = num_gpus
data_loader = DataLoader(dataset=FlowForwardBackwardDataset(
in_path_left, in_path_right),
batch_size=batch_size,
shuffle=False,
num_workers=2)
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model))
# model = model.module ## -> No idea why this was needed
model.to(DEVICE)
model.eval()
print("\nStart ...")
print(f"Running on {num_gpus} GPUs\n")
with torch.no_grad():
for img_str, img_left, img_right in tqdm(data_loader):
padder = InputPadder(img_left.shape)
img_left, img_right = padder.pad(img_left, img_right)
# get foward flow - left -> right
_, forward_flow_up = model(img_left,
img_right,
iters=20,
test_mode=True)
forward_flow = forward_flow_up.permute(0, 2, 3, 1).cpu().numpy()
# get backward flow - right -> left
_, backward_flow_up = model(img_right,
img_left,
iters=20,
test_mode=True)
backward_flow = backward_flow_up.permute(0, 2, 3, 1).cpu().numpy()
# Loop over the batches
for i in range(batch_size):
save_flow(forward_flow[i],
os.path.join(out_path_foward, img_str[i] + "_flo"))
save_flow(backward_flow[i],
os.path.join(out_path_backward, img_str[i] + "_flo"))
print("\nDone!")
def train(args):
model = nn.DataParallel(RAFT(args), device_ids=args.gpus)
print("Parameter Count: %d" % count_parameters(model))
#if args.restore_ckpt is not None:
# model.load_state_dict(torch.load(args.restore_ckpt), strict=False)
model.cuda()
model.train()
if args.stage != 'chairs':
model.module.freeze_bn()
train_loader = datasets_self.fetch_dataloader(args)
optimizer, scheduler = fetch_optimizer(args, model)
total_steps = 0
scaler = GradScaler(enabled=args.mixed_precision)
logger = Logger(model, scheduler)
add_noise = True
should_keep_training = True
while should_keep_training:
for i_batch, data_blob in enumerate(train_loader):
optimizer.zero_grad()
image1, image2, img1_orig, img2_orig = [
x.cuda() for x in data_blob
]
if args.add_noise:
stdv = np.random.uniform(0.0, 5.0)
image1 = (image1 +
stdv * torch.randn(*image1.shape).cuda()).clamp(
0.0, 255.0)
image2 = (image2 +
stdv * torch.randn(*image2.shape).cuda()).clamp(
0.0, 255.0)
print(image1.shape)
print(img1_orig.shape)
flow_predictions = model(image1, image2, iters=args.iters)
loss = sequence_loss(img1_orig, img2_orig, flow_predictions,
args.gamma)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
scaler.step(optimizer)
scheduler.step()
scaler.update()
#logger.push(metrics)
print('processing step ' + str(total_steps) + ', loss ' +
str(loss.item()))
if total_steps % VAL_FREQ == 0:
print('i=' + str(total_steps) + ' loss=' + str(loss))
PATH = args.save_ckpt + '/%d_%s.pth' % (total_steps + 1,
args.name)
print(PATH)
torch.save(model.state_dict(), PATH)
#results = {}
#for val_dataset in args.validation:
# if val_dataset == 'chairs':
# results.update(evaluate.validate_chairs(model.module))
# elif val_dataset == 'sintel':
# results.update(evaluate.validate_sintel(model.module))
# elif val_dataset == 'kitti':
# results.update(evaluate.validate_kitti(model.module))
# elif val_dataset == 'ESPRIT':
# results.update(evaluate.validate_ESPRIT(model.module))
#logger.write_dict(results)
model.train()
if args.stage != 'chairs':
model.module.freeze_bn()
total_steps += 1
if total_steps > args.num_steps:
should_keep_training = False
break
#logger.close()
PATH = args.save_ckpt + '/%s.pth' % args.name
torch.save(model.state_dict(), PATH)
return PATH
