457 lines
17 KiB
Python
457 lines
17 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
from timm.models.layers import trunc_normal_
|
|
from torch import nn
|
|
|
|
from .generator import OcclusionAwareGenerator
|
|
from .util import (
|
|
# AntiAliasInterpolation2d,
|
|
Hourglass,
|
|
batch_bilinear_sampler,
|
|
bilinear_sampler,
|
|
coords_grid,
|
|
kp2gaussian,
|
|
# make_coordinate_grid,
|
|
)
|
|
|
|
|
|
class CorrBlock:
|
|
def __init__(self, corr, num_levels=2, radius=3):
|
|
self.num_levels = num_levels
|
|
self.radius = radius
|
|
self.corr_pyramid = []
|
|
|
|
self.corr_pyramid.append(corr)
|
|
for i in range(self.num_levels - 1):
|
|
corr = F.avg_pool2d(corr, 2, stride=2)
|
|
self.corr_pyramid.append(corr)
|
|
|
|
def __call__(self, coords):
|
|
r = self.radius
|
|
coords = coords.permute(0, 2, 3, 1)
|
|
batch, h1, w1, _ = coords.shape
|
|
|
|
out_pyramid = []
|
|
for i in range(self.num_levels):
|
|
corr = self.corr_pyramid[i]
|
|
dx = torch.linspace(-r, r, 2 * r + 1, device=coords.device)
|
|
dy = torch.linspace(-r, r, 2 * r + 1, device=coords.device)
|
|
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
|
|
|
|
centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i
|
|
delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
|
|
coords_lvl = centroid_lvl + delta_lvl
|
|
|
|
if batch > 1 and h1 >= 128:
|
|
corr = batch_bilinear_sampler(
|
|
corr, coords_lvl, h=h1, w=w1, mini_batch=1
|
|
)
|
|
else:
|
|
corr = bilinear_sampler(corr, coords_lvl)
|
|
# corr = bilinear_sampler(corr, coords_lvl)
|
|
corr = corr.view(batch, h1, w1, -1)
|
|
out_pyramid.append(corr)
|
|
|
|
out = torch.cat(out_pyramid, dim=-1)
|
|
return out.permute(0, 3, 1, 2).contiguous().float()
|
|
|
|
|
|
class BasicMotionEncoder(nn.Module):
|
|
def __init__(self, num_levels=2, radius=3):
|
|
super(BasicMotionEncoder, self).__init__()
|
|
cor_planes = num_levels * (2 * radius + 1) ** 2
|
|
self.convc1 = nn.Conv2d(cor_planes, 128, 1, padding=0)
|
|
self.convc2 = nn.Conv2d(128, 96, 3, padding=1)
|
|
self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
|
|
self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
|
|
self.conv = nn.Conv2d(64 + 96, 128 - 2, 3, padding=1)
|
|
|
|
def forward(self, delta_flow, corr):
|
|
cor = F.relu(self.convc1(corr))
|
|
cor = F.relu(self.convc2(cor))
|
|
flo = F.relu(self.convf1(delta_flow))
|
|
flo = F.relu(self.convf2(flo))
|
|
|
|
cor_flo = torch.cat([cor, flo], dim=1)
|
|
out = F.relu(self.conv(cor_flo))
|
|
return torch.cat([out, delta_flow], dim=1)
|
|
|
|
|
|
class RefineFlow(nn.Module):
|
|
def __init__(self):
|
|
super(RefineFlow, self).__init__()
|
|
self.convc1 = nn.Conv2d(192, 128, 3, padding=1)
|
|
# self.convc2 = nn.Conv2d(128, 64, 3, padding=1)
|
|
self.conv1 = nn.Conv2d(256, 128, 3, padding=1)
|
|
self.conv2 = nn.Conv2d(128, 2, 3, padding=1)
|
|
self.convo1 = nn.Conv2d(256, 128, 3, padding=1)
|
|
self.convo2 = nn.Conv2d(128, 1, 3, padding=1)
|
|
|
|
def forward(self, m_f, warp_f):
|
|
c = F.relu(self.convc1(warp_f))
|
|
## c = F.relu(self.convc2(c))
|
|
inp = torch.cat([m_f, c], dim=1)
|
|
flow = self.conv2(F.relu(self.conv1(inp)))
|
|
occ = self.convo2(F.relu(self.convo1(inp)))
|
|
out = torch.cat([flow, occ], dim=1)
|
|
return out, inp
|
|
# return flow, inp
|
|
|
|
|
|
class RaftFlow(nn.Module):
|
|
def __init__(
|
|
self,
|
|
prior_only=False,
|
|
num_kp=10,
|
|
dim=256,
|
|
size=256,
|
|
generator=None,
|
|
driving_encoder=None,
|
|
source_encoder=None,
|
|
):
|
|
super(RaftFlow, self).__init__()
|
|
|
|
self.scale = dim**-0.5
|
|
self.size = size
|
|
## basic flow resolution of the prior motion model
|
|
self.h = size // 4
|
|
self.w = size // 4
|
|
self.prior_only = prior_only
|
|
self.generator = OcclusionAwareGenerator(**generator)
|
|
# self.down = AntiAliasInterpolation2d(3, 0.25)
|
|
|
|
## start iteration from resolution of size//2**5
|
|
channels = {
|
|
size // 32: 512,
|
|
size // 16: 512,
|
|
size // 8: 512,
|
|
size // 4: 256,
|
|
size // 2: 128,
|
|
size: 64,
|
|
# 512: 32,
|
|
}
|
|
self.total_iter = int(math.log(2**5, 2)) + 1
|
|
self.num_iter = int(math.log(2**5, 2)) + 1
|
|
self.basic_res_index = int(math.log(self.h // (size // 32), 2))
|
|
|
|
if not self.prior_only:
|
|
### driving and source structure encoder
|
|
self.kp = Hourglass(**driving_encoder)
|
|
self.kp_img = Hourglass(**source_encoder)
|
|
self.kp_head = nn.Conv2d(
|
|
in_channels=self.kp.out_filters,
|
|
out_channels=dim,
|
|
kernel_size=1,
|
|
padding=0,
|
|
)
|
|
self.kp_img_head = nn.Conv2d(
|
|
in_channels=self.kp_img.out_filters,
|
|
out_channels=dim,
|
|
kernel_size=1,
|
|
padding=0,
|
|
)
|
|
# self._make_position_embedding(64, 64, num_kp, 'sine-full')
|
|
# self.pos_embedding = rearrange(self.pos_embedding, 'b (h w) c -> b c h w', h=self.h, w=self.w)
|
|
self.pos_embedding = nn.Parameter(torch.zeros(1, num_kp, self.h, self.w))
|
|
trunc_normal_(self.pos_embedding, std=0.02)
|
|
### driving and source structure encoder
|
|
|
|
### flow updater
|
|
self.corr_enc = BasicMotionEncoder()
|
|
self.refine = RefineFlow()
|
|
## project different-scale warped features to the same dimension
|
|
self.to_context = nn.ModuleList()
|
|
for i in range(self.num_iter):
|
|
f_channel = channels[(size // 32) * (2**i)]
|
|
self.to_context.append(nn.Conv2d(f_channel, 192, 1, padding=0))
|
|
### flow updater
|
|
|
|
def forward(self, kp_s, kp_d, dense_motion, img, img_full):
|
|
feature = self.generator.encode(img_full)
|
|
if img is None:
|
|
img = self.down(img_full)
|
|
b, _, h, w = img.shape
|
|
|
|
out_warp_f = []
|
|
out_occlusion = []
|
|
out_flow = []
|
|
|
|
out_warp_f_c = []
|
|
out_occlusion_c = []
|
|
|
|
### prior only
|
|
if self.prior_only:
|
|
flow = dense_motion["deformation"]
|
|
occlusion = dense_motion["occlusion"]
|
|
for i in range(self.total_iter):
|
|
if flow.shape[2] != feature[i].shape[2]:
|
|
flow_res = F.interpolate(
|
|
flow.permute(0, 3, 1, 2),
|
|
size=feature[i].shape[2:],
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
occlusion_res = F.interpolate(
|
|
occlusion,
|
|
size=feature[i].shape[2:],
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
else:
|
|
flow_res = flow.permute(0, 3, 1, 2)
|
|
occlusion_res = occlusion
|
|
out_warp_f.append(
|
|
F.grid_sample(feature[i], flow_res.permute(0, 2, 3, 1))
|
|
)
|
|
out_occlusion.append(F.sigmoid(occlusion_res))
|
|
warp_img = F.grid_sample(img_full, flow_res.permute(0, 2, 3, 1))
|
|
out = self.generator.decode(out_warp_f, warp_img, out_occlusion)
|
|
for i in range(len(out_occlusion)):
|
|
out_occlusion[i] = F.interpolate(
|
|
out_occlusion[i],
|
|
size=self.size,
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
occlusion = torch.cat(out_occlusion, dim=3)
|
|
return out, warp_img, occlusion
|
|
### prior only
|
|
|
|
### compute correlation volume at basic resolution
|
|
kp_s = kp2gaussian(kp_s, (h, w), 0.1) + self.pos_embedding
|
|
kp_d = kp2gaussian(kp_d, (h, w), 0.1) + self.pos_embedding
|
|
fe_s = self.kp_img(torch.cat([kp_s, img], dim=1))
|
|
fe_d = self.kp(kp_d)
|
|
k_s = self.kp_img_head(fe_s)
|
|
q_d = self.kp_head(fe_d)
|
|
f_s = rearrange(k_s, "b c h w -> b (h w) c", h=self.h, w=self.w)
|
|
f_d = rearrange(q_d, "b c h w -> b (h w) c", h=self.h, w=self.w)
|
|
corr_volume = torch.einsum("bic,bjc->bij", f_d, f_s) * self.scale
|
|
### compute correlation volume at basic resolution
|
|
|
|
### prior motion initialization
|
|
id_grid = coords_grid(b, self.h, self.w, corr_volume.device)
|
|
init_flow = (self.h - 1) * (
|
|
dense_motion["deformation"].permute(0, 3, 1, 2) + 1
|
|
) / 2.0 - id_grid
|
|
init_occlusion = dense_motion["occlusion"]
|
|
### prior motion initialization
|
|
|
|
# ### NPMR only initialization
|
|
# id_grid = coords_grid(b, self.h, self.w, corr_volume.device)
|
|
# id_grid = rearrange(id_grid, 'b c h w -> b (h w) c', h=self.h, w=self.w)
|
|
# init_flow = torch.einsum('bij,bjc->bic', corr_volume.softmax(-1), id_grid)
|
|
# init_flow = rearrange(init_flow, 'b (h w) c-> b c h w', h=self.h, w=self.w)
|
|
# id_grid = rearrange(id_grid, 'b (h w) c -> b c h w', h=self.h, w=self.w)
|
|
# init_flow = init_flow - id_grid
|
|
# init_occlusion = None
|
|
# ### NPMR only initialization
|
|
|
|
## start iter resolution: size//32 = (size//4) // 8, where size//4 defines the basic flow resolution
|
|
flow = (
|
|
F.interpolate(
|
|
init_flow, scale_factor=1.0 / 8.0, mode="bilinear", align_corners=True
|
|
)
|
|
/ 8.0
|
|
)
|
|
occlusion = (
|
|
F.interpolate(
|
|
init_occlusion,
|
|
scale_factor=1.0 / 8.0,
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
if init_occlusion is not None
|
|
else init_occlusion
|
|
)
|
|
## reshape for pooling the driving image dimension, fitting different resolutions
|
|
corr_volume = rearrange(
|
|
corr_volume, "b (h w) n -> (b n) h w", h=self.h, w=self.w
|
|
).unsqueeze(1)
|
|
|
|
for i in range(self.total_iter):
|
|
## identity grid with current-iter resolution
|
|
id_grid = coords_grid(
|
|
b,
|
|
self.size // 32 * (2**i),
|
|
self.size // 32 * (2**i),
|
|
corr_volume.device,
|
|
)
|
|
flow_sample = flow
|
|
id_grid_sample = id_grid
|
|
|
|
### computing upsampled or downsampled structure correlation volume
|
|
if i < self.basic_res_index:
|
|
corr_volume_res = F.avg_pool2d(
|
|
corr_volume,
|
|
2 ** (self.basic_res_index - i),
|
|
stride=2 ** (self.basic_res_index - i),
|
|
)
|
|
scale = 2 ** (self.basic_res_index - i)
|
|
elif i == self.basic_res_index:
|
|
corr_volume_res = corr_volume
|
|
scale = 1
|
|
elif i > self.basic_res_index:
|
|
## for resolution larger than basic flow resolution, sampling with the basic flow resolution
|
|
corr_volume_res = corr_volume
|
|
scale = 0.5 ** (i - self.basic_res_index)
|
|
flow_sample = (
|
|
F.interpolate(
|
|
flow, size=self.h, mode="bilinear", align_corners=True
|
|
)
|
|
* scale
|
|
)
|
|
id_grid_sample = coords_grid(b, self.h, self.w, corr_volume.device)
|
|
scale = 1
|
|
### computing upsampled or downsampled structure correlation volume
|
|
|
|
if i < self.num_iter:
|
|
## reshape for pooling the source image dimension, thus to obtain correlation pyramids
|
|
corr_volume_res = rearrange(
|
|
corr_volume_res, "(b n) c h w -> (b h w) c n", n=self.h * self.w
|
|
)
|
|
corr_volume_res = rearrange(
|
|
corr_volume_res, "b c (p q) -> b c p q", p=self.h, q=self.w
|
|
)
|
|
## building correlation pyramids
|
|
corr_fn = CorrBlock(corr_volume_res)
|
|
## sampling correlation features from the pyramids
|
|
corr = corr_fn((flow_sample + id_grid_sample) * scale)
|
|
if i > self.basic_res_index:
|
|
## reshape from basic resolution to current resolution
|
|
corr = F.interpolate(
|
|
corr, size=flow.shape[2], mode="bilinear", align_corners=True
|
|
)
|
|
## non-prior-based motion feature
|
|
m_f = self.corr_enc(flow, corr)
|
|
|
|
warp_f = bilinear_sampler(
|
|
feature[i], (flow + id_grid).permute(0, 2, 3, 1)
|
|
)
|
|
warp_f = F.relu(self.to_context[i](warp_f))
|
|
|
|
d_flow, _ = self.refine(m_f, warp_f)
|
|
flow_w = flow + d_flow[:, 0:2, :, :]
|
|
d_occ = d_flow[:, 2:, :, :]
|
|
if occlusion is not None:
|
|
occlusion = occlusion + d_occ
|
|
else:
|
|
occlusion = d_occ
|
|
else:
|
|
flow_w = flow
|
|
|
|
out = bilinear_sampler(feature[i], (flow_w + id_grid).permute(0, 2, 3, 1))
|
|
out_occlusion.append(F.sigmoid(occlusion))
|
|
out_warp_f.append(out)
|
|
|
|
## coarse warping (prior motion flow)
|
|
if i != self.basic_res_index:
|
|
flow_res = F.interpolate(
|
|
dense_motion["deformation"].permute(0, 3, 1, 2),
|
|
size=feature[i].shape[2:],
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
occlusion_res = F.interpolate(
|
|
dense_motion["occlusion"],
|
|
size=feature[i].shape[2:],
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
else:
|
|
flow_res = dense_motion["deformation"].permute(0, 3, 1, 2)
|
|
occlusion_res = dense_motion["occlusion"]
|
|
out_warp_f_c.append(F.grid_sample(feature[i], flow_res.permute(0, 2, 3, 1)))
|
|
out_occlusion_c.append(F.sigmoid(occlusion_res))
|
|
## coarse warping (prior motion flow)
|
|
|
|
### updating flow by directly scaling initial flow
|
|
if i < self.num_iter - 1:
|
|
## scale between basic resolution and updated resolution
|
|
scale = 2 ** (self.basic_res_index - i) / 2.0
|
|
d_f = (
|
|
F.interpolate(
|
|
d_flow[:, 0:2, :, :],
|
|
scale_factor=2,
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
* 2
|
|
)
|
|
## flow refinement of current iteration
|
|
flow = (
|
|
d_f
|
|
+ F.interpolate(
|
|
init_flow,
|
|
size=self.size // 32 * (2 ** (i + 1)),
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
/ scale
|
|
)
|
|
## flow refinement of previous iterations
|
|
if i == 0:
|
|
d_f_pre = d_f
|
|
else:
|
|
flow = (
|
|
flow
|
|
+ F.interpolate(
|
|
d_f_pre, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
* 2
|
|
)
|
|
d_f_pre = (
|
|
d_f
|
|
+ F.interpolate(
|
|
d_f_pre, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
* 2
|
|
)
|
|
## occlusion refinement
|
|
d_occ = F.interpolate(
|
|
d_occ, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
occlusion = d_occ + F.interpolate(
|
|
init_occlusion,
|
|
size=self.size // 32 * (2 ** (i + 1)),
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
if i == 0:
|
|
d_occ_pre = d_occ
|
|
else:
|
|
occlusion = occlusion + F.interpolate(
|
|
d_occ_pre, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
d_occ_pre = d_occ + F.interpolate(
|
|
d_occ_pre, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
## for num_iter < total_iter
|
|
# elif i < self.total_iter-1:
|
|
# ## without refinement
|
|
# flow = F.interpolate(flow, scale_factor=2, mode='bilinear',align_corners=True) * 2
|
|
# occlusion = F.interpolate(occlusion, scale_factor=2, mode='bilinear',align_corners=True)
|
|
|
|
warp_img = bilinear_sampler(img_full, (flow + id_grid).permute(0, 2, 3, 1))
|
|
out = self.generator.decode(
|
|
out_warp_f, warp_img, out_occlusion, out_warp_f_c, out_occlusion_c
|
|
)
|
|
|
|
out_occlusion.append(F.sigmoid(init_occlusion))
|
|
out_occlusion_vis = []
|
|
for i in range(len(out_occlusion)):
|
|
out_occlusion_vis.append(
|
|
F.interpolate(
|
|
out_occlusion[i],
|
|
size=self.size,
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
)
|
|
occlusion = torch.cat(out_occlusion_vis, dim=3)
|
|
|
|
return out, warp_img, occlusion
|