1 /*
2 * Copyright (c) 2019 Paul B Mahol
3 *
4 * This file is part of FFmpeg.
5 *
6 * FFmpeg is free software; you can redistribute it and/or
7 * modify it under the terms of the GNU Lesser General Public
8 * License as published by the Free Software Foundation; either
9 * version 2.1 of the License, or (at your option) any later version.
10 *
11 * FFmpeg is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14 * Lesser General Public License for more details.
15 *
16 * You should have received a copy of the GNU Lesser General Public
17 * License along with FFmpeg; if not, write to the Free Software
18 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19 */
20
21 #include <pocketsphinx/pocketsphinx.h>
22
23 #include "libavutil/avassert.h"
24 #include "libavutil/avstring.h"
25 #include "libavutil/channel_layout.h"
26 #include "libavutil/opt.h"
27 #include "audio.h"
28 #include "avfilter.h"
29 #include "internal.h"
30
31 typedef struct ASRContext {
32 const AVClass *class;
33
34 int rate;
35 char *hmm;
36 char *dict;
37 char *lm;
38 char *lmctl;
39 char *lmname;
40 char *logfn;
41
42 ps_decoder_t *ps;
43 cmd_ln_t *config;
44
45 int utt_started;
46 } ASRContext;
47
48 #define OFFSET(x) offsetof(ASRContext, x)
49 #define FLAGS AV_OPT_FLAG_AUDIO_PARAM | AV_OPT_FLAG_FILTERING_PARAM
50 static const AVOption asr_options[] = {
51 { "rate", "set sampling rate", OFFSET(rate), AV_OPT_TYPE_INT, {.i64=16000}, 0, INT_MAX, .flags = FLAGS },
52 { "hmm", "set directory containing acoustic model files", OFFSET(hmm), AV_OPT_TYPE_STRING, {.str=NULL}, .flags = FLAGS },
53 { "dict", "set pronunciation dictionary", OFFSET(dict), AV_OPT_TYPE_STRING, {.str=NULL}, .flags = FLAGS },
54 { "lm", "set language model file", OFFSET(lm), AV_OPT_TYPE_STRING, {.str=NULL}, .flags = FLAGS },
55 { "lmctl", "set language model set", OFFSET(lmctl), AV_OPT_TYPE_STRING, {.str=NULL}, .flags = FLAGS },
56 { "lmname","set which language model to use", OFFSET(lmname), AV_OPT_TYPE_STRING, {.str=NULL}, .flags = FLAGS },
57 { "logfn", "set output for log messages", OFFSET(logfn), AV_OPT_TYPE_STRING, {.str="/dev/null"}, .flags = FLAGS },
58 { NULL }
59 };
60
61 AVFILTER_DEFINE_CLASS(asr);
62
filter_frame(AVFilterLink * inlink,AVFrame * in)63 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
64 {
65 AVFilterContext *ctx = inlink->dst;
66 AVDictionary **metadata = &in->metadata;
67 ASRContext *s = ctx->priv;
68 int have_speech;
69 const char *speech;
70
71 ps_process_raw(s->ps, (const int16_t *)in->data[0], in->nb_samples, 0, 0);
72 have_speech = ps_get_in_speech(s->ps);
73 if (have_speech && !s->utt_started)
74 s->utt_started = 1;
75 if (!have_speech && s->utt_started) {
76 ps_end_utt(s->ps);
77 speech = ps_get_hyp(s->ps, NULL);
78 if (speech != NULL)
79 av_dict_set(metadata, "lavfi.asr.text", speech, 0);
80 ps_start_utt(s->ps);
81 s->utt_started = 0;
82 }
83
84 return ff_filter_frame(ctx->outputs[0], in);
85 }
86
config_input(AVFilterLink * inlink)87 static int config_input(AVFilterLink *inlink)
88 {
89 AVFilterContext *ctx = inlink->dst;
90 ASRContext *s = ctx->priv;
91
92 ps_start_utt(s->ps);
93
94 return 0;
95 }
96
asr_init(AVFilterContext * ctx)97 static av_cold int asr_init(AVFilterContext *ctx)
98 {
99 ASRContext *s = ctx->priv;
100 const float frate = s->rate;
101 char *rate = av_asprintf("%f", frate);
102 const char *argv[] = { "-logfn", s->logfn,
103 "-hmm", s->hmm,
104 "-lm", s->lm,
105 "-lmctl", s->lmctl,
106 "-lmname", s->lmname,
107 "-dict", s->dict,
108 "-samprate", rate,
109 NULL };
110
111 s->config = cmd_ln_parse_r(NULL, ps_args(), 14, (char **)argv, 0);
112 av_free(rate);
113 if (!s->config)
114 return AVERROR(ENOMEM);
115
116 ps_default_search_args(s->config);
117 s->ps = ps_init(s->config);
118 if (!s->ps)
119 return AVERROR(ENOMEM);
120
121 return 0;
122 }
123
query_formats(AVFilterContext * ctx)124 static int query_formats(AVFilterContext *ctx)
125 {
126 ASRContext *s = ctx->priv;
127 int sample_rates[] = { s->rate, -1 };
128 int ret;
129
130 AVFilterFormats *formats = NULL;
131 AVFilterChannelLayouts *layout = NULL;
132
133 if ((ret = ff_add_format (&formats, AV_SAMPLE_FMT_S16 )) < 0 ||
134 (ret = ff_set_common_formats (ctx , formats )) < 0 ||
135 (ret = ff_add_channel_layout (&layout , AV_CH_LAYOUT_MONO )) < 0 ||
136 (ret = ff_set_common_channel_layouts (ctx , layout )) < 0 ||
137 (ret = ff_set_common_samplerates (ctx , ff_make_format_list(sample_rates) )) < 0)
138 return ret;
139
140 return 0;
141 }
142
asr_uninit(AVFilterContext * ctx)143 static av_cold void asr_uninit(AVFilterContext *ctx)
144 {
145 ASRContext *s = ctx->priv;
146
147 ps_free(s->ps);
148 s->ps = NULL;
149 cmd_ln_free_r(s->config);
150 s->config = NULL;
151 }
152
153 static const AVFilterPad asr_inputs[] = {
154 {
155 .name = "default",
156 .type = AVMEDIA_TYPE_AUDIO,
157 .filter_frame = filter_frame,
158 .config_props = config_input,
159 },
160 { NULL }
161 };
162
163 static const AVFilterPad asr_outputs[] = {
164 {
165 .name = "default",
166 .type = AVMEDIA_TYPE_AUDIO,
167 },
168 { NULL }
169 };
170
171 AVFilter ff_af_asr = {
172 .name = "asr",
173 .description = NULL_IF_CONFIG_SMALL("Automatic Speech Recognition."),
174 .priv_size = sizeof(ASRContext),
175 .priv_class = &asr_class,
176 .init = asr_init,
177 .uninit = asr_uninit,
178 .query_formats = query_formats,
179 .inputs = asr_inputs,
180 .outputs = asr_outputs,
181 };
182