xref: /btstack/test/sbc/sbc_encoder.py (revision 5c9bef5ba273d4d18de0f76f9370002de99dfae4)
1#!/usr/bin/env python
2import numpy as np
3import wave
4import struct
5import sys
6from sbc import *
7
8X = np.zeros(shape=(2,80), dtype = np.int16)
9
10
11def fetch_samples_for_next_sbc_frame(fin, frame):
12    nr_samples = frame.nr_blocks * frame.nr_subbands
13    raw_data = fin.readframes(nr_samples) # Returns byte data
14    fmt = "%ih" % (len(raw_data) / 2)
15    data = struct.unpack(fmt, raw_data)
16
17    if frame.nr_channels == 2:
18        for i in range(len(data)/2 - 1):
19            frame.pcm[0][i] = data[2*i]
20            frame.pcm[1][i] = data[2*i+1]
21    else:
22        for i in range(len(data)):
23            frame.pcm[0][i] = data[i]
24
25
26def sbc_frame_analysis(frame, ch, blk, C):
27    global X
28
29    M = frame.nr_subbands
30    L = 10 * M
31    M2 = 2*M
32    L2 = 2*L
33
34    Z = np.zeros(L)
35    Y = np.zeros(M2)
36    W = np.zeros(shape=(M, M2))
37    S = np.zeros(M)
38
39    for i in range(L-1, M-1, -1):
40        X[ch][i] = X[ch][i-M]
41    for i in range(M-1, -1, -1):
42        X[ch][i] = frame.EX[M-1-i]
43
44    for i in range(L):
45        Z[i] = X[ch][i] * C[i]
46
47    for i in range(M2):
48        for k in range(5):
49            Y[i] += Z[i+k*M2]
50
51    for i in range(M):
52        for k in range(M2):
53            W[i][k] = np.cos((i+0.5)*(k-M/2)*np.pi/M)
54            S[i] += W[i][k] * Y[k]
55
56    for sb in range(M):
57        frame.sb_sample[blk][ch][sb] = S[sb]
58
59def sbc_analysis(frame):
60    if frame.nr_subbands == 4:
61        C = Proto_4_40
62    elif frame.nr_subbands == 8:
63        C = Proto_8_80
64    else:
65        return -1
66
67    frame.sb_sample = np.ndarray(shape=(frame.nr_blocks, frame.nr_channels, frame.nr_subbands))
68    for ch in range(frame.nr_channels):
69        index = 0
70        for blk in range(frame.nr_blocks):
71            for sb in range(frame.nr_subbands):
72                frame.EX[sb] = np.int16(frame.pcm[ch][index])
73                index+=1
74            sbc_frame_analysis(frame, ch, blk, C)
75    return 0
76
77def sbc_encode(frame):
78    err = sbc_analysis(frame)
79    if err >= 0:
80        err = sbc_quantization(frame)
81    return err
82
83def sbc_quantization(frame):
84    frame.scale_factor, frame.scalefactor = calculate_scalefactors(frame.nr_blocks, frame.nr_channels, frame.nr_subbands, frame.sb_sample)
85    calculate_channel_mode(frame)
86    frame.bits = sbc_bit_allocation(frame)
87
88    # Reconstruct the Audio Samples
89    frame.levels = np.zeros(shape=(frame.nr_channels, frame.nr_subbands), dtype = np.int32)
90    for ch in range(frame.nr_channels):
91        for sb in range(frame.nr_subbands):
92            frame.levels[ch][sb] = (1 << frame.bits[ch][sb]) - 1 #pow(2.0, frame.bits[ch][sb]) - 1
93
94    frame.syncword = 156
95    frame.crc_check = calculate_crc(frame)
96
97    for blk in range(frame.nr_blocks):
98        for ch in range(frame.nr_channels):
99            for sb in range(frame.nr_subbands):
100                if frame.levels[ch][sb] > 0:
101                    SB = frame.sb_sample[blk][ch][sb]
102                    L  = frame.levels[ch][sb]
103                    SF = frame.scalefactor[ch][sb]
104                    frame.audio_sample[blk][ch][sb] = np.uint16(((SB * L / SF    + L) - 1.0)/2.0)
105                else:
106                    frame.audio_sample[blk][ch][sb] = 0
107
108    return 0
109
110def sbc_write_frame(fout, sbc_encoder_frame):
111    stream = frame_to_bitstream(sbc_encoder_frame)
112    barray = bytearray(stream)
113    fout.write(barray)
114
115if __name__ == "__main__":
116    usage = '''
117    Usage:      ./sbc_encoder.py input.wav blocks subbands bitpool allocation_method[0-LOUDNESS,1-SNR]
118    Example:    ./sbc_encoder.py fanfare.wav 16 4 31 0
119    '''
120    nr_blocks = 0
121    nr_subbands = 0
122
123
124    if (len(sys.argv) < 6):
125        print(usage)
126        sys.exit(1)
127    try:
128        infile = sys.argv[1]
129        if not infile.endswith('.wav'):
130            print(usage)
131            sys.exit(1)
132        sbcfile = infile.replace('.wav', '-encoded.sbc')
133
134        nr_blocks = int(sys.argv[2])
135        nr_subbands = int(sys.argv[3])
136        bitpool = int(sys.argv[4])
137        allocation_method = int(sys.argv[5])
138
139        fin = wave.open(infile, 'rb')
140        nr_channels = fin.getnchannels()
141        sampling_frequency = fin.getframerate()
142        nr_audio_frames = fin.getnframes()
143
144        subband_frame_count = 0
145        audio_frame_count = 0
146        nr_samples = nr_blocks * nr_subbands
147        fout = open(sbcfile, 'wb')
148        while audio_frame_count < nr_audio_frames:
149            if subband_frame_count % 200 == 0:
150                print("== Frame %d ==" % (subband_frame_count))
151
152            sbc_encoder_frame = SBCFrame(nr_blocks, nr_subbands, nr_channels, bitpool, sampling_frequency, allocation_method)
153            fetch_samples_for_next_sbc_frame(fin, sbc_encoder_frame)
154
155            sbc_encode(sbc_encoder_frame)
156            sbc_write_frame(fout, sbc_encoder_frame)
157
158            # if subband_frame_count == 0:
159            #     exit(0)
160            audio_frame_count += nr_samples
161            subband_frame_count += 1
162
163
164        fin.close()
165        fout.close()
166        print("DONE, WAV file %s encoded into SBC file %s " % (infile, sbcfile))
167
168
169    except IOError as e:
170        print(usage)
171        sys.exit(1)
172
173
174
175
176
177