This repository has been archived by the owner on May 28, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
decode.py
76 lines (52 loc) · 1.73 KB
/
decode.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import hashlib
import sys
import numpy as np
import random
from PIL import Image
MODES = {0: '1', 1: 'L', 2: 'RGB'}
CHANNELS = {'1': 1, 'L': 1, 'RGB': 3}
def decode(f, mode, filesize):
decoded = []
while f.tell() < filesize:
if mode in ('1', 'L'): # 2 byte count followed by 1 byte values
pixel = f.read(3)
else: # 2 byte count followed by 3 byte values
pixel = f.read(5)
count = pixel[:2]
value = pixel[2:]
decoded.append((count, value))
return decoded
in_filename = sys.argv[1]
with open(in_filename, 'rb') as f:
f.seek(0, 2)
filesize = f.tell()
f.seek(0)
magic_number = f.read(4)
assert magic_number == b'BOGO'
width = int.from_bytes(f.read(1), 'big')
height = int.from_bytes(f.read(1), 'big')
mode = MODES[int.from_bytes(f.read(1), 'big')]
checksum = f.read(16)
counts = decode(f, mode, filesize)
pixel_values = []
for count, value in counts:
pixel_values.extend([value] * int.from_bytes(count, 'big'))
while True:
random.shuffle(pixel_values)
md5 = hashlib.md5()
md5.update(b''.join(pixel_values))
if checksum == md5.digest():
break
# now we have the correct pixel value ordering, need to convert bytes to actual brightnesses
pixels = []
if mode in ('1', 'L'):
for pixel in pixel_values:
pixels.append(int.from_bytes(pixel, 'big'))
else:
for pixel in pixel_values:
rgb = (pixel[0], pixel[1], pixel[2])
pixels.append(rgb)
ndarray = np.asarray(pixels).astype(np.uint8)
ndarray = ndarray.reshape((height, width, CHANNELS[mode]))
img = Image.fromarray(ndarray, mode=mode)
img.save(sys.argv[2])