git.maemo.org Git - presencevnc/blob - libvnc/client_examples/vnc2mpg.c

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