I’m trying to build an UNET implementation to segment some cerebral images.
This is the implementation of the net:
class ConvBlock(nn.Module):
    def __init__(self,input,output):
        super(ConvBlock,self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(input, output, kernel_size=3,stride=1,padding=1,bias=True),
            nn.BatchNorm2d(output),
            nn.ReLU(inplace=True),
            nn.Conv2d(output, output, kernel_size=3,stride=1,padding=1,bias=True),
            nn.BatchNorm2d(output),
            nn.ReLU(inplace=True)
        )

    def forward(self,x):
        x = self.conv(x)
        return x

class Unet (nn.Module):
    def __init__(self,numberClasses = 2,dropout = 0.5, lastLayer = "SIGMOID"):
        super(Unet, self).__init__()

        self.CB1 = ConvBlock(1,64)
        self.CB2 = ConvBlock(64,128)
        self.CB3 = ConvBlock(128,256)
        self.CB4 = ConvBlock(256,512)

        self.CB5 = ConvBlock(512,1024)

        self.CB6 = ConvBlock(1024,512)
        self.CB7 = ConvBlock(512,256)
        self.CB8 = ConvBlock(256,128)
        self.CB9 = ConvBlock(128,64)

        self.pool = nn.MaxPool2d(2)

        self.dropout = nn.Dropout(dropout)

        self.UC1 = nn.ConvTranspose2d(1024,512,2,stride = 2)
        self.UC2 = nn.ConvTranspose2d(512,256,2,stride = 2)
        self.UC3 = nn.ConvTranspose2d(256,128,2,stride = 2)
        self.UC4 = nn.ConvTranspose2d(128,64,2,stride = 2)

        self.FC = nn.Conv2d(64,numberClasses,1)

        if lastLayer == "SIGMOID":
            self.lastLayer = nn.Sigmoid()
        elif lastLayer == "SOFTMAX":
            self.lastLayer = nn.Softmax(dim = 1)

    def forward(self,x):
        # DOWN SAMPLING
        C1 = self.CB1(x)
        x = C1.clone()
        x = self.pool(x)
        x = self.dropout(x)

        C2 = self.CB2(x)
        x = C2.clone()
        x = self.pool(x)
        x = self.dropout(x)

        C3 = self.CB3(x)
        x = C3.clone()
        x = self.pool(x)
        x = self.dropout(x)
     
        C4 = self.CB4(x)
        x = C4.clone()
        x = self.pool(x)
        x = self.dropout(x)

        x = self.CB5(x)


        # UP SAMPLING
        x = self.UC1(x)

        x = torch.cat((x,C4),dim=1)
        x = self.CB6(x)
        x = self.dropout(x)

        x = self.UC2(x)
        x = torch.cat((x,C3),dim=1)
        x = self.CB7(x)
        x = self.dropout(x)

        x = self.UC3(x)
        x = torch.cat((x,C2),dim=1)
        x = self.CB8(x)
        x = self.dropout(x)

        x = self.UC4(x)
        x = torch.cat((x,C1),dim=1)
        x = self.CB9(x)
       
        x = self.FC(x)

        x = self.lastLayer(x)
        
        return x   

I’m using a BCE loss with weighted classes, where class 0 is multiplied by 0.2 and 1 by 0.8. This is the best configuration I have found, but it still doesn’t work.
Here are some examples, comparing ground truth, expected segmentation and final results calculated by the net:

The shape seems to be correct, but I need better results, because I will count every single pixel later. If I keep training the model I get an empty output. What should I do? Apply some kind of morphological transformation?