Spaces:
Runtime error
Runtime error
File size: 2,859 Bytes
5c718d1 |
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 |
import os
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
import torch.nn as nn
from collections import defaultdict
import torchvision
import torch.nn.functional as F
from torch.utils.data.sampler import Sampler
class Block(nn.Module):
def __init__(self, in_ch, out_ch, padding='same'):
super().__init__()
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=padding)
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=padding)
def forward(self, x):
return self.conv2(self.relu(self.conv1(x)))
class Encoder(nn.Module):
def __init__(self, chs=(3,32,64,128,256)):
super().__init__()
self.enc_blocks = nn.ModuleList([Block(chs[i], chs[i+1]) for i in range(len(chs)-1)])
self.pool = nn.MaxPool2d(2)
def forward(self, x):
ftrs = []
for block in self.enc_blocks:
x = block(x)
ftrs.append(x)
x = self.pool(x)
return ftrs
class Decoder(nn.Module):
def __init__(self, chs=(256,128, 64, 32), aux_ch=70):
super().__init__()
upchs = tuple([chs[i]+aux_ch if i == 0 else chs[i] for i in range(len(chs))])
self.chs = chs
self.upchs = upchs
self.upconvs = nn.ModuleList([nn.ConvTranspose2d(upchs[i], upchs[i+1], 2, 2) for i in range(len(upchs)-1)])
self.dec_blocks = nn.ModuleList([Block(chs[i], chs[i+1]) for i in range(len(chs)-1)])
def forward(self, x, encoder_features):
for i in range(len(self.chs)-1):
# pdb.set_trace()
x = self.upconvs[i](x)
enc_ftrs = self.crop(encoder_features[i], x)
x = torch.cat([x, enc_ftrs], dim=1)
x = self.dec_blocks[i](x)
return x
def crop(self, enc_ftrs, x):
_, _, H, W = x.shape
enc_ftrs = torchvision.transforms.CenterCrop([H, W])(enc_ftrs)
return enc_ftrs
class AuxUNet(nn.Module):
# UNet with auxiliary feature at the bottom
def __init__(self, enc_chs=(3,32,64,128,256), dec_chs=(256,128, 64, 32), aux_ch=70, num_class=7, retain_dim=False, out_sz=(224,224)):
super().__init__()
self.encoder = Encoder(enc_chs)
self.decoder = Decoder(dec_chs, aux_ch)
self.head = nn.Conv2d(dec_chs[-1], num_class, 1)
self.retain_dim = retain_dim
def forward(self, x, aux):
# aux: auxiliary feature at the bottom
enc_ftrs = self.encoder(x)
enc_ftrs[-1] = torch.cat((enc_ftrs[-1], aux), 1)
out = self.decoder(enc_ftrs[::-1][0], enc_ftrs[::-1][1:])
out = self.head(out)
if self.retain_dim:
out = F.interpolate(out, out_sz)
return out |