def video_completion(args):
print('Pytorch 1.6.0 for flow prediction')
try:
RAFT_model = initialize_RAFT(args)
except:
RAFT_model = None
# Loads frames.
filename_list = glob.glob(os.path.join(args.path, '*.png')) + \
glob.glob(os.path.join(args.path, '*.jpg'))
# Obtains imgH, imgW and nFrame.
imgH, imgW = np.array(Image.open(filename_list[0])).shape[:2]
nFrame = len(filename_list)
# Loads video.
video = []
for filename in sorted(filename_list):
video.append(
torch.from_numpy(np.array(Image.open(filename)).astype(
np.uint8)).permute(2, 0, 1).float())
video = torch.stack(video, dim=0)
video = video.to('cuda')
# Calcutes the corrupted flow.
corrFlowF = calculate_flow(args, RAFT_model, video, 'forward')
corrFlowB = calculate_flow(args, RAFT_model, video, 'backward')
print('\nFinish Calculating flow.')
# Makes sure video is in BGR (opencv) format.
video = video.permute(2, 3, 1, 0).cpu().numpy()[:, :, ::-1, :] / 255.
if args.mode == 'video_extrapolation':
# Creates video and flow where the extrapolated region are missing.
video, corrFlowF, corrFlowB, flow_mask, mask_dilated, start_point, end_point = extrapolation(
args, video, corrFlowF, corrFlowB)
imgH, imgW = video.shape[:2]
# mask indicating the missing region in the video.
mask = np.tile(flow_mask[..., None], (1, 1, nFrame))
flow_mask = np.tile(flow_mask[..., None], (1, 1, nFrame))
else:
# Loads masks.
filename_list = glob.glob(os.path.join(args.path_mask, '*.png')) + \
glob.glob(os.path.join(args.path_mask, '*.jpg'))
mask = []
flow_mask = []
for filename in sorted(filename_list):
mask_img = np.array(Image.open(filename).convert('L'))
mask.append(mask_img)
# Dilate 15 pixel so that all known pixel is trustworthy
flow_mask_img = scipy.ndimage.binary_dilation(mask_img,
iterations=15)
# Close the small holes inside the foreground objects
flow_mask_img = cv2.morphologyEx(flow_mask_img.astype(np.uint8),
cv2.MORPH_CLOSE,
np.ones((21, 21),
np.uint8)).astype(np.bool)
flow_mask_img = scipy.ndimage.binary_fill_holes(
flow_mask_img).astype(np.bool)
flow_mask.append(flow_mask_img)
# mask indicating the missing region in the video.
mask = np.stack(mask, -1).astype(np.bool)
flow_mask = np.stack(flow_mask, -1).astype(np.bool)
if args.edge_guide:
print('Pytorch 1.6.0 for edge completion')
try:
from edgeconnect.networks import EdgeGenerator_
EdgeGenerator = EdgeGenerator_()
EdgeComp_ckpt = torch.load(args.edge_completion_model)
EdgeGenerator.load_state_dict(EdgeComp_ckpt['generator'])
EdgeGenerator.to(torch.device('cuda:0'))
EdgeGenerator.eval()
except:
EdgeGenerator = None
# Edge completion.
FlowF_edge = edge_completion(args, EdgeGenerator, corrFlowF, flow_mask,
'forward')
FlowB_edge = edge_completion(args, EdgeGenerator, corrFlowB, flow_mask,
'backward')
else:
FlowF_edge, FlowB_edge = None, None
# Completes the flow.
videoFlowF = complete_flow(args, corrFlowF, flow_mask, 'forward',
FlowF_edge)
videoFlowB = complete_flow(args, corrFlowB, flow_mask, 'backward',
FlowB_edge)
iter = 0
mask_tofill = mask
video_comp = video
try:
from spatial_inpaint import spatial_inpaint
from frame_inpaint import DeepFillv1
deepfill = DeepFillv1(
pretrained_model='/home/chengao/Weight/imagenet_deepfill.pth',
image_shape=[imgH, imgW])
except:
print('Please switch to Pytorch 0.4.0')
return
while ((np.sum(mask_tofill) > 0) and iter <= 5):
create_dir(os.path.join(args.outroot, 'frame_comp_' + str(iter)))
# Color propagation.
video_comp, mask_tofill, _ = get_flowNN(args, video_comp, mask_tofill,
videoFlowF, videoFlowB, None,
None)
for i in range(nFrame):
mask_tofill[:, :,
i] = scipy.ndimage.binary_dilation(mask_tofill[:, :,
i],
iterations=2)
img = video_comp[:, :, :, i] * 255
img[mask_tofill[:, :, i]] = [0, 255, 0]
cv2.imwrite(
os.path.join(args.outroot, 'frame_comp_' + str(iter),
'%05d.png' % i), img)
video_comp_ = (video_comp * 255).astype(np.uint8).transpose(
3, 0, 1, 2)[:, :, :, ::-1]
imageio.mimwrite(os.path.join(args.outroot, 'frame_comp_' + str(iter),
'video_extrapolation.mp4'),
video_comp_,
fps=12,
quality=8,
macro_block_size=1)
imageio.mimsave(os.path.join(args.outroot, 'frame_comp_' + str(iter),
'video_extrapolation.gif'),
video_comp_,
format='gif',
fps=12)
mask_tofill, video_comp = spatial_inpaint(deepfill, mask_tofill,
video_comp)
iter += 1
def video_completion(args):
# Flow model.
RAFT_model = initialize_RAFT(args)
# Loads frames.
filename_list = glob.glob(os.path.join(args.path, '*.png')) + \
glob.glob(os.path.join(args.path, '*.jpg'))
# Obtains imgH, imgW and nFrame.
imgH, imgW = np.array(Image.open(filename_list[0])).shape[:2]
nFrame = len(filename_list)
# Loads video.
video = []
for filename in sorted(filename_list):
video.append(
torch.from_numpy(
np.array(Image.open(filename)).astype(
np.uint8)[..., :3]).permute(2, 0, 1).float())
video = torch.stack(video, dim=0)
video = video.to('cuda')
# Calcutes the corrupted flow.
corrFlowF, corrFlowB, corrFlowNLF, corrFlowNLB = calculate_flow(
args, RAFT_model, video)
print('\nFinish flow prediction.')
# Makes sure video is in BGR (opencv) format.
video = video.permute(2, 3, 1, 0).cpu().numpy()[:, :, ::-1, :] / 255.
if args.mode == 'video_extrapolation':
# Creates video and flow where the extrapolated region are missing.
video, corrFlowF, corrFlowB, corrFlowNLF, corrFlowNLB, flow_mask, mask_dilated, start_point, end_point = extrapolation(
args, video, corrFlowF, corrFlowB, corrFlowNLF, corrFlowNLB)
imgH, imgW = video.shape[:2]
# mask indicating the missing region in the video.
mask = np.tile(flow_mask[..., None], (1, 1, nFrame))
flow_mask = np.tile(flow_mask[..., None], (1, 1, nFrame))
else:
# Loads masks.
filename_list = glob.glob(os.path.join(args.path_mask, '*.png')) + \
glob.glob(os.path.join(args.path_mask, '*.jpg'))
mask = []
flow_mask = []
for filename in sorted(filename_list):
mask_img = np.array(Image.open(filename).convert('L'))
mask.append(mask_img)
# Dilate 15 pixel so that all known pixel is trustworthy
flow_mask_img = scipy.ndimage.binary_dilation(mask_img,
iterations=15)
# Close the small holes inside the foreground objects
flow_mask_img = cv2.morphologyEx(flow_mask_img.astype(np.uint8),
cv2.MORPH_CLOSE,
np.ones((21, 21),
np.uint8)).astype(bool)
flow_mask_img = scipy.ndimage.binary_fill_holes(
flow_mask_img).astype(bool)
flow_mask.append(flow_mask_img)
# mask indicating the missing region in the video.
mask = np.stack(mask, -1).astype(bool)
flow_mask = np.stack(flow_mask, -1).astype(bool)
if args.edge_guide:
# Edge completion model.
EdgeGenerator = EdgeGenerator_()
EdgeComp_ckpt = torch.load(args.edge_completion_model)
EdgeGenerator.load_state_dict(EdgeComp_ckpt['generator'])
EdgeGenerator.to(torch.device('cuda:0'))
EdgeGenerator.eval()
# Edge completion.
FlowF_edge = edge_completion(args, EdgeGenerator, corrFlowF, flow_mask,
'forward')
FlowB_edge = edge_completion(args, EdgeGenerator, corrFlowB, flow_mask,
'backward')
print('\nFinish edge completion.')
else:
FlowF_edge, FlowB_edge = None, None
# Completes the flow.
videoFlowF = complete_flow(args, corrFlowF, flow_mask, 'forward',
FlowF_edge)
videoFlowB = complete_flow(args, corrFlowB, flow_mask, 'backward',
FlowB_edge)
if args.Nonlocal:
videoNonLocalFlowF = complete_flow(args, corrFlowNLF, flow_mask,
'nonlocal_forward', None)
videoNonLocalFlowB = complete_flow(args, corrFlowNLB, flow_mask,
'nonlocal_backward', None)
else:
videoNonLocalFlowF = None
videoNonLocalFlowB = None
print('\nFinish flow completion.')
iter = 0
mask_tofill = mask
video_comp = video
# Image inpainting model.
deepfill = DeepFillv1(pretrained_model=args.deepfill_model,
image_shape=[imgH, imgW])
# We iteratively complete the video.
while (np.sum(mask_tofill) > 0):
create_dir(os.path.join(args.outroot, 'frame_comp_' + str(iter)))
# Color propagation.
video_comp, mask_tofill, _ = get_flowNN(args, video_comp, mask_tofill,
videoFlowF, videoFlowB,
videoNonLocalFlowF,
videoNonLocalFlowB)
for i in range(nFrame):
mask_tofill[:, :,
i] = scipy.ndimage.binary_dilation(mask_tofill[:, :,
i],
iterations=2)
img = video_comp[:, :, :, i] * 255
# Green indicates the regions that are not filled yet.
img[mask_tofill[:, :, i]] = [0, 255, 0]
cv2.imwrite(
os.path.join(args.outroot, 'frame_comp_' + str(iter),
'%05d.png' % i), img)
# video_comp_ = (video_comp * 255).astype(np.uint8).transpose(3, 0, 1, 2)[:, :, :, ::-1]
# imageio.mimwrite(os.path.join(args.outroot, 'frame_comp_' + str(iter), 'intermediate_{0}.mp4'.format(str(iter))), video_comp_, fps=12, quality=8, macro_block_size=1)
# imageio.mimsave(os.path.join(args.outroot, 'frame_comp_' + str(iter), 'intermediate_{0}.gif'.format(str(iter))), video_comp_, format='gif', fps=12)
mask_tofill, video_comp = spatial_inpaint(deepfill, mask_tofill,
video_comp)
iter += 1
create_dir(os.path.join(args.outroot, 'frame_comp_' + 'final'))
video_comp_ = (video_comp * 255).astype(np.uint8).transpose(
3, 0, 1, 2)[:, :, :, ::-1]
for i in range(nFrame):
img = video_comp[:, :, :, i] * 255
cv2.imwrite(
os.path.join(args.outroot, 'frame_comp_' + 'final',
'%05d.png' % i), img)
imageio.mimwrite(os.path.join(args.outroot, 'frame_comp_' + 'final',
'final.mp4'),
video_comp_,
fps=12,
quality=8,
macro_block_size=1)