1939cb86b4df708f1cb2f68fe4af8e7720c5c666
[ffmpeg.git] / libavcodec / dvdsubenc.c
1 /*
2  * DVD subtitle encoding
3  * Copyright (c) 2005 Wolfram Gloger
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with FFmpeg; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21 #include "avcodec.h"
22 #include "bytestream.h"
23 #include "libavutil/avassert.h"
24 #include "libavutil/bprint.h"
25 #include "libavutil/imgutils.h"
26
27 typedef struct {
28     uint32_t global_palette[16];
29 } DVDSubtitleContext;
30
31 // ncnt is the nibble counter
32 #define PUTNIBBLE(val)\
33 do {\
34     if (ncnt++ & 1)\
35         *q++ = bitbuf | ((val) & 0x0f);\
36     else\
37         bitbuf = (val) << 4;\
38 } while(0)
39
40 static void dvd_encode_rle(uint8_t **pq,
41                            const uint8_t *bitmap, int linesize,
42                            int w, int h,
43                            const int cmap[256])
44 {
45     uint8_t *q;
46     unsigned int bitbuf = 0;
47     int ncnt;
48     int x, y, len, color;
49
50     q = *pq;
51
52     for (y = 0; y < h; ++y) {
53         ncnt = 0;
54         for(x = 0; x < w; x += len) {
55             color = bitmap[x];
56             for (len=1; x+len < w; ++len)
57                 if (bitmap[x+len] != color)
58                     break;
59             color = cmap[color];
60             av_assert0(color < 4);
61             if (len < 0x04) {
62                 PUTNIBBLE((len << 2)|color);
63             } else if (len < 0x10) {
64                 PUTNIBBLE(len >> 2);
65                 PUTNIBBLE((len << 2)|color);
66             } else if (len < 0x40) {
67                 PUTNIBBLE(0);
68                 PUTNIBBLE(len >> 2);
69                 PUTNIBBLE((len << 2)|color);
70             } else if (x+len == w) {
71                 PUTNIBBLE(0);
72                 PUTNIBBLE(0);
73                 PUTNIBBLE(0);
74                 PUTNIBBLE(color);
75             } else {
76                 if (len > 0xff)
77                     len = 0xff;
78                 PUTNIBBLE(0);
79                 PUTNIBBLE(len >> 6);
80                 PUTNIBBLE(len >> 2);
81                 PUTNIBBLE((len << 2)|color);
82             }
83         }
84         /* end of line */
85         if (ncnt & 1)
86             PUTNIBBLE(0);
87         bitmap += linesize;
88     }
89
90     *pq = q;
91 }
92
93 static int color_distance(uint32_t a, uint32_t b)
94 {
95     int r = 0, d, i;
96
97     for (i = 0; i < 32; i += 8) {
98         d = ((a >> i) & 0xFF) - ((b >> i) & 0xFF);
99         r += d * d;
100     }
101     return r;
102 }
103
104 /**
105  * Count colors used in a rectangle, quantizing alpha and grouping by
106  * nearest global palette entry.
107  */
108 static void count_colors(AVCodecContext *avctx, unsigned hits[33],
109                          const AVSubtitleRect *r)
110 {
111     DVDSubtitleContext *dvdc = avctx->priv_data;
112     unsigned count[256] = { 0 };
113     uint32_t *palette = (uint32_t *)r->pict.data[1];
114     uint32_t color;
115     int x, y, i, j, match, d, best_d, av_uninit(best_j);
116     uint8_t *p = r->pict.data[0];
117
118     for (y = 0; y < r->h; y++) {
119         for (x = 0; x < r->w; x++)
120             count[*(p++)]++;
121         p += r->pict.linesize[0] - r->w;
122     }
123     for (i = 0; i < 256; i++) {
124         if (!count[i]) /* avoid useless search */
125             continue;
126         color = palette[i];
127         /* 0: transparent, 1-16: semi-transparent, 17-33 opaque */
128         match = color < 0x33000000 ? 0 : color < 0xCC000000 ? 1 : 17;
129         if (match) {
130             best_d = INT_MAX;
131             for (j = 0; j < 16; j++) {
132                 d = color_distance(color & 0xFFFFFF, dvdc->global_palette[j]);
133                 if (d < best_d) {
134                     best_d = d;
135                     best_j = j;
136                 }
137             }
138             match += best_j;
139         }
140         hits[match] += count[i];
141     }
142 }
143
144 static void select_palette(AVCodecContext *avctx, int out_palette[4],
145                            int out_alpha[4], unsigned hits[33])
146 {
147     DVDSubtitleContext *dvdc = avctx->priv_data;
148     int i, j, bright, mult;
149     uint32_t color;
150     int selected[4] = { 0 };
151     uint32_t pseudopal[33] = { 0 };
152     uint32_t refcolor[3] = { 0x00000000, 0xFFFFFFFF, 0xFF000000 };
153
154     /* Bonus for transparent: if the rectangle fits tightly the text, the
155        background color can be quite rare, but it would be ugly without it */
156     hits[0] *= 16;
157     /* Bonus for bright colors */
158     for (i = 0; i < 16; i++) {
159         if (!(hits[1 + i] + hits[17 + i]))
160             continue; /* skip unused colors to gain time */
161         color = dvdc->global_palette[i];
162         bright = 0;
163         for (j = 0; j < 3; j++, color >>= 8)
164             bright += (color & 0xFF) < 0x40 || (color & 0xFF) >= 0xC0;
165         mult = 2 + FFMIN(bright, 2);
166         hits[ 1 + i] *= mult;
167         hits[17 + i] *= mult;
168     }
169
170     /* Select four most frequent colors */
171     for (i = 0; i < 4; i++) {
172         for (j = 0; j < 33; j++)
173             if (hits[j] > hits[selected[i]])
174                 selected[i] = j;
175         hits[selected[i]] = 0;
176     }
177
178     /* Order the colors like in most DVDs:
179        0: background, 1: foreground, 2: outline */
180     for (i = 0; i < 16; i++) {
181         pseudopal[ 1 + i] = 0x80000000 | dvdc->global_palette[i];
182         pseudopal[17 + i] = 0xFF000000 | dvdc->global_palette[i];
183     }
184     for (i = 0; i < 3; i++) {
185         int best_d = color_distance(refcolor[i], pseudopal[selected[i]]);
186         for (j = i + 1; j < 4; j++) {
187             int d = color_distance(refcolor[i], pseudopal[selected[j]]);
188             if (d < best_d) {
189                 FFSWAP(int, selected[i], selected[j]);
190                 best_d = d;
191             }
192         }
193     }
194
195     /* Output */
196     for (i = 0; i < 4; i++) {
197         out_palette[i] = selected[i] ? (selected[i] - 1) & 0xF : 0;
198         out_alpha  [i] = !selected[i] ? 0 : selected[i] < 17 ? 0x80 : 0xFF;
199     }
200 }
201
202 static void build_color_map(AVCodecContext *avctx, int cmap[],
203                             const uint32_t palette[],
204                             const int out_palette[], int const out_alpha[])
205 {
206     DVDSubtitleContext *dvdc = avctx->priv_data;
207     int i, j, d, best_d;
208     uint32_t pseudopal[4];
209
210     for (i = 0; i < 4; i++)
211         pseudopal[i] = (out_alpha[i] << 24) |
212                        dvdc->global_palette[out_palette[i]];
213     for (i = 0; i < 256; i++) {
214         best_d = INT_MAX;
215         for (j = 0; j < 4; j++) {
216             d = color_distance(pseudopal[j], palette[i]);
217             if (d < best_d) {
218                 cmap[i] = j;
219                 best_d = d;
220             }
221         }
222     }
223 }
224
225 static void copy_rectangle(AVSubtitleRect *dst, AVSubtitleRect *src, int cmap[])
226 {
227     int x, y;
228     uint8_t *p, *q;
229
230     p = src->pict.data[0];
231     q = dst->pict.data[0] + (src->x - dst->x) +
232                             (src->y - dst->y) * dst->pict.linesize[0];
233     for (y = 0; y < src->h; y++) {
234         for (x = 0; x < src->w; x++)
235             *(q++) = cmap[*(p++)];
236         p += src->pict.linesize[0] - src->w;
237         q += dst->pict.linesize[0] - src->w;
238     }
239 }
240
241 static int encode_dvd_subtitles(AVCodecContext *avctx,
242                                 uint8_t *outbuf, int outbuf_size,
243                                 const AVSubtitle *h)
244 {
245     DVDSubtitleContext *dvdc = avctx->priv_data;
246     uint8_t *q, *qq;
247     int offset1, offset2;
248     int i, rects = h->num_rects, ret;
249     unsigned global_palette_hits[33] = { 0 };
250     int cmap[256];
251     int out_palette[4];
252     int out_alpha[4];
253     AVSubtitleRect vrect;
254     uint8_t *vrect_data = NULL;
255     int x2, y2;
256
257     if (rects == 0 || h->rects == NULL)
258         return AVERROR(EINVAL);
259     vrect = *h->rects[0];
260
261     if (rects > 1) {
262         /* DVD subtitles can have only one rectangle: build a virtual
263            rectangle containing all actual rectangles.
264            The data of the rectangles will be copied later, when the palette
265            is decided, because the rectangles may have different palettes. */
266         int xmin = h->rects[0]->x, xmax = xmin + h->rects[0]->w;
267         int ymin = h->rects[0]->y, ymax = ymin + h->rects[0]->h;
268         for (i = 1; i < rects; i++) {
269             xmin = FFMIN(xmin, h->rects[i]->x);
270             ymin = FFMIN(ymin, h->rects[i]->y);
271             xmax = FFMAX(xmax, h->rects[i]->x + h->rects[i]->w);
272             ymax = FFMAX(ymax, h->rects[i]->y + h->rects[i]->h);
273         }
274         vrect.x = xmin;
275         vrect.y = ymin;
276         vrect.w = xmax - xmin;
277         vrect.h = ymax - ymin;
278         if ((ret = av_image_check_size(vrect.w, vrect.h, 0, avctx)) < 0)
279             return ret;
280
281         /* Count pixels outside the virtual rectangle as transparent */
282         global_palette_hits[0] = vrect.w * vrect.h;
283         for (i = 0; i < rects; i++)
284             global_palette_hits[0] -= h->rects[i]->w * h->rects[i]->h;
285     }
286
287     for (i = 0; i < rects; i++)
288         count_colors(avctx, global_palette_hits, h->rects[i]);
289     select_palette(avctx, out_palette, out_alpha, global_palette_hits);
290
291     if (rects > 1) {
292         if (!(vrect_data = av_calloc(vrect.w, vrect.h)))
293             return AVERROR(ENOMEM);
294         vrect.pict.data    [0] = vrect_data;
295         vrect.pict.linesize[0] = vrect.w;
296         for (i = 0; i < rects; i++) {
297             build_color_map(avctx, cmap, (uint32_t *)h->rects[i]->pict.data[1],
298                             out_palette, out_alpha);
299             copy_rectangle(&vrect, h->rects[i], cmap);
300         }
301         for (i = 0; i < 4; i++)
302             cmap[i] = i;
303     } else {
304         build_color_map(avctx, cmap, (uint32_t *)h->rects[0]->pict.data[1],
305                         out_palette, out_alpha);
306     }
307
308     av_log(avctx, AV_LOG_DEBUG, "Selected palette:");
309     for (i = 0; i < 4; i++)
310         av_log(avctx, AV_LOG_DEBUG, " 0x%06x@@%02x (0x%x,0x%x)",
311                dvdc->global_palette[out_palette[i]], out_alpha[i],
312                out_palette[i], out_alpha[i] >> 4);
313     av_log(avctx, AV_LOG_DEBUG, "\n");
314
315     // encode data block
316     q = outbuf + 4;
317     /* TODO reindent */
318         offset1 = q - outbuf;
319         // worst case memory requirement: 1 nibble per pixel..
320         if ((q - outbuf) + vrect.w * vrect.h / 2 + 17 + 21 > outbuf_size) {
321             av_log(NULL, AV_LOG_ERROR, "dvd_subtitle too big\n");
322             ret = AVERROR_BUFFER_TOO_SMALL;
323             goto fail;
324         }
325         dvd_encode_rle(&q, vrect.pict.data[0], vrect.w * 2,
326                        vrect.w, (vrect.h + 1) >> 1, cmap);
327         offset2 = q - outbuf;
328         dvd_encode_rle(&q, vrect.pict.data[0] + vrect.w, vrect.w * 2,
329                        vrect.w, vrect.h >> 1, cmap);
330
331     // set data packet size
332     qq = outbuf + 2;
333     bytestream_put_be16(&qq, q - outbuf);
334
335     // send start display command
336     bytestream_put_be16(&q, (h->start_display_time*90) >> 10);
337     bytestream_put_be16(&q, (q - outbuf) /*- 2 */ + 8 + 12 + 2);
338     *q++ = 0x03; // palette - 4 nibbles
339     *q++ = (out_palette[3] << 4) | out_palette[2];
340     *q++ = (out_palette[1] << 4) | out_palette[0];
341     *q++ = 0x04; // alpha - 4 nibbles
342     *q++ = (out_alpha[3] & 0xF0) | (out_alpha[2] >> 4);
343     *q++ = (out_alpha[1] & 0xF0) | (out_alpha[0] >> 4);
344
345     // 12 bytes per rect
346     /* TODO reindent */
347         x2 = vrect.x + vrect.w - 1;
348         y2 = vrect.y + vrect.h - 1;
349
350         *q++ = 0x05;
351         // x1 x2 -> 6 nibbles
352         *q++ = vrect.x >> 4;
353         *q++ = (vrect.x << 4) | ((x2 >> 8) & 0xf);
354         *q++ = x2;
355         // y1 y2 -> 6 nibbles
356         *q++ = vrect.y >> 4;
357         *q++ = (vrect.y << 4) | ((y2 >> 8) & 0xf);
358         *q++ = y2;
359
360         *q++ = 0x06;
361         // offset1, offset2
362         bytestream_put_be16(&q, offset1);
363         bytestream_put_be16(&q, offset2);
364
365     *q++ = 0x01; // start command
366     *q++ = 0xff; // terminating command
367
368     // send stop display command last
369     bytestream_put_be16(&q, (h->end_display_time*90) >> 10);
370     bytestream_put_be16(&q, (q - outbuf) - 2 /*+ 4*/);
371     *q++ = 0x02; // set end
372     *q++ = 0xff; // terminating command
373
374     qq = outbuf;
375     bytestream_put_be16(&qq, q - outbuf);
376
377     av_log(NULL, AV_LOG_DEBUG, "subtitle_packet size=%td\n", q - outbuf);
378     ret = q - outbuf;
379
380 fail:
381     av_free(vrect_data);
382     return ret;
383 }
384
385 static int dvdsub_init(AVCodecContext *avctx)
386 {
387     DVDSubtitleContext *dvdc = avctx->priv_data;
388     static const uint32_t default_palette[16] = {
389         0x000000, 0x0000FF, 0x00FF00, 0xFF0000,
390         0xFFFF00, 0xFF00FF, 0x00FFFF, 0xFFFFFF,
391         0x808000, 0x8080FF, 0x800080, 0x80FF80,
392         0x008080, 0xFF8080, 0x555555, 0xAAAAAA,
393     };
394     AVBPrint extradata;
395     int i, ret;
396
397     av_assert0(sizeof(dvdc->global_palette) == sizeof(default_palette));
398     memcpy(dvdc->global_palette, default_palette, sizeof(dvdc->global_palette));
399
400     av_bprint_init(&extradata, 0, 1);
401     av_bprintf(&extradata, "palette:");
402     for (i = 0; i < 16; i++)
403         av_bprintf(&extradata, " %06"PRIx32"%c",
404                    dvdc->global_palette[i] & 0xFFFFFF, i < 15 ? ',' : '\n');
405
406     if ((ret = av_bprint_finalize(&extradata, (char **)&avctx->extradata)) < 0)
407         return ret;
408     avctx->extradata_size = extradata.len;
409
410     return 0;
411 }
412
413 static int dvdsub_encode(AVCodecContext *avctx,
414                          unsigned char *buf, int buf_size, void *data)
415 {
416     //DVDSubtitleContext *s = avctx->priv_data;
417     AVSubtitle *sub = data;
418     int ret;
419
420     ret = encode_dvd_subtitles(avctx, buf, buf_size, sub);
421     return ret;
422 }
423
424 AVCodec ff_dvdsub_encoder = {
425     .name           = "dvdsub",
426     .type           = AVMEDIA_TYPE_SUBTITLE,
427     .id             = AV_CODEC_ID_DVD_SUBTITLE,
428     .init           = dvdsub_init,
429     .encode         = dvdsub_encode,
430     .long_name      = NULL_IF_CONFIG_SMALL("DVD subtitles"),
431     .priv_data_size = sizeof(DVDSubtitleContext),
432 };