def viz(img1, img2, flo, uv1=None, filename='flow.png'):
img_1 = img1[0].permute(1, 2, 0).cpu().numpy()
img_2 = img2[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo_image = flow_viz.flow_to_image(flo)
img_1_u = img_1.astype('uint8').copy()
img_2_u = img_2.astype('uint8').copy()
# Draw points if have
if uv1 != None:
uv2 = get_uv2(uv1, flo)
point_size = 1
point_color = (0, 255, 0) # BGR
thickness = 4 # 可以为 0 、4、8
for i in range(len(uv2)):
pt1 = uv1[i]
pt2 = uv2[i]
cv2.circle(img_1_u, pt1, 2, point_color, -1)
cv2.circle(img_2_u, pt2, 2, point_color, -1)
img_flo = np.concatenate([img_1_u, img_2_u, flo_image], axis=0)
# cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
cv2.imwrite(filename, img_flo)
def postprocess(flow, w, h):
flow = flow[-1][0]
flow = flow.permute(1, 2, 0).cpu().numpy() # [C,H,W] to [H,W,C]
flow_image = flow_viz.flow_to_image(flow)
flow_image = cv2.resize(flow_image, (w, h))
flow_image = cv2.cvtColor(flow_image, cv2.COLOR_RGB2BGR) # RGB to BGR
return flow_image
def viz(img, flo, i):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
#cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
#cv2.waitKey()
## Instead of showing the img , we write it out
folder = "output_img/demo"
if not os.path.exists(folder):
os.makedirs(folder)
if i < 10:
cv2.imwrite('output_img/demo/0' + str(i) + ".jpg", img_flo[:, :,
[2, 1, 0]])
else:
cv2.imwrite('output_img/demo/' + str(i) + ".jpg", img_flo[:, :,
[2, 1, 0]])
def display(image1, image2, flow, flow_gt, name):
image1 = image1.permute(1, 2, 0).cpu().numpy()
image2 = image2.permute(1, 2, 0).cpu().numpy()
flow = flow.permute(1, 2, 0).cpu().numpy()
flow_image = flow_viz.flow_to_image(flow)
flow_gt = flow_gt.permute(1, 2, 0).cpu().numpy()
flow_gt_image = flow_viz.flow_to_image(flow_gt)
im2save = np.concatenate([image1, image2, flow_image, flow_gt_image],
axis=1)
error = np.sum(np.abs(flow - flow_gt)**2, axis=2)**0.5
Image.fromarray((im2save).astype(np.uint8)).save(name)
plt.figure(), plt.imshow(error, cmap='gray'), plt.colorbar(), plt.savefig(
name.replace(".png", "_error.png"), bbox_inches='tight')
def create_kitti_submission(model,
iters=24,
output_path='kitti_submission',
write_png=False):
""" Create submission for the Sintel leaderboard """
model.eval()
test_dataset = datasets.KITTI(split='testing', aug_params=None)
if not os.path.exists(output_path):
os.makedirs(output_path)
if write_png:
out_path_png = output_path + '_png'
if not os.path.exists(out_path_png):
os.makedirs(out_path_png)
for test_id in range(len(test_dataset)):
image1, image2, (frame_id, ) = test_dataset[test_id]
padder = InputPadder(image1.shape, mode='kitti')
image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda())
_, flow_pr = model(image1, image2, iters=iters, test_mode=True)
flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
if write_png:
output_filename_png = os.path.join(out_path_png, frame_id + '.png')
cv2.imwrite(output_filename_png, flow_viz.flow_to_image(flow))
output_filename = os.path.join(output_path, frame_id)
frame_utils.writeFlowKITTI(output_filename, flow)
def vizproject(savename, img1, img2, flo):
u = flo[:,:,0]
v = flo[:,:,1]
refimg = img1
curimg = img2
imgproj = np.zeros(img1.shape)
reverseflowu = np.zeros(u.shape)
reverseflowv = np.zeros(v.shape)
for i in range(img1.shape[0]):
for j in range(img1.shape[1]):
ii = int(i + u[i,j])
jj = int(j + v[i,j])
if (ii>=0) and (jj>=0) and (ii<img1.shape[0]) and (jj<img1.shape[1]):
if (u[i,j]**2 + v[i,j]**2) > (reverseflowu[i,j]**2 + reverseflowv[i,j]**2):
imgproj[i,j,:] = curimg[ii,jj,:]
reverseflowu[i,j] = u[i,j]
reverseflowv[i,j] = v[i,j]
# map flow to rgb image
errormap = refimg - imgproj + 128
flo = flow_viz.flow_to_image(flo)
uimage = np.repeat(u[:,:,np.newaxis]*10, 3, axis=2)
vimage = np.repeat(v[:,:,np.newaxis]*10, 3, axis=2)
# img_flo1 = np.concatenate([refimg, 0.5*(refimg + imgproj)], axis=0)
# img_flo2 = np.concatenate([curimg, 0.5*(refimg +curimg)], axis=0)
img_flo1 = np.concatenate([refimg, uimage], axis=0)
img_flo2 = np.concatenate([curimg, vimage], axis=0)
img_flo = np.concatenate([img_flo1, img_flo2], axis=1)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
# cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
imgout = np.array(img_flo[:, :, [2,1,0]], dtype='uint8')
cv2.imwrite(savename , imgout) # cv2.waitKey()
def viz(img, img2, flo_fwd, flo_bwd, idx):
img = img[0].permute(1, 2, 0).cpu().numpy()
img2 = img2[0].permute(1, 2, 0).cpu().numpy()
flo_fwd = flo_fwd[0].permute(1, 2, 0).cpu().numpy()
flo_bwd = flo_bwd[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo_fwd = flow_viz.flow_to_image(flo_fwd)
flo_bwd = flow_viz.flow_to_image(flo_bwd)
img_flo = np.concatenate([img, img2, flo_fwd, flo_bwd], axis=0)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
cv2.imwrite(f"flow_img/{idx}.png", img_flo[:, :, [2, 1, 0]])
print(f"flow_img/{idx}.png")
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
flo = Image.fromarray(flo)
flo.save('example.png')
def viz(img, flo, out_fpath):
img = img[0].permute(1, 2, 0).cpu().numpy()
# flo = flo[0].permute(1,2,0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
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)
def evaluate_davis(model, iters=32):
""" Peform validation using the Sintel (train) split """
model.eval()
val_dataset = datasets.DAVISDataset(split='train')
for val_id in tqdm(range(len(val_dataset))):
image1, image2, image_paths = val_dataset[val_id]
image1 = image1[None].cuda()
image2 = image2[None].cuda()
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
_, flow_pr = model(image1, image2,
iters=iters, test_mode=True)
forward_flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
_, flow_pr = model(image2, image1,
iters=iters, test_mode=True)
backward_flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
# find out result storing paths
fpath = image_paths[0]
ind = fpath.rfind("/")
name = fpath[ind + 1:fpath.rfind(".")]
folder_path = fpath[:ind]
flow_folder = folder_path.replace("JPEGImages", "Flows")
flowviz_folder = folder_path.replace("JPEGImages", "FlowVizs")
flow_path = os.path.join(flow_folder, f"forward_{name}.flo")
flowviz_path = os.path.join(flowviz_folder, f"forward_{name}.png")
if not os.path.exists(flow_folder):
os.makedirs(flow_folder)
if not os.path.exists(flowviz_folder):
os.makedirs(flowviz_folder)
frame_utils.writeFlow(flow_path, forward_flow)
Image.fromarray(flow_viz.flow_to_image(forward_flow)).save(
open(flowviz_path, "wb"), format="PNG")
flow_path = os.path.join(flow_folder, f"backward_{name}.flo")
flowviz_path = os.path.join(flowviz_folder, f"backward_{name}.png")
frame_utils.writeFlow(flow_path, backward_flow)
Image.fromarray(flow_viz.flow_to_image(backward_flow)).save(
open(flowviz_path, "wb"), format="PNG")
def viz_opt_file(img, flo, out_fpath):
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()
# cv2.imwrite(out_fpath, img_flo[:, :, [2,1,0]]/255.0)
cv2.imwrite(out_fpath, img_flo[:, :, [2, 1, 0]])
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
import matplotlib.pyplot as plt
plt.imshow(img_flo / 255.0)
plt.show()
def viz(savename, img, flo):
img = img[0].permute(1,2,0).cpu().numpy()
flo = flo[0].permute(1,2,0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
# cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
imgout = np.array(img_flo[:, :, [2,1,0]], dtype='uint8')
cv2.imwrite(savename , imgout)
def display(image1, image2, flow):
image1 = image1.permute(1, 2, 0).cpu().numpy() / 255.0
image2 = image2.permute(1, 2, 0).cpu().numpy() / 255.0
flow = flow.permute(1, 2, 0).cpu().numpy()
flow_image = flow_viz.flow_to_image(flow)
flow_image = cv2.resize(flow_image, (image1.shape[1], image1.shape[0]))
cv2.imshow('image1', image1[..., ::-1])
cv2.imshow('image2', image2[..., ::-1])
cv2.imshow('flow', flow_image[..., ::-1])
cv2.waitKey()
def viz(img, flo, num):
img = img[0].permute(1,2,0).cpu().numpy()
flo = flo[0].permute(1,2,0).cpu().numpy()
print(flo[1,1])
np.save('flo.npy', flo)
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
print(flo[1,1])
cv2.imwrite('./out/%d.png'%num, flo)
print('./out/%d.png is done'%num)
img_flo = np.concatenate([img, flo], axis=0)
def viz(img, flo, i):
# Convert from rgb to bgr
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
cv2_imshow(img_flo[:, :, [2, 1, 0]] / 255.0)
filename = '{}.jpg'.format(i)
cv2.imwrite(filename, img_flo[:, :, [2, 1, 0]])
cv2.waitKey()
def viz(img, flo, i):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
plt.imshow(img_flo / 255.0)
plt.savefig(f'/home/sharif/Documents/RAFT/test_vis/{i}.png')
# clear plt
plt.clf()
plt.cla()
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
cv2.imshow('image', img_flo[:, :, [2, 1, 0]] / 255.0)
cv2.waitKey()
def toimage(self):
if self.type == 'raft':
image = flow_viz.flow_to_image(self.result)
return image
if self.type == 'farneback':
magnitude, angle = cv2.cartToPolar(self.result[..., 0],
self.result[..., 1])
mask = self.mask
mask[..., 1] = 255
mask[..., 0] = angle * 180 / np.pi / 2
mask[..., 2] = cv2.normalize(magnitude, None, 0, 255,
cv2.NORM_MINMAX)
image = cv2.cvtColor(mask, cv2.COLOR_HSV2BGR)
return image
def create_sintel_submission(model,
iters=32,
warm_start=False,
output_path='sintel_submission',
write_png=False):
""" Create submission for the Sintel leaderboard """
model.eval()
for dstype in ['clean', 'final']:
test_dataset = datasets.MpiSintel(split='test',
aug_params=None,
dstype=dstype)
flow_prev, sequence_prev = None, None
for test_id in range(len(test_dataset)):
image1, image2, (sequence, frame) = test_dataset[test_id]
if sequence != sequence_prev:
flow_prev = None
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1[None].cuda(),
image2[None].cuda())
flow_low, flow_pr = model(image1,
image2,
iters=iters,
flow_init=flow_prev,
test_mode=True)
flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
if warm_start:
flow_prev = forward_interpolate(flow_low[0])[None].cuda()
output_dir = os.path.join(output_path, dstype, sequence)
output_file = os.path.join(output_dir,
'frame%04d.flo' % (frame + 1))
if write_png:
output_dir_png = os.path.join(output_path + '_png', dstype,
sequence)
output_file_png = os.path.join(output_dir_png,
'frame%04d.png' % (frame + 1))
if not os.path.exists(output_dir_png):
os.makedirs(output_dir_png)
cv2.imwrite(output_file_png, flow_viz.flow_to_image(flow))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
frame_utils.writeFlow(output_file, flow)
sequence_prev = sequence
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 infer(model, seq_img_dir, suffix, iters=24, backward_flow=True):
if backward_flow:
flow_img_dir = os.path.join(seq_img_dir,
'../flow_backward_img_{}'.format(suffix))
flow_np_dir = os.path.join(seq_img_dir,
'../flow_backward_np_{}'.format(suffix))
# flow_np_save_path = os.path.join(seq_img_dir, '../flow_backward_{}.npy'.format(suffix))
else:
flow_img_dir = os.path.join(seq_img_dir,
'../flow_forward_img_{}'.format(suffix))
flow_np_dir = os.path.join(seq_img_dir,
'../flow_forward_np_{}'.format(suffix))
# flow_np_save_path = os.path.join(seq_img_dir, '../flow_forward_{}.npy'.format(suffix))
if not os.path.exists(flow_img_dir):
os.makedirs(flow_img_dir)
if not os.path.exists(flow_np_dir):
os.makedirs(flow_np_dir)
model.eval()
dataset = datasets.InferVideoDataset(seq_img_dir,
backward_flow=backward_flow)
# flow_list, flow_img_list = [], []
for val_id in tqdm.tqdm(range(len(dataset))):
image1, image2, path1, path2 = dataset[val_id]
image1 = image1[None].cuda()
image2 = image2[None].cuda()
padder = InputPadder(image1.shape, mode='sintel')
image1, image2 = padder.pad(image1, image2)
flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True)
flow = padder.unpad(flow_pr[0]).cpu()
# map flow to rgb image
# pdb.set_trace()
# flow = flow[0].permute(1,2,0).cpu().numpy()
flow = flow.permute(1, 2, 0).cpu().numpy()
flow_img = flow_viz.flow_to_image(flow)
# flow_list.append(flow)
# flow_img_list.append(flow_img)
imageio.imwrite(os.path.join(flow_img_dir,
path1.split('/')[-1]), flow_img)
np.save(
os.path.join(flow_np_dir,
path1.split('/')[-1].split('.')[0] + '.npy'), flow)
def viz(img, flo, count, out_dir):
# print(f' viz called with count: {count}')
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
# img_flo = np.concatenate([img, flo], axis=0)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
# import pdb
# pdb.set_trace()
# cv2.imwrite(f'/content/demo/{count}.png', img_flo[:, :, [2,1,0]]/255.0)
cv2.imwrite(str(Path(out_dir) / f'{str(count).zfill(5)}.png'), flo)
def vis(npy_dir, output_dir):
npy_dir = Path(npy_dir)
output_dir = Path(output_dir)
npy_files = list(npy_dir.glob('*.npy'))
for i, npy_file in enumerate(npy_files):
f = str(npy_file)
of = np.load(f)
of = torch.from_numpy(of)
of = of[0].permute(1, 2, 0).numpy()
of = flow_viz.flow_to_image(of)
img = Image.fromarray(of)
output_f = output_dir / npy_file.stem
output_f = str(output_f) + '.jpg'
img.save(output_f)
if i % 20 == 0: print(f'{i}/{len(npy_files)}')
def vizualize_flow(img, flo, save, counter):
# permute the channels and change device is necessary
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
flo = cv2.cvtColor(flo, cv2.COLOR_RGB2BGR)
# concatenate, save and show images
img_flo = np.concatenate([img, flo], axis=0)
if save:
cv2.imwrite(f"demo_frames/frame_{str(counter)}.jpg", img_flo)
cv2.imshow("Optical Flow", img_flo / 255.0)
k = cv2.waitKey(25) & 0xFF
if k == 27:
return False
return True
def viz(img, flo):
img = img[0].permute(1,2,0).cpu().numpy()
flo = flo[0].permute(1,2,0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
#img_flo = np.concatenate([img, flo], axis=0)
# import matplotlib.pyplot as plt
plt.subplot(221);plt.imshow(img / 255.0)
plt.subplot(222);plt.imshow(flo / 255.0)
if flo.ndim == 3:
flo = flo[:,:,0]
plt.subplot(223);plt.imshow(flo / 255.0, cmap='gray')
plt.subplot(224); plt.imshow(flo / 255.0, cmap = plt.cm.gray_r)
plt.show()
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
#flo = np.concatenate([img, flo], axis=0)
flo = np.concatenate([img, flo], axis=0)
# shape = flo.shape
# result = np.zeros(shape)
# print(result)
# for x in range(0, flo[:, , ]):
# for y in range(0, shape[1]):
# for z in range(0, shape[2]):
# if flo[x, y, z] <= 200:
# result[x, y, z] = 0
# flo = np.maximum(flo, 210)
out = []
a = 0
# if flo[[255,255,255]] in flo:
# a = a+1
# print(a)
# for i in range(250, 255):
# out.append(i)
# if out in flo:
# flo = 255
# flo = flo[:, :, [2, 1, 0]] / 255.0
# else:
# flo = flo[:, :, [2, 1, 0]] / 255.0
# print('total_numebr:{}'.format(flo.shape[0]*flo.shape[1]))
# # if [255,255,255] in flo:
# # a = a + 1
# # flo.Counter([255,255,255])
# print(sum(x.Counter([255,255,255]) for x in flo))
# print(flo)
# print((flo[:, :, [2, 1, 0]] / 255.0))
return flo[:, :, [2, 1, 0]] / 255.0
def write_vid(args):
# Setup model
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model))
model = model.module
model.to(DEVICE)
model.eval()
if args.write_location == None:
raise ValueError('You must provide a place to write the video.')
# Setup file write location from arguments
write_location = os.path.join(args.write_location, 'processedVideo.mp4')
with torch.no_grad():
# Get video and read to cv2 object
vid = cv2.VideoCapture(args.path)
baseVideoInformation = {}
# Retrieve data to write out later
baseVideoInformation["Channels"] = 3
baseVideoInformation["FPS"] = int(vid.get(cv2.CAP_PROP_FPS))
baseVideoInformation["Width"] = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
baseVideoInformation["Height"] = int(vid.get(
cv2.CAP_PROP_FRAME_HEIGHT))
baseVideoInformation["Frames"] = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
# Count frames read
count = 0
# Create video writer object
fourcc = cv2.VideoWriter_fourcc('M', 'P', 'E', 'G')
out = cv2.VideoWriter(write_location, fourcc,
baseVideoInformation["FPS"], (1024, 880))
# Time execution option
if (args.time == "True" or args.time == "true"):
tic = time.perf_counter()
while (vid.isOpened()):
ret, frame1 = vid.read()
count += 1
if not ret: # Faliure to get the next frame
print('Error on first frame.')
break
# Read and prepare the first frame ( convert BGR to RGB then resize, and reframe)
image1 = image_resize(cv2.cvtColor(frame1, cv2.COLOR_BGR2RGB),
(440, 1024)) # This gives [h], [w], [ch]
image1 = np.transpose(image1, axes=[
2, 0, 1
])[np.
newaxis, :, :, :] # We need to reframe to [ex], [ch], [h], [w]
image1 = torch.from_numpy(image1).to(DEVICE)
# Initialise lists
if (args.save_data):
flow_map = []
img_map = []
# Main processing loop
while (1):
ret, frame2 = vid.read()
count += 1
if not ret: # Faliure to get the next frame
print('Finished trying to retrieve frame {}, out of {}.'.
format(count, baseVideoInformation["Frames"]))
break
if not (count % 100):
print('Processing frame {}, out of {}.'.format(
count, baseVideoInformation["Frames"]))
# Read and process the following frames
image2 = image_resize(cv2.cvtColor(frame2, cv2.COLOR_BGR2RGB),
(440, 1024))
image2 = np.transpose(image2, axes=[2, 0,
1])[np.newaxis, :, :, :]
image2 = torch.from_numpy(image2).to(DEVICE)
# Pass a pair of images(to calculate flow from) to the model
flow_low, flow_up = model(image1,
image2,
iters=20,
test_mode=True)
### Write to file ###
image1 = image1[0].permute(1, 2, 0).cpu().numpy()
flow_up = flow_up[0].permute(1, 2, 0).cpu().numpy()
# Append data to a list to save later
if (args.save_data):
flow_map.append(flow_up)
img_map.append(image1)
# map flow to rgb image
flo = flow_viz.flow_to_image(flow_up, rad_max=args.scale)
img_flo = np.concatenate([image1, flo], axis=0)
# Reframe BGR back to RGB
img_write = img_flo[:, :, [2, 1, 0]]
out.write(img_write)
# Write image2 to image1
image1 = image2
# Time execution option
if (args.time == "True" or args.time == "true"):
toc = time.perf_counter()
print(f"Processing completed in {toc - tic:0.4f} seconds")
out.release()
cv2.destroyAllWindows()
print("Video resources relieved")
# Convert data lists to pickle and save
if (args.save_data):
flow_map = np.asarray(flow_map)
img_map = np.asarray(img_map)
# Extract relevant filename
filename = args.path.split(".")[0].split("/")[-1]
flow_name = "{}_flow_map.pkl".format(filename)
img_name = "{}_img_map.pkl".format(filename)
# save
with open(flow_name, 'wb') as f:
pickle.dump(flow_map, f)
with open(img_name, 'wb') as f:
pickle.dump(img_map, f)
break
current_img = frame_1
raft_img_1 = np.array(cv2.medianBlur(past_img,5)).astype(np.uint8)
raft_img_1 = torch.from_numpy(raft_img_1).permute(2, 0, 1).float()
raft_img_1 = raft_img_1[None].to(DEVICE)
raft_img_2 = np.array(cv2.medianBlur(current_img,5)).astype(np.uint8)
raft_img_2 = torch.from_numpy(raft_img_2).permute(2, 0, 1).float()
raft_img_2 = raft_img_2[None].to(DEVICE)
#padder = InputPadder(raft_img_1.shape)
#raft_img_2, raft_img_1 = padder.pad(raft_img_2, raft_img_1)
flow_low, flow_up = raft(raft_img_2, raft_img_1, iters=5, test_mode=True)
flow_up = flow_up[0].permute(1,2,0).detach().cpu().numpy()
flow_up = flow_viz.flow_to_image(flow_up)
#cv2.imwrite('{:06d}'.format(count) + '.png',flow_up)
merge = cv2.addWeighted(bbox, 1, flow_up, .5, 0)
merge_img1 = np.concatenate((frame_1,bbox), axis = 0)
merge_img2 = np.concatenate((flow_up, merge), axis = 0)
merge_img = np.concatenate((merge_img1,merge_img2), axis = 1)
cv2.imshow('image',merge_img)
cv2.waitKey(1)
#out.write(merge_img)
#cv2.imwrite('{:06d}'.format(count) + '.png',merge_img)
print('elapsed time: {}'.format(time.time()-start))
cap.release()
cv2.destroyAllWindows()
