datnguyentien204's picture
Upload 337 files
ce91ea1 verified
import torch.nn as nn
import torch.nn.functional as F
import torch
from layers import SaveFeature
import pretrainedmodels
from torchvision.models import resnet34, resnet50, resnet101, resnet152
from pathlib import Path
from torchvision.models.resnet import conv3x3, BasicBlock, Bottleneck
import skimage
from scipy import ndimage
import numpy as np
import torchvision.transforms as transforms
import cv2
from constant import IMAGENET_MEAN, IMAGENET_STD
device="cuda" if torch.cuda.is_available() else "cpu"
class UpBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, expansion=1):
super().__init__()
inplanes = inplanes * expansion
planes = planes * expansion
self.upconv = nn.ConvTranspose2d(inplanes, planes, 2, 2, 0)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv1 = conv3x3(inplanes, planes)
self.bn2 = nn.BatchNorm2d(planes)
def forward(self, u, x):
up = self.relu(self.bn1(self.upconv(u)))
out = torch.cat([x, up], dim=1) # cat along channel
out = self.relu(self.bn2(self.conv1(out)))
return out
class UpLayer(nn.Module):
def __init__(self, block, inplanes, planes, blocks):
super().__init__()
self.up = UpBlock(inplanes, planes, block.expansion)
layers = [block(planes * block.expansion, planes) for _ in range(1, blocks)]
self.conv = nn.Sequential(*layers)
def forward(self, u, x):
x = self.up(u, x)
x = self.conv(x)
return x
from pathlib import Path
class Unet(nn.Module):
tfm = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)
])
def __init__(self, trained=False, model_name=None):
super().__init__()
self.layers = [3, 4, 6]
self.block = Bottleneck
if trained:
assert model_name is not None
self.load_model(model_name)
else:
self.load_pretrained()
def cut_model(self, model, cut):
return list(model.children())[:cut]
def load_model(self, model_name):
resnet = resnet50(False)
self.backbone = nn.Sequential(*self.cut_model(resnet, 8))
self.init_head()
model_path = Path(__file__).parent / 'unet.h5'
state_dict = torch.load(model_path, map_location=torch.device(device))
self.load_state_dict(state_dict)
def load_pretrained(self, torch=False):
if torch:
resnet = resnet50(True)
else:
resnet = pretrainedmodels.__dict__['resnet50']()
self.backbone = nn.Sequential(*self.cut_model(resnet, 8))
self.init_head()
def init_head(self):
self.sfs = [SaveFeature(self.backbone[i]) for i in [2, 4, 5, 6]]
self.up_layer1 = UpLayer(self.block, 512, 256, self.layers[-1])
self.up_layer2 = UpLayer(self.block, 256, 128, self.layers[-2])
self.up_layer3 = UpLayer(self.block, 128, 64, self.layers[-3])
self.map = conv3x3(64 * self.block.expansion, 64) # 64e -> 64
self.conv = conv3x3(128, 64)
self.bn_conv = nn.BatchNorm2d(64)
self.up_conv = nn.ConvTranspose2d(64, 1, 2, 2, 0)
self.bn_up = nn.BatchNorm2d(1)
def forward(self, x):
x = F.relu(self.backbone(x))
x = self.up_layer1(x, self.sfs[3].features)
x = self.up_layer2(x, self.sfs[2].features)
x = self.up_layer3(x, self.sfs[1].features)
x = self.map(x)
x = F.interpolate(x, scale_factor=2)
x = torch.cat([self.sfs[0].features, x], dim=1)
x = F.relu(self.bn_conv(self.conv(x)))
x = F.relu(self.bn_up(self.up_conv(x)))
return x
def close(self):
for sf in self.sfs:
sf.remove()
def segment(self, image):
"""
image: cropped CXR PIL Image (h, w, 3)
"""
kernel = np.ones((10, 10))
iw, ih = image.size
image_tensor = self.tfm(image).unsqueeze(0).to(next(self.parameters()).device)
with torch.no_grad():
py = torch.sigmoid(self(image_tensor))
py = (py[0].cpu() > 0.5).type(torch.FloatTensor) # 1, 256, 256
mask = py[0].numpy()
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
mask = cv2.resize(mask, (iw, ih))
slice_y, slice_x = ndimage.find_objects(mask, 1)[0]
h, w = slice_y.stop - slice_y.start, slice_x.stop - slice_x.start
nw, nh = int(w / .875), int(h / .875)
dw, dh = (nw - w) // 2, (nh - h) // 2
t = max(slice_y.start - dh, 0)
l = max(slice_x.start - dw, 0)
b = min(slice_y.stop + dh, ih)
r = min(slice_x.stop + dw, iw)
return (t, l, b, r), mask