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)