def discover_images(data_directory, n_images=120, v_images=30, extension='png', randomize=0): """ Find available images and split them into training / validation sets. :param data_directory: directory :param n_images: number of training images :param v_images: number of validation images :param extension: file extension :param randomize: whether to shuffle files before the split """ files = fsutil.listdir(data_directory, '.*\\.{}$'.format(extension)) logger.debug(f'{data_directory}: in total {len(files)} files available') if randomize: np.random.seed(randomize) np.random.shuffle(files) if n_images == 0 and v_images == -1: v_images = len(files) if n_images == -1 and v_images == 0: n_images = len(files) if len(files) >= n_images + v_images: val_files = files[n_images:(n_images + v_images)] files = files[0:n_images] else: raise ValueError('Not enough images!') return files, val_files
def manipulation_metrics(nip_models, cameras, root_dir=ROOT_DIRNAME): """ Returns a dataframe with aggregated metrics from manipulation classification (NIP-specific). """ nip_models = [nip_models] if type(nip_models) is str else nip_models cameras = cameras or fsutil.listdir(root_dir, '.', dirs_only=True) if any(cam not in autodetect_cameras(root_dir) for cam in cameras): raise ValueError('The list of cameras does not match the auto-detected list of available models: {}'.format(cameras)) df = pd.DataFrame(columns=['camera', 'nip', 'ln', 'source', 'psnr', 'ssim', 'accuracy']) for camera in cameras: nip_models = nip_models or fsutil.listdir(os.path.join(root_dir, camera), '.', dirs_only=True) for nip in nip_models: find_dir = os.path.join(root_dir, camera, nip) experiment_dirs = fsutil.listdir(os.path.join(find_dir), '.*', dirs_only=True) for ed in experiment_dirs: exp_dir = os.path.join(find_dir, ed) jsons_files = sorted(str(f) for f in Path(exp_dir).glob('**/training.json')) for jf in jsons_files: with open(jf) as f: data = json.load(f) df = df.append({'camera': camera, 'nip': nip, 'ln': ed, 'source': jf.replace(find_dir, '').replace('training.json', ''), 'psnr': utils.get(data, 'nip.performance.psnr.validation')[-1], 'ssim': utils.get(data, 'nip.performance.ssim.validation')[-1], 'accuracy': utils.get(data, 'forensics.performance.accuracy.validation')[-1] }, ignore_index=True) return df
def autodetect_cameras(dirname): """ Returns a list of known cameras (based on available NIP). """ counter = 5 while counter > 0 and not os.path.exists(os.path.join(dirname, 'models', 'nip')): dirname = os.path.split(dirname)[0] counter -= 1 if counter == 0: raise ValueError('The {} directory does not seem to be a valid results directory'.format(dirname)) return fsutil.listdir(os.path.join(dirname, 'models', 'nip'), '.*', dirs_only=True)
def develop_images(camera, pipeline, n_images=0, root_dir='./data', model_dir='nip', dev_dir='developed', nip_params=None): if pipeline not in supported_pipelines: raise ValueError( 'Unsupported pipeline model ({})! Available models: {}'.format( pipeline, ', '.join(supported_pipelines))) dir_models = os.path.join(root_dir, model_dir) nip_directory = os.path.join(root_dir, 'raw', 'training_data', camera) out_directory = os.path.join(root_dir, 'raw', dev_dir, camera, pipeline) raw_directory = os.path.join(root_dir, 'raw', 'images', camera) if not os.path.exists(nip_directory): raise IOError('Directory not found! {}'.format(nip_directory)) if not os.path.exists(out_directory): os.makedirs(out_directory) # Lazy loading of remaining dependencies to ensure responsiveness of the CLI import numpy as np import imageio import tqdm from helpers import raw from models import pipelines print('Camera: {}'.format(camera)) print('Pipeline: {}'.format(pipeline)) print('NIP Models: {}'.format(dir_models)) print('NIP Training Directory: {}'.format(nip_directory)) print('Out Directory: {}'.format(out_directory)) # %% Process Bayer stacks with the given pipeline npy_filenames = fsutil.listdir(nip_directory, '.*\.{}$'.format(extensions)) log.info('Camera {} matched {:,} Bayer stacks'.format( camera, len(npy_filenames))) manual_dev_settings = { 'use_srgb': True, 'use_gamma': True, 'brightness': None } # Setup the NIP model if pipeline.endswith('Net'): sess = tf.Session() model = getattr(pipelines, pipeline)(sess, tf.get_default_graph(), loss_metric='L2', **nip_params) model.load_model(camera, out_directory_root=dir_models) # Limit the number of images if n_images > 0: npy_filenames = npy_filenames[:n_images] for npy_file in tqdm.tqdm(npy_filenames, ncols=120, desc='Developing ({}/{})'.format( camera, pipeline)): # Find the original RAW file (for standard pipelines.py) raw_file = os.path.join(raw_directory, os.path.splitext(npy_file)[0]) raw_found = False for extension in raw_extensions: if os.path.exists(raw_file + extension): raw_file = raw_file + extension raw_found = True break if not raw_found: raise RuntimeError( 'RAW file not found for Bayer stack: {}'.format(npy_file)) out_png = os.path.join(out_directory, os.path.splitext(npy_file)[0] + '.png') if not os.path.exists(out_png): # Process with the desired pipeline if pipeline == 'libRAW': rgb = raw.process_auto(raw_file) elif pipeline == 'Python': rgb = 255 * raw.process(raw_file, **manual_dev_settings) rgb = rgb.astype(np.uint8) else: # Find the cached Bayer stack bayer_file = os.path.join(nip_directory, npy_file) bayer_stack = np.load(bayer_file).astype( np.float32) / (2**16 - 1) rgb = 255 * model.process(bayer_stack).squeeze() rgb = rgb.astype(np.uint8) imageio.imwrite(out_png, rgb.astype(np.uint8))
def prepare_training_set(camera, target_pipeline, dev_settings, n_images=150, root_dir='./data/'): if target_pipeline not in ['auto', 'manual']: raise ValueError('Unsupported target pipeline!') raw_directory = os.path.join(root_dir, 'raw', 'images', camera) out_directory = os.path.join(root_dir, 'raw', 'training_data', camera) if not os.path.exists(raw_directory): log.error('Directory not found! {}'.format(raw_directory)) sys.exit(2) if not os.path.exists(out_directory): os.makedirs(out_directory) print('RAW Directory: {}'.format(raw_directory)) print('Out Directory: {}'.format(out_directory)) # List RAW files and find the ones with horizontal orientation raw_filenames = fsutil.listdir(raw_directory, '.*\.{}$'.format(EXTENSIONS)) log.info('Camera {} matched {:,} RAW images'.format(camera, len(raw_filenames))) raw_filenames_selected = [] for nef_file in raw_filenames: with open(os.path.join(raw_directory, nef_file), 'rb') as f: tags = exifread.process_file(f, details=False, stop_tag='Image Orientation') orientation = tags['Image Orientation'].printable log.info('{} -> {}'.format(nef_file, orientation)) if orientation.startswith('Horizontal'): raw_filenames_selected.append(nef_file) if len(raw_filenames_selected) >= n_images: break log.info('Collected {} landscape-oriented photos for training'.format(len(raw_filenames_selected))) if len(raw_filenames_selected) < n_images: log.error('Not enough horizontal images! Found {} but expected {}.'.format(len(raw_filenames_selected), n_images)) dev_settings = dev_settings or {'use_srgb': True, 'use_gamma': True, 'brightness': None} # Iterate over RAW files and produce: # 1. RGGB Bayer stacks (H/2, W/2, 4) # 2. RGB Optimization target (H, W, 3) for nef_file in tqdm.tqdm(raw_filenames_selected, ncols=120, desc='Preparing train. data ({})'.format(camera)): out_npy = os.path.join(out_directory, os.path.splitext(nef_file)[0] + '.npy') out_png = os.path.join(out_directory, os.path.splitext(nef_file)[0] + '.png') try: if not os.path.exists(out_npy): image_bayer = raw.unpack(os.path.join(raw_directory, nef_file))[0] image_bayer = ((2**16 - 1) * image_bayer).astype(np.uint16) np.save(out_npy, image_bayer) if not os.path.exists(out_png): if target_pipeline == 'auto': rgb = raw.process_auto(os.path.join(raw_directory, nef_file)) elif target_pipeline == 'manual': rgb = 255 * raw.process(os.path.join(raw_directory, nef_file), **dev_settings) else: raise ValueError('Unsupported develop mode!') imageio.imwrite(out_png, rgb.astype(np.uint8)) except Exception as error: log.error('RAW Processing failed for file: {}'.format(nef_file)) log.error(error) sys.exit(2) sys.exit(0)
def develop_image(pipeline, camera=None, batch=None, image=None, patch_size=0, patches=2, root_dir='./data', pipeline_args=None): """ Display a patch developed by a neural imaging pipeline. """ if camera is not None: supported_cameras = fsutil.listdir( os.path.join(root_dir, 'models', 'nip'), '.*') if camera not in supported_cameras: raise ValueError( 'Camera data not found ({})! Available cameras: {}'.format( camera, ', '.join(supported_cameras))) root_dirname = os.path.join(root_dir, 'models', 'nip', camera) data_dirname = os.path.join(root_dir, 'raw', 'training_data', camera) if patch_size != 0 and (patch_size < 4 or patch_size > 2048): raise ValueError('Patch size seems to be invalid!') # Lazy imports to minimize delay for invalid command line parameters import numpy as np import imageio as io import matplotlib.pyplot as plt import tensorflow as tf from models import pipelines # Construct the NIP model --------------------------------------------------------------------- if os.path.isdir(pipeline): # Restore a NIP model from a training log model = tfmodel.restore(pipeline, pipelines) else: # Construct the NIP model from class name (and optional arguments) if pipeline_args is None: model = getattr(pipelines, pipeline)() else: model = getattr(pipelines, pipeline)(**pipeline_args) loaded_model = False candidate_dirs = [ os.path.join(root_dirname, model.model_code), os.path.join(root_dirname) ] for candidate in candidate_dirs: if os.path.isdir(candidate): model.load_model(candidate) loaded_model = True break if not loaded_model: raise FileNotFoundError( f'Could not find the corresponding model: {candidate_dirs}') # Load image(s) ------------------------------------------------------------------------------- if image is None and batch is not None: print('Loading a batch of {} images'.format(batch)) data = dataset.Dataset(data_dirname, n_images=0, v_images=batch, val_rgb_patch_size=patch_size or 256, val_n_patches=patches) sample_x, sample_y = data.next_validation_batch( 0, data.count_validation) with open('config/cameras.json') as f: cameras = json.load(f) cfa, srgb = cameras[camera]['cfa'], np.array( cameras[camera]['srgb']) elif image is not None: print('Loading a RAW image {}'.format(image)) sample_x, cfa, srgb, _ = raw.unpack(image, expand=True) sample_y = raw.process(image, brightness=None, expand=True) if isinstance(model, pipelines.ClassicISP): print('Configuring ISP to CFA: {} & sRGB {}'.format( cfa, srgb.round(2).tolist())) model.set_cfa_pattern(cfa) model.set_srgb_conversion(srgb) sample_Y = model.process(sample_x).numpy() if patch_size > 0: xx = (sample_y.shape[2] - patch_size) // 2 yy = (sample_y.shape[1] - patch_size) // 2 sample_y = sample_y[:, yy:yy + patch_size, xx:xx + patch_size, :] sample_Y = sample_Y[:, yy:yy + patch_size, xx:xx + patch_size, :] psnrs = metrics.psnr(sample_y, sample_Y) ssims = metrics.ssim(sample_y, sample_Y) print('sample x: {}'.format(sample_x.shape)) print('sample y: {}'.format(sample_y.shape)) print('sample Y: {}'.format(sample_Y.shape)) # Plot images --------------------------------------------------------------------------------- if len(sample_y) > 1: sample_y = plots.thumbnails(sample_y, batch, True) sample_Y = plots.thumbnails(sample_Y, batch, True) else: sample_y = sample_y.squeeze() sample_Y = sample_Y.squeeze() print('thumbnails: {}'.format(sample_y.shape)) ncols = 1 if sample_y.shape[1] > sample_y.shape[0] else 2 nrows = 2 if ncols == 1 else 1 fig, axes = plt.subplots(nrows, ncols) plots.image(sample_Y, '{}, PSNR={:.1f} dB, SSIM={:.2f} : {{}}'.format( model.model_code, float(psnrs.mean()), float(ssims.mean())), axes=axes[0]) plots.image(sample_y, 'Target RGB images () : {}', axes=axes[1]) plt.show() plt.close()
def display_results(args): sns.set('paper', font_scale=1, style="ticks") plot = helpers.utils.match_option(args.plot, supported_plots) if not os.path.isdir(args.dir): raise FileNotFoundError('Directory {} not found!'.format(args.dir)) print('Results from: {}'.format(args.dir)) print('Matched plotting command: {}'.format(plot)) postfix = [ fsutil.split(args.dir)[-1], ','.join(args.nips) if args.nips is not None else None, ','.join(args.cameras) if args.cameras is not None else None, ] postfix = '-'.join(x for x in postfix if x is not None) if plot in ['ssim', 'psnr', 'accuracy']: df = results_data.manipulation_metrics(args.nips, args.cameras, root_dir=args.dir) sns.catplot(x='ln', y=plot, col='camera', row='nip', data=df, kind='box') save_df(df, args.df, 'manipulation_metrics-{}.csv'.format(postfix)) plt.show() return if plot == 'scatter-psnr' or plot == 'scatter-ssim': df = results_data.manipulation_metrics(args.nips, args.cameras, root_dir=args.dir) if len(df) == 0: print('ERROR No results found!') sys.exit(2) print(df) g = sns.relplot(x=plot.split('-')[-1], y='accuracy', hue='ln', col='camera', row='nip', data=df, palette=sns.color_palette("Set2", len(df['ln'].unique()))) save_df(df, args.df, 'manipulation_metrics-{}.csv'.format(postfix)) plt.show() return if plot == 'progress': cases = [] if args.cameras is None: args.cameras = fsutil.listdir(args.dir, '.', dirs_only=True) for cam in args.cameras: nip_models = args.nips or fsutil.listdir(os.path.join(args.dir, cam), '.', dirs_only=True) for nip in nip_models: reg_path = os.path.join(args.dir, cam, nip) if args.regularization: # If given, use specified regularization strengths reg_list = args.regularization else: # Otherwise, auto-detect available scenarios reg_list = fsutil.listdir(reg_path, '.*', dirs_only=True) if len(reg_list) > 4: indices = np.linspace(0, len(reg_list)-1, 4).astype(np.int32) reg_list = [reg_list[i] for i in indices] print('! warning - too many experiments to show - sampling: {}'.format(reg_list)) for reg in reg_list: for r in fsutil.listdir(os.path.join(reg_path, reg), '[0-9]+', dirs_only=True): print('* found scenario {}'.format((cam, nip, reg, int(r)))) cases.append((cam, nip, reg, int(r))) df, labels = results_data.manipulation_progress(cases, root_dir=args.dir) save_df(df, args.df, 'progress-{}.csv'.format(postfix)) for col in ['psnr', 'accuracy']: if len(df[col].dropna()) > 0: sns.relplot(x="step", y=col, hue='exp', row='nip', col='camera', style='exp', kind="line", legend="full", aspect=2, height=3, data=df) plt.show() return if plot == 'conf' or plot == 'conf-tex': if isinstance(args.nips, list): if len(args.nips) > 1: print('WARNING Only one NIP will be used for this plot!') args.nips = args.nips[0] conf = results_data.confusion_data(args.run, root_dir=args.dir) if len(conf) == 0: print('ERROR No results found!') return tex_output = plot == 'conf-tex' plot_data = not tex_output if len(conf.keys()) < 20 else False if plot_data: images_x = np.ceil(np.sqrt(len(conf))) images_y = np.ceil(len(conf) / images_x) f_size = 3 fig = plt.figure(figsize=(images_x*f_size, images_y*f_size)) for i, (k, c) in enumerate(conf.items()): data = (100*c['data']).round(0) labels = c['labels'] if tex_output: print(results_data.confusion_to_text(data, labels, k, 'tex')) else: print(results_data.confusion_to_text(data, labels, k, 'txt')) if plot_data: acc = np.mean(np.diag(data)) ax = fig.add_subplot(images_y, images_x, i+1) sns.heatmap(data, annot=True, fmt=".0f", linewidths=.5, xticklabels=[x[0] for x in labels], yticklabels=labels) ax.set_title('{} : acc={:.1f}'.format(k, acc)) if plot_data: plt.tight_layout() plt.show() return if plot == 'df': print('Searching for "training.json" in', args.dir) df = results_data.manipulation_summary(args.dir) if len(df) > 0: if False: print(df.groupby('scenario').mean().to_string()) else: gb = df.groupby('scenario') counts = gb.size().to_frame(name='reps') print(counts.join(gb.agg('mean')).reset_index().to_string()) save_df(df, args.df, 'summary-{}.csv'.format(postfix)) return if plot == 'auto': print('Searching for "training.json" in', args.dir) df = results_data.manipulation_summary(args.dir) df = df.sort_values('scenario') guessed_names = {} # Guess scenario components = df['scenario'].str.split("/", expand=True) for i in components: # Try to guess the column name based on content template = 'scenario:{}'.format(i) if components.iloc[0, i].endswith('Net'): guessed_names[template] = 'nip' elif components.iloc[0, i].startswith('ln-'): guessed_names[template] = 'nip reg.' elif components.iloc[0, i].startswith('lc-'): guessed_names[template] = 'dcn reg.' elif set(components.iloc[:, i].unique()) == {'4k', '8k', '16k'}: guessed_names[template] = 'dcn' elif all([re.match('^[0-9]{2,3}$', x) for x in components.iloc[:, i].unique()]): guessed_names[template] = 'jpeg' else: guessed_names[template] = template df[guessed_names[template]] = components[i] df['scenario'] = fsutil.strip_prefix(df['scenario']) mapping = {} mapping_targets = ['col', 'col', 'hue', 'style', 'size'] mapping_id = 0 # Choose the feature with most unique values as x axis uniques = [len(df[guessed_names['scenario:{}'.format(i)]].unique()) for i in components] x_feature = np.argmax(uniques) for i in components: if i == x_feature: continue if len(df[guessed_names['scenario:{}'.format(i)]].unique()) > 1: mapping[mapping_targets[mapping_id]] = guessed_names['scenario:{}'.format(i)] mapping_id += 1 sns.catplot(x=guessed_names['scenario:{}'.format(x_feature)], y='accuracy', data=df, kind='box', **mapping) # sns.catplot(x='scenario:0', y='dcn_ssim', data=df, kind='box', **mapping) # sns.scatterplot(x='dcn_ssim', y='accuracy', data=df) plt.show() if len(df) > 0: gb = df.groupby('scenario') counts = gb.size().to_frame(name='reps') print(counts.join(gb.agg('mean')).reset_index().to_string()) return raise RuntimeError('No plot matched! Available plots {}'.format(', '.join(supported_plots)))
def compare_nips(model_a_dirname, model_b_dirname, camera=None, image=None, patch_size=128, root_dirname='./data', output_dir=None, model_a_args=None, model_b_args=None, extras=False): """ Display a comparison of two variants of a neural imaging pipeline. :param camera: camera name (e.g., 'Nikon D90') :param model_a_dirname: directory with the first variant of the model :param model_b_dirname: directory with the second variant of the model :param ps: patch size (patch will be taken from the middle) :param image_id: index of the test image :param root_dir: root data directory :param output_dir: set an output directory if the figure should be saved (matplotlib2tikz will be used) """ # Lazy imports to minimize delay for invalid command line parameters import re import inspect import imageio as io import matplotlib.pyplot as plt import tensorflow as tf from models import pipelines, tfmodel from helpers import raw, loading supported_cameras = fsutil.listdir( os.path.join(root_dirname, 'models', 'nip'), '.*') supported_pipelines = pipelines.supported_models if patch_size > 0 and (patch_size < 8 or patch_size > 2048): raise ValueError('Patch size seems to be invalid!') if camera is not None and camera not in supported_cameras: raise ValueError( 'Camera data not found ({})! Available cameras: {}'.format( camera, ', '.join(supported_cameras))) # Check if the image is an integer try: image = int(image) except: pass # Construct the NIP models if os.path.isdir(model_a_dirname): # Restore a NIP model from a training log model_a = tfmodel.restore(model_a_dirname, pipelines) else: # Construct the NIP model from class name (and optional arguments) if model_a_args is None: model_a = getattr(pipelines, model_a_dirname)() else: model_a = getattr(pipelines, model_a_dirname)(**model_a_args) model_a.load_model(os.path.join(root_dirname, model_a.model_code)) if os.path.isdir(model_b_dirname): # Restore a NIP model from a training log model_b = tfmodel.restore(model_b_dirname, pipelines) else: # Construct the NIP model from class name (and optional arguments) if model_b_args is None: model_b = getattr(pipelines, model_b_dirname)() else: model_b = getattr(pipelines, model_b_dirname)(**model_b_args) model_b.load_model(os.path.join(root_dirname, model_b.model_code)) print('ISP-A: {}'.format(model_a.summary())) print('ISP-B: {}'.format(model_b.summary())) # Load sample data if isinstance(image, int) and camera is not None: data_dirname = os.path.join(root_dirname, 'raw', 'training_data', camera) files = fsutil.listdir(data_dirname, '.*\.png') files = files[image:image + 1] print('Loading image {} from the training set: {}'.format( image, files)) data = loading.load_images(files, data_dirname) sample_x, sample_y = data['x'].astype(np.float32) / ( 2**16 - 1), data['y'].astype(np.float32) / (2**8 - 1) with open('config/cameras.json') as f: cameras = json.load(f) cfa, srgb = cameras[camera]['cfa'], np.array( cameras[camera]['srgb']) image = files[0] elif image is not None: print('Loading a RAW image {}'.format(image)) sample_x, cfa, srgb, _ = raw.unpack(image, expand=True) sample_y = raw.process(image, brightness=None, expand=True) image = os.path.split(image)[-1] if isinstance(model_a, pipelines.ClassicISP): print('Configuring ISP-A to CFA: {} & sRGB {}'.format( cfa, srgb.round(2).tolist())) model_a.set_cfa_pattern(cfa) model_a.set_srgb_conversion(srgb) if isinstance(model_b, pipelines.ClassicISP): print('Configuring ISP-B to CFA: {} & sRGB {}'.format( cfa, srgb.round(2).tolist())) model_b.set_cfa_pattern(cfa) model_b.set_srgb_conversion(srgb) # Develop images sample_ya = model_a.process(sample_x).numpy() sample_yb = model_b.process(sample_x).numpy() if patch_size > 0: print('Cropping a {p}x{p} patch from the middle'.format(p=patch_size)) xx = (sample_x.shape[2] - patch_size // 2) // 2 yy = (sample_x.shape[1] - patch_size // 2) // 2 sample_x = sample_x[:, yy:yy + patch_size, xx:xx + patch_size, :] sample_y = sample_y[:, 2 * yy:2 * (yy + patch_size), 2 * xx:2 * (xx + patch_size), :] sample_ya = sample_ya[:, 2 * yy:2 * (yy + patch_size), 2 * xx:2 * (xx + patch_size), :] sample_yb = sample_yb[:, 2 * yy:2 * (yy + patch_size), 2 * xx:2 * (xx + patch_size), :] # Plot images fig = imdiff.compare_ab_ref(sample_y, sample_ya, sample_yb, fig=plt.figure(), extras=extras) if output_dir is not None: from tikzplotlib import save as tikz_save dcomp = [ x for x in fsutil.split(model_b_dirname) if re.match('(ln-.*|[0-9]{3})', x) ] tikz_save('{}/examples_{}_{}_{}_{}.tex'.format(output_dir, camera, image, model_a_dirname, model_b_dirname), figureheight='8cm', figurewidth='8cm', strict=False) else: fig.tight_layout() fig.show(fig) fig.suptitle('{}, A={}, B={}'.format(image, model_a.model_code, model_b.model_code)) plt.show() plt.close(fig)