File size: 3,925 Bytes
f9b628d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import torch
import torch.nn as nn
import torch.nn.functional as F
class DoubleConv(nn.Module):
def __init__(self, in_ch: int, out_ch: int):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.conv(x)
class Down(nn.Module):
def __init__(self, in_ch: int, out_ch: int):
super().__init__()
self.pool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_ch, out_ch),
)
def forward(self, x):
return self.pool_conv(x)
class Up(nn.Module):
def __init__(self, in_ch: int, out_ch: int):
super().__init__()
self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride=2)
self.conv = DoubleConv(in_ch, out_ch)
def forward(self, x1, x2):
x1 = self.up(x1)
dy = x2.size(2) - x1.size(2)
dx = x2.size(3) - x1.size(3)
x1 = F.pad(x1, [dx // 2, dx - dx // 2, dy // 2, dy - dy // 2])
return self.conv(torch.cat([x2, x1], dim=1))
class UNet(nn.Module):
def __init__(self, in_channels: int = 1, n_classes: int = 3,
base_filters: int = 64, dropout: float = 0.1):
super().__init__()
f = base_filters
self.inc = DoubleConv(in_channels, f)
self.down1 = Down(f, f * 2)
self.down2 = Down(f * 2, f * 4)
self.down3 = Down(f * 4, f * 8)
self.down4 = Down(f * 8, f * 16)
self.drop = nn.Dropout2d(dropout)
self.up1 = Up(f * 16, f * 8)
self.up2 = Up(f * 8, f * 4)
self.up3 = Up(f * 4, f * 2)
self.up4 = Up(f * 2, f)
self.outc = nn.Conv2d(f, n_classes, 1)
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.drop(self.down4(x4))
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
return self.outc(x)
class DiceLoss(nn.Module):
def __init__(self, smooth: float = 1.0):
super().__init__()
self.smooth = smooth
def forward(self, pred, target):
pred = F.softmax(pred, dim=1)
oh = F.one_hot(target, pred.size(1)).permute(0, 3, 1, 2).float()
inter = (pred * oh).sum(dim=(2, 3))
union = pred.sum(dim=(2, 3)) + oh.sum(dim=(2, 3))
return 1.0 - ((2.0 * inter + self.smooth) / (union + self.smooth)).mean()
class CombinedLoss(nn.Module):
"""
0.5 * Dice + 0.5 * weighted CrossEntropy
Weights: background=0.1 disc=1.5 cup=3.0
"""
def __init__(self):
super().__init__()
self.dice = DiceLoss()
self._w = torch.tensor([0.2, 1.5, 6.0])
self.ce = nn.CrossEntropyLoss(weight=self._w)
def to(self, device):
super().to(device)
self._w = self._w.to(device)
self.ce = nn.CrossEntropyLoss(weight=self._w)
return self
def forward(self, pred, target):
return 0.5 * self.dice(pred, target) + 0.5 * self.ce(pred, target)
def calculate_dice(pred, target) -> float:
import numpy as np
p = pred.astype(bool)
t = target.astype(bool)
return float((2.0 * (p & t).sum() + 1e-5) / (p.sum() + t.sum() + 1e-5))
|