[GNUnet-SVN] r9891 - Extractor/src/plugins

[gnunet]

Author: holindho
Date: 2009-12-24 16:26:29 +0100 (Thu, 24 Dec 2009)
New Revision: 9891

Modified:
   Extractor/src/plugins/thumbnailffmpeg_extractor.c
Log:
clean up and remove (fix) oop-only and force-kill requirements


Modified: Extractor/src/plugins/thumbnailffmpeg_extractor.c
===================================================================
--- Extractor/src/plugins/thumbnailffmpeg_extractor.c   2009-12-23 23:19:19 UTC 
(rev 9890)
+++ Extractor/src/plugins/thumbnailffmpeg_extractor.c   2009-12-24 15:26:29 UTC 
(rev 9891)
@@ -46,6 +46,7 @@
 
 void __attribute__ ((constructor)) ffmpeg_lib_init (void)
 {
+  printf("av_register_all\n");
   av_register_all ();
 }
 
@@ -56,16 +57,152 @@
 const char *
 EXTRACTOR_thumbnailffmpeg_options ()
 {
-  return "force-kill;oop-only;close-stderr";
+  return "close-stderr";
 }
 
 const char *
 EXTRACTOR_thumbnail_options ()
 {
-  return "force-kill;oop-only;close-stderr";
+  return "close-stderr";
 }
 
+/*
+ * Rescale and encode a png thumbnail
+ * on success, fills in output_data and returns the number of bytes used
+ */
+static size_t create_thumbnail(
+  int src_width, int src_height, int src_stride[],
+  enum PixelFormat src_pixfmt, uint8_t *src_data[],
+  int dst_width, int dst_height,
+  uint8_t **output_data, size_t output_max_size)
+{
+  AVCodecContext *encoder_codec_ctx = NULL;
+  AVCodec *encoder_codec = NULL;
+  struct SwsContext *scaler_ctx = NULL;
+  int sws_flags = SWS_BILINEAR;
+  AVFrame *dst_frame = NULL;
+  uint8_t *dst_buffer = NULL;
+  uint8_t *encoder_output_buffer = NULL;
+  size_t encoder_output_buffer_size;
+  int err;
 
+  encoder_codec = avcodec_find_encoder_by_name ("png");
+  if (encoder_codec == NULL)
+    {
+#if DEBUG
+      fprintf (stderr,
+              "Couldn't find a PNG encoder\n");
+#endif
+      return 0;
+    }
+
+  /* NOTE: the scaler will be used even if the src and dst image dimensions
+   * match, because the scaler will also perform colour space conversion */
+  scaler_ctx =
+    sws_getContext (src_width, src_height, src_pixfmt,
+                    dst_width, dst_height, PIX_FMT_RGB24, 
+                    sws_flags, NULL, NULL, NULL);
+  if (scaler_ctx == NULL)
+    {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to get a scaler context\n");
+#endif
+      return 0;
+    }
+
+  dst_frame = avcodec_alloc_frame ();
+  if (dst_frame == NULL)
+    {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to allocate the destination image frame\n");
+#endif
+      sws_freeContext(scaler_ctx);
+      return 0;
+    }
+  dst_buffer =
+    av_malloc (avpicture_get_size (PIX_FMT_RGB24, dst_width, dst_height));
+  if (dst_buffer == NULL)
+    {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to allocate the destination image buffer\n");
+#endif
+      av_free (dst_frame);
+      sws_freeContext(scaler_ctx);
+      return 0;
+    }
+  avpicture_fill ((AVPicture *) dst_frame, dst_buffer,
+                  PIX_FMT_RGB24, dst_width, dst_height);
+      
+  sws_scale (scaler_ctx,
+             src_data, 
+             src_stride,
+             0, src_height, 
+             dst_frame->data, 
+             dst_frame->linesize);
+
+  encoder_output_buffer_size = output_max_size;
+  encoder_output_buffer = av_malloc (encoder_output_buffer_size);
+  if (encoder_output_buffer == NULL)
+    {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to allocate the encoder output buffer\n");
+#endif
+      av_free (dst_buffer);
+      av_free (dst_frame);
+      sws_freeContext(scaler_ctx);
+      return 0;
+    }
+
+  encoder_codec_ctx = avcodec_alloc_context ();
+  if (encoder_codec_ctx == NULL)
+    {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to allocate the encoder codec context\n");
+#endif
+      av_free (encoder_output_buffer);
+      av_free (dst_buffer);
+      av_free (dst_frame);
+      sws_freeContext(scaler_ctx);
+      return 0;
+    }
+  encoder_codec_ctx->width = dst_width;
+  encoder_codec_ctx->height = dst_height;
+  encoder_codec_ctx->pix_fmt = PIX_FMT_RGB24;
+
+  if (avcodec_open (encoder_codec_ctx, encoder_codec) < 0)
+    {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to open the encoder\n");
+#endif
+      av_free (encoder_codec_ctx);
+      av_free (encoder_output_buffer);
+      av_free (dst_buffer);
+      av_free (dst_frame);
+      sws_freeContext(scaler_ctx);
+      return 0;
+    }
+
+  err = avcodec_encode_video (encoder_codec_ctx,
+                              encoder_output_buffer,
+                              encoder_output_buffer_size, dst_frame);
+
+  avcodec_close (encoder_codec_ctx);
+  av_free (encoder_codec_ctx);
+  av_free (dst_buffer);
+  av_free (dst_frame);
+  sws_freeContext(scaler_ctx);
+
+  *output_data = encoder_output_buffer;
+
+  return err < 0 ? 0 : err;
+}
+
 int 
 EXTRACTOR_thumbnailffmpeg_extract (const unsigned char *data,
                                   size_t size,
@@ -75,22 +212,14 @@
 {
   AVProbeData pd; 
   AVPacket packet;
-  AVInputFormat *m_pInputFormat;
-  ByteIOContext *bio;
-  struct AVFormatContext *fmt;
+  AVInputFormat *input_format;
+  int input_format_nofileflag;
+  ByteIOContext *bio_ctx;
+  struct AVFormatContext *format_ctx;
   AVCodecContext *codec_ctx;
   AVCodec *codec = NULL;
   AVFrame *frame = NULL;
-  AVFrame *thumb_frame = NULL;
-  uint8_t *encoder_output_buffer = NULL;
-  AVCodecContext *enc_codec_ctx = NULL;
-  AVCodec *enc_codec = NULL;
-  const AVFrame *tframe;
-  struct SwsContext *scaler_ctx = NULL;
-  uint8_t *thumb_buffer = NULL;
-  int sws_flags = SWS_BILINEAR;
-  int64_t ts;
-  size_t encoder_output_buffer_size;
+  uint8_t *encoded_thumbnail;
   int video_stream_index;
   int thumb_width;
   int thumb_height;
@@ -105,42 +234,59 @@
   fprintf (stderr,
           "ffmpeg starting\n");
 #endif
-  pd.buf_size = size;
+  /* probe format
+   * initial try with a smaller probe size for efficiency */
+  pd.filename = "";
   pd.buf = (void *) data; 
-  pd.filename = "__url_prot";
- 
-  if (NULL == (m_pInputFormat = av_probe_input_format(&pd, 1))) 
+  pd.buf_size = 128*1024 > size ? size : 128*1024;
+RETRY_PROBE: 
+  if (NULL == (input_format = av_probe_input_format(&pd, 1))) 
     {
 #if DEBUG
       fprintf (stderr,
               "Failed to probe input format\n");
 #endif
+      if (pd.buf_size != size) /* retry probe once with full data size */
+        {
+         pd.buf_size = size;
+          goto RETRY_PROBE;
+        }
       return 0;
     }
-  m_pInputFormat->flags |= AVFMT_NOFILE;
-  bio = NULL; 
-  url_open_buf(&bio, pd.buf, pd.buf_size, URL_RDONLY);
-  bio->is_streamed = 1;  
-  if ((av_open_input_stream(&fmt, bio, pd.filename, m_pInputFormat, NULL)) < 0)
+  input_format_nofileflag = input_format->flags & AVFMT_NOFILE;
+  input_format->flags |= AVFMT_NOFILE;
+  bio_ctx = NULL; 
+  pd.buf_size = size;
+  url_open_buf(&bio_ctx, pd.buf, pd.buf_size, URL_RDONLY);
+  bio_ctx->is_streamed = 1;  
+  if ((av_open_input_stream(&format_ctx, bio_ctx, pd.filename, input_format, 
NULL)) < 0)
     {
-      url_close_buf (bio);
+ #if DEBUG
       fprintf (stderr,
               "Failed to open input stream\n");
+#endif
+      url_close_buf (bio_ctx);
+      if (!input_format_nofileflag)
+        input_format->flags ^= AVFMT_NOFILE;
       return 0;
     }
-  if (0 > av_find_stream_info (fmt))
+  if (0 > av_find_stream_info (format_ctx))
     {
-      av_close_input_stream (fmt);
-      url_close_buf (bio);
+ #if DEBUG
       fprintf (stderr,
-              "Failed to read stream info\n");
+               "Failed to read stream info\n");
+#endif
+      av_close_input_stream (format_ctx);
+      url_close_buf (bio_ctx);
+      if (!input_format_nofileflag)
+        input_format->flags ^= AVFMT_NOFILE;
       return 0;
     }
 
   codec_ctx = NULL;
-  for (i=0; i<fmt->nb_streams; i++)
+  for (i=0; i<format_ctx->nb_streams; i++)
     {
-      codec_ctx = fmt->streams[i]->codec;
+      codec_ctx = format_ctx->streams[i]->codec;
       if (codec_ctx->codec_type != CODEC_TYPE_VIDEO)
        continue;
       codec = avcodec_find_decoder (codec_ctx->codec_id);
@@ -160,39 +306,46 @@
        (codec_ctx->width == 0) || 
        (codec_ctx->height == 0) )
     {
+#if DEBUG
+      fprintf (stderr,
+               "No video streams or no suitable codec found\n");
+#endif
       if (codec_ctx != NULL)
        avcodec_close (codec_ctx);
-      av_close_input_stream (fmt);
-      url_close_buf (bio);
-      fprintf (stderr,
-              "No video codec found\n");
+      av_close_input_stream (format_ctx);
+      url_close_buf (bio_ctx);
+      if (!input_format_nofileflag)
+        input_format->flags ^= AVFMT_NOFILE;
       return 0;
     }
 
   frame = avcodec_alloc_frame ();
   if (frame == NULL)
     {
+#if DEBUG
+      fprintf (stderr,
+               "Failed to allocate frame\n");
+#endif
       if (codec != NULL)
        avcodec_close (codec_ctx);
-      av_close_input_stream (fmt);
-      url_close_buf (bio);
-      fprintf (stderr,
-              "Failed to allocate frame\n");
+      av_close_input_stream (format_ctx);
+      url_close_buf (bio_ctx);
+      if (!input_format_nofileflag)
+        input_format->flags ^= AVFMT_NOFILE;
       return 0;
     }
 #if DEBUG
-  if (fmt->duration == AV_NOPTS_VALUE)
+  if (format_ctx->duration == AV_NOPTS_VALUE)
     fprintf (stderr,
-            "duration unknown\n");
+            "Duration unknown\n");
   else
     fprintf (stderr,
-            "duration: %lld\n", 
-            fmt->duration);      
+            "Duration: %lld\n", 
+            format_ctx->duration);      
 #endif
-  /* TODO: if duration is known seek to to some better place(?) */
-  ts = 10;                  /* s */
-  ts = ts * AV_TIME_BASE;
-  err = av_seek_frame (fmt, -1, ts, 0);
+  /* TODO: if duration is known, seek to some better place,
+   * but use 10 sec into stream for now */
+  err = av_seek_frame (format_ctx, -1, 10 * AV_TIME_BASE, 0);
   if (err >= 0)        
     avcodec_flush_buffers (codec_ctx);        
   frame_finished = 0;
@@ -205,7 +358,7 @@
     {
       while (1)
         {
-          err = av_read_frame (fmt, &packet);
+          err = av_read_frame (format_ctx, &packet);
           if (err < 0)
             break;
           if (packet.stream_index == video_stream_index)
@@ -225,15 +378,18 @@
     }
   if (!frame_finished)
     {
+      fprintf (stderr,
+              "Failed to decode a complete frame\n");
       if (frame != NULL)
        av_free (frame);
-      av_close_input_stream (fmt);
-      free (bio);
-      fprintf (stderr,
-              "Failed to seek to frame\n");
+      av_close_input_stream (format_ctx);
+      free (bio_ctx);
+      if (!input_format_nofileflag)
+        input_format->flags ^= AVFMT_NOFILE;
       return 0;
     }
 
+  /* calculate the thumbnail dimensions, taking pixel aspect into account */
   sar_num = codec_ctx->sample_aspect_ratio.num;
   sar_den = codec_ctx->sample_aspect_ratio.den;
   if (sar_num <= 0 || sar_den <= 0)
@@ -241,143 +397,49 @@
       sar_num = 1;
       sar_den = 1;
     }
-  if ( (codec_ctx->width < THUMBSIZE) &&
-       (codec_ctx->height < THUMBSIZE) )
+  if ((codec_ctx->width * sar_num) / sar_den > codec_ctx->height)
     {
-      /* no resize */
-      thumb_width = codec_ctx->width;
-      thumb_height = codec_ctx->height;
-      tframe = frame;
+      thumb_width = THUMBSIZE;
+      thumb_height = (thumb_width * codec_ctx->height) /
+                     ((codec_ctx->width * sar_num) / sar_den);
     }
   else
     {
-      /* need resize */
-      if ((codec_ctx->width * sar_num) / sar_den > codec_ctx->height)
-       {
-         thumb_width = THUMBSIZE;
-         thumb_height = (thumb_width * codec_ctx->height) /
-           ((codec_ctx->width * sar_num) / sar_den);
-       }
-      else
-       {
-         thumb_height = THUMBSIZE;
-         thumb_width = (thumb_height *
-                        ((codec_ctx->width * sar_num) / sar_den)) /
-           codec_ctx->height;
-       }
-      if (thumb_width < 8)
-       thumb_width = 8;
-      if (thumb_height < 1)
-       thumb_height = 1;
-#if DEBUG
-      fprintf (stderr,
-              "thumb dim: %d %d\n", 
-              thumb_width,
-              thumb_height);
-#endif
-
-      scaler_ctx =
-       sws_getContext (codec_ctx->width, codec_ctx->height, codec_ctx->pix_fmt,
-                       thumb_width, thumb_height, PIX_FMT_RGB24, sws_flags, 
NULL,
-                       NULL, NULL);
-      if (scaler_ctx == NULL)
-       {
-#if DEBUG
-         fprintf (stderr,
-                  "failed to alloc scaler\n");
-#endif
-         goto out;
-       }
-      thumb_frame = avcodec_alloc_frame ();
-      thumb_buffer =
-       av_malloc (avpicture_get_size (PIX_FMT_RGB24, thumb_width, 
thumb_height));
-      if (thumb_frame == NULL || thumb_buffer == NULL)
-       {
-#if DEBUG
-         fprintf (stderr,
-                  "failed to alloc thumb frame\n");
-#endif
-         goto out;
-       }
-      avpicture_fill ((AVPicture *) thumb_frame, thumb_buffer,
-                     PIX_FMT_RGB24, thumb_width, thumb_height);
-      
-      sws_scale (scaler_ctx,
-                frame->data, 
-                frame->linesize,
-                0, codec_ctx->height, 
-                thumb_frame->data, 
-                thumb_frame->linesize);
-      tframe = thumb_frame;  
+      thumb_height = THUMBSIZE;
+      thumb_width = (thumb_height *
+                    ((codec_ctx->width * sar_num) / sar_den)) /
+      codec_ctx->height;
     }
-
-  encoder_output_buffer_size = MAX_THUMB_SIZE;
-  encoder_output_buffer = av_malloc (encoder_output_buffer_size);
-  if (encoder_output_buffer == NULL)
-    {
+  if (thumb_width < 8)
+    thumb_width = 8;
+  if (thumb_height < 1)
+    thumb_height = 1;
 #if DEBUG
-      fprintf (stderr,
-              "couldn't alloc encoder output buf\n");
+  fprintf (stderr,
+           "Thumbnail dimensions: %d %d\n", 
+           thumb_width, thumb_height);
 #endif
-      goto out;
-    }
 
-  enc_codec = avcodec_find_encoder_by_name ("png");
-  if (enc_codec == NULL)
-    {
-#if DEBUG
-      fprintf (stderr,
-              "couldn't find encoder\n");
-#endif
-      goto out;
-    }
-  enc_codec_ctx = avcodec_alloc_context ();
-  enc_codec_ctx->width = thumb_width;
-  enc_codec_ctx->height = thumb_height;
-  enc_codec_ctx->pix_fmt = PIX_FMT_RGB24;
+  err = create_thumbnail (codec_ctx->width, codec_ctx->height,
+                          frame->linesize, codec_ctx->pix_fmt, frame->data,
+                          thumb_width, thumb_height,
+                          &encoded_thumbnail, MAX_THUMB_SIZE);
 
-  if (avcodec_open (enc_codec_ctx, enc_codec) < 0)
-    {
-#if DEBUG
-      fprintf (stderr,
-              "couldn't open encoder\n");
-#endif
-      enc_codec = NULL;
-      goto out;
-    }
-
-  err = avcodec_encode_video (enc_codec_ctx,
-                              encoder_output_buffer,
-                              encoder_output_buffer_size, tframe);
   if (err > 0)
     ret = proc (proc_cls,
                "thumbnailffmpeg",
                EXTRACTOR_METATYPE_THUMBNAIL,
                EXTRACTOR_METAFORMAT_BINARY,
                "image/png",
-               (const char*) encoder_output_buffer,
+               (const char*) encoded_thumbnail,
                err);
-out:
-  if (enc_codec != NULL)
-    avcodec_close (enc_codec_ctx);
-  if (enc_codec_ctx != NULL)
-    av_free (enc_codec_ctx);
-  if (encoder_output_buffer != NULL)
-    av_free (encoder_output_buffer);
-  if (scaler_ctx != NULL)
-    sws_freeContext(scaler_ctx);
-  if (codec_ctx != NULL)
-    avcodec_close (codec_ctx);
-  if (fmt != NULL)
-    av_close_input_stream (fmt);
-  if (frame != NULL)
-    av_free (frame);
-  if (thumb_buffer != NULL)
-    av_free (thumb_buffer);
-  if (thumb_frame != NULL)
-    av_free (thumb_frame);
-  if (bio != NULL)
-    url_close_buf (bio);
+
+  avcodec_close (codec_ctx);
+  av_close_input_stream (format_ctx);
+  if (!input_format_nofileflag)
+    input_format->flags ^= AVFMT_NOFILE;
+  av_free (frame);
+  url_close_buf (bio_ctx);
   return ret;
 }