Add the pic.py (de)compressor from pokemon-reverse-engineering-tools

$ tools/pic.py decompress x.pic  # creates x.2bpp
$ tools/gfx.py png x.2bpp        # creates x.png
$ tools/gfx.py 2bpp x.png        # creates x.2bpp
$ tools/pic.py compress x.2bpp   # creates x.pic

1 files changed, 491 insertions, 0 deletions

diff --git a/tools/pic.py b/tools/pic.py
new file mode 100644
index 00000000..afc34c90
--- /dev/null
+++ b/tools/pic.py
@@ -0,0 +1,491 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+
+"""
+A library for use with compressed monster and trainer pics in pokered.
+"""
+from __future__ import absolute_import
+from __future__ import division
+
+import os
+import sys
+import argparse
+from math import sqrt
+
+from pokemontools import gfx
+
+
+def bitflip(x, n):
+    r = 0
+    while n:
+        r = (r << 1) | (x & 1)
+        x >>= 1
+        n -= 1
+    return r
+
+
+class Decompressor:
+    """
+    pokered pic decompression.
+
+    Ported to python 2.7 from the python 3 code at https://github.com/magical/pokemon-sprites-rby.
+    """
+
+    table1 = [(2 << i) - 1 for i in range(16)]
+    table2 = [
+        [0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5, 0xf, 0xe, 0xc, 0xd, 0x8, 0x9, 0xb, 0xa],
+        [0xf, 0xe, 0xc, 0xd, 0x8, 0x9, 0xb, 0xa, 0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5], # prev ^ 0xf
+        [0x0, 0x8, 0xc, 0x4, 0xe, 0x6, 0x2, 0xa, 0xf, 0x7, 0x3, 0xb, 0x1, 0x9, 0xd, 0x5],
+        [0xf, 0x7, 0x3, 0xb, 0x1, 0x9, 0xd, 0x5, 0x0, 0x8, 0xc, 0x4, 0xe, 0x6, 0x2, 0xa], # prev ^ 0xf
+    ]
+    table3 = [bitflip(i, 4) for i in range(16)]
+
+    tilesize = 8
+
+
+    def __init__(self, f, mirror=False, planar=True):
+        self.bs = fbitstream(f)
+        self.mirror = mirror
+        self.planar = planar
+        self.data = None
+
+    def decompress(self):
+        rams = [[], []]
+
+        self.sizex  = self._readint(4) * self.tilesize
+        self.sizey  = self._readint(4)
+
+        self.size = self.sizex * self.sizey
+
+        self.ramorder = self._readbit()
+
+        r1 = self.ramorder
+        r2 = self.ramorder ^ 1
+
+        self._fillram(rams[r1])
+        mode = self._readbit()
+        if mode:
+            mode += self._readbit()
+        self._fillram(rams[r2])
+
+        rams[0] = bytearray(bitgroups_to_bytes(rams[0]))
+        rams[1] = bytearray(bitgroups_to_bytes(rams[1]))
+
+        if mode == 0:
+            self._decode(rams[0])
+            self._decode(rams[1])
+        elif mode == 1:
+            self._decode(rams[r1])
+            self._xor(rams[r1], rams[r2])
+        elif mode == 2:
+            self._decode(rams[r2], mirror=False)
+            self._decode(rams[r1])
+            self._xor(rams[r1], rams[r2])
+        else:
+            raise Exception("Invalid deinterlace mode!")
+
+        data = []
+        if self.planar:
+            for a, b in zip(rams[0], rams[1]):
+                data += [a, b]
+            self.data = bytearray(data)
+        else:
+            for a, b in zip(bitstream(rams[0]), bitstream(rams[1])):
+                data.append(a | (b << 1))
+            self.data = bitgroups_to_bytes(data)
+
+    def _fillram(self, ram):
+        mode = ['rle', 'data'][self._readbit()]
+        size = self.size * 4
+        while len(ram) < size:
+            if mode == 'rle':
+                self._read_rle_chunk(ram)
+                mode = 'data'
+            elif mode == 'data':
+                self._read_data_chunk(ram, size)
+                mode = 'rle'
+        if len(ram) > size:
+            #ram = ram[:size]
+            raise ValueError(size, len(ram))
+
+        ram[:] = self._deinterlace_bitgroups(ram)
+
+    def _read_rle_chunk(self, ram):
+
+        i = 0
+        while self._readbit():
+            i += 1
+
+        n = self.table1[i]
+        a = self._readint(i + 1)
+        n += a
+
+        for i in range(n):
+            ram.append(0)
+
+    def _read_data_chunk(self, ram, size):
+        while 1:
+            bitgroup = self._readint(2)
+            if bitgroup == 0:
+                break
+            ram.append(bitgroup)
+
+            if size <= len(ram):
+                break
+
+    def _decode(self, ram, mirror=None):
+        if mirror is None:
+            mirror = self.mirror
+
+        for x in range(self.sizex):
+            bit = 0
+            for y in range(self.sizey):
+                i = y * self.sizex + x
+                a = (ram[i] >> 4) & 0xf
+                b = ram[i] & 0xf
+
+                a = self.table2[bit][a]
+                bit = a & 1
+                if mirror:
+                    a = self.table3[a]
+
+                b = self.table2[bit][b]
+                bit = b & 1
+                if mirror:
+                    b = self.table3[b]
+
+                ram[i] = (a << 4) | b
+
+    def _xor(self, ram1, ram2, mirror=None):
+        if mirror is None:
+            mirror = self.mirror
+
+        for i in range(len(ram2)):
+            if mirror:
+                a = (ram2[i] >> 4) & 0xf
+                b = ram2[i] & 0xf
+                a = self.table3[a]
+                b = self.table3[b]
+                ram2[i] = (a << 4) | b
+
+            ram2[i] ^= ram1[i]
+
+    def _deinterlace_bitgroups(self, bits):
+        l = []
+        for y in range(self.sizey):
+            for x in range(self.sizex):
+                i = 4 * y * self.sizex + x
+                for j in range(4):
+                    l.append(bits[i])
+                    i += self.sizex
+        return l
+
+
+    def _readbit(self):
+        return next(self.bs)
+
+    def _readint(self, count):
+        return readint(self.bs, count)
+
+
+def fbitstream(f):
+    while 1:
+        char = f.read(1)
+        if not char:
+            break
+        byte = ord(char)
+
+        for i in range(7, -1, -1):
+            yield (byte >> i) & 1
+
+def bitstream(b):
+    for byte in b:
+        for i in range(7, -1, -1):
+            yield (byte >> i) & 1
+
+def readint(bs, count):
+    n = 0
+    while count:
+        n <<= 1
+        n |= next(bs)
+        count -= 1
+    return n
+
+def bitgroups_to_bytes(bits):
+    l = []
+    for i in range(0, len(bits) - 3, 4):
+        n = ((bits[i + 0] << 6)
+           | (bits[i + 1] << 4)
+           | (bits[i + 2] << 2)
+           | (bits[i + 3] << 0))
+        l.append(n)
+    return bytearray(l)
+
+
+def bytes_to_bits(bytelist):
+    return list(bitstream(bytelist))
+
+
+class Compressor:
+    """
+    pokered pic compression.
+
+    Adapted from stag019's C compressor.
+    """
+
+    table1 = [(2 << i) - 1 for i in range(16)]
+    table2 = [
+        [0x0, 0x1, 0x3, 0x2, 0x6, 0x7, 0x5, 0x4, 0xc, 0xd, 0xf, 0xe, 0xa, 0xb, 0x9, 0x8],
+        [0x8, 0x9, 0xb, 0xa, 0xe, 0xf, 0xd, 0xc, 0x4, 0x5, 0x7, 0x6, 0x2, 0x3, 0x1, 0x0], # reverse
+    ]
+    table3 = [bitflip(i, 4) for i in range(16)]
+
+    def __init__(self, image, width=None, height=None):
+        self.image = bytearray(image)
+        self.size = len(self.image)
+
+        planar_tile = 8 * 8 // 4
+        tile_size = self.size // planar_tile
+        if   height    and not width:  width  = tile_size // height
+        elif width     and not height: height = tile_size // width
+        elif not width and not height: width = height = int(sqrt(tile_size))
+        self.width, self.height = width, height
+
+    def compress(self):
+        """
+        Compress the image five times (twice for each mode, except 0)
+        and use the smallest one (in bits).
+        """
+        rams = [[],[]]
+        datas = []
+
+        for mode in range(3):
+
+            # Order is redundant for mode 0.
+
+            # While this seems like an optimization,
+            # it's actually required for 1:1 compression
+            # to the original compressed pics.
+
+            # This appears to be the algorithm
+            # that Game Freak's compressor used.
+
+            # Using order 0 instead of 1 breaks this feature.
+
+            for order in range(2):
+                if mode == 0 and order == 0:
+                    continue
+                for i in range(2):
+                    rams[i] = self.image[i::2]
+                self._interpret_compress(rams, mode, order)
+                datas += [(self.data[:], int(self.which_bit))]
+
+        # Pick the smallest pic, measured in bits.
+        datas = sorted(datas, key=lambda data_bit: (len(data_bit[0]), -data_bit[1]))
+        self.data, self.which_bit = datas[0]
+
+    def _interpret_compress(self, rams, mode, order):
+        self.data = []
+        self.which_bit = 0
+
+        r1 = order
+        r2 = order ^ 1
+
+        if mode == 0:
+            self._encode(rams[1])
+            self._encode(rams[0])
+        elif mode == 1:
+            self._xor(rams[r1], rams[r2])
+            self._encode(rams[r1])
+        elif mode == 2:
+            self._xor(rams[r1], rams[r2])
+            self._encode(rams[r1])
+            self._encode(rams[r2], mirror=False)
+        else:
+            raise Exception('invalid interlace mode!')
+
+        self._writeint(self.height, 4)
+        self._writeint(self.width,  4)
+
+        self._writebit(order)
+
+        self._fillram(rams[r1])
+        if mode == 0:
+            self._writebit(0)
+        else:
+            self._writebit(1)
+            self._writebit(mode - 1)
+        self._fillram(rams[r2])
+
+    def _fillram(self, ram):
+        rle = 0
+        nums = 0
+        bitgroups = []
+
+        for x in range(self.width):
+            for bit in range(0, 8, 2):
+                byte = x * self.height * 8
+                for y in range(self.height * 8):
+                    bitgroup = (ram[byte] >> (6 - bit)) & 3
+                    if bitgroup == 0:
+                        if rle == 0:
+                            self._writebit(0)
+                        elif rle == 1:
+                            nums += 1
+                        else:
+                            self._data_packet(bitgroups)
+                            self._writebit(0)
+                            self._writebit(0)
+                        rle = 1
+                        bitgroups = []
+                    else:
+                        if rle == 0:
+                            self._writebit(1)
+                        elif rle == 1:
+                            self._rle(nums)
+                        rle = -1
+                        bitgroups += [bitgroup]
+                        nums = 0
+                    byte += 1
+
+        if rle == 1:
+            self._rle(nums)
+        else:
+            self._data_packet(bitgroups)
+
+    def _data_packet(self, bitgroups):
+        for bitgroup in bitgroups:
+            self._writebit((bitgroup >> 1) & 1)
+            self._writebit((bitgroup >> 0) & 1)
+
+    def _rle(self, nums):
+        nums += 1
+
+        # Get the previous power of 2.
+        # Deriving the bitcount from that seems to be
+        # faster on average than using the lookup table.
+        v = nums
+        v += 1
+        v |= v >> 1
+        v |= v >> 2
+        v |= v >> 4
+        v |= v >> 8
+        v |= v >> 16
+        v -= v >> 1
+        v -= 1
+        number = nums - v
+
+        bitcount = -1
+        while v:
+            v >>= 1
+            bitcount += 1
+
+        for j in range(bitcount):
+            self._writebit(1)
+        self._writebit(0)
+        for j in range(bitcount, -1, -1):
+            self._writebit((number >> j) & 1)
+
+    def _encode(self, ram, mirror=None):
+        a = b = 0
+        for i in range(len(ram)):
+            j = i // self.height
+            j += i % self.height * self.width * 8
+            if i % self.height == 0:
+                b = 0
+
+            a = (ram[j] >> 4) & 0xf
+            table = b & 1
+            code_1 = self.table2[table][a]
+
+            b = ram[j] & 0xf
+            table = a & 1
+            code_2 = self.table2[table][b]
+
+            ram[j] = (code_1 << 4) | code_2
+
+    def _xor(self, ram1, ram2):
+        for i in range(len(ram2)):
+            ram2[i] ^= ram1[i]
+
+    def _writebit(self, bit):
+        self.which_bit -= 1
+        if self.which_bit == -1:
+            self.which_bit = 7
+            self.data += [0]
+        if bit: self.data[-1] |= bit << self.which_bit
+
+    def _writeint(self, num, size=None):
+        bits = []
+        if size:
+            for i in range(size):
+                bits += [num & 1]
+                num >>= 1
+        else:
+            while num > 0:
+                bits += [num & 1]
+                num >>= 1
+        for bit in reversed(bits):
+            self._writebit(bit)
+
+
+def decompress(f, offset=None, mirror=False):
+    """
+    Decompress a pic given a file object. Return a planar 2bpp image.
+
+    Optional: offset (for roms).
+    """
+    if offset is not None:
+        f.seek(offset)
+    dcmp = Decompressor(f, mirror=mirror)
+    dcmp.decompress()
+    return dcmp.data
+
+
+def compress(f):
+    """
+    Compress a planar 2bpp into a pic.
+    """
+    comp = Compressor(f)
+    comp.compress()
+    return comp.data
+
+
+def decompress_file(filename):
+    """
+    Decompress a pic given a filename.
+    Export the resulting planar 2bpp image to
+    """
+    pic = open(filename, 'rb')
+    image = decompress(pic)
+    image = gfx.transpose_tiles(image)
+    image = bytearray(image)
+    output_filename = os.path.splitext(filename)[0] + '.2bpp'
+    with open(output_filename, 'wb') as out:
+        out.write(image)
+
+def compress_file(filename):
+    image = open(filename, 'rb').read()
+    image = gfx.transpose_tiles(image)
+    pic = compress(image)
+    pic = bytearray(pic)
+    output_filename = os.path.splitext(filename)[0] + '.pic'
+    with open(output_filename, 'wb') as out:
+        out.write(pic)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('mode')
+    ap.add_argument('filenames', nargs='*')
+    args = ap.parse_args()
+
+    for filename in args.filenames:
+        if args.mode == 'decompress':
+            decompress_file(filename)
+        elif args.mode == 'compress':
+            compress_file(filename)
+
+if __name__ == '__main__':
+    main()
+