# 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
84.
85.         self.size = self.sizex * self.sizey
86.
88.
89.         r1 = self.ramorder
90.         r2 = self.ramorder ^ 1
91.
92.         self._fillram(rams[r1])
94.         if mode:
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):
126.         size = self.size * 4
127.         while len(ram) < size:
128.             if mode == 'rle':
130.                 mode = 'data'
131.             elif mode == 'data':
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.
141.
142.         i = 0
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.
154.         while 1:
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.
212.         return next(self.bs)
213.
216.
217.
218. def fbitstream(f):
219.     while 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.
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):
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()