Advertisement
TerusTheBird

pic.py standalone

Apr 25th, 2020
598
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/bin/env python3
  2. # coding: utf-8
  3.  
  4. """
  5. A library for use with compressed monster and trainer pics in pokered.
  6. """
  7. from __future__ import absolute_import
  8. from __future__ import division
  9.  
  10. import os
  11. import sys
  12. import argparse
  13. from math import sqrt
  14.  
  15. def connect(tiles):
  16.     """
  17.     Combine 8x8 tiles into a 2bpp image.
  18.     """
  19.     return [byte for tile in tiles for byte in tile]
  20.  
  21. def transpose(tiles, width=None):
  22.     """
  23.     Transpose a tile arrangement along line y=-x.
  24.  
  25.       00 01 02 03 04 05     00 06 0c 12 18 1e
  26.       06 07 08 09 0a 0b     01 07 0d 13 19 1f
  27.       0c 0d 0e 0f 10 11 <-> 02 08 0e 14 1a 20
  28.       12 13 14 15 16 17     03 09 0f 15 1b 21
  29.       18 19 1a 1b 1c 1d     04 0a 10 16 1c 22
  30.       1e 1f 20 21 22 23     05 0b 11 17 1d 23
  31.  
  32.       00 01 02 03     00 04 08
  33.       04 05 06 07 <-> 01 05 09
  34.       08 09 0a 0b     02 06 0a
  35.                       03 07 0b
  36.     """
  37.     if width == None:
  38.         width = int(sqrt(len(tiles))) # assume square image
  39.     tiles = sorted(enumerate(tiles), key= lambda i_tile: i_tile[0] % width)
  40.     return [tile for i, tile in tiles]
  41.  
  42. def transpose_tiles(image, width=None):
  43.     return connect(transpose(get_tiles(image), width))
  44.  
  45. def bitflip(x, n):
  46.     r = 0
  47.     while n:
  48.         r = (r << 1) | (x & 1)
  49.         x >>= 1
  50.         n -= 1
  51.     return r
  52.  
  53.  
  54. class Decompressor:
  55.     """
  56.     pokered pic decompression.
  57.  
  58.     Ported to python 2.7 from the python 3 code at https://github.com/magical/pokemon-sprites-rby.
  59.     """
  60.  
  61.     table1 = [(2 << i) - 1 for i in range(16)]
  62.     table2 = [
  63.         [0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5, 0xf, 0xe, 0xc, 0xd, 0x8, 0x9, 0xb, 0xa],
  64.         [0xf, 0xe, 0xc, 0xd, 0x8, 0x9, 0xb, 0xa, 0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5], # prev ^ 0xf
  65.         [0x0, 0x8, 0xc, 0x4, 0xe, 0x6, 0x2, 0xa, 0xf, 0x7, 0x3, 0xb, 0x1, 0x9, 0xd, 0x5],
  66.         [0xf, 0x7, 0x3, 0xb, 0x1, 0x9, 0xd, 0x5, 0x0, 0x8, 0xc, 0x4, 0xe, 0x6, 0x2, 0xa], # prev ^ 0xf
  67.     ]
  68.     table3 = [bitflip(i, 4) for i in range(16)]
  69.  
  70.     tilesize = 8
  71.  
  72.  
  73.     def __init__(self, f, mirror=False, planar=True):
  74.         self.bs = fbitstream(f)
  75.         self.mirror = mirror
  76.         self.planar = planar
  77.         self.data = None
  78.  
  79.     def decompress(self):
  80.         rams = [[], []]
  81.  
  82.         self.sizex  = self._readint(4) * self.tilesize
  83.         self.sizey  = self._readint(4)
  84.  
  85.         self.size = self.sizex * self.sizey
  86.  
  87.         self.ramorder = self._readbit()
  88.  
  89.         r1 = self.ramorder
  90.         r2 = self.ramorder ^ 1
  91.  
  92.         self._fillram(rams[r1])
  93.         mode = self._readbit()
  94.         if mode:
  95.             mode += self._readbit()
  96.         self._fillram(rams[r2])
  97.  
  98.         rams[0] = bytearray(bitgroups_to_bytes(rams[0]))
  99.         rams[1] = bytearray(bitgroups_to_bytes(rams[1]))
  100.  
  101.         if mode == 0:
  102.             self._decode(rams[0])
  103.             self._decode(rams[1])
  104.         elif mode == 1:
  105.             self._decode(rams[r1])
  106.             self._xor(rams[r1], rams[r2])
  107.         elif mode == 2:
  108.             self._decode(rams[r2], mirror=False)
  109.             self._decode(rams[r1])
  110.             self._xor(rams[r1], rams[r2])
  111.         else:
  112.             raise Exception("Invalid deinterlace mode!")
  113.  
  114.         data = []
  115.         if self.planar:
  116.             for a, b in zip(rams[0], rams[1]):
  117.                 data += [a, b]
  118.             self.data = bytearray(data)
  119.         else:
  120.             for a, b in zip(bitstream(rams[0]), bitstream(rams[1])):
  121.                 data.append(a | (b << 1))
  122.             self.data = bitgroups_to_bytes(data)
  123.  
  124.     def _fillram(self, ram):
  125.         mode = ['rle', 'data'][self._readbit()]
  126.         size = self.size * 4
  127.         while len(ram) < size:
  128.             if mode == 'rle':
  129.                 self._read_rle_chunk(ram)
  130.                 mode = 'data'
  131.             elif mode == 'data':
  132.                 self._read_data_chunk(ram, size)
  133.                 mode = 'rle'
  134.         if len(ram) > size:
  135.             #ram = ram[:size]
  136.             raise ValueError(size, len(ram))
  137.  
  138.         ram[:] = self._deinterlace_bitgroups(ram)
  139.  
  140.     def _read_rle_chunk(self, ram):
  141.  
  142.         i = 0
  143.         while self._readbit():
  144.             i += 1
  145.  
  146.         n = self.table1[i]
  147.         a = self._readint(i + 1)
  148.         n += a
  149.  
  150.         for i in range(n):
  151.             ram.append(0)
  152.  
  153.     def _read_data_chunk(self, ram, size):
  154.         while 1:
  155.             bitgroup = self._readint(2)
  156.             if bitgroup == 0:
  157.                 break
  158.             ram.append(bitgroup)
  159.  
  160.             if size <= len(ram):
  161.                 break
  162.  
  163.     def _decode(self, ram, mirror=None):
  164.         if mirror is None:
  165.             mirror = self.mirror
  166.  
  167.         for x in range(self.sizex):
  168.             bit = 0
  169.             for y in range(self.sizey):
  170.                 i = y * self.sizex + x
  171.                 a = (ram[i] >> 4) & 0xf
  172.                 b = ram[i] & 0xf
  173.  
  174.                 a = self.table2[bit][a]
  175.                 bit = a & 1
  176.                 if mirror:
  177.                     a = self.table3[a]
  178.  
  179.                 b = self.table2[bit][b]
  180.                 bit = b & 1
  181.                 if mirror:
  182.                     b = self.table3[b]
  183.  
  184.                 ram[i] = (a << 4) | b
  185.  
  186.     def _xor(self, ram1, ram2, mirror=None):
  187.         if mirror is None:
  188.             mirror = self.mirror
  189.  
  190.         for i in range(len(ram2)):
  191.             if mirror:
  192.                 a = (ram2[i] >> 4) & 0xf
  193.                 b = ram2[i] & 0xf
  194.                 a = self.table3[a]
  195.                 b = self.table3[b]
  196.                 ram2[i] = (a << 4) | b
  197.  
  198.             ram2[i] ^= ram1[i]
  199.  
  200.     def _deinterlace_bitgroups(self, bits):
  201.         l = []
  202.         for y in range(self.sizey):
  203.             for x in range(self.sizex):
  204.                 i = 4 * y * self.sizex + x
  205.                 for j in range(4):
  206.                     l.append(bits[i])
  207.                     i += self.sizex
  208.         return l
  209.  
  210.  
  211.     def _readbit(self):
  212.         return next(self.bs)
  213.  
  214.     def _readint(self, count):
  215.         return readint(self.bs, count)
  216.  
  217.  
  218. def fbitstream(f):
  219.     while 1:
  220.         char = f.read(1)
  221.         if not char:
  222.             break
  223.         byte = ord(char)
  224.  
  225.         for i in range(7, -1, -1):
  226.             yield (byte >> i) & 1
  227.  
  228. def bitstream(b):
  229.     for byte in b:
  230.         for i in range(7, -1, -1):
  231.             yield (byte >> i) & 1
  232.  
  233. def readint(bs, count):
  234.     n = 0
  235.     while count:
  236.         n <<= 1
  237.         n |= next(bs)
  238.         count -= 1
  239.     return n
  240.  
  241. def bitgroups_to_bytes(bits):
  242.     l = []
  243.     for i in range(0, len(bits) - 3, 4):
  244.         n = ((bits[i + 0] << 6)
  245.            | (bits[i + 1] << 4)
  246.            | (bits[i + 2] << 2)
  247.            | (bits[i + 3] << 0))
  248.         l.append(n)
  249.     return bytearray(l)
  250.  
  251.  
  252. def bytes_to_bits(bytelist):
  253.     return list(bitstream(bytelist))
  254.  
  255.  
  256. class Compressor:
  257.     """
  258.     pokered pic compression.
  259.  
  260.     Adapted from stag019's C compressor.
  261.     """
  262.  
  263.     table1 = [(2 << i) - 1 for i in range(16)]
  264.     table2 = [
  265.         [0x0, 0x1, 0x3, 0x2, 0x6, 0x7, 0x5, 0x4, 0xc, 0xd, 0xf, 0xe, 0xa, 0xb, 0x9, 0x8],
  266.         [0x8, 0x9, 0xb, 0xa, 0xe, 0xf, 0xd, 0xc, 0x4, 0x5, 0x7, 0x6, 0x2, 0x3, 0x1, 0x0], # reverse
  267.     ]
  268.     table3 = [bitflip(i, 4) for i in range(16)]
  269.  
  270.     def __init__(self, image, width=None, height=None):
  271.         self.image = bytearray(image)
  272.         self.size = len(self.image)
  273.  
  274.         planar_tile = 8 * 8 // 4
  275.         tile_size = self.size // planar_tile
  276.         if   height    and not width:  width  = tile_size // height
  277.         elif width     and not height: height = tile_size // width
  278.         elif not width and not height: width = height = int(sqrt(tile_size))
  279.         self.width, self.height = width, height
  280.  
  281.     def compress(self):
  282.         """
  283.         Compress the image five times (twice for each mode, except 0)
  284.         and use the smallest one (in bits).
  285.         """
  286.         rams = [[],[]]
  287.         datas = []
  288.  
  289.         for mode in range(3):
  290.  
  291.             # Order is redundant for mode 0.
  292.  
  293.             # While this seems like an optimization,
  294.             # it's actually required for 1:1 compression
  295.             # to the original compressed pics.
  296.  
  297.             # This appears to be the algorithm
  298.             # that Game Freak's compressor used.
  299.  
  300.             # Using order 0 instead of 1 breaks this feature.
  301.  
  302.             for order in range(2):
  303.                 if mode == 0 and order == 0:
  304.                     continue
  305.                 for i in range(2):
  306.                     rams[i] = self.image[i::2]
  307.                 self._interpret_compress(rams, mode, order)
  308.                 datas += [(self.data[:], int(self.which_bit))]
  309.  
  310.         # Pick the smallest pic, measured in bits.
  311.         datas = sorted(datas, key=lambda data_bit: (len(data_bit[0]), -data_bit[1]))
  312.         self.data, self.which_bit = datas[0]
  313.  
  314.     def _interpret_compress(self, rams, mode, order):
  315.         self.data = []
  316.         self.which_bit = 0
  317.  
  318.         r1 = order
  319.         r2 = order ^ 1
  320.  
  321.         if mode == 0:
  322.             self._encode(rams[1])
  323.             self._encode(rams[0])
  324.         elif mode == 1:
  325.             self._xor(rams[r1], rams[r2])
  326.             self._encode(rams[r1])
  327.         elif mode == 2:
  328.             self._xor(rams[r1], rams[r2])
  329.             self._encode(rams[r1])
  330.             self._encode(rams[r2], mirror=False)
  331.         else:
  332.             raise Exception('invalid interlace mode!')
  333.  
  334.         self._writeint(self.height, 4)
  335.         self._writeint(self.width,  4)
  336.  
  337.         self._writebit(order)
  338.  
  339.         self._fillram(rams[r1])
  340.         if mode == 0:
  341.             self._writebit(0)
  342.         else:
  343.             self._writebit(1)
  344.             self._writebit(mode - 1)
  345.         self._fillram(rams[r2])
  346.  
  347.     def _fillram(self, ram):
  348.         rle = 0
  349.         nums = 0
  350.         bitgroups = []
  351.  
  352.         for x in range(self.width):
  353.             for bit in range(0, 8, 2):
  354.                 byte = x * self.height * 8
  355.                 for y in range(self.height * 8):
  356.                     bitgroup = (ram[byte] >> (6 - bit)) & 3
  357.                     if bitgroup == 0:
  358.                         if rle == 0:
  359.                             self._writebit(0)
  360.                         elif rle == 1:
  361.                             nums += 1
  362.                         else:
  363.                             self._data_packet(bitgroups)
  364.                             self._writebit(0)
  365.                             self._writebit(0)
  366.                         rle = 1
  367.                         bitgroups = []
  368.                     else:
  369.                         if rle == 0:
  370.                             self._writebit(1)
  371.                         elif rle == 1:
  372.                             self._rle(nums)
  373.                         rle = -1
  374.                         bitgroups += [bitgroup]
  375.                         nums = 0
  376.                     byte += 1
  377.  
  378.         if rle == 1:
  379.             self._rle(nums)
  380.         else:
  381.             self._data_packet(bitgroups)
  382.  
  383.     def _data_packet(self, bitgroups):
  384.         for bitgroup in bitgroups:
  385.             self._writebit((bitgroup >> 1) & 1)
  386.             self._writebit((bitgroup >> 0) & 1)
  387.  
  388.     def _rle(self, nums):
  389.         nums += 1
  390.  
  391.         # Get the previous power of 2.
  392.         # Deriving the bitcount from that seems to be
  393.         # faster on average than using the lookup table.
  394.         v = nums
  395.         v += 1
  396.         v |= v >> 1
  397.         v |= v >> 2
  398.         v |= v >> 4
  399.         v |= v >> 8
  400.         v |= v >> 16
  401.         v -= v >> 1
  402.         v -= 1
  403.         number = nums - v
  404.  
  405.         bitcount = -1
  406.         while v:
  407.             v >>= 1
  408.             bitcount += 1
  409.  
  410.         for j in range(bitcount):
  411.             self._writebit(1)
  412.         self._writebit(0)
  413.         for j in range(bitcount, -1, -1):
  414.             self._writebit((number >> j) & 1)
  415.  
  416.     def _encode(self, ram, mirror=None):
  417.         a = b = 0
  418.         for i in range(len(ram)):
  419.             j = i // self.height
  420.             j += i % self.height * self.width * 8
  421.             if i % self.height == 0:
  422.                 b = 0
  423.  
  424.             a = (ram[j] >> 4) & 0xf
  425.             table = b & 1
  426.             code_1 = self.table2[table][a]
  427.  
  428.             b = ram[j] & 0xf
  429.             table = a & 1
  430.             code_2 = self.table2[table][b]
  431.  
  432.             ram[j] = (code_1 << 4) | code_2
  433.  
  434.     def _xor(self, ram1, ram2):
  435.         for i in range(len(ram2)):
  436.             ram2[i] ^= ram1[i]
  437.  
  438.     def _writebit(self, bit):
  439.         self.which_bit -= 1
  440.         if self.which_bit == -1:
  441.             self.which_bit = 7
  442.             self.data += [0]
  443.         if bit: self.data[-1] |= bit << self.which_bit
  444.  
  445.     def _writeint(self, num, size=None):
  446.         bits = []
  447.         if size:
  448.             for i in range(size):
  449.                 bits += [num & 1]
  450.                 num >>= 1
  451.         else:
  452.             while num > 0:
  453.                 bits += [num & 1]
  454.                 num >>= 1
  455.         for bit in reversed(bits):
  456.             self._writebit(bit)
  457.  
  458.  
  459. def decompress(f, offset=None, mirror=False):
  460.     """
  461.     Decompress a pic given a file object. Return a planar 2bpp image.
  462.  
  463.     Optional: offset (for roms).
  464.     """
  465.     if offset is not None:
  466.         f.seek(offset)
  467.     dcmp = Decompressor(f, mirror=mirror)
  468.     dcmp.decompress()
  469.     return dcmp.data
  470.  
  471.  
  472. def compress(f):
  473.     """
  474.     Compress a planar 2bpp into a pic.
  475.     """
  476.     comp = Compressor(f)
  477.     comp.compress()
  478.     return comp.data
  479.  
  480.  
  481. def decompress_file(filename):
  482.     """
  483.     Decompress a pic given a filename.
  484.     Export the resulting planar 2bpp image to
  485.     """
  486.     pic = open(filename, 'rb')
  487.     image = decompress(pic)
  488.     image = transpose_tiles(image)
  489.     image = bytearray(image)
  490.     output_filename = os.path.splitext(filename)[0] + '.2bpp'
  491.     with open(output_filename, 'wb') as out:
  492.         out.write(image)
  493.  
  494. def compress_file(filename):
  495.     image = open(filename, 'rb').read()
  496.     image = transpose_tiles(image)
  497.     pic = compress(image)
  498.     pic = bytearray(pic)
  499.     output_filename = os.path.splitext(filename)[0] + '.pic'
  500.     with open(output_filename, 'wb') as out:
  501.         out.write(pic)
  502.  
  503. if __name__ == '__main__':
  504.     ap = argparse.ArgumentParser()
  505.     ap.add_argument('mode')
  506.     ap.add_argument('filenames', nargs='*')
  507.     args = ap.parse_args()
  508.  
  509.     for filename in args.filenames:
  510.         if args.mode == 'decompress':
  511.             decompress_file(filename)
  512.         elif args.mode == 'compress':
  513.             compress_file(filename)
  514. #
Advertisement
RAW Paste Data Copied
Advertisement