def __init__(self,
data_directory,
*,
randomize=2468,
load='xy',
n_images=120,
v_images=30,
val_rgb_patch_size=128,
val_n_patches=1,
val_discard='flat-aggressive'):
"""
Represents a [RAW-]RGB dataset for training imaging pipelines. The class preloads full resolution images and
samples from them when requesting training batches. (Validation images are sampled upon creation.) Patch
selection takes care of proper alignment between RAW images (represented as half-size 4-channel RGGB stacks) and
their corresponding rendered RGB versions. Selection can be controlled to prefer certain types of patches (not
strictly enforced). The following DISCARD modes are available:
- flat - attempts to discard flat patches based on patch variance (not strict)
- flat-aggressive - a more aggressive version that avoids patches with variance < 0.01
- dark-n-textured - avoid dark (mean < 0.35) and textured patches (variance > 0.005)
Usage examples:
---------------
# Load a RAW -> RGB dataset
data = Dataset('data/raw/training_data/D90')
batch_raw, batch_rgb = data.next_training_batch(0, 10, 128, 'flat-aggressive')
# Load RGB only dataset
data = Dataset('data/rgb/native12k/', load='y')
batch_rgb = data.next_training_batch(0, 10, 128, 'flat-aggressive')
:param data_directory: directory path with RAW-RGB pairs (*.npy & *.png) or only RGB images (*.png)
:param randomize: randomization seed
:param load: what data to load: 'xy' load RAW+RGB, 'x' load RAW only, 'y' load RGB only
:param n_images: number of training images (full resolution)
:param v_images: number of validation images (patches sampled upon creation)
:param val_rgb_patch_size: validation patch size
:param val_n_patches: number of validation patches to load per full-resolution image
:param val_discard: patch discard mode (for validation data)
"""
if not any(load == allowed for allowed in ['xy', 'x', 'y']):
raise ValueError('Invalid X/Y data requested!')
if not os.path.isdir(data_directory):
if '/' in data_directory or '\\' in data_directory:
raise ValueError(
f'Cannot find the data directory: {data_directory}')
if os.path.isdir(
os.path.join('data/raw/training_data/', data_directory)):
data_directory = os.path.join('data/raw/training_data/',
data_directory)
elif os.path.isdir(os.path.join('data/rgb/', data_directory)):
data_directory = os.path.join('data/rgb/', data_directory)
else:
raise ValueError(
f'Cannot find the data directory: {data_directory}')
self.files = {}
self._loaded_data = load
self._data_directory = data_directory
self._counts = (n_images, v_images, val_n_patches)
self._val_discard = 'flat-aggressive'
self.files['training'], self.files[
'validation'] = loading.discover_images(data_directory,
randomize=randomize,
n_images=n_images,
v_images=v_images)
self.data = {
'training':
loading.load_images(self.files['training'],
data_directory,
load=load),
'validation':
loading.load_patches(self.files['validation'],
data_directory,
patch_size=val_rgb_patch_size // 2,
n_patches=val_n_patches,
load=load,
discard=val_discard)
}
if 'y' in self.data['training']:
self.H, self.W = self.data['training']['y'].shape[1:3]
else:
self.H, self.W = (2 * dim
for dim in self.data['training']['x'].shape[1:3])
def get_bpg_df(directory,
write_files=False,
effective_bytes=True,
force_calc=False):
"""
Compute and return (as Pandas DF) the rate distortion curve for BPG. The result is saved
as a CSV file in the source directory. If the file exists, the DF is loaded and returned.
The files are saved using the reference codec: https://bellard.org/bpg/
"""
files, _ = loading.discover_images(directory, n_images=-1, v_images=0)
batch_x = loading.load_images(files, directory, load='y')
batch_x = batch_x['y'].astypre(np.float32) / (2**8 - 1)
quality_levels = np.arange(10, 40, 1)
df_jpeg_path = os.path.join(directory, 'bpg.csv')
if os.path.isfile(df_jpeg_path) and not force_calc:
logger.info('Restoring BPG stats from {}'.format(df_jpeg_path))
df = pd.read_csv(df_jpeg_path, index_col=False)
else:
df = pd.DataFrame(columns=[
'image_id', 'filename', 'codec', 'quality', 'ssim', 'psnr',
'msssim', 'msssim_db', 'bytes', 'bpp'
])
with tqdm.tqdm(total=len(files) * len(quality_levels),
ncols=120,
desc='BPG') as pbar:
for image_id, filename in enumerate(files):
# Read the original image
image = batch_x[image_id]
for qi, q in enumerate(quality_levels):
# Compress to BPG
# Save as temporary file
imageio.imwrite('/tmp/image.png',
(255 * image).astype(np.uint8))
bpp_path = bpg_helpers.bpg_compress(
'/tmp/image.png', q, '/tmp')
image_compressed = imageio.imread(
bpg_helpers.decode_bpg_to_png(bpp_path)).astype(
np.float) / (2**8 - 1)
if effective_bytes:
bpp = bpg_helpers.bpp_of_bpg_image(bpp_path)
image_bytes = round(bpp * image.shape[0] *
image.shape[1] / 8)
else:
image_bytes = os.stat(bpp_path).st_size
bpp = 8 * image_bytes / image.shape[0] / image.shape[1]
if write_files:
image_dir = os.path.join(directory,
os.path.splitext(filename)[0])
if not os.path.isdir(image_dir):
os.makedirs(image_dir)
image_path = os.path.join(image_dir,
'bpg_q{:03d}.png'.format(q))
imageio.imwrite(image_path,
(255 * image_compressed).astype(
np.uint8))
msssim_value = msssim(image, image_compressed, MAX=1).real
df = df.append(
{
'image_id':
image_id,
'filename':
filename,
'codec':
'bpg',
'quality':
q,
'ssim':
compare_ssim(image,
image_compressed,
multichannel=True,
data_range=1),
'psnr':
compare_psnr(image, image_compressed,
data_range=1),
'msssim':
msssim_value,
'msssim_db':
-10 * np.log10(1 - msssim_value),
'bytes':
image_bytes,
'bpp':
bpp
},
ignore_index=True)
pbar.set_postfix(image_id=image_id, quality=q)
pbar.update(1)
df.to_csv(df_jpeg_path, index=False)
return df
def get_dcn_df(directory,
model_directory,
write_files=False,
force_calc=False):
"""
Compute and return (as Pandas DF) the rate distortion curve for the learned DCN codec.
The result is saved as a CSV file in the source directory. If the file exists, the DF
is loaded and returned.
"""
# Discover test files
files, _ = loading.discover_images(directory, n_images=-1, v_images=0)
batch_x = loading.load_images(files, directory, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
# Create a new table for the DCN
df = pd.DataFrame(columns=[
'image_id', 'filename', 'model_dir', 'codec', 'ssim', 'psnr', 'msssim',
'msssim_db', 'entropy', 'bytes', 'bpp', 'layers', 'quantization',
'entropy_reg', 'codebook', 'latent', 'latent_shape', 'n_features'
])
# Discover available models
model_dirs = list(Path(model_directory).glob('**/progress.json'))
logger.info('Found {} models'.format(len(model_dirs)))
df_path = os.path.join(
directory, 'dcn-{}.csv'.format(
[x for x in fsutil.split(model_directory) if len(x) > 0][-1]))
if os.path.isfile(df_path) and not force_calc:
logger.info('Restoring DCN stats from {}'.format(df_path))
df = pd.read_csv(df_path, index_col=False)
else:
for model_dir in model_dirs:
logger.info('Processing model dir: {}'.format(model_dir))
dcn = codec.restore(
os.path.split(str(model_dir))[0], batch_x.shape[1])
# Dump compressed images
for image_id, filename in enumerate(files):
try:
batch_y, image_bytes = codec.simulate_compression(
batch_x[image_id:image_id + 1], dcn)
batch_z = dcn.compress(batch_x[image_id:image_id + 1])
entropy = helpers.stats.entropy(batch_z,
dcn.get_codebook())
except Exception as e:
logger.error(
'Error while processing {} with {} : {}'.format(
filename, dcn.model_code, e))
raise e
if write_files:
image_dir = os.path.join(directory,
os.path.splitext(filename)[0])
if not os.path.isdir(image_dir):
os.makedirs(image_dir)
image_path = os.path.join(
image_dir,
dcn.model_code.replace('/', '-') + '.png')
imageio.imwrite(image_path,
(255 * batch_y[0]).astype(np.uint8))
msssim_value = msssim(batch_x[image_id], batch_y[0],
MAX=1).real
df = df.append(
{
'image_id':
image_id,
'filename':
filename,
'model_dir':
os.path.relpath(
os.path.split(str(model_dir))[0],
model_directory).replace(dcn.scoped_name, ''),
'codec':
dcn.model_code,
'ssim':
compare_ssim(batch_x[image_id],
batch_y[0],
multichannel=True,
data_range=1),
'psnr':
compare_psnr(
batch_x[image_id], batch_y[0], data_range=1),
'msssim':
msssim_value,
'msssim_db':
-10 * np.log10(1 - msssim_value),
'entropy':
entropy,
'bytes':
image_bytes,
'bpp':
8 * image_bytes / batch_x[image_id].shape[0] /
batch_x[image_id].shape[1],
'layers':
dcn.n_layers if 'n_layers' in dcn._h else None,
'quantization':
'{}-{:.0f}bpf'.format(dcn._h.rounding, dcn.latent_bpf),
'entropy_reg':
dcn.entropy_weight,
'codebook':
dcn._h.rounding,
'latent':
dcn.n_latent,
'latent_shape':
'{}x{}x{}'.format(*dcn.latent_shape[1:]),
'n_features':
dcn.latent_shape[-1]
},
ignore_index=True)
df.to_csv(df_path, index=False)
return df
def get_jpeg2k_df(directory,
write_files=False,
effective_bytes=True,
force_calc=False):
"""
Compute and return (as Pandas DF) the rate distortion curve for JPEG 2000. The result is saved
as a CSV file in the source directory. If the file exists, the DF is loaded and returned.
Files are saved as JPEG using glymur.
"""
files, _ = loading.discover_images(directory, n_images=-1, v_images=0)
batch_x = loading.load_images(files, directory, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
# Get trade-off for JPEG
quality_levels = np.arange(25, 45, 1)
df_jpeg_path = os.path.join(directory, 'jpeg2000.csv')
if os.path.isfile(df_jpeg_path) and not force_calc:
logger.info('Restoring JPEG 2000 stats from {}'.format(df_jpeg_path))
df = pd.read_csv(df_jpeg_path, index_col=False)
else:
df = pd.DataFrame(columns=[
'image_id', 'filename', 'codec', 'quality', 'ssim', 'psnr',
'msssim', 'msssim_db', 'bytes', 'bpp'
])
with tqdm.tqdm(total=len(files) * len(quality_levels),
ncols=120,
desc='JP2k') as pbar:
for image_id, filename in enumerate(files):
# Read the original image
image = batch_x[image_id]
for qi, q in enumerate(quality_levels):
# TODO Use Glymur to save JPEG 2000 images to a temp file
image_np = (255 * image.clip(0, 1)).astype(np.uint8)
glymur.Jp2k('/tmp/image.jp2', data=image_np, psnr=[q])
if effective_bytes:
image_bytes = jpeg_helpers.jp2bytes('/tmp/image.jp2')
else:
image_bytes = os.path.getsize('/tmp/image.jp2')
image_compressed = imageio.imread('/tmp/image.jp2').astype(
np.float) / (2**8 - 1)
# TODO Use Pillow to save JPEG 2000 images to a memory buffer
# TODO This has been disabled = their implementation seems to be invalid
# with io.BytesIO() as output:
# image_pillow = PIL.Image.fromarray((255*image.clip(0, 1)).astype(np.uint8))
# image_pillow.save(output, format='jpeg2000', quality_layers=[q])
# image_compressed = imageio.imread(output.getvalue()).astype(np.float) / (2**8 - 1)
# image_bytes = len(output.getvalue())
if write_files:
image_dir = os.path.join(directory,
os.path.splitext(filename)[0])
if not os.path.isdir(image_dir):
os.makedirs(image_dir)
image_path = os.path.join(
image_dir, 'jp2_q{:.1f}dB.png'.format(q))
imageio.imwrite(image_path,
(255 * image_compressed).astype(
np.uint8))
msssim_value = msssim(image, image_compressed, MAX=1).real
df = df.append(
{
'image_id':
image_id,
'filename':
filename,
'codec':
'jpeg2000',
'quality':
q,
'ssim':
compare_ssim(image,
image_compressed,
multichannel=True,
data_range=1),
'psnr':
compare_psnr(image, image_compressed,
data_range=1),
'msssim':
msssim_value,
'msssim_db':
-10 * np.log10(1 - msssim_value),
'bytes':
image_bytes,
'bpp':
8 * image_bytes / image.shape[0] / image.shape[1]
},
ignore_index=True)
pbar.set_postfix(image_id=image_id, quality=q)
pbar.update(1)
df.to_csv(df_jpeg_path, index=False)
return df
def get_jpeg_df(directory,
write_files=False,
effective_bytes=True,
force_calc=False):
"""
Compute and return (as Pandas DF) the rate distortion curve for JPEG. The result is saved
as a CSV file in the source directory. If the file exists, the DF is loaded and returned.
Files are saved as JPEG using imageio.
"""
files, _ = loading.discover_images(directory, n_images=-1, v_images=0)
batch_x = loading.load_images(files, directory, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
# Get trade-off for JPEG
quality_levels = np.arange(95, 5, -5)
df_jpeg_path = os.path.join(directory, 'jpeg.csv')
if os.path.isfile(df_jpeg_path) and not force_calc:
logger.info('Restoring JPEG stats from {}'.format(df_jpeg_path))
df = pd.read_csv(df_jpeg_path, index_col=False)
else:
df = pd.DataFrame(columns=[
'image_id', 'filename', 'codec', 'quality', 'ssim', 'psnr',
'msssim', 'msssim_db', 'bytes', 'bpp'
])
with tqdm.tqdm(total=len(files) * len(quality_levels),
ncols=120,
desc='JPEG') as pbar:
for image_id, filename in enumerate(files):
# Read the original image
image = batch_x[image_id]
for qi, q in enumerate(quality_levels):
# Compress images and get effective bytes (only image data - no headers)
image_compressed, image_bytes = jpeg_helpers.compress_batch(
image, q, effective=effective_bytes)
if write_files:
image_dir = os.path.join(directory,
os.path.splitext(filename)[0])
if not os.path.isdir(image_dir):
os.makedirs(image_dir)
image_path = os.path.join(image_dir,
'jpeg_q{:03d}.png'.format(q))
imageio.imwrite(image_path,
(255 * image_compressed).astype(
np.uint8))
msssim_value = msssim(image, image_compressed, MAX=1).real
df = df.append(
{
'image_id':
image_id,
'filename':
filename,
'codec':
'jpeg',
'quality':
q,
'ssim':
compare_ssim(image,
image_compressed,
multichannel=True,
data_range=1),
'psnr':
compare_psnr(image, image_compressed,
data_range=1),
'msssim':
msssim_value,
'msssim_db':
-10 * np.log10(1 - msssim_value),
'bytes':
image_bytes,
'bpp':
8 * image_bytes / image.shape[0] / image.shape[1]
},
ignore_index=True)
pbar.set_postfix(image_id=image_id, quality=q)
pbar.update(1)
df.to_csv(os.path.join(directory, 'jpeg.csv'), index=False)
return df
def get_diff_jpeg_df(directory,
write_files=False,
effective_bytes=True,
force_calc=False):
"""
Compute and return (as Pandas DF) the rate distortion curve for diff JPEG. The result is saved
as a CSV file in the source directory. If the file exists, the DF is loaded and returned.
Files are saved as JPEG using imagemagick
"""
files, _ = loading.discover_images(directory, n_images=-1, v_images=0)
batch_x = loading.load_images(files, directory, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
# Get trade-off for JPEG
alpha_range = np.arange(1000, 100, -100)
quality_levels = np.arange(95, 5, -5)
df_jpeg_path = os.path.join(directory, 'diff_jpeg.csv')
if os.path.isfile(df_jpeg_path) and not force_calc:
print('Restoring diff JPEG stats from {}'.format(df_jpeg_path))
df = pd.read_csv(df_jpeg_path, index_col=False)
else:
df = pd.DataFrame(columns=[
'image_id', 'filename', 'codec', 'quality', 'ssim', 'psnr',
'msssim', 'msssim_db', 'bytes', 'bpp'
])
data = dataset.Dataset(directory,
n_images=1,
v_images=0,
val_n_patches=1,
load='y')
batch_rgb = data.next_training_batch(0, 1, 64)
codec_differentiable = jpeg.JPEG(50, codec='soft', trainable=True)
with tqdm.tqdm(total=len(files) * len(quality_levels),
ncols=120,
desc='JPEG') as pbar:
for image_id, filename in enumerate(files):
for q_i, q in enumerate(quality_levels):
codec_differentiable.quality = q
# Read the original image
image = batch_x[image_id]
luma, chroma = codec_differentiable.train_q_table(
batch_rgb, alpha=1, beta=100, n_times=20)
with open(
'data/diff_jpeg/image1/quantization-table.xml', 'r'
) as input_file, open(
'data/diff_jpeg/image1/quantization-table-diff.xml',
'w') as output_file:
for line in input_file:
xml_data = input_file.readlines()
index_1 = 43
for i in range(0, len(luma.numpy())):
q_table_this_line_luma = luma.numpy()[i]
xml_data[index_1 + i] = " " + str(
int(q_table_this_line_luma[0])) + ", " + str(
int(q_table_this_line_luma[1])
) + ", " + str(int(
q_table_this_line_luma[2])) + ", " + str(
int(q_table_this_line_luma[3])
) + ", " + str(
int(q_table_this_line_luma[4])
) + ", " + str(
int(q_table_this_line_luma[5])
) + ", " + str(
int(q_table_this_line_luma[6])
) + ", " + str(
int(q_table_this_line_luma[7])
) + ", " + "\n"
index_2 = 68
for i in range(0, len(chroma.numpy())):
q_table_this_line_chroma = chroma.numpy()[i]
xml_data[index_2 + i] = " " + str(
int(q_table_this_line_chroma[0])) + ", " + str(
int(q_table_this_line_chroma[1])
) + ", " + str(int(
q_table_this_line_chroma[2])) + ", " + str(
int(q_table_this_line_chroma[3])
) + ", " + str(
int(q_table_this_line_chroma[4])
) + ", " + str(
int(q_table_this_line_chroma[5])
) + ", " + str(
int(q_table_this_line_chroma[6])
) + ", " + str(
int(q_table_this_line_chroma[7])
) + ", " + "\n"
output_file.writelines(xml_data)
# Compress images and get effective bytes (only image data - no headers)
os.system(
'magick convert -quality ' + str(50) +
' -define jpeg:q-table=data/diff_jpeg/image1/quantization-table-diff.xml'
+ ' ' + directory + '/' + filename + ' ' + directory +
'/' + filename + '_' + str(50) + '_compressed.jpeg')
if effective_bytes:
with open(
directory + '/' + filename + "_" + str(50) +
'_compressed.jpeg', 'rb') as fh:
buf = io.BytesIO(fh.read())
image_bytes = JPEGMarkerStats(
buf.getvalue()).get_effective_bytes()
else:
image_bytes = os.path.getsize(directory + '/' +
filename + "_" +
str(50) +
'_compressed.jpeg')
image = imageio.imread(directory + '/' + filename).astype(
np.float) / (2**8 - 1)
image_compressed = imageio.imread(
directory + '/' + filename + "_" + str(50) +
'_compressed.jpeg').astype(np.float) / (2**8 - 1)
image_compressed_path = directory + '/' + filename + "_" + str(
50) + '_compressed.jpeg'
image_path = directory + '/' + filename
if not write_files:
os.remove(directory + '/' + filename + "_" + str(50) +
'_compressed.jpeg')
msssim_value = msssim(image, image_compressed, MAX=1).real
df = df.append(
{
'image_id':
image_id,
'filename':
filename,
'codec':
'jpeg',
'quality':
q,
'ssim':
compare_ssim(image,
image_compressed,
multichannel=True,
data_range=1),
'psnr':
compare_psnr(image, image_compressed,
data_range=1),
'msssim':
msssim_value,
'msssim_db':
-10 * np.log10(1 - msssim_value),
#'perceptual similarity': perceptual_similarity,
'bytes':
image_bytes,
'bpp':
8 * image_bytes / image.shape[0] / image.shape[1]
},
ignore_index=True)
pbar.set_postfix(image_id=image_id, quality=q)
pbar.update(1)
df.to_csv(os.path.join(directory, 'diff_jpeg.csv'), index=False)
return df
def get_jpeg_imagemagick_df(directory,
write_files=False,
effective_bytes=True,
force_calc=False):
"""
Compute and return (as Pandas DF) the rate distortion curve for JPEG. The result is saved
as a CSV file in the source directory. If the file exists, the DF is loaded and returned.
Files are saved as JPEG using imageio.
"""
files, _ = loading.discover_images(directory, n_images=-1, v_images=0)
batch_x = loading.load_images(files, directory, load='y')
batch_x = batch_x['y'].astype(np.float32) / (2**8 - 1)
# Get trade-off for JPEG
quality_levels = np.arange(95, 5, -5)
df_jpeg_path = os.path.join(directory, 'jpeg.csv')
if os.path.isfile(df_jpeg_path) and not force_calc:
logger.info('Restoring JPEG stats from {}'.format(df_jpeg_path))
df = pd.read_csv(df_jpeg_path, index_col=False)
else:
df = pd.DataFrame(columns=[
'image_id', 'filename', 'codec', 'quality', 'ssim', 'psnr',
'msssim', 'msssim_db', 'bytes', 'bpp'
])
with tqdm.tqdm(total=len(files) * len(quality_levels),
ncols=120,
desc='JPEG') as pbar:
for image_id, filename in enumerate(files):
# Read the original image
image = batch_x[image_id]
for qi, q in enumerate(quality_levels):
os.system(
'magick convert -quality ' + str(q) +
' -define jpeg:q-table=data/diff_jpeg/image1/quantization-table.xml'
+ ' ' + directory + '/' + filename + ' ' + directory +
'/' + filename + '_' + str(q) + '_compressed.jpeg')
if effective_bytes:
with open(
directory + '/' + filename + "_" + str(q) +
'_compressed.jpeg', 'rb') as fh:
buf = io.BytesIO(fh.read())
image_bytes = JPEGMarkerStats(
buf.getvalue()).get_effective_bytes()
else:
image_bytes = os.path.getsize(directory + '/' +
filename + "_" + str(q) +
'_compressed.jpeg')
image = imageio.imread(directory + '/' + filename).astype(
np.float) / (2**8 - 1)
image_compressed = imageio.imread(
directory + '/' + filename + "_" + str(q) +
'_compressed.jpeg').astype(np.float) / (2**8 - 1)
image_compressed_path = directory + '/' + filename + "_" + str(
50) + '_compressed.jpeg'
image_path = directory + '/' + filename
if not write_files:
os.remove(directory + '/' + filename + "_" + str(q) +
'_compressed.jpeg')
msssim_value = msssim(image, image_compressed, MAX=1).real
df = df.append(
{
'image_id':
image_id,
'filename':
filename,
'codec':
'jpeg',
'quality':
q,
'ssim':
compare_ssim(image,
image_compressed,
multichannel=True,
data_range=1),
'psnr':
compare_psnr(image, image_compressed,
data_range=1),
'msssim':
msssim_value,
'msssim_db':
-10 * np.log10(1 - msssim_value),
'bytes':
image_bytes,
'bpp':
8 * image_bytes / image.shape[0] / image.shape[1]
},
ignore_index=True)
pbar.set_postfix(image_id=image_id, quality=q)
pbar.update(1)
df.to_csv(os.path.join(directory, 'jpeg.csv'), index=False)
return df
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 = helpers.utils.match_option(args.plot, supported_plots)
if args.plot == 'batch':
model, stats = codec.restore(args.dcn,
args.patch_size,
fetch_stats=True)
print('Training stats:', stats)
data = dataset.Dataset(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_images(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(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_images(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(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!')
