zpwiki/pages/students/2016/lukas_pokryvka/dp2021/lungCancer/model/model.py

225 lines
7.4 KiB
Python

import math
import random
from collections import namedtuple
import torch
from torch import nn as nn
import torch.nn.functional as F
from util.logconf import logging
from util.unet import UNet
log = logging.getLogger(__name__)
# log.setLevel(logging.WARN)
# log.setLevel(logging.INFO)
log.setLevel(logging.DEBUG)
class UNetWrapper(nn.Module):
def __init__(self, **kwargs):
super().__init__()
self.input_batchnorm = nn.BatchNorm2d(kwargs['in_channels'])
self.unet = UNet(**kwargs)
self.final = nn.Sigmoid()
self._init_weights()
def _init_weights(self):
init_set = {
nn.Conv2d,
nn.Conv3d,
nn.ConvTranspose2d,
nn.ConvTranspose3d,
nn.Linear,
}
for m in self.modules():
if type(m) in init_set:
nn.init.kaiming_normal_(
m.weight.data, mode='fan_out', nonlinearity='relu', a=0
)
if m.bias is not None:
fan_in, fan_out = \
nn.init._calculate_fan_in_and_fan_out(m.weight.data)
bound = 1 / math.sqrt(fan_out)
nn.init.normal_(m.bias, -bound, bound)
# nn.init.constant_(self.unet.last.bias, -4)
# nn.init.constant_(self.unet.last.bias, 4)
def forward(self, input_batch):
bn_output = self.input_batchnorm(input_batch)
un_output = self.unet(bn_output)
fn_output = self.final(un_output)
return fn_output
class SegmentationAugmentation(nn.Module):
def __init__(
self, flip=None, offset=None, scale=None, rotate=None, noise=None
):
super().__init__()
self.flip = flip
self.offset = offset
self.scale = scale
self.rotate = rotate
self.noise = noise
def forward(self, input_g, label_g):
transform_t = self._build2dTransformMatrix()
transform_t = transform_t.expand(input_g.shape[0], -1, -1)
transform_t = transform_t.to(input_g.device, torch.float32)
affine_t = F.affine_grid(transform_t[:,:2],
input_g.size(), align_corners=False)
augmented_input_g = F.grid_sample(input_g,
affine_t, padding_mode='border',
align_corners=False)
augmented_label_g = F.grid_sample(label_g.to(torch.float32),
affine_t, padding_mode='border',
align_corners=False)
if self.noise:
noise_t = torch.randn_like(augmented_input_g)
noise_t *= self.noise
augmented_input_g += noise_t
return augmented_input_g, augmented_label_g > 0.5
def _build2dTransformMatrix(self):
transform_t = torch.eye(3)
for i in range(2):
if self.flip:
if random.random() > 0.5:
transform_t[i,i] *= -1
if self.offset:
offset_float = self.offset
random_float = (random.random() * 2 - 1)
transform_t[2,i] = offset_float * random_float
if self.scale:
scale_float = self.scale
random_float = (random.random() * 2 - 1)
transform_t[i,i] *= 1.0 + scale_float * random_float
if self.rotate:
angle_rad = random.random() * math.pi * 2
s = math.sin(angle_rad)
c = math.cos(angle_rad)
rotation_t = torch.tensor([
[c, -s, 0],
[s, c, 0],
[0, 0, 1]])
transform_t @= rotation_t
return transform_t
# MaskTuple = namedtuple('MaskTuple', 'raw_dense_mask, dense_mask, body_mask, air_mask, raw_candidate_mask, candidate_mask, lung_mask, neg_mask, pos_mask')
#
# class SegmentationMask(nn.Module):
# def __init__(self):
# super().__init__()
#
# self.conv_list = nn.ModuleList([
# self._make_circle_conv(radius) for radius in range(1, 8)
# ])
#
# def _make_circle_conv(self, radius):
# diameter = 1 + radius * 2
#
# a = torch.linspace(-1, 1, steps=diameter)**2
# b = (a[None] + a[:, None])**0.5
#
# circle_weights = (b <= 1.0).to(torch.float32)
#
# conv = nn.Conv2d(1, 1, kernel_size=diameter, padding=radius, bias=False)
# conv.weight.data.fill_(1)
# conv.weight.data *= circle_weights / circle_weights.sum()
#
# return conv
#
#
# def erode(self, input_mask, radius, threshold=1):
# conv = self.conv_list[radius - 1]
# input_float = input_mask.to(torch.float32)
# result = conv(input_float)
#
# # log.debug(['erode in ', radius, threshold, input_float.min().item(), input_float.mean().item(), input_float.max().item()])
# # log.debug(['erode out', radius, threshold, result.min().item(), result.mean().item(), result.max().item()])
#
# return result >= threshold
#
# def deposit(self, input_mask, radius, threshold=0):
# conv = self.conv_list[radius - 1]
# input_float = input_mask.to(torch.float32)
# result = conv(input_float)
#
# # log.debug(['deposit in ', radius, threshold, input_float.min().item(), input_float.mean().item(), input_float.max().item()])
# # log.debug(['deposit out', radius, threshold, result.min().item(), result.mean().item(), result.max().item()])
#
# return result > threshold
#
# def fill_cavity(self, input_mask):
# cumsum = input_mask.cumsum(-1)
# filled_mask = (cumsum > 0)
# filled_mask &= (cumsum < cumsum[..., -1:])
# cumsum = input_mask.cumsum(-2)
# filled_mask &= (cumsum > 0)
# filled_mask &= (cumsum < cumsum[..., -1:, :])
#
# return filled_mask
#
#
# def forward(self, input_g, raw_pos_g):
# gcc_g = input_g + 1
#
# with torch.no_grad():
# # log.info(['gcc_g', gcc_g.min(), gcc_g.mean(), gcc_g.max()])
#
# raw_dense_mask = gcc_g > 0.7
# dense_mask = self.deposit(raw_dense_mask, 2)
# dense_mask = self.erode(dense_mask, 6)
# dense_mask = self.deposit(dense_mask, 4)
#
# body_mask = self.fill_cavity(dense_mask)
# air_mask = self.deposit(body_mask & ~dense_mask, 5)
# air_mask = self.erode(air_mask, 6)
#
# lung_mask = self.deposit(air_mask, 5)
#
# raw_candidate_mask = gcc_g > 0.4
# raw_candidate_mask &= air_mask
# candidate_mask = self.erode(raw_candidate_mask, 1)
# candidate_mask = self.deposit(candidate_mask, 1)
#
# pos_mask = self.deposit((raw_pos_g > 0.5) & lung_mask, 2)
#
# neg_mask = self.deposit(candidate_mask, 1)
# neg_mask &= ~pos_mask
# neg_mask &= lung_mask
#
# # label_g = (neg_mask | pos_mask).to(torch.float32)
# label_g = (pos_mask).to(torch.float32)
# neg_g = neg_mask.to(torch.float32)
# pos_g = pos_mask.to(torch.float32)
#
# mask_dict = {
# 'raw_dense_mask': raw_dense_mask,
# 'dense_mask': dense_mask,
# 'body_mask': body_mask,
# 'air_mask': air_mask,
# 'raw_candidate_mask': raw_candidate_mask,
# 'candidate_mask': candidate_mask,
# 'lung_mask': lung_mask,
# 'neg_mask': neg_mask,
# 'pos_mask': pos_mask,
# }
#
# return label_g, neg_g, pos_g, lung_mask, mask_dict