def view_results():
model = torch.load('../temp/exp_results/model_epoch_100.pkl', map_location='cpu')
config = json.load(open('../config.json'))
model.eval()
data_manager = DataManager(config)
val_loader, _ = data_manager.get_train_eval_dataloaders()
for idx, (inputs, labels) in enumerate(val_loader, 0):
print(idx)
if config['use_cuda']:
inputs = inputs.cuda().float()
labels = labels.cuda().float()
model = model.cuda()
else:
inputs = inputs.float()
labels = labels.float()
predictions = model(inputs)
predictions = predictions.data.cpu().numpy()
print('predictions', np.unique(predictions[0]))
labels = labels.cpu().numpy()
print('labels', np.unique(labels[0]))
mask = create_mask(labels[0], palette=config['palette'])
img = Image.fromarray(mask.astype(np.uint8))
img.save('gt' + str(idx) + '.jpg')
predictions = np.argmax(predictions[0], axis=0)
mask = create_mask(predictions, palette=config['palette'])
img = Image.fromarray(mask.astype(np.uint8))
img.save('pred_' + '_' + str(idx) + '.jpg')
def main():
config = json.load(open('./config.json'))
experiment_name = 'Fully Conv'
vis_legend = ['Training Loss', 'Eval Loss']
visualizer = VisdomVisualizer(experiment_name, vis_legend, config=config)
model = UNetV3(n_channels=config['n_channels'], n_classes=config['n_classes'])
model.apply(init_weights)
weights = torch.ones(21)
if config['use_cuda'] is True:
weights = weights.cuda()
weights[0] = 0.07
weights[15] = 0.5
criterion = SegmentationLosses(weight=weights).build_loss(mode='ce')
optimizer = optim.SGD(model.parameters(), lr=0.00001, momentum=0.9)
data_manager = DataManager(config)
train_loader, val_loader = data_manager.get_train_eval_dataloaders()
trainer = Trainer(model, train_loader, val_loader, criterion, optimizer, visualizer, experiment_name, config)
trainer.train()