1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 // Functions to write audio in WAV format.
17
18 #include <math.h>
19 #include <string.h>
20 #include <algorithm>
21
22 #include "absl/base/casts.h"
23 #include "tensorflow/core/lib/core/coding.h"
24 #include "tensorflow/core/lib/core/errors.h"
25 #include "tensorflow/core/lib/wav/wav_io.h"
26 #include "tensorflow/core/platform/byte_order.h"
27 #include "tensorflow/core/platform/logging.h"
28 #include "tensorflow/core/platform/macros.h"
29
30 namespace tensorflow {
31 namespace wav {
32 namespace {
33
34 struct TF_PACKED RiffChunk {
35 char chunk_id[4];
36 char chunk_data_size[4];
37 char riff_type[4];
38 };
39 static_assert(sizeof(RiffChunk) == 12, "TF_PACKED does not work.");
40
41 struct TF_PACKED FormatChunk {
42 char chunk_id[4];
43 char chunk_data_size[4];
44 char compression_code[2];
45 char channel_numbers[2];
46 char sample_rate[4];
47 char bytes_per_second[4];
48 char bytes_per_frame[2];
49 char bits_per_sample[2];
50 };
51 static_assert(sizeof(FormatChunk) == 24, "TF_PACKED does not work.");
52
53 struct TF_PACKED DataChunk {
54 char chunk_id[4];
55 char chunk_data_size[4];
56 };
57 static_assert(sizeof(DataChunk) == 8, "TF_PACKED does not work.");
58
59 struct TF_PACKED WavHeader {
60 RiffChunk riff_chunk;
61 FormatChunk format_chunk;
62 DataChunk data_chunk;
63 };
64 static_assert(sizeof(WavHeader) ==
65 sizeof(RiffChunk) + sizeof(FormatChunk) + sizeof(DataChunk),
66 "TF_PACKED does not work.");
67
68 constexpr char kRiffChunkId[] = "RIFF";
69 constexpr char kRiffType[] = "WAVE";
70 constexpr char kFormatChunkId[] = "fmt ";
71 constexpr char kDataChunkId[] = "data";
72
FloatToInt16Sample(float data)73 inline int16 FloatToInt16Sample(float data) {
74 constexpr float kMultiplier = 1.0f * (1 << 15);
75 return std::min<float>(std::max<float>(roundf(data * kMultiplier), kint16min),
76 kint16max);
77 }
78
Int16SampleToFloat(int16 data)79 inline float Int16SampleToFloat(int16 data) {
80 constexpr float kMultiplier = 1.0f / (1 << 15);
81 return data * kMultiplier;
82 }
83
84 } // namespace
85
86 // Handles moving the data index forward, validating the arguments, and avoiding
87 // overflow or underflow.
IncrementOffset(int old_offset,size_t increment,size_t max_size,int * new_offset)88 Status IncrementOffset(int old_offset, size_t increment, size_t max_size,
89 int* new_offset) {
90 if (old_offset < 0) {
91 return errors::InvalidArgument("Negative offsets are not allowed: ",
92 old_offset);
93 }
94 if (old_offset > max_size) {
95 return errors::InvalidArgument("Initial offset is outside data range: ",
96 old_offset);
97 }
98 *new_offset = old_offset + increment;
99 if (*new_offset > max_size) {
100 return errors::InvalidArgument("Data too short when trying to read string");
101 }
102 // See above for the check that the input offset is positive. If it's negative
103 // here then it means that there's been an overflow in the arithmetic.
104 if (*new_offset < 0) {
105 return errors::InvalidArgument("Offset too large, overflowed: ",
106 *new_offset);
107 }
108 return Status::OK();
109 }
110
ExpectText(const string & data,const string & expected_text,int * offset)111 Status ExpectText(const string& data, const string& expected_text,
112 int* offset) {
113 int new_offset;
114 TF_RETURN_IF_ERROR(
115 IncrementOffset(*offset, expected_text.size(), data.size(), &new_offset));
116 const string found_text(data.begin() + *offset, data.begin() + new_offset);
117 if (found_text != expected_text) {
118 return errors::InvalidArgument("Header mismatch: Expected ", expected_text,
119 " but found ", found_text);
120 }
121 *offset = new_offset;
122 return Status::OK();
123 }
124
ReadString(const string & data,int expected_length,string * value,int * offset)125 Status ReadString(const string& data, int expected_length, string* value,
126 int* offset) {
127 int new_offset;
128 TF_RETURN_IF_ERROR(
129 IncrementOffset(*offset, expected_length, data.size(), &new_offset));
130 *value = string(data.begin() + *offset, data.begin() + new_offset);
131 *offset = new_offset;
132 return Status::OK();
133 }
134
EncodeAudioAsS16LEWav(const float * audio,size_t sample_rate,size_t num_channels,size_t num_frames,string * wav_string)135 Status EncodeAudioAsS16LEWav(const float* audio, size_t sample_rate,
136 size_t num_channels, size_t num_frames,
137 string* wav_string) {
138 constexpr size_t kFormatChunkSize = 16;
139 constexpr size_t kCompressionCodePcm = 1;
140 constexpr size_t kBitsPerSample = 16;
141 constexpr size_t kBytesPerSample = kBitsPerSample / 8;
142 constexpr size_t kHeaderSize = sizeof(WavHeader);
143
144 if (audio == nullptr) {
145 return errors::InvalidArgument("audio is null");
146 }
147 if (wav_string == nullptr) {
148 return errors::InvalidArgument("wav_string is null");
149 }
150 if (sample_rate == 0 || sample_rate > kuint32max) {
151 return errors::InvalidArgument("sample_rate must be in (0, 2^32), got: ",
152 sample_rate);
153 }
154 if (num_channels == 0 || num_channels > kuint16max) {
155 return errors::InvalidArgument("num_channels must be in (0, 2^16), got: ",
156 num_channels);
157 }
158 if (num_frames == 0) {
159 return errors::InvalidArgument("num_frames must be positive.");
160 }
161
162 const size_t bytes_per_second = sample_rate * kBytesPerSample * num_channels;
163 const size_t num_samples = num_frames * num_channels;
164 const size_t data_size = num_samples * kBytesPerSample;
165 const size_t file_size = kHeaderSize + num_samples * kBytesPerSample;
166 const size_t bytes_per_frame = kBytesPerSample * num_channels;
167
168 // WAV represents the length of the file as a uint32 so file_size cannot
169 // exceed kuint32max.
170 if (file_size > kuint32max) {
171 return errors::InvalidArgument(
172 "Provided channels and frames cannot be encoded as a WAV.");
173 }
174
175 wav_string->resize(file_size);
176 char* data = &wav_string->at(0);
177 WavHeader* header = absl::bit_cast<WavHeader*>(data);
178
179 // Fill RIFF chunk.
180 auto* riff_chunk = &header->riff_chunk;
181 memcpy(riff_chunk->chunk_id, kRiffChunkId, 4);
182 core::EncodeFixed32(riff_chunk->chunk_data_size, file_size - 8);
183 memcpy(riff_chunk->riff_type, kRiffType, 4);
184
185 // Fill format chunk.
186 auto* format_chunk = &header->format_chunk;
187 memcpy(format_chunk->chunk_id, kFormatChunkId, 4);
188 core::EncodeFixed32(format_chunk->chunk_data_size, kFormatChunkSize);
189 core::EncodeFixed16(format_chunk->compression_code, kCompressionCodePcm);
190 core::EncodeFixed16(format_chunk->channel_numbers, num_channels);
191 core::EncodeFixed32(format_chunk->sample_rate, sample_rate);
192 core::EncodeFixed32(format_chunk->bytes_per_second, bytes_per_second);
193 core::EncodeFixed16(format_chunk->bytes_per_frame, bytes_per_frame);
194 core::EncodeFixed16(format_chunk->bits_per_sample, kBitsPerSample);
195
196 // Fill data chunk.
197 auto* data_chunk = &header->data_chunk;
198 memcpy(data_chunk->chunk_id, kDataChunkId, 4);
199 core::EncodeFixed32(data_chunk->chunk_data_size, data_size);
200
201 // Write the audio.
202 data += kHeaderSize;
203 for (size_t i = 0; i < num_samples; ++i) {
204 int16 sample = FloatToInt16Sample(audio[i]);
205 core::EncodeFixed16(&data[i * kBytesPerSample],
206 static_cast<uint16>(sample));
207 }
208 return Status::OK();
209 }
210
DecodeLin16WaveAsFloatVector(const string & wav_string,std::vector<float> * float_values,uint32 * sample_count,uint16 * channel_count,uint32 * sample_rate)211 Status DecodeLin16WaveAsFloatVector(const string& wav_string,
212 std::vector<float>* float_values,
213 uint32* sample_count, uint16* channel_count,
214 uint32* sample_rate) {
215 int offset = 0;
216 TF_RETURN_IF_ERROR(ExpectText(wav_string, kRiffChunkId, &offset));
217 uint32 total_file_size;
218 TF_RETURN_IF_ERROR(ReadValue<uint32>(wav_string, &total_file_size, &offset));
219 TF_RETURN_IF_ERROR(ExpectText(wav_string, kRiffType, &offset));
220 TF_RETURN_IF_ERROR(ExpectText(wav_string, kFormatChunkId, &offset));
221 uint32 format_chunk_size;
222 TF_RETURN_IF_ERROR(
223 ReadValue<uint32>(wav_string, &format_chunk_size, &offset));
224 if ((format_chunk_size != 16) && (format_chunk_size != 18)) {
225 return errors::InvalidArgument(
226 "Bad file size for WAV: Expected 16 or 18, but got", format_chunk_size);
227 }
228 uint16 audio_format;
229 TF_RETURN_IF_ERROR(ReadValue<uint16>(wav_string, &audio_format, &offset));
230 if (audio_format != 1) {
231 return errors::InvalidArgument(
232 "Bad audio format for WAV: Expected 1 (PCM), but got", audio_format);
233 }
234 TF_RETURN_IF_ERROR(ReadValue<uint16>(wav_string, channel_count, &offset));
235 if (*channel_count < 1) {
236 return errors::InvalidArgument(
237 "Bad number of channels for WAV: Expected at least 1, but got ",
238 *channel_count);
239 }
240 TF_RETURN_IF_ERROR(ReadValue<uint32>(wav_string, sample_rate, &offset));
241 uint32 bytes_per_second;
242 TF_RETURN_IF_ERROR(ReadValue<uint32>(wav_string, &bytes_per_second, &offset));
243 uint16 bytes_per_sample;
244 TF_RETURN_IF_ERROR(ReadValue<uint16>(wav_string, &bytes_per_sample, &offset));
245 // Confusingly, bits per sample is defined as holding the number of bits for
246 // one channel, unlike the definition of sample used elsewhere in the WAV
247 // spec. For example, bytes per sample is the memory needed for all channels
248 // for one point in time.
249 uint16 bits_per_sample;
250 TF_RETURN_IF_ERROR(ReadValue<uint16>(wav_string, &bits_per_sample, &offset));
251 if (bits_per_sample != 16) {
252 return errors::InvalidArgument(
253 "Can only read 16-bit WAV files, but received ", bits_per_sample);
254 }
255 const uint32 expected_bytes_per_sample =
256 ((bits_per_sample * *channel_count) + 7) / 8;
257 if (bytes_per_sample != expected_bytes_per_sample) {
258 return errors::InvalidArgument(
259 "Bad bytes per sample in WAV header: Expected ",
260 expected_bytes_per_sample, " but got ", bytes_per_sample);
261 }
262 const uint32 expected_bytes_per_second = bytes_per_sample * *sample_rate;
263 if (bytes_per_second != expected_bytes_per_second) {
264 return errors::InvalidArgument(
265 "Bad bytes per second in WAV header: Expected ",
266 expected_bytes_per_second, " but got ", bytes_per_second,
267 " (sample_rate=", *sample_rate, ", bytes_per_sample=", bytes_per_sample,
268 ")");
269 }
270 if (format_chunk_size == 18) {
271 // Skip over this unused section.
272 offset += 2;
273 }
274
275 bool was_data_found = false;
276 while (offset < wav_string.size()) {
277 string chunk_id;
278 TF_RETURN_IF_ERROR(ReadString(wav_string, 4, &chunk_id, &offset));
279 uint32 chunk_size;
280 TF_RETURN_IF_ERROR(ReadValue<uint32>(wav_string, &chunk_size, &offset));
281 if (chunk_size > std::numeric_limits<int32>::max()) {
282 return errors::InvalidArgument(
283 "WAV data chunk '", chunk_id, "' is too large: ", chunk_size,
284 " bytes, but the limit is ", std::numeric_limits<int32>::max());
285 }
286 if (chunk_id == kDataChunkId) {
287 if (was_data_found) {
288 return errors::InvalidArgument("More than one data chunk found in WAV");
289 }
290 was_data_found = true;
291 *sample_count = chunk_size / bytes_per_sample;
292 const uint32 data_count = *sample_count * *channel_count;
293 int unused_new_offset = 0;
294 // Validate that the data exists before allocating space for it
295 // (prevent easy OOM errors).
296 TF_RETURN_IF_ERROR(IncrementOffset(offset, sizeof(int16) * data_count,
297 wav_string.size(),
298 &unused_new_offset));
299 float_values->resize(data_count);
300 for (int i = 0; i < data_count; ++i) {
301 int16 single_channel_value = 0;
302 TF_RETURN_IF_ERROR(
303 ReadValue<int16>(wav_string, &single_channel_value, &offset));
304 (*float_values)[i] = Int16SampleToFloat(single_channel_value);
305 }
306 } else {
307 offset += chunk_size;
308 }
309 }
310 if (!was_data_found) {
311 return errors::InvalidArgument("No data chunk found in WAV");
312 }
313 return Status::OK();
314 }
315
316 } // namespace wav
317 } // namespace tensorflow
318