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
43 AVFrame *picture, *tmp_picture;
44 uint8_t *video_outbuf;
45 int frame_count, video_outbuf_size;
46
47 /* add a video output stream */
add_video_stream(AVFormatContext * oc,int codec_id,int w,int h)48 AVStream *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
alloc_picture(int pix_fmt,int width,int height)99 AVFrame *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
open_video(AVFormatContext * oc,AVStream * st)119 void 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
write_video_frame(AVFormatContext * oc,AVStream * st)171 void 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
close_video(AVFormatContext * oc,AVStream * st)233 void 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
245 static const char *filename;
246 static AVOutputFormat *fmt;
247 static AVFormatContext *oc;
248 static AVStream *video_st;
249 static double video_pts;
250
movie_open(int w,int h)251 static 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
movie_close()285 static 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
309 static rfbBool quit=FALSE;
signal_handler(int signal)310 static void signal_handler(int signal) {
311 fprintf(stderr,"Cleaning up.\n");
312 quit=TRUE;
313 }
314
315 /**************************************************************/
316 /* VNC callback functions */
resize(rfbClient * client)317 static 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
update(rfbClient * client,int x,int y,int w,int h)332 static void update(rfbClient* client,int x,int y,int w,int h) {
333 }
334
335 /**************************************************************/
336 /* media file output */
337
main(int argc,char ** argv)338 int 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