You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
29 lines
896 B
Python
29 lines
896 B
Python
from torch import nn
|
|
|
|
class LRASPP(nn.Module):
|
|
def __init__(self, in_channels, out_channels):
|
|
super().__init__()
|
|
self.aspp1 = nn.Sequential(
|
|
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
|
nn.BatchNorm2d(out_channels),
|
|
nn.ReLU(True)
|
|
)
|
|
self.aspp2 = nn.Sequential(
|
|
nn.AdaptiveAvgPool2d(1),
|
|
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward_single_frame(self, x):
|
|
return self.aspp1(x) * self.aspp2(x)
|
|
|
|
def forward_time_series(self, x):
|
|
B, T = x.shape[:2]
|
|
x = self.forward_single_frame(x.flatten(0, 1)).unflatten(0, (B, T))
|
|
return x
|
|
|
|
def forward(self, x):
|
|
if x.ndim == 5:
|
|
return self.forward_time_series(x)
|
|
else:
|
|
return self.forward_single_frame(x) |