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()