import torch
import numpy as np
import cv2
import torch.nn.functional as F
from src.utils import *
import sys
sys.path.append("./src/ebsynth/deps/gmflow/")
from gmflow.geometry import flow_warp
"""
==========================================================================
* warp_tensor(): warp and fuse tensors based on optical flow and mask
* get_single_mapping_ind(): get pixel index correspondence between two frames
* get_mapping_ind(): get pixel index correspondence between consecutive frames within a batch
==========================================================================
"""
@torch.no_grad()
def warp_tensor(sample, flows, occs, saliency, unet_chunk_size):
"""
Warp images or features based on optical flow
Fuse the warped imges or features based on occusion masks and saliency map
"""
scale = sample.shape[2] * 1.0 / flows[0].shape[2]
kernel = int(1 / scale)
bwd_flow_ = F.interpolate(flows[1] * scale, scale_factor=scale, mode='bilinear')
bwd_occ_ = F.max_pool2d(occs[1].unsqueeze(1), kernel_size=kernel) # (N-1)*1*H1*W1
if scale == 1:
bwd_occ_ = Dilate(kernel_size=13, device=sample.device)(bwd_occ_)
fwd_flow_ = F.interpolate(flows[0] * scale, scale_factor=scale, mode='bilinear')
fwd_occ_ = F.max_pool2d(occs[0].unsqueeze(1), kernel_size=kernel) # (N-1)*1*H1*W1
if scale == 1:
fwd_occ_ = Dilate(kernel_size=13, device=sample.device)(fwd_occ_)
scale2 = sample.shape[2] * 1.0 / saliency.shape[2]
saliency = F.interpolate(saliency, scale_factor=scale2, mode='bilinear')
latent = sample.to(torch.float32)
video_length = sample.shape[0] // unet_chunk_size
warp_saliency = flow_warp(saliency, bwd_flow_)
warp_saliency_ = flow_warp(saliency[0:1], fwd_flow_[video_length-1:video_length])
for j in range(unet_chunk_size):
for ii in range(video_length-1):
i = video_length * j + ii
warped_image = flow_warp(latent[i:i+1], bwd_flow_[ii:ii+1])
mask = (1 - bwd_occ_[ii:ii+1]) * saliency[ii+1:ii+2] * warp_saliency[ii:ii+1]
latent[i+1:i+2] = latent[i+1:i+2] * (1-mask) + warped_image * mask
i = video_length * j
ii = video_length - 1
warped_image = flow_warp(latent[i:i+1], fwd_flow_[ii:ii+1])
mask = (1 - fwd_occ_[ii:ii+1]) * saliency[ii:ii+1] * warp_saliency_
latent[ii+i:ii+i+1] = latent[ii+i:ii+i+1] * (1-mask) + warped_image * mask
return latent.to(sample.dtype)
@torch.no_grad()
def get_single_mapping_ind(bwd_flow, bwd_occ, imgs, scale=1.0):
"""
FLATTEN: Optical fLow-guided attention (Temoporal-guided attention)
Find the correspondence between every pixels in a pair of frames
[input]
bwd_flow: 1*2*H*W
bwd_occ: 1*H*W i.e., f2 = warp(f1, bwd_flow) * bwd_occ
imgs: 2*3*H*W i.e., [f1,f2]
[output]
mapping_ind: pixel index correspondence
unlinkedmask: indicate whether a pixel has no correspondence
i.e., f2 = f1[mapping_ind] * unlinkedmask
"""
flows = F.interpolate(bwd_flow, scale_factor=1./scale, mode='bilinear')[0][[1,0]] / scale # 2*H*W
_, H, W = flows.shape
masks = torch.logical_not(F.interpolate(bwd_occ[None], scale_factor=1./scale, mode='bilinear') > 0.5)[0] # 1*H*W
frames = F.interpolate(imgs, scale_factor=1./scale, mode='bilinear').view(2, 3, -1) # 2*3*HW
grid = torch.stack(torch.meshgrid([torch.arange(H), torch.arange(W)]), dim=0).to(flows.device) # 2*H*W
warp_grid = torch.round(grid + flows)
mask = torch.logical_and(torch.logical_and(torch.logical_and(torch.logical_and(warp_grid[0] >= 0, warp_grid[0] < H),
warp_grid[1] >= 0), warp_grid[1] < W), masks[0]).view(-1) # HW
warp_grid = warp_grid.view(2, -1) # 2*HW
warp_ind = (warp_grid[0] * W + warp_grid[1]).to(torch.long) # HW
mapping_ind = torch.zeros_like(warp_ind) - 1 # HW
for f0ind, f1ind in enumerate(warp_ind):
if mask[f0ind]:
if mapping_ind[f1ind] == -1:
mapping_ind[f1ind] = f0ind
else:
targetv = frames[0,:,f1ind]
pref0ind = mapping_ind[f1ind]
prev = frames[1,:,pref0ind]
v = frames[1,:,f0ind]
if ((prev - targetv)**2).mean() > ((v - targetv)**2).mean():
mask[pref0ind] = False
mapping_ind[f1ind] = f0ind
else:
mask[f0ind] = False
unusedind = torch.arange(len(mask)).to(mask.device)[~mask]
unlinkedmask = mapping_ind == -1
mapping_ind[unlinkedmask] = unusedind
return mapping_ind, unlinkedmask
@torch.no_grad()
def get_mapping_ind(bwd_flows, bwd_occs, imgs, scale=1.0):
"""
FLATTEN: Optical fLow-guided attention (Temoporal-guided attention)
Find pixel correspondence between every consecutive frames in a batch
[input]
bwd_flow: (N-1)*2*H*W
bwd_occ: (N-1)*H*W
imgs: N*3*H*W
[output]
fwd_mappings: N*1*HW
bwd_mappings: N*1*HW
flattn_mask: HW*1*N*N
i.e., imgs[i,:,fwd_mappings[i]] corresponds to imgs[0]
i.e., imgs[i,:,fwd_mappings[i]][:,bwd_mappings[i]] restore the original imgs[i]
"""
N, H, W = imgs.shape[0], int(imgs.shape[2] // scale), int(imgs.shape[3] // scale)
iterattn_mask = torch.ones(H*W, N, N, dtype=torch.bool).to(imgs.device)
for i in range(len(imgs)-1):
one_mask = torch.ones(N, N, dtype=torch.bool).to(imgs.device)
one_mask[:i+1,i+1:] = False
one_mask[i+1:,:i+1] = False
mapping_ind, unlinkedmask = get_single_mapping_ind(bwd_flows[i:i+1], bwd_occs[i:i+1], imgs[i:i+2], scale)
if i == 0:
fwd_mapping = [torch.arange(len(mapping_ind)).to(mapping_ind.device)]
bwd_mapping = [torch.arange(len(mapping_ind)).to(mapping_ind.device)]
iterattn_mask[unlinkedmask[fwd_mapping[-1]]] = torch.logical_and(iterattn_mask[unlinkedmask[fwd_mapping[-1]]], one_mask)
fwd_mapping += [mapping_ind[fwd_mapping[-1]]]
bwd_mapping += [torch.sort(fwd_mapping[-1])[1]]
fwd_mappings = torch.stack(fwd_mapping, dim=0).unsqueeze(1)
bwd_mappings = torch.stack(bwd_mapping, dim=0).unsqueeze(1)
return fwd_mappings, bwd_mappings, iterattn_mask.unsqueeze(1)