ggerganov commited on
Commit
e004a9e
·
1 Parent(s): 6f2f468

main : refactor subtitle output

Browse files
Files changed (1) hide show
  1. main.cpp +63 -50
main.cpp CHANGED
@@ -141,6 +141,66 @@ void whisper_print_usage(int argc, char ** argv, const whisper_params & params)
141
  fprintf(stderr, "\n");
142
  }
143
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
144
  int main(int argc, char ** argv) {
145
  whisper_params params;
146
 
@@ -291,66 +351,19 @@ int main(int argc, char ** argv) {
291
  // output to text file
292
  if (params.output_txt) {
293
  const auto fname_txt = fname_inp + ".txt";
294
- std::ofstream fout_txt(fname_txt);
295
- if (!fout_txt.is_open()) {
296
- fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_txt.c_str());
297
- return 8;
298
- }
299
-
300
- fprintf(stderr, "%s: saving output to '%s.txt'\n", __func__, fname_inp.c_str());
301
-
302
- const int n_segments = whisper_full_n_segments(ctx);
303
- for (int i = 0; i < n_segments; ++i) {
304
- const char * text = whisper_full_get_segment_text(ctx, i);
305
- fout_txt << text;
306
- }
307
  }
308
 
309
  // output to VTT file
310
  if (params.output_vtt) {
311
  const auto fname_vtt = fname_inp + ".vtt";
312
- std::ofstream fout_vtt(fname_vtt);
313
- if (!fout_vtt.is_open()) {
314
- fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_vtt.c_str());
315
- return 9;
316
- }
317
-
318
- fprintf(stderr, "%s: saving output to '%s.vtt'\n", __func__, fname_inp.c_str());
319
-
320
- fout_vtt << "WEBVTT\n\n";
321
-
322
- const int n_segments = whisper_full_n_segments(ctx);
323
- for (int i = 0; i < n_segments; ++i) {
324
- const char * text = whisper_full_get_segment_text(ctx, i);
325
- const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
326
- const int64_t t1 = whisper_full_get_segment_t1(ctx, i);
327
-
328
- fout_vtt << to_timestamp(t0) << " --> " << to_timestamp(t1) << "\n";
329
- fout_vtt << text << "\n\n";
330
- }
331
  }
332
 
333
  // output to SRT file
334
  if (params.output_srt) {
335
  const auto fname_srt = fname_inp + ".srt";
336
- std::ofstream fout_srt(fname_srt);
337
- if (!fout_srt.is_open()) {
338
- fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_srt.c_str());
339
- return 10;
340
- }
341
-
342
- fprintf(stderr, "%s: saving output to '%s.srt'\n", __func__, fname_inp.c_str());
343
-
344
- const int n_segments = whisper_full_n_segments(ctx);
345
- for (int i = 0; i < n_segments; ++i) {
346
- const char * text = whisper_full_get_segment_text(ctx, i);
347
- const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
348
- const int64_t t1 = whisper_full_get_segment_t1(ctx, i);
349
-
350
- fout_srt << i + 1 + params.offset_n << "\n";
351
- fout_srt << to_timestamp(t0) << " --> " << to_timestamp(t1) << "\n";
352
- fout_srt << text << "\n\n";
353
- }
354
  }
355
  }
356
  }
 
141
  fprintf(stderr, "\n");
142
  }
143
 
144
+ bool output_txt(struct whisper_context * ctx, const char * fname) {
145
+ std::ofstream fout(fname);
146
+ if (!fout.is_open()) {
147
+ fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname);
148
+ return false;
149
+ }
150
+
151
+ fprintf(stderr, "%s: saving output to '%s'\n", __func__, fname);
152
+
153
+ const int n_segments = whisper_full_n_segments(ctx);
154
+ for (int i = 0; i < n_segments; ++i) {
155
+ const char * text = whisper_full_get_segment_text(ctx, i);
156
+ fout << text;
157
+ }
158
+
159
+ return true;
160
+ }
161
+
162
+ bool output_vtt(struct whisper_context * ctx, const char * fname) {
163
+ std::ofstream fout(fname);
164
+ if (!fout.is_open()) {
165
+ fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname);
166
+ return 9;
167
+ }
168
+
169
+ fprintf(stderr, "%s: saving output to '%s'\n", __func__, fname);
170
+
171
+ fout << "WEBVTT\n\n";
172
+
173
+ const int n_segments = whisper_full_n_segments(ctx);
174
+ for (int i = 0; i < n_segments; ++i) {
175
+ const char * text = whisper_full_get_segment_text(ctx, i);
176
+ const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
177
+ const int64_t t1 = whisper_full_get_segment_t1(ctx, i);
178
+
179
+ fout << to_timestamp(t0) << " --> " << to_timestamp(t1) << "\n";
180
+ fout << text << "\n\n";
181
+ }
182
+
183
+ return true;
184
+ }
185
+
186
+ bool output_srt(struct whisper_context * ctx, const char * fname) {
187
+ std::ofstream fout(fname);
188
+ if (!fout.is_open()) {
189
+ fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname);
190
+ return false;
191
+ }
192
+
193
+ fprintf(stderr, "%s: saving output to '%s'\n", __func__, fname);
194
+
195
+ const int n_segments = whisper_full_n_segments(ctx);
196
+ for (int i = 0; i < n_segments; ++i) {
197
+ const char * text = whisper_full_get_segment_text(ctx, i);
198
+ fout << text;
199
+ }
200
+
201
+ return true;
202
+ }
203
+
204
  int main(int argc, char ** argv) {
205
  whisper_params params;
206
 
 
351
  // output to text file
352
  if (params.output_txt) {
353
  const auto fname_txt = fname_inp + ".txt";
354
+ output_txt(ctx, fname_txt.c_str());
 
 
 
 
 
 
 
 
 
 
 
 
355
  }
356
 
357
  // output to VTT file
358
  if (params.output_vtt) {
359
  const auto fname_vtt = fname_inp + ".vtt";
360
+ output_vtt(ctx, fname_vtt.c_str());
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
361
  }
362
 
363
  // output to SRT file
364
  if (params.output_srt) {
365
  const auto fname_srt = fname_inp + ".srt";
366
+ output_srt(ctx, fname_srt.c_str());
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
367
  }
368
  }
369
  }