def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
unet = builder.build_unet(num_class=args.num_class,
arch=args.arch_unet,
weights=args.weights_unet1)
crit = DualLoss(mode='val')
sm = SegmentationModule(crit, unet)
test_augs = ComposeTest([PaddingCenterCropTest(256)])
ac17 = AC17(root=args.data_root,
augmentations=test_augs,
img_norm=args.img_norm)
loader_val = data.DataLoader(ac17,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
sm.cuda()
# Main loop
evaluate(sm, loader_val, args)
print('Evaluation Done!')
def main(args):
# Network Builders
builder = ModelBuilder()
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = DualLoss(mode="train")
segmentation_module = SegmentationModule(crit, unet)
test_augs = Compose([PaddingCenterCrop(256)])
print("ready to load data")
dataset_val = LungData(
root=args.data_root,
split='test',
k_split=args.k_split,
augmentations=test_augs)
loader_val = data.DataLoader(
dataset_val,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
print(len(loader_val))
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit) if args.unet == False else (unet, crit)
optimizers = create_optimizers(nets, args)
'''
# Start the webapp: user update a dcm file, output the predicted segmentation pic of it
inp = gradio.inputs.DcmUpload(preprocessing_fn=preprocess)
#inp = gradio.inputs.ImageUpload(preprocessing_fn=preprocess)
io = gradio.Interface(inputs=inp, outputs="image", model_type="lung_seg", model=segmentation_module, args=args)
io.launch(validate=False)
'''
iou, loss = eval(loader_val, segmentation_module, args, crit)
print('Evaluation Done!')
''''''
#build model
builder = ModelBuilder()
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = DualLoss(mode="train")
segmentation_module = SegmentationModule(crit, unet)
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
print("ready to load data")
import train
def main(args):
# Network Builders
builder = ModelBuilder()
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = DualLoss(mode="train")
segmentation_module = SegmentationModule(crit, unet)
train_augs = Compose([PaddingCenterCrop(256), RandomHorizontallyFlip(), RandomVerticallyFlip(), RandomRotate(180)])
test_augs = Compose([PaddingCenterCrop(256)])
# Dataset and Loader
dataset_train = AC17( #Loads 3D volumes
root=args.data_root,
split='train',
k_split=args.k_split,
augmentations=train_augs,
img_norm=args.img_norm)
ac17_train = load2D(dataset_train, split='train', deform=True) #Dataloader for 2D slices. Requires 3D loader.
loader_train = data.DataLoader(
ac17_train,
batch_size=args.batch_size_per_gpu,
shuffle=True,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
dataset_val = AC17(
root=args.data_root,
split='val',
k_split=args.k_split,
augmentations=test_augs,
img_norm=args.img_norm)
ac17_val = load2D(dataset_val, split='val', deform=False)
loader_val = data.DataLoader(
ac17_val,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit) if args.unet == False else (unet, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': [], 'jaccard': []}}
best_val = {'epoch_1': 0, 'mIoU_1': 0,
'epoch_2': 0, 'mIoU_2': 0,
'epoch_3': 0, 'mIoU_3': 0,
'epoch' : 0, 'mIoU': 0}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, loader_train, optimizers, history, epoch, args)
iou, loss = eval(loader_val, segmentation_module, args, crit)
#checkpointing
ckpted = False
if loss < 0.215:
ckpted = True
if iou[0] > best_val['mIoU_1']:
best_val['epoch_1'] = epoch
best_val['mIoU_1'] = iou[0]
ckpted = True
if iou[1] > best_val['mIoU_2']:
best_val['epoch_2'] = epoch
best_val['mIoU_2'] = iou[1]
ckpted = True
if iou[2] > best_val['mIoU_3']:
best_val['epoch_3'] = epoch
best_val['mIoU_3'] = iou[2]
ckpted = True
if (iou[0]+iou[1]+iou[2])/3 > best_val['mIoU']:
best_val['epoch'] = epoch
best_val['mIoU'] = (iou[0]+iou[1]+iou[2])/3
ckpted = True
if epoch % 50 == 0:
checkpoint(nets, history, args, epoch)
continue
if epoch == args.num_epoch:
checkpoint(nets, history, args, epoch)
continue
if epoch < 15:
ckpted = False
if ckpted == False:
continue
else:
checkpoint(nets, history, args, epoch)
continue
print()
print('Training Done!')