def eval_ssim(args, params): assert args.dataset == 'cityscapes' # not implemented otherwise if args.direction == 'label2photo': ref_dir = os.path.join(params['data_folder']['real'], 'val') else: ref_dir = os.path.join(params['data_folder']['segment'], 'val') if args.model == 'pix2pix': syn_dir = f'/Midgard/home/sorkhei/glow2/checkpoints/cityscapes/256x256/pix2pix/label2photo/eval_{args.sampling_round}/val_imgs/pix2pix_cityscapes_label2photo/test_46/images' ssim_file = f'/Midgard/home/sorkhei/glow2/checkpoints/cityscapes/256x256/pix2pix/label2photo/ssim.txt' file_paths = helper.files_with_suffix(syn_dir, '_leftImg8bit.png') # print(syn_dir) # print('len file paths:', len(file_paths)) # print('is dir:', os.path.isdir('/Midgard/home/sorkhei/glow2/checkpoints/cityscapes/256x256/pix2pix/label2photo')) else: syn_dir = helper.extend_path( helper.compute_paths(args, params)['val_path'], args.sampling_round) ssim_file = os.path.join( helper.compute_paths(args, params)['eval_path'], 'ssim.txt') file_paths = helper.absolute_paths( syn_dir) # the absolute paths of all synthesized images # get the score ssim_score = compute_ssim(file_paths, ref_dir) with open(ssim_file, 'a') as file: file.write(f'Round {args.sampling_round}: {ssim_score}\n') print( f'In [eval_ssim]: writing ssim score of {ssim_score} to: "{ssim_file}"' ) print(f'In [eval_ssim]: all done.')
def evaluate_model(args, params): loader_params = { 'batch_size': params['batch_size'], 'shuffle': False, 'num_workers': params['num_workers'] } batch_size = params['batch_size'] _, _, test_loader = data_handler.init_data_loaders(params, loader_params) total_predicted = np.zeros((batch_size, 14)) total_labels = np.zeros((batch_size, 14)) path_to_load = helper.compute_paths( args, params)['save_path'] # compute path from args and params net = networks.init_unified_net(args.model, params) net, _ = helper.load_model(path_to_load, args.epoch, net, optimizer=None, resume_train=False) print( f'In [evaluate_model]: loading the model done from: "{path_to_load}"') print( f'In [evaluate_model]: starting evaluation with {len(test_loader)} batches' ) with torch.no_grad(): for i_batch, batch in enumerate(test_loader): img_batch = batch['image'].to(device).float() label_batch = batch['label'].to(device).float() pred = net(img_batch) if i_batch > 0: total_predicted = np.append(total_predicted, pred.cpu().detach().numpy(), axis=0) total_labels = np.append(total_labels, label_batch.cpu().detach().numpy(), axis=0) else: total_predicted = pred.cpu().detach().numpy() total_labels = label_batch.cpu().detach().numpy() if i_batch % 50 == 0: print( f'In [evaluate_model]: prediction done for batch {i_batch}' ) results_path = helper.compute_paths(args, params)['results_path'] helper.make_dir_if_not_exists(results_path) # plot roc print(f'In [evaluate_model]: starting plotting ROC...') plotting.plot_roc(total_predicted, total_labels, pathology_names, results_path)
def diverse_samples(args, params): """ command: python3 main.py --local --run diverse --image_name strasbourg_000001_061472 --temp 0.9 --model improved_so_large --last_optim_step 276000 \ --img_size 512 1024 --dataset cityscapes --direction label2photo \ --n_block 4 --n_flow 10 10 10 10 --do_lu --reg_factor 0.0001 --grad_checkpoint """ optim_step = args.last_optim_step temperature = args.temp n_samples = 10 img_size = [512, 1024] n_blocks = args.n_block image_name = args.image_name image_cond_path = f'/local_storage/datasets/moein/cityscapes/gtFine_trainvaltest/gtFine/train/{image_name.split("_")[0]}/{image_name}_gtFine_color.png' image_gt_path = f'/local_storage/datasets/moein/cityscapes/leftImg8bit_trainvaltest/leftImg8bit/train/{image_name.split("_")[0]}/{image_name}_leftImg8bit.png' model_paths = helper.compute_paths(args, params) output_folder = os.path.join(model_paths['diverse_path'], image_name, f'temp={temperature}') print(f'out folder: {output_folder}') helper.make_dir_if_not_exists(output_folder) shutil.copyfile(image_cond_path, os.path.join(output_folder, 'segment.png')) shutil.copyfile(image_gt_path, os.path.join(output_folder, 'real.png')) print('Copy segment and real images: done') model = models.init_model(args, params) model = helper.load_checkpoint(model_paths['checkpoints_path'], optim_step, model, optimizer=None, resume_train=False)[0] step = 2 for i in range(0, 10, step): paths_list = [os.path.join(output_folder, f'sample_{(i + 1) + j}.png') for j in range(step)] experiments.take_multiple_samples(model, n_blocks, temperature, step, img_size, image_cond_path, paths_list)
def evaluate_fcn(args, params): assert args.dataset == 'cityscapes' and params['img_size'] == [ 256, 256 ] # not supported otherwise for now # specify the paths if args.direction == 'label2photo' and args.gt: # evaluating ground-truth images all_paths = { 'resized_path': '/Midgard/home/sorkhei/glow2/data/cityscapes/resized/val', 'eval_results': '/Midgard/home/sorkhei/glow2/gt_eval_results' } else: all_paths = helper.compute_paths(args, params) syn_dir = extend_path(all_paths['val_path'], args.sampling_round) # val_imgs_1, val_imgs_2, ... eval_dir = all_paths['eval_path'] print(f'In [evaluate_city_fcn]: results will be read from: "{syn_dir}"') # evaluation if args.direction == 'label2photo': eval_real_imgs_with_temp( base_data_folder=params['data_folder']['base'], synthesized_dir=syn_dir, save_dir=eval_dir, sampling_round=args.sampling_round) else: eval_segmentations_with_temp( synthesized_dir=syn_dir, reference_dir=params['data_folder']['segment'], base_data_folder=params['data_folder']['base'], save_dir=eval_dir, sampling_round=args.sampling_round) print(f'In [eval_city_with_temp]: evaluation done')
def retrieve_data(sess, hps, args, params): train_tfrecords_file = params['tfrecords_file']['train'] val_tfrecords_file = params['tfrecords_file']['val'] train_iter = data_io.read_tfrecords(train_tfrecords_file, args.dataset, args.direction, hps.batch_size, is_training=True) valid_iter = data_io.read_tfrecords(val_tfrecords_file, args.dataset, args.direction, hps.batch_size, is_training=False) first_batch = make_first_batch(sess, train_iter) dataset_name, direction = args.dataset, args.direction assert dataset_name == 'cityscapes' # get from helper run_mode = 'infer' if args.exp else 'train' if run_mode == 'infer': if args.direction == 'label2photo': conditions_paths = helper.get_all_data_folder_images( path=params['data_folder']['segment'], partition='val', image_type='segment') else: conditions_paths = helper.get_all_data_folder_images( path=params['data_folder']['real'], partition='val', image_type='real') cond_list = create_conditions(conditions_paths, also_visual_grid=False) print( f'In [retrieve_data]: conditions_paths for inference is of len: {len(conditions_paths)}' ) else: conditions_paths = globals.seg_conds_abs_paths if direction == 'label2photo' else globals.real_conds_abs_path cond_list, visual_grid = create_conditions(conditions_paths) # save condition if not exists samples_path = helper.compute_paths(args, params)['samples_path'] helper.make_dir_if_not_exists(samples_path) path = os.path.join(samples_path, 'conditions.png') if not os.path.isfile(path): Image.fromarray(visual_grid).save(path) print(f'In [retrieve_data]: saved conditions to: "{path}"') # make a list of dicts containing both image path and image array conditions = [{ 'image_path': conditions_paths[i], 'image_array': cond_list[i] } for i in range(len(cond_list))] return train_iter, valid_iter, first_batch, conditions
def init_hps_for_dual_glow(args, params): class hps(object): # a helper class just for carrying attributes among functions pass # running params hps.inference = False if args.resume_train or args.exp: hps.restore_path = os.path.join(helper.compute_paths(args, params)['checkpoints_path'], f"step={args.last_optim_step}.ckpt") else: hps.restore_path = None # would be checkpoints path + step when resume training or inference # batch and image size hps.batch_size = 1 hps.input_size = [256, 256, 3] hps.output_size = [256, 256, 3] hps.n_bits_x = 8 # model config hps.n_levels, hps.depth = params['n_block'], params['n_flow'] hps.n_l = 1 # mlp basic layers, default: 1 by the paper hps.flow_permutation = 2 # 0: reverse (RealNVP), 1: shuffle, 2: invconv (Glow)" hps.flow_coupling = 1 # 0: additive, 1: affine hps.width = 512 # Width of hidden layers - default by the paper hps.eps_std = .7 # other model configs hps.ycond = False # Use y conditioning - default by the paper hps.learntop = True # Learn spatial prior hps.n_y = 1 # always 1 in the original code hps.ycond_loss_type = 'l2' # loss type of y inferred from z_in - default by the paper - not used by us as we do not have y conditioning # training config hps.train_its = 2000000 if args.max_step is None else args.max_step # training iterations hps.val_its = 500 # 500 val iterations so we get full validation result with batch size 1 (val set size is 500) hps.val_freq = 1000 # get val result every 1000 iterations hps.sample_freq = 500 if args.sample_freq is None else args.sample_freq hps.direct_iterator = True # default by the paper hps.weight_lambda = 0.001 # Weight of log p(x_o|x_u) in weighted loss, default by the paper hps.weight_y = 0.01 # Weight of log p(y|x) in weighted loss, default by the paper - not used by us as we do not have y conditioning # adam params hps.optimizer = 'adam' hps.gradient_checkpointing = 1 # default hps.beta1 = .9 hps.beta2 = .999 hps.lr = 0.0001 hps.weight_decay = 1. # Switched off by default hps.polyak_epochs = 1 # default by the code - not used by us return hps
def run_training(args, params): # print run info helper.print_info(args, params, model=None, which_info='params') # setting comet tracker tracker = None if args.use_comet: tracker = init_comet(args, params) print("In [run_training]: Comet experiment initialized...") if 'dual_glow' in args.model: models.train_dual_glow(args, params, tracker) else: model = models.init_model(args, params) optimizer = optim.Adam(model.parameters()) reverse_cond = data_handler.retrieve_rev_cond(args, params, run_mode='train') train_configs = trainer.init_train_configs(args) # resume training if args.resume_train: optim_step = args.last_optim_step checkpoints_path = helper.compute_paths(args, params)['checkpoints_path'] model, optimizer, _, lr = load_checkpoint(checkpoints_path, optim_step, model, optimizer) if lr is None: # if not saved in checkpoint lr = params['lr'] trainer.train(args, params, train_configs, model, optimizer, lr, tracker, resume=True, last_optim_step=optim_step, reverse_cond=reverse_cond) # train from scratch else: lr = params['lr'] trainer.train(args, params, train_configs, model, optimizer, lr, tracker, reverse_cond=reverse_cond)
def train(args, params, model, optimizer, tracker=None): loader_params = { 'batch_size': params['batch_size'], 'shuffle': params['shuffle'], 'num_workers': params['num_workers'] } train_loader, val_loader, _ = data_handler.init_data_loaders( params, loader_params) epoch = 0 max_epochs = params['max_epochs'] while epoch < max_epochs: print(f'{"=" * 40} In epoch: {epoch} {"=" * 40}') print(f'Training on {len(train_loader)} batches...') # each for loop for one epoch for i_batch, batch in enumerate(train_loader): # converting the labels batch to from Long tensor to Float tensor (otherwise won't work on GPU) img_batch = batch['image'].to(device).float() label_batch = batch['label'].to(device).float() # making gradients zero in each optimization step optimizer.zero_grad() # getting the network prediction and computing the loss pred = model(img_batch) train_loss = loss.compute_wcel(fx=pred, labels=label_batch) # if i_batch % 50 == 0: print( f'Batch: {i_batch} - train loss: {round(train_loss.item(), 3)}' ) # tracking the metrics using comet in each iteration track('train_loss', round(train_loss.item(), 3), args, tracker) # backward and optimization step train_loss.backward() optimizer.step() # save checkpoint after epoch save_path = helper.compute_paths(args, params)['save_path'] helper.make_dir_if_not_exists(save_path) helper.save_model(save_path, epoch, model, optimizer, train_loss) # save the model every epoch val_loss = loss.compute_val_loss(model, val_loader) # track val loss print(f'In [train]: epoch={epoch} - val_loss = {val_loss}') track('val_loss', val_loss, args, tracker) epoch += 1
def create_rev_cond(args, params, fixed_conds, also_save=True): trans = loader.init_transformers(params['img_size']) # only supports reading from train data now images = read_imgs(params['data_folder'], split='train', fixed_conds=fixed_conds) conds_as_list = [] real_as_list = [] for image in images: h, w = params['img_size'] # get half of the image corresponding to either the photo or the map photo = trans( Image.open(image) )[:, :, : w] # if args.direction == 'photo2map' else trans(Image.open(image))[:, :, w:] the_map = trans(Image.open(image))[:, :, w:] if args.direction == 'map2photo': cond = the_map real = photo else: cond = photo real = the_map conds_as_list.append(cond) real_as_list.append(real) conds_as_tensor = torch.stack(conds_as_list, dim=0).to(device) real_as_tensor = torch.stack(real_as_list, dim=0).to(device) if also_save: save_path = helper.compute_paths(args, params)['samples_path'] helper.make_dir_if_not_exists(save_path) if 'conditions.png' not in os.listdir(save_path): utils.save_image(conds_as_tensor, os.path.join(save_path, 'conditions.png')) print(f'In [create_rev_cond]: saved reverse conditions') if 'real.png' not in os.listdir(save_path): utils.save_image(real_as_tensor, os.path.join(save_path, 'real.png')) print(f'In [create_rev_cond]: saved real images') print( f'In [create_rev_cond]: returning cond_as_tensor of shape: {conds_as_tensor.shape}' ) return conds_as_tensor
def init_and_load(args, params, run_mode): checkpoints_path = helper.compute_paths(args, params)['checkpoints_path'] optim_step = args.last_optim_step model = init_model(args, params) if run_mode == 'infer': model, _, _, _ = helper.load_checkpoint(checkpoints_path, optim_step, model, None, resume_train=False) print(f'In [init_and_load]: returned model for inference') return model else: # train optimizer = optim.Adam(model.parameters(), lr=params['lr']) print(f'In [init_and_load]: returned model and optimizer for training') model, optimizer, _, lr = helper.load_checkpoint(checkpoints_path, optim_step, model, optimizer, resume_train=True) return model, optimizer, lr
def prepare_city_reverse_cond(args, params, run_mode='train'): samples_path = helper.compute_paths(args, params)['samples_path'] save_path = samples_path if run_mode == 'train' else None # no need for reverse_cond at inference time direction = args.direction segmentations, _, real_imgs, boundaries = _create_cond( args, params, fixed_conds=real_conds_abs_path, save_path=save_path, direction=direction) return { 'segment': segmentations, 'boundary': boundaries, 'real': real_imgs }
def create_rev_cond(args, params, also_save=True): trans = util.init_transformer(params['img_size']) # if args.direction == 'daylight2night': names = util.fixed_conds[args.direction] conds = [trans(Image.open(f"{params['paths']['data_folder']}/{name}")) for name in names] # else: # raise NotImplementedError conditions = torch.stack(conds, dim=0).to(device) # (B, C, H, W) if also_save: save_path = helper.compute_paths(args, params)['samples_path'] helper.make_dir_if_not_exists(save_path) if 'conditions.png' not in os.listdir(save_path): utils.save_image(conditions, f'{save_path}/conditions.png') print(f'In [create_rev_cond]: saved reverse conditions') return conditions
def compute_ssim_old(args, params): """ images should be np arrays of shape (H, W, C). This function should be called only after the inference has been done on validation set. This function computes SSIM for the temperature specified in params['temperature']. :return: """ # =========== init validation loader loader_params = {'batch_size': 1, 'shuffle': False, 'num_workers': 0} _, val_loader = data_handler.init_city_loader( data_folder=params['data_folder'], image_size=(params['img_size']), remove_alpha=True, # removing the alpha channel loader_params=loader_params) print( f'In [compute_ssim]: init val data loader of len {len(val_loader)}: done \n' ) # ============= computing the validation path for generated images if params['img_size'] != [256, 256]: raise NotImplementedError( 'Now only supports 256x256 images' ) # should use paths['resized_path'] probably? paths = helper.compute_paths(args, params) val_path = paths['val_path'] print(f'In [compute_ssim]: val_path to read from: \n"{val_path}" \n') # ============= SSIM calculation for all the images in validation set ssim_vals = [] for i_batch, batch in enumerate( val_loader): # for every single image (batch size 1) # reading reference (ground truth) and generated image ref_img = batch['real'].cpu().data.squeeze(dim=0).permute( 1, 2, 0).numpy() # (C, H, W) -> (H, W, C) pure_name = batch['real_path'][0].split('/')[ -1] # the city name of the image # compute corresponding path in generated images and read the image gen_img_path = val_path + f'/{pure_name}' gen_img = city_transforms(Image.open(gen_img_path)).cpu().data.permute( 1, 2, 0).numpy() # ============= computing SSIM pixel_range = 1 # with data range 1 since pixels can take a value between 0 to 1 ssim_val = ssim(ref_img, gen_img, multichannel=True, data_range=pixel_range) # expects (H, W, C) ordering ssim_vals.append(ssim_val) if i_batch % 100 == 0: print(f'In [compute_ssim]: evaluated {i_batch} images') # ssim_score = round(np.mean(ssim_vals), 2) ssim_score = np.mean(ssim_vals) print(f'In [compute_ssim]: ssim score on validation set: {ssim_score}') # ============= append result to ssim.txt eval_path_base = paths['eval_path_base'] with open(f'{eval_path_base}/ssim.txt', 'a') as ssim_file: string = f'temp = {params["temperature"]}: {ssim_score} \n' ssim_file.write(string) print(f'In [compute_ssim]: ssim score appended to ssim.txt')
def train(args, params, train_configs, model, optimizer, current_lr, comet_tracker=None, resume=False, last_optim_step=0, reverse_cond=None): # getting data loaders train_loader, val_loader = data_handler.init_data_loaders(args, params) # adjusting optim step optim_step = last_optim_step + 1 if resume else 1 max_optim_steps = params['iter'] paths = helper.compute_paths(args, params) if resume: print( f'In [train]: resuming training from optim_step={optim_step} - max_step: {max_optim_steps}' ) # optimization loop while optim_step < max_optim_steps: # after each epoch, adjust learning rate accordingly current_lr = adjust_lr( current_lr, initial_lr=params['lr'], step=optim_step, epoch_steps=len( train_loader)) # now only supports with batch size 1 for param_group in optimizer.param_groups: param_group['lr'] = current_lr print( f'In [train]: optimizer learning rate adjusted to: {current_lr}\n') for i_batch, batch in enumerate(train_loader): if optim_step > max_optim_steps: print( f'In [train]: reaching max_step or lr is zero. Terminating...' ) return # ============ terminate training if max steps reached begin_time = time.time() # forward pass left_batch, right_batch, extra_cond_batch = data_handler.extract_batches( batch, args) forward_output = forward_and_loss(args, params, model, left_batch, right_batch, extra_cond_batch) # regularize left loss if train_configs['reg_factor'] is not None: loss = train_configs['reg_factor'] * forward_output[ 'loss_left'] + forward_output['loss_right'] # regularized else: loss = forward_output['loss'] metrics = {'loss': loss} # also add left and right loss if available if 'loss_left' in forward_output.keys(): metrics.update({ 'loss_right': forward_output['loss_right'], 'loss_left': forward_output['loss_left'] }) # backward pass and optimizer step model.zero_grad() loss.backward() optimizer.step() print(f'In [train]: Step: {optim_step} => loss: {loss.item():.3f}') # validation loss if (params['monitor_val'] and optim_step % params['val_freq'] == 0) or current_lr == 0: val_loss_mean, _ = calc_val_loss(args, params, model, val_loader) metrics['val_loss'] = val_loss_mean print( f'====== In [train]: val_loss mean: {round(val_loss_mean, 3)}' ) # tracking metrics if args.use_comet: for key, value in metrics.items(): # track all metric values comet_tracker.track_metric(key, round(value.item(), 3), optim_step) # saving samples if (optim_step % params['sample_freq'] == 0) or current_lr == 0: samples_path = paths['samples_path'] helper.make_dir_if_not_exists(samples_path) sampled_images = models.take_samples(args, params, model, reverse_cond) utils.save_image( sampled_images, f'{samples_path}/{str(optim_step).zfill(6)}.png', nrow=10) print( f'\nIn [train]: Sample saved at iteration {optim_step} to: \n"{samples_path}"\n' ) # saving checkpoint if (optim_step > 0 and optim_step % params['checkpoint_freq'] == 0) or current_lr == 0: checkpoints_path = paths['checkpoints_path'] helper.make_dir_if_not_exists(checkpoints_path) helper.save_checkpoint(checkpoints_path, optim_step, model, optimizer, loss, current_lr) print("In [train]: Checkpoint saved at iteration", optim_step, '\n') optim_step += 1 end_time = time.time() print(f'Iteration took: {round(end_time - begin_time, 2)}') helper.show_memory_usage() print('\n') if current_lr == 0: print( 'In [train]: current_lr = 0, terminating the training...') sys.exit(0)
def infer_on_set(args, params, split_set='val'): """ The model name and paths should be equivalent to the name used in the resize_for_fcn function in evaluation.third_party.prepare.py module. :param split_set: :param args: :param params: :return: """ with torch.no_grad(): paths = helper.compute_paths(args, params) checkpt_path, val_path, train_vis_path = paths[ 'checkpoints_path'], paths['val_path'], paths['train_vis'] save_path = val_path if split_set == 'val' else train_vis_path save_path = helper.extend_path(save_path, args.sampling_round) print( f'In [infer_on_validation_set]:\n====== checkpt_path: "{checkpt_path}" \n====== save_path: "{save_path}" \n' ) helper.make_dir_if_not_exists(save_path) # init and load model model = models.init_and_load(args, params, run_mode='infer') print( f'In [infer_on_validation_set]: init model and load checkpoint: done' ) # validation loader batch_size = params['batch_size'] loader_params = { 'batch_size': batch_size, 'shuffle': False, 'num_workers': 0 } train_loader, val_loader = data_handler.init_city_loader( data_folder=params['data_folder'], image_size=(params['img_size']), loader_params=loader_params) loader = val_loader if split_set == 'val' else train_loader print('In [infer_on_validation_set]: using loader of len:', len(loader)) # inference on validation set print( f'In [infer_on_validation_set]: starting inference on {split_set} set' ) for i_batch, batch in enumerate(loader): img_batch = batch['real'].to(device) segment_batch = batch['segment'].to(device) boundary_batch = batch['boundary'].to( device) if args.use_bmaps else None real_paths = batch[ 'real_path'] # list: used to save samples with the same name as original images seg_paths = batch['segment_path'] if args.direction == 'label2photo': reverse_cond = { 'segment': segment_batch, 'boundary': boundary_batch } elif args.direction == 'photo2label': # 'photo2label' reverse_cond = {'real': img_batch} # sampling from model # using batch.shape[0] is safer than batch_size since the last batch may be different in size samples = models.take_samples(args, params, model, reverse_cond, n_samples=segment_batch.shape[0]) # save inferred images separately paths_list = real_paths if args.direction == 'label2photo' else seg_paths helper.save_one_by_one(samples, paths_list, save_path) print( f'In [infer_on_validation_set]: done for the {i_batch}th batch out of {len(loader)} batches (batch size: {segment_batch.shape[0]})' ) print( f'In [infer_on_validation_set]: all done. Inferred images could be found at: {save_path} \n' )
def run_model(mode, args, params, hps, sess, model, conditions, tracker): sess.graph.finalize() # inference on validation set if mode == 'infer': # val_path = helper.compute_paths(args, params)['val_path'] paths = helper.compute_paths(args, params) val_path = helper.extend_path(paths['val_path'], args.sampling_round) helper.make_dir_if_not_exists(val_path) take_sample(model, conditions, val_path, mode, direction=args.direction, iteration=None) print(f'In [run_model]: validation samples saved to: "{val_path}"') # training else: paths = helper.compute_paths(args, params) checkpoint_path = paths['checkpoints_path'] samples_path = paths['samples_path'] helper.make_dir_if_not_exists(checkpoint_path) helper.make_dir_if_not_exists(samples_path) # compute_conditional_bpd(model, args.last_optim_step, hps) iteration = 0 if not args.resume_train else args.last_optim_step + 1 while iteration <= hps.train_its: lr = hps.lr train_results = model.train( lr) # returns [local_loss, bits_x_u, bits_x_o, bits_y] train_loss = round(train_results[0], 3) print(f'Step {iteration} - train loss: {train_loss}') # take sample if iteration % hps.sample_freq == 0: take_sample(model, conditions, samples_path, mode, iteration) # track train loss if tracker is not None: tracker.track_metric('train_loss', round(train_loss, 3), iteration) # compute val loss if iteration % hps.val_freq == 0: test_results = [] print('Computing validation loss...') for _ in range(hps.val_its ): # one loop over all all validation examples test_results += [model.test()] test_results = np.mean(np.asarray(test_results), axis=0) # get the mean of val loss val_loss = round(test_results[0], 3) print(f'Step {iteration} - validation loss: {val_loss}') # save checkpoint path = os.path.join(checkpoint_path, f"step={iteration}.ckpt") model.save(path) print(f'Checkpoint saved to: "{path}"') # track val loss if tracker is not None: tracker.track_metric('val_loss', round(val_loss, 3), iteration) iteration += 1
def evaluate_iobb(args, params, img_name=None, img_disease=None): """ Goes through every image in the BBox csv file, and calculates IoBB. Results are stored in data/iobb.npy together with image name and disease type. :param img_name: If given, only load one image and plot bbox and heatmap. :param img_disease: Which disease to plot bbox for, used together with img_name. """ pathologies = { 'Atelectasis': 0, 'Cardiomegaly': 1, 'Effusion': 4, 'Infiltrate': 8, 'Mass': 9, 'Nodule': 10, 'Pneumonia': 12, 'Pneumothorax': 13 } ############################################ ### Load model checkpoint and get heatmap. ############################################ path_to_load = helper.compute_paths(args, params)['save_path'] net = networks.init_unified_net(args.model, params) net, _ = helper.load_model(path_to_load, args.epoch, net, optimizer=None, resume_train=False) # Get bbox_data from csv file. f = open('../data/BBox_List_2017.csv', 'rt') reader = csv.reader(f) rows = list(reader)[1:] # ignoring the first row because it is the titles # A list of tuples (img_name, disease_index, iobb, num_bboxes) results = [] for i, img_data in enumerate(rows): # Make sure image exists. file_path = f'../data/extracted/images/{img_data[0]}' if not os.path.isfile(file_path): continue # If only loading one image, check if this row contains the img and correct disease, otherwise continue. if img_name is not None: if img_data[0] != img_name: continue if img_disease is not None and pathologies[ img_data[1]] != img_disease: continue ############################################ ### Load image and turn into tensor. ############################################ xray_img = Image.open(file_path) rgb = Image.new('RGB', xray_img.size) rgb.paste(xray_img) preprocess = transforms.Compose([ transforms.Resize(256), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) img_tensor = preprocess(rgb) img_tensor.unsqueeze_(0) # Get heatmap for correct disease and turn into numpy array. heatmap = net.forward(img_tensor, return_cam=True) disease = pathologies[img_data[1]] heatmap = heatmap[0, disease].numpy() ground_truth = BBox(float(img_data[2]), float(img_data[3]), float(img_data[4]), float(img_data[5])) # Save results if evaluating all images, not just plotting one. if img_name is None: iobb, num = evaluate_single_bbox(ground_truth, heatmap, iobb=True) results.append((img_data[0], disease, iobb, num)) print(f'#{i}, n:{num}, {iobb}') else: iobb, _ = evaluate_single_bbox(ground_truth, heatmap, iobb=True, xray_img=xray_img) print(f'iobb: {iobb}') break if img_name is None: # Order results by IoBB value. results = sorted(results, key=lambda x: x[2], reverse=True) results = np.array(results) # Save as numpy array. np.save(f'../data/iou_blob_{args.model}.npy', results) # Save as txt with open(f'../data/iou_blob_{args.model}.txt', 'w') as txt: for i in range(results.shape[0]): txt.write( f'{float(results[i][2])}, {int(results[i][1])}, {int(results[i][3])}, {results[i][0]}\n' ) f.close() # Close csv file.
def prep_for_sampling(args, params, img_name, additional_info): """ :param args: :param params: :param img_name: the (path of the) real image whose segmentation which will be used for conditioning. :param additional_info: :return: """ # ========== specifying experiment path paths = helper.compute_paths(args, params, additional_info) if additional_info['exp_type'] == 'random_samples': experiment_path = paths['random_samples_path'] elif additional_info['exp_type'] == 'new_cond': experiment_path = paths['new_cond_path'] else: raise NotImplementedError # ========== make the condition a single image fixed_conds = [img_name] # ========== create condition and save it to experiment path # no need to save for new_cond type path_to_save = None if additional_info[ 'exp_type'] == 'new_cond' else experiment_path seg_batch, _, real_batch, boundary_batch = data_handler._create_cond( params, fixed_conds=fixed_conds, save_path=path_to_save) # (1, C, H, W) # ========== duplicate condition for n_samples times (not used by all exp_modes) seg_batch_dup = seg_batch.repeat( (params['n_samples'], 1, 1, 1)) # duplicated: (n_samples, C, H, W) boundary_dup = boundary_batch.repeat((params['n_samples'], 1, 1, 1)) real_batch_dup = real_batch.repeat((params['n_samples'], 1, 1, 1)) if additional_info['exp_type'] == 'random_samples': seg_rev_cond = seg_batch_dup # (n_samples, C, H, W) - duplicate for random samples bmap_rev_cond = boundary_dup real_rev_cond = real_batch_dup elif additional_info['exp_type'] == 'new_cond': seg_rev_cond = seg_batch # (1, C, H, W) - no duplicate needed bmap_rev_cond = boundary_batch else: raise NotImplementedError # ========== create reverse cond if not args.use_bmaps: rev_cond = {'segment': seg_rev_cond, 'boundary': None} elif args.use_bmaps: rev_cond = {'segment': seg_rev_cond, 'boundary': bmap_rev_cond} elif args.direction == 'photo2label': rev_cond = {'real': real_rev_cond} else: raise NotImplementedError # ========== specifying paths for saving samples if additional_info['exp_type'] == 'random_samples': exp_path = paths['random_samples_path'] save_paths = [ f'{exp_path}/sample {i + 1}.png' for i in range(params['n_samples']) ] elif additional_info['exp_type'] == 'new_cond': save_paths = experiment_path else: raise NotImplementedError return save_paths, rev_cond, real_batch