-
Notifications
You must be signed in to change notification settings - Fork 1
/
decoder_api.py
123 lines (118 loc) · 3.68 KB
/
decoder_api.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
import cv2
import numpy as np
import decoder_utils
def JPEG_decoder(data: np.ndarray, show=False) -> np.ndarray:
assert data.dtype == np.uint8
nowIdx = decoder_utils.decode_head(data)
if nowIdx is None or data[-2] != 0xff or data[-1] != 0xd9:
print("Error: unsupported jpeg format")
exit(1)
data = data[nowIdx:-2]
if decoder_utils.d == 1:
YCbCr = np.zeros((decoder_utils.h, decoder_utils.w), np.uint8)
last_block_dc = [0]
else:
YCbCr = np.zeros((decoder_utils.h, decoder_utils.w, decoder_utils.d), np.uint8)
last_block_dc = [0, 0, 0]
last_block_dc = [0, 0, 0]
nowLeft = 0
nowRight = 0
nowx = 0
nowy = 0
nowIdx, buffer = decoder_utils.load_buffer(data, '', 0)
while 1:
if show:
print("%5.2f%%" % (nowIdx / len(data)*100), end="")
if decoder_utils.d == 1:
nowblock_decoded = np.zeros((8, 8), np.int16)
else:
nowblock_decoded = np.zeros((8, 8, decoder_utils.d), np.int16)
for i in range(1, decoder_utils.d+1):
nowblock = np.zeros((64), np.int16)
nowblockIdx = 0
# decode DC
while 1:
try:
var = decoder_utils.DC_color_huffman_dict[i][buffer[nowLeft:nowRight+1]]
if var == 0:
num = 0
else:
num = int(buffer[nowRight+1:nowRight+var+1], 2)
if buffer[nowRight+1] == '0':
num = -((1<<var) - 1) + num
nowblock[nowblockIdx] = num + last_block_dc[i-1]
nowblockIdx += 1
nowLeft = nowRight + var + 1
nowRight = nowLeft
if nowRight + 32 >= len(buffer):
nowIdx, buffer = decoder_utils.load_buffer(data, buffer, nowIdx)
break
except KeyError:
nowRight += 1
if nowRight + 32 >= len(buffer):
nowIdx, buffer = decoder_utils.load_buffer(data, buffer, nowIdx)
last_block_dc[i-1] = nowblock[0]
buffer = buffer[nowRight:]
nowLeft = 0
nowRight = 0
# decode AC
while 1:
try:
var = decoder_utils.AC_color_huffman_dict[i][buffer[nowLeft:nowRight+1]]
num0 = var >> 4
num1 = var & 0xf
nowblockIdx += num0
if num0 == 15 and num1 == 0:
nowblockIdx += 1
elif num0 == 0 and num1 == 0:
nowLeft = nowRight + 1
nowRight = nowLeft
break
else:
num = int(buffer[nowRight+1:nowRight+num1+1], 2)
if buffer[nowRight+1] == '0':
num = -((1<<num1) - 1) + num
nowblock[nowblockIdx] = num
nowblockIdx += 1
nowLeft = nowRight + num1 + 1
nowRight = nowLeft
if nowRight + 32 >= len(buffer):
nowIdx, buffer = decoder_utils.load_buffer(data, buffer, nowIdx)
if nowblockIdx == 64:
break
assert nowblockIdx < 64
except KeyError:
nowRight += 1
if nowRight + 32 >= len(buffer):
nowIdx, buffer = decoder_utils.load_buffer(data, buffer, nowIdx)
buffer = buffer[nowRight:]
nowLeft = 0
nowRight = 0
nowblock_dezz = decoder_utils.zz2block(nowblock)
nowblock_dequt = nowblock_dezz * decoder_utils.color_quant_dict[i]
nowblock_idct = cv2.idct(nowblock_dequt.astype(np.float32)).round() + 128
nowblock_idct[nowblock_idct < 0] = 0
nowblock_idct[nowblock_idct > 255] = 255
if decoder_utils.d == 3:
nowblock_decoded[:, :, i-1] = nowblock_idct.astype(np.uint8)
if show:
print("\b\b\b\b\b\b\b\b\b\b", end="")
if decoder_utils.d == 3:
YCbCr[nowx:nowx+8, nowy:nowy+8, :] = nowblock_decoded
else:
YCbCr[nowx:nowx+8, nowy:nowy+8] = nowblock_idct.astype(np.uint8)
nowy += 8
if nowy >= decoder_utils.w:
nowy = 0
nowx += 8
if nowx >= decoder_utils.h:
break
if decoder_utils.d == 3:
YCrCb = np.zeros((decoder_utils.h, decoder_utils.w, decoder_utils.d), np.uint8)
YCrCb[:, :, 0] = YCbCr[:, :, 0]
YCrCb[:, :, 1] = YCbCr[:, :, 2]
YCrCb[:, :, 2] = YCbCr[:, :, 1]
BGR = cv2.cvtColor(YCrCb, cv2.COLOR_YCrCb2BGR)
return BGR
else:
return YCbCr