def main(): # setting output directory out_dir = cfg.train.out_dir if not os.path.exists(out_dir): os.makedirs(out_dir) else: print('Folder already exists. Are you sure you want to overwrite results?') print('Debug') # put a break point here print('Configuration:') print(cfg) ## Data loaders # Training data loader cfg.train.mode = 'train' ds_train = get_dataset( cfg.train.mode) # validation data loader cfg.train.mode = 'test' ds_test = get_dataset(cfg.train.mode) # ds print('Data loaders have been prepared!') ## Getting the model cfg.train.mode = 'train' # segmentation network net = Unet_class() device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) print('Network loaded. Starting training...') # weights for # building, road, BG my_weight = torch.from_numpy(np.asarray([1, 2, 0.5])).type('torch.cuda.FloatTensor') criterion = torch.nn.CrossEntropyLoss(weight=my_weight ) l1_maps = torch.nn.L1Loss(reduction='sum') # optimizer optim = torch.optim.Adam(list(net.parameters()) , lr=cfg.train.learning_rate, weight_decay=cfg.train.learning_rate_decay) # learning rate scheduler scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=5, gamma=0.5) ep = cfg.train.num_epochs # number of epochs # loss logs loss_train = 9999.0*np.ones(ep) temp_train_loss = 0 loss_val = 9999.0*np.ones(ep) # training the network for epoch in range(ep): running_loss = 0.0 running_ctr = 0 # switch model to training mode, clear gradient accumulators net.train() optim.zero_grad() scheduler.step() # update learning rate t1 = datetime.now() for i, data in enumerate(ds_train, 0): optim.zero_grad() # reading images images = data[0].type('torch.cuda.FloatTensor') # labels labels = data[1].type('torch.cuda.LongTensor') # occluded images occluded_imgs = data[2] # computing unweighted average base_map =0.0 * occluded_imgs[0].type('torch.cuda.FloatTensor') # initialize with zero for j in range(len(occluded_imgs)): base_map = base_map + occluded_imgs[j].type('torch.cuda.FloatTensor') # avoiding inline operation i.e. += base_map = base_map / np.float(len(occluded_imgs)) predicted = net(base_map) loss = criterion(predicted, labels) loss.backward() optim.step() # print statistics running_loss += loss.item() running_ctr += 1 if i %25 ==0: t2 = datetime.now() delta = t2 - t1 t_print = delta.total_seconds() temp_train_loss = running_loss/25.0 print('[%d, %5d out of %5d] loss: %f, time = %f' % (epoch + 1, i + 1, len(ds_train) , running_loss / running_ctr, t_print )) iou_build, iou_road, iou_bg = IoU(predicted, labels) print('building IoU = ' + str(iou_build) + ', road IoU = ' + str(iou_road) + ', background IoU = ' + str(iou_bg) ) basemap_error = l1_maps(base_map, images) print('L1 error (base map, true image) = ' + str(basemap_error.item())) running_loss = 0.0 running_ctr = 0 t1 = t2 # at the end of every epoch, calculating val loss net.eval() val_loss = 0 with torch.no_grad(): for i, data in enumerate(ds_test, 0): # get clean images images = data[0].type('torch.cuda.FloatTensor') # labels labels = data[1].type('torch.cuda.LongTensor') # occluded images occluded_imgs = data[2] base_map = 0.0 * occluded_imgs[0].type('torch.cuda.FloatTensor') # initialize with zero for j in range(len(occluded_imgs)): base_map = base_map + occluded_imgs[j].type( 'torch.cuda.FloatTensor') # avoiding inline operation i.e. += base_map = base_map / np.float(len(occluded_imgs)) predicted = net(base_map) loss = criterion(predicted, labels) # for cross entropy # val loss val_loss += loss.item() # print statistics val_loss = val_loss /len(ds_test) print('End of epoch ' + str(epoch + 1) + '. Val loss is ' + str(val_loss)) print('Following stats are only for the last batch of the test set:') iou_build, iou_road, iou_bg = IoU(predicted, labels) print('building IoU = ' + str(iou_build) + ', road IoU = ' + str(iou_road) + ', background IoU = ' + str( iou_bg)) basemap_error = l1_maps(base_map, images) print('L1 error (base map, true image) = ' + str(basemap_error.item())) # Model check point if val_loss < np.min(loss_val, axis=0): model_path = os.path.join(out_dir, "trained_model_checkpoint.pth") torch.save(net, model_path) # segmentation network # saving losses loss_val[epoch] = val_loss loss_train[epoch] = temp_train_loss temp_train_loss = 0 # setting additive losses to zero print('Training finished') # saving model model_path = os.path.join(out_dir, "trained_model_end.pth") torch.save(net, model_path) # Saving logs log_name = os.path.join(out_dir, "logging.txt") with open(log_name, 'w') as result_file: result_file.write('Logging... \n') result_file.write('Validation loss ') result_file.write(str(loss_val)) result_file.write('\nTraining loss ') result_file.write(str(loss_train)) print('Model saved') # saving loss curves a = loss_val.cpu().detach().numpy() b = loss_train.cpu().detach().numpy() # print(a.shape) print(a[0:epoch]) plt.figure() plt.plot(b[0:epoch]) plt.plot(a[0:epoch]) plt.title('Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend(['Training loss', 'Validation Loss']) fname1 = str('loss.png') plt.savefig(os.path.join(out_dir, fname1), bbox_inches='tight') print('All done!')
def main(): out_dir = cfg.train.out_dir if not os.path.exists(out_dir): raise ValueError( 'The folder does not exist. Make sure to set the correct folder variable cfg.train.out_dir in config.py' ) if os.path.exists(os.path.join(out_dir, 'qual_results')): raise ValueError( 'The validation folder image_results already exists. Delete the folder if those results are not needed' ) else: os.makedirs(os.path.join(out_dir, 'qual_results')) qual_net = torch.load( os.path.join(out_dir, "trained_basemap_checkpoint.pth")) print('Network loaded...') print(cfg) ## Data loader # only test/validation set is needed if eval_potsdam == True: cfg.train.mode = 'test_potsdam' else: cfg.train.mode = 'test' ds_test = get_dataset(cfg.train.mode) print('Data loaders have been prepared!') qual_net.eval() ctr = 0 with torch.no_grad(): for i, data in enumerate(ds_test, 0): # reading clean images images = data[0].type('torch.cuda.FloatTensor') # occluded images occluded_imgs = data[2] # initializing the quality scores of all images q_pre = torch.zeros( occluded_imgs[0].shape[0], len(occluded_imgs), occluded_imgs[0].shape[1], occluded_imgs[0].shape[2]).type('torch.cuda.FloatTensor') for j in range( len(occluded_imgs)): # compute all the quality masks q_now = qual_net( occluded_imgs[j].type('torch.cuda.FloatTensor')) q_pre[:, j, :, :] = q_now[:, 0, :, :] # do the softmax across quality masks dimension q_final = F.softmax(1 * q_pre, dim=1) # make the final basemap base_map = 0.0 * occluded_imgs[0].type( 'torch.cuda.FloatTensor') # initialization with zero for j in range( len(occluded_imgs)): # compute all the quality masks image_now = occluded_imgs[j].type('torch.cuda.FloatTensor') base_map = base_map + q_final[:, j, :, :].view( q_now.shape).permute(0, 2, 3, 1) * image_now # computing unweigted average as baseline average_image = 0.0 * occluded_imgs[0].type( 'torch.cuda.FloatTensor') # initialize with zero for j in range(len(occluded_imgs)): average_image = average_image + occluded_imgs[j].type( 'torch.cuda.FloatTensor' ) # avoiding inline operation i.e. += average_image = average_image / np.float(len(occluded_imgs)) num_fig = np.minimum(base_map.shape[0], 18) plt.ioff() # save results of the last batch for k in range(num_fig): # target output plt.figure() plt.imshow(images[k, :, :, :].detach().cpu().numpy()) plt.axis('off') fname1 = str(str(ctr) + '_target' + '.png') plt.savefig(os.path.join(out_dir, 'qual_results', fname1), bbox_inches='tight') # basemap plt.figure() plt.imshow(base_map[k, :, :, :].detach().cpu().numpy()) plt.axis('off') fname1 = str(str(ctr) + '_out_basemap' + '.png') plt.savefig(os.path.join(out_dir, 'qual_results', fname1), bbox_inches='tight') plt.figure() plt.imshow(base_map[k, :, :, :].detach().cpu().numpy()) plt.axis('off') fname1 = str(str(ctr) + '_out_basemap' + '.png') plt.savefig(os.path.join(out_dir, 'qual_results', fname1), bbox_inches='tight') # baseline plt.figure() plt.imshow(average_image[k, :, :, :].detach().cpu().numpy()) plt.axis('off') fname1 = str(str(ctr) + '_out_average' + '.png') plt.savefig(os.path.join(out_dir, 'qual_results', fname1), bbox_inches='tight') # input images for j in range(len(occluded_imgs)): plt.figure() plt.imshow(occluded_imgs[j][k, :, :, :]) plt.axis('off') fname1 = str(str(ctr) + '_image' + str(j) + '.png') plt.savefig(os.path.join(out_dir, 'qual_results', fname1), bbox_inches='tight') # quality masks for j in range(len(occluded_imgs)): plt.figure() plt.imshow(q_final[k, j, :, :].detach().cpu().numpy()) plt.axis('off') fname1 = str(str(ctr) + '_mask' + str(j) + '.png') plt.savefig(os.path.join(out_dir, 'qual_results', fname1), bbox_inches='tight') ctr += 1
def main(): ROOT_DIR = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) with open(os.path.join(ROOT_DIR, 'config.json')) as json_file: config = json.load(json_file) out_dir = os.path.join(ROOT_DIR, 'models') if not os.path.exists(out_dir): raise ValueError( 'The folder does not exist. Make sure the training has been completed before running this script' ) segment_net = Unet_class() latest_model = max(glob(out_dir + "/*")) trained_model = glob(os.path.join(out_dir, latest_model) + "/*.pth") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # loading segmentation net if len(trained_model) == 2: # segment_net = torch.load(os.path.join(out_dir, latest_model, trained_model[1])) segment_net.load_state_dict( torch.load(os.path.join(out_dir, latest_model, trained_model[1]))) else: # segment_net = torch.load(os.path.join(out_dir, latest_model, trained_model[0])) segment_net.load_state_dict( torch.load(os.path.join(out_dir, latest_model, trained_model[0]))) segment_net.eval() segment_net.to(device) print('Network loaded...') print(config) ## getting the dataset mode = 'test' ds_test = get_dataset(mode, config) print('Data loaders have been prepared!') # Initialize metrics iou_build = 0 iou_road = 0 iou_bg = 0 mIoU = 0 fwIou = 0 acc = 0 with torch.no_grad(): for t in range(repeat_times): # evaluate everything 10 times for i, data in enumerate(ds_test, 0): images = data[0].type('torch.FloatTensor') # reading images # labels labels = data[1].type('torch.LongTensor') # segmentation performance predicted = segment_net(images) i1, i2, i3, i4, i5, i6 = IoU(predicted, labels, extra=True) iou_build += i1 iou_road += i2 iou_bg += i3 mIoU += i4 fwIou += i5 acc += i6 print('Completed ' + str(t) + 'out of ' + str(repeat_times)) # average of segmentation numbers iou_build /= (len(ds_test) * repeat_times) iou_road /= (len(ds_test) * repeat_times) iou_bg /= (len(ds_test) * repeat_times) mIoU /= (len(ds_test) * repeat_times) fwIou /= (len(ds_test) * repeat_times) acc /= (len(ds_test) * repeat_times) print('Building IoU on test set = ' + str(iou_build)) print('Road IoU on test set = ' + str(iou_road)) print('BG IoU on test set = ' + str(iou_bg)) print('mIoU on test set = ' + str(mIoU)) print('Frequency weighted IoU on test set = ' + str(fwIou)) print('Pixel accuracy on test set = ' + str(acc)) fname = os.path.join(ROOT_DIR, 'reports', now.strftime("%Y%m%d-%H%M%S")) if not os.path.exists(fname): os.makedirs(fname) fname = os.path.join(fname, 'eval_results.txt') # saving results on disk with open(fname, 'w') as result_file: result_file.write('Logging... \n') result_file.write('\nBuilding IoU on test set = ') result_file.write(str(iou_build)) result_file.write('\nRoad IoU on test set = ') result_file.write(str(iou_road)) result_file.write('\nBG IoU on test set = ') result_file.write(str(iou_bg)) result_file.write('\nMean IoU on test set = ') result_file.write(str(mIoU)) result_file.write('\nfrequency weighted IoU on test set = ') result_file.write(str(fwIou)) result_file.write('\nPixel accuracy on test set = ') result_file.write(str(acc)) print('All done. Results saved in reports directory')
from torch.autograd import Variable from loss import loss_function from utils import Gaussian_Conv_Update, apply_warp, af_plus_loss, get_grid # setup checkpoint directory out_dir = cfg.train.out_dir if not os.path.exists(out_dir): os.makedirs(out_dir) else: print('Folder already exists. Are you sure you want to overwrite results?') print('Configuration: \n', cfg) # dataloaders train_loader = get_dataset(dataset_name=cfg.data.name, mode='train') test_loader = get_dataset(dataset_name=cfg.data.name, mode='test') print('Data loaders have been prepared!') # networks af_plus = get_network('AF_plus') fds = get_network('FDS') fusion = get_network('GAF') feature_loss = get_network('feature_loss') discrim = get_network('discriminator') gan_loss = get_network('gan_loss') # load weights of a trained AF++ af_plus.load_state_dict( torch.load(os.path.join('checkpoints', 'af_plus_dict.pth'))) af_plus.eval()
def main(): out_dir = cfg.train.out_dir if not os.path.exists(out_dir): raise ValueError( 'The folder does not exist. Make sure to set the correct folder variable cfg.train.out_dir in config.py' ) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if eval_baseline == False: # quality net based method qual_net = torch.load( os.path.join( out_dir, "trained_basemap_checkpoint.pth")) # loading quality net qual_net.eval() qual_net.to(device) # loading segmentation net segment_net = torch.load( os.path.join(out_dir, "trained_model_checkpoint.pth")) segment_net.eval() segment_net.to(device) print('Network loaded...') print(cfg) ## getting the dataset if eval_potsdam == True: cfg.train.mode = 'test_potsdam' # Potsdam else: cfg.train.mode = 'test' # Val set of Berlin ds_test = get_dataset(cfg.train.mode) print('Data loaders have been prepared!') # Metrics for pixel-wise comparison of fused images to respective clean images l1_error_abs = torch.nn.L1Loss( reduction='sum') # This is dependent on image size l1_error_mean = torch.nn.L1Loss( reduction='mean') # This is reported in the paper # Initializae metrics abs_error = 0 mean_error = 0 iou_build = 0 iou_road = 0 iou_bg = 0 mIoU = 0 fwIou = 0 acc = 0 with torch.no_grad(): for t in range(repeat_times): # evaluate everything 10 times for i, data in enumerate(ds_test, 0): images = data[0].type( 'torch.cuda.FloatTensor') # reading images # labels labels = data[1].type('torch.cuda.LongTensor') occluded_imgs = data[2] if eval_baseline == False: q_pre = torch.zeros(occluded_imgs[0].shape[0], len(occluded_imgs), occluded_imgs[0].shape[1], occluded_imgs[0].shape[2]).type( 'torch.cuda.FloatTensor') for j in range(len( occluded_imgs)): # compute all the quality masks q_now = qual_net( occluded_imgs[j].type('torch.cuda.FloatTensor')) q_pre[:, j, :, :] = q_now[:, 0, :, :] # do the softmax across quality masks dimension q_final = F.softmax(q_pre, dim=1) # make the final basemap base_map = 0.0 * occluded_imgs[0].type( 'torch.cuda.FloatTensor') # initialization with zero for j in range( len(occluded_imgs) ): # synthesizing fused image by combining images, weighted by quality scores image_now = occluded_imgs[j].type( 'torch.cuda.FloatTensor') base_map = base_map + q_final[:, j, :, :].view( q_now.shape).permute(0, 2, 3, 1) * image_now if eval_baseline == True: # Evaluating baseline? # the following code is for Baseline (average) ONLY base_map = 0.0 * occluded_imgs[0].type( 'torch.cuda.FloatTensor') # initialize with zero for j in range(len(occluded_imgs)): base_map = base_map + occluded_imgs[j].type( 'torch.cuda.FloatTensor' ) # avoiding inline operation i.e. += base_map = base_map / np.float(len(occluded_imgs)) loss_abs = l1_error_abs(base_map, images) loss_mean = l1_error_mean(base_map, images) abs_error += loss_abs.item() mean_error += loss_mean.item() # segmentation performance predicted = segment_net(base_map) i1, i2, i3, i4, i5, i6 = IoU(predicted, labels, extra=True) iou_build += i1 iou_road += i2 iou_bg += i3 mIoU += i4 fwIou += i5 acc += i6 print('Completed ' + str(t) + 'out of ' + str(repeat_times)) # computing average abs_error /= (len(ds_test) * repeat_times) mean_error /= (len(ds_test) * repeat_times) # average of segmentation numbers iou_build /= (len(ds_test) * repeat_times) iou_road /= (len(ds_test) * repeat_times) iou_bg /= (len(ds_test) * repeat_times) mIoU /= (len(ds_test) * repeat_times) fwIou /= (len(ds_test) * repeat_times) acc /= (len(ds_test) * repeat_times) print('Mean error on test set = ' + str(mean_error)) print('Absolute error on test set = ' + str(abs_error)) print('Building IoU on test set = ' + str(iou_build)) print('Road IoU on test set = ' + str(iou_road)) print('BG IoU on test set = ' + str(iou_bg)) print('mIoU on test set = ' + str(mIoU)) print('Frequency weighted IoU on test set = ' + str(fwIou)) print('Pixel accuracy on test set = ' + str(acc)) if eval_potsdam == True: if eval_baseline: n1 = str('eval_result_Potsdam_baseline_multiple.txt') fname = os.path.join(out_dir, n1) else: n1 = str('eval_result_Potsdam_multiple.txt') fname = os.path.join(out_dir, n1) else: fname = os.path.join(out_dir, 'eval_result_Berlin.txt') # saving results on disk with open(fname, 'w') as result_file: result_file.write('Logging... \n') result_file.write('Mean error on test set = ') result_file.write(str(mean_error)) result_file.write('\nAbsolute error on test set = ') result_file.write(str(abs_error)) result_file.write('\nBuilding IoU on test set = ') result_file.write(str(iou_build)) result_file.write('\nRoad IoU on test set = ') result_file.write(str(iou_road)) result_file.write('\nBG IoU on test set = ') result_file.write(str(iou_bg)) result_file.write('\nMean IoU on test set = ') result_file.write(str(mIoU)) result_file.write('\nfrequency weighted IoU on test set = ') result_file.write(str(fwIou)) result_file.write('\nPixel accuracy on test set = ') result_file.write(str(acc)) print('All done. Results saved in eval_result.txt in output directory')
def main(): ROOT_DIR = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) with open(os.path.join(ROOT_DIR, 'config.json')) as json_file: config = json.load(json_file) out_dir = os.path.join(ROOT_DIR, 'models', now.strftime("%Y%m%d-%H%M%S")) # setting output directory if not os.path.exists(out_dir): os.makedirs(out_dir) else: print('Folder already exists. overwriting the results......') print('Configuration:') print(config['train']) ## Data loaders mode = 'train' # Training data loader ds_train = get_dataset(mode, config) mode = 'test' # validation data loader ds_test = get_dataset(mode, config) print('Data loaders have been prepared!') ## Model mode = 'train' net = Unet_class() device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) print('Network loaded. Starting training...') # weights for # building, road, BG - To balance the classes my_weight = torch.from_numpy(np.asarray([1, 2, 0.5])).type('torch.FloatTensor') criterion = torch.nn.CrossEntropyLoss(weight=my_weight) optim = torch.optim.Adam( net.parameters(), lr=config['train']['learning_rate'], weight_decay=config['train']['learning_rate_decay']) # learning rate scheduler scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=5, gamma=0.5) ep = config['train']['num_epochs'] # number of epochs # loss logs loss_train = 9999.0 * np.ones(ep) temp_train_loss = 0 loss_val = 9999.0 * np.ones(ep) # training the network for epoch in range(ep): running_loss = 0.0 running_ctr = 0 # switch model to training mode, clear gradient accumulators net.train() # scheduler.step() # update learning rate optim.zero_grad() t1 = datetime.now() for i, data in enumerate(ds_train, 0): optim.zero_grad() # reading images images = data[0].type('torch.FloatTensor') # labels labels = data[1].type('torch.LongTensor') predicted = net(images) loss = criterion(predicted, labels) loss.backward() optim.step() # print statistics running_loss += loss.item() running_ctr += 1 if i % 25 == 0: t2 = datetime.now() delta = t2 - t1 t_print = delta.total_seconds() temp_train_loss = running_loss / 25.0 print('[%d, %5d out of %5d] loss: %f, time = %f' % (epoch + 1, i + 1, len(ds_train), running_loss / running_ctr, t_print)) iou_build, iou_road, iou_bg = IoU(predicted, labels) print('building IoU = ' + str(iou_build) + ', road IoU = ' + str(iou_road) + ', background IoU = ' + str(iou_bg)) running_loss = 0.0 running_ctr = 0 t1 = t2 net.eval() scheduler.step() # update learning rate val_loss = 0 with torch.no_grad(): for i, data in enumerate(ds_test, 0): # reading images images = data[0].type('torch.FloatTensor') # labels labels = data[1].type('torch.LongTensor') predicted = net(images) loss = criterion(predicted, labels) # Val loss val_loss += loss.item() # print statistics val_loss = val_loss / len(ds_test) print('End of epoch ' + str(epoch + 1) + '. Val loss is ' + str(val_loss)) print( 'Following stats are only for the last batch of the test set:') iou_build, iou_road, iou_bg = IoU(predicted, labels) print('building IoU = ' + str(iou_build) + ', road IoU = ' + str(iou_road) + ', background IoU = ' + str(iou_bg)) # Model check point if val_loss < np.min(loss_val, axis=0): model_path = os.path.join(out_dir, "trained_model_checkpoint.pth") # torch.save(net, model_path) torch.save(net.state_dict(), model_path) print('Model saved at epoch ' + str(epoch + 1)) # saving losses loss_val[epoch] = val_loss loss_train[epoch] = temp_train_loss temp_train_loss = 0 # setting additive losses to zero print('Training finished') # saving model model_path = os.path.join(out_dir, "trained_model_end.pth") # torch.save(net, model_path) torch.save(net.state_dict(), model_path) print('Model saved') # Saving logs in a text file in the output directory log_name = os.path.join(out_dir, "logging.txt") with open(log_name, 'w') as result_file: result_file.write('Logging... \n') result_file.write('Validation loss ') result_file.write(str(loss_val)) result_file.write('\nTraining loss ') result_file.write(str(loss_train)) # saving loss curves a = loss_val b = loss_train plt.figure() plt.plot(b) plt.plot(a) plt.title('Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend(['Training loss', 'Validation Loss']) fname1 = str('loss.png') plt.savefig(os.path.join(out_dir, fname1), bbox_inches='tight') print('Training finished!!!')