def pop(message):
message, n_compressed_bits = len_codec.pop(message)
cbits_codec = cs.repeat(codec, n_compressed_bits[0])
message, compressed_bits = cbits_codec.pop(message)
compressed_bits = np.squeeze(compressed_bits).astype(np.uint8)
bytes_buffer = b'FLIF' + bytes(compressed_bits)
process = subprocess.run(decode_command.split(),
input=bytes_buffer,
capture_output=True)
if process.returncode != 0:
raise Exception(f"flif decode failed: {process.stderr}")
im_buffer = np.frombuffer(process.stdout, dtype=np.uint8)
image = cv2.imdecode(im_buffer, flags=1) # this gives in hwc
return message, inverse_im_transform(image)
def test_substack():
n_data = 100
prec = 4
head, tail = cs.base_message((4, 4))
head = np.split(head, 2)
message = head, tail
data = rng.randint(1 << prec, size=(n_data, 2, 4), dtype='uint64')
view_fun = lambda h: h[0]
append, pop = cs.substack(cs.repeat(cs.Uniform(prec), n_data), view_fun)
message_ = append(message, data)
np.testing.assert_array_equal(message_[0][1], message[0][1])
message_, data_ = pop(message_)
np.testing.assert_equal(message, message_)
np.testing.assert_equal(data, data_)
def test_flatten_unflatten(shape, depth=1000):
np.random.seed(0)
p = 8
bits = np.random.randint(1 << p, size=(depth, ) + shape, dtype=np.uint64)
message = cs.base_message(shape)
other_bits_push, _ = cs.repeat(cs.Uniform(p), depth)
message = other_bits_push(message, bits)
flattened = cs.flatten(message)
reconstructed = cs.unflatten(flattened, shape)
assert_message_equal(message, reconstructed)
def test_reshape_head(old_shape, new_shape, depth=1000):
np.random.seed(0)
p = 8
bits = np.random.randint(1 << p, size=(depth,) + old_shape, dtype=np.uint64)
message = cs.empty_message(old_shape)
other_bits_push, _ = cs.repeat(cs.Uniform(p), depth)
message = other_bits_push(message, bits)
resized = cs.reshape_head(message, new_shape)
reconstructed = cs.reshape_head(resized, old_shape)
assert_message_equal(message, reconstructed)
def rvae_variable_size_codec(codec_from_shape,
latent_from_image_shape,
image_count,
dimensions=4,
dimension_bits=16,
previous_dims=0):
size_codec = cs.repeat(cs.Uniform(dimension_bits), dimensions)
def push(message, symbol):
"""push sizes and array in alternating order"""
assert len(symbol.shape) == dimensions
codec = codec_from_shape(symbol.shape)
head_size = np.prod(latent_from_image_shape(symbol.shape)) + np.prod(
symbol.shape)
message = cs.reshape_head(message, (head_size, ))
message = codec.push(message, symbol)
message = cs.reshape_head(message, (1, ))
message = size_codec.push(message, np.array(symbol.shape))
return message
def pop(message):
message, size = size_codec.pop(message)
# TODO make codec 0 dimensional:
size = np.array(size)[:, 0]
assert size.shape == (dimensions, )
size = size.astype(np.int)
head_size = np.prod(latent_from_image_shape(size)) + np.prod(size)
codec = codec_from_shape(tuple(size))
message = cs.reshape_head(message, (head_size, ))
message, symbol = codec.pop(message)
message = cs.reshape_head(message, (1, ))
return message, symbol
return rvae_serial_with_progress([cs.Codec(push, pop)] * image_count,
previous_dims)
def push(message, image):
"""expects image to be chw"""
image = im_transform(image).astype(np.uint8)
success, im_buffer = cv2.imencode(".ppm", image)
process = subprocess.run(encode_command.split(),
input=im_buffer.tobytes(),
capture_output=True)
if process.returncode != 0:
raise Exception(f"flif encode failed: {process.stderr}")
compressed_bytes = process.stdout
# take off the 'FLIF' magic header
compressed_bytes = compressed_bytes[4:]
# can also remove RGB interlaced byte and bytes per chan (next two bytes)
# then there are 3 varints for width, height and number of frames
# https://flif.info/spec.html for details
compressed_bits = list(compressed_bytes) # list of uint8s
n_compressed_bits = len(compressed_bits)
cbits_codec = cs.repeat(codec, n_compressed_bits)
message = cbits_codec.push(message, compressed_bits)
message = len_codec.push(message, np.uint64(n_compressed_bits))
return message
def test_repeat():
precision = 4
n_data = 7
shape = (2, 3, 5)
data = rng.randint(1 << precision, size=(n_data, ) + shape, dtype="uint64")
check_codec(shape, cs.repeat(cs.Uniform(precision), n_data), data)
def compress_samples(model, hparams, step=tf.constant(0), decode=False):
model.set_compression()
test_set = utils.load_training_files_tfrecords(record_pattern=os.path.join(
hparams['tfrecords_dir'], hparams['train_files'] + '*'))
datapoints = list(test_set.unbatch().batch(
hparams['compress_batch_size']).take(hparams['n_compress_datapoint']))
num_pixels = hparams['n_compress_datapoint'] * hparams[
'compress_batch_size'] * hparams['segment_length']
## Load Codec
waveglow_append, waveglow_pop = cs.repeat(
Waveglow_codec(model=model, hparams=hparams),
hparams['n_compress_datapoint'])
## Encode
encode_t0 = time.time()
init_message = cs.empty_message(shape=(hparams['compress_batch_size'],
hparams['segment_length'] // 4))
# Encode the audio samples
message = waveglow_append(init_message, datapoints)
flat_message = cs.flatten(message)
encode_t = time.time() - encode_t0
tf.print("All encoded in {:.2f}s.".format(encode_t))
original_len = 16 * hparams['n_compress_datapoint'] * hparams[
'segment_length']
message_len = 32 * len(flat_message)
tf.print("Used {} bits.".format(message_len))
tf.print("This is {:.2f} bits per pixel.".format(message_len / num_pixels))
tf.print("Compression ratio : {:.2f}".format(original_len / message_len))
tf.summary.scalar(name='bits_per_dim',
data=message_len / num_pixels,
step=step)
tf.summary.scalar(name='compression_ratio',
data=original_len / message_len,
step=step)
if decode:
## Decode
decode_t0 = time.time()
message = cs.unflatten(flat_message,
shape=(hparams['compress_batch_size'],
hparams['segment_length'] // 4))
message, datapoints_ = waveglow_pop(message)
decode_t = time.time() - decode_t0
print('All decoded in {:.2f}s.'.format(decode_t))
datacompare = [
data['wav'].numpy()[..., np.newaxis] for data in datapoints
]
np.testing.assert_equal(datacompare, datapoints_)
np.testing.assert_equal(message, init_message)
model.set_training()
def run_bbans(hps):
from autograd.builtins import tuple as ag_tuple
from rvae.resnet_codec import ResNetVAE
hps.num_gpus = 1
hps.batch_size = 1
batch_size = hps.batch_size
hps.eval_batch_size = batch_size
n_flif = hps.n_flif
_, datasets = images(hps)
datasets = datasets if isinstance(datasets, list) else [datasets]
test_images = [
np.array([image]).astype('uint64') for dataset in datasets
for image in dataset
]
n_batches = len(test_images) // batch_size
test_images = [
np.concatenate(test_images[i * batch_size:(i + 1) * batch_size],
axis=0) for i in range(n_batches)
]
flif_images = test_images[:n_flif]
vae_images = test_images[n_flif:]
num_dims = np.sum([batch.size for batch in test_images])
flif_dims = np.sum([batch.size
for batch in flif_images]) if flif_images else 0
prior_precision = 10
obs_precision = 24
q_precision = 18
@lru_cache(maxsize=1)
def codec_from_shape(shape):
print("Creating codec for shape " + str(shape))
hps.image_size = (shape[2], shape[3])
z_shape = latent_shape(hps)
z_size = np.prod(z_shape)
graph = tf.Graph()
with graph.as_default():
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
x = tf.placeholder(tf.float32, shape, 'x')
model = CVAE1(hps, "eval", x)
stepwise_model = LayerwiseCVAE(model)
saver = tf.train.Saver(model.avg_dict)
config = tf.ConfigProto(allow_soft_placement=True,
intra_op_parallelism_threads=4,
inter_op_parallelism_threads=4)
sess = tf.Session(config=config, graph=graph)
saver.restore(sess, restore_path())
run_all_contexts, run_top_prior, runs_down_prior, run_top_posterior, runs_down_posterior, \
run_reconstruction = stepwise_model.get_model_parts_as_numpy_functions(sess)
# Setup codecs
def vae_view(head):
return ag_tuple(
(np.reshape(head[:z_size],
z_shape), np.reshape(head[z_size:], shape)))
obs_codec = lambda h, z1: cs.Logistic_UnifBins(*run_reconstruction(
h, z1),
obs_precision,
bin_prec=8,
bin_lb=-0.5,
bin_ub=0.5)
return cs.substack(
ResNetVAE(run_all_contexts, run_top_posterior, runs_down_posterior,
run_top_prior, runs_down_prior, obs_codec,
prior_precision, q_precision), vae_view)
is_fixed = not hps.compression_always_variable and \
(len(set([dataset[0].shape[-2:] for dataset in datasets])) == 1)
fixed_size_codec = lambda: cs.repeat(codec_from_shape(vae_images[0].shape),
len(vae_images))
variable_codec_including_sizes = lambda: rvae_variable_size_codec(
codec_from_shape,
latent_from_image_shape=latent_from_image_shape(hps),
image_count=len(vae_images),
previous_dims=flif_dims)
variable_known_sizes_codec = lambda: rvae_variable_known_size_codec(
codec_from_image_shape=codec_from_shape,
latent_from_image_shape=latent_from_image_shape(hps),
shapes=[i.shape for i in vae_images],
previous_dims=flif_dims)
variable_size_codec = \
variable_known_sizes_codec if hps.compression_exclude_sizes else variable_codec_including_sizes
codec = fixed_size_codec if is_fixed else variable_size_codec
vae_push, vae_pop = codec()
np.seterr(divide='raise')
if n_flif:
print('Using FLIF to encode initial images...')
flif_push, flif_pop = cs.repeat(cs.repeat(FLIF, batch_size), n_flif)
message = cs.empty_message((1, ))
message = flif_push(message, flif_images)
else:
print('Creating a random initial message...')
message = cs.random_message(hps.initial_bits, (1, ))
init_head_shape = (np.prod(image_shape(hps)) +
np.prod(latent_shape(hps)) if is_fixed else 1, )
message = cs.reshape_head(message, init_head_shape)
print("Encoding with VAE...")
encode_t0 = time.time()
message = vae_push(message, vae_images)
encode_t = time.time() - encode_t0
print("All encoded in {:.2f}s".format(encode_t))
flat_message = cs.flatten(message)
message_len = 32 * len(flat_message)
print("Used {} bits.".format(message_len))
print("This is {:.2f} bits per dim.".format(message_len / num_dims))
if n_flif == 0:
extra_bits = message_len - 32 * hps.initial_bits
print('Extra bits: {}'.format(extra_bits))
print('This is {:.2f} bits per dim.'.format(extra_bits / num_dims))
print('Decoding with VAE...')
decode_t0 = time.time()
message = cs.unflatten(flat_message, init_head_shape)
message, decoded_vae_images = vae_pop(message)
message = cs.reshape_head(message, (1, ))
decode_t = time.time() - decode_t0
print('All decoded in {:.2f}s'.format(decode_t))
assert len(vae_images) == len(decoded_vae_images), (
len(vae_images), len(decoded_vae_images))
for test_image, decoded_image in zip(vae_images, decoded_vae_images):
np.testing.assert_equal(test_image, decoded_image)
if n_flif:
print('Decoding with FLIF...')
message, decoded_flif_images = flif_pop(message)
for test_image, decoded_image in zip(flif_images, decoded_flif_images):
np.testing.assert_equal(test_image, decoded_image)
assert cs.is_empty(message)
## Setup codecs
# VAE codec
model = BinaryVAE(hidden_dim=100, latent_dim=40)
model.load_state_dict(torch.load('vae_params'))
rec_net = torch_fun_to_numpy_fun(model.encode)
gen_net = torch_fun_to_numpy_fun(model.decode)
obs_codec = lambda p: cs.Bernoulli(p, bernoulli_precision)
def vae_view(head):
return ag_tuple((np.reshape(head[:latent_size], latent_shape),
np.reshape(head[latent_size:], obs_shape)))
vae_append, vae_pop = cs.repeat(cs.substack(
bb_ans.VAE(gen_net, rec_net, obs_codec, prior_precision, q_precision),
vae_view), num_batches)
## Load mnist images
images = datasets.MNIST(sys.argv[1], train=False, download=True).data.numpy()
images = np.uint64(rng.random_sample(np.shape(images)) < images / 255.)
images = np.split(np.reshape(images, (num_images, -1)), num_batches)
## Encode
# Initialize message with some 'extra' bits
encode_t0 = time.time()
init_message = cs.base_message(obs_size + latent_size)
# Encode the mnist images
message, = vae_append(init_message, images)
np.shape(images[0]) + (256,),
obs_elem_idxs,
obs_elem_codec)
def pop_(msg):
msg, (data, _) = pop(msg)
return msg, data
return append, pop_
# Setup codecs
def vae_view(head):
return ag_tuple((np.reshape(head[:latent_size], latent_shape),
np.reshape(head[latent_size:], (batch_size,))))
vae_append, vae_pop = cs.repeat(cs.substack(
bb_ans.VAE(gen_net, rec_net, obs_codec, prior_precision, q_precision),
vae_view), num_batches)
# Codec for adding extra bits to the start of the chain (necessary for bits
# back).
p = prior_precision
other_bits_depth = 10
other_bits_append, _ = cs.substack(cs.repeat(codecs.Uniform(p), other_bits_depth),
lambda h: vae_view(h)[0])
## Encode
# Initialize message with some 'extra' bits
encode_t0 = time.time()
init_message = vrans.x_init(batch_size + latent_size)
other_bits = rng.randint(1 << p, size=(other_bits_depth,) + latent_shape, dtype=np.uint64)
