Skip to content

Model architecture

_model

SmallUNet architecture — compact U-Net with base width 16.

ConvBlock 

ConvBlock(in_channels, out_channels)

Bases: Module

Source code in metrics_petri/_model.py 
19
20
21
22
23
24
25
26
def __init__(self, in_channels: int, out_channels: int) -> None:
    super().__init__()
    self.block = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
        nn.ReLU(inplace=True),
    )

block instance-attribute 

block = nn.Sequential(
    nn.Conv2d(
        in_channels,
        out_channels,
        kernel_size=3,
        padding=1,
        bias=False,
    ),
    nn.ReLU(inplace=True),
    nn.Conv2d(
        out_channels,
        out_channels,
        kernel_size=3,
        padding=1,
        bias=False,
    ),
    nn.ReLU(inplace=True),
)

forward 

forward(x)
Source code in metrics_petri/_model.py 
28
29
def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.block(x)

DownBlock 

DownBlock(in_channels, out_channels)

Bases: Module

Source code in metrics_petri/_model.py 
33
34
35
36
def __init__(self, in_channels: int, out_channels: int) -> None:
    super().__init__()
    self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
    self.conv = ConvBlock(in_channels, out_channels)

pool instance-attribute 

pool = nn.MaxPool2d(kernel_size=2, stride=2)

conv instance-attribute 

conv = ConvBlock(in_channels, out_channels)

forward 

forward(x)
Source code in metrics_petri/_model.py 
38
39
def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.conv(self.pool(x))

UpBlock 

UpBlock(in_channels, out_channels)

Bases: Module

Source code in metrics_petri/_model.py 
43
44
45
46
def __init__(self, in_channels: int, out_channels: int) -> None:
    super().__init__()
    self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
    self.conv = ConvBlock(in_channels, out_channels)

up instance-attribute 

up = nn.Upsample(
    scale_factor=2, mode="bilinear", align_corners=False
)

conv instance-attribute 

conv = ConvBlock(in_channels, out_channels)

forward 

forward(x, skip)
Source code in metrics_petri/_model.py 
48
49
50
51
52
53
def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
    x = self.up(x)
    if x.shape[-2:] != skip.shape[-2:]:
        x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False)
    x = torch.cat([skip, x], dim=1)  # skip first, matching training convention
    return self.conv(x)

SmallUNet 

SmallUNet(in_channels=3, out_channels=1, base_channels=16)

Bases: Module

Compact U-Net with base width 16 for low-data segmentation.

Source code in metrics_petri/_model.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def __init__(
    self,
    in_channels: int = 3,
    out_channels: int = 1,
    base_channels: int = 16,
) -> None:
    super().__init__()
    c1, c2, c3, c4 = (
        base_channels,
        base_channels * 2,
        base_channels * 4,
        base_channels * 8,
    )
    bn = base_channels * 16

    self.enc1 = ConvBlock(in_channels, c1)
    self.enc2 = DownBlock(c1, c2)
    self.enc3 = DownBlock(c2, c3)
    self.enc4 = DownBlock(c3, c4)
    self.bottleneck = DownBlock(c4, bn)

    self.up4 = UpBlock(bn + c4, c4)
    self.up3 = UpBlock(c4 + c3, c3)
    self.up2 = UpBlock(c3 + c2, c2)
    self.up1 = UpBlock(c2 + c1, c1)

    self.head = nn.Conv2d(c1, out_channels, kernel_size=1)
    self.activation = nn.Sigmoid()

enc1 instance-attribute 

enc1 = ConvBlock(in_channels, c1)

enc2 instance-attribute 

enc2 = DownBlock(c1, c2)

enc3 instance-attribute 

enc3 = DownBlock(c2, c3)

enc4 instance-attribute 

enc4 = DownBlock(c3, c4)

bottleneck instance-attribute 

bottleneck = DownBlock(c4, bn)

up4 instance-attribute 

up4 = UpBlock(bn + c4, c4)

up3 instance-attribute 

up3 = UpBlock(c4 + c3, c3)

up2 instance-attribute 

up2 = UpBlock(c3 + c2, c2)

up1 instance-attribute 

up1 = UpBlock(c2 + c1, c1)

head instance-attribute 

head = nn.Conv2d(c1, out_channels, kernel_size=1)

activation instance-attribute 

activation = nn.Sigmoid()

forward 

forward(x)
Source code in metrics_petri/_model.py 
88
89
90
91
92
93
94
95
96
97
98
99
def forward(self, x: torch.Tensor) -> torch.Tensor:
    s1 = self.enc1(x)
    s2 = self.enc2(s1)
    s3 = self.enc3(s2)
    s4 = self.enc4(s3)
    b = self.bottleneck(s4)
    x = self.up4(b, s4)
    x = self.up3(x, s3)
    x = self.up2(x, s2)
    x = self.up1(x, s1)
    x = self.head(x)
    return self.activation(x)