-
Notifications
You must be signed in to change notification settings - Fork 0
/
chunks.py
143 lines (115 loc) · 4.19 KB
/
chunks.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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
from collections import namedtuple
import zlib
import struct
import math
from errors import InvalidPNG
from crc import CRC
def split_into_chunks(file_bytes, chunks=[]):
''' chunks should look like this:
{'length': int, 'type': str, 'data': bytes, 'crc': int}
'''
if not file_bytes:
return chunks
chunk_data_length = int.from_bytes(file_bytes[0:4], 'big')
data_end = 8 + chunk_data_length # 8 bytes for the length and the type
chunk_end = data_end + 4 # The 4 bytes represent the crc
chunk_type = file_bytes[4:8].decode('utf-8')
chunk_data = file_bytes[8:data_end]
chunck_crc = file_bytes[data_end:chunk_end]
chunk = {
'length': chunk_data_length,
'type': chunk_type,
'data': chunk_data,
'crc': chunck_crc
}
chunks.append(chunk)
return split_into_chunks(file_bytes[chunk_end:], chunks)
def parse_ihdr_data(ihdr_chunk):
ihdr_data = ihdr_chunk.get('data')
Header = namedtuple(
'Header',
['width', 'height', 'bit_depth', 'color_type','compression_type',
'filter_type', 'interlace_type'])
return Header(
width=int.from_bytes(ihdr_data[0:4], 'big'),
height=int.from_bytes(ihdr_data[4:8], 'big'),
bit_depth=ihdr_data[8],
color_type=ihdr_data[9],
compression_type=ihdr_data[10],
filter_type = ihdr_data[11],
interlace_type = ihdr_data[12]
)
def create_chunk(data, type):
chunk_type = bytes(type, 'utf-8')
chunk_length = len(data).to_bytes(4, 'big')
buf = chunk_type + data
crc = CRC().create(buf)
return chunk_length + chunk_type + data + crc
# def add_padding(image_bytes, total):
# ''' add_padding : Bytes Int -> Bytes
# add_padding adds null bytes at random throughout the data
# '''
# import random
# return image_bytes
def create_image_data(image_bytes, ratio, bytes_per_pixel):
''' create_image_data: Bytes Fraction Int -> Bytes
creates the Byte string containing the image scanlines at a given ratio
'''
total_pixels = math.ceil(len(image_bytes) / bytes_per_pixel)
con_number = math.sqrt(total_pixels / (ratio.numerator * ratio.denominator))
height = math.ceil(con_number * ratio.denominator)
width = math.ceil(con_number * ratio.numerator)
bytes_per_row = width * bytes_per_pixel
# Adding a single null byte to the start of every row as the filter type 0
# TODO: Optimize filter types
scan_rows = [
bytes(1) + image_bytes[i * bytes_per_row:(i + 1) * bytes_per_row]
for i in range(height)
]
# filling in padding on last row
scan_rows[height - 1] = scan_rows[height - 1] + bytes(
bytes_per_row - len(scan_rows[height - 1]))
return zlib.compress(b''.join(scan_rows))
def create_ihdr_data(
width,
height,
bit_depth=8,
color_type=2,
compression_type=0,
interlace_type=0):
# Filter type is ALWAYS 0
filter_type=0
if bit_depth != 8:
raise NotImplementedError(
'Sorry, bit depths other than 8 have not yet been implemented')
width_b = width.to_bytes(4, 'big')
height_b = height.to_bytes(4, 'big')
bit_depth_b = bit_depth.to_bytes(1, 'big')
color_type_b = color_type.to_bytes(1, 'big')
compression_type_b = compression_type.to_bytes(1, 'big')
filter_type_b = filter_type.to_bytes(1, 'big')
interlace_type_b = interlace_type.to_bytes(1, 'big')
data = (
width_b + height_b + bit_depth_b + color_type_b +
compression_type_b + filter_type_b + interlace_type_b)
return data
def parse_header_and_data(chunks):
''' Returns the IHDR and IDAT from the chunks
'''
try:
image_header = parse_ihdr_data([
chunk for chunk in chunks
if chunk.get('type') == 'IHDR'
][0])
except IndexError as err:
raise InvalidPNG('Missing IHDR chunk')
try:
image_data = zlib.decompress(
b''.join([
idat.get('data')
for idat in [
chunk for chunk in chunks
if chunk.get('type') == 'IDAT']]))
except IndexError as err:
raise InvalidPNG('Invalid IDAT chunks')
return image_header, image_data