Skip to content

Commit

Permalink
ref ggerganov#17 : add options to output result to file
Browse files Browse the repository at this point in the history 
Support for:

- plain text
- VTT
- SRT
  • Loading branch information
@ggerganov
ggerganov committed Oct 8, 2022
1 parent 4c4ab71 commit 8c7c018
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 8 deletions.
98 changes: 91 additions & 7 deletions main.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#define DR_WAV_IMPLEMENTATION
#include "dr_wav.h"

#include <fstream>
#include <cstdio>
#include <string>
#include <thread>
Expand Down Expand Up @@ -32,6 +33,9 @@ struct whisper_params {

bool verbose = false;
bool translate = false;
bool output_txt = false;
bool output_vtt = false;
bool output_srt = false;
bool print_special_tokens = false;
bool no_timestamps = false;

Expand Down Expand Up @@ -69,6 +73,12 @@ bool whisper_params_parse(int argc, char ** argv, whisper_params & params) {
whisper_print_usage(argc, argv, params);
exit(0);
}
} else if (arg == "-otxt" || arg == "--output-txt") {
params.output_txt = true;
} else if (arg == "-ovtt" || arg == "--output-vtt") {
params.output_vtt = true;
} else if (arg == "-osrt" || arg == "--output-srt") {
params.output_srt = true;
} else if (arg == "-ps" || arg == "--print_special") {
params.print_special_tokens = true;
} else if (arg == "-nt" || arg == "--no_timestamps") {
Expand Down Expand Up @@ -101,6 +111,8 @@ void whisper_print_usage(int argc, char ** argv, const whisper_params & params)
fprintf(stderr, " -o N, --offset N offset in milliseconds (default: %d)\n", params.offset_ms);
fprintf(stderr, " -v, --verbose verbose output\n");
fprintf(stderr, " --translate translate from source language to english\n");
fprintf(stderr, " -otxt, --output-txt output result in a text file\n");
fprintf(stderr, " -ovtt, --output-vtt output result in a vtt file\n");
fprintf(stderr, " -ps, --print_special print special tokens\n");
fprintf(stderr, " -nt, --no_timestamps do not print timestamps\n");
fprintf(stderr, " -l LANG, --language LANG spoken language (default: %s)\n", params.language.c_str());
Expand All @@ -123,7 +135,7 @@ int main(int argc, char ** argv) {
if (params.fname_inp.empty()) {
fprintf(stderr, "error: no input files specified\n");
whisper_print_usage(argc, argv, params);
return 1;
return 2;
}

// whisper init
Expand All @@ -140,22 +152,22 @@ int main(int argc, char ** argv) {
if (!drwav_init_file(&wav, fname_inp.c_str(), NULL)) {
fprintf(stderr, "%s: failed to open WAV file '%s' - check your input\n", argv[0], fname_inp.c_str());
whisper_print_usage(argc, argv, {});
return 2;
return 3;
}

if (wav.channels != 1 && wav.channels != 2) {
fprintf(stderr, "%s: WAV file '%s' must be mono or stereo\n", argv[0], fname_inp.c_str());
return 3;
return 4;
}

if (wav.sampleRate != WHISPER_SAMPLE_RATE) {
fprintf(stderr, "%s: WAV file '%s' must be 16 kHz\n", argv[0], fname_inp.c_str());
return 4;
return 5;
}

if (wav.bitsPerSample != 16) {
fprintf(stderr, "%s: WAV file '%s' must be 16-bit\n", argv[0], fname_inp.c_str());
return 5;
return 6;
}

int n = wav.totalPCMFrameCount;
Expand Down Expand Up @@ -193,9 +205,11 @@ int main(int argc, char ** argv) {
params.language.c_str(),
params.translate ? "translate" : "transcribe",
params.no_timestamps ? 0 : 1);

printf("\n");
}


// run the inference
{
whisper_full_params wparams = whisper_full_default_params(WHISPER_DECODE_GREEDY);
Expand All @@ -211,10 +225,10 @@ int main(int argc, char ** argv) {

if (whisper_full(ctx, wparams, pcmf32.data(), pcmf32.size()) != 0) {
fprintf(stderr, "%s: failed to process audio\n", argv[0]);
return 6;
return 7;
}

// print result;
// print result
if (!wparams.print_realtime) {
printf("\n");

Expand All @@ -233,6 +247,76 @@ int main(int argc, char ** argv) {
}
}
}

printf("\n");

// output to text file
if (params.output_txt) {

const auto fname_txt = fname_inp + ".txt";
std::ofstream fout_txt(fname_txt);
if (!fout_txt.is_open()) {
fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_txt.c_str());
return 8;
}

printf("%s: saving output to '%s.txt'\n", __func__, fname_inp.c_str());

const int n_segments = whisper_full_n_segments(ctx);
for (int i = 0; i < n_segments; ++i) {
const char * text = whisper_full_get_segment_text(ctx, i);
fout_txt << text;
}
}

// output to VTT file
if (params.output_vtt) {

const auto fname_vtt = fname_inp + ".vtt";
std::ofstream fout_vtt(fname_vtt);
if (!fout_vtt.is_open()) {
fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_vtt.c_str());
return 9;
}

printf("%s: saving output to '%s.vtt'\n", __func__, fname_inp.c_str());

fout_vtt << "WEBVTT\n\n";

const int n_segments = whisper_full_n_segments(ctx);
for (int i = 0; i < n_segments; ++i) {
const char * text = whisper_full_get_segment_text(ctx, i);
const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
const int64_t t1 = whisper_full_get_segment_t1(ctx, i);

fout_vtt << to_timestamp(t0) << " --> " << to_timestamp(t1) << "\n";
fout_vtt << text << "\n\n";
}
}

// output to SRT file
if (params.output_srt) {

const auto fname_srt = fname_inp + ".srt";
std::ofstream fout_srt(fname_srt);
if (!fout_srt.is_open()) {
fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_srt.c_str());
return 10;
}

printf("%s: saving output to '%s.srt'\n", __func__, fname_inp.c_str());

const int n_segments = whisper_full_n_segments(ctx);
for (int i = 0; i < n_segments; ++i) {
const char * text = whisper_full_get_segment_text(ctx, i);
const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
const int64_t t1 = whisper_full_get_segment_t1(ctx, i);

fout_srt << i + 1 << "\n";
fout_srt << to_timestamp(t0) << " --> " << to_timestamp(t1) << "\n";
fout_srt << text << "\n\n";
}
}
}
}

Expand Down
2 changes: 1 addition & 1 deletion whisper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2242,7 +2242,7 @@ whisper_token whisper_token_transcribe() {
void whisper_print_timings(struct whisper_context * ctx) {
const int64_t t_end_us = ggml_time_us();

printf("\n\n");
printf("\n");
printf("%s: load time = %8.2f ms\n", __func__, ctx->t_load_us/1000.0f);
printf("%s: mel time = %8.2f ms\n", __func__, ctx->t_mel_us/1000.0f);
printf("%s: sample time = %8.2f ms\n", __func__, ctx->t_sample_us/1000.0f);
Expand Down

0 comments on commit 8c7c018

Please sign in to comment.