def main():
parser = argparse.ArgumentParser(
description='Test a neural imaging pipeline')
parser.add_argument('plot',
help='Plot type ({})'.format(
', '.join(supported_plots)))
parser.add_argument(
'--data',
dest='data',
action='store',
default='./data/rgb/clic256/',
help='directory with training & validation images (png)')
parser.add_argument('--images',
dest='images',
action='store',
default=10,
type=int,
help='number of images to test')
parser.add_argument('--image',
dest='image_id',
action='store',
default=1,
type=int,
help='ID of the image to load')
parser.add_argument('--patch',
dest='patch_size',
action='store',
default=128,
type=int,
help='training patch size')
parser.add_argument('--dcn',
dest='dcn',
action='store',
help='directory with a trained DCN model')
args = parser.parse_args()
# Match the current
args.plot = coreutils.match_option(args.plot, supported_plots)
if args.plot == 'batch':
model, stats = codec.restore_model(args.dcn,
args.patch_size,
fetch_stats=True)
print('Training stats:', stats)
data = dataset.IPDataset(args.data,
load='y',
n_images=0,
v_images=args.images,
val_rgb_patch_size=args.patch_size)
batch_x = data.next_validation_batch(0, args.images)
fig = show_example(model, batch_x)
plt.show()
plt.close()
elif args.plot == 'jpeg-match-ssim':
files, _ = loading.discover_files(args.data, n_images=-1, v_images=0)
files = files[args.image_id:args.image_id + 1]
batch_x = loading.load_images(files, args.data, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
model = codec.restore_model(args.dcn, batch_x.shape[1])
fig = match_jpeg(model, batch_x, match='ssim')
plt.show()
plt.close()
elif args.plot == 'jpeg-match-bpp':
files, _ = loading.discover_files(args.data, n_images=-1, v_images=0)
files = files[args.image_id:args.image_id + 1]
batch_x = loading.load_images(files, args.data, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
model = codec.restore_model(args.dcn, batch_x.shape[1])
fig = match_jpeg(model, batch_x, match='bpp')
plt.show()
plt.close()
elif args.plot == 'jpg-trade-off':
df = ratedistortion.get_jpeg_df(args.data, write_files=True)
print(df.to_string())
elif args.plot == 'jp2-trade-off':
df = ratedistortion.get_jpeg2k_df(args.data, write_files=True)
print(df.to_string())
elif args.plot == 'dcn-trade-off':
df = ratedistortion.get_dcn_df(args.data, args.dcn, write_files=False)
print(df.to_string())
elif args.plot == 'bpg-trade-off':
df = ratedistortion.get_bpg_df(args.data, write_files=False)
print(df.to_string())
else:
print('Error: Unknown plot!')
def main():
parser = argparse.ArgumentParser(
description='Train a neural imaging pipeline')
# Parameters related to the training data
parser.add_argument(
'--data',
dest='data',
action='store',
default='./data/rgb/32k',
help='directory with training & validation images (png)')
parser.add_argument(
'--split',
dest='split',
action='store',
default='16000:800:2',
help=
'data split with #training:#validation:#validation_patches - e.g., 16000:800:2'
)
parser.add_argument('--patch',
dest='patch_size',
action='store',
default=128,
type=int,
help='training patch size')
# Parameters of the DCN
parser.add_argument('--dcn',
dest='dcn',
action='store',
default='TwitterDCN',
help='specific DCN class name')
parser.add_argument(
'--params',
dest='dcn_params',
action='append',
help='Extra parameters for DCN constructor (JSON string)')
parser.add_argument('--param_list',
dest='dcn_param_list',
default=None,
help='CSV file with DCN configurations')
# General
parser.add_argument('--out',
dest='out_dir',
action='store',
default='./data/models/dcn/playground',
help='output directory for storing trained models')
parser.add_argument('--epochs',
dest='epochs',
action='store',
default=1500,
type=int,
help='maximum number of training epochs')
parser.add_argument(
'--v_schedule',
dest='validation_schedule',
action='store',
default=100,
type=int,
help='Validation schedule - evaluate the model every v_schedule epochs'
)
parser.add_argument('--lr',
dest='learning_rate',
action='store',
default=1e-4,
type=float,
help='learning rate')
parser.add_argument('--v_train',
dest='validation_is_training',
action='store_true',
default=False,
help='Use the model in training mode while testing')
parser.add_argument(
'--no_aug',
dest='no_aug',
action='store_true',
default=False,
help='disable data augmentation (flipping + gamma correction)')
parser.add_argument(
'--resume',
dest='resume',
action='store_true',
default=False,
help='Resume training from last checkpoint, if possible')
parser.add_argument('--dry',
dest='dry',
action='store_true',
default=False,
help='Dry run (no training - only does model setup)')
parser.add_argument(
'--group',
dest='run_group',
action='store',
type=int,
default=None,
help='Specify run group (sub-selects scenarios for running)')
parser.add_argument(
'--fill',
dest='fill',
action='store',
default=None,
help='Path of the extended scenarios table with appended result columns'
)
args = parser.parse_args()
if not args.dcn:
print('A DCN needs to be specified!')
parser.print_usage()
sys.exit(1)
parameters = pd.DataFrame(
columns=['scenario', 'label', 'active', 'run_group'])
try:
if args.dcn_params is not None:
for p in args.dcn_params:
cli_params = json.loads(p.replace('\'', '"'))
cli_params['scenario'] = np.nan
cli_params['label'] = 'command-line'
cli_params['active'] = True
cli_params['run_group'] = np.nan
parameters = parameters.append(cli_params, ignore_index=True)
if args.dcn_param_list is not None:
parameters = parameters.append(pd.read_csv(args.dcn_param_list),
ignore_index=True,
sort=True)
except json.decoder.JSONDecodeError as e:
print('WARNING', 'JSON parsing error: ', e)
sys.exit(2)
# Round the number of epochs to align with the sampling rate
args.epochs = int(
np.ceil(args.epochs / args.validation_schedule) *
args.validation_schedule) + 1
training_spec = {
'seed': 1234,
'dataset': args.data,
'n_images': int(args.split.split(':')[0]),
'v_images': int(args.split.split(':')[1]),
'valid_patches': int(args.split.split(':')[2]),
'n_epochs': args.epochs,
'batch_size': 50,
'patch_size': args.patch_size,
'sample_dropout': False,
'learning_rate': args.learning_rate,
'learning_rate_reduction_schedule': 1000,
'learning_rate_reduction_factor': 0.5,
'validation_schedule': args.validation_schedule,
'convergence_threshold': 1e-5,
'current_epoch': 0,
'validation_is_training': args.validation_is_training,
'augmentation_probs': {
'resize': 0.0,
'flip_h': 0.0 if args.no_aug else 0.5,
'flip_v': 0.0 if args.no_aug else 0.5,
'gamma': 0.0 if args.no_aug else 0.5,
}
}
if np.sum(parameters['active'] == True) == 0:
parameters.append({
'name': 'default',
'active': True
},
ignore_index=True)
print('DCN model: {}'.format(args.dcn))
if args.run_group is not None:
parameters = parameters[parameters['run_group'] == args.run_group]
parameters = parameters[parameters['active']].drop(
columns=['active', 'run_group'])
print('# DCN parameter list [{} active configs]:\n'.format(
len(parameters)))
print(parameters)
print('\n# Training Spec:')
for key, value in training_spec.items():
print(' {:50s}: {}'.format(key, value))
# Load the dataset
if not args.dry:
print('\n# Dataset:')
np.random.seed(training_spec['seed'])
data = dataset.IPDataset(
args.data,
n_images=training_spec['n_images'],
v_images=training_spec['v_images'],
load='y',
val_rgb_patch_size=training_spec['patch_size'],
val_n_patches=training_spec['valid_patches'])
for key in ['Training', 'Validation']:
print('{:>16s} [{:5.1f} GB] : Y -> {} '.format(
'{} data'.format(key), coreutils.mem(data[key.lower()]['y']),
data[key.lower()]['y'].shape),
flush=True)
model_log = {}
# If requested, add columns to include results
parameters['ssim'] = np.nan
parameters['entropy'] = np.nan
parameters['loss'] = np.nan
print('\n# Training:\n')
for counter, (index, params) in enumerate(
parameters.drop(columns=['scenario', 'label']).iterrows()):
print('## Scenario {} - {} / {}'.format(index, counter + 1,
len(parameters)))
# Create TF session and graph
graph = tf.Graph()
sess = tf.Session(graph=graph)
# Create a DCN according to the spec
dcn_params = {
k: v
for k, v in params.to_dict().items() if not utils.is_nan(v)
}
dcn_params['default_val_is_train'] = training_spec[
'validation_is_training']
dcn = getattr(compression,
args.dcn)(sess,
graph,
None,
patch_size=training_spec['patch_size'],
**dcn_params)
model_code = dcn.model_code
if model_code in model_log:
print(
'WARNING - model {} already registered by scenario {}'.format(
model_code, index))
model_log[model_code].append(index)
else:
model_log[model_code] = [index]
if not args.dry:
train_dcn({'dcn': dcn}, training_spec, data, args.out_dir)
# Fill the table with results, if requested
if args.fill is not None:
results_json = os.path.join(args.out_dir, dcn.model_code,
dcn.scoped_name, 'progress.json')
if os.path.isfile(results_json):
with open(results_json) as f:
results = json.load(f)
parameters.loc[
index,
'ssim'] = results['performance']['ssim']['validation'][-1]
parameters.loc[
index,
'loss'] = results['performance']['loss']['validation'][-1]
parameters.loc[index, 'entropy'] = results['performance'][
'entropy']['training'][-1]
# Cleanup
sess.close()
del graph
if args.fill is not None:
if args.fill == '-':
print(parameters.to_string())
elif args.fill.endswith('.csv'):
print('Saving the results to {}'.format(args.fill))
parameters.to_csv(args.fill, index=False)
else:
raise ValueError(
'Invalid value for the output results file: {}'.format(
args.fill))
if args.dry:
print('List of instantiated models [{}]:'.format(len(model_log)))
for index, key in enumerate(sorted(model_log.keys())):
print('{} {:3d}. {} -> {}'.format(
' ' if len(model_log[key]) == 1 else '!', index, key,
model_log[key]))
def batch_training(nip_model,
camera_names=None,
root_directory=None,
loss_metric='L2',
jpeg_quality=50,
jpeg_mode='sin',
use_pretrained=True,
end_repetition=10,
start_repetition=0,
n_epochs=1001,
nip_directory=None,
split='120:30:4',
regularization_strengths=None):
"""
Repeat training for multiple NIP regularization strengths.
"""
training = {
'use_pretrained_nip': use_pretrained,
'n_epochs': n_epochs,
'patch_size': 128,
'batch_size': 20,
'sampling_rate': 50,
'learning_rate': 1e-4,
'n_images': int(split.split(':')[0]),
'v_images': int(split.split(':')[1]),
'val_n_patches': int(split.split(':')[2])
}
if root_directory is None:
raise FileNotFoundError(
'Invalid root directory: {}'.format(root_directory))
if not os.path.isdir(root_directory):
os.makedirs(root_directory)
if nip_directory is None or not os.path.isdir(nip_directory):
raise FileNotFoundError(
'Invalid NIP snapshots directory: {}'.format(nip_directory))
# Experiment setup
camera_names = camera_names or [
'Nikon D90', 'Nikon D7000', 'Canon EOS 5D', 'Canon EOS 40D'
]
if regularization_strengths is None or len(regularization_strengths) == 0:
regularization_strengths = [
0, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 0.1, 0.25, 0.5, 1
]
else:
regularization_strengths = [float(x) for x in regularization_strengths]
# Construct the TF model
tf_ops, distribution = construct_models(nip_model,
distribution_jpeg=jpeg_quality,
loss_metric=loss_metric,
jpeg_approx=jpeg_mode)
for camera_name in camera_names:
print('\n# Loading data for camera {}'.format(camera_name))
training['camera_name'] = camera_name
# Load the dataset for the given camera
data_directory = os.path.join('./data/raw/nip_training_data/',
camera_name)
# Find available images
data = dataset.IPDataset(data_directory,
n_images=training['n_images'],
v_images=training['v_images'],
load='xy',
val_rgb_patch_size=2 * training['patch_size'],
val_n_patches=training['val_n_patches'])
# Repeat evaluation
for rep in range(start_repetition, end_repetition):
for reg in regularization_strengths:
training['nip_weight'] = reg
training['run_number'] = rep
train_manipulation_nip(tf_ops, training, distribution, data, {
'root': root_directory,
'nip_snapshots': nip_directory
})
def main():
parser = argparse.ArgumentParser(description='Test manipulation detection (FAN) on RGB images')
group = parser.add_argument_group('General settings')
group.add_argument('--patch', dest='patch', action='store', default=64, type=int,
help='patch size')
group.add_argument('--patches', dest='patches', action='store', default=1, type=int,
help='number of validation patches')
group.add_argument('--data', dest='data', action='store', default='./data/rgb/32k',
help='directory with test RGB images')
group = parser.add_argument_group('Training session selection')
group.add_argument('--dir', dest='dir', action='store', default='./data/m/7-raw',
help='directory with training sessions')
group.add_argument('--re', dest='re', action='store', default=None,
help='regular expression to filter training sessions')
group = parser.add_argument_group('Override training settings')
group.add_argument('--jpeg', dest='jpeg', action='store', default=None, type=int,
help='Override JPEG quality level (distribution channel)')
group.add_argument('--dcn', dest='dcn_model', action='store', default=None,
help='DCN compression model path')
group.add_argument('--manip', dest='manipulations', action='store', default=None,
help='Included manipulations, e.g., : {}'.format('sharpen,jpeg,resample,gaussian'))
args = parser.parse_args()
# Split manipulations
if args.manipulations is not None:
args.manipulations = args.manipulations.strip().split(',')
json_files = sorted(str(f) for f in Path(args.dir).glob('**/training.json'))
if len(json_files) == 0:
sys.exit(0)
# Load training / validation data
data = dataset.IPDataset(args.data, n_images=0, v_images=-1, load='y', val_rgb_patch_size=2 * args.patch, val_n_patches=args.patches)
print('Found {} candidate training sessions ({})'.format(len(json_files), args.dir))
print('Data: {}'.format(data.description))
for filename in json_files:
if args.re is None or re.findall(args.re, filename):
with open(filename) as f:
training_log = json.load(f)
# Setup manipulations
if args.manipulations is None:
manipulations = eval(coreutils.getkey(training_log, 'summary/Classes'))
# TODO More elegant solution needed
manipulations.remove('native')
if 'jpg' in manipulations:
manipulations.append('jpeg')
manipulations.remove('jpg')
else:
manipulations = args.manipulations
accuracy = coreutils.getkey(training_log, 'forensics/validation/accuracy')[-1]
compression = coreutils.getkey(training_log, 'summary/Channel Compression')
downsampling = coreutils.getkey(training_log, 'summary/Channel Downsampling')
# Setup compression
if compression.startswith('jpeg'):
# Override from CLI arguments or read from training log
if args.jpeg is not None:
jpeg = args.jpeg
else:
jpeg = int(re.findall('\(([0-9]+),', compression)[0])
dcn_model = None
else:
jpeg = None
if args.dcn_model is not None:
# Override from CLI arguments
dcn_model = args.dcn_model
else:
# Otherwise, read from the training log
if 'dcn' in coreutils.getkey(training_log, 'summary/Joint optimization'):
# If the DCN is trainable, load the fine-tuned model
dcn_model = os.path.join(os.path.split(filename)[0], 'models')
else:
# If not trainable, load a baseline DCN model
compression_params = eval(coreutils.getkey(training_log, 'summary/Channel Compression Parameters'))
dcn_model = compression_params['dirname']
print('\n> {}'.format(filename))
print('Compression: {}'.format(compression))
print('Downsampling: {}'.format(downsampling))
conf_mat, labels = validate_fan(os.path.join(os.path.split(filename)[0], 'models'), manipulations, data, args.patch, dcn_model, downsampling, jpeg)
print(results_data.confusion_to_text((100*conf_mat).round(0), labels, filename, 'txt'))
print('Accuracy: {:.2f} // Expected {:.2f}'.format(np.mean(np.diag(conf_mat)), accuracy))
print(';{};{:.4f};{:.4f}'.format(os.path.split(filename)[0], np.mean(np.diag(conf_mat)), accuracy))
def main():
parser = argparse.ArgumentParser(
description='Train a neural imaging pipeline')
parser.add_argument('--cam', dest='camera', action='store', help='camera')
parser.add_argument('--nip',
dest='nips',
action='append',
choices=["INet", "UNet", "DNet", "OctUNet", "UNet3D"],
help='add NIP for training (repeat if needed)')
parser.add_argument('--out',
dest='out_dir',
action='store',
default='./checkpoint/nip_model_snapshots',
help='output directory for storing trained NIP models')
parser.add_argument(
'--data',
dest='data_dir',
action='store',
default='../../datasets/raw/nip_training_data/',
help='input directory with training data (.npy and .png pairs)')
parser.add_argument('--patch',
dest='patch_size',
action='store',
default=512,
type=int,
help='training patch size (RGB)')
parser.add_argument('--epochs',
dest='epochs',
action='store',
default=25000,
type=int,
help='maximum number of training epochs')
parser.add_argument('--batch',
dest='batch_size',
action='store',
default=20,
type=int,
help='training batch size')
parser.add_argument(
'--params',
dest='nip_params',
default=None,
help='Extra parameters for NIP constructor (JSON string)')
parser.add_argument(
'--resume',
dest='resume',
action='store_true',
default=False,
help='Resume training from last checkpoint, if possible')
parser.add_argument(
'--split',
dest='split',
action='store',
default='270:30:1',
help=
'data split with #training:#validation:#validation_patches - e.g., 120:30:1'
)
parser.add_argument('--ext',
dest='extension',
action='store',
default='png',
help='file extension of rgb images - e.g., png, JPG')
args = parser.parse_args()
if not args.camera:
print('A camera needs to be specified!')
parser.print_usage()
sys.exit(1)
if not args.nips:
print('At least one NIP needs to be specified!')
parser.print_usage()
sys.exit(1)
data_directory = os.path.join(args.data_dir, args.camera)
out_directory_root = args.out_dir
try:
if args.nip_params is not None:
args.nip_params = json.loads(args.nip_params.replace('\'', '"'))
except json.decoder.JSONDecodeError:
print('WARNING', 'JSON parsing error for: ',
args.nip_params.replace('\'', '"'))
sys.exit(2)
print('## Parameters summary')
print('Camera : {}'.format(args.camera))
print('NIPs : {}'.format(args.nips))
print('Params : {}'.format(args.nip_params))
print('Input : {}'.format(data_directory))
print('Output : {}'.format(out_directory_root))
print('Resume : {}'.format(args.resume))
# Load training and validation data
training_spec = {
'seed': 1234,
'n_images': int(args.split.split(':')[0]),
'v_images': int(args.split.split(':')[1]),
'valid_patches': int(args.split.split(':')[2]),
'valid_patch_size': 512,
}
np.random.seed(training_spec['seed'])
# Load and summarize the training data
data = dataset.IPDataset(
data_directory,
n_images=training_spec['n_images'],
v_images=training_spec['v_images'],
load='xy',
val_rgb_patch_size=training_spec['valid_patch_size'],
val_n_patches=training_spec['valid_patches'],
rgb_extension=args.extension)
for key in ['Training', 'Validation']:
print('{:>16s} [{:5.1f} GB] : X -> {}, Y -> {} '.format(
'{} data'.format(key),
coreutils.mem(data[key.lower()]['x']) +
coreutils.mem(data[key.lower()]['y']),
data[key.lower()]['x'].shape, data[key.lower()]['y'].shape),
flush=True)
# Lazy loading to prevent delays in basic CLI interaction
from models import pipelines
import tensorflow as tf
# Train the Desired NIP Models
for pipe in args.nips:
if not issubclass(getattr(pipelines, pipe), pipelines.NIPModel):
supported_nips = [
x for x in dir(pipelines)
if x != 'NIPModel' and type(getattr(pipelines, x)) is type
and issubclass(getattr(pipelines, x), pipelines.NIPModel)
]
raise ValueError(
'Invalid NIP model ({})! Available NIPs: ({})'.format(
pipe, supported_nips))
args.nip_params = args.nip_params or {}
tf.reset_default_graph()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1)
config = tf.ConfigProto(allow_soft_placement=True,
gpu_options=gpu_options)
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
model = getattr(pipelines, pipe)(sess,
tf.get_default_graph(),
loss_metric='L1',
**args.nip_params)
model.sess.run(tf.global_variables_initializer())
train_nip_model(model,
args.camera,
args.epochs,
validation_loss_threshold=1e-5,
patch_size=args.patch_size,
resume=args.resume,
sampling_rate=1000,
batch_size=args.batch_size,
learning_rate=1e-4,
data=data,
out_directory_root=args.out_dir)
sess.close()
return
def batch_training(nip_model,
camera_names=None,
root_directory=None,
loss_metric='L2',
trainables=None,
jpeg_quality=None,
jpeg_mode='soft',
manipulations=None,
dcn_model=None,
downsampling='pool',
end_repetition=10,
start_repetition=0,
n_epochs=1001,
patch=128,
use_pretrained=True,
lambdas_nip=None,
lambdas_dcn=None,
nip_directory=None,
split='120:30:4'):
"""
Repeat training for multiple NIP regularization strengths.
"""
if nip_model is None:
raise FileNotFoundError('NIP model not specified!')
if nip_directory is None or not os.path.isdir(nip_directory):
raise FileNotFoundError(
'Invalid NIP snapshots directory: {}'.format(nip_directory))
if root_directory is None:
raise FileNotFoundError(
'Invalid root directory: {}'.format(root_directory))
if not os.path.isdir(root_directory):
os.makedirs(root_directory)
# Lazy loading to minimize delays when checking cli parameters
from training.manipulation import construct_models, train_manipulation_nip
camera_names = camera_names or [
'Nikon D90', 'Nikon D7000', 'Canon EOS 5D', 'Canon EOS 40D'
]
training = {
'use_pretrained_nip': use_pretrained,
'n_epochs': n_epochs,
'patch_size': patch,
'batch_size': 20,
'sampling_rate': 50,
'learning_rate': 1e-4,
'n_images': int(split.split(':')[0]),
'v_images': int(split.split(':')[1]),
'val_n_patches': int(split.split(':')[2]),
}
# Setup trainable elements and regularization ----------------------------------------------------------------------
trainables = trainables if trainables is not None else set()
for tr in trainables:
if tr not in {'nip', 'dcn'}:
raise ValueError(
'Invalid specifier of trainable elements: only nip, dcn allowed!'
)
training['trainable'] = trainables
if lambdas_nip is None or len(lambdas_nip) == 0:
lambdas_nip = [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 0.1, 0.25, 0.5, 1
] if 'nip' in trainables else [0]
else:
lambdas_nip = [float(x) for x in lambdas_nip]
if lambdas_dcn is None or len(lambdas_dcn) == 0:
lambdas_dcn = [0.1, 0.05, 0.01, 0.005, 0.001
] if 'dcn' in trainables else [0]
else:
lambdas_dcn = [float(x) for x in lambdas_dcn]
if downsampling not in ['pool', 'bilinear', 'none']:
raise ValueError('Unsupported channel down-sampling')
if dcn_model is None and jpeg_quality is None:
jpeg_quality = 50
# Define the distribution channel ----------------------------------------------------------------------------------
compression_params = {}
if jpeg_quality is not None:
compression_params['quality'] = jpeg_quality
compression_params['rounding_approximation'] = jpeg_mode
else:
if dcn_model in codec.dcn_presets:
dcn_model = codec.dcn_presets[dcn_model]
compression_params['dirname'] = dcn_model
if jpeg_quality is not None:
compression = 'jpeg'
elif dcn_model is not None:
compression = 'dcn'
else:
compression = 'none'
# Parse manipulations
manipulations = manipulations or [
'sharpen', 'resample', 'gaussian', 'jpeg'
]
distribution_spec = {
'downsampling': downsampling,
'compression': compression,
'compression_params': compression_params
}
# Construct the TF model
tf_ops, distribution = construct_models(nip_model,
patch_size=training['patch_size'],
trainable=trainables,
distribution=distribution_spec,
manipulations=manipulations,
loss_metric=loss_metric)
for camera_name in camera_names:
print('\n# Loading data for {}'.format(camera_name))
training['camera_name'] = camera_name
# Load the dataset
if nip_model == 'ONet':
# TODO Dirty hack - if the NIP model is the dummy empty model, load RGB images only
data_directory = os.path.join(root_directory, 'rgb', camera_name)
patch_mul = 2
load = 'y'
else:
# Otherwise, load (RAW, RGB) pairs for a specific camera
data_directory = os.path.join(root_directory, 'raw',
'training_data', camera_name)
patch_mul = 2
load = 'xy'
# If the target root directory has no training images, fallback to use the default root
if not os.path.isdir(data_directory):
print(
'WARNING Training images not found in the target root directory - using default root as image source'
)
data_directory = data_directory.replace(root_directory, 'data/')
data_directory = data_directory.replace('//', '/')
# Find available images
data = dataset.IPDataset(data_directory,
n_images=training['n_images'],
v_images=training['v_images'],
load=load,
val_rgb_patch_size=patch_mul *
training['patch_size'],
val_n_patches=training['val_n_patches'])
# data = dataset.IPDataset(data_directory, n_images=training['n_images'], v_images=training['v_images'], load=load, val_rgb_patch_size=training['patch_size'], val_n_patches=training['val_n_patches'])
# Repeat evaluation
for rep in range(start_repetition, end_repetition):
for lr in lambdas_nip:
for lc in lambdas_dcn:
training['lambda_nip'] = lr
training['lambda_dcn'] = lc
training['run_number'] = rep
train_manipulation_nip(tf_ops, training, distribution,
data, {
'root': root_directory,
'nip_snapshots': nip_directory
})