1/**
2 * @example vnc2mpg.c
3 * Simple movie writer for vnc; based on Libavformat API example from FFMPEG
4 *
5 * Copyright (c) 2003 Fabrice Bellard, 2004 Johannes E. Schindelin
6 *
7 * Permission is hereby granted, free of charge, to any person obtaining a copy
8 * of this software and associated documentation files (the "Software"), to deal
9 * in the Software without restriction, including without limitation the rights
10 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11 * copies of the Software, and to permit persons to whom the Software is
12 * furnished to do so, subject to the following conditions:
13 *
14 * The above copyright notice and this permission notice shall be included in
15 * all copies or substantial portions of the Software.
16 *
17 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
20 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
23 * THE SOFTWARE.
24 */
25#include <stdlib.h>
26#include <stdio.h>
27#include <string.h>
28#include <signal.h>
29#include <math.h>
30
31#ifndef M_PI
32#define M_PI 3.1415926535897931
33#endif
34
35#include "avformat.h"
36#include <rfb/rfbclient.h>
37
38#define STREAM_FRAME_RATE 25 /* 25 images/s */
39
40/**************************************************************/
41/* video output */
42
43AVFrame *picture, *tmp_picture;
44uint8_t *video_outbuf;
45int frame_count, video_outbuf_size;
46
47/* add a video output stream */
48AVStream *add_video_stream(AVFormatContext *oc, int codec_id, int w, int h)
49{
50    AVCodecContext *c;
51    AVStream *st;
52
53    st = av_new_stream(oc, 0);
54    if (!st) {
55        fprintf(stderr, "Could not alloc stream\n");
56        exit(1);
57    }
58
59#if LIBAVFORMAT_BUILD<4629
60    c = &st->codec;
61#else
62    c = st->codec;
63#endif
64    c->codec_id = codec_id;
65    c->codec_type = CODEC_TYPE_VIDEO;
66
67    /* put sample parameters */
68    c->bit_rate = 800000;
69    /* resolution must be a multiple of two */
70    c->width = w;
71    c->height = h;
72    /* frames per second */
73#if LIBAVCODEC_BUILD<4754
74    c->frame_rate = STREAM_FRAME_RATE;
75    c->frame_rate_base = 1;
76#else
77    c->time_base.den = STREAM_FRAME_RATE;
78    c->time_base.num = 1;
79    c->pix_fmt = PIX_FMT_YUV420P;
80#endif
81    c->gop_size = 12; /* emit one intra frame every twelve frames at most */
82    if (c->codec_id == CODEC_ID_MPEG2VIDEO) {
83        /* just for testing, we also add B frames */
84        c->max_b_frames = 2;
85    }
86    if (c->codec_id == CODEC_ID_MPEG1VIDEO){
87        /* needed to avoid using macroblocks in which some coeffs overflow
88           this doesnt happen with normal video, it just happens here as the
89           motion of the chroma plane doesnt match the luma plane */
90        c->mb_decision=2;
91    }
92    /* some formats want stream headers to be seperate */
93    if(!strcmp(oc->oformat->name, "mp4") || !strcmp(oc->oformat->name, "mov") || !strcmp(oc->oformat->name, "3gp"))
94        c->flags |= CODEC_FLAG_GLOBAL_HEADER;
95
96    return st;
97}
98
99AVFrame *alloc_picture(int pix_fmt, int width, int height)
100{
101    AVFrame *picture;
102    uint8_t *picture_buf;
103    int size;
104
105    picture = avcodec_alloc_frame();
106    if (!picture)
107        return NULL;
108    size = avpicture_get_size(pix_fmt, width, height);
109    picture_buf = malloc(size);
110    if (!picture_buf) {
111        av_free(picture);
112        return NULL;
113    }
114    avpicture_fill((AVPicture *)picture, picture_buf,
115                   pix_fmt, width, height);
116    return picture;
117}
118
119void open_video(AVFormatContext *oc, AVStream *st)
120{
121    AVCodec *codec;
122    AVCodecContext *c;
123
124#if LIBAVFORMAT_BUILD<4629
125    c = &st->codec;
126#else
127    c = st->codec;
128#endif
129
130    /* find the video encoder */
131    codec = avcodec_find_encoder(c->codec_id);
132    if (!codec) {
133        fprintf(stderr, "codec not found\n");
134        exit(1);
135    }
136
137    /* open the codec */
138    if (avcodec_open(c, codec) < 0) {
139        fprintf(stderr, "could not open codec\n");
140        exit(1);
141    }
142
143    video_outbuf = NULL;
144    if (!(oc->oformat->flags & AVFMT_RAWPICTURE)) {
145        /* allocate output buffer */
146        /* XXX: API change will be done */
147        video_outbuf_size = 200000;
148        video_outbuf = malloc(video_outbuf_size);
149    }
150
151    /* allocate the encoded raw picture */
152    picture = alloc_picture(c->pix_fmt, c->width, c->height);
153    if (!picture) {
154        fprintf(stderr, "Could not allocate picture\n");
155        exit(1);
156    }
157
158    /* if the output format is not RGB565, then a temporary RGB565
159       picture is needed too. It is then converted to the required
160       output format */
161    tmp_picture = NULL;
162    if (c->pix_fmt != PIX_FMT_RGB565) {
163        tmp_picture = alloc_picture(PIX_FMT_RGB565, c->width, c->height);
164        if (!tmp_picture) {
165            fprintf(stderr, "Could not allocate temporary picture\n");
166            exit(1);
167        }
168    }
169}
170
171void write_video_frame(AVFormatContext *oc, AVStream *st)
172{
173    int out_size, ret;
174    AVCodecContext *c;
175    AVFrame *picture_ptr;
176
177#if LIBAVFORMAT_BUILD<4629
178    c = &st->codec;
179#else
180    c = st->codec;
181#endif
182
183        if (c->pix_fmt != PIX_FMT_RGB565) {
184            /* as we only generate a RGB565 picture, we must convert it
185               to the codec pixel format if needed */
186            img_convert((AVPicture *)picture, c->pix_fmt,
187                        (AVPicture *)tmp_picture, PIX_FMT_RGB565,
188                        c->width, c->height);
189        }
190	picture_ptr = picture;
191
192
193    if (oc->oformat->flags & AVFMT_RAWPICTURE) {
194        /* raw video case. The API will change slightly in the near
195           futur for that */
196        AVPacket pkt;
197        av_init_packet(&pkt);
198
199        pkt.flags |= PKT_FLAG_KEY;
200        pkt.stream_index= st->index;
201        pkt.data= (uint8_t *)picture_ptr;
202        pkt.size= sizeof(AVPicture);
203
204        ret = av_write_frame(oc, &pkt);
205    } else {
206        /* encode the image */
207        out_size = avcodec_encode_video(c, video_outbuf, video_outbuf_size, picture_ptr);
208        /* if zero size, it means the image was buffered */
209        if (out_size != 0) {
210            AVPacket pkt;
211            av_init_packet(&pkt);
212
213            pkt.pts= c->coded_frame->pts;
214            if(c->coded_frame->key_frame)
215                pkt.flags |= PKT_FLAG_KEY;
216            pkt.stream_index= st->index;
217            pkt.data= video_outbuf;
218            pkt.size= out_size;
219
220            /* write the compressed frame in the media file */
221            ret = av_write_frame(oc, &pkt);
222        } else {
223            ret = 0;
224        }
225    }
226    if (ret != 0) {
227        fprintf(stderr, "Error while writing video frame\n");
228        exit(1);
229    }
230    frame_count++;
231}
232
233void close_video(AVFormatContext *oc, AVStream *st)
234{
235    avcodec_close(st->codec);
236    av_free(picture->data[0]);
237    av_free(picture);
238    if (tmp_picture) {
239        av_free(tmp_picture->data[0]);
240        av_free(tmp_picture);
241    }
242    av_free(video_outbuf);
243}
244
245static const char *filename;
246static AVOutputFormat *fmt;
247static AVFormatContext *oc;
248static AVStream *video_st;
249static double video_pts;
250
251static int movie_open(int w, int h) {
252    if (fmt->video_codec != CODEC_ID_NONE) {
253        video_st = add_video_stream(oc, fmt->video_codec, w, h);
254    } else
255	    return 1;
256
257    /* set the output parameters (must be done even if no
258       parameters). */
259    if (av_set_parameters(oc, NULL) < 0) {
260        fprintf(stderr, "Invalid output format parameters\n");
261        return 2;
262    }
263
264    dump_format(oc, 0, filename, 1);
265
266    /* now that all the parameters are set, we can open the audio and
267       video codecs and allocate the necessary encode buffers */
268    if (video_st)
269        open_video(oc, video_st);
270
271    /* open the output file, if needed */
272    if (!(fmt->flags & AVFMT_NOFILE)) {
273        if (url_fopen(&oc->pb, filename, URL_WRONLY) < 0) {
274            fprintf(stderr, "Could not open '%s'\n", filename);
275            return 3;
276        }
277    }
278
279    /* write the stream header, if any */
280    av_write_header(oc);
281
282    return 0;
283}
284
285static int movie_close() {
286    int i;
287
288     /* close each codec */
289    close_video(oc, video_st);
290
291    /* write the trailer, if any */
292    av_write_trailer(oc);
293
294    /* free the streams */
295    for(i = 0; i < oc->nb_streams; i++) {
296        av_freep(&oc->streams[i]);
297    }
298
299    if (!(fmt->flags & AVFMT_NOFILE)) {
300        /* close the output file */
301        url_fclose(&oc->pb);
302    }
303
304    /* free the stream */
305    av_free(oc);
306
307}
308
309static rfbBool quit=FALSE;
310static void signal_handler(int signal) {
311	fprintf(stderr,"Cleaning up.\n");
312	quit=TRUE;
313}
314
315/**************************************************************/
316/* VNC callback functions */
317static rfbBool resize(rfbClient* client) {
318	static rfbBool first=TRUE;
319	if(!first) {
320		movie_close();
321		perror("I don't know yet how to change resolutions!\n");
322	}
323	movie_open(client->width, client->height);
324	signal(SIGINT,signal_handler);
325	if(tmp_picture)
326		client->frameBuffer=tmp_picture->data[0];
327	else
328		client->frameBuffer=picture->data[0];
329	return TRUE;
330}
331
332static void update(rfbClient* client,int x,int y,int w,int h) {
333}
334
335/**************************************************************/
336/* media file output */
337
338int main(int argc, char **argv)
339{
340    time_t stop=0;
341    rfbClient* client;
342    int i,j;
343
344    /* get a vnc client structure (don't connect yet). */
345    client = rfbGetClient(5,3,2);
346    client->format.redShift=11; client->format.redMax=31;
347    client->format.greenShift=5; client->format.greenMax=63;
348    client->format.blueShift=0; client->format.blueMax=31;
349
350    /* initialize libavcodec, and register all codecs and formats */
351    av_register_all();
352
353    if(!strncmp(argv[argc-1],":",1) ||
354	!strncmp(argv[argc-1],"127.0.0.1",9) ||
355	!strncmp(argv[argc-1],"localhost",9))
356	    client->appData.encodingsString="raw";
357
358    filename=0;
359    for(i=1;i<argc;i++) {
360	    j=i;
361	    if(argc>i+1 && !strcmp("-o",argv[i])) {
362		    filename=argv[2];
363		    j+=2;
364	    } else if(argc>i+1 && !strcmp("-t",argv[i])) {
365		    stop=time(0)+atoi(argv[i+1]);
366		    j+=2;
367	    }
368	    if(j>i) {
369		    argc-=j-i;
370		    memmove(argv+i,argv+j,(argc-i)*sizeof(char*));
371		    i--;
372	    }
373    }
374
375
376    /* auto detect the output format from the name. default is
377       mpeg. */
378    fmt = filename?guess_format(NULL, filename, NULL):0;
379    if (!fmt) {
380        printf("Could not deduce output format from file extension: using MPEG.\n");
381        fmt = guess_format("mpeg", NULL, NULL);
382    }
383    if (!fmt) {
384        fprintf(stderr, "Could not find suitable output format\n");
385        exit(1);
386    }
387
388    /* allocate the output media context */
389    oc = av_alloc_format_context();
390    if (!oc) {
391        fprintf(stderr, "Memory error\n");
392        exit(1);
393    }
394    oc->oformat = fmt;
395    snprintf(oc->filename, sizeof(oc->filename), "%s", filename);
396
397    /* add the audio and video streams using the default format codecs
398       and initialize the codecs */
399    video_st = NULL;
400
401    /* open VNC connection */
402    client->MallocFrameBuffer=resize;
403    client->GotFrameBufferUpdate=update;
404    if(!rfbInitClient(client,&argc,argv)) {
405        printf("usage: %s [-o output_file] [-t seconds] server:port\n"
406	       "Shoot a movie from a VNC server.\n", argv[0]);
407        exit(1);
408    }
409    if(client->serverPort==-1)
410      client->vncRec->doNotSleep = TRUE; /* vncrec playback */
411
412     /* main loop */
413
414    while(!quit) {
415	int i=WaitForMessage(client,1000000/STREAM_FRAME_RATE);
416	if(i<0) {
417		movie_close();
418		return 0;
419	}
420	if(i)
421		if(!HandleRFBServerMessage(client))
422			quit=TRUE;
423	else {
424	        /* compute current audio and video time */
425               	video_pts = (double)video_st->pts.val * video_st->time_base.num / video_st->time_base.den;
426
427        	/* write interleaved audio and video frames */
428	        write_video_frame(oc, video_st);
429	}
430	if(stop!=0 && stop<time(0))
431		quit=TRUE;
432    }
433
434    movie_close();
435    return 0;
436}
437