def prepare_dataloader(args, input_dir, output_dir, batch_size=1):
# `batch_size` must be 1 for images of different shapes
input_images = glob.glob(os.path.join(input_dir, '*.jpg'))
input_images += glob.glob(os.path.join(input_dir, '*.png'))
assert len(input_images) > 0, 'No valid image files found in supplied directory!'
print('Input images')
pprint(input_images)
eval_loader = datasets.get_dataloaders('evaluation', root=input_dir, batch_size=batch_size,
logger=None, shuffle=False, normalize=args.normalize_input_image)
utils.makedirs(output_dir)
return eval_loader
def end_of_epoch_metrics(args, model, data_loader, device, logger):
model.eval()
old_mode = model.model_mode
#model.set_model_mode(ModelModes.EVALUATION)
model.training = False
classi_acc_total = []
n, N = 0, len(data_loader.dataset)
input_filenames_total = list()
output_filenames_total = list()
q_bpp_total, q_bpp_total_attained, LPIPS_total = torch.Tensor(
N), torch.Tensor(N), torch.Tensor(N)
SSIM_total, PSNR_total = torch.Tensor(N), torch.Tensor(N)
comp_loss_total, classi_loss_total, classi_acc_total1 = torch.Tensor(
N), torch.Tensor(N), torch.Tensor(N)
with torch.no_grad():
thisIndx = 0
for idx1, (dataAll, yAll) in enumerate(tqdm(data_loader), 0):
dataAll = dataAll.to(device, dtype=torch.float)
yAll = yAll.to(device)
losses, intermediates = model(dataAll,
yAll,
return_intermediates=True,
writeout=True)
classi_acc = losses['classi_acc']
classi_acc_total.append(classi_acc.item())
# Reproducibility
make_deterministic()
perceptual_loss_fn = ps.PerceptualLoss(model='net-lin',
net='alex',
use_gpu=torch.cuda.is_available())
# Build probability tables
logger.info('Building hyperprior probability tables...')
model.Hyperprior.hyperprior_entropy_model.build_tables()
logger.info('All tables built.')
max_value = 255.
SSIM_func = metrics.SSIM(data_range=max_value)
utils.makedirs(args.output_dir)
logger.info('Starting compression...')
start_time = time.time()
with torch.no_grad():
thisIndx = 0
for idx1, (dataAll, yAll) in enumerate(tqdm(data_loader), 0):
dataAll = dataAll.to(device, dtype=torch.float)
yAll = yAll.to(device)
#if idx1 > 2:
# break
B = dataAll.size(0)
for idxB in range(B):
data = dataAll[idxB, :, :, :]
data = data.unsqueeze(0)
y = yAll[idxB]
y = y.unsqueeze(0)
model.set_model_mode(old_mode)
model.training = False
losses = model(data, y, train_generator=False)
compression_loss = losses['compression']
if model.use_classiOnly is True:
classi_loss = losses['classi']
classi_acc = losses['classi_acc']
model.set_model_mode(ModelModes.EVALUATION)
model.training = False
# Perform entropy coding
q_bpp_attained, compressed_output = model.compress(data,
silent=True)
if args.save is True:
compression_utils.save_compressed_format(
compressed_output,
out_path=os.path.join(args.output_dir,
"compressed.hfc"))
reconstruction = model.decompress(compressed_output)
q_bpp = compressed_output.total_bpp
if args.normalize_input_image is True:
# [-1., 1.] -> [0., 1.]
data = (data + 1.) / 2.
perceptual_loss = perceptual_loss_fn.forward(reconstruction,
data,
normalize=True)
# [0., 1.] -> [0., 255.]
psnr = metrics.psnr(reconstruction.cpu().numpy() * max_value,
data.cpu().numpy() * max_value, max_value)
ms_ssim = SSIM_func(reconstruction * max_value,
data * max_value)
PSNR_total[thisIndx] = torch.Tensor(psnr)
SSIM_total[thisIndx] = ms_ssim.data
q_bpp_per_im = float(q_bpp.item()) if type(
q_bpp) == torch.Tensor else float(q_bpp)
fname = os.path.join(
args.output_dir,
"{}_RECON_{:.3f}bpp.png".format(thisIndx, q_bpp_per_im))
torchvision.utils.save_image(reconstruction,
fname,
normalize=True)
output_filenames_total.append(fname)
q_bpp_total[thisIndx] = q_bpp.data if type(
q_bpp) == torch.Tensor else q_bpp
q_bpp_total_attained[thisIndx] = q_bpp_attained.data if type(
q_bpp_attained) == torch.Tensor else q_bpp_attained
LPIPS_total[thisIndx] = perceptual_loss.data
comp_loss_total[thisIndx] = compression_loss.data
if model.use_classiOnly is True:
classi_loss_total[thisIndx] = classi_loss.data
classi_acc_total1[thisIndx] = classi_acc.data
thisIndx = thisIndx + 1
logger.info(
f'BPP: mean={q_bpp_total.mean(dim=0):.3f}, std={q_bpp_total.std(dim=0):.3f}'
)
logger.info(
f'BPPA: mean={q_bpp_total_attained.mean(dim=0):.3f}, std={q_bpp_total_attained.std(dim=0):.3f}'
)
logger.info(
f'LPIPS: mean={LPIPS_total.mean(dim=0):.3f}, std={LPIPS_total.std(dim=0):.3f}'
)
logger.info(
f'PSNR: mean={PSNR_total.mean(dim=0):.3f}, std={PSNR_total.std(dim=0):.3f}'
)
logger.info(
f'SSIM: mean={SSIM_total.mean(dim=0):.3f}, std={SSIM_total.std(dim=0):.3f}'
)
logger.info(
f'CompLoss: mean={comp_loss_total.mean(dim=0):.3f}, std={comp_loss_total.std(dim=0):.3f}'
)
logger.info(
f'ClassiLoss: mean={classi_loss_total.mean(dim=0):.3f}, std={classi_loss_total.std(dim=0):.3f}'
)
logger.info(
f'ClassiAcc1: mean={classi_acc_total1.mean(dim=0):.3f}, std={classi_acc_total1.std(dim=0):.3f}'
)
logger.info(f'ClassiAcc2: mean={np.mean(classi_acc_total):.3f}')
#df = pd.DataFrame([input_filenames_total, output_filenames_total]).T
#df.columns = ['input_filename', 'output_filename']
#df['bpp_original'] = bpp_total.cpu().numpy()
#df['q_bpp'] = q_bpp_total.cpu().numpy()
#df['LPIPS'] = LPIPS_total.cpu().numpy()
#df['PSNR'] = PSNR_total.cpu().numpy()
#df['MS_SSIM'] = MS_SSIM_total.cpu().numpy()
#df_path = os.path.join(args.output_dir, 'compression_metrics.h5')
#df.to_hdf(df_path, key='df')
#pprint(df)
#logger.info('Complete. Reconstructions saved to {}. Output statistics saved to {}'.format(args.output_dir, df_path))
delta_t = time.time() - start_time
logger.info('Time elapsed: {:.3f} s'.format(delta_t))
logger.info('Rate: {:.3f} Images / s:'.format(float(N) / delta_t))
model.set_model_mode(old_mode)
def compress_and_decompress(args):
# Reproducibility
make_deterministic()
perceptual_loss_fn = ps.PerceptualLoss(model='net-lin', net='alex', use_gpu=torch.cuda.is_available())
# Load model
device = utils.get_device()
logger = utils.logger_setup(logpath=os.path.join(args.image_dir, 'logs'), filepath=os.path.abspath(__file__))
loaded_args, model, _ = utils.load_model(args.ckpt_path, logger, device, model_mode=ModelModes.EVALUATION,
current_args_d=None, prediction=True, strict=False)
# Override current arguments with recorded
dictify = lambda x: dict((n, getattr(x, n)) for n in dir(x) if not (n.startswith('__') or 'logger' in n))
loaded_args_d, args_d = dictify(loaded_args), dictify(args)
loaded_args_d.update(args_d)
args = utils.Struct(**loaded_args_d)
logger.info(loaded_args_d)
# Build probability tables
logger.info('Building hyperprior probability tables...')
model.Hyperprior.hyperprior_entropy_model.build_tables()
logger.info('All tables built.')
eval_loader = datasets.get_dataloaders('evaluation', root=args.image_dir, batch_size=args.batch_size,
logger=logger, shuffle=False, normalize=args.normalize_input_image)
n, N = 0, len(eval_loader.dataset)
input_filenames_total = list()
output_filenames_total = list()
bpp_total, q_bpp_total, LPIPS_total = torch.Tensor(N), torch.Tensor(N), torch.Tensor(N)
utils.makedirs(args.output_dir)
logger.info('Starting compression...')
start_time = time.time()
with torch.no_grad():
for idx, (data, bpp, filenames) in enumerate(tqdm(eval_loader), 0):
data = data.to(device, dtype=torch.float)
B = data.size(0)
input_filenames_total.extend(filenames)
if args.reconstruct is True:
# Reconstruction without compression
reconstruction, q_bpp = model(data, writeout=False)
else:
# Perform entropy coding
compressed_output = model.compress(data)
if args.save is True:
assert B == 1, 'Currently only supports saving single images.'
compression_utils.save_compressed_format(compressed_output,
out_path=os.path.join(args.output_dir, f"{filenames[0]}_compressed.hfc"))
reconstruction = model.decompress(compressed_output)
q_bpp = compressed_output.total_bpp
if args.normalize_input_image is True:
# [-1., 1.] -> [0., 1.]
data = (data + 1.) / 2.
perceptual_loss = perceptual_loss_fn.forward(reconstruction, data, normalize=True)
for subidx in range(reconstruction.shape[0]):
if B > 1:
q_bpp_per_im = float(q_bpp.cpu().numpy()[subidx])
else:
q_bpp_per_im = float(q_bpp.item()) if type(q_bpp) == torch.Tensor else float(q_bpp)
fname = os.path.join(args.output_dir, "{}_RECON_{:.3f}bpp.png".format(filenames[subidx], q_bpp_per_im))
torchvision.utils.save_image(reconstruction[subidx], fname, normalize=True)
output_filenames_total.append(fname)
bpp_total[n:n + B] = bpp.data
q_bpp_total[n:n + B] = q_bpp.data if type(q_bpp) == torch.Tensor else q_bpp
LPIPS_total[n:n + B] = perceptual_loss.data
n += B
df = pd.DataFrame([input_filenames_total, output_filenames_total]).T
df.columns = ['input_filename', 'output_filename']
df['bpp_original'] = bpp_total.cpu().numpy()
df['q_bpp'] = q_bpp_total.cpu().numpy()
df['LPIPS'] = LPIPS_total.cpu().numpy()
df_path = os.path.join(args.output_dir, 'compression_metrics.h5')
df.to_hdf(df_path, key='df')
pprint(df)
logger.info('Complete. Reconstructions saved to {}. Output statistics saved to {}'.format(args.output_dir, df_path))
delta_t = time.time() - start_time
logger.info('Time elapsed: {:.3f} s'.format(delta_t))
logger.info('Rate: {:.3f} Images / s:'.format(float(N) / delta_t))