Spaces:
Sleeping
Sleeping
import sys | |
import os | |
import warnings | |
sys.path.append(os.path.abspath(os.path.join(__file__, "..", ".."))) | |
sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", "FGT"))) | |
sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", "LAFC"))) | |
warnings.filterwarnings("ignore") | |
import cvbase | |
from torchvision.transforms import ToTensor | |
from tool.get_flowNN_gradient import get_flowNN_gradient | |
from tool.utils.Poisson_blend_img import Poisson_blend_img | |
from tool.utils.region_fill import regionfill | |
from importlib import import_module | |
import yaml | |
from RAFT import RAFT | |
from RAFT import utils | |
import torch.nn.functional as F2 | |
import torchvision.transforms.functional as F | |
from skimage.feature import canny | |
import scipy.ndimage | |
from PIL import Image | |
import imageio | |
import torch | |
import numpy as np | |
import copy | |
import glob | |
import cv2 | |
import argparse | |
def to_tensor(img): | |
img = Image.fromarray(img) | |
img_t = F.to_tensor(img).float() | |
return img_t | |
def diffusion(flows, masks): | |
flows_filled = [] | |
for i in range(flows.shape[0]): | |
flow, mask = flows[i], masks[i] | |
flow_filled = np.zeros(flow.shape) | |
flow_filled[:, :, 0] = regionfill(flow[:, :, 0], mask[:, :, 0]) | |
flow_filled[:, :, 1] = regionfill(flow[:, :, 1], mask[:, :, 0]) | |
flows_filled.append(flow_filled) | |
return flows_filled | |
def np2tensor(array, near="c"): | |
if isinstance(array, list): | |
array = np.stack(array, axis=0) # [t, h, w, c] | |
if near == "c": | |
array = ( | |
torch.from_numpy(np.transpose(array, (3, 0, 1, 2))).unsqueeze(0).float() | |
) # [1, c, t, h, w] | |
elif near == "t": | |
array = torch.from_numpy(np.transpose(array, (0, 3, 1, 2))).unsqueeze(0).float() | |
else: | |
raise ValueError(f"Unknown near type: {near}") | |
return array | |
def tensor2np(array): | |
array = torch.stack(array, dim=-1).squeeze(0).permute(1, 2, 0, 3).cpu().numpy() | |
return array | |
def gradient_mask(mask): | |
gradient_mask = np.logical_or.reduce( | |
( | |
mask, | |
np.concatenate( | |
(mask[1:, :], np.zeros((1, mask.shape[1]), dtype=np.bool)), axis=0 | |
), | |
np.concatenate( | |
(mask[:, 1:], np.zeros((mask.shape[0], 1), dtype=np.bool)), axis=1 | |
), | |
) | |
) | |
return gradient_mask | |
def indicesGen(pivot, interval, frames, t): | |
singleSide = frames // 2 | |
results = [] | |
for i in range(-singleSide, singleSide + 1): | |
index = pivot + interval * i | |
if index < 0: | |
index = abs(index) | |
if index > t - 1: | |
index = 2 * (t - 1) - index | |
results.append(index) | |
return results | |
def get_ref_index(f, neighbor_ids, length, ref_length, num_ref): | |
ref_index = [] | |
if num_ref == -1: | |
for i in range(0, length, ref_length): | |
if i not in neighbor_ids: | |
ref_index.append(i) | |
else: | |
start_idx = max(0, f - ref_length * (num_ref // 2)) | |
end_idx = min(length, f + ref_length * (num_ref // 2)) | |
for i in range(start_idx, end_idx + 1, ref_length): | |
if i not in neighbor_ids: | |
if len(ref_index) > num_ref: | |
break | |
ref_index.append(i) | |
return ref_index | |
def save_flows(output, videoFlowF, videoFlowB): | |
create_dir(os.path.join(output, "completed_flow", "forward_flo")) | |
create_dir(os.path.join(output, "completed_flow", "backward_flo")) | |
create_dir(os.path.join(output, "completed_flow", "forward_png")) | |
create_dir(os.path.join(output, "completed_flow", "backward_png")) | |
N = videoFlowF.shape[-1] | |
for i in range(N): | |
forward_flow = videoFlowF[..., i] | |
backward_flow = videoFlowB[..., i] | |
forward_flow_vis = cvbase.flow2rgb(forward_flow) | |
backward_flow_vis = cvbase.flow2rgb(backward_flow) | |
cvbase.write_flow( | |
forward_flow, | |
os.path.join( | |
output, "completed_flow", "forward_flo", "{:05d}.flo".format(i) | |
), | |
) | |
cvbase.write_flow( | |
backward_flow, | |
os.path.join( | |
output, "completed_flow", "backward_flo", "{:05d}.flo".format(i) | |
), | |
) | |
imageio.imwrite( | |
os.path.join( | |
output, "completed_flow", "forward_png", "{:05d}.png".format(i) | |
), | |
forward_flow_vis, | |
) | |
imageio.imwrite( | |
os.path.join( | |
output, "completed_flow", "backward_png", "{:05d}.png".format(i) | |
), | |
backward_flow_vis, | |
) | |
def save_fgcp(output, frames, masks): | |
create_dir(os.path.join(output, "prop_frames")) | |
create_dir(os.path.join(output, "masks_left")) | |
create_dir(os.path.join(output, "prop_frames_npy")) | |
create_dir(os.path.join(output, "masks_left_npy")) | |
assert len(frames) == masks.shape[2] | |
for i in range(len(frames)): | |
cv2.imwrite( | |
os.path.join(output, "prop_frames", "%05d.png" % i), frames[i] * 255.0 | |
) | |
cv2.imwrite( | |
os.path.join(output, "masks_left", "%05d.png" % i), masks[:, :, i] * 255.0 | |
) | |
np.save( | |
os.path.join(output, "prop_frames_npy", "%05d.npy" % i), frames[i] * 255.0 | |
) | |
np.save( | |
os.path.join(output, "masks_left_npy", "%05d.npy" % i), | |
masks[:, :, i] * 255.0, | |
) | |
def create_dir(dir): | |
"""Creates a directory if not exist.""" | |
if not os.path.exists(dir): | |
os.makedirs(dir) | |
def initialize_RAFT(args, device): | |
"""Initializes the RAFT model.""" | |
model = torch.nn.DataParallel(RAFT(args)) | |
if not torch.cuda.is_available(): | |
model.load_state_dict(torch.load(args.raft_model, map_location=lambda storage, loc: storage)) | |
else: | |
model.load_state_dict(torch.load(args.raft_model)) | |
model = model.module | |
model.to(device) | |
model.eval() | |
return model | |
def initialize_LAFC(args, device): | |
print(args.lafc_ckpts) | |
assert len(os.listdir(args.lafc_ckpts)) == 2 | |
checkpoint, config_file = ( | |
glob.glob(os.path.join(args.lafc_ckpts, "*.tar"))[0], | |
glob.glob(os.path.join(args.lafc_ckpts, "*.yaml"))[0], | |
) | |
with open(config_file, "r") as f: | |
configs = yaml.full_load(f) | |
model = configs["model"] | |
pkg = import_module("LAFC.models.{}".format(model)) | |
model = pkg.Model(configs) | |
if not torch.cuda.is_available(): | |
state = torch.load( | |
checkpoint, map_location=lambda storage, loc: storage | |
) | |
else: | |
state = torch.load( | |
checkpoint, map_location=lambda storage, loc: storage.cuda(device) | |
) | |
model.load_state_dict(state["model_state_dict"]) | |
model = model.to(device) | |
return model, configs | |
def initialize_FGT(args, device): | |
assert len(os.listdir(args.fgt_ckpts)) == 2 | |
checkpoint, config_file = ( | |
glob.glob(os.path.join(args.fgt_ckpts, "*.tar"))[0], | |
glob.glob(os.path.join(args.fgt_ckpts, "*.yaml"))[0], | |
) | |
with open(config_file, "r") as f: | |
configs = yaml.full_load(f) | |
model = configs["model"] | |
net = import_module("FGT.models.{}".format(model)) | |
model = net.Model(configs).to(device) | |
if not torch.cuda.is_available(): | |
state = torch.load( | |
checkpoint, map_location=lambda storage, loc: storage | |
) | |
else: | |
state = torch.load( | |
checkpoint, map_location=lambda storage, loc: storage.cuda(device) | |
) | |
model.load_state_dict(state["model_state_dict"]) | |
return model, configs | |
def calculate_flow(args, model, video, mode): | |
"""Calculates optical flow.""" | |
if mode not in ["forward", "backward"]: | |
raise NotImplementedError | |
imgH, imgW = args.imgH, args.imgW | |
Flow = np.empty(((imgH, imgW, 2, 0)), dtype=np.float32) | |
if args.vis_flows: | |
create_dir(os.path.join(args.outroot, "flow", mode + "_flo")) | |
create_dir(os.path.join(args.outroot, "flow", mode + "_png")) | |
with torch.no_grad(): | |
for i in range(video.shape[0] - 1): | |
print( | |
"Calculating {0} flow {1:2d} <---> {2:2d}".format(mode, i, i + 1), | |
"\r", | |
end="", | |
) | |
if mode == "forward": | |
# Flow i -> i + 1 | |
image1 = video[i, None] | |
image2 = video[i + 1, None] | |
elif mode == "backward": | |
# Flow i + 1 -> i | |
image1 = video[i + 1, None] | |
image2 = video[i, None] | |
else: | |
raise NotImplementedError | |
_, flow = model(image1, image2, iters=20, test_mode=True) | |
flow = flow[0].permute(1, 2, 0).cpu().numpy() | |
# resize optical flows | |
h, w = flow.shape[:2] | |
if h != imgH or w != imgW: | |
flow = cv2.resize(flow, (imgW, imgH), cv2.INTER_LINEAR) | |
flow[:, :, 0] *= float(imgW) / float(w) | |
flow[:, :, 1] *= float(imgH) / float(h) | |
Flow = np.concatenate((Flow, flow[..., None]), axis=-1) | |
if args.vis_flows: | |
# Flow visualization. | |
flow_img = utils.flow_viz.flow_to_image(flow) | |
flow_img = Image.fromarray(flow_img) | |
# Saves the flow and flow_img. | |
flow_img.save( | |
os.path.join(args.outroot, "flow", mode + "_png", "%05d.png" % i) | |
) | |
utils.frame_utils.writeFlow( | |
os.path.join(args.outroot, "flow", mode + "_flo", "%05d.flo" % i), | |
flow, | |
) | |
return Flow | |
def extrapolation(args, video_ori, corrFlowF_ori, corrFlowB_ori): | |
"""Prepares the data for video extrapolation.""" | |
imgH, imgW, _, nFrame = video_ori.shape | |
# Defines new FOV. | |
imgH_extr = int(args.H_scale * imgH) | |
imgW_extr = int(args.W_scale * imgW) | |
imgH_extr = imgH_extr - imgH_extr % 4 | |
imgW_extr = imgW_extr - imgW_extr % 4 | |
H_start = int((imgH_extr - imgH) / 2) | |
W_start = int((imgW_extr - imgW) / 2) | |
# Generates the mask for missing region. | |
flow_mask = np.ones(((imgH_extr, imgW_extr)), dtype=np.bool) | |
flow_mask[H_start : H_start + imgH, W_start : W_start + imgW] = 0 | |
mask_dilated = gradient_mask(flow_mask) | |
# Extrapolates the FOV for video. | |
video = np.zeros(((imgH_extr, imgW_extr, 3, nFrame)), dtype=np.float32) | |
video[H_start : H_start + imgH, W_start : W_start + imgW, :, :] = video_ori | |
for i in range(nFrame): | |
print("Preparing frame {0}".format(i), "\r", end="") | |
video[:, :, :, i] = ( | |
cv2.inpaint( | |
(video[:, :, :, i] * 255).astype(np.uint8), | |
flow_mask.astype(np.uint8), | |
3, | |
cv2.INPAINT_TELEA, | |
).astype(np.float32) | |
/ 255.0 | |
) | |
# Extrapolates the FOV for flow. | |
corrFlowF = np.zeros(((imgH_extr, imgW_extr, 2, nFrame - 1)), dtype=np.float32) | |
corrFlowB = np.zeros(((imgH_extr, imgW_extr, 2, nFrame - 1)), dtype=np.float32) | |
corrFlowF[H_start : H_start + imgH, W_start : W_start + imgW, :] = corrFlowF_ori | |
corrFlowB[H_start : H_start + imgH, W_start : W_start + imgW, :] = corrFlowB_ori | |
return ( | |
video, | |
corrFlowF, | |
corrFlowB, | |
flow_mask, | |
mask_dilated, | |
(W_start, H_start), | |
(W_start + imgW, H_start + imgH), | |
) | |
def complete_flow(config, flow_model, flows, flow_masks, mode, device): | |
if mode not in ["forward", "backward"]: | |
raise NotImplementedError(f"Error flow mode {mode}") | |
flow_masks = np.moveaxis(flow_masks, -1, 0) # [N, H, W] | |
flows = np.moveaxis(flows, -1, 0) # [N, H, W, 2] | |
if len(flow_masks.shape) == 3: | |
flow_masks = flow_masks[:, :, :, np.newaxis] | |
if mode == "forward": | |
flow_masks = flow_masks[0:-1] | |
else: | |
flow_masks = flow_masks[1:] | |
num_flows, flow_interval = config["num_flows"], config["flow_interval"] | |
diffused_flows = diffusion(flows, flow_masks) | |
flows = np2tensor(flows) | |
flow_masks = np2tensor(flow_masks) | |
diffused_flows = np2tensor(diffused_flows) | |
flows = flows.to(device) | |
flow_masks = flow_masks.to(device) | |
diffused_flows = diffused_flows.to(device) | |
t = diffused_flows.shape[2] | |
filled_flows = [None] * t | |
pivot = num_flows // 2 | |
for i in range(t): | |
indices = indicesGen(i, flow_interval, num_flows, t) | |
print("Indices: ", indices, "\r", end="") | |
cand_flows = flows[:, :, indices] | |
cand_masks = flow_masks[:, :, indices] | |
inputs = diffused_flows[:, :, indices] | |
pivot_mask = cand_masks[:, :, pivot] | |
pivot_flow = cand_flows[:, :, pivot] | |
with torch.no_grad(): | |
output_flow = flow_model(inputs, cand_masks) | |
if isinstance(output_flow, tuple) or isinstance(output_flow, list): | |
output_flow = output_flow[0] | |
comp = output_flow * pivot_mask + pivot_flow * (1 - pivot_mask) | |
if filled_flows[i] is None: | |
filled_flows[i] = comp | |
assert None not in filled_flows | |
return filled_flows | |
def read_flow(flow_dir, video): | |
nFrame, _, imgH, imgW = video.shape | |
Flow = np.empty(((imgH, imgW, 2, 0)), dtype=np.float32) | |
flows = sorted(glob.glob(os.path.join(flow_dir, "*.flo"))) | |
for flow in flows: | |
flow_data = cvbase.read_flow(flow) | |
h, w = flow_data.shape[:2] | |
flow_data = cv2.resize(flow_data, (imgW, imgH), cv2.INTER_LINEAR) | |
flow_data[:, :, 0] *= float(imgW) / float(w) | |
flow_data[:, :, 1] *= float(imgH) / float(h) | |
Flow = np.concatenate((Flow, flow_data[..., None]), axis=-1) | |
return Flow | |
def norm_flows(flows): | |
assert len(flows.shape) == 5, "FLow shape: {}".format(flows.shape) | |
flattened_flows = flows.flatten(3) | |
flow_max = torch.max(flattened_flows, dim=-1, keepdim=True)[0] | |
flows = flows / flow_max.unsqueeze(-1) | |
return flows | |
def save_results(outdir, comp_frames): | |
out_dir = os.path.join(outdir, "frames") | |
if not os.path.exists(out_dir): | |
os.makedirs(out_dir) | |
for i in range(len(comp_frames)): | |
out_path = os.path.join(out_dir, "{:05d}.png".format(i)) | |
cv2.imwrite(out_path, comp_frames[i][:, :, ::-1]) | |
def video_inpainting(args, imgArr, imgMaskArr): | |
#device = torch.device("cuda:{}".format(args.gpu)) | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
print(args) | |
if args.opt is not None: | |
with open(args.opt, "r") as f: | |
opts = yaml.full_load(f) | |
for k in opts.keys(): | |
if k in args: | |
setattr(args, k, opts[k]) | |
print(args) | |
# Flow model. | |
RAFT_model = initialize_RAFT(args, device) | |
# LAFC (flow completion) | |
LAFC_model, LAFC_config = initialize_LAFC(args, device) | |
# FGT | |
FGT_model, FGT_config = initialize_FGT(args, device) | |
# 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 = args.imgH, args.imgW | |
# nFrame = len(filename_list) | |
nFrame = len(imgArr) | |
if imgH < 350: | |
flowH, flowW = imgH * 2, imgW * 2 | |
else: | |
flowH, flowW = imgH, imgW | |
# Load video. | |
video, video_flow = [], [] | |
if args.mode == "watermark_removal": | |
maskname_list = glob.glob(os.path.join(args.path_mask, "*.png")) + glob.glob( | |
os.path.join(args.path_mask, "*.jpg") | |
) | |
assert len(filename_list) == len(maskname_list) | |
for filename, maskname in zip(sorted(filename_list), sorted(maskname_list)): | |
frame = ( | |
torch.from_numpy(np.array(Image.open(filename)).astype(np.uint8)) | |
.permute(2, 0, 1) | |
.float() | |
.unsqueeze(0) | |
) | |
mask = ( | |
torch.from_numpy(np.array(Image.open(maskname)).astype(np.uint8)) | |
.permute(2, 0, 1) | |
.float() | |
.unsqueeze(0) | |
) | |
mask[mask > 0] = 1 | |
frame = frame * (1 - mask) | |
frame = F2.upsample( | |
frame, size=(imgH, imgW), mode="bilinear", align_corners=False | |
) | |
frame_flow = F2.upsample( | |
frame, size=(flowH, flowW), mode="bilinear", align_corners=False | |
) | |
video.append(frame) | |
video_flow.append(frame_flow) | |
else: | |
"""for filename in sorted(filename_list): | |
frame = torch.from_numpy(np.array(Image.open(filename)).astype(np.uint8)).permute(2, 0, 1).float().unsqueeze(0) | |
frame = F2.upsample(frame, size=(imgH, imgW), mode='bilinear', align_corners=False) | |
frame_flow = F2.upsample(frame, size=(flowH, flowW), mode='bilinear', align_corners=False) | |
video.append(frame) | |
video_flow.append(frame_flow)""" | |
for im in imgArr: | |
frame = ( | |
torch.from_numpy(np.array(im).astype(np.uint8)) | |
.permute(2, 0, 1) | |
.float() | |
.unsqueeze(0) | |
) | |
frame = F2.upsample( | |
frame, size=(imgH, imgW), mode="bilinear", align_corners=False | |
) | |
frame_flow = F2.upsample( | |
frame, size=(flowH, flowW), mode="bilinear", align_corners=False | |
) | |
video.append(frame) | |
video_flow.append(frame_flow) | |
video = torch.cat(video, dim=0) # [n, c, h, w] | |
video_flow = torch.cat(video_flow, dim=0) | |
gts = video.clone() | |
video = video.to(device) | |
video_flow = video_flow.to(device) | |
# Calcutes the corrupted flow. | |
forward_flows = calculate_flow( | |
args, RAFT_model, video_flow, "forward" | |
) # [B, C, 2, N] | |
backward_flows = calculate_flow(args, RAFT_model, video_flow, "backward") | |
# Makes sure video is in BGR (opencv) format. | |
video = ( | |
video.permute(2, 3, 1, 0).cpu().numpy()[:, :, ::-1, :] / 255.0 | |
) # np array -> [h, w, c, N] (0~1) | |
if args.mode == "video_extrapolation": | |
# Creates video and flow where the extrapolated region are missing. | |
( | |
video, | |
forward_flows, | |
backward_flows, | |
flow_mask, | |
mask_dilated, | |
start_point, | |
end_point, | |
) = extrapolation(args, video, forward_flows, backward_flows) | |
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)) | |
mask_dilated = np.tile(mask_dilated[..., 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 = [] | |
mask_dilated = [] | |
flow_mask = [] | |
"""for filename in sorted(filename_list): | |
mask_img = np.array(Image.open(filename).convert('L')) | |
mask_img = cv2.resize(mask_img, dsize=(imgW, imgH), interpolation=cv2.INTER_NEAREST) | |
if args.flow_mask_dilates > 0: | |
flow_mask_img = scipy.ndimage.binary_dilation(mask_img, iterations=args.flow_mask_dilates) | |
else: | |
flow_mask_img = mask_img | |
flow_mask.append(flow_mask_img) | |
if args.frame_dilates > 0: | |
mask_img = scipy.ndimage.binary_dilation(mask_img, iterations=args.frame_dilates) | |
mask.append(mask_img) | |
mask_dilated.append(gradient_mask(mask_img))""" | |
for f_mask in imgMaskArr: | |
mask_img = np.array(f_mask) | |
mask_img = cv2.resize( | |
mask_img, dsize=(imgW, imgH), interpolation=cv2.INTER_NEAREST | |
) | |
if args.flow_mask_dilates > 0: | |
flow_mask_img = scipy.ndimage.binary_dilation( | |
mask_img, iterations=args.flow_mask_dilates | |
) | |
else: | |
flow_mask_img = mask_img | |
flow_mask.append(flow_mask_img) | |
if args.frame_dilates > 0: | |
mask_img = scipy.ndimage.binary_dilation( | |
mask_img, iterations=args.frame_dilates | |
) | |
mask.append(mask_img) | |
mask_dilated.append(gradient_mask(mask_img)) | |
# mask indicating the missing region in the video. | |
mask = np.stack(mask, -1).astype(np.bool) # [H, W, C, N] | |
mask_dilated = np.stack(mask_dilated, -1).astype(np.bool) | |
flow_mask = np.stack(flow_mask, -1).astype(np.bool) | |
# Completes the flow. | |
videoFlowF = complete_flow( | |
LAFC_config, LAFC_model, forward_flows, flow_mask, "forward", device | |
) | |
videoFlowB = complete_flow( | |
LAFC_config, LAFC_model, backward_flows, flow_mask, "backward", device | |
) | |
videoFlowF = tensor2np(videoFlowF) | |
videoFlowB = tensor2np(videoFlowB) | |
print("\nFinish flow completion.") | |
if args.vis_completed_flows: | |
save_flows(args.outroot, videoFlowF, videoFlowB) | |
# Prepare gradients | |
gradient_x = np.empty(((imgH, imgW, 3, 0)), dtype=np.float32) | |
gradient_y = np.empty(((imgH, imgW, 3, 0)), dtype=np.float32) | |
for indFrame in range(nFrame): | |
img = video[:, :, :, indFrame] | |
img[mask[:, :, indFrame], :] = 0 | |
img = ( | |
cv2.inpaint( | |
(img * 255).astype(np.uint8), | |
mask[:, :, indFrame].astype(np.uint8), | |
3, | |
cv2.INPAINT_TELEA, | |
).astype(np.float32) | |
/ 255.0 | |
) | |
gradient_x_ = np.concatenate( | |
(np.diff(img, axis=1), np.zeros((imgH, 1, 3), dtype=np.float32)), axis=1 | |
) | |
gradient_y_ = np.concatenate( | |
(np.diff(img, axis=0), np.zeros((1, imgW, 3), dtype=np.float32)), axis=0 | |
) | |
gradient_x = np.concatenate( | |
(gradient_x, gradient_x_.reshape(imgH, imgW, 3, 1)), axis=-1 | |
) | |
gradient_y = np.concatenate( | |
(gradient_y, gradient_y_.reshape(imgH, imgW, 3, 1)), axis=-1 | |
) | |
gradient_x[mask_dilated[:, :, indFrame], :, indFrame] = 0 | |
gradient_y[mask_dilated[:, :, indFrame], :, indFrame] = 0 | |
gradient_x_filled = gradient_x | |
gradient_y_filled = gradient_y | |
mask_gradient = mask_dilated | |
video_comp = video | |
# Gradient propagation. | |
gradient_x_filled, gradient_y_filled, mask_gradient = get_flowNN_gradient( | |
args, | |
gradient_x_filled, | |
gradient_y_filled, | |
mask, | |
mask_gradient, | |
videoFlowF, | |
videoFlowB, | |
None, | |
None, | |
) | |
# if there exist holes in mask, Poisson blending will fail. So I did this trick. I sacrifice some value. Another solution is to modify Poisson blending. | |
for indFrame in range(nFrame): | |
mask_gradient[:, :, indFrame] = scipy.ndimage.binary_fill_holes( | |
mask_gradient[:, :, indFrame] | |
).astype(np.bool) | |
# After one gradient propagation iteration | |
# gradient --> RGB | |
frameBlends = [] | |
for indFrame in range(nFrame): | |
print("Poisson blending frame {0:3d}".format(indFrame)) | |
if mask[:, :, indFrame].sum() > 0: | |
try: | |
frameBlend, UnfilledMask = Poisson_blend_img( | |
video_comp[:, :, :, indFrame], | |
gradient_x_filled[:, 0 : imgW - 1, :, indFrame], | |
gradient_y_filled[0 : imgH - 1, :, :, indFrame], | |
mask[:, :, indFrame], | |
mask_gradient[:, :, indFrame], | |
) | |
except: | |
frameBlend, UnfilledMask = ( | |
video_comp[:, :, :, indFrame], | |
mask[:, :, indFrame], | |
) | |
frameBlend = np.clip(frameBlend, 0, 1.0) | |
tmp = ( | |
cv2.inpaint( | |
(frameBlend * 255).astype(np.uint8), | |
UnfilledMask.astype(np.uint8), | |
3, | |
cv2.INPAINT_TELEA, | |
).astype(np.float32) | |
/ 255.0 | |
) | |
frameBlend[UnfilledMask, :] = tmp[UnfilledMask, :] | |
video_comp[:, :, :, indFrame] = frameBlend | |
mask[:, :, indFrame] = UnfilledMask | |
frameBlend_ = copy.deepcopy(frameBlend) | |
# Green indicates the regions that are not filled yet. | |
frameBlend_[mask[:, :, indFrame], :] = [0, 1.0, 0] | |
else: | |
frameBlend_ = video_comp[:, :, :, indFrame] | |
frameBlends.append(frameBlend_) | |
if args.vis_prop: | |
save_fgcp(args.outroot, frameBlends, mask) | |
video_length = len(frameBlends) | |
for i in range(len(frameBlends)): | |
frameBlends[i] = frameBlends[i][:, :, ::-1] | |
frames_first = np2tensor(frameBlends, near="t").to(device) | |
mask = np.moveaxis(mask, -1, 0) | |
mask = mask[:, :, :, np.newaxis] | |
masks = np2tensor(mask, near="t").to(device) | |
normed_frames = frames_first * 2 - 1 | |
comp_frames = [None] * video_length | |
ref_length = args.step | |
num_ref = args.num_ref | |
neighbor_stride = args.neighbor_stride | |
videoFlowF = np.moveaxis(videoFlowF, -1, 0) | |
videoFlowF = np.concatenate([videoFlowF, videoFlowF[-1:, ...]], axis=0) | |
flows = np2tensor(videoFlowF, near="t") | |
flows = norm_flows(flows).to(device) | |
for f in range(0, video_length, neighbor_stride): | |
neighbor_ids = [ | |
i | |
for i in range( | |
max(0, f - neighbor_stride), min(video_length, f + neighbor_stride + 1) | |
) | |
] | |
ref_ids = get_ref_index(f, neighbor_ids, video_length, ref_length, num_ref) | |
print(f, len(neighbor_ids), len(ref_ids)) | |
selected_frames = normed_frames[:, neighbor_ids + ref_ids] | |
selected_masks = masks[:, neighbor_ids + ref_ids] | |
masked_frames = selected_frames * (1 - selected_masks) | |
selected_flows = flows[:, neighbor_ids + ref_ids] | |
with torch.no_grad(): | |
filled_frames = FGT_model(masked_frames, selected_flows, selected_masks) | |
filled_frames = (filled_frames + 1) / 2 | |
filled_frames = filled_frames.cpu().permute(0, 2, 3, 1).numpy() * 255 | |
for i in range(len(neighbor_ids)): | |
idx = neighbor_ids[i] | |
valid_frame = frames_first[0, idx].cpu().permute(1, 2, 0).numpy() * 255.0 | |
valid_mask = masks[0, idx].cpu().permute(1, 2, 0).numpy() | |
comp = np.array(filled_frames[i]).astype(np.uint8) * valid_mask + np.array( | |
valid_frame | |
).astype(np.uint8) * (1 - valid_mask) | |
if comp_frames[idx] is None: | |
comp_frames[idx] = comp | |
else: | |
comp_frames[idx] = ( | |
comp_frames[idx].astype(np.float32) * 0.5 | |
+ comp.astype(np.float32) * 0.5 | |
) | |
if args.vis_frame: | |
save_results(args.outroot, comp_frames) | |
create_dir(args.outroot) | |
for i in range(len(comp_frames)): | |
comp_frames[i] = comp_frames[i].astype(np.uint8) | |
imageio.mimwrite( | |
os.path.join(args.outroot, args.outfilename), comp_frames, fps=args.out_fps, quality=8 | |
) | |
print(f"Done, please check your result in {args.outroot} ") | |
def main(args): | |
assert args.mode in ( | |
"object_removal", | |
"video_extrapolation", | |
"watermark_removal", | |
), ( | |
"Accepted modes: 'object_removal', 'video_extrapolation', and 'watermark_removal', but input is %s" | |
) % args.mode | |
video_inpainting(args) | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
"--opt", | |
default="configs/object_removal.yaml", | |
help="Please select your config file for inference", | |
) | |
# video completion | |
parser.add_argument( | |
"--mode", | |
default="object_removal", | |
choices=["object_removal", "watermark_removal", "video_extrapolation"], | |
help="modes: object_removal / video_extrapolation", | |
) | |
parser.add_argument( | |
"--path", default="/myData/davis_resized/walking", help="dataset for evaluation" | |
) | |
parser.add_argument( | |
"--path_mask", | |
default="/myData/dilateAnnotations_4/walking", | |
help="mask for object removal", | |
) | |
parser.add_argument( | |
"--outroot", default="quick_start/walking3", help="output directory" | |
) | |
parser.add_argument( | |
"--consistencyThres", | |
dest="consistencyThres", | |
default=5, | |
type=float, | |
help="flow consistency error threshold", | |
) | |
parser.add_argument("--alpha", dest="alpha", default=0.1, type=float) | |
parser.add_argument("--Nonlocal", dest="Nonlocal", default=False, type=bool) | |
# RAFT | |
parser.add_argument( | |
"--raft_model", | |
default="../LAFC/flowCheckPoint/raft-things.pth", | |
help="restore checkpoint", | |
) | |
parser.add_argument("--small", action="store_true", help="use small model") | |
parser.add_argument( | |
"--mixed_precision", action="store_true", help="use mixed precision" | |
) | |
parser.add_argument( | |
"--alternate_corr", | |
action="store_true", | |
help="use efficent correlation implementation", | |
) | |
# LAFC | |
parser.add_argument("--lafc_ckpts", type=str, default="../LAFC/checkpoint") | |
# FGT | |
parser.add_argument("--fgt_ckpts", type=str, default="../FGT/checkpoint") | |
# extrapolation | |
parser.add_argument( | |
"--H_scale", dest="H_scale", default=2, type=float, help="H extrapolation scale" | |
) | |
parser.add_argument( | |
"--W_scale", dest="W_scale", default=2, type=float, help="W extrapolation scale" | |
) | |
# Image basic information | |
parser.add_argument("--imgH", type=int, default=256) | |
parser.add_argument("--imgW", type=int, default=432) | |
parser.add_argument("--flow_mask_dilates", type=int, default=8) | |
parser.add_argument("--frame_dilates", type=int, default=0) | |
parser.add_argument("--gpu", type=int, default=0) | |
# FGT inference parameters | |
parser.add_argument("--step", type=int, default=10) | |
parser.add_argument("--num_ref", type=int, default=-1) | |
parser.add_argument("--neighbor_stride", type=int, default=5) | |
# visualization | |
parser.add_argument( | |
"--vis_flows", action="store_true", help="Visualize the initialized flows" | |
) | |
parser.add_argument( | |
"--vis_completed_flows", | |
action="store_true", | |
help="Visualize the completed flows", | |
) | |
parser.add_argument( | |
"--vis_prop", | |
action="store_true", | |
help="Visualize the frames after stage-I filling (flow guided content propagation)", | |
) | |
parser.add_argument("--vis_frame", action="store_true", help="Visualize frames") | |
args = parser.parse_args() | |
main(args) | |