00001
00025 #include <stdlib.h>
00026 #include <stdio.h>
00027 #include <string.h>
00028 #include <signal.h>
00029 #include <math.h>
00030
00031 #ifndef M_PI
00032 #define M_PI 3.1415926535897931
00033 #endif
00034
00035 #include "avformat.h"
00036 #include <rfb/rfbclient.h>
00037
00038 #define STREAM_FRAME_RATE 25
00039
00040
00041
00042
00043 AVFrame *picture, *tmp_picture;
00044 uint8_t *video_outbuf;
00045 int frame_count, video_outbuf_size;
00046
00047
00048 AVStream *add_video_stream(AVFormatContext *oc, int codec_id, int w, int h)
00049 {
00050 AVCodecContext *c;
00051 AVStream *st;
00052
00053 st = av_new_stream(oc, 0);
00054 if (!st) {
00055 fprintf(stderr, "Could not alloc stream\n");
00056 exit(1);
00057 }
00058
00059 #if LIBAVFORMAT_BUILD<4629
00060 c = &st->codec;
00061 #else
00062 c = st->codec;
00063 #endif
00064 c->codec_id = codec_id;
00065 c->codec_type = CODEC_TYPE_VIDEO;
00066
00067
00068 c->bit_rate = 800000;
00069
00070 c->width = w;
00071 c->height = h;
00072
00073 #if LIBAVCODEC_BUILD<4754
00074 c->frame_rate = STREAM_FRAME_RATE;
00075 c->frame_rate_base = 1;
00076 #else
00077 c->time_base.den = STREAM_FRAME_RATE;
00078 c->time_base.num = 1;
00079 c->pix_fmt = PIX_FMT_YUV420P;
00080 #endif
00081 c->gop_size = 12;
00082 if (c->codec_id == CODEC_ID_MPEG2VIDEO) {
00083
00084 c->max_b_frames = 2;
00085 }
00086 if (c->codec_id == CODEC_ID_MPEG1VIDEO){
00087
00088
00089
00090 c->mb_decision=2;
00091 }
00092
00093 if(!strcmp(oc->oformat->name, "mp4") || !strcmp(oc->oformat->name, "mov") || !strcmp(oc->oformat->name, "3gp"))
00094 c->flags |= CODEC_FLAG_GLOBAL_HEADER;
00095
00096 return st;
00097 }
00098
00099 AVFrame *alloc_picture(int pix_fmt, int width, int height)
00100 {
00101 AVFrame *picture;
00102 uint8_t *picture_buf;
00103 int size;
00104
00105 picture = avcodec_alloc_frame();
00106 if (!picture)
00107 return NULL;
00108 size = avpicture_get_size(pix_fmt, width, height);
00109 picture_buf = malloc(size);
00110 if (!picture_buf) {
00111 av_free(picture);
00112 return NULL;
00113 }
00114 avpicture_fill((AVPicture *)picture, picture_buf,
00115 pix_fmt, width, height);
00116 return picture;
00117 }
00118
00119 void open_video(AVFormatContext *oc, AVStream *st)
00120 {
00121 AVCodec *codec;
00122 AVCodecContext *c;
00123
00124 #if LIBAVFORMAT_BUILD<4629
00125 c = &st->codec;
00126 #else
00127 c = st->codec;
00128 #endif
00129
00130
00131 codec = avcodec_find_encoder(c->codec_id);
00132 if (!codec) {
00133 fprintf(stderr, "codec not found\n");
00134 exit(1);
00135 }
00136
00137
00138 if (avcodec_open(c, codec) < 0) {
00139 fprintf(stderr, "could not open codec\n");
00140 exit(1);
00141 }
00142
00143 video_outbuf = NULL;
00144 if (!(oc->oformat->flags & AVFMT_RAWPICTURE)) {
00145
00146
00147 video_outbuf_size = 200000;
00148 video_outbuf = malloc(video_outbuf_size);
00149 }
00150
00151
00152 picture = alloc_picture(c->pix_fmt, c->width, c->height);
00153 if (!picture) {
00154 fprintf(stderr, "Could not allocate picture\n");
00155 exit(1);
00156 }
00157
00158
00159
00160
00161 tmp_picture = NULL;
00162 if (c->pix_fmt != PIX_FMT_RGB565) {
00163 tmp_picture = alloc_picture(PIX_FMT_RGB565, c->width, c->height);
00164 if (!tmp_picture) {
00165 fprintf(stderr, "Could not allocate temporary picture\n");
00166 exit(1);
00167 }
00168 }
00169 }
00170
00171 void write_video_frame(AVFormatContext *oc, AVStream *st)
00172 {
00173 int out_size, ret;
00174 AVCodecContext *c;
00175 AVFrame *picture_ptr;
00176
00177 #if LIBAVFORMAT_BUILD<4629
00178 c = &st->codec;
00179 #else
00180 c = st->codec;
00181 #endif
00182
00183 if (c->pix_fmt != PIX_FMT_RGB565) {
00184
00185
00186 img_convert((AVPicture *)picture, c->pix_fmt,
00187 (AVPicture *)tmp_picture, PIX_FMT_RGB565,
00188 c->width, c->height);
00189 }
00190 picture_ptr = picture;
00191
00192
00193 if (oc->oformat->flags & AVFMT_RAWPICTURE) {
00194
00195
00196 AVPacket pkt;
00197 av_init_packet(&pkt);
00198
00199 pkt.flags |= PKT_FLAG_KEY;
00200 pkt.stream_index= st->index;
00201 pkt.data= (uint8_t *)picture_ptr;
00202 pkt.size= sizeof(AVPicture);
00203
00204 ret = av_write_frame(oc, &pkt);
00205 } else {
00206
00207 out_size = avcodec_encode_video(c, video_outbuf, video_outbuf_size, picture_ptr);
00208
00209 if (out_size != 0) {
00210 AVPacket pkt;
00211 av_init_packet(&pkt);
00212
00213 pkt.pts= c->coded_frame->pts;
00214 if(c->coded_frame->key_frame)
00215 pkt.flags |= PKT_FLAG_KEY;
00216 pkt.stream_index= st->index;
00217 pkt.data= video_outbuf;
00218 pkt.size= out_size;
00219
00220
00221 ret = av_write_frame(oc, &pkt);
00222 } else {
00223 ret = 0;
00224 }
00225 }
00226 if (ret != 0) {
00227 fprintf(stderr, "Error while writing video frame\n");
00228 exit(1);
00229 }
00230 frame_count++;
00231 }
00232
00233 void close_video(AVFormatContext *oc, AVStream *st)
00234 {
00235 avcodec_close(st->codec);
00236 av_free(picture->data[0]);
00237 av_free(picture);
00238 if (tmp_picture) {
00239 av_free(tmp_picture->data[0]);
00240 av_free(tmp_picture);
00241 }
00242 av_free(video_outbuf);
00243 }
00244
00245 static const char *filename;
00246 static AVOutputFormat *fmt;
00247 static AVFormatContext *oc;
00248 static AVStream *video_st;
00249 static double video_pts;
00250
00251 static int movie_open(int w, int h) {
00252 if (fmt->video_codec != CODEC_ID_NONE) {
00253 video_st = add_video_stream(oc, fmt->video_codec, w, h);
00254 } else
00255 return 1;
00256
00257
00258
00259 if (av_set_parameters(oc, NULL) < 0) {
00260 fprintf(stderr, "Invalid output format parameters\n");
00261 return 2;
00262 }
00263
00264 dump_format(oc, 0, filename, 1);
00265
00266
00267
00268 if (video_st)
00269 open_video(oc, video_st);
00270
00271
00272 if (!(fmt->flags & AVFMT_NOFILE)) {
00273 if (url_fopen(&oc->pb, filename, URL_WRONLY) < 0) {
00274 fprintf(stderr, "Could not open '%s'\n", filename);
00275 return 3;
00276 }
00277 }
00278
00279
00280 av_write_header(oc);
00281
00282 return 0;
00283 }
00284
00285 static int movie_close() {
00286 int i;
00287
00288
00289 close_video(oc, video_st);
00290
00291
00292 av_write_trailer(oc);
00293
00294
00295 for(i = 0; i < oc->nb_streams; i++) {
00296 av_freep(&oc->streams[i]);
00297 }
00298
00299 if (!(fmt->flags & AVFMT_NOFILE)) {
00300
00301 url_fclose(&oc->pb);
00302 }
00303
00304
00305 av_free(oc);
00306
00307 }
00308
00309 static rfbBool quit=FALSE;
00310 static void signal_handler(int signal) {
00311 fprintf(stderr,"Cleaning up.\n");
00312 quit=TRUE;
00313 }
00314
00315
00316
00317 static rfbBool resize(rfbClient* client) {
00318 static rfbBool first=TRUE;
00319 if(!first) {
00320 movie_close();
00321 perror("I don't know yet how to change resolutions!\n");
00322 }
00323 movie_open(client->width, client->height);
00324 signal(SIGINT,signal_handler);
00325 if(tmp_picture)
00326 client->frameBuffer=tmp_picture->data[0];
00327 else
00328 client->frameBuffer=picture->data[0];
00329 return TRUE;
00330 }
00331
00332 static void update(rfbClient* client,int x,int y,int w,int h) {
00333 }
00334
00335
00336
00337
00338 int main(int argc, char **argv)
00339 {
00340 time_t stop=0;
00341 rfbClient* client;
00342 int i,j;
00343
00344
00345 client = rfbGetClient(5,3,2);
00346 client->format.redShift=11; client->format.redMax=31;
00347 client->format.greenShift=5; client->format.greenMax=63;
00348 client->format.blueShift=0; client->format.blueMax=31;
00349
00350
00351 av_register_all();
00352
00353 if(!strncmp(argv[argc-1],":",1) ||
00354 !strncmp(argv[argc-1],"127.0.0.1",9) ||
00355 !strncmp(argv[argc-1],"localhost",9))
00356 client->appData.encodingsString="raw";
00357
00358 filename=0;
00359 for(i=1;i<argc;i++) {
00360 j=i;
00361 if(argc>i+1 && !strcmp("-o",argv[i])) {
00362 filename=argv[2];
00363 j+=2;
00364 } else if(argc>i+1 && !strcmp("-t",argv[i])) {
00365 stop=time(0)+atoi(argv[i+1]);
00366 j+=2;
00367 }
00368 if(j>i) {
00369 argc-=j-i;
00370 memmove(argv+i,argv+j,(argc-i)*sizeof(char*));
00371 i--;
00372 }
00373 }
00374
00375
00376
00377
00378 fmt = filename?guess_format(NULL, filename, NULL):0;
00379 if (!fmt) {
00380 printf("Could not deduce output format from file extension: using MPEG.\n");
00381 fmt = guess_format("mpeg", NULL, NULL);
00382 }
00383 if (!fmt) {
00384 fprintf(stderr, "Could not find suitable output format\n");
00385 exit(1);
00386 }
00387
00388
00389 oc = av_alloc_format_context();
00390 if (!oc) {
00391 fprintf(stderr, "Memory error\n");
00392 exit(1);
00393 }
00394 oc->oformat = fmt;
00395 snprintf(oc->filename, sizeof(oc->filename), "%s", filename);
00396
00397
00398
00399 video_st = NULL;
00400
00401
00402 client->MallocFrameBuffer=resize;
00403 client->GotFrameBufferUpdate=update;
00404 if(!rfbInitClient(client,&argc,argv)) {
00405 printf("usage: %s [-o output_file] [-t seconds] server:port\n"
00406 "Shoot a movie from a VNC server.\n", argv[0]);
00407 exit(1);
00408 }
00409 if(client->serverPort==-1)
00410 client->vncRec->doNotSleep = TRUE;
00411
00412
00413
00414 while(!quit) {
00415 int i=WaitForMessage(client,1000000/STREAM_FRAME_RATE);
00416 if(i<0) {
00417 movie_close();
00418 return 0;
00419 }
00420 if(i)
00421 if(!HandleRFBServerMessage(client))
00422 quit=TRUE;
00423 else {
00424
00425 video_pts = (double)video_st->pts.val * video_st->time_base.num / video_st->time_base.den;
00426
00427
00428 write_video_frame(oc, video_st);
00429 }
00430 if(stop!=0 && stop<time(0))
00431 quit=TRUE;
00432 }
00433
00434 movie_close();
00435 return 0;
00436 }