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)