1#!/usr/bin/env python3
2#
3# Copyright 2022 Google LLC
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import numpy as np
19import scipy.signal as signal
20import scipy.io.wavfile as wavfile
21import struct
22import argparse
23
24import lc3
25import tables as T, appendix_c as C
26
27import attdet, ltpf
28import mdct, energy, bwdet, sns, tns, spec
29import bitstream
30
31### ------------------------------------------------------------------------ ###
32
33class Encoder:
34
35    def __init__(self, dt_ms, sr_hz):
36
37        dt = { 7.5: T.DT_7M5, 10: T.DT_10M }[dt_ms]
38
39        sr = {  8000: T.SRATE_8K , 16000: T.SRATE_16K, 24000: T.SRATE_24K,
40               32000: T.SRATE_32K, 48000: T.SRATE_48K }[sr_hz]
41
42        self.ne = T.NE[dt][sr]
43
44        self.attdet = attdet.AttackDetector(dt, sr)
45        self.ltpf = ltpf.LtpfAnalysis(dt, sr)
46
47        self.mdct = mdct.MdctForward(dt, sr)
48        self.energy = energy.EnergyBand(dt, sr)
49        self.bwdet = bwdet.BandwidthDetector(dt, sr)
50        self.sns = sns.SnsAnalysis(dt, sr)
51        self.tns = tns.TnsAnalysis(dt)
52        self.spec = spec.SpectrumAnalysis(dt, sr)
53
54    def analyse(self, x, nbytes):
55
56        att = self.attdet.run(nbytes, x)
57
58        pitch_present = self.ltpf.run(x)
59
60        x = self.mdct.run(x)[:self.ne]
61
62        (e, nn_flag) = self.energy.compute(x)
63        if nn_flag:
64            self.ltpf.disable()
65
66        bw = self.bwdet.run(e)
67
68        x = self.sns.run(e, att, x)
69
70        x = self.tns.run(x, bw, nn_flag, nbytes)
71
72        (xq, lastnz, x) = self.spec.run(bw, nbytes,
73            self.bwdet.get_nbits(), self.ltpf.get_nbits(),
74            self.sns.get_nbits(), self.tns.get_nbits(), x)
75
76        return pitch_present
77
78    def encode(self, pitch_present, nbytes):
79
80        b = bitstream.BitstreamWriter(nbytes)
81
82        self.bwdet.store(b)
83
84        self.spec.store(b)
85
86        self.tns.store(b)
87
88        b.write_bit(pitch_present)
89
90        self.sns.store(b)
91
92        if pitch_present:
93            self.ltpf.store(b)
94
95        self.spec.encode(b)
96
97        return b.terminate()
98
99    def run(self, x, nbytes):
100
101        pitch_present = self.analyse(x, nbytes)
102
103        data = self.encode(pitch_present, nbytes)
104
105        return data
106
107### ------------------------------------------------------------------------ ###
108
109def check_appendix_c(dt):
110
111    ok = True
112
113    enc_c = lc3.setup_encoder(int(T.DT_MS[dt] * 1000), 16000)
114
115    for i in range(len(C.X_PCM[dt])):
116
117        data = lc3.encode(enc_c, C.X_PCM[dt][i], C.NBYTES[dt])
118        ok = ok and data == C.BYTES_AC[dt][i]
119
120    return ok
121
122def check():
123
124    ok = True
125
126    for dt in range(T.NUM_DT):
127        ok = ok and check_appendix_c(dt)
128
129    return ok
130
131### ------------------------------------------------------------------------ ###
132
133def dump(data):
134    for i in range(0, len(data), 20):
135        print(''.join('{:02x} '.format(x)
136            for x in data[i:min(i+20, len(data))] ))
137
138if __name__ == "__main__":
139
140    parser = argparse.ArgumentParser(description='LC3 Encoder Test Framework')
141    parser.add_argument('wav_file',
142        help='Input wave file', type=argparse.FileType('r'))
143    parser.add_argument('--bitrate',
144        help='Bitrate in bps', type=int, required=True)
145    parser.add_argument('--dt',
146        help='Frame duration in ms', type=float, default=10)
147    parser.add_argument('--pyout',
148        help='Python output file', type=argparse.FileType('w'))
149    parser.add_argument('--cout',
150        help='C output file', type=argparse.FileType('w'))
151    args = parser.parse_args()
152
153    if args.bitrate < 16000 or args.bitrate > 320000:
154        raise ValueError('Invalid bitate %d bps' % args.bitrate)
155
156    if args.dt not in (7.5, 10):
157        raise ValueError('Invalid frame duration %.1f ms' % args.dt)
158
159    (sr_hz, pcm) = wavfile.read(args.wav_file.name)
160    if sr_hz not in (8000, 16000, 24000, 320000, 48000):
161        raise ValueError('Unsupported input samplerate: %d' % sr_hz)
162    if pcm.ndim != 1:
163        raise ValueError('Only single channel wav file supported')
164
165    ### Setup ###
166
167    enc = Encoder(args.dt, sr_hz)
168    enc_c = lc3.setup_encoder(int(args.dt * 1000), sr_hz)
169
170    frame_samples = int((args.dt * sr_hz) / 1000)
171    frame_nbytes = int((args.bitrate * args.dt) / (1000 * 8))
172
173    ### File Header ###
174
175    f_py = open(args.pyout.name, 'wb') if args.pyout else None
176    f_c  = open(args.cout.name , 'wb') if args.cout  else None
177
178    header = struct.pack('=HHHHHHHI', 0xcc1c, 18,
179        sr_hz // 100, args.bitrate // 100, 1, int(args.dt * 100), 0, len(pcm))
180
181    for f in (f_py, f_c):
182        if f: f.write(header)
183
184    ### Encoding loop ###
185
186    if len(pcm) % frame_samples > 0:
187        pcm = np.append(pcm, np.zeros(frame_samples - (len(pcm) % frame_samples)))
188
189    for i in range(0, len(pcm), frame_samples):
190
191        print('Encoding frame %d' % (i // frame_samples), end='\r')
192
193        frame_pcm = pcm[i:i+frame_samples]
194
195        data = enc.run(frame_pcm, frame_nbytes)
196        data_c = lc3.encode(enc_c, frame_pcm, frame_nbytes)
197
198        for f in (f_py, f_c):
199            if f: f.write(struct.pack('=H', frame_nbytes))
200
201        if f_py: f_py.write(data)
202        if f_c: f_c.write(data_c)
203
204    print('done ! %16s' % '')
205
206    ### Terminate ###
207
208    for f in (f_py, f_c):
209        if f: f.close()
210
211
212### ------------------------------------------------------------------------ ###
213