def main():
nClasses = args.nClasses
train_batch_size = 16
val_batch_size = 16
epochs = 50
img_height = 256
img_width = 256
root_path = '../../datasets/segmentation/'
mode = 'seg' if nClasses == 2 else 'parse'
train_file = './data/{}_train.txt'.format(mode)
val_file = './data/{}_test.txt'.format(mode)
if args.model == 'unet':
model = unet.Unet(nClasses,
input_height=img_height,
input_width=img_width)
elif args.model == 'segnet':
model = segnet.SegNet(nClasses,
input_height=img_height,
input_width=img_width)
else:
raise ValueError(
'Does not support {}, only supports unet and segnet now'.format(
args.model))
model.compile(loss='categorical_crossentropy',
optimizer=Adam(lr=1e-4),
metrics=['accuracy'])
model.summary()
train = segdata_generator.generator(root_path, train_file,
train_batch_size, nClasses, img_height,
img_width)
val = segdata_generator.generator(root_path,
val_file,
val_batch_size,
nClasses,
img_height,
img_width,
train=False)
if not os.path.exists('./results/'):
os.mkdir('./results')
save_file = './weights/{}_seg_weights.h5'.format(args.model) if nClasses == 2 \
else './weights/{}_parse_weights.h5'.format(args.model)
checkpoint = ModelCheckpoint(save_file,
monitor='val_acc',
save_best_only=True,
save_weights_only=True,
verbose=1)
history = model.fit_generator(
train,
steps_per_epoch=12706 // train_batch_size,
validation_data=val,
validation_steps=5000 // val_batch_size,
epochs=epochs,
callbacks=[checkpoint],
)
plot_history(history, './results/', args.model)
save_history(history, './results/', args.model)
def loader(model, num_classes):
if model == 'LinkNet18':
return linknet.LinkNet18(num_classes=num_classes) #albu
if model == 'LinkNet34':
return linknet.LinkNet34(num_classes=num_classes)
if model == 'LinkNet50':
return linknet.LinkNet50(num_classes=num_classes)
if model == 'UNet11':
return unet_models.UNet11(num_classes=num_classes,
num_filters=32) #ternaus
if model == 'UnetVgg11':
return vgg_unet.UnetVgg11(n_classes=num_classes, num_filters=64,
v=2) #asanakoy
if model == 'Vgg11a':
return vgg_unet.Vgg11a(n_classes=num_classes, num_filters=32, v=1)
if model == 'Unet4':
return unet.Unet4(n_classes=num_classes)
if model == 'Unet5':
return unet.Unet5(n_classes=num_classes)
if model == 'Unet7':
return unet.Unet7(n_classes=num_classes)
if model == 'UNarrow':
return unet.UNarrow(n_classes=num_classes)
if model == 'Unet':
return unet.Unet(n_classes=num_classes)
if model == 'UNet256_3x3':
return unet1.UNet256_3x3(in_shape=(3, 256, 256),
num_classes=num_classes)
def predict_segmentation():
n_classes = 2
images_path = '/home/deep/datasets/'
val_file = './data/seg_test.txt'
input_height = 256
input_width = 256
if args.model == 'unet':
m = unet.Unet(n_classes,
input_height=input_height,
input_width=input_width)
elif args.model == 'segnet':
m = segnet.SegNet(n_classes,
input_height=input_height,
input_width=input_width)
else:
raise ValueError('Do not support {}'.format(args.model))
m.load_weights("./results/{}_weights.h5".format(args.model))
m.compile(loss='categorical_crossentropy',
optimizer='adam',
metrics=['accuracy'])
colors = np.array([[0, 0, 0], [255, 255, 255]])
i = 0
for x, y in generator(images_path, val_file, 1, n_classes, input_height,
input_width):
pr = m.predict(x)[0]
pr = pr.reshape((input_height, input_width, n_classes)).argmax(axis=2)
seg_img = np.zeros((input_height, input_width, 3))
for c in range(n_classes):
seg_img[:, :,
0] += ((pr[:, :] == c) * (colors[c][0])).astype('uint8')
seg_img[:, :,
1] += ((pr[:, :] == c) * (colors[c][1])).astype('uint8')
seg_img[:, :,
2] += ((pr[:, :] == c) * (colors[c][2])).astype('uint8')
cv2.imshow('test', seg_img)
cv2.imwrite('./output/{}.jpg'.format(i), seg_img)
i += 1
cv2.waitKey(30)
from models import unet
from compression import decode_weights
from utils.seg_data import generator
if __name__ == '__main__':
nClasses = 2
img_height = 256
img_width = 256
root_path = '../../datasets/'
val_file = './data/seg_test.txt'
model = unet.Unet(nClasses, input_height=img_height, input_width=img_width)
model.summary()
weights = decode_weights('./results/compressed_unet_weights.h5')
for i in range(len(model.layers)):
if model.layers[i].name in weights:
weight = [w for w in weights[model.layers[i].name]]
model.layers[i].set_weights(weight)
model.compile(loss='categorical_crossentropy',
optimizer='adam',
metrics=['accuracy'])
val = generator(root_path,
val_file,
16,
nClasses,
img_height,
img_width,
train=False)
score = model.evaluate_generator(val, steps=5000 // 16)
print('val loss: {}'.format(score[0]))
def main():
nClasses = 2
batch_size = 16
epochs = 50
fine_tune_epochs = 20
img_height = 256
img_width = 256
root_path = '../../datasets/'
mode = 'seg' if nClasses == 2 else 'parse'
train_file = './data/{}_train.txt'.format(mode)
val_file = './data/{}_test.txt'.format(mode)
model = unet.Unet(nClasses, input_height=img_height, input_width=img_width)
model.compile(loss='categorical_crossentropy',
optimizer=Adam(lr=1e-4),
metrics=['accuracy'])
model.summary()
train = seg_data.generator(root_path, train_file, batch_size, nClasses,
img_height, img_width)
val = seg_data.generator(root_path,
val_file,
batch_size,
nClasses,
img_height,
img_width,
train=False)
if not os.path.exists('./results/'):
os.mkdir('./results')
history = model.fit_generator(train,
steps_per_epoch=12706 // batch_size,
validation_data=val,
validation_steps=5000 // batch_size,
epochs=epochs)
save_history(history, './results/', 'unet')
model.save_weights('./results/Unet_weights.h5')
# prune weights
# save masks for weight layers
masks = {}
layer_count = 0
# not compress first convolution layer
first_conv = True
for layer in model.layers:
weight = layer.get_weights()
if len(weight) >= 2:
if not first_conv:
w = deepcopy(weight)
tmp, mask = prune_weights(w[0],
compress_rate=args.compress_rate)
masks[layer_count] = mask
w[0] = tmp
layer.set_weights(w)
else:
first_conv = False
layer_count += 1
# evaluate model after pruning
score = model.evaluate_generator(val, steps=5000 // batch_size)
print('val loss: {}'.format(score[0]))
print('val acc: {}'.format(score[1]))
# fine-tune
for i in range(fine_tune_epochs):
for X, Y in seg_data.generator(root_path, train_file, batch_size,
nClasses, img_height, img_width):
# train on each batch
model.train_on_batch(X, Y)
# apply masks
for layer_id in masks:
w = model.layers[layer_id].get_weights()
w[0] = w[0] * masks[layer_id]
model.layers[layer_id].set_weights(w)
score = model.evaluate_generator(seg_data.generator(root_path,
val_file,
batch_size,
nClasses,
img_height,
img_width,
train=False),
steps=5000 // batch_size)
print('val loss: {}'.format(score[0]))
print('val acc: {}'.format(score[1]))
# save compressed weights
compressed_name = './results/compressed_unet_weights'
save_compressed_weights(model, compressed_name)
parse.add_argument('--dtype', type=str, default='float16')
args = parse.parse_args()
K.set_floatx(args.dtype)
n_classes = args.nClasses
images_path = '../../datasets/segmentation/'
val_file = './data/seg_test.txt' if n_classes == 2 else './data/parse_test.txt'
weights_file = './weights/{}_seg_weights.h5'.format(args.model) if n_classes == 2 \
else './weights/{}_parse_weights.h5'.format(args.model)
input_height = 256
input_width = 256
if args.model == 'unet':
m = unet.Unet(n_classes,
input_height=input_height,
input_width=input_width)
elif args.model == 'segnet':
m = segnet.SegNet(n_classes,
input_height=input_height,
input_width=input_width)
else:
raise ValueError('Do not support {}'.format(args.model))
m.load_weights(weights_file.format(args.model))
m.compile(loss='categorical_crossentropy',
optimizer='adam',
metrics=['accuracy'])
print('Start evaluating..')
pbdr = tqdm(total=5000)
num_workers=config.train.numWorkers,
drop_last=True)
if db_test_synthetic_list:
assert (config.train.testBatchSize <= len(db_test_synthetic)), \
('testBatchSize ({}) cannot be more than the ' +
'number of images in test dataset: {}').format(config.train.testBatchSize, len(db_test_synthetic_list))
testSyntheticLoader = DataLoader(db_test_synthetic,
batch_size=config.train.testBatchSize,
shuffle=True,
num_workers=config.train.numWorkers,
drop_last=True)
###################### ModelBuilder #############################
if config.train.model == 'unet':
model = unet.Unet(num_classes=config.train.numClasses)
else:
raise ValueError(
'Invalid model "{}" in config file. Must be one of ["unet"]'.format(
config.train.model))
if config.train.continueTraining:
print(
'Transfer Learning enabled. Model State to be loaded from a prev checkpoint...'
)
if not os.path.isfile(config.train.pathPrevCheckpoint):
raise ValueError(
'Invalid path to the given weights file for transfer learning.\
The file {} does not exist'.format(
config.train.pathPrevCheckpoint))
def __init__(self, opt):
super(DDPModel, self).__init__()
self.gpu_ids = opt.gpu_ids
self.device = torch.device(
'cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device(
'cpu') # get device name: CPU or GPU
self.save_dir = os.path.join(
opt.checkpoints_dir,
opt.name) # save all the checkpoints to save_dir
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
print('Directory created: %s' % self.save_dir)
# define schedule
if opt.beta_schedule == 'cosine':
""" Cosine Schedule
@inproceedings{
anonymous2021improved,
title={Improved Denoising Diffusion Probabilistic Models},
author={Anonymous},
booktitle={Submitted to International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=-NEXDKk8gZ},
note={under review}
}
"""
s = 0.008
x = np.linspace(0, opt.num_timesteps - 1, opt.num_timesteps - 1)
alphas_cumprod = np.cos(
((x / opt.num_timesteps) + s) / (1 + s) * np.pi * 0.5)**2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
alphas_cumprod[alphas_cumprod > 0.999999] = 0.999999
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
betas = np.clip(1 - (alphas_cumprod / alphas_cumprod_prev),
a_min=0,
a_max=0.999)
alphas = 1. - betas
else:
betas = self.get_beta_schedule(opt.beta_schedule, opt.beta_start,
opt.beta_end, opt.num_timesteps)
alphas = 1. - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
assert alphas_cumprod_prev.shape == betas.shape
assert isinstance(
betas, np.ndarray) and (betas >= 0).all() and (betas <= 1).all()
timesteps, = betas.shape
self.num_timesteps = int(timesteps)
self.betas = torch.tensor(betas)
self.alphas_cumprod = torch.tensor(alphas_cumprod)
self.alphas_cumprod_prev = torch.tensor(alphas_cumprod_prev)
# calculations for diffusion q(x_t | x_{t-1}) and others
self.sqrt_alphas_cumprod = torch.tensor(np.sqrt(alphas_cumprod),
dtype=torch.float32,
device=self.device)
self.sqrt_one_minus_alphas_cumprod = torch.tensor(
np.sqrt(1. - alphas_cumprod),
dtype=torch.float32,
device=self.device)
self.log_one_minus_alphas_cumprod = torch.tensor(
np.log(1. - alphas_cumprod),
dtype=torch.float32,
device=self.device)
self.sqrt_recip_alphas_cumprod = torch.tensor(np.sqrt(1. /
alphas_cumprod),
dtype=torch.float32,
device=self.device)
self.sqrt_recipm1_alphas_cumprod = torch.tensor(
np.sqrt(1. / alphas_cumprod - 1),
dtype=torch.float32,
device=self.device)
# calculations for posterior q(x_{t-1} | x_t, x_0)
posterior_variance = betas * (1. - alphas_cumprod_prev) / (
1. - alphas_cumprod)
# above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
self.posterior_variance = torch.tensor(posterior_variance,
dtype=torch.float32,
device=self.device)
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
self.posterior_log_variance_clipped = torch.tensor(np.log(
np.maximum(posterior_variance, 1e-20)),
dtype=torch.float32,
device=self.device)
self.posterior_mean_coef1 = torch.tensor(
betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod),
dtype=torch.float32,
device=self.device)
self.posterior_mean_coef2 = torch.tensor(
(1. - alphas_cumprod_prev) * np.sqrt(alphas) /
(1. - alphas_cumprod),
dtype=torch.float32,
device=self.device)
# setup denoise model
model = []
if opt.block_size != 1:
model += [utils.SpaceToDepth(opt.block_size)]
model += [
unet.Unet(opt.input_nc,
opt.input_nc,
num_middles=1,
ngf=opt.ngf,
norm=opt.norm,
activation=opt.activation,
use_dropout=opt.dropout,
use_attention=opt.attention,
device=self.device)
]
if opt.block_size != 1:
model += [utils.SpaceToDepth(opt.block_size)]
self.denoise_model = utils.init_net(nn.Sequential(*model),
opt.init_type, opt.init_gain,
opt.gpu_ids)
if opt.phase == 'train':
# setup optimizer, visualizer, and learning rate scheduler
self.optimizer = torch.optim.Adam(self.denoise_model.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if 'mse' in opt.loss_type:
self.loss_criteria = nn.MSELoss()
elif 'l1' in opt.loss_type:
self.loss_criteria = nn.L1Loss()
else:
raise NotImplementedError(opt.loss_type)
# set prediction function
if 'noisepred' in opt.loss_type:
self.pred_fn = DDPModel._noisepred
else:
raise NotImplementedError(opt.loss_type)
self.loss_type = opt.loss_type
self.visualizer = visualizer.Visualizer(opt)
self.scheduler = utils.get_scheduler(self.optimizer, opt)
self.lr_policy = opt.lr_policy
else:
self.image_size = (opt.batch_size, opt.input_nc, opt.load_size,
opt.load_size)
self.denoise_model.train(False)
# set prediction function
if 'noisepred' in opt.loss_type:
self.pred_fn = self.predict_start_from_noise
else:
raise NotImplementedError(opt.loss_type)
if opt.phase == 'interpolate':
self.mix_rate = opt.mix_rate
def unetPredict(BASE, gpu=False):
'''
Outputs segmentations based on trained unet model.
'''
with warnings.catch_warnings():
warnings.simplefilter("ignore")
fxn()
print('-------- UNet Segmentation --------')
dtype = torch.FloatTensor
if gpu:
dtype = torch.cuda.FloatTensor
Scan_Folder = os.path.join(BASE, 'Scans')
scans_all = [
s for s in os.listdir(Scan_Folder)
if (s.endswith('nii.gz') or s.endswith('nii') and not 'MNI152' in s)
]
#[s for s in os.listdir(Scan_Folder) if (s.endswith('nii.gz') and not 'MNI152' in s)]
N1 = 128
N = 256
batch_size = 1
def NPH_data_gen():
N_pred = len(scans_all)
X = np.empty([batch_size, 1, N, N, N1])
while True:
inds = range(len(scans_all))
i = 0
while i <= (N_pred - batch_size):
for j in range(batch_size):
scanname = scans_all[inds[i + j]]
img_nib = nib.load(os.path.join(Scan_Folder, scanname))
img = img_nib.get_data()
img_info = (img.shape, img_nib.affine)
img[np.where(img < -1000)] = -1000
temp_img = resize(img,
output_shape=[N, N, N1],
preserve_range=True,
mode='constant',
order=1,
anti_aliasing=True)
X[j, 0] = temp_img
i += 1
X -= 100
X /= 100
yield scanname, img_info, torch.from_numpy(X).type(dtype)
data_gen = NPH_data_gen()
''' Loading unet '''
import models
from models import criterions, unet
unet_model = 'source/unet_ce_hard_per_im_s8841_all'
ckpt = 'model_last.tar'
unet = unet.Unet()
unet.cpu()
if gpu:
unet.cuda()
model_file = os.path.join(unet_model, ckpt)
if gpu:
device = 'cuda'
else:
device = 'cpu'
checkpoint = torch.load(model_file, map_location=torch.device(device))
unet.load_state_dict(checkpoint['state_dict'])
'''Unet Definition'''
def normalization(planes, norm='gn'):
if norm == 'bn':
m = nn.BatchNorm3d(planes)
elif norm == 'gn':
m = nn.GroupNorm(4, planes)
elif norm == 'in':
m = nn.InstanceNorm3d(planes)
else:
raise ValueError(
'normalization type {} is not supported'.format(norm))
return m
class ConvD(nn.Module):
def __init__(self,
inplanes,
planes,
dropout=0.0,
norm='gn',
first=False):
super(ConvD, self).__init__()
self.first = first
self.maxpool = nn.MaxPool3d(2, 2)
self.dropout = dropout
self.relu = nn.ReLU(inplace=True)
self.conv1 = nn.Conv3d(inplanes, planes, 3, 1, 1, bias=False)
self.bn1 = normalization(planes, norm)
self.conv2 = nn.Conv3d(planes, planes, 3, 1, 1, bias=False)
self.bn2 = normalization(planes, norm)
self.conv3 = nn.Conv3d(planes, planes, 3, 1, 1, bias=False)
self.bn3 = normalization(planes, norm)
def forward(self, x):
if not self.first:
x = self.maxpool(x)
x = self.bn1(self.conv1(x))
y = self.relu(self.bn2(self.conv2(x)))
if self.dropout > 0:
y = F.dropout3d(y, self.dropout)
y = self.bn3(self.conv3(x))
return self.relu(x + y)
class ConvU(nn.Module):
def __init__(self, planes, norm='gn', first=False):
super(ConvU, self).__init__()
self.first = first
if not self.first:
self.conv1 = nn.Conv3d(2 * planes, planes, 3, 1, 1, bias=False)
self.bn1 = normalization(planes, norm)
self.conv2 = nn.Conv3d(planes, planes // 2, 1, 1, 0, bias=False)
self.bn2 = normalization(planes // 2, norm)
self.conv3 = nn.Conv3d(planes, planes, 3, 1, 1, bias=False)
self.bn3 = normalization(planes, norm)
self.relu = nn.ReLU(inplace=True)
def forward(self, x, prev):
# final output is the localization layer
if not self.first:
x = self.relu(self.bn1(self.conv1(x)))
y = F.interpolate(x,
scale_factor=2,
mode='trilinear',
align_corners=False)
y = self.relu(self.bn2(self.conv2(y)))
y = torch.cat([prev, y], 1)
y = self.relu(self.bn3(self.conv3(y)))
return y
# End Unet definition
import copy
net = copy.deepcopy(unet)
net.convd1.conv1 = ConvD(1, 16, 0.5, 'gn', first=True)
net.convd1.conv1.weight = nn.Parameter(
unet.convd1.conv1.weight[:, 1, :, :, :].unsqueeze(1))
net.seg3 = nn.Conv3d(128, 3, kernel_size=(1, 1, 1), stride=(1, 1, 1))
net.seg2 = nn.Conv3d(64, 3, kernel_size=(1, 1, 1), stride=(1, 1, 1))
net.seg1 = nn.Conv3d(32, 3, kernel_size=(1, 1, 1), stride=(1, 1, 1))
net.seg3.weight = nn.Parameter(unet.seg3.weight[0:3, :, :, :, :])
net.seg2.weight = nn.Parameter(unet.seg2.weight[0:3, :, :, :, :])
net.seg1.weight = nn.Parameter(unet.seg1.weight[0:3, :, :, :, :])
net.cpu()
if gpu:
net.cuda()
reload_path = os.path.join(BASE, 'unet_model.pt')
net.load_state_dict(
torch.load(reload_path, map_location=torch.device(device)))
del unet
output_path = os.path.join(BASE, 'UNet_Outputs')
if not os.path.exists(output_path):
os.mkdir(output_path)
''' Predict '''
for k in range(len(scans_all)):
name, info, inputs = next(data_gen)
print(name)
save_name = os.path.join(
output_path, name[:name.find('.nii.gz')] + '.segmented.nii.gz')
if os.path.exists(save_name):
continue
inputs_np = inputs.cpu().detach().numpy()
in_mean = np.mean(inputs_np)
in_std = np.std(inputs_np)
output = net(inputs)
n, c = output.shape[:2]
N = output.shape[3]
output = output.view(n, c, -1)
lsoftmax = nn.LogSoftmax(dim=1)
output = lsoftmax(output)
output = output.argmax(dim=1)
if n == 1:
output = output.view(N, N, N1)
else:
output = output.view(n, N, N, N1)
output = output.cpu().detach().numpy()
output = resize(output,
info[0],
preserve_range=True,
mode='constant',
order=0,
anti_aliasing=None)
out_img = nib.Nifti1Image(output, info[1])
nib.save(out_img, save_name)
images_path = '../../datasets/segmentation/'
val_file = './data/seg_test.txt' if n_classes == 2 else './data/parse_test.txt'
weights_file = './weights/{}_seg_weights.h5'.format(args.model) if n_classes == 2 \
else './weights/{}_parse_weights.h5'.format(args.model)
input_height = 256
input_width = 256
weights = weight_loader(weights_file)
X = tf.placeholder(tf.float32, [None, 256, 256, 3])
Y = tf.placeholder(tf.float32, [None, 256 * 256, args.nClasses])
with tf.device('/gpu:0'):
if args.model == 'unet':
logits = unet.Unet(X,
weights,
n_classes,
input_height=input_height,
input_width=input_width)
else:
raise ValueError('Do not support {}'.format(args.model))
prediction = tf.nn.softmax(logits)
print('Start evaluating..')
pbdr = tqdm(total=5000)
iou = [0. for _ in range(1, n_classes)]
count = [0. for _ in range(1, n_classes)]
with tf.Session() as sess:
for x, y in generator(images_path,
val_file,
1,
n_classes,
previous best {:.4f}'.format(epoch, lr, best_pred))
self.epoch = epoch
assert lr >= 0
self._adjust_lr(lr, optimizer)
def _adjust_lr(self, lr, optimizer):
optimizer.param_groups[0]['lr'] = lr
if __name__ == "__main__":
import torch
import sys
sys.path.append('./')
import models.unet as unet
model = unet.Unet()
model.cuda()
mode = 'poly'
base_lr = 0.1
num_epochs = 100
training_params = [{
'params': model.parameters(),
'lr': base_lr,
'momentum': 0.9
}]
optimizer = torch.optim.SGD(training_params)
i = 3
epoch = 5
best_pred = 0.87
lrsche = LRScheduler(mode, base_lr, num_epochs, iters_per_epoch=20)
testLoader_synthetic = DataLoader(db_test_synthetic,
batch_size=config.eval.batchSize,
shuffle=False,
num_workers=config.eval.numWorkers,
drop_last=False)
if db_test_list_real:
db_test_real = torch.utils.data.ConcatDataset(db_test_list_real)
testLoader_real = DataLoader(db_test_real,
batch_size=config.eval.batchSize,
shuffle=False,
num_workers=config.eval.numWorkers,
drop_last=False)
###################### ModelBuilder #############################
model = unet.Unet(num_classes=config_checkpoint.train.numClasses)
model.load_state_dict(CHECKPOINT['model_state_dict'])
# Enable Multi-GPU training
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()
### Select Loss Func ###
criterion = nn.CrossEntropyLoss(size_average=False, reduce=True)
