import torch
import torch.nn as nn
from segmentation_models_pytorch.deeplabv3 import DeepLabV3Plus


class SegGenerator(nn.Module):
    def __init__(self):
        super(SegGenerator, self).__init__()
        self.model = DeepLabV3Plus(in_channels=3, classes=3, activation=None, encoder_weights=None)

    def forward(self, mask):
        return self.model(mask)


class SegDiscriminator(nn.Module):
    def __init__(self):
        super(SegDiscriminator, self).__init__()
        self.model = DeepLabV3Plus(in_channels=3, classes=1, activation="sigmoid", encoder_weights=None)

    def forward(self, mask):
        # adds nosie to inputs, see https://arxiv.org/pdf/1701.04862.pdf
        return self.model(mask + torch.normal(torch.zeros_like(mask), torch.ones_like(mask) / 10))
        # return self.model(mask)