def validate(val_loader, model, epoch, output_writers): global args batch_time = AverageMeter() flow2_EPEs = AverageMeter() # switch to evaluate mode model.eval() # end = time.time(); runtime_count = 0; sum_runtime = 0; end = time.time() for i, (input, target) in enumerate(val_loader): target = target.to(device) # concatnate the tensor input = torch.cat(input,1).to(device) # compute output start_time = time.time() if args.graymodel: output = model(torch.cat([input[:,0:1,:,:], input[:,3:4,:,:]],1)) else: output = model(input) runtime = (time.time()-start_time) #sum_runtime += runtime # if runtime_count == 0: # sum_runtime = 0 # else: print ('AvgRunTime: ', sum_runtime/runtime_count) runtime_count += 1 sum_runtime += runtime # print (runtime) flow2_EPE = args.div_flow*realEPE(output, target, sparse=args.sparse) # record EPE flow2_EPEs.update(flow2_EPE.item(), target.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i < len(output_writers): # log first output of first batches if epoch == 0: mean_values = torch.tensor([0.45,0.432,0.411], dtype=input.dtype).view(3,1,1) output_writers[i].add_image('GroundTruth', flow2rgb(args.div_flow * target[0], max_value=10), 0) output_writers[i].add_image('Inputs', (input[0,:3].cpu() + mean_values).clamp(0,1), 0) output_writers[i].add_image('Inputs', (input[0,3:].cpu() + mean_values).clamp(0,1), 1) output_writers[i].add_image('FlowNet Outputs', flow2rgb(args.div_flow * output[0], max_value=10), epoch) if i % args.print_freq == 0: print('Test: [{0}/{1}]\t Time {2}\t EPE {3}' .format(i, len(val_loader), batch_time, flow2_EPEs)) print(' * EPE {:.3f}'.format(flow2_EPEs.avg)) print ('sum_runtime', sum_runtime/runtime_count, runtime_count) return flow2_EPEs.avg
def validate(val_loader, model, epoch, output_writers): global args batch_time = AverageMeter() flow2_EPEs = AverageMeter() # switch to evaluate mode model.eval() end = time.time() for i, (input, target) in enumerate(val_loader): target = target.to(device) input = torch.cat(input, 1).to(device) # compute output output = model(input) flow2_EPE = args.div_flow * realEPE(output, target, sparse=args.sparse) # record EPE flow2_EPEs.update(flow2_EPE.item(), target.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i < len(output_writers): # log first output of first batches if epoch == 0: mean_values = torch.tensor([0.45, 0.432, 0.411], dtype=input.dtype).view(3, 1, 1) output_writers[i].add_image( 'GroundTruth', flow2rgb(args.div_flow * target[0], max_value=10), 0) output_writers[i].add_image( 'Inputs', (input[0, :3].cpu() + mean_values).clamp(0, 1), 0) output_writers[i].add_image( 'Inputs', (input[0, 3:].cpu() + mean_values).clamp(0, 1), 1) output_writers[i].add_image( 'FlowNet Outputs', flow2rgb(args.div_flow * output[0], max_value=10), epoch) if i % args.print_freq == 0: print('Test: [{0}/{1}]\t Time {2}\t EPE {3}'.format( i, len(val_loader), batch_time, flow2_EPEs)) print(' * EPE {:.3f}'.format(flow2_EPEs.avg)) return flow2_EPEs.avg
def run(img1, img2, bidirectional=False, upsampling='bilinear', div_flow=20, max_flow=None, to_rgb=False): if model is None: setup() input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0,0,0], std=[255,255,255]), transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1]) ]) img1 = input_transform(np.array(img1)) img2 = input_transform(np.array(img2)) input_var = torch.cat([img1, img2]).unsqueeze(0) if bidirectional: inverted_input_var = torch.cat([img2, img1]).unsqueeze(0) input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.to(device) flow = model(input_var) if upsampling is not None: assert upsampling in ['nearest', 'bilinear'], \ 'Upsampling mode {} not recognized'.format(upsampling) flow = torch.nn.functional.interpolate( flow, size=img1.size()[-2:], mode=upsampling, align_corners=False) if to_rgb: rgb_flow = flow2rgb(div_flow * flow[0], max_value=max_flow) rgb_flow = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0) return rgb_flow else: flow = (div_flow * flow[0]).detach().cpu().numpy().transpose(1,2,0) return flow
def validate(val_loader, model, epoch, output_writers): global args batch_time = AverageMeter() flow2_EPEs = AverageMeter() # switch to evaluate mode model.eval() end = time.time() for i, data in enumerate(val_loader): if args.data_loader == "torch": (input, target) = data if args.data_loader == "dali": input = [ data[0]["images"][:, 0:3, :, :], data[0]["images"][:, 3:6, :, :], ] target = data[0]["flow"] if args.show_val_images: for k in range(len(input[0].cpu())): f, axarr = plt.subplots(2, 2) axarr[0, 0].imshow(np.moveaxis(np.array(input[0].cpu()[k]), 0, 2)) axarr[0, 1].imshow(np.moveaxis(np.array(input[1].cpu()[k]), 0, 2)) axarr[1, 0].imshow( np.moveaxis( flow2rgb(args.div_flow * np.squeeze(target.cpu()[k]), max_value=10), 0, 2, )) plt.show() target = target.to(device) input = torch.cat(input, 1).to(device) # compute output output = model(input) # [0], input[1]) flow2_EPE = args.div_flow * realEPE(output, target, sparse=args.sparse) # record EPE flow2_EPEs.update(flow2_EPE.item(), target.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i < len(output_writers): # log first output of first batches if epoch == 0: # input = torch.cat(input, 1).to(device) mean_values = torch.tensor([0.45, 0.432, 0.411], dtype=input.dtype).view(3, 1, 1) output_writers[i].add_image( "GroundTruth", flow2rgb(args.div_flow * target[0], max_value=10), 0) output_writers[i].add_image( "Inputs", (input[0, :3].cpu() + mean_values).clamp(0, 1), 0) output_writers[i].add_image( "Inputs", (input[0, 3:].cpu() + mean_values).clamp(0, 1), 1) output_writers[i].add_image( "FlowNet Outputs", flow2rgb(args.div_flow * output[0], max_value=10), epoch, ) if i % args.print_freq == 0: print("Test: [{0}/{1}]\t Time {2}\t EPE {3}".format( i, len(val_loader), batch_time, flow2_EPEs)) print(" * EPE {:.3f}".format(flow2_EPEs.avg)) return flow2_EPEs.avg
def train(train_loader, model, optimizer, epoch, train_writer): global n_iter, args batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() flow2_EPEs = AverageMeter() if args.data_loader == "torch": epoch_size = (len(train_loader) if args.epoch_size == 0 else min( len(train_loader), args.epoch_size)) if args.data_loader == "dali": epoch_size = (9999 if args.epoch_size == 0 else min( len(train_loader), args.epoch_size)) # switch to train mode model.train() end = time.time() for i, data in enumerate(train_loader): if args.data_loader == "torch": (input, target) = data if args.data_loader == "dali": input = [ data[0]["images"][:, 0:3, :, :], data[0]["images"][:, 3:6, :, :], ] target = data[0]["flow"] if args.show_train_images: for k in range(len(input[0].cpu())): f, axarr = plt.subplots(2, 2) axarr[0, 0].imshow(np.moveaxis(np.array(input[0].cpu()[k]), 0, 2)) axarr[0, 1].imshow(np.moveaxis(np.array(input[1].cpu()[k]), 0, 2)) axarr[1, 0].imshow( np.moveaxis( flow2rgb(args.div_flow * np.squeeze(target.cpu()[k]), max_value=10), 0, 2, )) plt.show() # measure data loading time data_time.update(time.time() - end) target = target.to(device) input = torch.cat(input, 1).to(device) # compute output output = model(input) # [0], input[1]) if args.sparse: # Since Target pooling is not very precise when sparse, # take the highest resolution prediction and upsample it instead of downsampling target h, w = target.size()[-2:] output = [F.interpolate(output[0], (h, w)), *output[1:]] loss = multiscaleEPE(output, target, weights=args.multiscale_weights, sparse=args.sparse) flow2_EPE = args.div_flow * realEPE( output, target, sparse=args.sparse) # output[0] if using multi-scale loss # record loss and EPE losses.update(loss.item(), target.size(0)) train_writer.add_scalar("train_loss", loss.item(), n_iter) flow2_EPEs.update(flow2_EPE.item(), target.size(0)) # compute gradient and do optimization step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i % args.print_freq == 0: print( "Epoch: [{0}][{1}/{2}]\t Time {3}\t Data {4}\t Loss {5}\t EPE {6}" .format(epoch, i, epoch_size, batch_time, data_time, losses, flow2_EPEs)) n_iter += 1 if i >= epoch_size: break return losses.avg, flow2_EPEs.avg
def validate(test_loader, model, epoch, output_writers): global args, image_resize, sp_threshold d_label = h5py.File(gt_file, 'r') gt_temp = np.float32(d_label['davis']['left']['flow_dist']) gt_ts_temp = np.float64(d_label['davis']['left']['flow_dist_ts']) d_label = None d_set = h5py.File(testfile, 'r') gray_image = d_set['davis']['left']['image_raw'] batch_time = AverageMeter() # switch to evaluate mode model.eval() end = time.time() batch_size_v = 4 sp_threshold = 0.5 AEE_sum = 0. AEE_sum_sum = 0. AEE_sum_gt = 0. AEE_sum_sum_gt = 0. percent_AEE_sum = 0. iters = 0. scale = 1 for i, data in enumerate(test_loader, 0): former_inputs_on, former_inputs_off, latter_inputs_on, latter_inputs_off, st_time, ed_time = data if torch.sum(former_inputs_on + former_inputs_off) > 0: input_representation = torch.zeros( former_inputs_on.size(0), batch_size_v, image_resize, image_resize, former_inputs_on.size(3)).float() for b in range(batch_size_v): if b == 0: input_representation[:, 0, :, :, :] = former_inputs_on elif b == 1: input_representation[:, 1, :, :, :] = former_inputs_off elif b == 2: input_representation[:, 2, :, :, :] = latter_inputs_on elif b == 3: input_representation[:, 3, :, :, :] = latter_inputs_off # compute output input_representation = input_representation.to(device) output = model(input_representation.type(torch.cuda.FloatTensor), image_resize, sp_threshold) # pred_flow = output pred_flow = np.zeros((image_resize, image_resize, 2)) output_temp = output.cpu() pred_flow[:, :, 0] = cv2.resize(np.array(output_temp[0, 0, :, :]), (image_resize, image_resize), interpolation=cv2.INTER_LINEAR) pred_flow[:, :, 1] = cv2.resize(np.array(output_temp[0, 1, :, :]), (image_resize, image_resize), interpolation=cv2.INTER_LINEAR) U_gt_all = np.array(gt_temp[:, 0, :, :]) V_gt_all = np.array(gt_temp[:, 1, :, :]) U_gt, V_gt = estimate_corresponding_gt_flow( U_gt_all, V_gt_all, gt_ts_temp, np.array(st_time), np.array(ed_time)) gt_flow = np.stack((U_gt, V_gt), axis=2) # ----------- Visualization if epoch < 0: mask_temp = former_inputs_on + former_inputs_off + latter_inputs_on + latter_inputs_off mask_temp = torch.sum(torch.sum(mask_temp, 0), 2) mask_temp_np = np.squeeze(np.array(mask_temp)) > 0 spike_image = mask_temp spike_image[spike_image > 0] = 255 if args.render: cv2.imshow('Spike Image', np.array(spike_image, dtype=np.uint8)) gray = cv2.resize(gray_image[i], (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) if args.render: cv2.imshow('Gray Image', cv2.cvtColor(gray, cv2.COLOR_BGR2RGB)) out_temp = np.array(output_temp.cpu().detach()) x_flow = cv2.resize( np.array(out_temp[0, 0, :, :]), (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) y_flow = cv2.resize( np.array(out_temp[0, 1, :, :]), (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) flow_rgb = flow_viz_np(x_flow, y_flow) if args.render: cv2.imshow('Predicted Flow Output', cv2.cvtColor(flow_rgb, cv2.COLOR_BGR2RGB)) gt_flow_x = cv2.resize( gt_flow[:, :, 0], (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) gt_flow_y = cv2.resize( gt_flow[:, :, 1], (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) gt_flow_large = flow_viz_np(gt_flow_x, gt_flow_y) if args.render: cv2.imshow('GT Flow', cv2.cvtColor(gt_flow_large, cv2.COLOR_BGR2RGB)) masked_x_flow = cv2.resize( np.array(out_temp[0, 0, :, :] * mask_temp_np), (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) masked_y_flow = cv2.resize( np.array(out_temp[0, 1, :, :] * mask_temp_np), (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) flow_rgb_masked = flow_viz_np(masked_x_flow, masked_y_flow) if args.render: cv2.imshow( 'Masked Predicted Flow', cv2.cvtColor(flow_rgb_masked, cv2.COLOR_BGR2RGB)) gt_flow_cropped = gt_flow[2:-2, 45:-45] gt_flow_masked_x = cv2.resize( gt_flow_cropped[:, :, 0] * mask_temp_np, (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) gt_flow_masked_y = cv2.resize( gt_flow_cropped[:, :, 1] * mask_temp_np, (scale * image_resize, scale * image_resize), interpolation=cv2.INTER_LINEAR) gt_masked_flow = flow_viz_np(gt_flow_masked_x, gt_flow_masked_y) if args.render: cv2.imshow('GT Masked Flow', cv2.cvtColor(gt_masked_flow, cv2.COLOR_BGR2RGB)) cv2.waitKey(1) image_size = pred_flow.shape full_size = gt_flow.shape xsize = full_size[1] ysize = full_size[0] xcrop = image_size[1] ycrop = image_size[0] xoff = (xsize - xcrop) // 2 yoff = (ysize - ycrop) // 2 gt_flow = gt_flow[yoff:-yoff, xoff:-xoff, :] AEE, percent_AEE, n_points, AEE_sum_temp, AEE_gt, AEE_sum_temp_gt = flow_error_dense( gt_flow, pred_flow, (torch.sum(torch.sum(torch.sum(input_representation, dim=0), dim=0), dim=2)).cpu(), is_car=False) AEE_sum = AEE_sum + args.div_flow * AEE AEE_sum_sum = AEE_sum_sum + AEE_sum_temp AEE_sum_gt = AEE_sum_gt + args.div_flow * AEE_gt AEE_sum_sum_gt = AEE_sum_sum_gt + AEE_sum_temp_gt percent_AEE_sum += percent_AEE # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i < len(output_writers): # log first output of first batches # if epoch == 0: # mean_values = torch.tensor([0.411,0.432,0.45], dtype=input_representation.dtype).view(3,1,1) output_writers[i].add_image( 'FlowNet Outputs', flow2rgb(args.div_flow * output[0], max_value=10), epoch) iters += 1 print('-------------------------------------------------------') print( 'Mean AEE: {:.2f}, sum AEE: {:.2f}, Mean AEE_gt: {:.2f}, sum AEE_gt: {:.2f}, mean %AEE: {:.2f}, # pts: {:.2f}' .format(AEE_sum / iters, AEE_sum_sum / iters, AEE_sum_gt / iters, AEE_sum_sum_gt / iters, percent_AEE_sum / iters, n_points)) print('-------------------------------------------------------') gt_temp = None return AEE_sum / iters
def main(): global args, save_path args = parser.parse_args() if args.output_value == 'both': output_string = "raw output and RGB visualization" elif args.output_value == 'raw': output_string = "raw output" elif args.output_value == 'vis': output_string = "RGB visualization" print("=> will save " + output_string) data_dir = Path(args.data) print("=> fetching img pairs in '{}'".format(args.data)) if args.output is None: save_path = data_dir / 'flow' else: save_path = Path(args.output) print('=> will save everything to {}'.format(save_path)) save_path.makedirs_p() # Data loading code input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]), transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1]) ]) img_pairs = [] for ext in args.img_exts: test_files = data_dir.files('*0.{}'.format(ext)) for file in test_files: img_pair = file.parent / (file.namebase[:-1] + '1.{}'.format(ext)) if img_pair.isfile(): img_pairs.append([file, img_pair]) print('{} samples found'.format(len(img_pairs))) # create model network_data = torch.load(Path(args.pretrained), map_location=torch.device('cpu')) print("=> using pre-trained model '{}'".format(network_data['arch'])) model = models.__dict__[network_data['arch']](network_data).to(device) model.eval() cudnn.benchmark = True print("load model success!") if 'div_flow' in network_data.keys(): args.div_flow = network_data['div_flow'] for (img1_file, img2_file) in tqdm(img_pairs): img1 = input_transform(imread(img1_file)) img2 = input_transform(imread(img2_file)) input_var = torch.cat([img1, img2]).unsqueeze(0) print("read image success!") if args.bidirectional: # feed inverted pair along with normal pair inverted_input_var = torch.cat([img2, img1]).unsqueeze(0) input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.to(device) print("prepare to compute flow!") # compute output output = model(input_var) print("flow computing success!") if args.upsampling is not None: output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False) for suffix, flow_output in zip(['flow', 'inv_flow'], output): filename = save_path / '{}{}'.format(img1_file.namebase[:-1], suffix) if args.output_value in ['vis', 'both']: rgb_flow = flow2rgb(args.div_flow * flow_output, max_value=args.max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0) imwrite(filename + '.png', to_save) if args.output_value in ['raw', 'both']: # Make the flow map a HxWx3 array as in .flo files to_save = (args.div_flow * flow_output).cpu().numpy().transpose(1, 2, 0) np.save(filename + '.npy', to_save)
input_var = input_var.cuda() # compute output output = model(input_var) outs.append(np.array(output.cpu().detach())) # Upsample output = F.interpolate(output, size=img1.size()[-2:], mode='bilinear', align_corners=False) # Convert to an RGBim1. for suffix, flow_output in zip(['flow', 'inv_flow'], output): rgb_flow = flow2rgb(div_flow * flow_output, max_value=max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0) combined = np.zeros([nY, nX * 2, nC], dtype=np.uint8) combined[:, :nX, :] = mp4_ims[aa][:, :, [2, 1, 0]] # Invert back for BGR combined[:, nX:, :] = to_save export_ims.append(combined) # Save for list # Write out the MP4 fn_out = os.path.join('/'.join(mp4_fn.split('/')[:-1]), 'flowflownet_' + mp4_fn.split('/')[-1]) fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(fn_out, fourcc, 30.0, (int(2 * nX), int(nY)))
def generate_flow(self): input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]), transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1]) ]) # create model network_data = torch.load(self.model_dir, map_location='cpu') print("=> using pre-trained model '{}'".format(network_data['arch'])) model = models.__dict__[network_data['arch']](network_data).cuda() model.eval() cudnn.benchmark = True # if 'div_flow' in network_data.keys(): # args.div_flow = network_data['div_flow'] mp4_list = glob.glob(os.path.join(self.args.input_dir, "*.mp4")) print(self.args.input_dir) print(mp4_list) for mp4_fn in mp4_list: print(mp4_fn) mp4_ims = mp4_load(mp4_fn) n_im = len(mp4_ims) # Number of images export_ims = [] [nY, nX, nC] = mp4_ims[0].shape outs = [] #frame_ints = list(range(n_im)) for ii in tqdm(range(n_im - args.window)): aa = ii bb = ii + args.window img1 = input_transform(mp4_ims[aa]) img2 = input_transform(mp4_ims[bb]) input_var = torch.cat([img1, img2]).unsqueeze(0) # if args.bidirectional: # feed inverted pair along with normal pair inverted_input_var = torch.cat([img2, img1]).unsqueeze(0) input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.cuda() # compute output output = model(input_var) outs.append(np.array(output.cpu().detach())) # Upsample output = F.interpolate(output, size=img1.size()[-2:], mode='bilinear', align_corners=False) # Convert to an RGBim1. for suffix, flow_output in zip(['flow', 'inv_flow'], output): rgb_flow = flow2rgb(args.div_flow * flow_output, max_value=args.max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose( 1, 2, 0) combined = np.zeros([nY, args.join * nX, nC], dtype=np.uint8) if args.join == 2: combined[:, :nX, :] = mp4_ims[aa][:, :, [ 2, 1, 0 ]] # Invert back for BGR combined[:, nX:, :] = to_save else: combined[:, :, :] = to_save export_ims.append(combined) # Save for list self.save_flow(mp4_fn, export_ims, mp4_ims[0].shape)
def main(): #0. assert args global args, save_path args = parser.parse_args() if args.output_value == 'both': output_string = "raw output and RGB visualization" elif args.output_value == 'raw': output_string = "raw output" elif args.output_value == 'vis': output_string = "RGB visualization" print("=> will save " + output_string) data_dir = Path(args.data) print("=> fetching img pairs in '{}'".format(args.data)) if args.output is None: save_path = data_dir / 'flow' else: save_path = Path(args.output) print('=> will save everything to {}'.format(save_path)) save_path.makedirs_p() #1. Data loading code input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]), transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1]) ]) img_pairs = [] for ext in args.img_exts: test_files = data_dir.files('*1.{}'.format(ext)) for file in test_files: img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext)) if img_pair.isfile(): img_pairs.append([file, img_pair]) print('{} samples found'.format(len(img_pairs))) #2. create model network_data = torch.load(args.pretrained) print("=> using pre-trained model '{}'".format(network_data['arch'])) model = models.__dict__[network_data['arch']](network_data).to( device ) #models = {module}<module 'models' from '/home/roit/ws/flownet_pt/models/__init__.py'> model.eval() #flownetc or flownets cudnn.benchmark = True if 'div_flow' in network_data.keys(): args.div_flow = network_data['div_flow'] for (img1_file, img2_file) in tqdm(img_pairs): img1 = imread(img1_file) img2 = imread(img2_file) img1 = input_transform(img1)[:3, :, :] img2 = input_transform(img2)[:3, :, :] input_var = torch.cat([img1, img2 ]).unsqueeze(0) #input_var={Tensor}, [1,6,h,w] if args.bidirectional: # feed inverted pair along with normal pair inverted_input_var = torch.cat([img2, img1]).unsqueeze( 0) # 如果用flownetc 到时候再分开 input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.to(device) # compute output output = model(input_var) #outTensor [1,2,h/4,w/4] if args.upsampling is not None: output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False) out1 = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False) for suffix, flow_output in zip(['flow', 'inv_flow'], output): filename = save_path / '{}{}'.format(img1_file.namebase[:-1], suffix) if args.output_value in ['vis', 'both']: tmp = args.div_flow * flow_output rgb_flow = flow2rgb(tmp, max_value=args.max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0) imwrite(filename + '.png', to_save) if args.output_value in ['raw', 'both']: # Make the flow map a HxWx2 array as in .flo files to_save = (args.div_flow * flow_output).cpu().numpy().transpose(1, 2, 0) np.save(filename + '.npy', to_save)
def main(): global args, save_path args = parser.parse_args() data_dir = args.data print("=> fetching img pairs in '{}'".format(args.data)) save_path = args.output print('=> will save everything to {}'.format(save_path)) # Data loading code input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]), transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1]) ]) img_pairs = [] names = [] # for ext in args.img_exts: # test_files = data_dir.files('*1.{}'.format(ext)) # for file in test_files: # img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext)) # if img_pair.isfile(): # img_pairs.append([file, img_pair]) # print(img_pairs) canon_im1num = [ 1, 1, 2, 10, 10, 11, 21, 21, 21, 21, 22, 22, 22, 24, 25, 31, 31, 31, 32, 32, 33, 41, 43, 51, 54 ] # 25 canon_im2num = [ 2, 3, 3, 11, 12, 12, 22, 24, 25, 26, 23, 24, 27, 26, 27, 32, 33, 34, 33, 34, 34, 42, 45, 52, 55 ] # im1 = '%s/1/(1).JPG'%data_dir # im2 = '%s/2/(1).JPG'%data_dir # img_pairs.append([im1,im2]) s = 0 for i in range(0, 25): foldernum1 = canon_im1num[i] foldernum2 = canon_im2num[i] rawnum = len(glob.glob('../../data/canon/%d/*.CR2' % (foldernum1))) for j in range(rawnum): image_path1 = '../../data/canon/%d/(%d).CR2' % (foldernum1, j + 1) image_path2 = '../../data/canon/%d/(%d).CR2' % (foldernum2, j + 1) img_pairs.append([image_path1, image_path2]) n = '%d-%d_%d' % (foldernum1, foldernum2, j + 1) names.append(n) # create model network_data = torch.load(args.pretrained) print("=> using pre-trained model '{}'".format(network_data['arch'])) model = models.__dict__[network_data['arch']](network_data).to(device) model.eval() cudnn.benchmark = True if 'div_flow' in network_data.keys(): args.div_flow = network_data['div_flow'] count = -1 for (img1_file, img2_file) in tqdm(img_pairs): raw1 = rawpy.imread(img1_file) raw2 = rawpy.imread(img2_file) im1 = raw1.raw_image_visible.astype(np.float32) im1 = (im1 - 2047) / (16383 - 2047) ratio = 0.3 / np.mean(im1) im1 = np.minimum(np.maximum(im1 * ratio, 0.0), 1.0) im2 = raw2.raw_image_visible.astype(np.float32) im2 = (im2 - 2047) / (16383 - 2047) ratio = 0.3 / np.mean(im2) im2 = np.minimum(np.maximum(im2 * ratio, 0.0), 1.0) # print(ratio) im1 = np.expand_dims(im1, axis=2) H = im1.shape[0] W = im1.shape[1] image1 = np.concatenate( ( im1[0:H:2, 0:W:2, :], #r (im1[0:H:2, 1:W:2, :] + im1[1:H:2, 0:W:2, :]) / 2.0, #g im1[1:H:2, 1:W:2, :]), axis=2) #b im2 = np.expand_dims(im2, axis=2) H = im2.shape[0] W = im2.shape[1] image2 = np.concatenate( ( im2[0:H:2, 0:W:2, :], #r (im2[0:H:2, 1:W:2, :] + im2[1:H:2, 0:W:2, :]) / 2.0, #g im2[1:H:2, 1:W:2, :]), axis=2) #b i1 = crop(image1, 1920, 2944, 2) i2 = crop(image2, 1920, 2944, 2) i1 = (i1 * 255).astype('uint8') i2 = (i2 * 255).astype('uint8') i1 = cv.fastNlMeansDenoisingColored(i1, None, 10, 10, 7, 21) i2 = cv.fastNlMeansDenoisingColored(i2, None, 10, 10, 7, 21) img1 = input_transform(i1) img2 = input_transform(i2) # print(i1.shape) # print(img1.shape) input_var = torch.cat([img1, img2]).unsqueeze(0) # if args.bidirectional: # # feed inverted pair along with normal pair # inverted_input_var = torch.cat([img2, img1]).unsqueeze(0) # input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.to(device) # compute output output = model(input_var) if args.upsampling is not None: output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False) count += 1 for suffix, flow_output in zip(['flow', 'inv_flow'], output): filename = '%s/%s' % (save_path, names[count]) print(filename) # if args.output_value in['vis', 'both']: rgb_flow = flow2rgb(flow_output, max_value=args.max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0) imwrite(filename + '.png', to_save) # if args.output_value in ['raw', 'both']: # # Make the flow map a HxWx2 array as in .flo files to_save = (flow_output).cpu().numpy().transpose(1, 2, 0) # print(to_save.shape) flow_write(to_save, filename + '.flo')
def main(): global args, save_path args = parser.parse_args() save_path = args.output print('=> will save everything to {}'.format(save_path)) # Data loading code input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0,0,0], std=[255,255,255]), transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1]) ]) img_pairs = [] names = [] # for ext in args.img_exts: # test_files = data_dir.files('*1.{}'.format(ext)) # for file in test_files: # img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext)) # if img_pair.isfile(): # img_pairs.append([file, img_pair]) # print(img_pairs) canon_im1num = [1,1,2,10,10,11,21,21,21,21,22,22,22,24,25,31,31,31,32,32,33,41,43,51,54] # 25 canon_im2num = [2,3,3,11,12,12,22,24,25,26,23,24,27,26,27,32,33,34,33,34,34,42,45,52,55] sony_im1num = [1,1,1,2,11,11,12,15] # 8 sony_im2num = [2,3,4,4,12,13,13,16] # im1 = '%s/1/(1).JPG'%data_dir # im2 = '%s/2/(1).JPG'%data_dir # img_pairs.append([im1,im2]) s = 0 for i in range(0,8): foldernum1 = sony_im1num[i] foldernum2 = sony_im2num[i] rawnum = len(glob.glob('../../data/sid_Sony/%d/*.png'%(foldernum1))) for j in range(rawnum): image_path1 = '../../data/nlm/%d_%d.jpg'%(foldernum1,j+1) image_path2 = '../../data/nlm/%d_%d.jpg'%(foldernum2,j+1) img_pairs.append([image_path1, image_path2]) n = '%d-%d_%d'%(foldernum1,foldernum2,j+1) names.append(n) # create model network_data = torch.load(args.pretrained) print("=> using pre-trained model '{}'".format(network_data['arch'])) model = models.__dict__[network_data['arch']](network_data).to(device) model.eval() cudnn.benchmark = True if 'div_flow' in network_data.keys(): args.div_flow = network_data['div_flow'] count = -1 for (img1_file, img2_file) in tqdm(img_pairs): img1 = input_transform(imread(img1_file)) img2 = input_transform(imread(img2_file)) # print(i1.shape) # print(img1.shape) input_var = torch.cat([img1, img2]).unsqueeze(0) # if args.bidirectional: # # feed inverted pair along with normal pair # inverted_input_var = torch.cat([img2, img1]).unsqueeze(0) # input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.to(device) # compute output output = model(input_var) if args.upsampling is not None: output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False) count += 1 for suffix, flow_output in zip(['flow', 'inv_flow'], output): filename = '%s/%s'%(save_path,names[count]) print(filename) # if args.output_value in['vis', 'both']: rgb_flow = flow2rgb(flow_output, max_value=args.max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0) imwrite(filename + '.png', to_save) # if args.output_value in ['raw', 'both']: # # Make the flow map a HxWx2 array as in .flo files to_save = (flow_output).cpu().numpy().transpose(1,2,0) # print(to_save.shape) flow_write(to_save, filename+'.flo')
def main(): global args, save_path args = parser.parse_args() if args.output_value == 'both': output_string = "raw output and RGB visualization" elif args.output_value == 'raw': output_string = "raw output" elif args.output_value == 'vis': output_string = "RGB visualization" print("=> will save " + output_string) data_dir = Path(args.data) print("=> fetching img pairs in '{}'".format(args.data)) if args.output is None: save_path = data_dir / 'flow' else: save_path = Path(args.output) print('=> will save everything to {}'.format(save_path)) save_path.makedirs_p() # Data loading code input_transform = transforms.Compose([ flow_transforms.ArrayToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]), transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1]) ]) # Make sure the folder has extracted images in the format of "name + num + extension", where num start from 0, 1, 2, ... # Currently only test on MPI-Sintel, which comprises .png images img_pairs = [] img_list = os.listdir(data_dir) img_list = sorted(img_list) #img_list.sort() img_list = img_list[1:] print(img_list) #for ext in args.img_exts: for ind in range(len(img_list) - 1): img1 = img_list[ind] print("img1: ", img1) #test_files = data_dir.files('*1.{}'.format(ext)) img1_name, ext = img1.split('.') basename, num = img1_name.split('_') img2 = '.'.join( ['_'.join([basename, str(int(num) + 1).zfill(4)]), ext]) print("img 2: ", img2) img_pairs.append([data_dir + '/' + img1, data_dir + '/' + img2]) #print("Test files: ") #print(test_files) #for file in test_files: #img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext)) #print("file: ") #print(file) #print("img_pair: ") #print(img_pair) #if img_pair.isfile(): # img_pairs.append([file, img_pair]) print('{} samples found'.format(len(img_pairs))) #print(img_pairs) # create model network_data = torch.load(args.pretrained) print("=> using pre-trained model '{}'".format(network_data['arch'])) model = models.__dict__[network_data['arch']](network_data).to(device) model.eval() cudnn.benchmark = True if 'div_flow' in network_data.keys(): args.div_flow = network_data['div_flow'] output = data_dir + '/flow/' + 'output.mp4' print(output) fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output, fourcc, 10.0, (640, 480)) for (img1_file, img2_file) in tqdm(img_pairs): img1 = input_transform(imread(img1_file)) img2 = input_transform(imread(img2_file)) input_var = torch.cat([img1, img2]).unsqueeze(0) if args.bidirectional: # feed inverted pair along with normal pair inverted_input_var = torch.cat([img2, img1]).unsqueeze(0) input_var = torch.cat([input_var, inverted_input_var]) input_var = input_var.to(device) # compute output output = model(input_var) if args.upsampling is not None: output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False) for suffix, flow_output in zip(['flow', 'inv_flow'], output): print(img1_file.namebase) filename = save_path / '{}{}'.format(img1_file.namebase, suffix) print(filename) if args.output_value in ['vis', 'both']: rgb_flow = flow2rgb(args.div_flow * flow_output, max_value=args.max_flow) to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0) imwrite(filename + '.png', to_save) out.write(to_save) cv2.imshow('video', to_save) if (cv2.waitKey(1) & 0xFF) == ord('q'): break if args.output_value in ['raw', 'both']: # Make the flow map a HxWx2 array as in .flo files to_save = (args.div_flow * flow_output).cpu().numpy().transpose(1, 2, 0) np.save(filename + '.npy', to_save) out.release() cv2.destroyAllWindows()
def main(): global args, save_path args = parser.parse_args() input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.233], std=[0.255]), ]) ref_list = glob.glob(os.path.join(args.data, '*_img1.tif')) mov_list = glob.glob(os.path.join(args.data, '*_img2.tif')) # create model args = parser.parse_args() args.lambda_list = [0.15, 0.85, 1, 0.1] args.linear_resolution = 128 args.nonlinear_resolution = 512 model = Baseline(18, 18, args) state_dict = torch.load(args.pretrained, map_location='cpu') model.load_state_dict(state_dict['state_dict'], strict=True) model = model.cuda() model.eval() cudnn.benchmark = True i = 0 for ref_path, mov_path in zip(ref_list, mov_list): ref = input_transform(Image.open(ref_path)) mov = input_transform(Image.open(mov_path)) input = torch.cat([ref, mov], dim=0).unsqueeze(0).cuda() _, warp1, warp2, flow = model(input) timg, simg = input.chunk(2, dim=1) timg = (timg * 0.255 + 0.233) * 255 simg = (simg * 0.255 + 0.233) * 255 warp1 = (warp1 * 0.255 + 0.233) * 255 warp2 = (warp2 * 0.255 + 0.233) * 255 img = warp1.squeeze().cpu().numpy().astype(np.uint8) path = '../eval2/' + str(i) + '_warp1.tif' imsave(path, img) img = warp2.squeeze().cpu().numpy().astype(np.uint8) path = '../eval2/' + str(i) + '_warp2.tif' imsave(path, img) img = simg.squeeze().cpu().numpy().astype(np.uint8) path = '../eval2/' + str(i) + '_src.tif' imsave(path, img) img = timg.squeeze().cpu().numpy().astype(np.uint8) path = '../eval2/' + str(i) + '_tgt.tif' imsave(path, img) flow = flow.squeeze().cpu().numpy() rgb_flow = flow2rgb(10 * flow, max_value=None) image = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0) path = '../eval2/' + str(i) + '_flow.png' plt.imsave(path, image) image = image[50:461, 50:461] path = '../eval2/' + str(i) + '_flowcrop.png' plt.imsave(path, image) i = i + 1