def test_range_encoder_decoder():
"""
Additional tests to check whether RangeDecoder reproduces symbols encoded by RangeEncoder.
"""
random.seed(0)
filepath = mkstemp()[1]
for _ in range(20):
numSymbols = np.random.randint(1, 6)
cumFreq = [0] + np.cumsum(np.random.randint(1, 10,
size=numSymbols)).tolist()
data = np.random.randint(numSymbols,
size=np.random.randint(20)).tolist()
encoder = RangeEncoder(filepath)
encoder.encode(data, cumFreq)
encoder.close()
decoder = RangeDecoder(filepath)
dataRec = decoder.decode(len(data), cumFreq)
decoder.close()
assert data == dataRec
os.remove(filepath)
def test_range_decoder_fuzz():
"""
Test random inputs to the decoder.
"""
random.seed(827)
randomState = np.random.RandomState(827)
for _ in range(10):
# generate random frequency table
numSymbols = random.randint(1, 20)
maxFreq = random.randint(2, 100)
cumFreq = np.cumsum(randomState.randint(1, maxFreq, size=numSymbols))
cumFreq = [0] + [int(i) for i in cumFreq] # convert numpy.int64 to int
# decode random symbols
decoder = RangeDecoder('/dev/urandom')
decoder.decode(100, cumFreq)
def decompress(input_bin_path, input_res_path, output_img_path, ckp_dir, tau):
with tf.device('/cpu:0'):
# Load bin and res
string = tf.placeholder(tf.string, [1])
side_string = tf.placeholder(tf.string, [1])
x_shape = tf.placeholder(tf.int32, [2])
y_shape = tf.placeholder(tf.int32, [2])
z_shape = tf.placeholder(tf.int32, [2])
with open(input_bin_path, "rb") as f:
packed = tfc.PackedTensors(f.read())
tensors = [string, side_string, x_shape, y_shape, z_shape]
arrays = packed.unpack(tensors)
# instantiate model
decoder = nll_codec.Decoder(192)
hyper_decoder_sigma = nll_codec.HyperDecoder(192)
hyper_decoder_mu = nll_codec.HyperDecoder(192)
entropy_parameters_sigma = nll_codec.EntropyParameters(192)
entropy_parameters_mu = nll_codec.EntropyParameters(192)
entropy_bottleneck = tfc.EntropyBottleneck(dtype=tf.float32)
res_compressor = nll_codec.ResidualCompressor(128, 5)
masked_conv = nll_codec.MaskedConv2d("A", 64, (5, 5), padding="VALID")
res_compressor_cond = bc.ResidualCompressor_cond(128, 5)
# build decoder
z_shape = tf.concat([z_shape, [192]], axis=0)
z_hat_decode = entropy_bottleneck.decompress(
side_string, z_shape,
channels=192) # decode z (including dequantization)
psi_sigma = hyper_decoder_sigma(z_hat_decode)
psi_mu = hyper_decoder_mu(z_hat_decode)
sigma = entropy_parameters_sigma(psi_sigma)
mu = entropy_parameters_mu(psi_mu)
sigma = sigma[:, :y_shape[0], :y_shape[1], :]
mu = mu[:, :y_shape[0], :y_shape[1], :]
scale_table = np.exp(
np.linspace(np.log(SCALE_MIN), np.log(SCALE_MAX), SCALES_LEVELS))
conditional_bottleneck = tfc.GaussianConditional(sigma,
scale_table,
mean=mu,
dtype=tf.float32)
y_hat_decode = conditional_bottleneck.decompress(
string) # decode y (including dequantization)
x_hat, res_prior = decoder(y_hat_decode)
x_hat = x_hat[:, :x_shape[0], :x_shape[1], :]
x_hat = tf.clip_by_value(x_hat, 0, 1)
x_hat = tf.math.floor(x_hat * 255 + 0.5)
res_prior = res_prior[:, :x_shape[0], :x_shape[1], :]
tau_list = tf.constant([int(tau - 1)], tf.int32)
cond = tf.one_hot(tau_list, 5)
num_pixels = tf.cast(tf.reduce_prod(x_shape[:-1]), dtype=tf.float32)
res_q_patch = tf.placeholder(dtype=tf.float32, shape=(1, 5, 5, 3))
res_prior_channel_num = 64
res_prior_patch = tf.placeholder(dtype=tf.float32,
shape=(1, 1, 1,
res_prior_channel_num))
res_q_vector = tf.placeholder(dtype=tf.float32, shape=(1, 1, 1, 3))
bin_sz = 2 * tau + 1
pmf_length = int(510 // bin_sz + 1)
pmf_end = (255 // bin_sz) * bin_sz
context = masked_conv(res_q_patch)
res_prior_context = tf.concat([res_prior_patch, context], 3)
bias_correction = True
if bias_correction and int(tau) > 0:
res_mu, res_log_sigma, res_pi, res_lambda = res_compressor_cond(
res_prior_context, cond)
else:
res_mu, res_log_sigma, res_pi, res_lambda = res_compressor(
res_prior_context)
res_mu_tiled = tf.tile(res_mu,
tf.constant([pmf_length, 1, 1, 1], tf.int32))
res_log_sigma_tiled = tf.tile(
res_log_sigma, tf.constant([pmf_length, 1, 1, 1], tf.int32))
res_pi_tiled = tf.tile(res_pi,
tf.constant([pmf_length, 1, 1, 1], tf.int32))
res_lambda_tiled = tf.tile(
res_lambda, tf.constant([pmf_length, 1, 1, 1], tf.int32))
res_bottleneck = lmm.LogisticMixtureModel(res_mu_tiled,
res_log_sigma_tiled,
res_pi_tiled,
res_lambda_tiled)
res_pmf = res_bottleneck.pmf_tau(res_q_vector, tau)
with tf.Session() as sess:
latest = tf.train.latest_checkpoint(checkpoint_dir=ckp_dir)
tf.train.Saver().restore(sess, save_path=latest)
# lossy image decoding
print("Lossy Image Decoding Start.")
res_prior_out, x_out, num_pixels_out, x_shape_out = sess.run(
[res_prior, x_hat, num_pixels, x_shape],
feed_dict=dict(zip(tensors, arrays)))
print("Lossy Image Decoding Finish.")
k_sz = 5
pad_sz = 2
x_h, x_w = x_shape_out
x_c = 3
res_q_dec_padded = np.zeros(
(1, x_h + 2 * pad_sz, x_w + 2 * pad_sz, x_c))
decoder = RangeDecoder(input_res_path)
print('Residual Decoding Start.')
for h_idx in range(x_h):
for w_idx in range(x_w):
res_q_extracted = res_q_dec_padded[:, h_idx:h_idx + k_sz,
w_idx:w_idx + k_sz, :]
res_prior_extracted = res_prior_out[:, h_idx,
w_idx, :].reshape(
1, 1, 1,
res_prior_channel_num
)
for c_idx in range(x_c):
res_q_vector_extracted = res_q_dec_padded[:, h_idx +
pad_sz,
w_idx +
pad_sz, :].reshape(
1, 1, 1,
3)
res_pmf_out = sess.run(res_pmf,
feed_dict={
res_q_patch:
res_q_extracted,
res_prior_patch:
res_prior_extracted,
res_q_vector:
res_q_vector_extracted
})
c_pmf = res_pmf_out[:, 0, 0, c_idx]
c_pmf = np.clip(c_pmf, 1.0 / 65025, 1.0)
c_pmf = c_pmf / np.sum(c_pmf)
cumFreq = np.floor(
np.append([0.], np.cumsum(c_pmf)) * 65536. +
0.5).astype(np.int32).tolist()
dataRec = decoder.decode(1, cumFreq)
res_q_dec_padded[0, h_idx + pad_sz, w_idx + pad_sz,
c_idx] = dataRec[0] * bin_sz - pmf_end
print("Decode Finish.")
decoder.close()
res_q_dec = res_q_dec_padded[:, pad_sz:x_h + pad_sz,
pad_sz:x_w + pad_sz, :]
x_rec = np.clip(np.squeeze(x_out + res_q_dec), 0, 255)
im = Image.fromarray(np.uint8(x_rec))
im.save(output_img_path)
return x_rec
def test_range_decoder():
"""
Tests whether RangeDecoder reproduces symbols encoded by RangeEncoder.
"""
random.seed(558)
filepath = mkstemp()[1]
# encoding one sequence should require 1 byte
cumFreq0 = [0, 4, 6, 8]
cumFreq1 = [0, 2, 5, 7, 10, 14]
data0 = [random.randint(0, len(cumFreq0) - 2) for _ in range(10)]
data1 = [random.randint(0, len(cumFreq1) - 2) for _ in range(17)]
encoder = RangeEncoder(filepath)
encoder.encode(data0, cumFreq0)
encoder.encode(data1, cumFreq1)
encoder.close()
decoder = RangeDecoder(filepath)
dataRec0 = decoder.decode(len(data0), cumFreq0)
dataRec1 = decoder.decode(len(data1), cumFreq1)
decoder.close()
# encoded and decoded data should be the same
assert data0 == dataRec0
assert data1 == dataRec1
# make sure reference counting is implemented correctly (call to getrefcount increases it by 1)
assert sys.getrefcount(dataRec0) == 2
assert sys.getrefcount(dataRec1) == 2
decoder = RangeDecoder(filepath)
with pytest.raises(ValueError):
# invalid frequency table
decoder.decode(len(data0), [])
with pytest.raises(ValueError):
# invalid frequency table
decoder.decode(len(data0), [0])
assert decoder.decode(0, cumFreq0) == []
os.remove(filepath)
from range_coder import RangeEncoder, RangeDecoder, prob_to_cum_freq import os data = [2, 0, 1, 0, 0, 0, 1, 2, 2] prob = [0.5, 0.2, 0.3] # convert probabilities to cumulative integer frequency table cumFreq = prob_to_cum_freq(prob, resolution=4) print(cumFreq) filepath="output.txt" # encode data encoder = RangeEncoder(filepath) encoder.encode(data, cumFreq) encoder.close() # decode data decoder = RangeDecoder(filepath) dataRec = decoder.decode(len(data), cumFreq) decoder.close() print(os.stat(filepath)) print (dataRec)
def decompress(input, output, num_filters, checkpoint_dir):
"""Decompresses an image by a fast implementation."""
start = time.time()
tf.set_random_seed(1)
tf.reset_default_graph()
with tf.device('/cpu:0'):
print(input)
# Read the shape information and compressed string from the binary file.
fileobj = open(input, mode='rb')
x_shape = np.frombuffer(fileobj.read(4), dtype=np.uint16)
length, minmax = np.frombuffer(fileobj.read(4), dtype=np.uint16)
num = np.frombuffer(fileobj.read(16), dtype=np.uint8)
string = fileobj.read(length)
fileobj.close()
flag = np.unpackbits(num)
non_zero_idx = np.squeeze(np.where(flag == 1))
# Get x_pad_shape, y_shape, z_shape
pad_size = 64
x_pad_shape = [1] + [
int(math.ceil(x_shape[0] / pad_size) * pad_size)
] + [int(math.ceil(x_shape[1] / pad_size) * pad_size)] + [3]
y_shape = [1] + [x_pad_shape[1] // 16] + [x_pad_shape[2] // 16
] + [num_filters]
z_shape = [y_shape[1] // 4] + [y_shape[2] // 4] + [num_filters]
# Add a batch dimension, then decompress and transform the image back.
strings = tf.expand_dims(string, 0)
entropy_bottleneck = tfc.EntropyBottleneck(dtype=tf.float32)
z_tilde = entropy_bottleneck.decompress(strings,
z_shape,
channels=num_filters)
phi = hyper_synthesis(z_tilde, num_filters)
# Transform the quantized image back (if requested).
tiny_y = tf.placeholder(dtype=tf.float32,
shape=[1] + [5] + [5] + [num_filters])
tiny_phi = tf.placeholder(dtype=tf.float32,
shape=[1] + [5] + [5] + [num_filters * 2])
_, _, means, variances, probs = entropy_parameter(tiny_phi,
tiny_y,
num_filters,
training=False)
# Decode the x_hat usign the decoded y
y_hat = tf.placeholder(dtype=tf.float32, shape=y_shape)
x_hat = synthesis_transform(y_hat, num_filters)
# Remove batch dimension, and crop away any extraneous padding on the bottom or right boundaries.
x_hat = x_hat[0, :int(x_shape[0]), :int(x_shape[1]), :]
# Write reconstructed image out as a PNG file.
op = save_image(output, x_hat)
# Load the latest model checkpoint, and perform the above actions.
with tf.Session() as sess:
#latest = tf.train.latest_checkpoint(checkpoint_dir=checkpoint_dir)
latest = "models/model-1399000" #lambda = 14
print(latest)
tf.train.Saver().restore(sess, save_path=latest)
phi_value = sess.run(phi)
print("INFO: start decoding y")
print(time.time() - start)
decoder = RangeDecoder(input[:-4] + '.bin')
samples = np.arange(0, minmax * 2 + 1)
TINY = 1e-10
# Fast implementation to decode the y_hat
kernel_size = 5
pad_size = (kernel_size - 1) // 2
decoded_y = np.zeros([1] + [y_shape[1] + kernel_size - 1] +
[y_shape[2] + kernel_size - 1] +
[num_filters])
padded_phi = np.pad(phi_value, ((0, 0), (pad_size, pad_size),
(pad_size, pad_size), (0, 0)),
'constant',
constant_values=((0., 0.), (0., 0.), (0., 0.),
(0., 0.)))
for h_idx in range(y_shape[1]):
for w_idx in range(y_shape[2]):
y_means, y_variances, y_probs = \
sess.run([means, variances, probs], \
feed_dict={tiny_y: decoded_y[:, h_idx: h_idx+kernel_size, w_idx:w_idx+kernel_size, :], \
tiny_phi: padded_phi[:, h_idx: h_idx+kernel_size, w_idx:w_idx+kernel_size, :]})
for i in range(len(non_zero_idx)):
ch_idx = non_zero_idx[i]
mu = y_means[0, pad_size, pad_size, ch_idx, :] + minmax
sigma = y_variances[0, pad_size, pad_size, ch_idx, :]
weight = y_probs[0, pad_size, pad_size, ch_idx, :]
pmf = (0.5 * (1 + scipy.special.erf((samples + 0.5 - mu[0]) / ((sigma[0] + TINY) * 2 ** 0.5))) - \
0.5 * (1 + scipy.special.erf((samples - 0.5 - mu[0]) / ((sigma[0] + TINY) * 2 ** 0.5)))) * weight[0] + \
(0.5 * (1 + scipy.special.erf((samples + 0.5 - mu[1]) / ((sigma[1] + TINY) * 2 ** 0.5))) - \
0.5 * (1 + scipy.special.erf((samples - 0.5 - mu[1]) / ((sigma[1] + TINY) * 2 ** 0.5)))) * weight[1] +\
(0.5 * (1 + scipy.special.erf((samples + 0.5 - mu[2]) / ((sigma[2] + TINY) * 2 ** 0.5))) - \
0.5 * (1 + scipy.special.erf((samples - 0.5 - mu[2]) / ((sigma[2] + TINY) * 2 ** 0.5)))) * weight[2]
pmf_clip = np.clip(pmf, 1.0 / 65536, 1.0)
pmf_clip = np.round(pmf_clip / np.sum(pmf_clip) *
65536)
cdf = list(np.add.accumulate(pmf_clip))
cdf = [0] + [int(i) for i in cdf]
decoded_y[0, h_idx + pad_size, w_idx + pad_size,
ch_idx] = decoder.decode(1, cdf)[0] - minmax
decoded_y = decoded_y[:, pad_size:y_shape[1] + pad_size,
pad_size:y_shape[2] + pad_size, :]
sess.run(op, feed_dict={y_hat: decoded_y})
end = time.time()
print("Time (s): {:0.3f}".format(end - start))