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