296 lines
9.9 KiB
Python
296 lines
9.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .refine import Unet
|
|
from .warplayer import warp
|
|
|
|
|
|
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
|
return nn.Sequential(
|
|
nn.Conv2d(
|
|
in_planes,
|
|
out_planes,
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
bias=True,
|
|
),
|
|
nn.PReLU(out_planes),
|
|
)
|
|
|
|
|
|
class Head(nn.Module):
|
|
def __init__(self, in_planes, scale, c, in_else=17):
|
|
super(Head, self).__init__()
|
|
self.upsample = nn.Sequential(nn.PixelShuffle(2), nn.PixelShuffle(2))
|
|
self.scale = scale
|
|
self.conv = nn.Sequential(
|
|
conv(in_planes * 2 // (4 * 4) + in_else, c),
|
|
conv(c, c),
|
|
conv(c, 5),
|
|
)
|
|
|
|
def forward(self, motion_feature, x, flow):
|
|
motion_feature = self.upsample(motion_feature)
|
|
if self.scale != 4:
|
|
x = F.interpolate(
|
|
x, scale_factor=4.0 / self.scale, mode="bilinear", align_corners=False
|
|
)
|
|
if flow is not None:
|
|
if self.scale != 4:
|
|
flow = (
|
|
F.interpolate(
|
|
flow,
|
|
scale_factor=4.0 / self.scale,
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
* 4.0
|
|
/ self.scale
|
|
)
|
|
x = torch.cat((x, flow), 1)
|
|
x = self.conv(torch.cat([motion_feature, x], 1))
|
|
if self.scale != 4:
|
|
x = F.interpolate(
|
|
x, scale_factor=self.scale // 4, mode="bilinear", align_corners=False
|
|
)
|
|
flow = x[:, :4] * (self.scale // 4)
|
|
else:
|
|
flow = x[:, :4]
|
|
mask = x[:, 4:5]
|
|
return flow, mask
|
|
|
|
|
|
class IFBlock(nn.Module):
|
|
def __init__(self, in_planes, c, scale):
|
|
super(IFBlock, self).__init__()
|
|
self.conv0 = nn.Sequential(
|
|
conv(in_planes, c // 2, 3, 2, 1),
|
|
conv(c // 2, c, 3, 2, 1),
|
|
)
|
|
self.convblock = nn.Sequential(
|
|
conv(c, c),
|
|
conv(c, c),
|
|
conv(c, c),
|
|
conv(c, c),
|
|
conv(c, c),
|
|
conv(c, c),
|
|
conv(c, c),
|
|
conv(c, c),
|
|
)
|
|
self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
|
|
self.scale = scale
|
|
|
|
def forward(self, x, flow):
|
|
scale = self.scale
|
|
if scale != 1:
|
|
x = F.interpolate(
|
|
x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
|
|
)
|
|
flow = (
|
|
F.interpolate(
|
|
flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
|
|
)
|
|
* 1.0
|
|
/ scale
|
|
)
|
|
x = torch.cat((x, flow), 1)
|
|
x = self.conv0(x)
|
|
x = self.convblock(x) + x
|
|
tmp = self.lastconv(x)
|
|
tmp = F.interpolate(
|
|
tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False
|
|
)
|
|
flow = tmp[:, :4] * scale * 2
|
|
mask = tmp[:, 4:5]
|
|
return flow, mask
|
|
|
|
|
|
class MultiScaleFlow(nn.Module):
|
|
def __init__(self, backbone, **kargs):
|
|
super(MultiScaleFlow, self).__init__()
|
|
self.flow_num_stage = len(kargs["hidden_dims"])
|
|
self.local_num = kargs["local_num"]
|
|
self.feature_bone = backbone
|
|
self.block = nn.ModuleList(
|
|
[
|
|
Head(
|
|
kargs["embed_dims"][-1 - i],
|
|
kargs["scales"][-1 - i],
|
|
kargs["hidden_dims"][-1 - i],
|
|
7 if i == 0 else 18,
|
|
)
|
|
for i in range(self.flow_num_stage)
|
|
]
|
|
)
|
|
self.local_block = nn.ModuleList(
|
|
[
|
|
IFBlock(17, c=kargs["local_hidden_dims"], scale=2 - i)
|
|
for i in range(self.local_num)
|
|
]
|
|
)
|
|
self.unet = Unet(kargs["c"] * 2, kargs["M"])
|
|
|
|
def warp_features(self, xs, flow):
|
|
y0 = []
|
|
y1 = []
|
|
B = xs[0].size(0) // 2
|
|
for x in xs:
|
|
y0.append(warp(x[:B], flow[:, 0:2]))
|
|
y1.append(warp(x[B:], flow[:, 2:4]))
|
|
flow = (
|
|
F.interpolate(
|
|
flow,
|
|
scale_factor=0.5,
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
recompute_scale_factor=False,
|
|
)
|
|
* 0.5
|
|
)
|
|
return y0, y1
|
|
|
|
def calculate_flow(self, imgs, timestep, local=False, af=None):
|
|
img0, img1 = imgs[:, :3], imgs[:, 3:6]
|
|
B = img0.size(0)
|
|
flow, mask = None, None
|
|
if af is None:
|
|
af = self.feature_bone(img0, img1)
|
|
timestep = (img0[:, :1].clone() * 0 + 1) * timestep
|
|
for i in range(self.flow_num_stage):
|
|
if flow is not None:
|
|
warped_img0 = warp(img0, flow[:, :2])
|
|
warped_img1 = warp(img1, flow[:, 2:4])
|
|
flow_, mask_ = self.block[i](
|
|
torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
|
|
torch.cat(
|
|
(img0, img1, warped_img0, warped_img1, mask, timestep), 1
|
|
),
|
|
flow,
|
|
)
|
|
flow = flow + flow_
|
|
mask = mask + mask_
|
|
else:
|
|
flow, mask = self.block[i](
|
|
torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
|
|
torch.cat((img0, img1, timestep), 1),
|
|
None,
|
|
)
|
|
|
|
if local:
|
|
for i in range(self.local_num):
|
|
warped_img0 = warp(img0, flow[:, :2])
|
|
warped_img1 = warp(img1, flow[:, 2:4])
|
|
|
|
flow_d, mask_d = self.local_block[i](
|
|
torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow
|
|
)
|
|
flow = flow + flow_d
|
|
mask = mask + mask_d
|
|
|
|
return flow, mask
|
|
|
|
def coraseWarp_and_Refine(self, imgs, af, flow, mask):
|
|
img0, img1 = imgs[:, :3], imgs[:, 3:6]
|
|
warped_img0 = warp(img0, flow[:, :2])
|
|
warped_img1 = warp(img1, flow[:, 2:4])
|
|
c0, c1 = self.warp_features(af, flow)
|
|
tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
|
|
res = tmp[:, :3] * 2 - 1
|
|
mask_ = torch.sigmoid(mask)
|
|
merged = warped_img0 * mask_ + warped_img1 * (1 - mask_)
|
|
pred = torch.clamp(merged + res, 0, 1)
|
|
return pred
|
|
|
|
def forward(self, x, local=False, timestep=0.5, scale=0):
|
|
if scale > 0:
|
|
x_o = x
|
|
x = F.interpolate(
|
|
x, scale_factor=scale, mode="bilinear", align_corners=False
|
|
)
|
|
|
|
img0, img1 = x[:, :3], x[:, 3:6]
|
|
B = x.size(0)
|
|
flow_list = []
|
|
merged = []
|
|
mask_list = []
|
|
warped_img0 = img0
|
|
warped_img1 = img1
|
|
flow = None
|
|
af = self.feature_bone(img0, img1)
|
|
timestep = (x[:, :1].clone() * 0 + 1) * (
|
|
timestep.float() if not isinstance(timestep, float) else timestep
|
|
)
|
|
for i in range(self.flow_num_stage):
|
|
if flow is not None:
|
|
flow_d, mask_d = self.block[i](
|
|
torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
|
|
torch.cat(
|
|
(img0, img1, warped_img0, warped_img1, mask, timestep), # noqa: F821
|
|
1, # noqa: F821
|
|
),
|
|
flow,
|
|
)
|
|
flow = flow + flow_d
|
|
mask = mask + mask_d # noqa: F821
|
|
else:
|
|
flow, mask = self.block[i](
|
|
torch.cat([af[-1 - i][:B], af[-1 - i][B:]], 1),
|
|
torch.cat((img0, img1, timestep), 1),
|
|
None,
|
|
)
|
|
mask_list.append(torch.sigmoid(mask))
|
|
flow_list.append(flow)
|
|
warped_img0 = warp(img0, flow[:, :2])
|
|
warped_img1 = warp(img1, flow[:, 2:4])
|
|
merged.append(warped_img0 * mask_list[i] + warped_img1 * (1 - mask_list[i]))
|
|
|
|
if scale > 0:
|
|
img0, img1 = x_o[:, :3], x_o[:, 3:6]
|
|
af1 = self.feature_bone(img0, img1)
|
|
scale = img0.shape[3] / flow.shape[3]
|
|
flow = (
|
|
F.interpolate(
|
|
flow, scale_factor=scale, mode="bilinear", align_corners=False
|
|
)
|
|
* scale
|
|
)
|
|
mask = F.interpolate(
|
|
mask, scale_factor=scale, mode="bilinear", align_corners=False
|
|
)
|
|
mask_ = torch.sigmoid(mask)
|
|
warped_img0 = warp(img0, flow[:, :2])
|
|
warped_img1 = warp(img1, flow[:, 2:4])
|
|
merged.append(warped_img0 * mask_ + warped_img1 * (1 - mask_))
|
|
|
|
if local:
|
|
flow_list = []
|
|
merged = []
|
|
mask_list = []
|
|
|
|
for i in range(self.local_num):
|
|
flow_d, mask_d = self.local_block[i](
|
|
torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow
|
|
)
|
|
flow = flow + flow_d
|
|
mask = mask + mask_d
|
|
|
|
mask_list.append(torch.sigmoid(mask))
|
|
flow_list.append(flow)
|
|
warped_img0 = warp(img0, flow[:, :2])
|
|
warped_img1 = warp(img1, flow[:, 2:4])
|
|
merged.append(
|
|
warped_img0 * mask_list[i] + warped_img1 * (1 - mask_list[i])
|
|
)
|
|
|
|
if scale:
|
|
c0, c1 = self.warp_features(af1, flow)
|
|
else:
|
|
c0, c1 = self.warp_features(af, flow)
|
|
tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
|
|
res = tmp[:, :3] * 2 - 1
|
|
pred = torch.clamp(merged[-1] + res, 0, 1)
|
|
return flow_list, mask_list, merged, pred
|