| """PyTorch implementation of RainNet v1.0.""" |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class RainNet(nn.Module): |
| """The 20-convolution encoder-decoder from Ayzel et al. (2020).""" |
|
|
| def __init__(self, in_channels=4, out_channels=1, dropout=0.5, mode="regression"): |
| super().__init__() |
| if mode != "regression": |
| raise ValueError("This package implements RainNet regression mode only") |
| self.in_channels = in_channels |
| self.pool = nn.MaxPool2d(2) |
| self.upsample = nn.Upsample(scale_factor=2, mode="nearest") |
| self.relu = nn.ReLU(inplace=True) |
|
|
| self.conv1f = nn.Conv2d(in_channels, 64, 3, padding=1) |
| self.conv1s = nn.Conv2d(64, 64, 3, padding=1) |
| self.conv2f = nn.Conv2d(64, 128, 3, padding=1) |
| self.conv2s = nn.Conv2d(128, 128, 3, padding=1) |
| self.conv3f = nn.Conv2d(128, 256, 3, padding=1) |
| self.conv3s = nn.Conv2d(256, 256, 3, padding=1) |
| self.conv4f = nn.Conv2d(256, 512, 3, padding=1) |
| self.conv4s = nn.Conv2d(512, 512, 3, padding=1) |
| self.drop4 = nn.Dropout(dropout) |
| self.conv5f = nn.Conv2d(512, 1024, 3, padding=1) |
| self.conv5s = nn.Conv2d(1024, 1024, 3, padding=1) |
| self.drop5 = nn.Dropout(dropout) |
|
|
| self.conv6f = nn.Conv2d(1536, 512, 3, padding=1) |
| self.conv6s = nn.Conv2d(512, 512, 3, padding=1) |
| self.conv7f = nn.Conv2d(768, 256, 3, padding=1) |
| self.conv7s = nn.Conv2d(256, 256, 3, padding=1) |
| self.conv8f = nn.Conv2d(384, 128, 3, padding=1) |
| self.conv8s = nn.Conv2d(128, 128, 3, padding=1) |
| self.conv9f = nn.Conv2d(192, 64, 3, padding=1) |
| self.conv9s = nn.Conv2d(64, 64, 3, padding=1) |
| self.conv9out = nn.Conv2d(64, 2, 3, padding=1) |
| self.output = nn.Conv2d(2, out_channels, 1) |
| self.apply(self._initialize) |
|
|
| @staticmethod |
| def _initialize(module): |
| if isinstance(module, nn.Conv2d): |
| nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu") |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def forward(self, x): |
| if x.ndim != 4: |
| raise ValueError(f"RainNet expects BCHW input, got shape {tuple(x.shape)}") |
| if x.shape[1] != self.in_channels: |
| raise ValueError(f"RainNet expects {self.in_channels} channels, got {x.shape[1]}") |
| if x.shape[-2] % 16 or x.shape[-1] % 16: |
| raise ValueError( |
| f"Spatial dimensions must be divisible by 16, got {tuple(x.shape[-2:])}" |
| ) |
|
|
| c1 = self.relu(self.conv1s(self.relu(self.conv1f(x)))) |
| c2 = self.relu(self.conv2s(self.relu(self.conv2f(self.pool(c1))))) |
| c3 = self.relu(self.conv3s(self.relu(self.conv3f(self.pool(c2))))) |
| c4 = self.relu(self.conv4s(self.relu(self.conv4f(self.pool(c3))))) |
| c5 = self.relu(self.conv5s(self.relu(self.conv5f(self.pool(self.drop4(c4)))))) |
| c5 = self.drop5(c5) |
|
|
| c6 = torch.cat((self.upsample(c5), c4), dim=1) |
| c6 = self.relu(self.conv6s(self.relu(self.conv6f(c6)))) |
| c7 = torch.cat((self.upsample(c6), c3), dim=1) |
| c7 = self.relu(self.conv7s(self.relu(self.conv7f(c7)))) |
| c8 = torch.cat((self.upsample(c7), c2), dim=1) |
| c8 = self.relu(self.conv8s(self.relu(self.conv8f(c8)))) |
| c9 = torch.cat((self.upsample(c8), c1), dim=1) |
| c9 = self.relu(self.conv9s(self.relu(self.conv9f(c9)))) |
| return self.output(self.relu(self.conv9out(c9))) |
|
|
|
|
| def build_rainnet(in_channels=4, out_channels=1, dropout=0.5, mode="regression"): |
| return RainNet(in_channels, out_channels, dropout, mode) |
|
|
