def complete_flow(args, corrFlow, flow_mask, mode, edge=None): """Completes flow. """ if mode not in ['forward', 'backward']: raise NotImplementedError imgH, imgW, _, nFrame = corrFlow.shape if os.path.isdir(os.path.join(args.outroot, 'flow_comp', mode + '_flo')): compFlow = np.empty(((imgH, imgW, 2, 0)), dtype=np.float32) for flow_name in sorted( glob.glob( os.path.join(args.outroot, 'flow_comp', mode + '_flo', '*.flo'))): flow = utils.frame_utils.readFlow(flow_name) compFlow = np.concatenate((compFlow, flow[..., None]), axis=-1) return compFlow create_dir(os.path.join(args.outroot, 'flow_comp', mode + '_flo')) create_dir(os.path.join(args.outroot, 'flow_comp', mode + '_png')) compFlow = np.zeros(((imgH, imgW, 2, nFrame)), dtype=np.float32) for i in range(nFrame): print("Completing {0} flow {1:2d} <---> {2:2d}".format(mode, i, i + 1)) flow = corrFlow[:, :, :, i] flow_mask_img = flow_mask[:, :, i] if mode == 'forward' else flow_mask[:, :, i + 1] flow_mask_gradient_img = gradient_mask(flow_mask_img) if edge is not None: # imgH x (imgW - 1 + 1) x 2 gradient_x = np.concatenate( (np.diff(flow, axis=1), np.zeros( (imgH, 1, 2), dtype=np.float32)), axis=1) # (imgH - 1 + 1) x imgW x 2 gradient_y = np.concatenate( (np.diff(flow, axis=0), np.zeros( (1, imgW, 2), dtype=np.float32)), axis=0) # concatenate gradient_x and gradient_y gradient = np.concatenate((gradient_x, gradient_y), axis=2) # We can trust the gradient outside of flow_mask_gradient_img # We assume the gradient within flow_mask_gradient_img is 0. gradient[flow_mask_gradient_img, :] = 0 # Complete the flow imgSrc_gy = gradient[:, :, 2:4] imgSrc_gy = imgSrc_gy[0:imgH - 1, :, :] imgSrc_gx = gradient[:, :, 0:2] imgSrc_gx = imgSrc_gx[:, 0:imgW - 1, :] compFlow[:, :, :, i] = Poisson_blend(flow, imgSrc_gx, imgSrc_gy, flow_mask_img, edge) else: flow[:, :, 0] = rf.regionfill(flow[:, :, 0], flow_mask_img) flow[:, :, 1] = rf.regionfill(flow[:, :, 1], flow_mask_img) compFlow[:, :, :, i] = flow # Flow visualization. flow_img = utils.flow_viz.flow_to_image(compFlow[:, :, :, i]) flow_img = Image.fromarray(flow_img) # Saves the flow and flow_img. flow_img.save( os.path.join(args.outroot, 'flow_comp', mode + '_png', '%05d.png' % i)) utils.frame_utils.writeFlow( os.path.join(args.outroot, 'flow_comp', mode + '_flo', '%05d.flo' % i), compFlow[:, :, :, i]) return compFlow
def __getitem__(self, idx): flow_dir = self.data_items[idx][0] video_class_no = self.data_items[idx][1] if self.config.get_mask: mask_dir = self.data_items[idx][2] if self.isTest: output_dirs = self.data_items[idx][-1] flow_set = [] mask_set = [] flow_mask_cat_set = [] flow_masked_set = [] if self.config.MASK_MODE == 'bbox': tmp_bbox = im.random_bbox(self.config) tmp_mask = im.bbox2mask(self.config, tmp_bbox) tmp_mask = tmp_mask[0, 0, :, :] fix_mask = np.expand_dims(tmp_mask, axis=2) elif self.config.MASK_MODE == 'mid-bbox': tmp_mask = im.mid_bbox_mask(self.config) tmp_mask = tmp_mask[0, 0, :, :] fix_mask = np.expand_dims(tmp_mask, axis=2) for i in range(11): tmp_flow = cvb.read_flow(flow_dir[i]) if self.config.get_mask: tmp_mask = cv2.imread(mask_dir[i], cv2.IMREAD_UNCHANGED) tmp_mask = self._mask_tf(tmp_mask) else: if self.config.FIX_MASK: tmp_mask = fix_mask.copy() else: tmp_bbox = im.random_bbox(self.config) tmp_mask = im.bbox2mask(self.config, tmp_bbox) tmp_mask = tmp_mask[0, 0, :, :] tmp_mask = np.expand_dims(tmp_mask, axis=2) tmp_flow = self._flow_tf(tmp_flow) tmp_flow_masked = tmp_flow * (1. - tmp_mask) if self.config.INITIAL_HOLE: tmp_flow_resized = cv2.resize(tmp_flow, (self.size[1] // 2, self.size[0] // 2)) tmp_mask_resized = cv2.resize(tmp_mask, (self.size[1] // 2, self.size[0] // 2), cv2.INTER_NEAREST) tmp_flow_masked_small = tmp_flow_resized tmp_flow_masked_small[:, :, 0] = rf.regionfill(tmp_flow_resized[:, :, 0], tmp_mask_resized) tmp_flow_masked_small[:, :, 1] = rf.regionfill(tmp_flow_resized[:, :, 1], tmp_mask_resized) tmp_flow_masked = tmp_flow_masked + \ tmp_mask * cv2.resize(tmp_flow_masked_small, (self.size[1], self.size[0])) flow_masked_set.append(tmp_flow_masked) flow_set.append(tmp_flow) mask_set.append(tmp_mask) mask_set.append(tmp_mask) tmp_flow_mask_cat = np.concatenate((tmp_flow_masked, tmp_mask), axis=2) flow_mask_cat_set.append(tmp_flow_mask_cat) flow_mask_cat = np.concatenate(flow_mask_cat_set, axis=2) flow_masked = np.concatenate(flow_masked_set, axis=2) gt_flow = np.concatenate(flow_set, axis=2) mask = np.concatenate(mask_set, axis=2) flow_mask_cat = torch.from_numpy(flow_mask_cat).permute(2, 0, 1).contiguous().float() flow_masked = torch.from_numpy(flow_masked).permute(2, 0, 1).contiguous().float() gt_flow = torch.from_numpy(gt_flow).permute(2, 0, 1).contiguous().float() mask = torch.from_numpy(mask).permute(2, 0, 1).contiguous().float() if self.isTest: return flow_mask_cat, flow_masked, gt_flow, mask, output_dirs return flow_mask_cat, flow_masked, gt_flow, mask
def complete_flow(args, corrFlow, flow_mask, mode, edge=None): """Completes flow. """ if mode not in [ 'forward', 'backward', 'nonlocal_forward', 'nonlocal_backward' ]: raise NotImplementedError sh = corrFlow.shape imgH = sh[0] imgW = sh[1] nFrame = sh[-1] create_dir(os.path.join(args.outroot, 'flow_comp', mode + '_flo')) create_dir(os.path.join(args.outroot, 'flow_comp', mode + '_png')) compFlow = np.zeros(((sh)), dtype=np.float32) for i in range(nFrame): print("Completing {0} flow {1:2d} <---> {2:2d}".format(mode, i, i + 1), '\r', end='') flow = corrFlow[..., i] if mode == 'forward': flow_mask_img = flow_mask[:, :, i] flow_mask_gradient_img = gradient_mask(flow_mask_img) elif mode == 'backward': flow_mask_img = flow_mask[:, :, i + 1] flow_mask_gradient_img = gradient_mask(flow_mask_img) else: # nonlocal_backward assert edge == None if edge is not None: # imgH x (imgW - 1 + 1) x 2 gradient_x = np.concatenate( (np.diff(flow, axis=1), np.zeros( (imgH, 1, 2), dtype=np.float32)), axis=1) # (imgH - 1 + 1) x imgW x 2 gradient_y = np.concatenate( (np.diff(flow, axis=0), np.zeros( (1, imgW, 2), dtype=np.float32)), axis=0) # concatenate gradient_x and gradient_y gradient = np.concatenate((gradient_x, gradient_y), axis=2) # We can trust the gradient outside of flow_mask_gradient_img # We assume the gradient within flow_mask_gradient_img is 0. gradient[flow_mask_gradient_img, :] = 0 # Complete the flow imgSrc_gy = gradient[:, :, 2:4] imgSrc_gy = imgSrc_gy[0:imgH - 1, :, :] imgSrc_gx = gradient[:, :, 0:2] imgSrc_gx = imgSrc_gx[:, 0:imgW - 1, :] compFlow[:, :, :, i] = Poisson_blend(flow, imgSrc_gx, imgSrc_gy, flow_mask_img, edge[:, :, i]) else: if mode == 'forward' or mode == 'backward': flow[:, :, 0] = rf.regionfill(flow[:, :, 0], flow_mask_img) flow[:, :, 1] = rf.regionfill(flow[:, :, 1], flow_mask_img) compFlow[:, :, :, i] = flow elif mode == 'nonlocal_forward': flow[:, :, 0, 0] = rf.regionfill(flow[:, :, 0, 0], flow_mask[:, :, i]) flow[:, :, 1, 0] = rf.regionfill(flow[:, :, 1, 0], flow_mask[:, :, i]) flow[:, :, 0, 1] = rf.regionfill(flow[:, :, 0, 1], flow_mask[:, :, i]) flow[:, :, 1, 1] = rf.regionfill(flow[:, :, 1, 1], flow_mask[:, :, i]) flow[:, :, 0, 2] = rf.regionfill(flow[:, :, 0, 2], flow_mask[:, :, i]) flow[:, :, 1, 2] = rf.regionfill(flow[:, :, 1, 2], flow_mask[:, :, i]) else: flow[:, :, 0, 0] = rf.regionfill(flow[:, :, 0, 0], flow_mask[:, :, 0]) flow[:, :, 1, 0] = rf.regionfill(flow[:, :, 1, 0], flow_mask[:, :, 0]) flow[:, :, 0, 1] = rf.regionfill(flow[:, :, 0, 1], flow_mask[:, :, nFrame // 2]) flow[:, :, 1, 1] = rf.regionfill(flow[:, :, 1, 1], flow_mask[:, :, nFrame // 2]) flow[:, :, 0, 2] = rf.regionfill(flow[:, :, 0, 2], flow_mask[:, :, nFrame - 1]) flow[:, :, 1, 2] = rf.regionfill(flow[:, :, 1, 2], flow_mask[:, :, nFrame - 1]) # # Flow visualization. # flow_img = utils.flow_viz.flow_to_image(compFlow[:, :, :, i]) # flow_img = Image.fromarray(flow_img) # # # Saves the flow and flow_img. # flow_img.save(os.path.join(args.outroot, 'flow_comp', mode + '_png', '%05d.png'%i)) # utils.frame_utils.writeFlow(os.path.join(args.outroot, 'flow_comp', mode + '_flo', '%05d.flo'%i), compFlow[:, :, :, i]) return compFlow
def complete_flow(args, corrFlow, flow_mask, mode, minbbox_tl, minbbox_br, edge=None): """Completes flow. """ if mode not in ['forward', 'backward']: raise NotImplementedError imgH, imgW, _, nFrame = corrFlow.shape # if os.path.isdir(os.path.join(args.outroot, 'flow_comp', mode + '_flo')): # compFlow = np.empty(((imgH, imgW, 2, 0)), dtype=np.float32) # for flow_name in sorted(glob.glob(os.path.join(args.outroot, 'flow_comp', mode + '_flo', '*.flo'))): # print("Loading {0}".format(flow_name), '\r', end='') # flow = utils.frame_utils.readFlow(flow_name) # compFlow = np.concatenate((compFlow, flow[..., None]), axis=-1) # return compFlow # create_dir(os.path.join(args.outroot, 'flow_comp', mode + '_flo')) # create_dir(os.path.join(args.outroot, 'flow_comp', mode + '_png')) compFlow = corrFlow.copy() for i in range(nFrame): print("Completing {0} flow {1:2d} <---> {2:2d}".format(mode, i, i + 1), '\r', end='') flow = corrFlow[:, :, :, i] if mode == 'forward': flow_crop = flow[minbbox_tl[i][1]:minbbox_br[i][1], minbbox_tl[i][0]:minbbox_br[i][0], :] flow_mask_img = flow_mask[:, :, i] flow_mask_img_crop = flow_mask_img[ minbbox_tl[i][1]:minbbox_br[i][1], minbbox_tl[i][0]:minbbox_br[i][0]] else: flow_crop = flow[minbbox_tl[i + 1][1]:minbbox_br[i + 1][1], minbbox_tl[i + 1][0]:minbbox_br[i + 1][0], :] flow_mask_img = flow_mask[:, :, i + 1] flow_mask_img_crop = flow_mask_img[ minbbox_tl[i + 1][1]:minbbox_br[i + 1][1], minbbox_tl[i + 1][0]:minbbox_br[i + 1][0]] # cv2.imwrite("./flow_mask_img_crop.png",flow_mask_img_crop*255) flow_mask_gradient_img = gradient_mask(flow_mask_img) flow_mask_gradient_img_crop = gradient_mask(flow_mask_img_crop) if edge is not None: # imgH x (imgW - 1 + 1) x 2 gradient_x = np.concatenate( (np.diff(flow_crop, axis=1), np.zeros((flow_crop.shape[0], 1, 2), dtype=np.float32)), axis=1) # (imgH - 1 + 1) x imgW x 2 gradient_y = np.concatenate( (np.diff(flow_crop, axis=0), np.zeros((1, flow_crop.shape[1], 2), dtype=np.float32)), axis=0) # concatenate gradient_x and gradient_y gradient = np.concatenate((gradient_x, gradient_y), axis=2) # We can trust the gradient outside of flow_mask_gradient_img # We assume the gradient within flow_mask_gradient_img is 0. gradient[flow_mask_gradient_img_crop, :] = 0 # Complete the flow imgSrc_gy = gradient[:, :, 2:4] imgSrc_gy = imgSrc_gy[0:flow_crop.shape[0] - 1, :, :] imgSrc_gx = gradient[:, :, 0:2] imgSrc_gx = imgSrc_gx[:, 0:flow_crop.shape[1] - 1, :] if mode == 'forward': edge_crop = edge[minbbox_tl[i][1]:minbbox_br[i][1], minbbox_tl[i][0]:minbbox_br[i][0], i] else: edge_crop = edge[minbbox_tl[i + 1][1]:minbbox_br[i + 1][1], minbbox_tl[i + 1][0]:minbbox_br[i + 1][0], i] compFlow_crop = Poisson_blend(flow_crop, imgSrc_gx, imgSrc_gy, flow_mask_img_crop, edge_crop) #return original size if mode == 'forward': compFlow[minbbox_tl[i][1]:minbbox_br[i][1], minbbox_tl[i][0]:minbbox_br[i][0], :, i] = compFlow_crop else: compFlow[minbbox_tl[i + 1][1]:minbbox_br[i + 1][1], minbbox_tl[i + 1][0]:minbbox_br[i + 1][0], :, i] = compFlow_crop else: flow[:, :, 0] = rf.regionfill(flow[:, :, 0], flow_mask_img) flow[:, :, 1] = rf.regionfill(flow[:, :, 1], flow_mask_img) compFlow[:, :, :, i] = flow ## Flow visualization. # flow_img = utils.flow_viz.flow_to_image(compFlow[:, :, :, i]) # flow_img = Image.fromarray(flow_img) ## Saves the flow and flow_img. # flow_img.save(os.path.join(args.outroot, 'flow_comp', mode + '_png', '%05d.png'%i)) # utils.frame_utils.writeFlow(os.path.join(args.outroot, 'flow_comp', mode + '_flo', '%05d.flo'%i), compFlow[:, :, :, i]) return compFlow
def complete_flow(args, dataset, frame): """Completes flow. """ zflow, zflow_comp, zflow_masks = dataset['flow'], dataset[ 'flow_comp'], dataset['flow_masks'] if valid_checksum(dataset, dataset['flow_comp'][..., frame], 2, frame): return imgH, imgW = zflow.shape[:2] #compFlow = np.zeros(((imgH, imgW, 2, nFrame)), dtype=np.float32) for mode in range(2): flow = zflow[:, :, :, mode, frame] flow_mask_img = zflow_masks[:, :, frame + mode] flow_mask_gradient_img = gradient_mask(flow_mask_img) if args.edge_guide: zedges = dataset['edges'] # imgH x (imgW - 1 + 1) x 2 gradient_x = np.concatenate( (np.diff(flow, axis=1), np.zeros( (imgH, 1, 2), dtype=np.float32)), axis=1) # (imgH - 1 + 1) x imgW x 2 gradient_y = np.concatenate( (np.diff(flow, axis=0), np.zeros( (1, imgW, 2), dtype=np.float32)), axis=0) # concatenate gradient_x and gradient_y gradient = np.concatenate((gradient_x, gradient_y), axis=2) # We can trust the gradient outside of flow_mask_gradient_img # We assume the gradient within flow_mask_gradient_img is 0. gradient[flow_mask_gradient_img, :] = 0 # Complete the flow imgSrc_gy = gradient[:, :, 2:4] imgSrc_gy = imgSrc_gy[0:imgH - 1, :, :] imgSrc_gx = gradient[:, :, 0:2] imgSrc_gx = imgSrc_gx[:, 0:imgW - 1, :] zflow_comp[:, :, :, mode, frame] = Poisson_blend(flow, imgSrc_gx, imgSrc_gy, flow_mask_img, zedges[:, :, mode, frame]) else: flow[:, :, 0] = rf.regionfill(flow[:, :, 0], flow_mask_img) flow[:, :, 1] = rf.regionfill(flow[:, :, 1], flow_mask_img) zflow_comp[:, :, :, mode, frame] = flow # Flow visualization. flow_img = utils.flow_viz.flow_to_image(zflow_comp[:, :, :, mode, frame]) #flow_img = Image.fromarray(flow_img) # Saves the flow and flow_img. path_mode = 'forward' if mode == 0 else 'backward' cv2.imwrite( os.path.join(args.outroot, 'flow_comp', path_mode + '_png', '%09d.png' % frame), flow_img) save_checksum(dataset, dataset['flow_comp'][..., frame], 2, frame)