compression_example.py

from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
import zlib

def generate_const_image_bw(nx, ny):
    im = np.zeros((nx, ny), dtype='uint8')
    return im

def generate_checkerboard_image_bw(nx, ny):
    im = np.zeros(nx*ny, dtype='uint8')
    im[::2] = 1
    return im.reshape(nx, ny)

def generate_random_image_bw(nx, ny):
    rand_matrix = np.random.randint(0, 2, size=(nx, ny), dtype='uint8')
    return rand_matrix

def get_comp_size_bytes(im, complevel=9):
    compr_str = zlib.compress(im.astype('uint8').tostring(), complevel)
    return len(compr_str)

if __name__ == "__main__":
    n_list = np.arange(16, 512, 8)
    im_const_size_list = []
    im_check_size_list = []
    im_rand_size_list = []
    for n in n_list:
        const_image = generate_const_image_bw(n, n)
        im_size = get_comp_size_bytes(const_image)
        im_const_size_list.append(im_size)
        #
        check_image = generate_checkerboard_image_bw(n, n)
        im_size = get_comp_size_bytes(check_image)
        im_check_size_list.append(im_size)
        #
        rand_image = generate_random_image_bw(n, n)
        im_size = get_comp_size_bytes(rand_image)
        im_rand_size_list.append(im_size)

    fig0 = plt.figure()
    ax0 = fig0.add_subplot(111)
    ax0.plot(n_list ** 2, n_list ** 2, color='k', label='raw')
    ax0.plot(n_list ** 2, im_const_size_list, color='g', label='comp. constant')
    ax0.plot(n_list ** 2, im_check_size_list, color='b', label='comp. checkerboard')
    ax0.plot(n_list ** 2, im_rand_size_list, color='r', label='comp. random')

    ax0.set_yscale('log')
    ax0.set_xscale('log')
    ax0.set_ylabel('bytes')
    ax0.set_xlabel('raw bytes')
    ax0.legend(loc='best')
    fig0.tight_layout()
    fig0.savefig('comparison.pdf')