def write_imgs(model, dset, device, s_batch=1, model_name='model'): """ Writing collages to 'show_cases/predictions'. A collage feature the original image the L channel and the predicted image in that order from left to right. :param model: model to be used :param dset: data set, containing images in 'show_cases/imgs' :param device: GPU device :param s_batch: size of batch :param model_name: name of model """ data_iter = data.RAMDataSetIter(dset, s_batch=s_batch, shuffle=False) with torch.no_grad(): model.to(device) model.eval() it_count = 0 prog_bar = tqdm.tqdm(total=len(dset), bar_format=f'{Fore.GREEN}{model_name} show case images {Fore.RESET} ' + "{l_bar}%s{bar}%s{r_bar}" % (Fore.GREEN, Fore.RESET)) while True: try: batch, targets = data_iter.next() except StopIteration: break batch = batch.to(device) targets = targets.to(device) y = torch.nn.functional.hardtanh(model(batch), min_val=0, max_val=1).cpu().data.numpy() paths = dset.paths[it_count:it_count+len(y)] for i, pred_ab in enumerate(y): # lightness channel L = batch[i].cpu().data.numpy() bw = np.repeat(L*255., 3, axis=0).transpose(1, 2, 0).astype(np.uint8) # true a, b channels img_ab = targets[i].cpu().data.numpy() # ture image img_Lab = np.vstack([L, img_ab]) * 255. img_Lab = img_Lab.transpose(1, 2, 0).astype(np.uint8) img_BGR = cv2.cvtColor(img_Lab, cv2.COLOR_LAB2BGR) # estimated image pred_Lab = np.vstack([L, pred_ab]) * 255. pred_Lab = pred_Lab.transpose(1, 2, 0).astype(np.uint8) pred_BGR = cv2.cvtColor(pred_Lab, cv2.COLOR_LAB2BGR) # summary: ture image, lightness, estimated image summary = np.hstack([img_BGR, bw, pred_BGR]) # channels: true channels, estimated channels channels = np.hstack([np.vstack(img_ab), np.vstack(pred_ab)]) * 255. channels = channels.astype(np.uint8) # saving summary and channels cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}.png'), summary) cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}_channels.png'), channels) it_count += 1 prog_bar.update(len(y)) model.cpu() batch.cpu() targets.cpu() prog_bar.close()
def __init__(self, model, dset, device, s_batch=32, model_name='model'): """ :param model: ConvNet model :param dset: data set to use in training :param device: GPU device :param s_batch: batch size :param model_name: name of model """ self.model = model self.dset = dset self.device = device self.s_batch = s_batch self.model_name = model_name self.print_c = Fore.GREEN if not os.path.isdir(join(f'logs/{self.model_name}')): os.makedirs(join(f'logs/{self.model_name}'))
def plot_progress(self, error_series, plt_imgs, plt_imgs_t): """ This method generates an informative plot, displaying the loss over time along logged images. :param error_series: loss values over time (/ steps) :param plt_imgs: logged images :param plt_imgs_t: time (/ step) at which respective images have been logged :return: """ fig, ax = plt.subplots(figsize=(160, 40)) ax.plot(np.arange(1, 1 + len(error_series)), error_series, alpha=0.9) for i, p_I in enumerate(plt_imgs): imagebox = OffsetImage(p_I, zoom=1.) x = plt_imgs_t[i] ab = AnnotationBbox(imagebox, (x, error_series[x - 1]), xybox=(0., 256.), xycoords='data', boxcoords="offset points", arrowprops={'arrowstyle': '->'}) ax.add_artist(ab) plt.savefig(join(f'logs/{self.model_name}/progress.png'), bbox_inches='tight', dpi=50) plt.close(fig)
hist = np.mean(hists, axis=0) return bounds, hist if __name__ == '__main__': # establishing GPU device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # evaluating the models accs = {} print( f'\n{Fore.BLUE} evaluating all models stored in {Fore.WHITE}models/{Fore.RESET}' ) for model_name in sorted( [p for p in os.listdir('models') if os.path.isdir(f'models/{p}')]): # loading the model and corresponding dataset with open(join(f'models/{model_name}/kwargs.json'), 'r') as f_kwargs: kwargs = json.load(f_kwargs) state_dict = torch.load(f'models/{model_name}/state_dict.pth', map_location=torch.device('cpu')) dset = data.RAMDataSet(**kwargs['dset']) model = nn.ConvNet(**kwargs['model']) model.load_state_dict(state_dict) # initializing the evaluator evaluator = Evaluator(model, dset, device, model_name=model_name) # storing and saving some discrete upper bounds, for a quick look/ preview accs[model_name] = evaluator.evaluate() print( f'{Fore.BLUE}accuracy {model_name}: {Fore.WHITE}{accs[model_name]}{Fore.RESET}' ) # constructing a plot visualizing the distribution of abs. differences, that lie within certain ranges bounds, hist = evaluator.evaluate_hist()
def log_step(self, y, batch, targets, str_it_count): """ Saving a sample of the batch. :param y: predicted ab channels :param batch: true L channels :param targets: true ab channels :param str_it_count: prefix of the saved images (embodies the current optimization step) :return: """ # index of sample for displayal k = np.random.randint(0, y.shape[0]) # BW images if y.shape[1] == 1: img = y[k].cpu().numpy()[0] img = np.repeat(img[:, :, None], 3, axis=2) orig = targets[k].cpu().data.numpy()[0] orig = np.repeat(orig[:, :, None], 3, axis=2) # logging images cv2.imwrite(join(f'logs/{self.model_name}/{str_it_count}_est.png'), img) cv2.imwrite( join(f'logs/{self.model_name}/{str_it_count}_orig.png'), orig) # assemble the plot iamge img = cv2.resize(img, (128, 128)) orig = cv2.resize(orig, (128, 128)) p_I = np.vstack([orig, img]) # Lab images else: ab = y[k].cpu().data.numpy() orig_ab = targets[k].cpu().data.numpy() L = batch[k].cpu().data.numpy() # estimated image and fmaps est_Lab = np.vstack([L, ab]) * 255. est_RGB = cv2.cvtColor( est_Lab.transpose(1, 2, 0).astype(np.uint8), cv2.COLOR_Lab2RGB) / 255. est_Lab = np.hstack(est_Lab[1:]) / 255. est_ab_BW = np.repeat(est_Lab[:, :, None], 3, axis=2) # original image and fmaps orig_Lab = np.vstack([L, orig_ab]) * 255. orig_RGB = cv2.cvtColor( orig_Lab.transpose(1, 2, 0).astype(np.uint8), cv2.COLOR_Lab2RGB) / 255. orig_Lab = np.hstack(orig_Lab[1:]) / 255. orig_ab_BW = np.repeat(orig_Lab[:, :, None], 3, axis=2) # logging images cv2.imwrite( join(f'logs/{self.model_name}/{str_it_count}_colour_est.png'), cv2.cvtColor((est_RGB * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) cv2.imwrite( join(f'logs/{self.model_name}/{str_it_count}_ab_est.png'), (est_ab_BW * 255).astype(np.uint8)) cv2.imwrite( join(f'logs/{self.model_name}/{str_it_count}_colour_orig.png'), cv2.cvtColor((orig_RGB * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) cv2.imwrite( join(f'logs/{self.model_name}/{str_it_count}_ab_orig.png'), (orig_ab_BW * 255).astype(np.uint8)) # assemble compressed visual summary (the plot image) est_RGB = cv2.resize(est_RGB, (128, 128)) orig_RGB = cv2.resize(orig_RGB, (128, 128)) est_ab_BW = cv2.resize(est_ab_BW, (2 * 128, 128)) orig_ab_BW = cv2.resize(orig_ab_BW, (2 * 128, 128)) p_I = np.vstack([ np.hstack([orig_RGB, orig_ab_BW]), np.hstack([est_RGB, est_ab_BW]) ]) return p_I
return ap if __name__ == '__main__': # checking for GPU device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # args parsing args = arg_parser().parse_args() print_info(args) w, h = args.w, args.h s_batch = args.s_batch iterations = args.iterations # loading the dataset if args.dset == 'daisy': dset_kwargs = dict( path=join('/media/ssd1/data_sets/flowers/daisy/train'), width=w, height=h, black_white=False, suffix='.jpg') dset = data.RAMDataSet(**dset_kwargs) dset_kwargs['path'] = join('/media/ssd1/data_sets/flowers/daisy/test') c_out = 2 else: raise NotImplementedError(f'dataset {args.dset} is unknown') # training models with 0-3 pooling operations for n_pooling in [0, 1, 2, 3][::-1]: # keeping track of the data set- and model parameters model_kwargs = dict(c_out=c_out, img_w=w, img_h=h,
# summary: ture image, lightness, estimated image summary = np.hstack([img_BGR, bw, pred_BGR]) # channels: true channels, estimated channels channels = np.hstack([np.vstack(img_ab), np.vstack(pred_ab)]) * 255. channels = channels.astype(np.uint8) # saving summary and channels cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}.png'), summary) cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}_channels.png'), channels) it_count += 1 prog_bar.update(len(y)) model.cpu() batch.cpu() targets.cpu() prog_bar.close() if __name__ == '__main__': device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'\n{Fore.BLUE}generating showcase images{Fore.RESET}') for model_name in sorted(os.listdir('models')): state_dict = torch.load(join(f'models/{model_name}/state_dict.pth'), map_location=torch.device('cpu')) with open(join(f'models/{model_name}/kwargs.json'), 'r') as f_kwargs: kwargs = json.load(f_kwargs) model = nn.ConvNet(**kwargs['model']) model.load_state_dict(state_dict) kwargs['dset']['path'] = join('show_cases/imgs/') kwargs['dset']['suffix'] = '.jpg' dset = data.RAMDataSet(**kwargs['dset']) write_imgs(model, dset, device, model_name=model_name)