def main():
# patches_file = "50_patches_dataset.h5"
# hf = h5py.File(patches_file, 'r')
# We obtain a list with all the IDs of the patches
# all_groups = list(hf)
# Dividing the dataset into train and validation. Shuffle has to be false otherwise the model might be trained
# on what was previously validation set and validated on what was previously train set.
# X_train, X_validation = train_test_split(all_groups, test_size=0.2, shuffle=False)
# print(X_train, X_validation)
# for testing
datapath = "Data/"
train_file = datapath + "patches_dataset_test.h5"
val_file = datapath + "val250.h5"
# Loader Parameters
params = {'batch_size': 2, 'shuffle': False, 'num_workers': 0}
train_dataset = PatchDataset(train_file, n_classes=3)
print(len(train_dataset))
val_dataset = PatchDataset(val_file, n_classes=3)
print(len(val_dataset))
train_loader = DataLoader(train_dataset, **params)
val_loader = DataLoader(val_dataset, **params)
loaders = {'train': train_loader, 'val': val_loader}
# Model and param
model = Modified3DUNet(in_channels=1, n_classes=3)
optimizer = optim.Adam(model.parameters())
max_epochs = 10
# Median foreground percentage = 0.2 (= class 1,2)
# Median cancer percentage = 0.01 (= class 2)
# Median pancreas percentage = 0.2 - 0.01 = 0.19 (= class 1)
# Median background percentage = 1-0.2 = 99.8 (=class 0)
# [99.8, 0.19, 0.01] => corresponding class weights = [1, 525, 9980]
# class_weights = torch.tensor([1., 525., 9980.])
# loss_criterion = GeneralizedDiceLoss(weight=class_weights)
# loss_criterion = WeightedCrossEntropyLoss(weight=class_weights)
weights = [1, 100, 500]
class_weights = torch.FloatTensor(weights)
loss_criterion = nn.CrossEntropyLoss(weight=class_weights)
# trainer = UNetTrainer(model, optimizer, loaders, max_epochs, loss_criterion=loss_criterion)
# trainer.train()
# Load from last epoch
checkpoint_trainer = UNetTrainer.load_checkpoint(
"WCEL_1_10_50_last_model",
model,
optimizer,
loaders,
max_epochs,
loss_criterion=loss_criterion)
pred = checkpoint_trainer.single_image_forward(val_dataset[0][0])
def main():
# setup environments and seeds
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
# setup networks
#Network = getattr(models, args.net)
#model = Network(**args.net_params)
model = Modified3DUNet(in_channels=1, n_classes=2, base_n_filter=16)
model = model.cuda()
'''optimizer = getattr(torch.optim, args.opt)(
model.parameters(), **args.opt_params)'''
optimizer = torch.optim.Adam(model.parameters(),
lr=0.001,
weight_decay=0.0001)
#optimizer = torch.optim.SGD(model.parameters(),lr = 0.1,momentum=0.9)
criterion = getattr(criterions, args.criterion)
msg = '-------------- New training session -----------------'
msg += '\n' + str(args)
logging.info(msg)
num_gpus = len(args.gpu.split(','))
args.batch_size *= num_gpus
args.workers *= num_gpus
args.opt_params['lr'] *= num_gpus
# create dataloaders
#Dataset = getattr(datasets, args.dataset)
dset = cell_training('/home/tom/Modified-3D-UNet-Pytorch/PNAS/')
train_loader = DataLoader(dset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
file_name_best = os.path.join(ckpts, 'cell/model_best.tar')
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train_loss = train(train_loader, model, criterion, optimizer, epoch)
# remember best lost and save checkpoint
ckpt = {
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict(),
'train_loss': train_loss,
}
file_name = os.path.join(ckpts, 'model_last.tar')
torch.save(ckpt, file_name)
msg = 'Epoch: {:02d} Train loss {:.4f}'.format(epoch + 1, train_loss)
logging.info(msg)
logging.getLogger('').addHandler(console)
global args, best_loss
best_loss = float('inf')
args = parser.parse_args()
#print os.environ['CUDA_VISIBLE_DEVICES']
dtype = torch.float
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# input = data.to(device)
# Loading the model
in_channels = 4
n_classes = 4
base_n_filter = 16
model = Modified3DUNet(in_channels, n_classes, base_n_filter).to(device)
#print args.data
# Split the training and testing dataset
test_size = 0.1
train_idx, test_idx = train_test_split(range(285), test_size = test_size)
train_data = load_dataset(train_idx)
test_data = load_dataset(test_idx)
#print all_data.keys()
# create your optimizer
#optimizer = optim.adam(net.parameteres(), lr=)
'''
shuffle=True,
num_workers=int(opt.workers))
test_dataset = PartDataset(root='..//Thingi10K//sdf_polar', train=False)
testdataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=opt.batchSize,
shuffle=True,
num_workers=int(opt.workers))
print(len(dataset), len(test_dataset))
try:
os.makedirs(opt.outf)
except OSError:
pass
classifier = Modified3DUNet(in_channels=1, n_classes=1)
if opt.model != '':
classifier.load_state_dict(torch.load(opt.model))
optimizer = optim.SGD(classifier.parameters(), lr=opt.lr, momentum=0.9)
classifier.cuda()
loss_fn = torch.nn.MSELoss(reduce=True, size_average=True)
num_batch = len(dataset) / opt.batchSize
for epoch in range(opt.nowepoch, opt.nowepoch + opt.nepoch):
for i, data in enumerate(dataloader, 0):
points, target = data
points, target = Variable(points), Variable(target)
def main():
start_time = time.time()
# setup environments and seeds
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# setup networks
#Network = getattr(models, args.net)
#model = Network(**args.net_params)
model = Modified3DUNet(in_channels=1, n_classes=2, base_n_filter=16)
#load_model
model_file = os.path.join(ckpts, 'model_last.tar')
print model_file
checkpoint = torch.load(model_file,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
model = model.cuda()
'''optimizer = getattr(torch.optim, args.opt)(
model.parameters(), **args.opt_params)'''
#optimizer = torch.optim.SGD(model.parameters(),lr = 0.1,momentum=0.9)
criterion = getattr(criterions, args.criterion)
num_gpus = len(args.gpu.split(','))
args.batch_size *= num_gpus
args.workers *= num_gpus
# create dataloaders
#Dataset = getattr(datasets, args.dataset)
dset = cell_testing_inter('/home/tom/data1_match/dataset4/')
print dset.__len__()
test_loader = DataLoader(dset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
model.eval()
torch.set_grad_enabled(False)
inputs = []
outputs = []
ground_truth = []
for i, sample in enumerate(test_loader):
input = sample['data']
img_size = sample['image_size']
print img_size[0]
file_name = sample['name']
file_name = str(file_name[0])
_, _, z, x, y = input.shape
seg = np.zeros((z, x, y))
target = sample['seg']
for j in range(z / 16):
#ground_truth.append(target)
#print file_name[0]
input_temp = input[0, 0, j * 16:(j + 1) * 16].float()
input_temp = input_temp[None, None, ...]
output_temp = nn.parallel.data_parallel(model, input_temp)
output_temp = output_temp.detach().cpu().numpy()
output_temp = output_temp[0]
seg_temp = output_temp.argmax(0)
seg[j * 16:(j + 1) * 16] = seg_temp
data = input.detach().numpy()
#print data.shape
data = data[0, 0, :, :, :]
data = (255 * data[0:5 * img_size[0]]).astype('uint8')
data_img = sitk.GetImageFromArray(data)
sitk.WriteImage(data_img, '/home/tom/result/' + file_name)
#outputs.append(output)
#output = output[0]
#print output.shape
#seg = output.argmax(0)
seg = (seg[0:5 * img_size[0]] * 255).astype('uint8')
seg = seg.astype('float32')
seg = seg / 255.0
seg = np.multiply(data, seg)
result = np.zeros(img_size)
result = seg[0:img_size[0], 0:512, 0:512]
'''
result = result/255
threshold = 0.05
result[result>0.06] = 1
result[result<=0.02] = 0
result = binary_closing(result)
gt = target[0]
gt = gt[0]
gt = gt[0:img_size[0],0:512,0:512]
gt = gt.numpy()
print ("precision:%f",Precision_img(result,gt))
print ("Recall:%f",Recall_img(result,gt))
print ("f1_score:%f",F1_score_img(result,gt))'''
result = (result * 255).astype('uint8')
seg = sitk.GetImageFromArray(result)
sitk.WriteImage(seg, '/home/tom/membrane/' + file_name + 'mem.tif')
print("running time %s" % (time.time() - start_time))
def main(model_path, cell_hist_datadir, prob_map_datadir):
props = readprops(os.path.join(model_path, 'cfg.txt'))
# setup environments and seeds
os.environ['CUDA_VISIBLE_DEVICES'] = props['gpu']
# setup networks
#Network = getattr(models, args.net)
#model = Network(**args.net_params)
model = Modified3DUNet(in_channels=1, n_classes=2, base_n_filter=16)
#load_model
model_file = os.path.join(model_path, 'model_last.tar')
checkpoint = torch.load(model_file,
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
model = model.cuda()
criterion = getattr(criterions, props['criterion'])
num_gpus = len(props['gpu'].split(','))
batch_size = int(props['batch_size']) * num_gpus
workers = int(props['workers']) * num_gpus
# create dataloaders
#Dataset = getattr(datasets, args.dataset)
dset = cell_testing_inter(cell_hist_datadir)
print dset.__len__()
test_loader = DataLoader(dset,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True)
model.eval()
torch.set_grad_enabled(False)
inputs = []
outputs = []
ground_truth = []
for i, sample in enumerate(test_loader):
input = sample['data']
img_size = sample['image_size']
print img_size[0]
file_name = sample['name']
print file_name
file_name = os.path.splitext(file_name[0])[0] # str(file_name[0])
_, _, z, x, y = input.shape
seg = np.zeros((z, x, y))
for j in range(z / 16):
target = sample['seg']
#ground_truth.append(target)
#print file_name[0]
input_temp = input[0, 0, j * 16:(j + 1) * 16].float()
input_temp = input_temp[None, None, ...]
output_temp = nn.parallel.data_parallel(model, input_temp)
output_temp = output_temp.detach().cpu().numpy()
output_temp = output_temp[0]
seg_temp = output_temp.argmax(0)
seg[j * 16:(j + 1) * 16] = seg_temp
data = input.detach().numpy()
#print data.shape
data = data[0, 0, :, :, :]
data = (255 * data[0:5 * img_size[0]]).astype('uint8')
#outputs.append(output)
#output = output[0]
#print output.shape
#seg = output.argmax(0)
prob_map = (seg[0:5 * img_size[0]] * 255).astype('uint8')
prob_map = prob_map.astype('float32')
prob_map = prob_map / 255.0
prob_map = np.multiply(data, prob_map)
prob_map = resize(
prob_map,
(prob_map.shape[0] / 5, prob_map.shape[1], prob_map.shape[2]))
prob_map_img = sitk.GetImageFromArray(prob_map.astype('uint8'))
sitk.WriteImage(prob_map_img,
prob_map_datadir + '/' + file_name + '-prob.tif')
train_dataset = CustomDataset(train_image_paths, train_mask_paths, train=True)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=1,
shuffle=True,
num_workers=12)
test_dataset = CustomDataset(test_image_paths, test_mask_paths, train=False)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=12)
in_channels = 1
n_classes = 2
base_n_filter = 16
model = Modified3DUNet(in_channels, n_classes, base_n_filter).cuda()
# weights = [0.01, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0,10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0,10.0,10.0,10.0,10.0,10.0]
weights = [0.1, 50.0]
class_weights = torch.FloatTensor(weights).cuda()
# loss_function = GeneralizedDiceLoss(weight = class_weights)
loss_function = nn.CrossEntropyLoss(weight=class_weights)
# loss_function = FocalLoss()
optimizer = optim.Adam(model.parameters())
# criterion = FocalLoss(num_class = 2)
epochs = 300
# model = nn.DataParallel(model, NumerofGPU)
for epoch in range(epochs):
model.train()
train_path = os.path.join(root_path, 'training_data')
test_path = os.path.join(
root_path, 'testing_data') # /media/NAS/NAS/tms_data/testing_data
model_save_path = os.path.join('xxx')
print('model_saved_path:', model_save_path)
## loada data
batch_size = 4
## build network
# Loading the model
in_channels = 6
out_channels = 3
base_n_filter = 16
initial_lr = 0.002
net = Modified3DUNet(in_channels, out_channels, base_n_filter)
net = net.float()
# select gpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
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
net = nn.DataParallel(net)
# put the model on GPUs by model.to(device)
net = net.to(device)
optimizer = RD.RAdam(net.parameters(),
lr=initial_lr,
betas=(0.9, 0.999),
eps=1e-08)
# set learning_rate strategry: Decay LR by a factor of 0.1 every 10 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)