xref: /btstack/test/sbc/sbc_encoder.py (revision c21a9c2f3621834ccdbca1efaf4d958002a4ef9c)
1*c21a9c2fSMilanka Ringwald#!/usr/bin/env python
2*c21a9c2fSMilanka Ringwaldimport numpy as np
3*c21a9c2fSMilanka Ringwaldimport wave
4*c21a9c2fSMilanka Ringwaldimport struct
5*c21a9c2fSMilanka Ringwaldimport sys
6*c21a9c2fSMilanka Ringwaldfrom sbc import *
7*c21a9c2fSMilanka Ringwald
8*c21a9c2fSMilanka RingwaldX = np.zeros(40)
9*c21a9c2fSMilanka Ringwald
10*c21a9c2fSMilanka Ringwalddef fetch_samples_for_next_sbc_frame(fin, nr_audio_frames, frame):
11*c21a9c2fSMilanka Ringwald    raw_data = fin.readframes(nr_audio_frames) # Returns byte data
12*c21a9c2fSMilanka Ringwald
13*c21a9c2fSMilanka Ringwald    total_samples = nr_audio_frames * frame.nr_channels
14*c21a9c2fSMilanka Ringwald    fmt = "%ih" % total_samples # read signed 2 byte shorts
15*c21a9c2fSMilanka Ringwald
16*c21a9c2fSMilanka Ringwald    frame.pcm =  np.array(struct.unpack(fmt, raw_data))
17*c21a9c2fSMilanka Ringwald    del raw_data
18*c21a9c2fSMilanka Ringwald
19*c21a9c2fSMilanka Ringwald    # channels = [ [] for ch in range(frame.nr_channels) ]
20*c21a9c2fSMilanka Ringwald    # for index, value in enumerate(integer_data):
21*c21a9c2fSMilanka Ringwald    #     bucket = index % nr_channels
22*c21a9c2fSMilanka Ringwald    #     channels[bucket].append(value)
23*c21a9c2fSMilanka Ringwald
24*c21a9c2fSMilanka Ringwald
25*c21a9c2fSMilanka Ringwalddef sbc_analyse(frame, ch, blk, C, debug):
26*c21a9c2fSMilanka Ringwald    global X
27*c21a9c2fSMilanka Ringwald    M = frame.nr_subbands
28*c21a9c2fSMilanka Ringwald    L = 10 * M
29*c21a9c2fSMilanka Ringwald    M2 = 2*M
30*c21a9c2fSMilanka Ringwald    L2 = 2*L
31*c21a9c2fSMilanka Ringwald
32*c21a9c2fSMilanka Ringwald    Z = np.zeros(L)
33*c21a9c2fSMilanka Ringwald    Y = np.zeros(M2)
34*c21a9c2fSMilanka Ringwald    W = np.zeros(shape=(M, M2))
35*c21a9c2fSMilanka Ringwald    S = np.zeros(M)
36*c21a9c2fSMilanka Ringwald
37*c21a9c2fSMilanka Ringwald    for i in range(L-1, M-1, -1):
38*c21a9c2fSMilanka Ringwald        X[i] = X[i-M]
39*c21a9c2fSMilanka Ringwald    for i in range(M-1, -1, -1):
40*c21a9c2fSMilanka Ringwald        X[i] = frame.EX[M-1-i]
41*c21a9c2fSMilanka Ringwald
42*c21a9c2fSMilanka Ringwald    for i in range(L):
43*c21a9c2fSMilanka Ringwald        Z[i] = X[i] * C[i]
44*c21a9c2fSMilanka Ringwald
45*c21a9c2fSMilanka Ringwald    for i in range(M2):
46*c21a9c2fSMilanka Ringwald        for k in range(5):
47*c21a9c2fSMilanka Ringwald            Y[i] += Z[i+k*8]
48*c21a9c2fSMilanka Ringwald
49*c21a9c2fSMilanka Ringwald    for i in range(M):
50*c21a9c2fSMilanka Ringwald        for k in range(M2):
51*c21a9c2fSMilanka Ringwald            W[i][k] = np.cos((i+0.5)*(k-2)*np.pi/M)
52*c21a9c2fSMilanka Ringwald            S[i] += W[i][k] * Y[k]
53*c21a9c2fSMilanka Ringwald
54*c21a9c2fSMilanka Ringwald    if debug:
55*c21a9c2fSMilanka Ringwald        #print "EX:", frame.EX
56*c21a9c2fSMilanka Ringwald        print "X:", X
57*c21a9c2fSMilanka Ringwald        print "Z:"
58*c21a9c2fSMilanka Ringwald        print "Y:", Y
59*c21a9c2fSMilanka Ringwald        print "W:", W
60*c21a9c2fSMilanka Ringwald        print "S:", S
61*c21a9c2fSMilanka Ringwald
62*c21a9c2fSMilanka Ringwald    for sb in range(M):
63*c21a9c2fSMilanka Ringwald        frame.sb_sample[blk][ch][sb] = S[sb]
64*c21a9c2fSMilanka Ringwald
65*c21a9c2fSMilanka Ringwald
66*c21a9c2fSMilanka Ringwalddef sbc_encode(frame,debug):
67*c21a9c2fSMilanka Ringwald    if frame.nr_subbands == 4:
68*c21a9c2fSMilanka Ringwald        proto_table = Proto_4_40
69*c21a9c2fSMilanka Ringwald    elif frame.nr_subbands == 8:
70*c21a9c2fSMilanka Ringwald        proto_table = Proto_8_80
71*c21a9c2fSMilanka Ringwald    else:
72*c21a9c2fSMilanka Ringwald        return -1
73*c21a9c2fSMilanka Ringwald
74*c21a9c2fSMilanka Ringwald    frame.sb_sample = np.ndarray(shape=(frame.nr_blocks, frame.nr_channels, frame.nr_subbands))
75*c21a9c2fSMilanka Ringwald
76*c21a9c2fSMilanka Ringwald    # channels = [ [] for ch in range(frame.nr_channels) ]
77*c21a9c2fSMilanka Ringwald    # for index, value in enumerate(frame.pcm):
78*c21a9c2fSMilanka Ringwald    #     bucket = index % frame.nr_channels
79*c21a9c2fSMilanka Ringwald    #     channels[bucket].append(value)
80*c21a9c2fSMilanka Ringwald
81*c21a9c2fSMilanka Ringwald    # print "encoder pcm ", frame.pcm
82*c21a9c2fSMilanka Ringwald
83*c21a9c2fSMilanka Ringwald    index = 0
84*c21a9c2fSMilanka Ringwald    for ch in range(frame.nr_channels):
85*c21a9c2fSMilanka Ringwald        for blk in range(frame.nr_blocks):
86*c21a9c2fSMilanka Ringwald            for sb in range(frame.nr_subbands):
87*c21a9c2fSMilanka Ringwald                frame.EX[sb] = frame.pcm[index] #channels[ch][blk * frame.nr_subbands + sb]
88*c21a9c2fSMilanka Ringwald                index+=1
89*c21a9c2fSMilanka Ringwald            sbc_analyse(frame, ch, blk, proto_table,debug)
90*c21a9c2fSMilanka Ringwald    sbc_quantization(frame)
91*c21a9c2fSMilanka Ringwald
92*c21a9c2fSMilanka Ringwald
93*c21a9c2fSMilanka Ringwalddef should_use_joint_coding(frame):
94*c21a9c2fSMilanka Ringwald    return False
95*c21a9c2fSMilanka Ringwald
96*c21a9c2fSMilanka Ringwalddef calculate_scalefactor(max_subbandsample):
97*c21a9c2fSMilanka Ringwald    x = 0
98*c21a9c2fSMilanka Ringwald    while True:
99*c21a9c2fSMilanka Ringwald        y = 1 << x + 1
100*c21a9c2fSMilanka Ringwald        if y > max_subbandsample:
101*c21a9c2fSMilanka Ringwald            break
102*c21a9c2fSMilanka Ringwald        x += 1
103*c21a9c2fSMilanka Ringwald    return (x,y)
104*c21a9c2fSMilanka Ringwald
105*c21a9c2fSMilanka Ringwald
106*c21a9c2fSMilanka Ringwalddef frame_to_bitstream(frame):
107*c21a9c2fSMilanka Ringwald    global bitstream, bitstream_bits_available
108*c21a9c2fSMilanka Ringwald    init_bitstream()
109*c21a9c2fSMilanka Ringwald
110*c21a9c2fSMilanka Ringwald    add_bits(frame.syncword, 8)
111*c21a9c2fSMilanka Ringwald    add_bits(frame.sampling_frequency, 2)
112*c21a9c2fSMilanka Ringwald    add_bits(frame.nr_blocks/4-1, 2)
113*c21a9c2fSMilanka Ringwald    add_bits(frame.channel_mode, 2)
114*c21a9c2fSMilanka Ringwald    add_bits(frame.allocation_method, 1)
115*c21a9c2fSMilanka Ringwald    add_bits(frame.nr_subbands/4-1, 1)
116*c21a9c2fSMilanka Ringwald    add_bits(frame.bitpool, 8)
117*c21a9c2fSMilanka Ringwald    add_bits(frame.crc_check, 8)
118*c21a9c2fSMilanka Ringwald
119*c21a9c2fSMilanka Ringwald    for sb in range(frame.nr_subbands):
120*c21a9c2fSMilanka Ringwald        add_bits(frame.join[sb],1)
121*c21a9c2fSMilanka Ringwald
122*c21a9c2fSMilanka Ringwald    for ch in range(frame.nr_channels):
123*c21a9c2fSMilanka Ringwald        for sb in range(frame.nr_subbands):
124*c21a9c2fSMilanka Ringwald            add_bits(frame.scale_factor[ch][sb], 4)
125*c21a9c2fSMilanka Ringwald
126*c21a9c2fSMilanka Ringwald    for blk in range(frame.nr_blocks):
127*c21a9c2fSMilanka Ringwald        for ch in range(frame.nr_channels):
128*c21a9c2fSMilanka Ringwald            for sb in range(frame.nr_subbands):
129*c21a9c2fSMilanka Ringwald                add_bits(frame.audio_sample[blk][ch][sb], frame.bits[ch][sb])
130*c21a9c2fSMilanka Ringwald
131*c21a9c2fSMilanka Ringwald    # bitstream_len = 16 + frame.nr_subbands + frame.nr_channels * frame.nr_subbands * 4
132*c21a9c2fSMilanka Ringwald    return bitstream
133*c21a9c2fSMilanka Ringwald
134*c21a9c2fSMilanka Ringwalddef sbc_quantization(frame):
135*c21a9c2fSMilanka Ringwald
136*c21a9c2fSMilanka Ringwald    frame.join = np.zeros(frame.nr_subbands, dtype = np.uint8)
137*c21a9c2fSMilanka Ringwald    if should_use_joint_coding(frame):
138*c21a9c2fSMilanka Ringwald        return
139*c21a9c2fSMilanka Ringwald
140*c21a9c2fSMilanka Ringwald    max_subbandsample = np.zeros(shape = (frame.nr_channels, frame.nr_subbands))
141*c21a9c2fSMilanka Ringwald
142*c21a9c2fSMilanka Ringwald    for blk in range(frame.nr_blocks):
143*c21a9c2fSMilanka Ringwald        for ch in range(frame.nr_channels):
144*c21a9c2fSMilanka Ringwald            for sb in range(frame.nr_subbands):
145*c21a9c2fSMilanka Ringwald                m = abs(frame.sb_sample[blk][ch][sb])
146*c21a9c2fSMilanka Ringwald                if max_subbandsample[ch][sb] < m:
147*c21a9c2fSMilanka Ringwald                    max_subbandsample[ch][sb] = m
148*c21a9c2fSMilanka Ringwald
149*c21a9c2fSMilanka Ringwald
150*c21a9c2fSMilanka Ringwald    for ch in range(frame.nr_channels):
151*c21a9c2fSMilanka Ringwald        for sb in range(frame.nr_subbands):
152*c21a9c2fSMilanka Ringwald            (frame.scale_factor[ch][sb], frame.scalefactor[ch][sb]) = calculate_scalefactor(max_subbandsample[ch][sb])
153*c21a9c2fSMilanka Ringwald
154*c21a9c2fSMilanka Ringwald    frame.bits = sbc_bit_allocation(frame)
155*c21a9c2fSMilanka Ringwald
156*c21a9c2fSMilanka Ringwald    # Reconstruct the Audio Samples
157*c21a9c2fSMilanka Ringwald    frame.levels = np.zeros(shape=(frame.nr_channels, frame.nr_subbands), dtype = np.int32)
158*c21a9c2fSMilanka Ringwald    for ch in range(frame.nr_channels):
159*c21a9c2fSMilanka Ringwald        for sb in range(frame.nr_subbands):
160*c21a9c2fSMilanka Ringwald            frame.levels[ch][sb] = pow(2.0, frame.bits[ch][sb]) - 1
161*c21a9c2fSMilanka Ringwald
162*c21a9c2fSMilanka Ringwald    frame.syncword = 156
163*c21a9c2fSMilanka Ringwald    frame.crc_check = calculate_crc(frame)
164*c21a9c2fSMilanka Ringwald
165*c21a9c2fSMilanka Ringwald    for blk in range(frame.nr_blocks):
166*c21a9c2fSMilanka Ringwald        for ch in range(frame.nr_channels):
167*c21a9c2fSMilanka Ringwald            for sb in range(frame.nr_subbands):
168*c21a9c2fSMilanka Ringwald                if frame.levels[ch][sb] > 0:
169*c21a9c2fSMilanka Ringwald                    SB = frame.sb_sample[blk][ch][sb]
170*c21a9c2fSMilanka Ringwald                    SF = frame.scalefactor[ch][sb]
171*c21a9c2fSMilanka Ringwald                    L  = frame.levels[ch][sb]
172*c21a9c2fSMilanka Ringwald                    AS = np.uint16(((SB * L / SF    + L) - 1.0)/2.0)
173*c21a9c2fSMilanka Ringwald                    frame.audio_sample[blk][ch][sb] = AS
174*c21a9c2fSMilanka Ringwald                else:
175*c21a9c2fSMilanka Ringwald                    frame.audio_sample[blk][ch][sb] = 0
176*c21a9c2fSMilanka Ringwald
177*c21a9c2fSMilanka Ringwald
178*c21a9c2fSMilanka Ringwald    return 0
179*c21a9c2fSMilanka Ringwald
180*c21a9c2fSMilanka Ringwald
181*c21a9c2fSMilanka Ringwald# usage = '''
182*c21a9c2fSMilanka Ringwald# Usage: ./sbc_encoder.py input.wav block_size nr_subbands bitpool
183*c21a9c2fSMilanka Ringwald# '''
184*c21a9c2fSMilanka Ringwald# nr_blocks = 0
185*c21a9c2fSMilanka Ringwald# nr_subbands = 0
186*c21a9c2fSMilanka Ringwald
187*c21a9c2fSMilanka Ringwald# if (len(sys.argv) < 5):
188*c21a9c2fSMilanka Ringwald#     print(usage)
189*c21a9c2fSMilanka Ringwald#     sys.exit(1)
190*c21a9c2fSMilanka Ringwald# try:
191*c21a9c2fSMilanka Ringwald#     infile = sys.argv[1]
192*c21a9c2fSMilanka Ringwald#     if not infile.endswith('.wav'):
193*c21a9c2fSMilanka Ringwald#         print(usage)
194*c21a9c2fSMilanka Ringwald#         sys.exit(1)
195*c21a9c2fSMilanka Ringwald#     nr_blocks = int(sys.argv[2])
196*c21a9c2fSMilanka Ringwald#     nr_subbands = int(sys.argv[3])
197*c21a9c2fSMilanka Ringwald#     bitpool = int(sys.argv[4])
198*c21a9c2fSMilanka Ringwald#     sbcfile = infile.replace('.wav', '-encoded.sbc')
199*c21a9c2fSMilanka Ringwald
200*c21a9c2fSMilanka Ringwald#     fin = wave.open(infile, 'rb')
201*c21a9c2fSMilanka Ringwald
202*c21a9c2fSMilanka Ringwald#     wav_nr_channels = fin.getnchannels()
203*c21a9c2fSMilanka Ringwald#     wav_sample_rate = fin.getframerate()
204*c21a9c2fSMilanka Ringwald#     wav_nr_frames = fin.getnframes()
205*c21a9c2fSMilanka Ringwald#     sbc_sampling_frequency = sbc_sampling_frequency_index(wav_sample_rate)
206*c21a9c2fSMilanka Ringwald
207*c21a9c2fSMilanka Ringwald#     sbc_frame_count = 0
208*c21a9c2fSMilanka Ringwald#     audio_frame_count = 0
209*c21a9c2fSMilanka Ringwald
210*c21a9c2fSMilanka Ringwald#     while audio_frame_count < wav_nr_frames:
211*c21a9c2fSMilanka Ringwald#         if sbc_frame_count % 200 == 0:
212*c21a9c2fSMilanka Ringwald#             print "== Frame %d ==" % (sbc_frame_count)
213*c21a9c2fSMilanka Ringwald
214*c21a9c2fSMilanka Ringwald#         sbc_encoder_frame = SBCFrame(nr_blocks, nr_subbands, wav_nr_channels, sbc_sampling_frequency, bitpool)
215*c21a9c2fSMilanka Ringwald
216*c21a9c2fSMilanka Ringwald#         wav_nr_audio_frames = sbc_encoder_frame.nr_blocks * sbc_encoder_frame.nr_subbands
217*c21a9c2fSMilanka Ringwald#         fetch_samples_for_next_sbc_frame(fin, wav_nr_audio_frames, sbc_encoder_frame)
218*c21a9c2fSMilanka Ringwald#         sbc_encode(sbc_encoder_frame)
219*c21a9c2fSMilanka Ringwald
220*c21a9c2fSMilanka Ringwald#         # stream = frame_to_bitstream(frame)
221*c21a9c2fSMilanka Ringwald#         audio_frame_count += wav_nr_audio_frames
222*c21a9c2fSMilanka Ringwald#         sbc_frame_count += 1
223*c21a9c2fSMilanka Ringwald
224*c21a9c2fSMilanka Ringwald#     # except TypeError:
225*c21a9c2fSMilanka Ringwald#     #     fin.close()
226*c21a9c2fSMilanka Ringwald#     #     print "DONE, WAV file %s encoded into SBC file %s ", (infile, sbcfile)
227*c21a9c2fSMilanka Ringwald
228*c21a9c2fSMilanka Ringwald#     #channels, num_audio_frames, wav_nr_channels, wav_sample_rate = read_waw_file(wavfile)
229*c21a9c2fSMilanka Ringwald
230*c21a9c2fSMilanka Ringwald
231*c21a9c2fSMilanka Ringwald# except IOError as e:
232*c21a9c2fSMilanka Ringwald#     print(usage)
233*c21a9c2fSMilanka Ringwald#     sys.exit(1)
234*c21a9c2fSMilanka Ringwald
235*c21a9c2fSMilanka Ringwald
236*c21a9c2fSMilanka Ringwald
237*c21a9c2fSMilanka Ringwald
238*c21a9c2fSMilanka Ringwald
239