generator

    #                    same_image_flag: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
    #     if first_flag:
    #         mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
    #         return mask, logit, painted_image

    #     if interact_flag:
    #         mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
    #         return mask, logit, painted_image

    #     mask, logit, painted_image = self.xmem.track(image, logit)
    #     return mask, logit, painted_image

    def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
        mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
        return mask, logit, painted_image

    # def interact(self, image: np.ndarray, same_image_flag: bool, points:np.ndarray, labels: np.ndarray, logits: np.ndarray=None, multimask=True):
    #     mask, logit, painted_image = self.samcontroler.interact_loop(image, same_image_flag, points, labels, logits, multimask)
    #     return mask, logit, painted_image

    def generator(self, images: list, template_mask:np.ndarray):

        masks = []
        logits = []
        painted_images = []
        for i in tqdm(range(len(images)), desc="Tracking image"):
            if i ==0:
                mask, logit, painted_image = self.xmem.track(images[i], template_mask)
                masks.append(mask)
                logits.append(logit)
                painted_images.append(painted_image)

            else:
                mask, logit, painted_image = self.xmem.track(images[i])
                masks.append(mask)
                logits.append(logit)
                painted_images.append(painted_image)
        return masks, logits, painted_images

def parse_augment():
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=str, default="cuda:0")
    parser.add_argument('--sam_model_type', type=str, default="vit_h")
    parser.add_argument('--port', type=int, default=6080, help="only useful when running gradio applications")
    parser.add_argument('--debug', action="store_true")
    parser.add_argument('--mask_save', default=False)
    args = parser.parse_args()

    if args.debug:
        print(args)
    return args

if __name__ == "__main__":
    masks = None
    logits = None
    painted_images = None


import os
import glob
from PIL import Image

import torch
import yaml
import cv2
import importlib
import numpy as np
from tqdm import tqdm

from inpainter.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('inpainter.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("inpainter/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, 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)]
			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 - (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[:, :num_tcb]
		if num_tca > 0:
			tca_imgs = imgs[:, -num_tca:]
			tca_masks = masks[:, -num_tca:]
			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[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 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 frames.shape[:3] == masks.shape, '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, mask_split, 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)

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
	mask_path = glob.glob(os.path.join('/ssd1/gaomingqi/test-sample13', '*.npy'))
	mask_path.sort()
	frames = np.load('/ssd1/gaomingqi/revenger.npy')
	save_path = '/ssd1/gaomingqi/results/inpainting/avengers_split'

	if not os.path.exists(save_path):
		os.mkdir(save_path)

	masks = []
	for ti, mid in enumerate(mask_path):
		masks.append(np.load(mid, allow_pickle=True))
		if ti > 1122:
			break

	masks = np.stack(masks[:len(frames)], 0)

	# ----------------------------------------------
	# how to use
	# ----------------------------------------------
	# 1/3: set checkpoint and device
	checkpoint = '/ssd1/gaomingqi/checkpoints/E2FGVI-HQ-CVPR22.pth'
	device = 'cuda:4'
	# 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[:300], masks[:300], ratio=0.6)   # numpy array, T, H, W, 3

	# 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'))

	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'))


BaseInpainter

import os
import glob
from PIL import Image

import torch
import yaml
import cv2
import importlib
import numpy as np
from tqdm import tqdm

from inpainter.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('inpainter.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("inpainter/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, 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)]
			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 - (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[:, :num_tcb]
		if num_tca > 0:
			tca_imgs = imgs[:, -num_tca:]
			tca_masks = masks[:, -num_tca:]
			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[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 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 frames.shape[:3] == masks.shape, '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, mask_split, 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)

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
	mask_path = glob.glob(os.path.join('/ssd1/gaomingqi/test-sample13', '*.npy'))
	mask_path.sort()
	frames = np.load('/ssd1/gaomingqi/revenger.npy')
	save_path = '/ssd1/gaomingqi/results/inpainting/avengers_split'

	if not os.path.exists(save_path):
		os.mkdir(save_path)

	masks = []


import os
import glob
from PIL import Image

import torch
import yaml
import cv2
import importlib
import numpy as np
from tqdm import tqdm

from inpainter.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('inpainter.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("inpainter/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):
		"""


	def __init__(self, E2FGVI_checkpoint, device) -> None:
		"""
		E2FGVI_checkpoint: checkpoint of inpainter (version hq, with multi-resolution support)
		"""
		net = importlib.import_module('inpainter.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("inpainter/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
		# --------------------


		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, 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)]
			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 - (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[:, :num_tcb]
		if num_tca > 0:
			tca_imgs = imgs[:, -num_tca:]
			tca_masks = masks[:, -num_tca:]
			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[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 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 frames.shape[:3] == masks.shape, '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


inpaint

				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 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 frames.shape[:3] == masks.shape, '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, mask_split, 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


				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, mask_split, 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)

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
	mask_path = glob.glob(os.path.join('/ssd1/gaomingqi/test-sample13', '*.npy'))
	mask_path.sort()
	frames = np.load('/ssd1/gaomingqi/revenger.npy')
	save_path = '/ssd1/gaomingqi/results/inpainting/avengers_split'

	if not os.path.exists(save_path):
		os.mkdir(save_path)

	masks = []


import cv2
import numpy as np

# resize frames
def resize_frames(frames, size=None):
    """
    size: (w, h)
    """
    if size is not None:
        frames = [cv2.resize(f, size) for f in frames]
        frames = np.stack(frames, 0)

    return frames

# resize frames
def resize_masks(masks, size=None):
    """
    size: (w, h)
    """
    if size is not None:
        masks = [np.expand_dims(cv2.resize(m, size), 2) for m in masks]
        masks = np.stack(masks, 0)

    return masks


''' Towards An End-to-End Framework for Video Inpainting
'''

import torch
import torch.nn as nn
import torch.nn.functional as F

from inpainter.model.modules.flow_comp import SPyNet
from inpainter.model.modules.feat_prop import BidirectionalPropagation, SecondOrderDeformableAlignment
from inpainter.model.modules.tfocal_transformer_hq import TemporalFocalTransformerBlock, SoftSplit, SoftComp
from inpainter.model.modules.spectral_norm import spectral_norm as _spectral_norm

class BaseNetwork(nn.Module):
    def __init__(self):
        super(BaseNetwork, self).__init__()

    def print_network(self):
        if isinstance(self, list):
            self = self[0]
        num_params = 0
        for param in self.parameters():
            num_params += param.numel()
        print(
            'Network [%s] was created. Total number of parameters: %.1f million. '
            'To see the architecture, do print(network).' %
            (type(self).__name__, num_params / 1000000))

    def init_weights(self, init_type='normal', gain=0.02):
        '''
        initialize network's weights
        init_type: normal | xavier | kaiming | orthogonal
        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
        '''
        def init_func(m):
            classname = m.__class__.__name__
            if classname.find('InstanceNorm2d') != -1:
                if hasattr(m, 'weight') and m.weight is not None:
                    nn.init.constant_(m.weight.data, 1.0)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)
            elif hasattr(m, 'weight') and (classname.find('Conv') != -1
                                           or classname.find('Linear') != -1):
                if init_type == 'normal':
                    nn.init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    nn.init.xavier_normal_(m.weight.data, gain=gain)
                elif init_type == 'xavier_uniform':
                    nn.init.xavier_uniform_(m.weight.data, gain=1.0)
                elif init_type == 'kaiming':
                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    nn.init.orthogonal_(m.weight.data, gain=gain)
                elif init_type == 'none':  # uses pytorch's default init method
                    m.reset_parameters()
                else:
                    raise NotImplementedError(
                        'initialization method [%s] is not implemented' %
                        init_type)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)

        self.apply(init_func)

        # propagate to children
        for m in self.children():
            if hasattr(m, 'init_weights'):
                m.init_weights(init_type, gain)

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.group = [1, 2, 4, 8, 1]
        self.layers = nn.ModuleList([
            nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, groups=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(640, 512, kernel_size=3, stride=1, padding=1, groups=2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(768, 384, kernel_size=3, stride=1, padding=1, groups=4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(640, 256, kernel_size=3, stride=1, padding=1, groups=8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 128, kernel_size=3, stride=1, padding=1, groups=1),
            nn.LeakyReLU(0.2, inplace=True)
        ])

    def forward(self, x):
        bt, c, _, _ = x.size()
        # h, w = h//4, w//4
        out = x
        for i, layer in enumerate(self.layers):
            if i == 8:
                x0 = out
                _, _, h, w = x0.size()
            if i > 8 and i % 2 == 0:
                g = self.group[(i - 8) // 2]
                x = x0.view(bt, g, -1, h, w)
                o = out.view(bt, g, -1, h, w)
                out = torch.cat([x, o], 2).view(bt, -1, h, w)
            out = layer(out)
        return out

class deconv(nn.Module):
    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel_size=3,
                 padding=0):
        super().__init__()
        self.conv = nn.Conv2d(input_channel,
                              output_channel,
                              kernel_size=kernel_size,
                              stride=1,
                              padding=padding)

    def forward(self, x):
        x = F.interpolate(x,
                          scale_factor=2,
                          mode='bilinear',
                          align_corners=True)
        return self.conv(x)

class InpaintGenerator(BaseNetwork):
    def __init__(self, init_weights=True):
        super(InpaintGenerator, self).__init__()
        channel = 256
        hidden = 512

        # encoder
        self.encoder = Encoder()

        # decoder
        self.decoder = nn.Sequential(
            deconv(channel // 2, 128, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            deconv(64, 64, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1))

        # feature propagation module
        self.feat_prop_module = BidirectionalPropagation(channel // 2)

        # soft split and soft composition
        kernel_size = (7, 7)
        padding = (3, 3)
        stride = (3, 3)
        output_size = (60, 108)
        t2t_params = {
            'kernel_size': kernel_size,
            'stride': stride,
            'padding': padding
        }
        self.ss = SoftSplit(channel // 2,
                            hidden,
                            kernel_size,
                            stride,
                            padding,
                            t2t_param=t2t_params)
        self.sc = SoftComp(channel // 2, hidden, kernel_size, stride, padding)

        n_vecs = 1
        for i, d in enumerate(kernel_size):
            n_vecs *= int((output_size[i] + 2 * padding[i] -
                           (d - 1) - 1) / stride[i] + 1)

        blocks = []
        depths = 8
        num_heads = [4] * depths
        window_size = [(5, 9)] * depths
        focal_windows = [(5, 9)] * depths
        focal_levels = [2] * depths
        pool_method = "fc"

        for i in range(depths):
            blocks.append(
                TemporalFocalTransformerBlock(dim=hidden,
                                              num_heads=num_heads[i],
                                              window_size=window_size[i],
                                              focal_level=focal_levels[i],
                                              focal_window=focal_windows[i],
                                              n_vecs=n_vecs,
                                              t2t_params=t2t_params,
                                              pool_method=pool_method))
        self.transformer = nn.Sequential(*blocks)

        if init_weights:
            self.init_weights()
            # Need to initial the weights of MSDeformAttn specifically
            for m in self.modules():
                if isinstance(m, SecondOrderDeformableAlignment):
                    m.init_offset()

        # flow completion network
        self.update_spynet = SPyNet()

    def forward_bidirect_flow(self, masked_local_frames):
        b, l_t, c, h, w = masked_local_frames.size()

        # compute forward and backward flows of masked frames
        masked_local_frames = F.interpolate(masked_local_frames.view(
            -1, c, h, w),
                                            scale_factor=1 / 4,
                                            mode='bilinear',
                                            align_corners=True,
                                            recompute_scale_factor=True)
        masked_local_frames = masked_local_frames.view(b, l_t, c, h // 4,
                                                       w // 4)
        mlf_1 = masked_local_frames[:, :-1, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        mlf_2 = masked_local_frames[:, 1:, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        pred_flows_forward = self.update_spynet(mlf_1, mlf_2)
        pred_flows_backward = self.update_spynet(mlf_2, mlf_1)

        pred_flows_forward = pred_flows_forward.view(b, l_t - 1, 2, h // 4,
                                                     w // 4)
        pred_flows_backward = pred_flows_backward.view(b, l_t - 1, 2, h // 4,
                                                       w // 4)

        return pred_flows_forward, pred_flows_backward

    def forward(self, masked_frames, num_local_frames):
        l_t = num_local_frames
        b, t, ori_c, ori_h, ori_w = masked_frames.size()

        # normalization before feeding into the flow completion module
        masked_local_frames = (masked_frames[:, :l_t, ...] + 1) / 2
        pred_flows = self.forward_bidirect_flow(masked_local_frames)

        # extracting features and performing the feature propagation on local features
        enc_feat = self.encoder(masked_frames.view(b * t, ori_c, ori_h, ori_w))
        _, c, h, w = enc_feat.size()
        fold_output_size = (h, w)
        local_feat = enc_feat.view(b, t, c, h, w)[:, :l_t, ...]
        ref_feat = enc_feat.view(b, t, c, h, w)[:, l_t:, ...]
        local_feat = self.feat_prop_module(local_feat, pred_flows[0],
                                           pred_flows[1])
        enc_feat = torch.cat((local_feat, ref_feat), dim=1)

        # content hallucination through stacking multiple temporal focal transformer blocks
        trans_feat = self.ss(enc_feat.view(-1, c, h, w), b, fold_output_size)
        trans_feat = self.transformer([trans_feat, fold_output_size])
        trans_feat = self.sc(trans_feat[0], t, fold_output_size)
        trans_feat = trans_feat.view(b, t, -1, h, w)
        enc_feat = enc_feat + trans_feat

        # decode frames from features
        output = self.decoder(enc_feat.view(b * t, c, h, w))
        output = torch.tanh(output)
        return output, pred_flows

# ######################################################################
#  Discriminator for Temporal Patch GAN
# ######################################################################

class Discriminator(BaseNetwork):
    def __init__(self,
                 in_channels=3,
                 use_sigmoid=False,
                 use_spectral_norm=True,
                 init_weights=True):
        super(Discriminator, self).__init__()
        self.use_sigmoid = use_sigmoid
        nf = 32

        self.conv = nn.Sequential(
            spectral_norm(
                nn.Conv3d(in_channels=in_channels,
                          out_channels=nf * 1,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=1,
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(64, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 1,
                          nf * 2,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(128, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 2,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv3d(nf * 4,
                      nf * 4,
                      kernel_size=(3, 5, 5),
                      stride=(1, 2, 2),
                      padding=(1, 2, 2)))

        if init_weights:
            self.init_weights()

    def forward(self, xs):
        # T, C, H, W = xs.shape (old)
        # B, T, C, H, W (new)
        xs_t = torch.transpose(xs, 1, 2)
        feat = self.conv(xs_t)
        if self.use_sigmoid:
            feat = torch.sigmoid(feat)
        out = torch.transpose(feat, 1, 2)  # B, T, C, H, W
        return out

def spectral_norm(module, mode=True):
    if mode:
        return _spectral_norm(module)
    return module


BaseNetwork

''' Towards An End-to-End Framework for Video Inpainting
'''

import torch
import torch.nn as nn
import torch.nn.functional as F

from inpainter.model.modules.flow_comp import SPyNet
from inpainter.model.modules.feat_prop import BidirectionalPropagation, SecondOrderDeformableAlignment
from inpainter.model.modules.tfocal_transformer_hq import TemporalFocalTransformerBlock, SoftSplit, SoftComp
from inpainter.model.modules.spectral_norm import spectral_norm as _spectral_norm

class BaseNetwork(nn.Module):
    def __init__(self):
        super(BaseNetwork, self).__init__()

    def print_network(self):
        if isinstance(self, list):
            self = self[0]
        num_params = 0
        for param in self.parameters():
            num_params += param.numel()
        print(
            'Network [%s] was created. Total number of parameters: %.1f million. '
            'To see the architecture, do print(network).' %
            (type(self).__name__, num_params / 1000000))

    def init_weights(self, init_type='normal', gain=0.02):
        '''
        initialize network's weights
        init_type: normal | xavier | kaiming | orthogonal
        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
        '''
        def init_func(m):
            classname = m.__class__.__name__
            if classname.find('InstanceNorm2d') != -1:
                if hasattr(m, 'weight') and m.weight is not None:
                    nn.init.constant_(m.weight.data, 1.0)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)
            elif hasattr(m, 'weight') and (classname.find('Conv') != -1
                                           or classname.find('Linear') != -1):
                if init_type == 'normal':
                    nn.init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    nn.init.xavier_normal_(m.weight.data, gain=gain)
                elif init_type == 'xavier_uniform':
                    nn.init.xavier_uniform_(m.weight.data, gain=1.0)
                elif init_type == 'kaiming':
                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    nn.init.orthogonal_(m.weight.data, gain=gain)
                elif init_type == 'none':  # uses pytorch's default init method
                    m.reset_parameters()
                else:
                    raise NotImplementedError(
                        'initialization method [%s] is not implemented' %
                        init_type)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)

        self.apply(init_func)

        # propagate to children
        for m in self.children():
            if hasattr(m, 'init_weights'):
                m.init_weights(init_type, gain)

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.group = [1, 2, 4, 8, 1]
        self.layers = nn.ModuleList([
            nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, groups=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(640, 512, kernel_size=3, stride=1, padding=1, groups=2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(768, 384, kernel_size=3, stride=1, padding=1, groups=4),


Encoder

                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    nn.init.orthogonal_(m.weight.data, gain=gain)
                elif init_type == 'none':  # uses pytorch's default init method
                    m.reset_parameters()
                else:
                    raise NotImplementedError(
                        'initialization method [%s] is not implemented' %
                        init_type)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)

        self.apply(init_func)

        # propagate to children
        for m in self.children():
            if hasattr(m, 'init_weights'):
                m.init_weights(init_type, gain)

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.group = [1, 2, 4, 8, 1]
        self.layers = nn.ModuleList([
            nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, groups=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(640, 512, kernel_size=3, stride=1, padding=1, groups=2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(768, 384, kernel_size=3, stride=1, padding=1, groups=4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(640, 256, kernel_size=3, stride=1, padding=1, groups=8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 128, kernel_size=3, stride=1, padding=1, groups=1),
            nn.LeakyReLU(0.2, inplace=True)
        ])

    def forward(self, x):
        bt, c, _, _ = x.size()
        # h, w = h//4, w//4
        out = x
        for i, layer in enumerate(self.layers):
            if i == 8:
                x0 = out
                _, _, h, w = x0.size()
            if i > 8 and i % 2 == 0:
                g = self.group[(i - 8) // 2]
                x = x0.view(bt, g, -1, h, w)
                o = out.view(bt, g, -1, h, w)
                out = torch.cat([x, o], 2).view(bt, -1, h, w)
            out = layer(out)
        return out

class deconv(nn.Module):
    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel_size=3,
                 padding=0):
        super().__init__()
        self.conv = nn.Conv2d(input_channel,
                              output_channel,
                              kernel_size=kernel_size,
                              stride=1,
                              padding=padding)

    def forward(self, x):
        x = F.interpolate(x,
                          scale_factor=2,
                          mode='bilinear',
                          align_corners=True)


deconv

            nn.LeakyReLU(0.2, inplace=True)
        ])

    def forward(self, x):
        bt, c, _, _ = x.size()
        # h, w = h//4, w//4
        out = x
        for i, layer in enumerate(self.layers):
            if i == 8:
                x0 = out
                _, _, h, w = x0.size()
            if i > 8 and i % 2 == 0:
                g = self.group[(i - 8) // 2]
                x = x0.view(bt, g, -1, h, w)
                o = out.view(bt, g, -1, h, w)
                out = torch.cat([x, o], 2).view(bt, -1, h, w)
            out = layer(out)
        return out

class deconv(nn.Module):
    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel_size=3,
                 padding=0):
        super().__init__()
        self.conv = nn.Conv2d(input_channel,
                              output_channel,
                              kernel_size=kernel_size,
                              stride=1,
                              padding=padding)

    def forward(self, x):
        x = F.interpolate(x,
                          scale_factor=2,
                          mode='bilinear',
                          align_corners=True)
        return self.conv(x)

class InpaintGenerator(BaseNetwork):
    def __init__(self, init_weights=True):
        super(InpaintGenerator, self).__init__()
        channel = 256
        hidden = 512

        # encoder
        self.encoder = Encoder()

        # decoder
        self.decoder = nn.Sequential(
            deconv(channel // 2, 128, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            deconv(64, 64, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1))


InpaintGenerator

    def __init__(self,
                 input_channel,
                 output_channel,
                 kernel_size=3,
                 padding=0):
        super().__init__()
        self.conv = nn.Conv2d(input_channel,
                              output_channel,
                              kernel_size=kernel_size,
                              stride=1,
                              padding=padding)

    def forward(self, x):
        x = F.interpolate(x,
                          scale_factor=2,
                          mode='bilinear',
                          align_corners=True)
        return self.conv(x)

class InpaintGenerator(BaseNetwork):
    def __init__(self, init_weights=True):
        super(InpaintGenerator, self).__init__()
        channel = 256
        hidden = 512

        # encoder
        self.encoder = Encoder()

        # decoder
        self.decoder = nn.Sequential(
            deconv(channel // 2, 128, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            deconv(64, 64, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1))

        # feature propagation module
        self.feat_prop_module = BidirectionalPropagation(channel // 2)

        # soft split and soft composition
        kernel_size = (7, 7)
        padding = (3, 3)
        stride = (3, 3)
        output_size = (60, 108)
        t2t_params = {
            'kernel_size': kernel_size,
            'stride': stride,
            'padding': padding
        }
        self.ss = SoftSplit(channel // 2,
                            hidden,
                            kernel_size,
                            stride,
                            padding,
                            t2t_param=t2t_params)
        self.sc = SoftComp(channel // 2, hidden, kernel_size, stride, padding)

        n_vecs = 1
        for i, d in enumerate(kernel_size):
            n_vecs *= int((output_size[i] + 2 * padding[i] -
                           (d - 1) - 1) / stride[i] + 1)

        blocks = []
        depths = 8
        num_heads = [4] * depths
        window_size = [(5, 9)] * depths
        focal_windows = [(5, 9)] * depths
        focal_levels = [2] * depths
        pool_method = "fc"

        for i in range(depths):
            blocks.append(
                TemporalFocalTransformerBlock(dim=hidden,
                                              num_heads=num_heads[i],
                                              window_size=window_size[i],
                                              focal_level=focal_levels[i],
                                              focal_window=focal_windows[i],
                                              n_vecs=n_vecs,
                                              t2t_params=t2t_params,
                                              pool_method=pool_method))
        self.transformer = nn.Sequential(*blocks)

        if init_weights:
            self.init_weights()
            # Need to initial the weights of MSDeformAttn specifically
            for m in self.modules():
                if isinstance(m, SecondOrderDeformableAlignment):
                    m.init_offset()

        # flow completion network
        self.update_spynet = SPyNet()

    def forward_bidirect_flow(self, masked_local_frames):
        b, l_t, c, h, w = masked_local_frames.size()

        # compute forward and backward flows of masked frames
        masked_local_frames = F.interpolate(masked_local_frames.view(
            -1, c, h, w),
                                            scale_factor=1 / 4,
                                            mode='bilinear',
                                            align_corners=True,
                                            recompute_scale_factor=True)
        masked_local_frames = masked_local_frames.view(b, l_t, c, h // 4,
                                                       w // 4)
        mlf_1 = masked_local_frames[:, :-1, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        mlf_2 = masked_local_frames[:, 1:, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        pred_flows_forward = self.update_spynet(mlf_1, mlf_2)
        pred_flows_backward = self.update_spynet(mlf_2, mlf_1)

        pred_flows_forward = pred_flows_forward.view(b, l_t - 1, 2, h // 4,
                                                     w // 4)
        pred_flows_backward = pred_flows_backward.view(b, l_t - 1, 2, h // 4,
                                                       w // 4)

        return pred_flows_forward, pred_flows_backward

    def forward(self, masked_frames, num_local_frames):
        l_t = num_local_frames
        b, t, ori_c, ori_h, ori_w = masked_frames.size()

        # normalization before feeding into the flow completion module
        masked_local_frames = (masked_frames[:, :l_t, ...] + 1) / 2
        pred_flows = self.forward_bidirect_flow(masked_local_frames)

        # extracting features and performing the feature propagation on local features
        enc_feat = self.encoder(masked_frames.view(b * t, ori_c, ori_h, ori_w))
        _, c, h, w = enc_feat.size()
        fold_output_size = (h, w)
        local_feat = enc_feat.view(b, t, c, h, w)[:, :l_t, ...]
        ref_feat = enc_feat.view(b, t, c, h, w)[:, l_t:, ...]
        local_feat = self.feat_prop_module(local_feat, pred_flows[0],
                                           pred_flows[1])
        enc_feat = torch.cat((local_feat, ref_feat), dim=1)

        # content hallucination through stacking multiple temporal focal transformer blocks
        trans_feat = self.ss(enc_feat.view(-1, c, h, w), b, fold_output_size)
        trans_feat = self.transformer([trans_feat, fold_output_size])
        trans_feat = self.sc(trans_feat[0], t, fold_output_size)
        trans_feat = trans_feat.view(b, t, -1, h, w)
        enc_feat = enc_feat + trans_feat

        # decode frames from features
        output = self.decoder(enc_feat.view(b * t, c, h, w))
        output = torch.tanh(output)
        return output, pred_flows

# ######################################################################
#  Discriminator for Temporal Patch GAN
# ######################################################################

class Discriminator(BaseNetwork):
    def __init__(self,
                 in_channels=3,
                 use_sigmoid=False,
                 use_spectral_norm=True,
                 init_weights=True):
        super(Discriminator, self).__init__()
        self.use_sigmoid = use_sigmoid
        nf = 32

        self.conv = nn.Sequential(
            spectral_norm(
                nn.Conv3d(in_channels=in_channels,


Discriminator

        enc_feat = torch.cat((local_feat, ref_feat), dim=1)

        # content hallucination through stacking multiple temporal focal transformer blocks
        trans_feat = self.ss(enc_feat.view(-1, c, h, w), b, fold_output_size)
        trans_feat = self.transformer([trans_feat, fold_output_size])
        trans_feat = self.sc(trans_feat[0], t, fold_output_size)
        trans_feat = trans_feat.view(b, t, -1, h, w)
        enc_feat = enc_feat + trans_feat

        # decode frames from features
        output = self.decoder(enc_feat.view(b * t, c, h, w))
        output = torch.tanh(output)
        return output, pred_flows

# ######################################################################
#  Discriminator for Temporal Patch GAN
# ######################################################################

class Discriminator(BaseNetwork):
    def __init__(self,
                 in_channels=3,
                 use_sigmoid=False,
                 use_spectral_norm=True,
                 init_weights=True):
        super(Discriminator, self).__init__()
        self.use_sigmoid = use_sigmoid
        nf = 32

        self.conv = nn.Sequential(
            spectral_norm(
                nn.Conv3d(in_channels=in_channels,
                          out_channels=nf * 1,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=1,
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(64, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 1,
                          nf * 2,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(128, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 2,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv3d(nf * 4,
                      nf * 4,
                      kernel_size=(3, 5, 5),
                      stride=(1, 2, 2),
                      padding=(1, 2, 2)))

        if init_weights:
            self.init_weights()

    def forward(self, xs):
        # T, C, H, W = xs.shape (old)
        # B, T, C, H, W (new)
        xs_t = torch.transpose(xs, 1, 2)
        feat = self.conv(xs_t)
        if self.use_sigmoid:
            feat = torch.sigmoid(feat)
        out = torch.transpose(feat, 1, 2)  # B, T, C, H, W
        return out

def spectral_norm(module, mode=True):
    if mode:
        return _spectral_norm(module)
    return module


            nn.Conv3d(nf * 4,
                      nf * 4,
                      kernel_size=(3, 5, 5),
                      stride=(1, 2, 2),
                      padding=(1, 2, 2)))

        if init_weights:
            self.init_weights()

    def forward(self, xs):
        # T, C, H, W = xs.shape (old)
        # B, T, C, H, W (new)
        xs_t = torch.transpose(xs, 1, 2)
        feat = self.conv(xs_t)
        if self.use_sigmoid:
            feat = torch.sigmoid(feat)
        out = torch.transpose(feat, 1, 2)  # B, T, C, H, W
        return out

def spectral_norm(module, mode=True):
    if mode:
        return _spectral_norm(module)
    return module


        # content hallucination through stacking multiple temporal focal transformer blocks
        trans_feat = self.ss(enc_feat.view(-1, c, h, w), b, fold_output_size)
        trans_feat = self.transformer([trans_feat, fold_output_size])
        trans_feat = self.sc(trans_feat[0], t, fold_output_size)
        trans_feat = trans_feat.view(b, t, -1, h, w)
        enc_feat = enc_feat + trans_feat

        # decode frames from features
        output = self.decoder(enc_feat.view(b * t, c, h, w))
        output = torch.tanh(output)
        return output, pred_flows

# ######################################################################
#  Discriminator for Temporal Patch GAN
# ######################################################################

class Discriminator(BaseNetwork):
    def __init__(self,
                 in_channels=3,
                 use_sigmoid=False,
                 use_spectral_norm=True,
                 init_weights=True):
        super(Discriminator, self).__init__()
        self.use_sigmoid = use_sigmoid
        nf = 32

        self.conv = nn.Sequential(
            spectral_norm(
                nn.Conv3d(in_channels=in_channels,
                          out_channels=nf * 1,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=1,
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(64, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 1,
                          nf * 2,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(128, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 2,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv3d(nf * 4,
                      nf * 4,
                      kernel_size=(3, 5, 5),
                      stride=(1, 2, 2),
                      padding=(1, 2, 2)))

        if init_weights:
            self.init_weights()

    def forward(self, xs):
        # T, C, H, W = xs.shape (old)
        # B, T, C, H, W (new)
        xs_t = torch.transpose(xs, 1, 2)
        feat = self.conv(xs_t)
        if self.use_sigmoid:
            feat = torch.sigmoid(feat)
        out = torch.transpose(feat, 1, 2)  # B, T, C, H, W
        return out

def spectral_norm(module, mode=True):
    if mode:
        return _spectral_norm(module)
    return module


''' Towards An End-to-End Framework for Video Inpainting
'''

import torch
import torch.nn as nn
import torch.nn.functional as F

from inpainter.model.modules.flow_comp import SPyNet
from inpainter.model.modules.feat_prop import BidirectionalPropagation, SecondOrderDeformableAlignment
from inpainter.model.modules.tfocal_transformer_hq import TemporalFocalTransformerBlock, SoftSplit, SoftComp
from inpainter.model.modules.spectral_norm import spectral_norm as _spectral_norm

class BaseNetwork(nn.Module):
    def __init__(self):
        super(BaseNetwork, self).__init__()

    def print_network(self):
        if isinstance(self, list):
            self = self[0]
        num_params = 0
        for param in self.parameters():
            num_params += param.numel()
        print(
            'Network [%s] was created. Total number of parameters: %.1f million. '
            'To see the architecture, do print(network).' %
            (type(self).__name__, num_params / 1000000))

    def init_weights(self, init_type='normal', gain=0.02):
        '''
        initialize network's weights
        init_type: normal | xavier | kaiming | orthogonal
        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
        '''
        def init_func(m):
            classname = m.__class__.__name__
            if classname.find('InstanceNorm2d') != -1:
                if hasattr(m, 'weight') and m.weight is not None:
                    nn.init.constant_(m.weight.data, 1.0)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)
            elif hasattr(m, 'weight') and (classname.find('Conv') != -1
                                           or classname.find('Linear') != -1):
                if init_type == 'normal':
                    nn.init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    nn.init.xavier_normal_(m.weight.data, gain=gain)


from inpainter.model.modules.feat_prop import BidirectionalPropagation, SecondOrderDeformableAlignment
from inpainter.model.modules.tfocal_transformer_hq import TemporalFocalTransformerBlock, SoftSplit, SoftComp
from inpainter.model.modules.spectral_norm import spectral_norm as _spectral_norm

class BaseNetwork(nn.Module):
    def __init__(self):
        super(BaseNetwork, self).__init__()

    def print_network(self):
        if isinstance(self, list):
            self = self[0]
        num_params = 0
        for param in self.parameters():
            num_params += param.numel()
        print(
            'Network [%s] was created. Total number of parameters: %.1f million. '
            'To see the architecture, do print(network).' %
            (type(self).__name__, num_params / 1000000))

    def init_weights(self, init_type='normal', gain=0.02):
        '''
        initialize network's weights
        init_type: normal | xavier | kaiming | orthogonal
        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
        '''
        def init_func(m):
            classname = m.__class__.__name__
            if classname.find('InstanceNorm2d') != -1:
                if hasattr(m, 'weight') and m.weight is not None:
                    nn.init.constant_(m.weight.data, 1.0)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)
            elif hasattr(m, 'weight') and (classname.find('Conv') != -1
                                           or classname.find('Linear') != -1):
                if init_type == 'normal':
                    nn.init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    nn.init.xavier_normal_(m.weight.data, gain=gain)
                elif init_type == 'xavier_uniform':
                    nn.init.xavier_uniform_(m.weight.data, gain=1.0)
                elif init_type == 'kaiming':
                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    nn.init.orthogonal_(m.weight.data, gain=gain)
                elif init_type == 'none':  # uses pytorch's default init method
                    m.reset_parameters()
                else:
                    raise NotImplementedError(
                        'initialization method [%s] is not implemented' %
                        init_type)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)

        self.apply(init_func)

        # propagate to children
        for m in self.children():
            if hasattr(m, 'init_weights'):
                m.init_weights(init_type, gain)

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.group = [1, 2, 4, 8, 1]
        self.layers = nn.ModuleList([
            nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, groups=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(640, 512, kernel_size=3, stride=1, padding=1, groups=2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(768, 384, kernel_size=3, stride=1, padding=1, groups=4),


forward

            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(
                nn.Conv3d(nf * 4,
                          nf * 4,
                          kernel_size=(3, 5, 5),
                          stride=(1, 2, 2),
                          padding=(1, 2, 2),
                          bias=not use_spectral_norm), use_spectral_norm),
            # nn.InstanceNorm2d(256, track_running_stats=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv3d(nf * 4,
                      nf * 4,
                      kernel_size=(3, 5, 5),
                      stride=(1, 2, 2),
                      padding=(1, 2, 2)))

        if init_weights:
            self.init_weights()

    def forward(self, xs):
        # T, C, H, W = xs.shape (old)
        # B, T, C, H, W (new)
        xs_t = torch.transpose(xs, 1, 2)
        feat = self.conv(xs_t)
        if self.use_sigmoid:
            feat = torch.sigmoid(feat)
        out = torch.transpose(feat, 1, 2)  # B, T, C, H, W
        return out

def spectral_norm(module, mode=True):
    if mode:
        return _spectral_norm(module)
    return module


                TemporalFocalTransformerBlock(dim=hidden,
                                              num_heads=num_heads[i],
                                              window_size=window_size[i],
                                              focal_level=focal_levels[i],
                                              focal_window=focal_windows[i],
                                              n_vecs=n_vecs,
                                              t2t_params=t2t_params,
                                              pool_method=pool_method))
        self.transformer = nn.Sequential(*blocks)

        if init_weights:
            self.init_weights()
            # Need to initial the weights of MSDeformAttn specifically
            for m in self.modules():
                if isinstance(m, SecondOrderDeformableAlignment):
                    m.init_offset()

        # flow completion network
        self.update_spynet = SPyNet()

    def forward_bidirect_flow(self, masked_local_frames):
        b, l_t, c, h, w = masked_local_frames.size()

        # compute forward and backward flows of masked frames
        masked_local_frames = F.interpolate(masked_local_frames.view(
            -1, c, h, w),
                                            scale_factor=1 / 4,
                                            mode='bilinear',
                                            align_corners=True,
                                            recompute_scale_factor=True)
        masked_local_frames = masked_local_frames.view(b, l_t, c, h // 4,
                                                       w // 4)
        mlf_1 = masked_local_frames[:, :-1, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        mlf_2 = masked_local_frames[:, 1:, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        pred_flows_forward = self.update_spynet(mlf_1, mlf_2)
        pred_flows_backward = self.update_spynet(mlf_2, mlf_1)

        pred_flows_forward = pred_flows_forward.view(b, l_t - 1, 2, h // 4,
                                                     w // 4)
        pred_flows_backward = pred_flows_backward.view(b, l_t - 1, 2, h // 4,
                                                       w // 4)

        return pred_flows_forward, pred_flows_backward

    def forward(self, masked_frames, num_local_frames):
        l_t = num_local_frames
        b, t, ori_c, ori_h, ori_w = masked_frames.size()

        # normalization before feeding into the flow completion module
        masked_local_frames = (masked_frames[:, :l_t, ...] + 1) / 2
        pred_flows = self.forward_bidirect_flow(masked_local_frames)

        # extracting features and performing the feature propagation on local features
        enc_feat = self.encoder(masked_frames.view(b * t, ori_c, ori_h, ori_w))
        _, c, h, w = enc_feat.size()
        fold_output_size = (h, w)
        local_feat = enc_feat.view(b, t, c, h, w)[:, :l_t, ...]
        ref_feat = enc_feat.view(b, t, c, h, w)[:, l_t:, ...]
        local_feat = self.feat_prop_module(local_feat, pred_flows[0],
                                           pred_flows[1])
        enc_feat = torch.cat((local_feat, ref_feat), dim=1)

        # content hallucination through stacking multiple temporal focal transformer blocks