def develop_images(camera, pipeline, n_images=0, root_dir='./data/raw/', model_dir='nip_model_snapshots', dev_dir='nip_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, 'nip_training_data', camera) out_directory = os.path.join(root_dir, dev_dir, camera, pipeline) raw_directory = os.path.join(root_dir, '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_api 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 = coreutils.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 'Net' in pipeline: sess = tf.Session() model = getattr(pipelines, pipeline)(sess, tf.get_default_graph(), loss_metric='L1', **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_api.process_auto(raw_file) elif pipeline == 'Python': rgb = 255 * raw_api.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/raw/', raw_dir=''): if target_pipeline not in ['auto', 'manual', 'camera']: raise ValueError('Unsupported target pipeline!') raw_directory = raw_dir out_directory = os.path.join(root_dir, 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 = coreutils.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_api.stacked_bayer(os.path.join( raw_directory, nef_file), camera, use_wb=True) 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_api.process_auto( os.path.join(raw_directory, nef_file)) elif target_pipeline == 'manual': rgb = 255 * raw_api.process( os.path.join(raw_directory, nef_file), **dev_settings) elif target_pipeline == 'camera': out_jpg = os.path.join( out_directory, os.path.splitext(nef_file)[0] + '.JPG') else: raise ValueError('Unsupported develop mode!') if target_pipeline != 'camera': 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)