import torch
from torch import nn
from torchvision import models


class DeePixBiS(nn.Module):
    def __init__(self, pretrained=True):
        super().__init__()
        weights = pretrained if pretrained else None
        dense = models.densenet161(weights=weights)
        features = list(dense.features.children())
        self.enc = nn.Sequential(*features[:8])
        self.dec = nn.Conv2d(384, 1, kernel_size=1, stride=1, padding=0)
        self.linear = nn.Linear(14 * 14, 1)

    def forward(self, x):
        enc = self.enc(x)
        dec = self.dec(enc)
        out_map = torch.sigmoid(dec)
        # print(out_map.shape)
        out = self.linear(out_map.view(-1, 14 * 14))
        out = torch.sigmoid(out)
        out = torch.flatten(out)
        return out_map, out