from __future__ import absolute_import, division, print_function, unicode_literals import os import cv2 import copy import numpy as np import scipy.io as sio from tool.utils.common_utils import interp, BFconsistCheck, \ FBconsistCheck, consistCheck, get_KeySourceFrame_flowNN_gradient def get_flowNN_gradient(args, gradient_x, gradient_y, mask_RGB, mask, videoFlowF, videoFlowB, videoNonLocalFlowF, videoNonLocalFlowB): # gradient_x: imgH x (imgW - 1 + 1) x 3 x nFrame # gradient_y: (imgH - 1 + 1) x imgW x 3 x nFrame # mask_RGB: imgH x imgW x nFrame # mask: imgH x imgW x nFrame # videoFlowF: imgH x imgW x 2 x (nFrame - 1) | [u, v] # videoFlowB: imgH x imgW x 2 x (nFrame - 1) | [u, v] # videoNonLocalFlowF: imgH x imgW x 2 x 3 x nFrame # videoNonLocalFlowB: imgH x imgW x 2 x 3 x nFrame if args.Nonlocal: num_candidate = 5 else: num_candidate = 2 imgH, imgW, nFrame = mask.shape numPix = np.sum(mask) # |--------------------| |--------------------| # | y | | v | # | x * | | u * | # | | | | # |--------------------| |--------------------| # sub: numPix * 3 | [y, x, t] # flowNN: numPix * 3 * 2 | [y, x, t], [BN, FN] # HaveFlowNN: imgH * imgW * nFrame * 2 # numPixInd: imgH * imgW * nFrame # consistencyMap: imgH * imgW * 5 * nFrame | [BN, FN, NL2, NL3, NL4] # consistency_uv: imgH * imgW * [BN, FN] * [u, v] * nFrame # sub: numPix * [y, x, t] | position of mising pixels sub = np.concatenate((np.where(mask == 1)[0].reshape(-1, 1), np.where(mask == 1)[1].reshape(-1, 1), np.where(mask == 1)[2].reshape(-1, 1)), axis=1) # flowNN: numPix * [y, x, t] * [BN, FN] | flow neighbors flowNN = np.ones((numPix, 3, 2)) * 99999 # * -1 HaveFlowNN = np.ones((imgH, imgW, nFrame, 2)) * 99999 HaveFlowNN[mask, :] = 0 numPixInd = np.ones((imgH, imgW, nFrame)) * -1 consistencyMap = np.zeros((imgH, imgW, num_candidate, nFrame)) consistency_uv = np.zeros((imgH, imgW, 2, 2, nFrame)) # numPixInd[y, x, t] gives the index of the missing pixel@[y, x, t] in sub, # i.e. which row. numPixInd[y, x, t] = idx; sub[idx, :] = [y, x, t] for idx in range(len(sub)): numPixInd[sub[idx, 0], sub[idx, 1], sub[idx, 2]] = idx # Initialization frameIndSetF = range(1, nFrame) frameIndSetB = range(nFrame - 2, -1, -1) # 1. Forward Pass (backward flow propagation) print('Forward Pass......') NN_idx = 0 # BN:0 for indFrame in frameIndSetF: # Bool indicator of missing pixels at frame t holepixPosInd = (sub[:, 2] == indFrame) # Hole pixel location at frame t, i.e. [y, x, t] holepixPos = sub[holepixPosInd, :] # Calculate the backward flow neighbor. Should be located at frame t-1 flowB_neighbor = copy.deepcopy(holepixPos) flowB_neighbor = flowB_neighbor.astype(np.float32) flowB_vertical = videoFlowB[:, :, 1, indFrame - 1] # t --> t-1 flowB_horizont = videoFlowB[:, :, 0, indFrame - 1] flowF_vertical = videoFlowF[:, :, 1, indFrame - 1] # t-1 --> t flowF_horizont = videoFlowF[:, :, 0, indFrame - 1] flowB_neighbor[:, 0] += flowB_vertical[holepixPos[:, 0], holepixPos[:, 1]] flowB_neighbor[:, 1] += flowB_horizont[holepixPos[:, 0], holepixPos[:, 1]] flowB_neighbor[:, 2] -= 1 # Round the backward flow neighbor location flow_neighbor_int = np.round(copy.deepcopy(flowB_neighbor)).astype(np.int32) # Chen: I should combine the following two operations together # Check the backward/forward consistency IsConsist, _ = BFconsistCheck(flowB_neighbor, flowF_vertical, flowF_horizont, holepixPos, args.consistencyThres) trueConsist = IsConsist[IsConsist == True] BFdiff, BF_uv = consistCheck(videoFlowF[:, :, :, indFrame - 1], videoFlowB[:, :, :, indFrame - 1]) # Check out-of-boundary # Last column and last row does not have valid gradient ValidPos = np.logical_and( np.logical_and(flow_neighbor_int[:, 0] >= 0, flow_neighbor_int[:, 0] < imgH - 1), np.logical_and(flow_neighbor_int[:, 1] >= 0, flow_neighbor_int[:, 1] < imgW - 1)) # Only work with pixels that are not out-of-boundary holepixPos = holepixPos[ValidPos, :] flowB_neighbor = flowB_neighbor[ValidPos, :] flow_neighbor_int = flow_neighbor_int[ValidPos, :] IsConsist = IsConsist[ValidPos] # For each missing pixel in holepixPos|[y, x, t], # we check its backward flow neighbor flowB_neighbor|[y', x', t-1]. # Case 1: If mask[round(y'), round(x'), t-1] == 0, # the backward flow neighbor of [y, x, t] is known. # [y', x', t-1] is the backward flow neighbor. # KnownInd: Among all backward flow neighbors, which pixel is known. KnownInd = mask[flow_neighbor_int[:, 0], flow_neighbor_int[:, 1], indFrame - 1] == 0 KnownIsConsist = np.logical_and(KnownInd, IsConsist) holepixPos_valid = holepixPos[KnownIsConsist] # We save backward flow neighbor flowB_neighbor in flowNN flowNN[numPixInd[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], indFrame].astype(np.int32), :, NN_idx] = \ flowB_neighbor[KnownIsConsist, :] # flowNN[np.where(holepixPosInd == 1)[0][ValidPos][KnownIsConsist], :, 0] = \ # flowB_neighbor[KnownIsConsist, :] # We mark [y, x, t] in HaveFlowNN as 1 HaveFlowNN[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], indFrame, NN_idx] = 1 # HaveFlowNN[:, :, :, 0] # 0: Backward flow neighbor can not be reached # 1: Backward flow neighbor can be reached # -1: Pixels that do not need to be completed consistency_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], NN_idx, 0, indFrame] = np.abs(BF_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], 0]) consistency_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], NN_idx, 1, indFrame] = np.abs(BF_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], 1]) # Case 2: If mask[round(y'), round(x'), t-1] == 1, # the pixel@[round(y'), round(x'), t-1] is also occluded. # We further check if we already assign a backward flow neighbor for the backward flow neighbor # If HaveFlowNN[round(y'), round(x'), t-1] == 0, # this is isolated pixel. Do nothing. # If HaveFlowNN[round(y'), round(x'), t-1] == 1, # we can borrow the value and refine it. UnknownInd = np.invert(KnownInd) # If we already assign a backward flow neighbor@[round(y'), round(x'), t-1] HaveNNInd = HaveFlowNN[flow_neighbor_int[:, 0], flow_neighbor_int[:, 1], indFrame - 1, NN_idx] == 1 # Unknown & IsConsist & HaveNNInd Valid_ = np.logical_and.reduce((UnknownInd, HaveNNInd, IsConsist)) refineVec = np.concatenate(( (flowB_neighbor[:, 0] - flow_neighbor_int[:, 0]).reshape(-1, 1), (flowB_neighbor[:, 1] - flow_neighbor_int[:, 1]).reshape(-1, 1), np.zeros((flowB_neighbor[:, 0].shape[0])).reshape(-1, 1)), 1) # Check if the transitive backward flow neighbor of [y, x, t] is known. # Sometimes after refinement, it is no longer known. flowNN_tmp = copy.deepcopy(flowNN[numPixInd[flow_neighbor_int[:, 0], flow_neighbor_int[:, 1], indFrame - 1].astype(np.int32), :, NN_idx] + refineVec[:, :]) flowNN_tmp = np.round(flowNN_tmp).astype(np.int32) # Check out-of-boundary. flowNN_tmp may be out-of-boundary ValidPos_ = np.logical_and( np.logical_and(flowNN_tmp[:, 0] >= 0, flowNN_tmp[:, 0] < imgH - 1), np.logical_and(flowNN_tmp[:, 1] >= 0, flowNN_tmp[:, 1] < imgW - 1)) # Change the out-of-boundary value to 0, in order to run mask[y,x,t] # in the next line. It won't affect anything as ValidPos_ is saved already flowNN_tmp[np.invert(ValidPos_), :] = 0 ValidNN = mask[flowNN_tmp[:, 0], flowNN_tmp[:, 1], flowNN_tmp[:, 2]] == 0 # Valid = np.logical_and.reduce((Valid_, ValidNN, ValidPos_)) Valid = np.logical_and.reduce((Valid_, ValidPos_)) # We save the transitive backward flow neighbor flowB_neighbor in flowNN flowNN[numPixInd[holepixPos[Valid, 0], holepixPos[Valid, 1], indFrame].astype(np.int32), :, NN_idx] = \ flowNN[numPixInd[flow_neighbor_int[Valid, 0], flow_neighbor_int[Valid, 1], indFrame - 1].astype(np.int32), :, NN_idx] + refineVec[Valid, :] # We mark [y, x, t] in HaveFlowNN as 1 HaveFlowNN[holepixPos[Valid, 0], holepixPos[Valid, 1], indFrame, NN_idx] = 1 consistency_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], NN_idx, 0, indFrame] = np.maximum(np.abs(BF_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], 0]), np.abs(consistency_uv[flow_neighbor_int[Valid, 0], flow_neighbor_int[Valid, 1], NN_idx, 0, indFrame - 1])) consistency_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], NN_idx, 1, indFrame] = np.maximum(np.abs(BF_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], 1]), np.abs(consistency_uv[flow_neighbor_int[Valid, 0], flow_neighbor_int[Valid, 1], NN_idx, 1, indFrame - 1])) consistencyMap[:, :, NN_idx, indFrame] = (consistency_uv[:, :, NN_idx, 0, indFrame] ** 2 + consistency_uv[:, :, NN_idx, 1, indFrame] ** 2) ** 0.5 print("Frame {0:3d}: {1:8d} + {2:8d} = {3:8d}" .format(indFrame, np.sum(HaveFlowNN[:, :, indFrame, NN_idx] == 1), np.sum(HaveFlowNN[:, :, indFrame, NN_idx] == 0), np.sum(HaveFlowNN[:, :, indFrame, NN_idx] != 99999))) # 2. Backward Pass (forward flow propagation) print('Backward Pass......') NN_idx = 1 # FN:1 for indFrame in frameIndSetB: # Bool indicator of missing pixels at frame t holepixPosInd = (sub[:, 2] == indFrame) # Hole pixel location at frame t, i.e. [y, x, t] holepixPos = sub[holepixPosInd, :] # Calculate the forward flow neighbor. Should be located at frame t+1 flowF_neighbor = copy.deepcopy(holepixPos) flowF_neighbor = flowF_neighbor.astype(np.float32) flowF_vertical = videoFlowF[:, :, 1, indFrame] # t --> t+1 flowF_horizont = videoFlowF[:, :, 0, indFrame] flowB_vertical = videoFlowB[:, :, 1, indFrame] # t+1 --> t flowB_horizont = videoFlowB[:, :, 0, indFrame] flowF_neighbor[:, 0] += flowF_vertical[holepixPos[:, 0], holepixPos[:, 1]] flowF_neighbor[:, 1] += flowF_horizont[holepixPos[:, 0], holepixPos[:, 1]] flowF_neighbor[:, 2] += 1 # Round the forward flow neighbor location flow_neighbor_int = np.round(copy.deepcopy(flowF_neighbor)).astype(np.int32) # Check the forawrd/backward consistency IsConsist, _ = FBconsistCheck(flowF_neighbor, flowB_vertical, flowB_horizont, holepixPos, args.consistencyThres) FBdiff, FB_uv = consistCheck(videoFlowB[:, :, :, indFrame], videoFlowF[:, :, :, indFrame]) # Check out-of-boundary # Last column and last row does not have valid gradient ValidPos = np.logical_and( np.logical_and(flow_neighbor_int[:, 0] >= 0, flow_neighbor_int[:, 0] < imgH - 1), np.logical_and(flow_neighbor_int[:, 1] >= 0, flow_neighbor_int[:, 1] < imgW - 1)) # Only work with pixels that are not out-of-boundary holepixPos = holepixPos[ValidPos, :] flowF_neighbor = flowF_neighbor[ValidPos, :] flow_neighbor_int = flow_neighbor_int[ValidPos, :] IsConsist = IsConsist[ValidPos] # Case 1: KnownInd = mask[flow_neighbor_int[:, 0], flow_neighbor_int[:, 1], indFrame + 1] == 0 KnownIsConsist = np.logical_and(KnownInd, IsConsist) flowNN[numPixInd[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], indFrame].astype(np.int32), :, NN_idx] = \ flowF_neighbor[KnownIsConsist, :] HaveFlowNN[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], indFrame, NN_idx] = 1 consistency_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], NN_idx, 0, indFrame] = np.abs(FB_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], 0]) consistency_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], NN_idx, 1, indFrame] = np.abs(FB_uv[holepixPos[KnownIsConsist, 0], holepixPos[KnownIsConsist, 1], 1]) # Case 2: UnknownInd = np.invert(KnownInd) HaveNNInd = HaveFlowNN[flow_neighbor_int[:, 0], flow_neighbor_int[:, 1], indFrame + 1, NN_idx] == 1 # Unknown & IsConsist & HaveNNInd Valid_ = np.logical_and.reduce((UnknownInd, HaveNNInd, IsConsist)) refineVec = np.concatenate(( (flowF_neighbor[:, 0] - flow_neighbor_int[:, 0]).reshape(-1, 1), (flowF_neighbor[:, 1] - flow_neighbor_int[:, 1]).reshape(-1, 1), np.zeros((flowF_neighbor[:, 0].shape[0])).reshape(-1, 1)), 1) # Check if the transitive backward flow neighbor of [y, x, t] is known. # Sometimes after refinement, it is no longer known. flowNN_tmp = copy.deepcopy(flowNN[numPixInd[flow_neighbor_int[:, 0], flow_neighbor_int[:, 1], indFrame + 1].astype(np.int32), :, NN_idx] + refineVec[:, :]) flowNN_tmp = np.round(flowNN_tmp).astype(np.int32) # Check out-of-boundary. flowNN_tmp may be out-of-boundary ValidPos_ = np.logical_and( np.logical_and(flowNN_tmp[:, 0] >= 0, flowNN_tmp[:, 0] < imgH - 1), np.logical_and(flowNN_tmp[:, 1] >= 0, flowNN_tmp[:, 1] < imgW - 1)) # Change the out-of-boundary value to 0, in order to run mask[y,x,t] # in the next line. It won't affect anything as ValidPos_ is saved already flowNN_tmp[np.invert(ValidPos_), :] = 0 ValidNN = mask[flowNN_tmp[:, 0], flowNN_tmp[:, 1], flowNN_tmp[:, 2]] == 0 # Valid = np.logical_and.reduce((Valid_, ValidNN, ValidPos_)) Valid = np.logical_and.reduce((Valid_, ValidPos_)) # We save the transitive backward flow neighbor flowB_neighbor in flowNN flowNN[numPixInd[holepixPos[Valid, 0], holepixPos[Valid, 1], indFrame].astype(np.int32), :, NN_idx] = \ flowNN[numPixInd[flow_neighbor_int[Valid, 0], flow_neighbor_int[Valid, 1], indFrame + 1].astype(np.int32), :, NN_idx] + refineVec[Valid, :] # We mark [y, x, t] in HaveFlowNN as 1 HaveFlowNN[holepixPos[Valid, 0], holepixPos[Valid, 1], indFrame, NN_idx] = 1 consistency_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], NN_idx, 0, indFrame] = np.maximum(np.abs(FB_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], 0]), np.abs(consistency_uv[flow_neighbor_int[Valid, 0], flow_neighbor_int[Valid, 1], NN_idx, 0, indFrame + 1])) consistency_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], NN_idx, 1, indFrame] = np.maximum(np.abs(FB_uv[holepixPos[Valid, 0], holepixPos[Valid, 1], 1]), np.abs(consistency_uv[flow_neighbor_int[Valid, 0], flow_neighbor_int[Valid, 1], NN_idx, 1, indFrame + 1])) consistencyMap[:, :, NN_idx, indFrame] = (consistency_uv[:, :, NN_idx, 0, indFrame] ** 2 + consistency_uv[:, :, NN_idx, 1, indFrame] ** 2) ** 0.5 print("Frame {0:3d}: {1:8d} + {2:8d} = {3:8d}" .format(indFrame, np.sum(HaveFlowNN[:, :, indFrame, NN_idx] == 1), np.sum(HaveFlowNN[:, :, indFrame, NN_idx] == 0), np.sum(HaveFlowNN[:, :, indFrame, NN_idx] != 99999))) # Interpolation gradient_x_BN = copy.deepcopy(gradient_x) gradient_y_BN = copy.deepcopy(gradient_y) gradient_x_FN = copy.deepcopy(gradient_x) gradient_y_FN = copy.deepcopy(gradient_y) for indFrame in range(nFrame): # Index of missing pixel whose backward flow neighbor is from frame indFrame SourceFmInd = np.where(flowNN[:, 2, 0] == indFrame) print("{0:8d} pixels are from source Frame {1:3d}" .format(len(SourceFmInd[0]), indFrame)) # The location of the missing pixel whose backward flow neighbor is # from frame indFrame flowNN[SourceFmInd, 0, 0], flowNN[SourceFmInd, 1, 0] if len(SourceFmInd[0]) != 0: # |--------------------| # | y | # | x * | # | | # |--------------------| # sub: numPix x 3 [y, x, t] # img: [y, x] # interp(img, x, y) gradient_x_BN[sub[SourceFmInd[0], :][:, 0], sub[SourceFmInd[0], :][:, 1], :, sub[SourceFmInd[0], :][:, 2]] = \ interp(gradient_x_BN[:, :, :, indFrame], flowNN[SourceFmInd, 1, 0].reshape(-1), flowNN[SourceFmInd, 0, 0].reshape(-1)) gradient_y_BN[sub[SourceFmInd[0], :][:, 0], sub[SourceFmInd[0], :][:, 1], :, sub[SourceFmInd[0], :][:, 2]] = \ interp(gradient_y_BN[:, :, :, indFrame], flowNN[SourceFmInd, 1, 0].reshape(-1), flowNN[SourceFmInd, 0, 0].reshape(-1)) assert(((sub[SourceFmInd[0], :][:, 2] - indFrame) <= 0).sum() == 0) for indFrame in range(nFrame - 1, -1, -1): # Index of missing pixel whose forward flow neighbor is from frame indFrame SourceFmInd = np.where(flowNN[:, 2, 1] == indFrame) print("{0:8d} pixels are from source Frame {1:3d}" .format(len(SourceFmInd[0]), indFrame)) if len(SourceFmInd[0]) != 0: gradient_x_FN[sub[SourceFmInd[0], :][:, 0], sub[SourceFmInd[0], :][:, 1], :, sub[SourceFmInd[0], :][:, 2]] = \ interp(gradient_x_FN[:, :, :, indFrame], flowNN[SourceFmInd, 1, 1].reshape(-1), flowNN[SourceFmInd, 0, 1].reshape(-1)) gradient_y_FN[sub[SourceFmInd[0], :][:, 0], sub[SourceFmInd[0], :][:, 1], :, sub[SourceFmInd[0], :][:, 2]] = \ interp(gradient_y_FN[:, :, :, indFrame], flowNN[SourceFmInd, 1, 1].reshape(-1), flowNN[SourceFmInd, 0, 1].reshape(-1)) assert(((indFrame - sub[SourceFmInd[0], :][:, 2]) <= 0).sum() == 0) # New mask mask_tofill = np.zeros((imgH, imgW, nFrame)).astype(np.bool) for indFrame in range(nFrame): if args.Nonlocal: consistencyMap[:, :, 2, indFrame], _ = consistCheck( videoNonLocalFlowB[:, :, :, 0, indFrame], videoNonLocalFlowF[:, :, :, 0, indFrame]) consistencyMap[:, :, 3, indFrame], _ = consistCheck( videoNonLocalFlowB[:, :, :, 1, indFrame], videoNonLocalFlowF[:, :, :, 1, indFrame]) consistencyMap[:, :, 4, indFrame], _ = consistCheck( videoNonLocalFlowB[:, :, :, 2, indFrame], videoNonLocalFlowF[:, :, :, 2, indFrame]) HaveNN = np.zeros((imgH, imgW, num_candidate)) if args.Nonlocal: HaveKeySourceFrameFlowNN, gradient_x_KeySourceFrameFlowNN, gradient_y_KeySourceFrameFlowNN = \ get_KeySourceFrame_flowNN_gradient(sub, indFrame, mask, videoNonLocalFlowB, videoNonLocalFlowF, gradient_x, gradient_y, args.consistencyThres) HaveNN[:, :, 2] = HaveKeySourceFrameFlowNN[:, :, 0] == 1 HaveNN[:, :, 3] = HaveKeySourceFrameFlowNN[:, :, 1] == 1 HaveNN[:, :, 4] = HaveKeySourceFrameFlowNN[:, :, 2] == 1 HaveNN[:, :, 0] = HaveFlowNN[:, :, indFrame, 0] == 1 HaveNN[:, :, 1] = HaveFlowNN[:, :, indFrame, 1] == 1 NotHaveNN = np.logical_and(np.invert(HaveNN.astype(np.bool)), np.repeat(np.expand_dims((mask[:, :, indFrame]), 2), num_candidate, axis=2)) if args.Nonlocal: HaveNN_sum = np.logical_or.reduce((HaveNN[:, :, 0], HaveNN[:, :, 1], HaveNN[:, :, 2], HaveNN[:, :, 3], HaveNN[:, :, 4])) else: HaveNN_sum = np.logical_or.reduce((HaveNN[:, :, 0], HaveNN[:, :, 1])) gradient_x_Candidate = np.zeros((imgH, imgW, 3, num_candidate)) gradient_y_Candidate = np.zeros((imgH, imgW, 3, num_candidate)) gradient_x_Candidate[:, :, :, 0] = gradient_x_BN[:, :, :, indFrame] gradient_y_Candidate[:, :, :, 0] = gradient_y_BN[:, :, :, indFrame] gradient_x_Candidate[:, :, :, 1] = gradient_x_FN[:, :, :, indFrame] gradient_y_Candidate[:, :, :, 1] = gradient_y_FN[:, :, :, indFrame] if args.Nonlocal: gradient_x_Candidate[:, :, :, 2] = gradient_x_KeySourceFrameFlowNN[:, :, :, 0] gradient_y_Candidate[:, :, :, 2] = gradient_y_KeySourceFrameFlowNN[:, :, :, 0] gradient_x_Candidate[:, :, :, 3] = gradient_x_KeySourceFrameFlowNN[:, :, :, 1] gradient_y_Candidate[:, :, :, 3] = gradient_y_KeySourceFrameFlowNN[:, :, :, 1] gradient_x_Candidate[:, :, :, 4] = gradient_x_KeySourceFrameFlowNN[:, :, :, 2] gradient_y_Candidate[:, :, :, 4] = gradient_y_KeySourceFrameFlowNN[:, :, :, 2] consistencyMap[:, :, :, indFrame] = np.exp( - consistencyMap[:, :, :, indFrame] / args.alpha) consistencyMap[NotHaveNN[:, :, 0], 0, indFrame] = 0 consistencyMap[NotHaveNN[:, :, 1], 1, indFrame] = 0 if args.Nonlocal: consistencyMap[NotHaveNN[:, :, 2], 2, indFrame] = 0 consistencyMap[NotHaveNN[:, :, 3], 3, indFrame] = 0 consistencyMap[NotHaveNN[:, :, 4], 4, indFrame] = 0 weights = (consistencyMap[HaveNN_sum, :, indFrame] * HaveNN[HaveNN_sum, :]) / ((consistencyMap[HaveNN_sum, :, indFrame] * HaveNN[HaveNN_sum, :]).sum(axis=1, keepdims=True)) # Fix the numerical issue. 0 / 0 fix = np.where((consistencyMap[HaveNN_sum, :, indFrame] * HaveNN[HaveNN_sum, :]).sum(axis=1, keepdims=True) == 0)[0] weights[fix, :] = HaveNN[HaveNN_sum, :][fix, :] / HaveNN[HaveNN_sum, :][fix, :].sum(axis=1, keepdims=True) # Fuse RGB channel independently gradient_x[HaveNN_sum, 0, indFrame] = \ np.sum(np.multiply(gradient_x_Candidate[HaveNN_sum, 0, :], weights), axis=1) gradient_x[HaveNN_sum, 1, indFrame] = \ np.sum(np.multiply(gradient_x_Candidate[HaveNN_sum, 1, :], weights), axis=1) gradient_x[HaveNN_sum, 2, indFrame] = \ np.sum(np.multiply(gradient_x_Candidate[HaveNN_sum, 2, :], weights), axis=1) gradient_y[HaveNN_sum, 0, indFrame] = \ np.sum(np.multiply(gradient_y_Candidate[HaveNN_sum, 0, :], weights), axis=1) gradient_y[HaveNN_sum, 1, indFrame] = \ np.sum(np.multiply(gradient_y_Candidate[HaveNN_sum, 1, :], weights), axis=1) gradient_y[HaveNN_sum, 2, indFrame] = \ np.sum(np.multiply(gradient_y_Candidate[HaveNN_sum, 2, :], weights), axis=1) mask_tofill[np.logical_and(np.invert(HaveNN_sum), mask[:, :, indFrame]), indFrame] = True return gradient_x, gradient_y, mask_tofill