mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-02-18 15:04:45 +08:00
345 lines
16 KiB
Python
345 lines
16 KiB
Python
import math
|
|
from functools import reduce
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class SoftSplit(nn.Module):
|
|
def __init__(self, channel, hidden, kernel_size, stride, padding):
|
|
super(SoftSplit, self).__init__()
|
|
self.kernel_size = kernel_size
|
|
self.stride = stride
|
|
self.padding = padding
|
|
self.t2t = nn.Unfold(kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=padding)
|
|
c_in = reduce((lambda x, y: x * y), kernel_size) * channel
|
|
self.embedding = nn.Linear(c_in, hidden)
|
|
|
|
def forward(self, x, b, output_size):
|
|
f_h = int((output_size[0] + 2 * self.padding[0] -
|
|
(self.kernel_size[0] - 1) - 1) / self.stride[0] + 1)
|
|
f_w = int((output_size[1] + 2 * self.padding[1] -
|
|
(self.kernel_size[1] - 1) - 1) / self.stride[1] + 1)
|
|
|
|
feat = self.t2t(x)
|
|
feat = feat.permute(0, 2, 1)
|
|
# feat shape [b*t, num_vec, ks*ks*c]
|
|
feat = self.embedding(feat)
|
|
# feat shape after embedding [b, t*num_vec, hidden]
|
|
feat = feat.view(b, -1, f_h, f_w, feat.size(2))
|
|
return feat
|
|
|
|
|
|
class SoftComp(nn.Module):
|
|
def __init__(self, channel, hidden, kernel_size, stride, padding):
|
|
super(SoftComp, self).__init__()
|
|
self.relu = nn.LeakyReLU(0.2, inplace=True)
|
|
c_out = reduce((lambda x, y: x * y), kernel_size) * channel
|
|
self.embedding = nn.Linear(hidden, c_out)
|
|
self.kernel_size = kernel_size
|
|
self.stride = stride
|
|
self.padding = padding
|
|
self.bias_conv = nn.Conv2d(channel,
|
|
channel,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1)
|
|
|
|
def forward(self, x, t, output_size):
|
|
b_, _, _, _, c_ = x.shape
|
|
x = x.view(b_, -1, c_)
|
|
feat = self.embedding(x)
|
|
b, _, c = feat.size()
|
|
feat = feat.view(b * t, -1, c).permute(0, 2, 1)
|
|
feat = F.fold(feat,
|
|
output_size=output_size,
|
|
kernel_size=self.kernel_size,
|
|
stride=self.stride,
|
|
padding=self.padding)
|
|
feat = self.bias_conv(feat)
|
|
return feat
|
|
|
|
|
|
class FusionFeedForward(nn.Module):
|
|
def __init__(self, dim, hidden_dim=1960, t2t_params=None):
|
|
super(FusionFeedForward, self).__init__()
|
|
# We set hidden_dim as a default to 1960
|
|
self.fc1 = nn.Sequential(nn.Linear(dim, hidden_dim))
|
|
self.fc2 = nn.Sequential(nn.GELU(), nn.Linear(hidden_dim, dim))
|
|
assert t2t_params is not None
|
|
self.t2t_params = t2t_params
|
|
self.kernel_shape = reduce((lambda x, y: x * y), t2t_params['kernel_size']) # 49
|
|
|
|
def forward(self, x, output_size):
|
|
n_vecs = 1
|
|
for i, d in enumerate(self.t2t_params['kernel_size']):
|
|
n_vecs *= int((output_size[i] + 2 * self.t2t_params['padding'][i] -
|
|
(d - 1) - 1) / self.t2t_params['stride'][i] + 1)
|
|
|
|
x = self.fc1(x)
|
|
b, n, c = x.size()
|
|
normalizer = x.new_ones(b, n, self.kernel_shape).view(-1, n_vecs, self.kernel_shape).permute(0, 2, 1)
|
|
normalizer = F.fold(normalizer,
|
|
output_size=output_size,
|
|
kernel_size=self.t2t_params['kernel_size'],
|
|
padding=self.t2t_params['padding'],
|
|
stride=self.t2t_params['stride'])
|
|
|
|
x = F.fold(x.view(-1, n_vecs, c).permute(0, 2, 1),
|
|
output_size=output_size,
|
|
kernel_size=self.t2t_params['kernel_size'],
|
|
padding=self.t2t_params['padding'],
|
|
stride=self.t2t_params['stride'])
|
|
|
|
x = F.unfold(x / normalizer,
|
|
kernel_size=self.t2t_params['kernel_size'],
|
|
padding=self.t2t_params['padding'],
|
|
stride=self.t2t_params['stride']).permute(
|
|
0, 2, 1).contiguous().view(b, n, c)
|
|
x = self.fc2(x)
|
|
return x
|
|
|
|
|
|
def window_partition(x, window_size, n_head):
|
|
"""
|
|
Args:
|
|
x: shape is (B, T, H, W, C)
|
|
window_size (tuple[int]): window size
|
|
Returns:
|
|
windows: (B, num_windows_h, num_windows_w, n_head, T, window_size, window_size, C//n_head)
|
|
"""
|
|
B, T, H, W, C = x.shape
|
|
x = x.view(B, T, H // window_size[0], window_size[0], W // window_size[1], window_size[1], n_head, C//n_head)
|
|
windows = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous()
|
|
return windows
|
|
|
|
class SparseWindowAttention(nn.Module):
|
|
def __init__(self, dim, n_head, window_size, pool_size=(4,4), qkv_bias=True, attn_drop=0., proj_drop=0.,
|
|
pooling_token=True):
|
|
super().__init__()
|
|
assert dim % n_head == 0
|
|
# key, query, value projections for all heads
|
|
self.key = nn.Linear(dim, dim, qkv_bias)
|
|
self.query = nn.Linear(dim, dim, qkv_bias)
|
|
self.value = nn.Linear(dim, dim, qkv_bias)
|
|
# regularization
|
|
self.attn_drop = nn.Dropout(attn_drop)
|
|
self.proj_drop = nn.Dropout(proj_drop)
|
|
# output projection
|
|
self.proj = nn.Linear(dim, dim)
|
|
self.n_head = n_head
|
|
self.window_size = window_size
|
|
self.pooling_token = pooling_token
|
|
if self.pooling_token:
|
|
ks, stride = pool_size, pool_size
|
|
self.pool_layer = nn.Conv2d(dim, dim, kernel_size=ks, stride=stride, padding=(0, 0), groups=dim)
|
|
self.pool_layer.weight.data.fill_(1. / (pool_size[0] * pool_size[1]))
|
|
self.pool_layer.bias.data.fill_(0)
|
|
# self.expand_size = tuple(i // 2 for i in window_size)
|
|
self.expand_size = tuple((i + 1) // 2 for i in window_size)
|
|
|
|
if any(i > 0 for i in self.expand_size):
|
|
# get mask for rolled k and rolled v
|
|
mask_tl = torch.ones(self.window_size[0], self.window_size[1])
|
|
mask_tl[:-self.expand_size[0], :-self.expand_size[1]] = 0
|
|
mask_tr = torch.ones(self.window_size[0], self.window_size[1])
|
|
mask_tr[:-self.expand_size[0], self.expand_size[1]:] = 0
|
|
mask_bl = torch.ones(self.window_size[0], self.window_size[1])
|
|
mask_bl[self.expand_size[0]:, :-self.expand_size[1]] = 0
|
|
mask_br = torch.ones(self.window_size[0], self.window_size[1])
|
|
mask_br[self.expand_size[0]:, self.expand_size[1]:] = 0
|
|
masrool_k = torch.stack((mask_tl, mask_tr, mask_bl, mask_br), 0).flatten(0)
|
|
self.register_buffer("valid_ind_rolled", masrool_k.nonzero(as_tuple=False).view(-1))
|
|
|
|
self.max_pool = nn.MaxPool2d(window_size, window_size, (0, 0))
|
|
|
|
|
|
def forward(self, x, mask=None, T_ind=None, attn_mask=None):
|
|
b, t, h, w, c = x.shape # 20 36
|
|
w_h, w_w = self.window_size[0], self.window_size[1]
|
|
c_head = c // self.n_head
|
|
n_wh = math.ceil(h / self.window_size[0])
|
|
n_ww = math.ceil(w / self.window_size[1])
|
|
new_h = n_wh * self.window_size[0] # 20
|
|
new_w = n_ww * self.window_size[1] # 36
|
|
pad_r = new_w - w
|
|
pad_b = new_h - h
|
|
# reverse order
|
|
if pad_r > 0 or pad_b > 0:
|
|
x = F.pad(x,(0, 0, 0, pad_r, 0, pad_b, 0, 0), mode='constant', value=0)
|
|
mask = F.pad(mask,(0, 0, 0, pad_r, 0, pad_b, 0, 0), mode='constant', value=0)
|
|
|
|
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
|
|
q = self.query(x)
|
|
k = self.key(x)
|
|
v = self.value(x)
|
|
win_q = window_partition(q.contiguous(), self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head)
|
|
win_k = window_partition(k.contiguous(), self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head)
|
|
win_v = window_partition(v.contiguous(), self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head)
|
|
# roll_k and roll_v
|
|
if any(i > 0 for i in self.expand_size):
|
|
(k_tl, v_tl) = map(lambda a: torch.roll(a, shifts=(-self.expand_size[0], -self.expand_size[1]), dims=(2, 3)), (k, v))
|
|
(k_tr, v_tr) = map(lambda a: torch.roll(a, shifts=(-self.expand_size[0], self.expand_size[1]), dims=(2, 3)), (k, v))
|
|
(k_bl, v_bl) = map(lambda a: torch.roll(a, shifts=(self.expand_size[0], -self.expand_size[1]), dims=(2, 3)), (k, v))
|
|
(k_br, v_br) = map(lambda a: torch.roll(a, shifts=(self.expand_size[0], self.expand_size[1]), dims=(2, 3)), (k, v))
|
|
|
|
(k_tl_windows, k_tr_windows, k_bl_windows, k_br_windows) = map(
|
|
lambda a: window_partition(a, self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head),
|
|
(k_tl, k_tr, k_bl, k_br))
|
|
(v_tl_windows, v_tr_windows, v_bl_windows, v_br_windows) = map(
|
|
lambda a: window_partition(a, self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head),
|
|
(v_tl, v_tr, v_bl, v_br))
|
|
rool_k = torch.cat((k_tl_windows, k_tr_windows, k_bl_windows, k_br_windows), 4).contiguous()
|
|
rool_v = torch.cat((v_tl_windows, v_tr_windows, v_bl_windows, v_br_windows), 4).contiguous() # [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
|
# mask out tokens in current window
|
|
rool_k = rool_k[:, :, :, :, self.valid_ind_rolled]
|
|
rool_v = rool_v[:, :, :, :, self.valid_ind_rolled]
|
|
roll_N = rool_k.shape[4]
|
|
rool_k = rool_k.view(b, n_wh*n_ww, self.n_head, t, roll_N, c // self.n_head)
|
|
rool_v = rool_v.view(b, n_wh*n_ww, self.n_head, t, roll_N, c // self.n_head)
|
|
win_k = torch.cat((win_k, rool_k), dim=4)
|
|
win_v = torch.cat((win_v, rool_v), dim=4)
|
|
else:
|
|
win_k = win_k
|
|
win_v = win_v
|
|
|
|
# pool_k and pool_v
|
|
if self.pooling_token:
|
|
pool_x = self.pool_layer(x.view(b*t, new_h, new_w, c).permute(0,3,1,2))
|
|
_, _, p_h, p_w = pool_x.shape
|
|
pool_x = pool_x.permute(0,2,3,1).view(b, t, p_h, p_w, c)
|
|
# pool_k
|
|
pool_k = self.key(pool_x).unsqueeze(1).repeat(1, n_wh*n_ww, 1, 1, 1, 1) # [b, n_wh*n_ww, t, p_h, p_w, c]
|
|
pool_k = pool_k.view(b, n_wh*n_ww, t, p_h, p_w, self.n_head, c_head).permute(0,1,5,2,3,4,6)
|
|
pool_k = pool_k.contiguous().view(b, n_wh*n_ww, self.n_head, t, p_h*p_w, c_head)
|
|
win_k = torch.cat((win_k, pool_k), dim=4)
|
|
# pool_v
|
|
pool_v = self.value(pool_x).unsqueeze(1).repeat(1, n_wh*n_ww, 1, 1, 1, 1) # [b, n_wh*n_ww, t, p_h, p_w, c]
|
|
pool_v = pool_v.view(b, n_wh*n_ww, t, p_h, p_w, self.n_head, c_head).permute(0,1,5,2,3,4,6)
|
|
pool_v = pool_v.contiguous().view(b, n_wh*n_ww, self.n_head, t, p_h*p_w, c_head)
|
|
win_v = torch.cat((win_v, pool_v), dim=4)
|
|
|
|
# [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
|
out = torch.zeros_like(win_q)
|
|
l_t = mask.size(1)
|
|
|
|
mask = self.max_pool(mask.view(b * l_t, new_h, new_w))
|
|
mask = mask.view(b, l_t, n_wh*n_ww)
|
|
mask = torch.sum(mask, dim=1) # [b, n_wh*n_ww]
|
|
for i in range(win_q.shape[0]):
|
|
### For masked windows
|
|
mask_ind_i = mask[i].nonzero(as_tuple=False).view(-1)
|
|
# mask out quary in current window
|
|
# [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
|
mask_n = len(mask_ind_i)
|
|
if mask_n > 0:
|
|
win_q_t = win_q[i, mask_ind_i].view(mask_n, self.n_head, t*w_h*w_w, c_head)
|
|
win_k_t = win_k[i, mask_ind_i]
|
|
win_v_t = win_v[i, mask_ind_i]
|
|
# mask out key and value
|
|
if T_ind is not None:
|
|
# key [n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
|
win_k_t = win_k_t[:, :, T_ind.view(-1)].view(mask_n, self.n_head, -1, c_head)
|
|
# value
|
|
win_v_t = win_v_t[:, :, T_ind.view(-1)].view(mask_n, self.n_head, -1, c_head)
|
|
else:
|
|
win_k_t = win_k_t.view(n_wh*n_ww, self.n_head, t*w_h*w_w, c_head)
|
|
win_v_t = win_v_t.view(n_wh*n_ww, self.n_head, t*w_h*w_w, c_head)
|
|
|
|
att_t = (win_q_t @ win_k_t.transpose(-2, -1)) * (1.0 / math.sqrt(win_q_t.size(-1)))
|
|
att_t = F.softmax(att_t, dim=-1)
|
|
att_t = self.attn_drop(att_t)
|
|
y_t = att_t @ win_v_t
|
|
|
|
out[i, mask_ind_i] = y_t.view(-1, self.n_head, t, w_h*w_w, c_head)
|
|
|
|
### For unmasked windows
|
|
unmask_ind_i = (mask[i] == 0).nonzero(as_tuple=False).view(-1)
|
|
# mask out quary in current window
|
|
# [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
|
win_q_s = win_q[i, unmask_ind_i]
|
|
win_k_s = win_k[i, unmask_ind_i, :, :, :w_h*w_w]
|
|
win_v_s = win_v[i, unmask_ind_i, :, :, :w_h*w_w]
|
|
|
|
att_s = (win_q_s @ win_k_s.transpose(-2, -1)) * (1.0 / math.sqrt(win_q_s.size(-1)))
|
|
att_s = F.softmax(att_s, dim=-1)
|
|
att_s = self.attn_drop(att_s)
|
|
y_s = att_s @ win_v_s
|
|
out[i, unmask_ind_i] = y_s
|
|
|
|
# re-assemble all head outputs side by side
|
|
out = out.view(b, n_wh, n_ww, self.n_head, t, w_h, w_w, c_head)
|
|
out = out.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous().view(b, t, new_h, new_w, c)
|
|
|
|
|
|
if pad_r > 0 or pad_b > 0:
|
|
out = out[:, :, :h, :w, :]
|
|
|
|
# output projection
|
|
out = self.proj_drop(self.proj(out))
|
|
return out
|
|
|
|
|
|
class TemporalSparseTransformer(nn.Module):
|
|
def __init__(self, dim, n_head, window_size, pool_size,
|
|
norm_layer=nn.LayerNorm, t2t_params=None):
|
|
super().__init__()
|
|
self.window_size = window_size
|
|
self.attention = SparseWindowAttention(dim, n_head, window_size, pool_size)
|
|
self.norm1 = norm_layer(dim)
|
|
self.norm2 = norm_layer(dim)
|
|
self.mlp = FusionFeedForward(dim, t2t_params=t2t_params)
|
|
|
|
def forward(self, x, fold_x_size, mask=None, T_ind=None):
|
|
"""
|
|
Args:
|
|
x: image tokens, shape [B T H W C]
|
|
fold_x_size: fold feature size, shape [60 108]
|
|
mask: mask tokens, shape [B T H W 1]
|
|
Returns:
|
|
out_tokens: shape [B T H W C]
|
|
"""
|
|
B, T, H, W, C = x.shape # 20 36
|
|
|
|
shortcut = x
|
|
x = self.norm1(x)
|
|
att_x = self.attention(x, mask, T_ind)
|
|
|
|
# FFN
|
|
x = shortcut + att_x
|
|
y = self.norm2(x)
|
|
x = x + self.mlp(y.view(B, T * H * W, C), fold_x_size).view(B, T, H, W, C)
|
|
|
|
return x
|
|
|
|
|
|
class TemporalSparseTransformerBlock(nn.Module):
|
|
def __init__(self, dim, n_head, window_size, pool_size, depths, t2t_params=None):
|
|
super().__init__()
|
|
blocks = []
|
|
for i in range(depths):
|
|
blocks.append(
|
|
TemporalSparseTransformer(dim, n_head, window_size, pool_size, t2t_params=t2t_params)
|
|
)
|
|
self.transformer = nn.Sequential(*blocks)
|
|
self.depths = depths
|
|
|
|
def forward(self, x, fold_x_size, l_mask=None, t_dilation=2):
|
|
"""
|
|
Args:
|
|
x: image tokens, shape [B T H W C]
|
|
fold_x_size: fold feature size, shape [60 108]
|
|
l_mask: local mask tokens, shape [B T H W 1]
|
|
Returns:
|
|
out_tokens: shape [B T H W C]
|
|
"""
|
|
assert self.depths % t_dilation == 0, 'wrong t_dilation input.'
|
|
T = x.size(1)
|
|
T_ind = [torch.arange(i, T, t_dilation) for i in range(t_dilation)] * (self.depths // t_dilation)
|
|
|
|
for i in range(0, self.depths):
|
|
x = self.transformer[i](x, fold_x_size, l_mask, T_ind[i])
|
|
|
|
return x
|