| import torch.nn as nn | |
| from .encoder import resnet34 | |
| from .decoder import DeepLabV3Decoder | |
| class DeepLabV3(nn.Module): | |
| def __init__(self, input_channels): | |
| super().__init__() | |
| self.encoder = resnet34(input_channels=input_channels) | |
| self.decoder = DeepLabV3Decoder(in_channels=128) | |
| def forward(self, x): | |
| feat = self.encoder(x) | |
| out = self.decoder(*feat) | |
| return out | |