db2d933ac0c4e10ae0cb957e211987a334a8bd22
[ffmpeg.git] / libav / img.c
1 /*
2  * Image format
3  * Copyright (c) 2000, 2001 Gerard Lantau.
4  *
5  * This program is free software; you can redistribute it and/or modify
6  * it under the terms of the GNU General Public License as published by
7  * the Free Software Foundation; either version 2 of the License, or
8  * (at your option) any later version.
9  *
10  * This program is distributed in the hope that it will be useful,
11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13  * GNU General Public License for more details.
14  *
15  * You should have received a copy of the GNU General Public License
16  * along with this program; if not, write to the Free Software
17  * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
18  */
19 #include <stdlib.h>
20 #include <stdio.h>
21 #include <unistd.h>
22 #include <fcntl.h>
23 #include <errno.h>
24 #include <string.h>
25 #include <math.h>
26
27 #include "avformat.h"
28
29 #define IMGFMT_YUV     1
30 #define IMGFMT_PGMYUV  2
31 #define IMGFMT_PGM     3
32
33 typedef struct {
34     int width;
35     int height;
36     int img_number;
37     int img_size;
38     int img_fmt;
39     int is_pipe;
40     char path[1024];
41 } VideoData;
42
43 static inline int pnm_space(int c)  
44 {
45     return (c==' ' || c=='\n' || c=='\r' || c=='\t');
46 }
47
48 static void pnm_get(ByteIOContext *f, char *str, int buf_size) 
49 {
50     char *s;
51     int c;
52     
53     do  {
54         c=get_byte(f);
55         if (c=='#')  {
56             do  {
57                 c=get_byte(f);
58             } while (c!='\n');
59             c=get_byte(f);
60         }
61     } while (pnm_space(c));
62     
63     s=str;
64     do  {
65         if (url_feof(f))
66             break;
67         if ((s - str)  < buf_size - 1)
68             *s++=c;
69         c=get_byte(f);
70     } while (!pnm_space(c));
71     *s = '\0';
72 }
73
74 static int pgm_read(VideoData *s, ByteIOContext *f, UINT8 *buf, int size, int is_yuv)
75 {
76     int width, height, i;
77     char buf1[32];
78     UINT8 *picture[3];
79
80     width = s->width;
81     height = s->height;
82
83     pnm_get(f, buf1, sizeof(buf1));
84     if (strcmp(buf1, "P5")) {
85         return -EIO;
86     }
87     pnm_get(f, buf1, sizeof(buf1));
88     pnm_get(f, buf1, sizeof(buf1));
89     pnm_get(f, buf1, sizeof(buf1));
90     
91     picture[0] = buf;
92     picture[1] = buf + width * height;
93     picture[2] = buf + width * height + (width * height / 4);
94     get_buffer(f, picture[0], width * height);
95     
96     height>>=1;
97     width>>=1;
98     if (is_yuv) {
99         for(i=0;i<height;i++) {
100             get_buffer(f, picture[1] + i * width, width);
101             get_buffer(f, picture[2] + i * width, width);
102         }
103     } else {
104         for(i=0;i<height;i++) {
105             memset(picture[1] + i * width, 128, width);
106             memset(picture[2] + i * width, 128, width);
107         }
108     }
109     return 0;
110 }
111
112 static int yuv_read(VideoData *s, const char *filename, UINT8 *buf, int size1)
113 {
114     ByteIOContext pb1, *pb = &pb1;
115     char fname[1024], *p;
116     int size;
117
118     size = s->width * s->height;
119     
120     strcpy(fname, filename);
121     p = strrchr(fname, '.');
122     if (!p || p[1] != 'Y')
123         return -EIO;
124
125     if (url_fopen(pb, fname, URL_RDONLY) < 0)
126         return -EIO;
127     
128     get_buffer(pb, buf, size);
129     url_fclose(pb);
130     
131     p[1] = 'U';
132     if (url_fopen(pb, fname, URL_RDONLY) < 0)
133         return -EIO;
134
135     get_buffer(pb, buf + size, size / 4);
136     url_fclose(pb);
137     
138     p[1] = 'V';
139     if (url_fopen(pb, fname, URL_RDONLY) < 0)
140         return -EIO;
141
142     get_buffer(pb, buf + size + (size / 4), size / 4);
143     url_fclose(pb);
144     return 0;
145 }
146
147 int img_read_packet(AVFormatContext *s1, AVPacket *pkt)
148 {
149     VideoData *s = s1->priv_data;
150     char filename[1024];
151     int ret;
152     ByteIOContext f1, *f;
153
154     snprintf(filename, sizeof(filename), s->path, s->img_number);
155     
156     if (!s->is_pipe) {
157         f = &f1;
158         if (url_fopen(f, filename, URL_RDONLY) < 0)
159             return -EIO;
160     } else {
161         f = &s1->pb;
162         if (url_feof(f))
163             return -EIO;
164     }
165
166     av_new_packet(pkt, s->img_size);
167     pkt->stream_index = 0;
168
169     switch(s->img_fmt) {
170     case IMGFMT_PGMYUV:
171         ret = pgm_read(s, f, pkt->data, pkt->size, 1);
172         break;
173     case IMGFMT_PGM:
174         ret = pgm_read(s, f, pkt->data, pkt->size, 0);
175         break;
176     case IMGFMT_YUV:
177         ret = yuv_read(s, filename, pkt->data, pkt->size);
178         break;
179     default:
180         return -EIO;
181     }
182     
183     if (!s->is_pipe) {
184         url_fclose(f);
185     }
186
187     if (ret < 0) {
188         av_free_packet(pkt);
189         return -EIO; /* signal EOF */
190     } else {
191         s->img_number++;
192         return 0;
193     }
194 }
195
196 static int sizes[][2] = {
197     { 640, 480 },
198     { 720, 480 },
199     { 720, 576 },
200     { 352, 288 },
201     { 352, 240 },
202     { 160, 128 },
203     { 512, 384 },
204     { 640, 352 },
205 };
206
207 static int infer_size(int *width_ptr, int *height_ptr, int size)
208 {
209     int i;
210
211     for(i=0;i<sizeof(sizes)/sizeof(sizes[0]);i++) {
212         if ((sizes[i][0] * sizes[i][1]) == size) {
213             *width_ptr = sizes[i][0];
214             *height_ptr = sizes[i][1];
215             return 0;
216         }
217     }
218     return -1;
219 }
220
221 static int img_read_header(AVFormatContext *s1, AVFormatParameters *ap)
222 {
223     VideoData *s;
224     int i, h;
225     char buf[1024];
226     char buf1[32];
227     ByteIOContext pb1, *f = &pb1;
228     AVStream *st;
229
230     s = malloc(sizeof(VideoData));
231     if (!s)
232         return -ENOMEM;
233
234     s1->priv_data = s;
235
236     s1->nb_streams = 1;
237     st = av_mallocz(sizeof(AVStream));
238     if (!st) {
239         free(s);
240         return -ENOMEM;
241     }
242     s1->streams[0] = st;
243     
244     strcpy(s->path, s1->filename);
245     s->img_number = 0;
246
247     /* find format */
248     if (s1->format->flags & AVFMT_NOFILE)
249         s->is_pipe = 0;
250     else
251         s->is_pipe = 1;
252         
253     if (s1->format == &pgmpipe_format ||
254         s1->format == &pgmyuv_format)
255         s->img_fmt = IMGFMT_PGMYUV;
256     else if (s1->format == &pgm_format)
257         s->img_fmt = IMGFMT_PGM;
258     else if (s1->format == &imgyuv_format)
259         s->img_fmt = IMGFMT_YUV;
260     else
261         goto fail;
262
263     if (!s->is_pipe) {
264         /* try to find the first image */
265         for(i=0;i<5;i++) {
266             snprintf(buf, sizeof(buf), s->path, s->img_number);
267             if (url_fopen(f, buf, URL_RDONLY) >= 0)
268                 break;
269             s->img_number++;
270         }
271         if (i == 5)
272             goto fail;
273     } else {
274         f = &s1->pb;
275     }
276     
277     /* find the image size */
278     /* XXX: use generic file format guessing, as mpeg */
279     switch(s->img_fmt) {
280     case IMGFMT_PGM:
281     case IMGFMT_PGMYUV:
282         pnm_get(f, buf1, sizeof(buf1));
283         pnm_get(f, buf1, sizeof(buf1));
284         s->width = atoi(buf1);
285         pnm_get(f, buf1, sizeof(buf1));
286         h = atoi(buf1);
287         if (s->img_fmt == IMGFMT_PGMYUV)
288             h = (h * 2) / 3;
289         s->height = h;
290         if (s->width <= 0 ||
291             s->height <= 0 ||
292             (s->width % 2) != 0 ||
293             (s->height % 2) != 0) {
294             goto fail1;
295         }
296         break;
297     case IMGFMT_YUV:
298         /* infer size by using the file size. */
299         {
300             int img_size;
301             URLContext *h;
302
303             /* XXX: hack hack */
304             h = url_fileno(f);
305             img_size = lseek((int)h->priv_data, 0, SEEK_END);
306             if (infer_size(&s->width, &s->height, img_size) < 0) {
307                 goto fail1;
308             }
309         }
310         break;
311     }
312
313     if (!s->is_pipe) {
314         url_fclose(f);
315     } else {
316         url_fseek(f, 0, SEEK_SET);
317     }
318     
319     s->img_size = (s->width * s->height * 3) / 2;
320
321     st->codec.codec_type = CODEC_TYPE_VIDEO;
322     st->codec.codec_id = CODEC_ID_RAWVIDEO;
323     st->codec.width = s->width;
324     st->codec.height = s->height;
325     st->codec.pix_fmt = PIX_FMT_YUV420P;
326     if (!ap || !ap->frame_rate)
327         st->codec.frame_rate = 25 * FRAME_RATE_BASE;
328     else
329         st->codec.frame_rate = ap->frame_rate;
330     
331     return 0;
332  fail1:
333     if (!s->is_pipe)
334         url_fclose(f);
335  fail:
336     free(s);
337     return -EIO;
338 }
339
340 static int img_read_close(AVFormatContext *s1)
341 {
342     VideoData *s = s1->priv_data;
343     free(s);
344     return 0;
345 }
346
347 /******************************************************/
348 /* image output */
349
350 int pgm_save(AVPicture *picture, int width, int height, ByteIOContext *pb, int is_yuv) 
351 {
352     int i, h;
353     char buf[100];
354     UINT8 *ptr, *ptr1, *ptr2;
355
356     h = height;
357     if (is_yuv)
358         h = (height * 3) / 2;
359     snprintf(buf, sizeof(buf), 
360              "P5\n%d %d\n%d\n",
361              width, h, 255);
362     put_buffer(pb, buf, strlen(buf));
363     
364     ptr = picture->data[0];
365     for(i=0;i<height;i++) {
366         put_buffer(pb, ptr, width);
367         ptr += picture->linesize[0];
368     }
369
370     if (is_yuv) {
371         height >>= 1;
372         width >>= 1;
373         ptr1 = picture->data[1];
374         ptr2 = picture->data[2];
375         for(i=0;i<height;i++) {
376             put_buffer(pb, ptr1, width);
377             put_buffer(pb, ptr2, width);
378             ptr1 += picture->linesize[1];
379             ptr2 += picture->linesize[2];
380         }
381     }
382     put_flush_packet(pb);
383     return 0;
384 }
385
386 static int yuv_save(AVPicture *picture, int width, int height, const char *filename)
387 {
388     ByteIOContext pb1, *pb = &pb1;
389     char fname[1024], *p;
390     int i, j;
391     UINT8 *ptr;
392     static char *ext = "YUV";
393
394     strcpy(fname, filename);
395     p = strrchr(fname, '.');
396     if (!p || p[1] != 'Y')
397         return -EIO;
398
399     for(i=0;i<3;i++) {
400         if (i == 1) {
401             width >>= 1;
402             height >>= 1;
403         }
404
405         p[1] = ext[i];
406         if (url_fopen(pb, fname, URL_WRONLY) < 0)
407             return -EIO;
408     
409         ptr = picture->data[i];
410         for(j=0;j<height;j++) {
411             put_buffer(pb, ptr, width);
412             ptr += picture->linesize[i];
413         }
414         put_flush_packet(pb);
415         url_fclose(pb);
416     }
417     return 0;
418 }
419
420 static int img_write_header(AVFormatContext *s)
421 {
422     VideoData *img;
423
424     img = av_mallocz(sizeof(VideoData));
425     if (!img)
426         return -1;
427     s->priv_data = img;
428     img->img_number = 1;
429     strcpy(img->path, s->filename);
430
431     /* find format */
432     if (s->format->flags & AVFMT_NOFILE)
433         img->is_pipe = 0;
434     else
435         img->is_pipe = 1;
436         
437     if (s->format == &pgmpipe_format ||
438         s->format == &pgmyuv_format)
439         img->img_fmt = IMGFMT_PGMYUV;
440     else if (s->format == &pgm_format)
441         img->img_fmt = IMGFMT_PGM;
442     else if (s->format == &imgyuv_format)
443         img->img_fmt = IMGFMT_YUV;
444     else
445         goto fail;
446     return 0;
447  fail:
448     free(img);
449     return -EIO;
450 }
451
452 static int img_write_packet(AVFormatContext *s, int stream_index,
453                             UINT8 *buf, int size)
454 {
455     VideoData *img = s->priv_data;
456     AVStream *st = s->streams[stream_index];
457     ByteIOContext pb1, *pb;
458     AVPicture picture;
459     int width, height, ret, size1;
460     char filename[1024];
461
462     width = st->codec.width;
463     height = st->codec.height;
464     size1 = (width * height * 3) / 2;
465     if (size != size1)
466         return -EIO;
467
468     picture.data[0] = buf;
469     picture.data[1] = picture.data[0] + width * height;
470     picture.data[2] = picture.data[1] + (width * height) / 4;
471     picture.linesize[0] = width;
472     picture.linesize[1] = width >> 1; 
473     picture.linesize[2] = width >> 1;
474     snprintf(filename, sizeof(filename), img->path, img->img_number);
475
476     if (!img->is_pipe) {
477         pb = &pb1;
478         if (url_fopen(pb, filename, URL_WRONLY) < 0)
479             return -EIO;
480     } else {
481         pb = &s->pb;
482     }
483     switch(img->img_fmt) {
484     case IMGFMT_PGMYUV:
485         ret = pgm_save(&picture, width, height, pb, 1);
486         break;
487     case IMGFMT_PGM:
488         ret = pgm_save(&picture, width, height, pb, 0);
489         break;
490     case IMGFMT_YUV:
491         ret = yuv_save(&picture, width, height, filename);
492         break;
493     }
494     if (!img->is_pipe) {
495         url_fclose(pb);
496     }
497
498     img->img_number++;
499     return 0;
500 }
501
502 static int img_write_trailer(AVFormatContext *s)
503 {
504     VideoData *img = s->priv_data;
505     free(img);
506     return 0;
507 }
508
509 AVFormat pgm_format = {
510     "pgm",
511     "pgm image format",
512     "",
513     "pgm",
514     CODEC_ID_NONE,
515     CODEC_ID_RAWVIDEO,
516     img_write_header,
517     img_write_packet,
518     img_write_trailer,
519
520     img_read_header,
521     img_read_packet,
522     img_read_close,
523     NULL,
524     AVFMT_NOFILE,
525 };
526
527 AVFormat pgmyuv_format = {
528     "pgmyuv",
529     "pgm with YUV content image format",
530     "",
531     "pgm",
532     CODEC_ID_NONE,
533     CODEC_ID_RAWVIDEO,
534     img_write_header,
535     img_write_packet,
536     img_write_trailer,
537
538     img_read_header,
539     img_read_packet,
540     img_read_close,
541     NULL,
542     AVFMT_NOFILE,
543 };
544
545 AVFormat imgyuv_format = {
546     ".Y.U.V",
547     ".Y.U.V format",
548     "",
549     "Y",
550     CODEC_ID_NONE,
551     CODEC_ID_RAWVIDEO,
552     img_write_header,
553     img_write_packet,
554     img_write_trailer,
555
556     img_read_header,
557     img_read_packet,
558     img_read_close,
559     NULL,
560     AVFMT_NOFILE,
561 };
562
563 AVFormat pgmpipe_format = {
564     "pgmpipe",
565     "PGM pipe format",
566     "",
567     "pgm",
568     CODEC_ID_NONE,
569     CODEC_ID_RAWVIDEO,
570     img_write_header,
571     img_write_packet,
572     img_write_trailer,
573
574     img_read_header,
575     img_read_packet,
576     img_read_close,
577     NULL,
578 };