SHOW:
|
|
- or go back to the newest paste.
| 1 | class BitWriter(): | |
| 2 | def __init__(self, consumer): | |
| 3 | self.consumer = consumer | |
| 4 | self.bit_buffer = 0 | |
| 5 | self.bit_buffer_len = 0 | |
| 6 | ||
| 7 | def write_bits(self, value, bits): | |
| 8 | # TODO catch invalid parameter | |
| 9 | self.bit_buffer |= value << self.bit_buffer_len | |
| 10 | self.bit_buffer_len += bits | |
| 11 | while self.bit_buffer_len >= 8: | |
| 12 | self.consumer.write(chr(self.bit_buffer & 0xff)) | |
| 13 | self.bit_buffer >>= 8 | |
| 14 | self.bit_buffer_len -= 8 | |
| 15 | ||
| 16 | def finish(self): | |
| 17 | if self.bit_buffer_len > 0: | |
| 18 | self.consumer.write(chr(self.bit_buffer & 0xff)) | |
| 19 | ||
| 20 | def lzw_encode_default(data, code_bits): | |
| 21 | """ | |
| 22 | @data: bytes (in py3) or something iterable with ints <= 255 | |
| 23 | @code_bits: the minimum_lzw_size (size of code in bits at start = code_bits +1) | |
| 24 | data should contain the frames pixel index data. | |
| 25 | The result is a string (or byte in py3). | |
| 26 | TODO: replace StringIO with BytesIo | |
| 27 | """ | |
| 28 | MAX_DICT_SIZE = 4096 | |
| 29 | buf = StringIO() | |
| 30 | writer = BitWriter(buf) | |
| 31 | ||
| 32 | if code_bits < 2 or code_bits > 8: | |
| 33 | raise Exception("Invalid code size");
| |
| 34 | ||
| 35 | code_size = code_bits +1 | |
| 36 | clear_code = 1 << code_bits | |
| 37 | eoi_code = clear_code +1 | |
| 38 | next_code = eoi_code +1 | |
| 39 | # We could use lists or arrays or something more memory friendly here i guess ? | |
| 40 | # Also we don't really need to recreate all the lists. Zeroing all fields should be enough ? | |
| 41 | - | root = [[None]*265]*(1 << code_size) |
| 41 | + | root = [[None]*256]*(1 << code_size) |
| 42 | writer.write_bits(clear_code, code_size) | |
| 43 | ||
| 44 | i = 1 | |
| 45 | last_index = data[0] | |
| 46 | while i < len(data) -1: | |
| 47 | # Loop till we find an unknown sequence | |
| 48 | while i < len(data) -1 and not root[last_index][data[i]] is None: | |
| 49 | - | last_index = data[i] |
| 49 | + | last_index = root[last_index][data[i]] |
| 50 | i += 1 | |
| 51 | ||
| 52 | writer.write_bits(last_index, code_size) | |
| 53 | root[last_index][data[i]] = next_code | |
| 54 | ||
| 55 | - | #// Add new dictionary entry |
| 55 | + | # Add new dictionary entry |
| 56 | if next_code < MAX_DICT_SIZE: | |
| 57 | - | #// But we must update the size and code bits for the decoder's sake |
| 57 | + | |
| 58 | code_size += 1 | |
| 59 | # TODO: clean this up # Should be way nicer to calc.. to tired TODO: | |
| 60 | root.extend( [ [None]*255 ] * ((1 << code_size) - (1 << (code_size-1))) ) | |
| 61 | - | root.extend( [ [ [None]*255 ] * ((1 << code_size) - (1 << (code_size-1))) ] ) |
| 61 | + | |
| 62 | next_code += 1 | |
| 63 | else: | |
| 64 | writer.write_bits(clear_code, code_size) | |
| 65 | code_size = code_bits +1 | |
| 66 | next_code = eoi_code +1 | |
| 67 | root = [[None]*265]*(1 << code_size) | |
| 68 | pass | |
| 69 | ||
| 70 | writer.write_bits(eoi_code, code_size) | |
| 71 | writer.finish() | |
| 72 | return buf.getvalue() |
