git.sesse.net Git - ffmpeg/blob - libavcodec/webvttenc.c

   1 /*
   2  * WebVTT subtitle encoder
   3  * Copyright (c) 2010  Aurelien Jacobs <aurel@gnuage.org>
   4  * Copyright (c) 2014  Aman Gupta <ffmpeg@tmm1.net>
   5  *
   6  * This file is part of FFmpeg.
   7  *
   8  * FFmpeg is free software; you can redistribute it and/or
   9  * modify it under the terms of the GNU Lesser General Public
  10  * License as published by the Free Software Foundation; either
  11  * version 2.1 of the License, or (at your option) any later version.
  12  *
  13  * FFmpeg is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16  * Lesser General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU Lesser General Public
  19  * License along with FFmpeg; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  21  */
  22
  23 #include <stdarg.h>
  24 #include "avcodec.h"
  25 #include "libavutil/avstring.h"
  26 #include "libavutil/bprint.h"
  27 #include "ass_split.h"
  28 #include "ass.h"
  29
  30 #define WEBVTT_STACK_SIZE 64
  31 typedef struct {
  32     AVCodecContext *avctx;
  33     ASSSplitContext *ass_ctx;
  34     AVBPrint buffer;
  35     unsigned timestamp_end;
  36     int count;
  37     char stack[WEBVTT_STACK_SIZE];
  38     int stack_ptr;
  39 } WebVTTContext;
  40
  41 #ifdef __GNUC__
  42 __attribute__ ((__format__ (__printf__, 2, 3)))
  43 #endif
  44 static void webvtt_print(WebVTTContext *s, const char *str, ...)
  45 {
  46     va_list vargs;
  47     va_start(vargs, str);
  48     av_vbprintf(&s->buffer, str, vargs);
  49     va_end(vargs);
  50 }
  51
  52 static int webvtt_stack_push(WebVTTContext *s, const char c)
  53 {
  54     if (s->stack_ptr >= WEBVTT_STACK_SIZE)
  55         return -1;
  56     s->stack[s->stack_ptr++] = c;
  57     return 0;
  58 }
  59
  60 static char webvtt_stack_pop(WebVTTContext *s)
  61 {
  62     if (s->stack_ptr <= 0)
  63         return 0;
  64     return s->stack[--s->stack_ptr];
  65 }
  66
  67 static int webvtt_stack_find(WebVTTContext *s, const char c)
  68 {
  69     int i;
  70     for (i = s->stack_ptr-1; i >= 0; i--)
  71         if (s->stack[i] == c)
  72             break;
  73     return i;
  74 }
  75
  76 static void webvtt_close_tag(WebVTTContext *s, char tag)
  77 {
  78     webvtt_print(s, "</%c>", tag);
  79 }
  80
  81 static void webvtt_stack_push_pop(WebVTTContext *s, const char c, int close)
  82 {
  83     if (close) {
  84         int i = c ? webvtt_stack_find(s, c) : 0;
  85         if (i < 0)
  86             return;
  87         while (s->stack_ptr != i)
  88             webvtt_close_tag(s, webvtt_stack_pop(s));
  89     } else if (webvtt_stack_push(s, c) < 0)
  90         av_log(s->avctx, AV_LOG_ERROR, "tag stack overflow\n");
  91 }
  92
  93 static void webvtt_style_apply(WebVTTContext *s, const char *style)
  94 {
  95     ASSStyle *st = ff_ass_style_get(s->ass_ctx, style);
  96     if (st) {
  97         if (st->bold != ASS_DEFAULT_BOLD) {
  98             webvtt_print(s, "<b>");
  99             webvtt_stack_push(s, 'b');
 100         }
 101         if (st->italic != ASS_DEFAULT_ITALIC) {
 102             webvtt_print(s, "<i>");
 103             webvtt_stack_push(s, 'i');
 104         }
 105         if (st->underline != ASS_DEFAULT_UNDERLINE) {
 106             webvtt_print(s, "<u>");
 107             webvtt_stack_push(s, 'u');
 108         }
 109     }
 110 }
 111
 112 static void webvtt_text_cb(void *priv, const char *text, int len)
 113 {
 114     WebVTTContext *s = priv;
 115     av_bprint_append_data(&s->buffer, text, len);
 116 }
 117
 118 static void webvtt_new_line_cb(void *priv, int forced)
 119 {
 120     webvtt_print(priv, "\n");
 121 }
 122
 123 static void webvtt_style_cb(void *priv, char style, int close)
 124 {
 125     if (style == 's') // strikethrough unsupported
 126         return;
 127
 128     webvtt_stack_push_pop(priv, style, close);
 129     if (!close)
 130         webvtt_print(priv, "<%c>", style);
 131 }
 132
 133 static void webvtt_cancel_overrides_cb(void *priv, const char *style)
 134 {
 135     webvtt_stack_push_pop(priv, 0, 1);
 136     webvtt_style_apply(priv, style);
 137 }
 138
 139 static void webvtt_end_cb(void *priv)
 140 {
 141     webvtt_stack_push_pop(priv, 0, 1);
 142 }
 143
 144 static const ASSCodesCallbacks webvtt_callbacks = {
 145     .text             = webvtt_text_cb,
 146     .new_line         = webvtt_new_line_cb,
 147     .style            = webvtt_style_cb,
 148     .color            = NULL,
 149     .font_name        = NULL,
 150     .font_size        = NULL,
 151     .alignment        = NULL,
 152     .cancel_overrides = webvtt_cancel_overrides_cb,
 153     .move             = NULL,
 154     .end              = webvtt_end_cb,
 155 };
 156
 157 static int webvtt_encode_frame(AVCodecContext *avctx,
 158                                unsigned char *buf, int bufsize, const AVSubtitle *sub)
 159 {
 160     WebVTTContext *s = avctx->priv_data;
 161     ASSDialog *dialog;
 162     int i;
 163
 164     av_bprint_clear(&s->buffer);
 165
 166     for (i=0; i<sub->num_rects; i++) {
 167         const char *ass = sub->rects[i]->ass;
 168
 169         if (sub->rects[i]->type != SUBTITLE_ASS) {
 170             av_log(avctx, AV_LOG_ERROR, "Only SUBTITLE_ASS type supported.\n");
 171             return AVERROR(EINVAL);
 172         }
 173
 174         dialog = ff_ass_split_dialog2(s->ass_ctx, ass);
 175         if (!dialog)
 176             return AVERROR(ENOMEM);
 177         webvtt_style_apply(s, dialog->style);
 178         ff_ass_split_override_codes(&webvtt_callbacks, s, dialog->text);
 179         ff_ass_free_dialog(&dialog);
 180     }
 181
 182     if (!av_bprint_is_complete(&s->buffer))
 183         return AVERROR(ENOMEM);
 184     if (!s->buffer.len)
 185         return 0;
 186
 187     if (s->buffer.len > bufsize) {
 188         av_log(avctx, AV_LOG_ERROR, "Buffer too small for ASS event.\n");
 189         return AVERROR_BUFFER_TOO_SMALL;
 190     }
 191     memcpy(buf, s->buffer.str, s->buffer.len);
 192
 193     return s->buffer.len;
 194 }
 195
 196 static int webvtt_encode_close(AVCodecContext *avctx)
 197 {
 198     WebVTTContext *s = avctx->priv_data;
 199     ff_ass_split_free(s->ass_ctx);
 200     av_bprint_finalize(&s->buffer, NULL);
 201     return 0;
 202 }
 203
 204 static av_cold int webvtt_encode_init(AVCodecContext *avctx)
 205 {
 206     WebVTTContext *s = avctx->priv_data;
 207     s->avctx = avctx;
 208     s->ass_ctx = ff_ass_split(avctx->subtitle_header);
 209     av_bprint_init(&s->buffer, 0, AV_BPRINT_SIZE_UNLIMITED);
 210     return s->ass_ctx ? 0 : AVERROR_INVALIDDATA;
 211 }
 212
 213 AVCodec ff_webvtt_encoder = {
 214     .name           = "webvtt",
 215     .long_name      = NULL_IF_CONFIG_SMALL("WebVTT subtitle"),
 216     .type           = AVMEDIA_TYPE_SUBTITLE,
 217     .id             = AV_CODEC_ID_WEBVTT,
 218     .priv_data_size = sizeof(WebVTTContext),
 219     .init           = webvtt_encode_init,
 220     .encode_sub     = webvtt_encode_frame,
 221     .close          = webvtt_encode_close,
 222 };