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)