E2FGVI
E2FGVI copied to clipboard
Published
20 hours ago •
MCG-NKU
MCG-NKU
How can I use only one mask files?
I want to remove wartermark from the video, the watermark is at the fixed postion, so can I use only one mask file for all the frames? if so how
Thanks in advance
Hi, in case someone else want it too, you can try the following code By the way, it's based on https://github.com/gaomingqi/Track-Anything/blob/master/inpainter/base_inpainter.py `import os import glob from PIL import Image
import torch import yaml import cv2 import importlib import numpy as np from tqdm import tqdm import torchvision from util.tensor_util import resize_frames, resize_masks
class BaseInpainter: def init(self, E2FGVI_checkpoint, device) -> None: """ E2FGVI_checkpoint: checkpoint of inpainter (version hq, with multi-resolution support) """ net = importlib.import_module('model.e2fgvi_hq') self.model = net.InpaintGenerator().to(device) self.model.load_state_dict(torch.load(E2FGVI_checkpoint, map_location=device)) self.model.eval() self.device = device # load configurations with open("config/config.yaml", 'r') as stream: config = yaml.safe_load(stream) self.neighbor_stride = config['neighbor_stride'] self.num_ref = config['num_ref'] self.step = config['step'] # config for E2FGVI with splits self.num_subset_frames = config['num_subset_frames'] self.num_external_ref = config['num_external_ref']
# sample reference frames from the whole video
def get_ref_index(self, f, neighbor_ids, length):
ref_index = []
if self.num_ref == -1:
for i in range(0, length, self.step):
if i not in neighbor_ids:
ref_index.append(i)
else:
start_idx = max(0, f - self.step * (self.num_ref // 2))
end_idx = min(length, f + self.step * (self.num_ref // 2))
for i in range(start_idx, end_idx + 1, self.step):
if i not in neighbor_ids:
if len(ref_index) > self.num_ref:
break
ref_index.append(i)
return ref_index
def inpaint_efficient(self, frames, masks, num_tcb, num_tca, dilate_radius=15, ratio=1):
"""
Perform Inpainting for video subsets
frames: numpy array, T, H, W, 3
masks: numpy array, T, H, W
num_tcb: constant, number of temporal context before, frames
num_tca: constant, number of temporal context after, frames
dilate_radius: radius when applying dilation on masks
ratio: down-sample ratio
Output:
inpainted_frames: numpy array, T, H, W, 3
"""
#assert frames.shape[:3] == masks.shape, 'different size between frames and masks'
#assert ratio > 0 and ratio <= 1, 'ratio must in (0, 1]'
# --------------------
# pre-processing
# --------------------
masks = masks.copy()
masks = np.clip(masks, 0, 1)
kernel = cv2.getStructuringElement(2, (dilate_radius, dilate_radius))
masks = np.stack([cv2.dilate(mask, kernel) for mask in masks], 0)
T = frames.shape[0]
masks = np.expand_dims(masks, axis=3) # expand to T, H, W, 1
# size: (w, h)
if ratio == 1:
size = None
binary_masks = masks
else:
print("Doesn't support ratio !!!!")
return null
# frames and binary_masks are numpy arrays
h, w = frames.shape[1:3]
video_length = T - (num_tca + num_tcb) # real video length
# convert to tensor
imgs = (torch.from_numpy(frames).permute(0, 3, 1, 2).contiguous().unsqueeze(0).float().div(255)) * 2 - 1
masks = torch.from_numpy(binary_masks).permute(0, 3, 1, 2).contiguous().unsqueeze(0)
imgs, masks = imgs.to(self.device), masks.to(self.device)
comp_frames = [None] * video_length
tcb_imgs = None
tca_imgs = None
tcb_masks = None
tca_masks = None
# --------------------
# end of pre-processing
# --------------------
# separate tc frames/masks from imgs and masks
if num_tcb > 0:
tcb_imgs = imgs[:, :num_tcb]
tcb_masks = masks
if num_tca > 0:
tca_imgs = imgs[:, -num_tca:]
tca_masks = masks
end_idx = -num_tca
else:
end_idx = T
imgs = imgs[:, num_tcb:end_idx]
#masks = masks[:, num_tcb:end_idx]
#binary_masks = binary_masks[num_tcb:end_idx] # only neighbor area are involved
frames = frames[num_tcb:end_idx] # only neighbor area are involved
for f in tqdm(range(0, video_length, self.neighbor_stride), desc='Inpainting image'):
neighbor_ids = [
i for i in range(max(0, f - self.neighbor_stride),
min(video_length, f + self.neighbor_stride + 1))
]
ref_ids = self.get_ref_index(f, neighbor_ids, video_length)
# selected_imgs = imgs[:1, neighbor_ids + ref_ids, :, :, :]
# selected_masks = masks[:1, neighbor_ids + ref_ids, :, :, :]
selected_imgs = imgs[:, neighbor_ids]
selected_masks = masks#[:, neighbor_ids]
# pad before
if tcb_imgs is not None:
selected_imgs = torch.concat([selected_imgs, tcb_imgs], dim=1)
selected_masks = torch.concat([selected_masks, tcb_masks], dim=1)
# integrate ref frames
selected_imgs = torch.concat([selected_imgs, imgs[:, ref_ids]], dim=1)
#selected_masks = torch.concat([selected_masks, masks[:, ref_ids]], dim=1)
# pad after
if tca_imgs is not None:
selected_imgs = torch.concat([selected_imgs, tca_imgs], dim=1)
#selected_masks = torch.concat([selected_masks, tca_masks], dim=1)
with torch.no_grad():
masked_imgs = selected_imgs * (1 - selected_masks)
mod_size_h = 60
mod_size_w = 108
h_pad = (mod_size_h - h % mod_size_h) % mod_size_h
w_pad = (mod_size_w - w % mod_size_w) % mod_size_w
masked_imgs = torch.cat(
[masked_imgs, torch.flip(masked_imgs, [3])],
3)[:, :, :, :h + h_pad, :]
masked_imgs = torch.cat(
[masked_imgs, torch.flip(masked_imgs, [4])],
4)[:, :, :, :, :w + w_pad]
pred_imgs, _ = self.model(masked_imgs, len(neighbor_ids))
pred_imgs = pred_imgs[:, :, :h, :w]
pred_imgs = (pred_imgs + 1) / 2
pred_imgs = pred_imgs.cpu().permute(0, 2, 3, 1).numpy() * 255
for i in range(len(neighbor_ids)):
idx = neighbor_ids[i]
img = pred_imgs[i].astype(np.uint8) * binary_masks[0] + frames[idx] * (
1 - binary_masks[0])
if comp_frames[idx] is None:
comp_frames[idx] = img
else:
comp_frames[idx] = comp_frames[idx].astype(
np.float32) * 0.5 + img.astype(np.float32) * 0.5
torch.cuda.empty_cache()
inpainted_frames = np.stack(comp_frames, 0)
return inpainted_frames.astype(np.uint8)
def inpaint(self, frames, masks, dilate_radius=15, ratio=1):
"""
Perform Inpainting for video subsets
frames: numpy array, T, H, W, 3
masks: numpy array, T, H, W
dilate_radius: radius when applying dilation on masks
ratio: down-sample ratio
Output:
inpainted_frames: numpy array, T, H, W, 3
"""
assert masks.shape[0] == 1, 'different size between frames and masks'
assert ratio > 0 and ratio <= 1, 'ratio must in (0, 1]'
# set num_subset_frames
num_subset_frames = self.num_subset_frames
# split frames into subsets
video_length = len(frames)
num_splits = video_length // num_subset_frames
id_splits = [[i*num_subset_frames, (i+1)*num_subset_frames] for i in range(num_splits)] # id splits
if num_splits == 0:
id_splits = [[0, video_length]]
# if remaining split > num_subset_frames/2, add a new split, else, append to the last split
if video_length - id_splits[-1][-1] > num_subset_frames / 3:
id_splits.append([num_splits*num_subset_frames, video_length])
else:
diff = video_length - id_splits[-1][-1]
id_splits = [[ids[0]+diff, ids[1]+diff] for ids in id_splits]
id_splits[0][0] = 0 # if OOM, let it happen at the begining :D
# if appending, convert the appended split to the FIRST one, avoiding OOM at last
# perform inpainting for each split
inpainted_splits = []
for id_split in id_splits:
video_split = frames[id_split[0]:id_split[1]]
#mask_split = masks[id_split[0]:id_split[1]]
# | id_before | ----- | id_split[0] | ----- | id_split[1] | ----- | id_after |
# for each split, consider its temporal context [-context_range] frames and [context_range] frames
id_before = max(0, id_split[0] - self.step * self.num_external_ref)
try:
tcb_frames = np.stack([frames[idb] for idb in range(id_before, (id_split[0]-self.step) + 1, self.step)], 0)
tcb_masks = np.stack([masks[idb] for idb in range(id_before, (id_split[0]-self.step) + 1, self.step)], 0)
num_tcb = len(tcb_frames)
except:
num_tcb = 0
id_after = min(video_length, id_split[1] + self.step * self.num_external_ref + 1)
try:
tca_frames = np.stack([frames[ida] for ida in range(id_split[1]+self.step, id_after, self.step)], 0)
tca_masks = np.stack([masks[ida] for ida in range(id_split[1]+self.step, id_after, self.step)], 0)
num_tca = len(tca_frames)
except:
num_tca = 0
# concatenate temporal context frames/masks with input frames/masks (for parallel pre-processing)
if num_tcb > 0:
video_split = np.concatenate([tcb_frames, video_split], 0)
#mask_split = np.concatenate([tcb_masks, mask_split], 0)
if num_tca > 0:
video_split = np.concatenate([video_split, tca_frames], 0)
#mask_split = np.concatenate([mask_split, tca_masks], 0)
torch.cuda.empty_cache()
# inpaint each split
inpainted_splits.append(self.inpaint_efficient(video_split, masks, num_tcb, num_tca, dilate_radius, ratio))
torch.cuda.empty_cache()
inpainted_frames = np.concatenate(inpainted_splits, 0)
return inpainted_frames.astype(np.uint8)
def inpaint_ori(self, frames, masks, dilate_radius=15, ratio=1):
"""
frames: numpy array, T, H, W, 3
masks: numpy array, T, H, W
dilate_radius: radius when applying dilation on masks
ratio: down-sample ratio
Output:
inpainted_frames: numpy array, T, H, W, 3
"""
assert frames.shape[:3] == masks.shape, 'different size between frames and masks'
assert ratio > 0 and ratio <= 1, 'ratio must in (0, 1]'
masks = masks.copy()
masks = np.clip(masks, 0, 1)
kernel = cv2.getStructuringElement(2, (dilate_radius, dilate_radius))
masks = np.stack([cv2.dilate(mask, kernel) for mask in masks], 0)
T, H, W = masks.shape
masks = np.expand_dims(masks, axis=3) # expand to T, H, W, 1
# size: (w, h)
if ratio == 1:
size = None
binary_masks = masks
else:
size = [int(W*ratio), int(H*ratio)]
size = [si+1 if si%2>0 else si for si in size] # only consider even values
# shortest side should be larger than 50
if min(size) < 50:
ratio = 50. / min(H, W)
size = [int(W*ratio), int(H*ratio)]
size = [160, 120]
binary_masks = resize_masks(masks, tuple(size))
frames = resize_frames(frames, tuple(size)) # T, H, W, 3
# frames and binary_masks are numpy arrays
h, w = frames.shape[1:3]
video_length = T
# convert to tensor
imgs = (torch.from_numpy(frames).permute(0, 3, 1, 2).contiguous().unsqueeze(0).float().div(255)) * 2 - 1
masks = torch.from_numpy(binary_masks).permute(0, 3, 1, 2).contiguous().unsqueeze(0)
imgs, masks = imgs.to(self.device), masks.to(self.device)
comp_frames = [None] * video_length
for f in tqdm(range(0, video_length, self.neighbor_stride), desc='Inpainting image'):
neighbor_ids = [
i for i in range(max(0, f - self.neighbor_stride),
min(video_length, f + self.neighbor_stride + 1))
]
ref_ids = self.get_ref_index(f, neighbor_ids, video_length)
selected_imgs = imgs[:1, neighbor_ids + ref_ids, :, :, :]
selected_masks = masks[:1, neighbor_ids + ref_ids, :, :, :]
with torch.no_grad():
masked_imgs = selected_imgs * (1 - selected_masks)
mod_size_h = 60
mod_size_w = 108
h_pad = (mod_size_h - h % mod_size_h) % mod_size_h
w_pad = (mod_size_w - w % mod_size_w) % mod_size_w
masked_imgs = torch.cat(
[masked_imgs, torch.flip(masked_imgs, [3])],
3)[:, :, :, :h + h_pad, :]
masked_imgs = torch.cat(
[masked_imgs, torch.flip(masked_imgs, [4])],
4)[:, :, :, :, :w + w_pad]
pred_imgs, _ = self.model(masked_imgs, len(neighbor_ids))
pred_imgs = pred_imgs[:, :, :h, :w]
pred_imgs = (pred_imgs + 1) / 2
pred_imgs = pred_imgs.cpu().permute(0, 2, 3, 1).numpy() * 255
for i in range(len(neighbor_ids)):
idx = neighbor_ids[i]
img = pred_imgs[i].astype(np.uint8) * binary_masks[idx] + frames[idx] * (
1 - binary_masks[idx])
if comp_frames[idx] is None:
comp_frames[idx] = img
else:
comp_frames[idx] = comp_frames[idx].astype(
np.float32) * 0.5 + img.astype(np.float32) * 0.5
torch.cuda.empty_cache()
inpainted_frames = np.stack(comp_frames, 0)
return inpainted_frames.astype(np.uint8)
def read_mask(mpath): masks = [] mnames = os.listdir(mpath) mnames.sort() for mp in mnames: m = Image.open(os.path.join(mpath, mp)) #m = m.resize(size, Image.NEAREST) m = np.array(m.convert('L')) m = np.array(m > 0).astype(np.uint8) m = cv2.dilate(m, cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3)), iterations=4)
masks.append(Image.fromarray(m * 255))
masks = np.stack(masks)
return masks
read frames from video
def read_frame_from_videos(vname, use_mp4 = False):
frames = []
if use_mp4:
vidcap = cv2.VideoCapture(vname)
success, image = vidcap.read()
count = 0
while success:
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
frames.append(image)
success, image = vidcap.read()
count += 1
else:
lst = os.listdir(vname)
lst.sort()
fr_lst = [vname + '/' + name for name in lst]
for fr in fr_lst:
image = cv2.imread(fr)
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
frames.append(image)
return frames
generate video after vos inference
def generate_video_from_frames(frames, output_path, fps=30): """ Generates a video from a list of frames.
Args:
frames (list of numpy arrays): The frames to include in the video.
output_path (str): The path to save the generated video.
fps (int, optional): The frame rate of the output video. Defaults to 30.
"""
# height, width, layers = frames[0].shape
# fourcc = cv2.VideoWriter_fourcc(*"mp4v")
# video = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
# print(output_path)
# for frame in frames:
# video.write(frame)
# video.release()
frames = torch.from_numpy(np.asarray(frames))
if not os.path.exists(os.path.dirname(output_path)):
os.makedirs(os.path.dirname(output_path))
torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
return output_path
if name == 'main':
# # davis-2017
# frame_path = glob.glob(os.path.join('/ssd1/gaomingqi/datasets/davis/JPEGImages/480p/parkour', '*.jpg'))
# frame_path.sort()
# mask_path = glob.glob(os.path.join('/ssd1/gaomingqi/datasets/davis/Annotations/480p/parkour', "*.png"))
# mask_path.sort()
# long and large video
frames = read_frame_from_videos("examples/schoolgirls.mp4",True)
video_length = len(frames)
frames = [np.array(f).astype(np.uint8) for f in frames]
masks = read_mask("examples/schoolgirls_mask")
frames = np.array(frames) # Convert list to numpy array
masks = np.array(masks) # Convert list to numpy array
# ----------------------------------------------
# how to use
# ----------------------------------------------
# 1/3: set checkpoint and device
checkpoint = 'release_model\\E2FGVI-HQ-CVPR22.pth'
device = 'cuda:0'
# 2/3: initialise inpainter
base_inpainter = BaseInpainter(checkpoint, device)
# 3/3: inpainting (frames: numpy array, T, H, W, 3; masks: numpy array, T, H, W)
# ratio: (0, 1], ratio for down sample, default value is 1
inpainted_frames = base_inpainter.inpaint(frames, masks) # numpy array, T, H, W, 3
generate_video_from_frames(inpainted_frames,"result/test.mp4")
# save
torch.cuda.empty_cache()
print('switch to ori')
# inpainted_frames = base_inpainter.inpaint_ori(frames[:50], masks[:50], ratio=0.1)
# save_path = '/ssd1/gaomingqi/results/inpainting/avengers'
# # ----------------------------------------------
# # end
# # ----------------------------------------------
# # save
# for ti, inpainted_frame in enumerate(inpainted_frames):
# frame = Image.fromarray(inpainted_frame).convert('RGB')
# frame.save(os.path.join(save_path, f'{ti:05d}.jpg'))
`