xref: /aosp_15_r20/external/webrtc/modules/audio_processing/agc2/rnn_vad/rnn_vad_tool.cc (revision d9f758449e529ab9291ac668be2861e7a55c2422)
1 /*
2  *  Copyright (c) 2018 The WebRTC project authors. All Rights Reserved.
3  *
4  *  Use of this source code is governed by a BSD-style license
5  *  that can be found in the LICENSE file in the root of the source
6  *  tree. An additional intellectual property rights grant can be found
7  *  in the file PATENTS.  All contributing project authors may
8  *  be found in the AUTHORS file in the root of the source tree.
9  */
10 
11 #include <array>
12 #include <string>
13 #include <vector>
14 
15 #include "absl/flags/flag.h"
16 #include "absl/flags/parse.h"
17 #include "common_audio/resampler/push_sinc_resampler.h"
18 #include "common_audio/wav_file.h"
19 #include "modules/audio_processing/agc2/cpu_features.h"
20 #include "modules/audio_processing/agc2/rnn_vad/common.h"
21 #include "modules/audio_processing/agc2/rnn_vad/features_extraction.h"
22 #include "modules/audio_processing/agc2/rnn_vad/rnn.h"
23 #include "rtc_base/logging.h"
24 #include "rtc_base/numerics/safe_compare.h"
25 
26 ABSL_FLAG(std::string, i, "", "Path to the input wav file");
27 ABSL_FLAG(std::string, f, "", "Path to the output features file");
28 ABSL_FLAG(std::string, o, "", "Path to the output VAD probabilities file");
29 
30 namespace webrtc {
31 namespace rnn_vad {
32 namespace test {
33 
main(int argc,char * argv[])34 int main(int argc, char* argv[]) {
35   absl::ParseCommandLine(argc, argv);
36   rtc::LogMessage::LogToDebug(rtc::LS_INFO);
37 
38   // Open wav input file and check properties.
39   const std::string input_wav_file = absl::GetFlag(FLAGS_i);
40   WavReader wav_reader(input_wav_file);
41   if (wav_reader.num_channels() != 1) {
42     RTC_LOG(LS_ERROR) << "Only mono wav files are supported";
43     return 1;
44   }
45   if (wav_reader.sample_rate() % 100 != 0) {
46     RTC_LOG(LS_ERROR) << "The sample rate rate must allow 10 ms frames.";
47     return 1;
48   }
49   RTC_LOG(LS_INFO) << "Input sample rate: " << wav_reader.sample_rate();
50 
51   // Init output files.
52   const std::string output_vad_probs_file = absl::GetFlag(FLAGS_o);
53   FILE* vad_probs_file = fopen(output_vad_probs_file.c_str(), "wb");
54   FILE* features_file = nullptr;
55   const std::string output_feature_file = absl::GetFlag(FLAGS_f);
56   if (!output_feature_file.empty()) {
57     features_file = fopen(output_feature_file.c_str(), "wb");
58   }
59 
60   // Initialize.
61   const int frame_size_10ms =
62       rtc::CheckedDivExact(wav_reader.sample_rate(), 100);
63   std::vector<float> samples_10ms;
64   samples_10ms.resize(frame_size_10ms);
65   std::array<float, kFrameSize10ms24kHz> samples_10ms_24kHz;
66   PushSincResampler resampler(frame_size_10ms, kFrameSize10ms24kHz);
67   const AvailableCpuFeatures cpu_features = GetAvailableCpuFeatures();
68   FeaturesExtractor features_extractor(cpu_features);
69   std::array<float, kFeatureVectorSize> feature_vector;
70   RnnVad rnn_vad(cpu_features);
71 
72   // Compute VAD probabilities.
73   while (true) {
74     // Read frame at the input sample rate.
75     const size_t read_samples =
76         wav_reader.ReadSamples(frame_size_10ms, samples_10ms.data());
77     if (rtc::SafeLt(read_samples, frame_size_10ms)) {
78       break;  // EOF.
79     }
80     // Resample input.
81     resampler.Resample(samples_10ms.data(), samples_10ms.size(),
82                        samples_10ms_24kHz.data(), samples_10ms_24kHz.size());
83 
84     // Extract features and feed the RNN.
85     bool is_silence = features_extractor.CheckSilenceComputeFeatures(
86         samples_10ms_24kHz, feature_vector);
87     float vad_probability =
88         rnn_vad.ComputeVadProbability(feature_vector, is_silence);
89     // Write voice probability.
90     RTC_DCHECK_GE(vad_probability, 0.f);
91     RTC_DCHECK_GE(1.f, vad_probability);
92     fwrite(&vad_probability, sizeof(float), 1, vad_probs_file);
93     // Write features.
94     if (features_file) {
95       const float float_is_silence = is_silence ? 1.f : 0.f;
96       fwrite(&float_is_silence, sizeof(float), 1, features_file);
97       if (is_silence) {
98         // Do not write uninitialized values.
99         feature_vector.fill(0.f);
100       }
101       fwrite(feature_vector.data(), sizeof(float), kFeatureVectorSize,
102              features_file);
103     }
104   }
105 
106   // Close output file(s).
107   fclose(vad_probs_file);
108   RTC_LOG(LS_INFO) << "VAD probabilities written to " << output_vad_probs_file;
109   if (features_file) {
110     fclose(features_file);
111     RTC_LOG(LS_INFO) << "features written to " << output_feature_file;
112   }
113 
114   return 0;
115 }
116 
117 }  // namespace test
118 }  // namespace rnn_vad
119 }  // namespace webrtc
120 
main(int argc,char * argv[])121 int main(int argc, char* argv[]) {
122   return webrtc::rnn_vad::test::main(argc, argv);
123 }
124