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()