utvideodec: Set colorspace by codec_tag.
[ffmpeg.git] / libavformat / yuv4mpeg.c
1 /*
2  * YUV4MPEG format
3  * Copyright (c) 2001, 2002, 2003 Fabrice Bellard
4  *
5  * This file is part of Libav.
6  *
7  * Libav 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  * Libav 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 Libav; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21
22 #include "libavutil/pixdesc.h"
23 #include "avformat.h"
24 #include "internal.h"
25
26 #define Y4M_MAGIC "YUV4MPEG2"
27 #define Y4M_FRAME_MAGIC "FRAME"
28 #define Y4M_LINE_MAX 256
29
30 struct frame_attributes {
31     int interlaced_frame;
32     int top_field_first;
33 };
34
35 #if CONFIG_YUV4MPEGPIPE_MUXER
36 static int yuv4_generate_header(AVFormatContext *s, char* buf)
37 {
38     AVStream *st;
39     int width, height;
40     int raten, rated, aspectn, aspectd, n;
41     char inter;
42     const char *colorspace = "";
43
44     st     = s->streams[0];
45     width  = st->codec->width;
46     height = st->codec->height;
47
48     av_reduce(&raten, &rated, st->codec->time_base.den,
49               st->codec->time_base.num, (1UL << 31) - 1);
50
51     aspectn = st->sample_aspect_ratio.num;
52     aspectd = st->sample_aspect_ratio.den;
53
54     if (aspectn == 0 && aspectd == 1)
55         aspectd = 0;  // 0:0 means unknown
56
57     inter = 'p'; /* progressive is the default */
58     if (st->codec->coded_frame && st->codec->coded_frame->interlaced_frame)
59         inter = st->codec->coded_frame->top_field_first ? 't' : 'b';
60
61     switch (st->codec->pix_fmt) {
62     case AV_PIX_FMT_GRAY8:
63         colorspace = " Cmono";
64         break;
65     case AV_PIX_FMT_YUV411P:
66         colorspace = " C411 XYSCSS=411";
67         break;
68     case AV_PIX_FMT_YUV420P:
69         switch (st->codec->chroma_sample_location) {
70         case AVCHROMA_LOC_TOPLEFT: colorspace = " C420paldv XYSCSS=420PALDV"; break;
71         case AVCHROMA_LOC_LEFT:    colorspace = " C420mpeg2 XYSCSS=420MPEG2"; break;
72         default:                   colorspace = " C420jpeg XYSCSS=420JPEG";   break;
73         }
74         break;
75     case AV_PIX_FMT_YUV422P:
76         colorspace = " C422 XYSCSS=422";
77         break;
78     case AV_PIX_FMT_YUV444P:
79         colorspace = " C444 XYSCSS=444";
80         break;
81     }
82
83     /* construct stream header, if this is the first frame */
84     n = snprintf(buf, Y4M_LINE_MAX, "%s W%d H%d F%d:%d I%c A%d:%d%s\n",
85                  Y4M_MAGIC, width, height, raten, rated, inter,
86                  aspectn, aspectd, colorspace);
87
88     return n;
89 }
90
91 static int yuv4_write_packet(AVFormatContext *s, AVPacket *pkt)
92 {
93     AVStream *st = s->streams[pkt->stream_index];
94     AVIOContext *pb = s->pb;
95     AVPicture *picture;
96     int* first_pkt = s->priv_data;
97     int width, height, h_chroma_shift, v_chroma_shift;
98     int i;
99     char buf2[Y4M_LINE_MAX + 1];
100     char buf1[20];
101     uint8_t *ptr, *ptr1, *ptr2;
102
103     picture = (AVPicture *)pkt->data;
104
105     /* for the first packet we have to output the header as well */
106     if (*first_pkt) {
107         *first_pkt = 0;
108         if (yuv4_generate_header(s, buf2) < 0) {
109             av_log(s, AV_LOG_ERROR,
110                    "Error. YUV4MPEG stream header write failed.\n");
111             return AVERROR(EIO);
112         } else {
113             avio_write(pb, buf2, strlen(buf2));
114         }
115     }
116
117     /* construct frame header */
118
119     snprintf(buf1, sizeof(buf1), "%s\n", Y4M_FRAME_MAGIC);
120     avio_write(pb, buf1, strlen(buf1));
121
122     width  = st->codec->width;
123     height = st->codec->height;
124
125     ptr = picture->data[0];
126     for (i = 0; i < height; i++) {
127         avio_write(pb, ptr, width);
128         ptr += picture->linesize[0];
129     }
130
131     if (st->codec->pix_fmt != AV_PIX_FMT_GRAY8) {
132         // Adjust for smaller Cb and Cr planes
133         av_pix_fmt_get_chroma_sub_sample(st->codec->pix_fmt, &h_chroma_shift,
134                                          &v_chroma_shift);
135         width  >>= h_chroma_shift;
136         height >>= v_chroma_shift;
137
138         ptr1 = picture->data[1];
139         ptr2 = picture->data[2];
140         for (i = 0; i < height; i++) {     /* Cb */
141             avio_write(pb, ptr1, width);
142             ptr1 += picture->linesize[1];
143         }
144         for (i = 0; i < height; i++) {     /* Cr */
145             avio_write(pb, ptr2, width);
146             ptr2 += picture->linesize[2];
147         }
148     }
149     avio_flush(pb);
150     return 0;
151 }
152
153 static int yuv4_write_header(AVFormatContext *s)
154 {
155     int *first_pkt = s->priv_data;
156
157     if (s->nb_streams != 1)
158         return AVERROR(EIO);
159
160     if (s->streams[0]->codec->codec_id != AV_CODEC_ID_RAWVIDEO) {
161         av_log(s, AV_LOG_ERROR, "ERROR: Only rawvideo supported.\n");
162         return AVERROR_INVALIDDATA;
163     }
164
165     if (s->streams[0]->codec->pix_fmt == AV_PIX_FMT_YUV411P) {
166         av_log(s, AV_LOG_ERROR, "Warning: generating rarely used 4:1:1 YUV "
167                "stream, some mjpegtools might not work.\n");
168     } else if ((s->streams[0]->codec->pix_fmt != AV_PIX_FMT_YUV420P) &&
169                (s->streams[0]->codec->pix_fmt != AV_PIX_FMT_YUV422P) &&
170                (s->streams[0]->codec->pix_fmt != AV_PIX_FMT_GRAY8)   &&
171                (s->streams[0]->codec->pix_fmt != AV_PIX_FMT_YUV444P)) {
172         av_log(s, AV_LOG_ERROR, "ERROR: yuv4mpeg only handles yuv444p, "
173                "yuv422p, yuv420p, yuv411p and gray pixel formats. "
174                "Use -pix_fmt to select one.\n");
175         return AVERROR(EIO);
176     }
177
178     *first_pkt = 1;
179     return 0;
180 }
181
182 AVOutputFormat ff_yuv4mpegpipe_muxer = {
183     .name              = "yuv4mpegpipe",
184     .long_name         = NULL_IF_CONFIG_SMALL("YUV4MPEG pipe"),
185     .mime_type         = "",
186     .extensions        = "y4m",
187     .priv_data_size    = sizeof(int),
188     .audio_codec       = AV_CODEC_ID_NONE,
189     .video_codec       = AV_CODEC_ID_RAWVIDEO,
190     .write_header      = yuv4_write_header,
191     .write_packet      = yuv4_write_packet,
192     .flags             = AVFMT_RAWPICTURE,
193 };
194 #endif
195
196 /* Header size increased to allow room for optional flags */
197 #define MAX_YUV4_HEADER 80
198 #define MAX_FRAME_HEADER 80
199
200 static int yuv4_read_header(AVFormatContext *s)
201 {
202     char header[MAX_YUV4_HEADER + 10];  // Include headroom for
203                                         // the longest option
204     char *tokstart, *tokend, *header_end;
205     int i;
206     AVIOContext *pb = s->pb;
207     int width = -1, height  = -1, raten   = 0,
208         rated =  0, aspectn =  0, aspectd = 0;
209     enum AVPixelFormat pix_fmt = AV_PIX_FMT_NONE, alt_pix_fmt = AV_PIX_FMT_NONE;
210     enum AVChromaLocation chroma_sample_location = AVCHROMA_LOC_UNSPECIFIED;
211     AVStream *st;
212     struct frame_attributes *s1 = s->priv_data;
213
214     for (i = 0; i < MAX_YUV4_HEADER; i++) {
215         header[i] = avio_r8(pb);
216         if (header[i] == '\n') {
217             header[i + 1] = 0x20;  // Add a space after last option.
218                                    // Makes parsing "444" vs "444alpha" easier.
219             header[i + 2] = 0;
220             break;
221         }
222     }
223     if (i == MAX_YUV4_HEADER)
224         return -1;
225     if (strncmp(header, Y4M_MAGIC, strlen(Y4M_MAGIC)))
226         return -1;
227
228     s1->interlaced_frame = 0;
229     s1->top_field_first = 0;
230     header_end = &header[i + 1]; // Include space
231     for (tokstart = &header[strlen(Y4M_MAGIC) + 1];
232          tokstart < header_end; tokstart++) {
233         if (*tokstart == 0x20)
234             continue;
235         switch (*tokstart++) {
236         case 'W': // Width. Required.
237             width    = strtol(tokstart, &tokend, 10);
238             tokstart = tokend;
239             break;
240         case 'H': // Height. Required.
241             height   = strtol(tokstart, &tokend, 10);
242             tokstart = tokend;
243             break;
244         case 'C': // Color space
245             if (strncmp("420jpeg", tokstart, 7) == 0) {
246                 pix_fmt = AV_PIX_FMT_YUV420P;
247                 chroma_sample_location = AVCHROMA_LOC_CENTER;
248             } else if (strncmp("420mpeg2", tokstart, 8) == 0) {
249                 pix_fmt = AV_PIX_FMT_YUV420P;
250                 chroma_sample_location = AVCHROMA_LOC_LEFT;
251             } else if (strncmp("420paldv", tokstart, 8) == 0) {
252                 pix_fmt = AV_PIX_FMT_YUV420P;
253                 chroma_sample_location = AVCHROMA_LOC_TOPLEFT;
254             } else if (strncmp("420", tokstart, 3) == 0) {
255                 pix_fmt = AV_PIX_FMT_YUV420P;
256                 chroma_sample_location = AVCHROMA_LOC_CENTER;
257             } else if (strncmp("411", tokstart, 3) == 0)
258                 pix_fmt = AV_PIX_FMT_YUV411P;
259             else if (strncmp("422", tokstart, 3) == 0)
260                 pix_fmt = AV_PIX_FMT_YUV422P;
261             else if (strncmp("444alpha", tokstart, 8) == 0 ) {
262                 av_log(s, AV_LOG_ERROR, "Cannot handle 4:4:4:4 "
263                        "YUV4MPEG stream.\n");
264                 return -1;
265             } else if (strncmp("444", tokstart, 3) == 0)
266                 pix_fmt = AV_PIX_FMT_YUV444P;
267             else if (strncmp("mono", tokstart, 4) == 0) {
268                 pix_fmt = AV_PIX_FMT_GRAY8;
269             } else {
270                 av_log(s, AV_LOG_ERROR, "YUV4MPEG stream contains an unknown "
271                        "pixel format.\n");
272                 return -1;
273             }
274             while (tokstart < header_end && *tokstart != 0x20)
275                 tokstart++;
276             break;
277         case 'I': // Interlace type
278             switch (*tokstart++){
279             case '?':
280                 break;
281             case 'p':
282                 s1->interlaced_frame = 0;
283                 break;
284             case 't':
285                 s1->interlaced_frame = 1;
286                 s1->top_field_first = 1;
287                 break;
288             case 'b':
289                 s1->interlaced_frame = 1;
290                 s1->top_field_first = 0;
291                 break;
292             case 'm':
293                 av_log(s, AV_LOG_ERROR, "YUV4MPEG stream contains mixed "
294                        "interlaced and non-interlaced frames.\n");
295                 return -1;
296             default:
297                 av_log(s, AV_LOG_ERROR, "YUV4MPEG has invalid header.\n");
298                 return -1;
299             }
300             break;
301         case 'F': // Frame rate
302             sscanf(tokstart, "%d:%d", &raten, &rated); // 0:0 if unknown
303             while (tokstart < header_end && *tokstart != 0x20)
304                 tokstart++;
305             break;
306         case 'A': // Pixel aspect
307             sscanf(tokstart, "%d:%d", &aspectn, &aspectd); // 0:0 if unknown
308             while (tokstart < header_end && *tokstart != 0x20)
309                 tokstart++;
310             break;
311         case 'X': // Vendor extensions
312             if (strncmp("YSCSS=", tokstart, 6) == 0) {
313                 // Older nonstandard pixel format representation
314                 tokstart += 6;
315                 if (strncmp("420JPEG", tokstart, 7) == 0)
316                     alt_pix_fmt = AV_PIX_FMT_YUV420P;
317                 else if (strncmp("420MPEG2", tokstart, 8) == 0)
318                     alt_pix_fmt = AV_PIX_FMT_YUV420P;
319                 else if (strncmp("420PALDV", tokstart, 8) == 0)
320                     alt_pix_fmt = AV_PIX_FMT_YUV420P;
321                 else if (strncmp("411", tokstart, 3) == 0)
322                     alt_pix_fmt = AV_PIX_FMT_YUV411P;
323                 else if (strncmp("422", tokstart, 3) == 0)
324                     alt_pix_fmt = AV_PIX_FMT_YUV422P;
325                 else if (strncmp("444", tokstart, 3) == 0)
326                     alt_pix_fmt = AV_PIX_FMT_YUV444P;
327             }
328             while (tokstart < header_end && *tokstart != 0x20)
329                 tokstart++;
330             break;
331         }
332     }
333
334     if (width == -1 || height == -1) {
335         av_log(s, AV_LOG_ERROR, "YUV4MPEG has invalid header.\n");
336         return -1;
337     }
338
339     if (pix_fmt == AV_PIX_FMT_NONE) {
340         if (alt_pix_fmt == AV_PIX_FMT_NONE)
341             pix_fmt = AV_PIX_FMT_YUV420P;
342         else
343             pix_fmt = alt_pix_fmt;
344     }
345
346     if (raten <= 0 || rated <= 0) {
347         // Frame rate unknown
348         raten = 25;
349         rated = 1;
350     }
351
352     if (aspectn == 0 && aspectd == 0) {
353         // Pixel aspect unknown
354         aspectd = 1;
355     }
356
357     st = avformat_new_stream(s, NULL);
358     if (!st)
359         return AVERROR(ENOMEM);
360     st->codec->width  = width;
361     st->codec->height = height;
362     av_reduce(&raten, &rated, raten, rated, (1UL << 31) - 1);
363     avpriv_set_pts_info(st, 64, rated, raten);
364     st->codec->pix_fmt                = pix_fmt;
365     st->codec->codec_type             = AVMEDIA_TYPE_VIDEO;
366     st->codec->codec_id               = AV_CODEC_ID_RAWVIDEO;
367     st->sample_aspect_ratio           = (AVRational){ aspectn, aspectd };
368     st->codec->chroma_sample_location = chroma_sample_location;
369
370     return 0;
371 }
372
373 static int yuv4_read_packet(AVFormatContext *s, AVPacket *pkt)
374 {
375     int i;
376     char header[MAX_FRAME_HEADER+1];
377     int packet_size, width, height, ret;
378     AVStream *st = s->streams[0];
379     struct frame_attributes *s1 = s->priv_data;
380
381     for (i = 0; i < MAX_FRAME_HEADER; i++) {
382         header[i] = avio_r8(s->pb);
383         if (header[i] == '\n') {
384             header[i + 1] = 0;
385             break;
386         }
387     }
388     if (s->pb->error)
389         return s->pb->error;
390     else if (s->pb->eof_reached)
391         return AVERROR_EOF;
392     else if (i == MAX_FRAME_HEADER)
393         return AVERROR_INVALIDDATA;
394
395     if (strncmp(header, Y4M_FRAME_MAGIC, strlen(Y4M_FRAME_MAGIC)))
396         return AVERROR_INVALIDDATA;
397
398     width  = st->codec->width;
399     height = st->codec->height;
400
401     packet_size = avpicture_get_size(st->codec->pix_fmt, width, height);
402     if (packet_size < 0)
403         return packet_size;
404
405     ret = av_get_packet(s->pb, pkt, packet_size);
406     if (ret < 0)
407         return ret;
408     else if (ret != packet_size)
409         return s->pb->eof_reached ? AVERROR_EOF : AVERROR(EIO);
410
411     if (st->codec->coded_frame) {
412         st->codec->coded_frame->interlaced_frame = s1->interlaced_frame;
413         st->codec->coded_frame->top_field_first  = s1->top_field_first;
414     }
415
416     pkt->stream_index = 0;
417     return 0;
418 }
419
420 static int yuv4_probe(AVProbeData *pd)
421 {
422     /* check file header */
423     if (strncmp(pd->buf, Y4M_MAGIC, sizeof(Y4M_MAGIC) - 1) == 0)
424         return AVPROBE_SCORE_MAX;
425     else
426         return 0;
427 }
428
429 #if CONFIG_YUV4MPEGPIPE_DEMUXER
430 AVInputFormat ff_yuv4mpegpipe_demuxer = {
431     .name           = "yuv4mpegpipe",
432     .long_name      = NULL_IF_CONFIG_SMALL("YUV4MPEG pipe"),
433     .priv_data_size = sizeof(struct frame_attributes),
434     .read_probe     = yuv4_probe,
435     .read_header    = yuv4_read_header,
436     .read_packet    = yuv4_read_packet,
437     .extensions     = "y4m",
438 };
439 #endif