import torch
import torch.nn as nn
import torch.nn.functional as F
from .warplayer import warp_features
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class DecoderBlock(nn.Module):
def __init__(self, in_planes, c=224, out_msgs=0, out_locals=0, block_nums=1, out_masks=1, out_local_flows=32, out_msgs_flows=32, out_feat_flows=0):
super(DecoderBlock, self).__init__()
self.conv0 = nn.Sequential(
nn.Conv2d(in_planes, c, 3, 2, 1),
nn.PReLU(c),
nn.Conv2d(c, c, 3, 2, 1),
nn.PReLU(c),
)
self.convblocks = nn.ModuleList()
for i in range(block_nums):
self.convblocks.append(nn.Sequential(
nn.Conv2d(c, c, 3, 1, 1),
nn.PReLU(c),
nn.Conv2d(c, c, 3, 1, 1),
nn.PReLU(c),
nn.Conv2d(c, c, 3, 1, 1),
nn.PReLU(c),
nn.Conv2d(c, c, 3, 1, 1),
nn.PReLU(c),
nn.Conv2d(c, c, 3, 1, 1),
nn.PReLU(c),
nn.Conv2d(c, c, 3, 1, 1),
nn.PReLU(c),
))
self.out_flows = 2
self.out_msgs = out_msgs
self.out_msgs_flows = out_msgs_flows if out_msgs > 0 else 0
self.out_locals = out_locals
self.out_local_flows = out_local_flows if out_locals > 0 else 0
self.out_masks = out_masks
self.out_feat_flows = out_feat_flows
self.conv_last = nn.Sequential(
nn.ConvTranspose2d(c, c, 4, 2, 1),
nn.PReLU(c),
nn.ConvTranspose2d(c, self.out_flows+self.out_msgs+self.out_msgs_flows +
self.out_locals+self.out_local_flows+self.out_masks+self.out_feat_flows, 4, 2, 1),
)
def forward(self, accumulated_flow, *other):
x = [accumulated_flow]
for each in other:
if each is not None:
assert(accumulated_flow.shape[-1] == each.shape[-1]), "decoder want {}, but get {}".format(
accumulated_flow.shape, each.shape)
x.append(each)
feat = self.conv0(torch.cat(x, dim=1))
for convblock1 in self.convblocks:
feat = convblock1(feat) + feat
feat = self.conv_last(feat)
prev = 0
flow = feat[:, prev:prev+self.out_flows, :, :]
prev += self.out_flows
message = feat[:, prev:prev+self.out_msgs,
:, :] if self.out_msgs > 0 else None
prev += self.out_msgs
message_flow = feat[:, prev:prev + self.out_msgs_flows,
:, :] if self.out_msgs_flows > 0 else None
prev += self.out_msgs_flows
local_message = feat[:, prev:prev + self.out_locals,
:, :] if self.out_locals > 0 else None
prev += self.out_locals
local_message_flow = feat[:, prev:prev+self.out_local_flows,
:, :] if self.out_local_flows > 0 else None
prev += self.out_local_flows
mask = torch.sigmoid(
feat[:, prev:prev+self.out_masks, :, :]) if self.out_masks > 0 else None
prev += self.out_masks
feat_flow = feat[:, prev:prev+self.out_feat_flows,
:, :] if self.out_feat_flows > 0 else None
prev += self.out_feat_flows
return flow, mask, message, message_flow, local_message, local_message_flow, feat_flow
class CINN(nn.Module):
def __init__(self, DIM_SHADER_REFERENCE, target_feature_chns=[512, 256, 128, 64, 64], feature_chns=[2048, 1024, 512, 256, 64], out_msgs_chn=[2048, 1024, 512, 256, 64, 64], out_locals_chn=[2048, 1024, 512, 256, 64, 0], block_num=[1, 1, 1, 1, 1, 2], block_chn_num=[224, 224, 224, 224, 224, 224]):
super(CINN, self).__init__()
self.in_msgs_chn = [0, *out_msgs_chn[:-1]]
self.in_locals_chn = [0, *out_locals_chn[:-1]]
self.decoder_blocks = nn.ModuleList()
self.feed_weighted = True
if self.feed_weighted:
in_planes = 2+2+DIM_SHADER_REFERENCE*2
else:
in_planes = 2+DIM_SHADER_REFERENCE
for each_target_feature_chns, each_feature_chns, each_out_msgs_chn, each_out_locals_chn, each_in_msgs_chn, each_in_locals_chn, each_block_num, each_block_chn_num in zip(target_feature_chns, feature_chns, out_msgs_chn, out_locals_chn, self.in_msgs_chn, self.in_locals_chn, block_num, block_chn_num):
self.decoder_blocks.append(
DecoderBlock(in_planes+each_target_feature_chns+each_feature_chns+each_in_locals_chn+each_in_msgs_chn, c=each_block_chn_num, block_nums=each_block_num, out_msgs=each_out_msgs_chn, out_locals=each_out_locals_chn, out_masks=2+each_out_locals_chn))
for i in range(len(feature_chns), len(out_locals_chn)):
#print("append extra block", i, "msg",
# out_msgs_chn[i], "local", out_locals_chn[i], "block", block_num[i])
self.decoder_blocks.append(
DecoderBlock(in_planes+self.in_msgs_chn[i]+self.in_locals_chn[i], c=block_chn_num[i], block_nums=block_num[i], out_msgs=out_msgs_chn[i], out_locals=out_locals_chn[i], out_masks=2+out_msgs_chn[i], out_feat_flows=0))
def apply_flow(self, mask, message, message_flow, local_message, local_message_flow, x_reference, accumulated_flow, each_x_reference_features=None, each_x_reference_features_flow=None):
if each_x_reference_features is not None:
size_from = each_x_reference_features
else:
size_from = x_reference
f_size = (size_from.shape[2], size_from.shape[3])
accumulated_flow = self.flow_rescale(
accumulated_flow, size_from)
# mask = warp_features(F.interpolate(
# mask, size=f_size, mode="bilinear"), accumulated_flow) if mask is not None else None
mask = F.interpolate(
mask, size=f_size, mode="bilinear") if mask is not None else None
message = F.interpolate(
message, size=f_size, mode="bilinear") if message is not None else None
message_flow = self.flow_rescale(
message_flow, size_from) if message_flow is not None else None
message = warp_features(
message, message_flow) if message_flow is not None else message
local_message = F.interpolate(
local_message, size=f_size, mode="bilinear") if local_message is not None else None
local_message_flow = self.flow_rescale(
local_message_flow, size_from) if local_message_flow is not None else None
local_message = warp_features(
local_message, local_message_flow) if local_message_flow is not None else local_message
warp_x_reference = warp_features(F.interpolate(
x_reference, size=f_size, mode="bilinear"), accumulated_flow)
each_x_reference_features_flow = self.flow_rescale(
each_x_reference_features_flow, size_from) if (each_x_reference_features is not None and each_x_reference_features_flow is not None) else None
warp_each_x_reference_features = warp_features(
each_x_reference_features, each_x_reference_features_flow) if each_x_reference_features_flow is not None else each_x_reference_features
return mask, message, local_message, warp_x_reference, accumulated_flow, warp_each_x_reference_features, each_x_reference_features_flow
def forward(self, x_target_features=[], x_reference=None, x_reference_features=[]):
y_flow = []
y_feat_flow = []
y_local_message = []
y_warp_x_reference = []
y_warp_x_reference_features = []
y_weighted_flow = []
y_weighted_mask = []
y_weighted_message = []
y_weighted_x_reference = []
y_weighted_x_reference_features = []
for pyrlevel, ifblock in enumerate(self.decoder_blocks):
stacked_wref = []
stacked_feat = []
stacked_anci = []
stacked_flow = []
stacked_mask = []
stacked_mesg = []
stacked_locm = []
stacked_feat_flow = []
for view_id in range(x_reference.shape[1]): # NMCHW
if pyrlevel == 0:
# create from zero flow
feat_ev = x_reference_features[pyrlevel][:,
view_id, :, :, :] if pyrlevel < len(x_reference_features) else None
accumulated_flow = torch.zeros_like(
feat_ev[:, :2, :, :]).to(device)
accumulated_feat_flow = torch.zeros_like(
feat_ev[:, :32, :, :]).to(device)
# domestic inputs
warp_x_reference = F.interpolate(x_reference[:, view_id, :, :, :], size=(
feat_ev.shape[-2], feat_ev.shape[-1]), mode="bilinear")
warp_x_reference_features = feat_ev
local_message = None
# federated inputs
weighted_flow = accumulated_flow if self.feed_weighted else None
weighted_wref = warp_x_reference if self.feed_weighted else None
weighted_message = None
else:
# resume from last layer
accumulated_flow = y_flow[-1][:, view_id, :, :, :]
accumulated_feat_flow = y_feat_flow[-1][:,
view_id, :, :, :] if y_feat_flow[-1] is not None else None
# domestic inputs
warp_x_reference = y_warp_x_reference[-1][:,
view_id, :, :, :]
warp_x_reference_features = y_warp_x_reference_features[-1][:,
view_id, :, :, :] if y_warp_x_reference_features[-1] is not None else None
local_message = y_local_message[-1][:, view_id, :,
:, :] if len(y_local_message) > 0 else None
# federated inputs
weighted_flow = y_weighted_flow[-1] if self.feed_weighted else None
weighted_wref = y_weighted_x_reference[-1] if self.feed_weighted else None
weighted_message = y_weighted_message[-1] if len(
y_weighted_message) > 0 else None
scaled_x_target = x_target_features[pyrlevel][:, :, :, :].detach() if pyrlevel < len(
x_target_features) else None
# compute flow
residual_flow, mask, message, message_flow, local_message, local_message_flow, residual_feat_flow = ifblock(
accumulated_flow, scaled_x_target, warp_x_reference, warp_x_reference_features, weighted_flow, weighted_wref, weighted_message, local_message)
accumulated_flow = residual_flow + accumulated_flow
accumulated_feat_flow = accumulated_flow
feat_ev = x_reference_features[pyrlevel+1][:,
view_id, :, :, :] if pyrlevel+1 < len(x_reference_features) else None
mask, message, local_message, warp_x_reference, accumulated_flow, warp_x_reference_features, accumulated_feat_flow = self.apply_flow(
mask, message, message_flow, local_message, local_message_flow, x_reference[:, view_id, :, :, :], accumulated_flow, feat_ev, accumulated_feat_flow)
stacked_flow.append(accumulated_flow)
if accumulated_feat_flow is not None:
stacked_feat_flow.append(accumulated_feat_flow)
stacked_mask.append(mask)
if message is not None:
stacked_mesg.append(message)
if local_message is not None:
stacked_locm.append(local_message)
stacked_wref.append(warp_x_reference)
if warp_x_reference_features is not None:
stacked_feat.append(warp_x_reference_features)
stacked_flow = torch.stack(stacked_flow, dim=1) # M*NCHW -> NMCHW
stacked_feat_flow = torch.stack(stacked_feat_flow, dim=1) if len(
stacked_feat_flow) > 0 else None
stacked_mask = torch.stack(
stacked_mask, dim=1)
stacked_mesg = torch.stack(stacked_mesg, dim=1) if len(
stacked_mesg) > 0 else None
stacked_locm = torch.stack(stacked_locm, dim=1) if len(
stacked_locm) > 0 else None
stacked_wref = torch.stack(stacked_wref, dim=1)
stacked_feat = torch.stack(stacked_feat, dim=1) if len(
stacked_feat) > 0 else None
stacked_anci = torch.stack(stacked_anci, dim=1) if len(
stacked_anci) > 0 else None
y_flow.append(stacked_flow)
y_feat_flow.append(stacked_feat_flow)
y_warp_x_reference.append(stacked_wref)
y_warp_x_reference_features.append(stacked_feat)
# compute normalized confidence
stacked_contrib = torch.nn.functional.softmax(stacked_mask, dim=1)
# torch.sum to remove temp dimension M from NMCHW --> NCHW
weighted_flow = torch.sum(
stacked_mask[:, :, 0:1, :, :] * stacked_contrib[:, :, 0:1, :, :] * stacked_flow, dim=1)
weighted_mask = torch.sum(
stacked_contrib[:, :, 0:1, :, :] * stacked_mask[:, :, 0:1, :, :], dim=1)
weighted_wref = torch.sum(
stacked_mask[:, :, 0:1, :, :] * stacked_contrib[:, :, 0:1, :, :] * stacked_wref, dim=1) if stacked_wref is not None else None
weighted_feat = torch.sum(
stacked_mask[:, :, 1:2, :, :] * stacked_contrib[:, :, 1:2, :, :] * stacked_feat, dim=1) if stacked_feat is not None else None
weighted_mesg = torch.sum(
stacked_mask[:, :, 2:, :, :] * stacked_contrib[:, :, 2:, :, :] * stacked_mesg, dim=1) if stacked_mesg is not None else None
y_weighted_flow.append(weighted_flow)
y_weighted_mask.append(weighted_mask)
if weighted_mesg is not None:
y_weighted_message.append(weighted_mesg)
if stacked_locm is not None:
y_local_message.append(stacked_locm)
y_weighted_message.append(weighted_mesg)
y_weighted_x_reference.append(weighted_wref)
y_weighted_x_reference_features.append(weighted_feat)
if weighted_feat is not None:
y_weighted_x_reference_features.append(weighted_feat)
return {
"y_last_remote_features": [weighted_mesg],
}
def flow_rescale(self, prev_flow, each_x_reference_features):
if prev_flow is None:
prev_flow = torch.zeros_like(
each_x_reference_features[:, :2]).to(device)
else:
up_scale_factor = each_x_reference_features.shape[-1] / \
prev_flow.shape[-1]
if up_scale_factor != 1:
prev_flow = F.interpolate(prev_flow, scale_factor=up_scale_factor, mode="bilinear",
align_corners=False, recompute_scale_factor=False) * up_scale_factor
return prev_flow