399 lines
12 KiB
Python
399 lines
12 KiB
Python
import copy
|
|
from typing import Optional
|
|
|
|
import paddle
|
|
import paddle.nn as nn
|
|
import paddle.nn.functional as F
|
|
|
|
from .extractor import BasicEncoder
|
|
from .position_encoding import build_position_encoding
|
|
from .weight_init import weight_init_
|
|
|
|
|
|
class attnLayer(nn.Layer):
|
|
def __init__(
|
|
self,
|
|
d_model,
|
|
nhead=8,
|
|
dim_feedforward=2048,
|
|
dropout=0.1,
|
|
activation="relu",
|
|
normalize_before=False,
|
|
):
|
|
super().__init__()
|
|
|
|
self.self_attn = nn.MultiHeadAttention(d_model, nhead, dropout=dropout)
|
|
self.multihead_attn_list = nn.LayerList(
|
|
[
|
|
copy.deepcopy(nn.MultiHeadAttention(d_model, nhead, dropout=dropout))
|
|
for i in range(2)
|
|
]
|
|
)
|
|
|
|
self.linear1 = nn.Linear(d_model, dim_feedforward)
|
|
self.dropout = nn.Dropout(p=dropout)
|
|
self.linear2 = nn.Linear(dim_feedforward, d_model)
|
|
|
|
self.norm1 = nn.LayerNorm(d_model)
|
|
self.norm2_list = nn.LayerList(
|
|
[copy.deepcopy(nn.LayerNorm(d_model)) for i in range(2)]
|
|
)
|
|
|
|
self.norm3 = nn.LayerNorm(d_model)
|
|
self.dropout1 = nn.Dropout(p=dropout)
|
|
self.dropout2_list = nn.LayerList(
|
|
[copy.deepcopy(nn.Dropout(p=dropout)) for i in range(2)]
|
|
)
|
|
self.dropout3 = nn.Dropout(p=dropout)
|
|
|
|
self.activation = _get_activation_fn(activation)
|
|
self.normalize_before = normalize_before
|
|
|
|
def with_pos_embed(self, tensor, pos: Optional[paddle.Tensor]):
|
|
return tensor if pos is None else tensor + pos
|
|
|
|
def forward_post(
|
|
self,
|
|
tgt,
|
|
memory_list,
|
|
tgt_mask=None,
|
|
memory_mask=None,
|
|
pos=None,
|
|
memory_pos=None,
|
|
):
|
|
q = k = self.with_pos_embed(tgt, pos)
|
|
tgt2 = self.self_attn(
|
|
q.transpose((1, 0, 2)),
|
|
k.transpose((1, 0, 2)),
|
|
value=tgt.transpose((1, 0, 2)),
|
|
attn_mask=tgt_mask,
|
|
)
|
|
tgt2 = tgt2.transpose((1, 0, 2))
|
|
tgt = tgt + self.dropout1(tgt2)
|
|
tgt = self.norm1(tgt)
|
|
|
|
for memory, multihead_attn, norm2, dropout2, m_pos in zip(
|
|
memory_list,
|
|
self.multihead_attn_list,
|
|
self.norm2_list,
|
|
self.dropout2_list,
|
|
memory_pos,
|
|
):
|
|
tgt2 = multihead_attn(
|
|
query=self.with_pos_embed(tgt, pos).transpose((1, 0, 2)),
|
|
key=self.with_pos_embed(memory, m_pos).transpose((1, 0, 2)),
|
|
value=memory.transpose((1, 0, 2)),
|
|
attn_mask=memory_mask,
|
|
).transpose((1, 0, 2))
|
|
|
|
tgt = tgt + dropout2(tgt2)
|
|
tgt = norm2(tgt)
|
|
|
|
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
|
|
tgt = tgt + self.dropout3(tgt2)
|
|
tgt = self.norm3(tgt)
|
|
|
|
return tgt
|
|
|
|
def forward_pre(
|
|
self,
|
|
tgt,
|
|
memory,
|
|
tgt_mask=None,
|
|
memory_mask=None,
|
|
pos=None,
|
|
memory_pos=None,
|
|
):
|
|
tgt2 = self.norm1(tgt)
|
|
|
|
q = k = self.with_pos_embed(tgt2, pos)
|
|
tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask)
|
|
tgt = tgt + self.dropout1(tgt2)
|
|
tgt2 = self.norm2(tgt)
|
|
|
|
tgt2 = self.multihead_attn(
|
|
query=self.with_pos_embed(tgt2, pos),
|
|
key=self.with_pos_embed(memory, memory_pos),
|
|
value=memory,
|
|
attn_mask=memory_mask,
|
|
)
|
|
tgt = tgt + self.dropout2(tgt2)
|
|
tgt2 = self.norm3(tgt)
|
|
|
|
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
|
|
tgt = tgt + self.dropout3(tgt2)
|
|
|
|
return tgt
|
|
|
|
def forward(
|
|
self,
|
|
tgt,
|
|
memory_list,
|
|
tgt_mask=None,
|
|
memory_mask=None,
|
|
pos=None,
|
|
memory_pos=None,
|
|
):
|
|
if self.normalize_before:
|
|
return self.forward_pre(
|
|
tgt,
|
|
memory_list,
|
|
tgt_mask,
|
|
memory_mask,
|
|
pos,
|
|
memory_pos,
|
|
)
|
|
return self.forward_post(
|
|
tgt,
|
|
memory_list,
|
|
tgt_mask,
|
|
memory_mask,
|
|
pos,
|
|
memory_pos,
|
|
)
|
|
|
|
|
|
def _get_clones(module, N):
|
|
return nn.LayerList([copy.deepcopy(module) for i in range(N)])
|
|
|
|
|
|
def _get_activation_fn(activation):
|
|
"""Return an activation function given a string"""
|
|
if activation == "relu":
|
|
return F.relu
|
|
if activation == "gelu":
|
|
return F.gelu
|
|
if activation == "glu":
|
|
return F.glu
|
|
raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
|
|
|
|
|
|
class TransDecoder(nn.Layer):
|
|
def __init__(self, num_attn_layers: int, hidden_dim: int = 128):
|
|
super(TransDecoder, self).__init__()
|
|
|
|
attn_layer = attnLayer(hidden_dim)
|
|
self.layers = _get_clones(attn_layer, num_attn_layers)
|
|
self.position_embedding = build_position_encoding(hidden_dim)
|
|
|
|
def forward(self, image: paddle.Tensor, query_embed: paddle.Tensor):
|
|
pos = self.position_embedding(
|
|
paddle.ones([image.shape[0], image.shape[2], image.shape[3]], dtype="bool")
|
|
)
|
|
|
|
b, c, h, w = image.shape
|
|
|
|
image = image.flatten(2).transpose(perm=[2, 0, 1])
|
|
pos = pos.flatten(2).transpose(perm=[2, 0, 1])
|
|
|
|
for layer in self.layers:
|
|
query_embed = layer(query_embed, [image], pos=pos, memory_pos=[pos, pos])
|
|
|
|
query_embed = query_embed.transpose(perm=[1, 2, 0]).reshape([b, c, h, w])
|
|
|
|
return query_embed
|
|
|
|
|
|
class TransEncoder(nn.Layer):
|
|
def __init__(self, num_attn_layers: int, hidden_dim: int = 128):
|
|
super(TransEncoder, self).__init__()
|
|
|
|
attn_layer = attnLayer(hidden_dim)
|
|
self.layers = _get_clones(attn_layer, num_attn_layers)
|
|
self.position_embedding = build_position_encoding(hidden_dim)
|
|
|
|
def forward(self, image: paddle.Tensor):
|
|
pos = self.position_embedding(
|
|
paddle.ones([image.shape[0], image.shape[2], image.shape[3]], dtype="bool")
|
|
)
|
|
|
|
b, c, h, w = image.shape
|
|
|
|
image = image.flatten(2).transpose(perm=[2, 0, 1])
|
|
pos = pos.flatten(2).transpose(perm=[2, 0, 1])
|
|
|
|
for layer in self.layers:
|
|
image = layer(image, [image], pos=pos, memory_pos=[pos, pos])
|
|
|
|
image = image.transpose(perm=[1, 2, 0]).reshape([b, c, h, w])
|
|
|
|
return image
|
|
|
|
|
|
class FlowHead(nn.Layer):
|
|
def __init__(self, input_dim=128, hidden_dim=256):
|
|
super(FlowHead, self).__init__()
|
|
|
|
self.conv1 = nn.Conv2D(input_dim, hidden_dim, 3, padding=1)
|
|
self.conv2 = nn.Conv2D(hidden_dim, 2, 3, padding=1)
|
|
self.relu = nn.ReLU()
|
|
|
|
def forward(self, x):
|
|
return self.conv2(self.relu(self.conv1(x)))
|
|
|
|
|
|
class UpdateBlock(nn.Layer):
|
|
def __init__(self, hidden_dim: int = 128):
|
|
super(UpdateBlock, self).__init__()
|
|
|
|
self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
|
|
|
|
self.mask = nn.Sequential(
|
|
nn.Conv2D(hidden_dim, 256, 3, padding=1),
|
|
nn.ReLU(),
|
|
nn.Conv2D(256, 64 * 9, 1, padding=0),
|
|
)
|
|
|
|
def forward(self, image, coords):
|
|
mask = 0.25 * self.mask(image)
|
|
dflow = self.flow_head(image)
|
|
coords = coords + dflow
|
|
return mask, coords
|
|
|
|
|
|
def coords_grid(batch, ht, wd):
|
|
coords = paddle.meshgrid(paddle.arange(end=ht), paddle.arange(end=wd))
|
|
coords = paddle.stack(coords[::-1], axis=0).astype(dtype="float32")
|
|
return coords[None].tile([batch, 1, 1, 1])
|
|
|
|
|
|
def upflow8(flow, mode="bilinear"):
|
|
new_size = 8 * flow.shape[2], 8 * flow.shape[3]
|
|
return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
|
|
|
|
|
class OverlapPatchEmbed(nn.Layer):
|
|
"""Image to Patch Embedding"""
|
|
|
|
def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
|
|
super().__init__()
|
|
|
|
img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
|
|
patch_size = (
|
|
patch_size if isinstance(patch_size, tuple) else (patch_size, patch_size)
|
|
)
|
|
|
|
self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
|
|
self.num_patches = self.H * self.W
|
|
|
|
self.proj = nn.Conv2D(
|
|
in_chans,
|
|
embed_dim,
|
|
patch_size,
|
|
stride=stride,
|
|
padding=(patch_size[0] // 2, patch_size[1] // 2),
|
|
)
|
|
self.norm = nn.LayerNorm(embed_dim)
|
|
|
|
self.apply(self._init_weights)
|
|
|
|
def _init_weights(self, m):
|
|
if isinstance(m, nn.Linear):
|
|
weight_init_(m, "trunc_normal_", std=0.02)
|
|
elif isinstance(m, nn.LayerNorm):
|
|
weight_init_(m, "Constant", value=1.0)
|
|
elif isinstance(m, nn.Conv2D):
|
|
weight_init_(
|
|
m.weight, "kaiming_normal_", mode="fan_out", nonlinearity="relu"
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.proj(x)
|
|
_, _, H, W = x.shape
|
|
x = x.flatten(2)
|
|
|
|
perm = list(range(x.ndim))
|
|
perm[1] = 2
|
|
perm[2] = 1
|
|
x = x.transpose(perm=perm)
|
|
|
|
x = self.norm(x)
|
|
|
|
return x, H, W
|
|
|
|
|
|
class GeoTr(nn.Layer):
|
|
def __init__(self):
|
|
super(GeoTr, self).__init__()
|
|
|
|
self.hidden_dim = hdim = 256
|
|
|
|
self.fnet = BasicEncoder(output_dim=hdim, norm_fn="instance")
|
|
|
|
self.encoder_block = [("encoder_block" + str(i)) for i in range(3)]
|
|
for i in self.encoder_block:
|
|
self.__setattr__(i, TransEncoder(2, hidden_dim=hdim))
|
|
|
|
self.down_layer = [("down_layer" + str(i)) for i in range(2)]
|
|
for i in self.down_layer:
|
|
self.__setattr__(i, nn.Conv2D(256, 256, 3, stride=2, padding=1))
|
|
|
|
self.decoder_block = [("decoder_block" + str(i)) for i in range(3)]
|
|
for i in self.decoder_block:
|
|
self.__setattr__(i, TransDecoder(2, hidden_dim=hdim))
|
|
|
|
self.up_layer = [("up_layer" + str(i)) for i in range(2)]
|
|
for i in self.up_layer:
|
|
self.__setattr__(
|
|
i, nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
|
|
)
|
|
|
|
self.query_embed = nn.Embedding(81, self.hidden_dim)
|
|
|
|
self.update_block = UpdateBlock(self.hidden_dim)
|
|
|
|
def initialize_flow(self, img):
|
|
N, _, H, W = img.shape
|
|
coodslar = coords_grid(N, H, W)
|
|
coords0 = coords_grid(N, H // 8, W // 8)
|
|
coords1 = coords_grid(N, H // 8, W // 8)
|
|
return coodslar, coords0, coords1
|
|
|
|
def upsample_flow(self, flow, mask):
|
|
N, _, H, W = flow.shape
|
|
|
|
mask = mask.reshape([N, 1, 9, 8, 8, H, W])
|
|
mask = F.softmax(mask, axis=2)
|
|
|
|
up_flow = F.unfold(8 * flow, [3, 3], paddings=1)
|
|
up_flow = up_flow.reshape([N, 2, 9, 1, 1, H, W])
|
|
|
|
up_flow = paddle.sum(mask * up_flow, axis=2)
|
|
up_flow = up_flow.transpose(perm=[0, 1, 4, 2, 5, 3])
|
|
|
|
return up_flow.reshape([N, 2, 8 * H, 8 * W])
|
|
|
|
def forward(self, image):
|
|
fmap = self.fnet(image)
|
|
fmap = F.relu(fmap)
|
|
|
|
fmap1 = self.__getattr__(self.encoder_block[0])(fmap)
|
|
fmap1d = self.__getattr__(self.down_layer[0])(fmap1)
|
|
fmap2 = self.__getattr__(self.encoder_block[1])(fmap1d)
|
|
fmap2d = self.__getattr__(self.down_layer[1])(fmap2)
|
|
fmap3 = self.__getattr__(self.encoder_block[2])(fmap2d)
|
|
|
|
query_embed0 = self.query_embed.weight.unsqueeze(1).tile([1, fmap3.shape[0], 1])
|
|
|
|
fmap3d_ = self.__getattr__(self.decoder_block[0])(fmap3, query_embed0)
|
|
fmap3du_ = (
|
|
self.__getattr__(self.up_layer[0])(fmap3d_)
|
|
.flatten(2)
|
|
.transpose(perm=[2, 0, 1])
|
|
)
|
|
fmap2d_ = self.__getattr__(self.decoder_block[1])(fmap2, fmap3du_)
|
|
fmap2du_ = (
|
|
self.__getattr__(self.up_layer[1])(fmap2d_)
|
|
.flatten(2)
|
|
.transpose(perm=[2, 0, 1])
|
|
)
|
|
fmap_out = self.__getattr__(self.decoder_block[2])(fmap1, fmap2du_)
|
|
|
|
coodslar, coords0, coords1 = self.initialize_flow(image)
|
|
coords1 = coords1.detach()
|
|
mask, coords1 = self.update_block(fmap_out, coords1)
|
|
flow_up = self.upsample_flow(coords1 - coords0, mask)
|
|
bm_up = coodslar + flow_up
|
|
|
|
return bm_up
|