png: Return slightly more meaningful error codes
[ffmpeg.git] / libavcodec / pngenc.c
1 /*
2  * PNG image format
3  * Copyright (c) 2003 Fabrice Bellard
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
22 #include "avcodec.h"
23 #include "internal.h"
24 #include "bytestream.h"
25 #include "huffyuvencdsp.h"
26 #include "png.h"
27
28 #include "libavutil/avassert.h"
29 #include "libavutil/libm.h"
30 #include "libavutil/opt.h"
31 #include "libavutil/color_utils.h"
32
33 #include <zlib.h>
34
35 #define IOBUF_SIZE 4096
36
37 typedef struct PNGEncContext {
38     AVClass *class;
39     HuffYUVEncDSPContext hdsp;
40
41     uint8_t *bytestream;
42     uint8_t *bytestream_start;
43     uint8_t *bytestream_end;
44
45     int filter_type;
46
47     z_stream zstream;
48     uint8_t buf[IOBUF_SIZE];
49     int dpi;                     ///< Physical pixel density, in dots per inch, if set
50     int dpm;                     ///< Physical pixel density, in dots per meter, if set
51 } PNGEncContext;
52
53 static void png_get_interlaced_row(uint8_t *dst, int row_size,
54                                    int bits_per_pixel, int pass,
55                                    const uint8_t *src, int width)
56 {
57     int x, mask, dst_x, j, b, bpp;
58     uint8_t *d;
59     const uint8_t *s;
60     static const int masks[] = {0x80, 0x08, 0x88, 0x22, 0xaa, 0x55, 0xff};
61
62     mask = masks[pass];
63     switch (bits_per_pixel) {
64     case 1:
65         memset(dst, 0, row_size);
66         dst_x = 0;
67         for (x = 0; x < width; x++) {
68             j = (x & 7);
69             if ((mask << j) & 0x80) {
70                 b = (src[x >> 3] >> (7 - j)) & 1;
71                 dst[dst_x >> 3] |= b << (7 - (dst_x & 7));
72                 dst_x++;
73             }
74         }
75         break;
76     default:
77         bpp = bits_per_pixel >> 3;
78         d = dst;
79         s = src;
80         for (x = 0; x < width; x++) {
81             j = x & 7;
82             if ((mask << j) & 0x80) {
83                 memcpy(d, s, bpp);
84                 d += bpp;
85             }
86             s += bpp;
87         }
88         break;
89     }
90 }
91
92 static void sub_png_paeth_prediction(uint8_t *dst, uint8_t *src, uint8_t *top,
93                                      int w, int bpp)
94 {
95     int i;
96     for (i = 0; i < w; i++) {
97         int a, b, c, p, pa, pb, pc;
98
99         a = src[i - bpp];
100         b = top[i];
101         c = top[i - bpp];
102
103         p  = b - c;
104         pc = a - c;
105
106         pa = abs(p);
107         pb = abs(pc);
108         pc = abs(p + pc);
109
110         if (pa <= pb && pa <= pc)
111             p = a;
112         else if (pb <= pc)
113             p = b;
114         else
115             p = c;
116         dst[i] = src[i] - p;
117     }
118 }
119
120 static void sub_left_prediction(PNGEncContext *c, uint8_t *dst, const uint8_t *src, int bpp, int size)
121 {
122     const uint8_t *src1 = src + bpp;
123     const uint8_t *src2 = src;
124     int x, unaligned_w;
125
126     memcpy(dst, src, bpp);
127     dst += bpp;
128     size -= bpp;
129     unaligned_w = FFMIN(32 - bpp, size);
130     for (x = 0; x < unaligned_w; x++)
131         *dst++ = *src1++ - *src2++;
132     size -= unaligned_w;
133     c->hdsp.diff_bytes(dst, src1, src2, size);
134 }
135
136 static void png_filter_row(PNGEncContext *c, uint8_t *dst, int filter_type,
137                            uint8_t *src, uint8_t *top, int size, int bpp)
138 {
139     int i;
140
141     switch (filter_type) {
142     case PNG_FILTER_VALUE_NONE:
143         memcpy(dst, src, size);
144         break;
145     case PNG_FILTER_VALUE_SUB:
146         sub_left_prediction(c, dst, src, bpp, size);
147         break;
148     case PNG_FILTER_VALUE_UP:
149         c->hdsp.diff_bytes(dst, src, top, size);
150         break;
151     case PNG_FILTER_VALUE_AVG:
152         for (i = 0; i < bpp; i++)
153             dst[i] = src[i] - (top[i] >> 1);
154         for (; i < size; i++)
155             dst[i] = src[i] - ((src[i - bpp] + top[i]) >> 1);
156         break;
157     case PNG_FILTER_VALUE_PAETH:
158         for (i = 0; i < bpp; i++)
159             dst[i] = src[i] - top[i];
160         sub_png_paeth_prediction(dst + i, src + i, top + i, size - i, bpp);
161         break;
162     }
163 }
164
165 static uint8_t *png_choose_filter(PNGEncContext *s, uint8_t *dst,
166                                   uint8_t *src, uint8_t *top, int size, int bpp)
167 {
168     int pred = s->filter_type;
169     av_assert0(bpp || !pred);
170     if (!top && pred)
171         pred = PNG_FILTER_VALUE_SUB;
172     if (pred == PNG_FILTER_VALUE_MIXED) {
173         int i;
174         int cost, bcost = INT_MAX;
175         uint8_t *buf1 = dst, *buf2 = dst + size + 16;
176         for (pred = 0; pred < 5; pred++) {
177             png_filter_row(s, buf1 + 1, pred, src, top, size, bpp);
178             buf1[0] = pred;
179             cost = 0;
180             for (i = 0; i <= size; i++)
181                 cost += abs((int8_t) buf1[i]);
182             if (cost < bcost) {
183                 bcost = cost;
184                 FFSWAP(uint8_t *, buf1, buf2);
185             }
186         }
187         return buf2;
188     } else {
189         png_filter_row(s, dst + 1, pred, src, top, size, bpp);
190         dst[0] = pred;
191         return dst;
192     }
193 }
194
195 static void png_write_chunk(uint8_t **f, uint32_t tag,
196                             const uint8_t *buf, int length)
197 {
198     uint32_t crc;
199     uint8_t tagbuf[4];
200
201     bytestream_put_be32(f, length);
202     crc = crc32(0, Z_NULL, 0);
203     AV_WL32(tagbuf, tag);
204     crc = crc32(crc, tagbuf, 4);
205     bytestream_put_be32(f, av_bswap32(tag));
206     if (length > 0) {
207         crc = crc32(crc, buf, length);
208         memcpy(*f, buf, length);
209         *f += length;
210     }
211     bytestream_put_be32(f, crc);
212 }
213
214 /* XXX: do filtering */
215 static int png_write_row(PNGEncContext *s, const uint8_t *data, int size)
216 {
217     int ret;
218
219     s->zstream.avail_in = size;
220     s->zstream.next_in  = data;
221     while (s->zstream.avail_in > 0) {
222         ret = deflate(&s->zstream, Z_NO_FLUSH);
223         if (ret != Z_OK)
224             return -1;
225         if (s->zstream.avail_out == 0) {
226             if (s->bytestream_end - s->bytestream > IOBUF_SIZE + 100)
227                 png_write_chunk(&s->bytestream,
228                                 MKTAG('I', 'D', 'A', 'T'), s->buf, IOBUF_SIZE);
229             s->zstream.avail_out = IOBUF_SIZE;
230             s->zstream.next_out  = s->buf;
231         }
232     }
233     return 0;
234 }
235
236 #define AV_WB32_PNG(buf, n) AV_WB32(buf, lrint((n) * 100000))
237 static int png_get_chrm(enum AVColorPrimaries prim,  uint8_t *buf)
238 {
239     double rx, ry, gx, gy, bx, by, wx = 0.3127, wy = 0.3290;
240     switch (prim) {
241         case AVCOL_PRI_BT709:
242             rx = 0.640; ry = 0.330;
243             gx = 0.300; gy = 0.600;
244             bx = 0.150; by = 0.060;
245             break;
246         case AVCOL_PRI_BT470M:
247             rx = 0.670; ry = 0.330;
248             gx = 0.210; gy = 0.710;
249             bx = 0.140; by = 0.080;
250             wx = 0.310; wy = 0.316;
251             break;
252         case AVCOL_PRI_BT470BG:
253             rx = 0.640; ry = 0.330;
254             gx = 0.290; gy = 0.600;
255             bx = 0.150; by = 0.060;
256             break;
257         case AVCOL_PRI_SMPTE170M:
258         case AVCOL_PRI_SMPTE240M:
259             rx = 0.630; ry = 0.340;
260             gx = 0.310; gy = 0.595;
261             bx = 0.155; by = 0.070;
262             break;
263         case AVCOL_PRI_BT2020:
264             rx = 0.708; ry = 0.292;
265             gx = 0.170; gy = 0.797;
266             bx = 0.131; by = 0.046;
267             break;
268         default:
269             return 0;
270     }
271
272     AV_WB32_PNG(buf     , wx); AV_WB32_PNG(buf + 4 , wy);
273     AV_WB32_PNG(buf + 8 , rx); AV_WB32_PNG(buf + 12, ry);
274     AV_WB32_PNG(buf + 16, gx); AV_WB32_PNG(buf + 20, gy);
275     AV_WB32_PNG(buf + 24, bx); AV_WB32_PNG(buf + 28, by);
276     return 1;
277 }
278
279 static int png_get_gama(enum AVColorTransferCharacteristic trc, uint8_t *buf)
280 {
281     double gamma = avpriv_get_gamma_from_trc(trc);
282     if (gamma <= 1e-6)
283         return 0;
284
285     AV_WB32_PNG(buf, 1.0 / gamma);
286     return 1;
287 }
288
289 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
290                         const AVFrame *pict, int *got_packet)
291 {
292     PNGEncContext *s       = avctx->priv_data;
293     const AVFrame *const p = pict;
294     int bit_depth, color_type, y, len, row_size, ret, is_progressive;
295     int bits_per_pixel, pass_row_size, enc_row_size;
296     int64_t max_packet_size;
297     int compression_level;
298     uint8_t *ptr, *top, *crow_buf, *crow;
299     uint8_t *crow_base       = NULL;
300     uint8_t *progressive_buf = NULL;
301     uint8_t *top_buf         = NULL;
302
303     is_progressive = !!(avctx->flags & CODEC_FLAG_INTERLACED_DCT);
304     switch (avctx->pix_fmt) {
305     case AV_PIX_FMT_RGBA64BE:
306         bit_depth = 16;
307         color_type = PNG_COLOR_TYPE_RGB_ALPHA;
308         break;
309     case AV_PIX_FMT_RGB48BE:
310         bit_depth = 16;
311         color_type = PNG_COLOR_TYPE_RGB;
312         break;
313     case AV_PIX_FMT_RGBA:
314         bit_depth  = 8;
315         color_type = PNG_COLOR_TYPE_RGB_ALPHA;
316         break;
317     case AV_PIX_FMT_RGB24:
318         bit_depth  = 8;
319         color_type = PNG_COLOR_TYPE_RGB;
320         break;
321     case AV_PIX_FMT_GRAY16BE:
322         bit_depth  = 16;
323         color_type = PNG_COLOR_TYPE_GRAY;
324         break;
325     case AV_PIX_FMT_GRAY8:
326         bit_depth  = 8;
327         color_type = PNG_COLOR_TYPE_GRAY;
328         break;
329     case AV_PIX_FMT_GRAY8A:
330         bit_depth = 8;
331         color_type = PNG_COLOR_TYPE_GRAY_ALPHA;
332         break;
333     case AV_PIX_FMT_YA16BE:
334         bit_depth = 16;
335         color_type = PNG_COLOR_TYPE_GRAY_ALPHA;
336         break;
337     case AV_PIX_FMT_MONOBLACK:
338         bit_depth  = 1;
339         color_type = PNG_COLOR_TYPE_GRAY;
340         break;
341     case AV_PIX_FMT_PAL8:
342         bit_depth  = 8;
343         color_type = PNG_COLOR_TYPE_PALETTE;
344         break;
345     default:
346         return -1;
347     }
348     bits_per_pixel = ff_png_get_nb_channels(color_type) * bit_depth;
349     row_size       = (avctx->width * bits_per_pixel + 7) >> 3;
350
351     s->zstream.zalloc = ff_png_zalloc;
352     s->zstream.zfree  = ff_png_zfree;
353     s->zstream.opaque = NULL;
354     compression_level = avctx->compression_level == FF_COMPRESSION_DEFAULT
355                       ? Z_DEFAULT_COMPRESSION
356                       : av_clip(avctx->compression_level, 0, 9);
357     ret = deflateInit2(&s->zstream, compression_level,
358                        Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY);
359     if (ret != Z_OK)
360         return -1;
361
362     enc_row_size    = deflateBound(&s->zstream, row_size);
363     max_packet_size = avctx->height * (int64_t)(enc_row_size +
364                                        ((enc_row_size + IOBUF_SIZE - 1) / IOBUF_SIZE) * 12)
365                       + FF_MIN_BUFFER_SIZE;
366     if (max_packet_size > INT_MAX)
367         return AVERROR(ENOMEM);
368     if ((ret = ff_alloc_packet2(avctx, pkt, max_packet_size)) < 0)
369         return ret;
370
371     s->bytestream_start =
372     s->bytestream       = pkt->data;
373     s->bytestream_end   = pkt->data + pkt->size;
374
375     crow_base = av_malloc((row_size + 32) << (s->filter_type == PNG_FILTER_VALUE_MIXED));
376     if (!crow_base) {
377         ret = AVERROR(ENOMEM);
378         goto the_end;
379     }
380     // pixel data should be aligned, but there's a control byte before it
381     crow_buf = crow_base + 15;
382     if (is_progressive) {
383         progressive_buf = av_malloc(row_size + 1);
384         top_buf = av_malloc(row_size + 1);
385         if (!progressive_buf || !top_buf) {
386             ret = AVERROR(ENOMEM);
387             goto the_end;
388         }
389     }
390
391     /* write png header */
392     AV_WB64(s->bytestream, PNGSIG);
393     s->bytestream += 8;
394
395     AV_WB32(s->buf, avctx->width);
396     AV_WB32(s->buf + 4, avctx->height);
397     s->buf[8]  = bit_depth;
398     s->buf[9]  = color_type;
399     s->buf[10] = 0; /* compression type */
400     s->buf[11] = 0; /* filter type */
401     s->buf[12] = is_progressive; /* interlace type */
402     png_write_chunk(&s->bytestream, MKTAG('I', 'H', 'D', 'R'), s->buf, 13);
403
404     /* write physical information */
405     if (s->dpm) {
406       AV_WB32(s->buf, s->dpm);
407       AV_WB32(s->buf + 4, s->dpm);
408       s->buf[8] = 1; /* unit specifier is meter */
409     } else {
410       AV_WB32(s->buf, avctx->sample_aspect_ratio.num);
411       AV_WB32(s->buf + 4, avctx->sample_aspect_ratio.den);
412       s->buf[8] = 0; /* unit specifier is unknown */
413     }
414     png_write_chunk(&s->bytestream, MKTAG('p', 'H', 'Y', 's'), s->buf, 9);
415
416     /* write colorspace information */
417     if (pict->color_primaries == AVCOL_PRI_BT709 &&
418         pict->color_trc == AVCOL_TRC_IEC61966_2_1) {
419         s->buf[0] = 1; /* rendering intent, relative colorimetric by default */
420         png_write_chunk(&s->bytestream, MKTAG('s', 'R', 'G', 'B'), s->buf, 1);
421     }
422
423     if (png_get_chrm(pict->color_primaries, s->buf))
424         png_write_chunk(&s->bytestream, MKTAG('c', 'H', 'R', 'M'), s->buf, 32);
425     if (png_get_gama(pict->color_trc, s->buf))
426         png_write_chunk(&s->bytestream, MKTAG('g', 'A', 'M', 'A'), s->buf, 4);
427
428     /* put the palette if needed */
429     if (color_type == PNG_COLOR_TYPE_PALETTE) {
430         int has_alpha, alpha, i;
431         unsigned int v;
432         uint32_t *palette;
433         uint8_t *alpha_ptr;
434
435         palette   = (uint32_t *)p->data[1];
436         ptr       = s->buf;
437         alpha_ptr = s->buf + 256 * 3;
438         has_alpha = 0;
439         for (i = 0; i < 256; i++) {
440             v     = palette[i];
441             alpha = v >> 24;
442             if (alpha != 0xff)
443                 has_alpha = 1;
444             *alpha_ptr++ = alpha;
445             bytestream_put_be24(&ptr, v);
446         }
447         png_write_chunk(&s->bytestream,
448                         MKTAG('P', 'L', 'T', 'E'), s->buf, 256 * 3);
449         if (has_alpha) {
450             png_write_chunk(&s->bytestream,
451                             MKTAG('t', 'R', 'N', 'S'), s->buf + 256 * 3, 256);
452         }
453     }
454
455     /* now put each row */
456     s->zstream.avail_out = IOBUF_SIZE;
457     s->zstream.next_out  = s->buf;
458     if (is_progressive) {
459         int pass;
460
461         for (pass = 0; pass < NB_PASSES; pass++) {
462             /* NOTE: a pass is completely omitted if no pixels would be
463              * output */
464             pass_row_size = ff_png_pass_row_size(pass, bits_per_pixel, avctx->width);
465             if (pass_row_size > 0) {
466                 top = NULL;
467                 for (y = 0; y < avctx->height; y++)
468                     if ((ff_png_pass_ymask[pass] << (y & 7)) & 0x80) {
469                         ptr = p->data[0] + y * p->linesize[0];
470                         FFSWAP(uint8_t *, progressive_buf, top_buf);
471                         png_get_interlaced_row(progressive_buf, pass_row_size,
472                                                bits_per_pixel, pass,
473                                                ptr, avctx->width);
474                         crow = png_choose_filter(s, crow_buf, progressive_buf,
475                                                  top, pass_row_size, bits_per_pixel >> 3);
476                         png_write_row(s, crow, pass_row_size + 1);
477                         top = progressive_buf;
478                     }
479             }
480         }
481     } else {
482         top = NULL;
483         for (y = 0; y < avctx->height; y++) {
484             ptr = p->data[0] + y * p->linesize[0];
485             crow = png_choose_filter(s, crow_buf, ptr, top,
486                                      row_size, bits_per_pixel >> 3);
487             png_write_row(s, crow, row_size + 1);
488             top = ptr;
489         }
490     }
491     /* compress last bytes */
492     for (;;) {
493         ret = deflate(&s->zstream, Z_FINISH);
494         if (ret == Z_OK || ret == Z_STREAM_END) {
495             len = IOBUF_SIZE - s->zstream.avail_out;
496             if (len > 0 && s->bytestream_end - s->bytestream > len + 100) {
497                 png_write_chunk(&s->bytestream, MKTAG('I', 'D', 'A', 'T'), s->buf, len);
498             }
499             s->zstream.avail_out = IOBUF_SIZE;
500             s->zstream.next_out  = s->buf;
501             if (ret == Z_STREAM_END)
502                 break;
503         } else {
504             ret = -1;
505             goto the_end;
506         }
507     }
508     png_write_chunk(&s->bytestream, MKTAG('I', 'E', 'N', 'D'), NULL, 0);
509
510     pkt->size   = s->bytestream - s->bytestream_start;
511     pkt->flags |= AV_PKT_FLAG_KEY;
512     *got_packet = 1;
513     ret         = 0;
514
515 the_end:
516     av_freep(&crow_base);
517     av_freep(&progressive_buf);
518     av_freep(&top_buf);
519     deflateEnd(&s->zstream);
520     return ret;
521 }
522
523 static av_cold int png_enc_init(AVCodecContext *avctx)
524 {
525     PNGEncContext *s = avctx->priv_data;
526
527     switch (avctx->pix_fmt) {
528     case AV_PIX_FMT_RGBA:
529         avctx->bits_per_coded_sample = 32;
530         break;
531     case AV_PIX_FMT_RGB24:
532         avctx->bits_per_coded_sample = 24;
533         break;
534     case AV_PIX_FMT_GRAY8:
535         avctx->bits_per_coded_sample = 0x28;
536         break;
537     case AV_PIX_FMT_MONOBLACK:
538         avctx->bits_per_coded_sample = 1;
539         break;
540     case AV_PIX_FMT_PAL8:
541         avctx->bits_per_coded_sample = 8;
542     }
543
544     avctx->coded_frame = av_frame_alloc();
545     if (!avctx->coded_frame)
546         return AVERROR(ENOMEM);
547
548     avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
549     avctx->coded_frame->key_frame = 1;
550
551     ff_huffyuvencdsp_init(&s->hdsp);
552
553     s->filter_type = av_clip(avctx->prediction_method,
554                              PNG_FILTER_VALUE_NONE,
555                              PNG_FILTER_VALUE_MIXED);
556     if (avctx->pix_fmt == AV_PIX_FMT_MONOBLACK)
557         s->filter_type = PNG_FILTER_VALUE_NONE;
558
559     if (s->dpi && s->dpm) {
560       av_log(avctx, AV_LOG_ERROR, "Only one of 'dpi' or 'dpm' options should be set\n");
561       return AVERROR(EINVAL);
562     } else if (s->dpi) {
563       s->dpm = s->dpi * 10000 / 254;
564     }
565
566     return 0;
567 }
568
569 static av_cold int png_enc_close(AVCodecContext *avctx)
570 {
571     av_frame_free(&avctx->coded_frame);
572     return 0;
573 }
574
575 #define OFFSET(x) offsetof(PNGEncContext, x)
576 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
577 static const AVOption options[] = {
578     {"dpi", "Set image resolution (in dots per inch)",  OFFSET(dpi), AV_OPT_TYPE_INT, {.i64 = 0}, 0, 0x10000, VE},
579     {"dpm", "Set image resolution (in dots per meter)", OFFSET(dpm), AV_OPT_TYPE_INT, {.i64 = 0}, 0, 0x10000, VE},
580     { NULL }
581 };
582
583 static const AVClass pngenc_class = {
584     .class_name = "PNG encoder",
585     .item_name  = av_default_item_name,
586     .option     = options,
587     .version    = LIBAVUTIL_VERSION_INT,
588 };
589
590 AVCodec ff_png_encoder = {
591     .name           = "png",
592     .long_name      = NULL_IF_CONFIG_SMALL("PNG (Portable Network Graphics) image"),
593     .type           = AVMEDIA_TYPE_VIDEO,
594     .id             = AV_CODEC_ID_PNG,
595     .priv_data_size = sizeof(PNGEncContext),
596     .init           = png_enc_init,
597     .close          = png_enc_close,
598     .encode2        = encode_frame,
599     .capabilities   = CODEC_CAP_FRAME_THREADS | CODEC_CAP_INTRA_ONLY,
600     .pix_fmts       = (const enum AVPixelFormat[]) {
601         AV_PIX_FMT_RGB24, AV_PIX_FMT_RGBA,
602         AV_PIX_FMT_RGB48BE, AV_PIX_FMT_RGBA64BE,
603         AV_PIX_FMT_PAL8,
604         AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY8A,
605         AV_PIX_FMT_GRAY16BE, AV_PIX_FMT_YA16BE,
606         AV_PIX_FMT_MONOBLACK, AV_PIX_FMT_NONE
607     },
608     .priv_class     = &pngenc_class,
609 };