mpegtsdemux: implement key_unit seeking for MPEG2 video

5 files changed, 230 insertions, 29 deletions

diff --git a/gst/mpegtsdemux/Makefile.am b/gst/mpegtsdemux/Makefile.am
index 184b588dd..7d3e66325 100644
--- a/gst/mpegtsdemux/Makefile.am
+++ b/gst/mpegtsdemux/Makefile.am
@@ -6,6 +6,7 @@ libgstmpegtsdemux_la_SOURCES = \
 	mpegtsbase.c	\
 	mpegtspacketizer.c \
 	mpegtsparse.c \
+	payload_parsers.c \
 	tsdemux.c
 
 libgstmpegtsdemux_la_CFLAGS = \
@@ -23,6 +24,7 @@ noinst_HEADERS = \
 	mpegtsbase.h	\
 	mpegtspacketizer.h \
 	mpegtsparse.h \
+	payload_parsers.h \
 	tsdemux.h
 
 Android.mk: Makefile.am $(BUILT_SOURCES)
@@ -37,4 +39,4 @@ Android.mk: Makefile.am $(BUILT_SOURCES)
 	           -ldl \
 	 -:PASSTHROUGH LOCAL_ARM_MODE:=arm \
 		       LOCAL_MODULE_PATH:='$$(TARGET_OUT)/lib/gstreamer-0.10' \
-	> $@
\ No newline at end of file
+	> $@
diff --git a/gst/mpegtsdemux/mpegtsbase.c b/gst/mpegtsdemux/mpegtsbase.c
index 7e272e02f..866f0d4cc 100644
--- a/gst/mpegtsdemux/mpegtsbase.c
+++ b/gst/mpegtsdemux/mpegtsbase.c
@@ -1272,8 +1272,7 @@ mpegts_base_handle_seek_event (MpegTSBase * base, GstPad * pad,
     gst_pad_push_event (base->sinkpad, gst_event_new_flush_stop ());
   }
 
-  if (flags & (GST_SEEK_FLAG_KEY_UNIT | GST_SEEK_FLAG_SEGMENT |
-          GST_SEEK_FLAG_SKIP)) {
+  if (flags & (GST_SEEK_FLAG_SEGMENT | GST_SEEK_FLAG_SKIP)) {
     GST_WARNING ("seek flags 0x%x are not supported", (int) flags);
     goto done;
   }
diff --git a/gst/mpegtsdemux/payload_parsers.c b/gst/mpegtsdemux/payload_parsers.c
new file mode 100644
index 000000000..b2bd0aa6f
--- /dev/null
+++ b/gst/mpegtsdemux/payload_parsers.c
@@ -0,0 +1,128 @@
+/*
+ * payload_parsers.c
+ * Copyright (C) 2011 Janne Grunau
+ *
+ * Authors:
+ *   Janne Grunau <janne.grunau@collabora.co.uk>
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Library General Public
+ * License as published by the Free Software Foundation; either
+ * version 2 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Library General Public License for more details.
+ *
+ * You should have received a copy of the GNU Library General Public
+ * License along with this library; if not, write to the
+ * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
+ * Boston, MA 02111-1307, USA.
+ */
+
+#include "payload_parsers.h"
+#include <gst/base/gstbitreader.h>
+
+#define PICTURE_START_CODE 0x00000100
+#define GROUP_START_CODE   0x000001B8
+
+
+typedef struct Mpeg2PictureHeader
+{
+  guint16 temporal_reference;
+  guint8 picture_coding_type;
+  guint16 vbv_delay;
+
+  /* picture_coding_type == 2 || picture_coding_type */
+  guint8 full_pel_forward_vector;
+  guint8 forward_f_code;
+
+  /* picture_coding_type == 3 */
+  guint8 full_pel_backward_vector;
+  guint8 backward_f_code;
+} Mpeg2PictureHeader;
+
+
+static guint8 *
+find_start_code (guint32 * start_code, guint8 * buffer, guint8 * buffer_end)
+{
+  if (G_UNLIKELY (buffer == NULL) || G_UNLIKELY (buffer_end == NULL)
+      || G_UNLIKELY (start_code == NULL))
+    return NULL;
+
+  while (buffer <= buffer_end) {
+
+    *start_code <<= 8;
+    *start_code |= *buffer++;
+
+    if ((*start_code & 0xffffff00) == 0x00000100)
+      return buffer;
+  }
+
+  return NULL;
+}
+
+static gboolean
+parse_mpeg2_picture_header (Mpeg2PictureHeader * hdr, guint8 * buffer,
+    guint8 * buffer_end)
+{
+  GstBitReader br = GST_BIT_READER_INIT (buffer, buffer_end - buffer);
+
+  if (gst_bit_reader_get_remaining (&br) < 40)
+    return FALSE;
+
+  hdr->temporal_reference = gst_bit_reader_get_bits_uint16_unchecked (&br, 10);
+  hdr->picture_coding_type = gst_bit_reader_get_bits_uint8_unchecked (&br, 3);
+  hdr->vbv_delay = gst_bit_reader_get_bits_uint16_unchecked (&br, 16);
+
+  if (hdr->picture_coding_type == 2 || hdr->picture_coding_type == 3) {
+    hdr->full_pel_forward_vector =
+        gst_bit_reader_get_bits_uint8_unchecked (&br, 1);
+    hdr->forward_f_code = gst_bit_reader_get_bits_uint8_unchecked (&br, 3);
+  }
+  if (hdr->picture_coding_type == 3) {
+    hdr->full_pel_backward_vector =
+        gst_bit_reader_get_bits_uint8_unchecked (&br, 1);
+    hdr->backward_f_code = gst_bit_reader_get_bits_uint8_unchecked (&br, 3);
+  }
+  return TRUE;
+}
+
+gboolean
+gst_tsdemux_has_mpeg2_keyframe (guint32 * state,
+    MpegTSPacketizerPacket * packet)
+{
+
+  //guint32 i = 0;
+  guint8 *data = packet->payload;
+  guint8 *data_end = packet->data_end;
+
+  GST_LOG ("state: 0x%08x", *state);
+
+  while (data <= data_end) {
+
+    data = find_start_code (state, data, data_end);
+
+    if (!data)
+      return FALSE;
+
+    GST_LOG ("found start code: 0x%08x", *state);
+
+    if (*state == GROUP_START_CODE) {
+      GST_DEBUG ("found group start code");
+      *state = 0xffffffff;
+      return TRUE;
+    } else if (*state == PICTURE_START_CODE) {
+      Mpeg2PictureHeader hdr = { 0 };
+      gboolean success;
+      *state = 0xffffffff;
+      success = parse_mpeg2_picture_header (&hdr, data, data_end);
+      GST_DEBUG ("found picture start code, %sparsed, picture coding type: %d",
+          success ? "" : "not ", hdr.picture_coding_type);
+      return success && hdr.picture_coding_type == 1;
+    }
+  }
+
+  return FALSE;
+}
diff --git a/gst/mpegtsdemux/payload_parsers.h b/gst/mpegtsdemux/payload_parsers.h
new file mode 100644
index 000000000..629bdd9e3
--- /dev/null
+++ b/gst/mpegtsdemux/payload_parsers.h
@@ -0,0 +1,29 @@
+/*
+ * payload_parsers.h
+ * Copyright (C) 2011 Janne Grunau
+ *
+ * Authors:
+ *   Janne Grunau <janne.grunau@collabora.co.uk>
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Library General Public
+ * License as published by the Free Software Foundation; either
+ * version 2 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Library General Public License for more details.
+ *
+ * You should have received a copy of the GNU Library General Public
+ * License along with this library; if not, write to the
+ * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
+ * Boston, MA 02111-1307, USA.
+ */
+
+#include "mpegtspacketizer.h"
+
+typedef gboolean (*payload_parse_keyframe) (guint32 *state, MpegTSPacketizerPacket * packet);
+
+gboolean
+gst_tsdemux_has_mpeg2_keyframe (guint32 *state, MpegTSPacketizerPacket * packet);
diff --git a/gst/mpegtsdemux/tsdemux.c b/gst/mpegtsdemux/tsdemux.c
index 83f7e983c..8b5bf23b8 100644
--- a/gst/mpegtsdemux/tsdemux.c
+++ b/gst/mpegtsdemux/tsdemux.c
@@ -37,6 +37,7 @@
 #include "gstmpegdesc.h"
 #include "gstmpegdefs.h"
 #include "mpegtspacketizer.h"
+#include "payload_parsers.h"
 
 /* latency in mseconds */
 #define TS_LATENCY 700
@@ -515,25 +516,34 @@ discont:
   return res;
 }
 
-
-/* performs a accurate seek to the last packet with pts < seektime */
+/* performs a accurate/key_unit seek */
 static GstFlowReturn
-gst_ts_demux_perform_accurate_seek (MpegTSBase * base, GstClockTime seektime,
-    TSPcrOffset * pcroffset, gint64 length, gint16 pid)
+gst_ts_demux_perform_auxiliary_seek (MpegTSBase * base, GstClockTime seektime,
+    TSPcrOffset * pcroffset, gint64 length, gint16 pid, GstSeekFlags flags,
+    payload_parse_keyframe auxiliary_seek_fn)
 {
   GstTSDemux *demux = (GstTSDemux *) base;
   GstFlowReturn res = GST_FLOW_ERROR;
   gboolean done = FALSE;
+  gboolean found_keyframe = FALSE, found_accurate = FALSE;
   GstBuffer *buf;
   MpegTSPacketizerPacket packet;
   MpegTSPacketizerPacketReturn pret;
   gint64 offset = pcroffset->offset;
   gint64 scan_offset = MIN (length, 50 * MPEGTS_MAX_PACKETSIZE);
+  guint32 state = 0xffffffff;
+  TSPcrOffset key_pos = { 0 };
 
+  GST_DEBUG ("auxiliary seek for %" GST_TIME_FORMAT " from offset: %"
+      G_GINT64_FORMAT " in %" G_GINT64_FORMAT " bytes for PID: %d "
+      "%s %s", GST_TIME_ARGS (seektime), pcroffset->offset, length, pid,
+      (flags & GST_SEEK_FLAG_ACCURATE) ? "accurate" : "",
+      (flags & GST_SEEK_FLAG_KEY_UNIT) ? "key_unit" : "");
 
-  GST_DEBUG ("accurate seek for %" GST_TIME_FORMAT " from offset: %"
-      G_GINT64_FORMAT " in %" G_GINT64_FORMAT " bytes for PID: %d",
-      GST_TIME_ARGS (seektime), pcroffset->offset, length, pid);
+  if ((flags & GST_SEEK_FLAG_KEY_UNIT) && !auxiliary_seek_fn) {
+    GST_ERROR ("key_unit seek for unkown video codec");
+    goto beach;
+  }
 
   mpegts_packetizer_flush (base->packetizer);
 
@@ -558,25 +568,52 @@ gst_ts_demux_perform_accurate_seek (MpegTSBase * base, GstClockTime seektime,
             " at offset: %" G_GINT64_FORMAT, packet.pid,
             GST_TIME_ARGS (packet.pcr), packet.offset);
 
-      if (packet.payload != NULL && packet.payload_unit_start_indicator
-          && packet.pid == pid) {
-        guint64 pts = 0;
+      if (packet.payload != NULL && packet.pid == pid) {
 
-        res = gst_ts_demux_parse_pes_header_pts (demux, &packet, &pts);
-        if (res == GST_FLOW_OK) {
-          GstClockTime time = calculate_gsttime (pcroffset, pts * 300);
+        if (packet.payload_unit_start_indicator) {
+          guint64 pts = 0;
+          res = gst_ts_demux_parse_pes_header_pts (demux, &packet, &pts);
+          if (res == GST_FLOW_OK) {
+            GstClockTime time = calculate_gsttime (pcroffset, pts * 300);
 
-          GST_DEBUG ("packet has PTS: %" GST_TIME_FORMAT,
+            GST_DEBUG ("packet has PTS: %" GST_TIME_FORMAT,
               GST_TIME_ARGS (time));
 
-          if (time <= seektime) {
-            pcroffset->gsttime = time;
-            pcroffset->pcr = packet.pcr;
-            pcroffset->offset = packet.offset;
+            if (time <= seektime) {
+              pcroffset->gsttime = time;
+              pcroffset->pcr = packet.pcr;
+              pcroffset->offset = packet.offset;
+            } else
+              found_accurate = TRUE;
           } else
-            done = TRUE;
-        } else
-          goto next;
+            goto next;
+          /* reset state for new packet */
+          state = 0xffffffff;
+        }
+
+        if (flags & GST_SEEK_FLAG_KEY_UNIT) {
+          gboolean is_keyframe = auxiliary_seek_fn (&state, &packet);
+          if (is_keyframe) {
+            found_keyframe = TRUE;
+            key_pos = *pcroffset;
+            GST_DEBUG ("found keyframe: time: %" GST_TIME_FORMAT " pcr: %"
+                GST_TIME_FORMAT " offset %" G_GINT64_FORMAT,
+                GST_TIME_ARGS (pcroffset->gsttime),
+                GST_TIME_ARGS (pcroffset->pcr), pcroffset->offset);
+          }
+        }
+      }
+      switch (flags & (GST_SEEK_FLAG_ACCURATE | GST_SEEK_FLAG_KEY_UNIT)) {
+        case GST_SEEK_FLAG_ACCURATE | GST_SEEK_FLAG_KEY_UNIT:
+          done = found_accurate && found_keyframe;
+          *pcroffset = key_pos;
+          break;
+        case GST_SEEK_FLAG_ACCURATE:
+          done = found_accurate;
+          break;
+        case GST_SEEK_FLAG_KEY_UNIT:
+          done = found_keyframe;
+          break;
       }
     next:
       mpegts_packetizer_clear_packet (base->packetizer, &packet);
@@ -584,6 +621,9 @@ gst_ts_demux_perform_accurate_seek (MpegTSBase * base, GstClockTime seektime,
     scan_offset += 50 * MPEGTS_MAX_PACKETSIZE;
   }
 
+  if (done)
+    res = GST_FLOW_OK;
+
 beach:
   mpegts_packetizer_flush (base->packetizer);
   return res;
@@ -741,13 +781,17 @@ gst_ts_demux_perform_seek (MpegTSBase * base, GstSegment * segment, guint16 pid)
 
   GST_DEBUG ("seeking finished after %d loops", loop_cnt);
 
-  if (segment->flags & GST_SEEK_FLAG_ACCURATE) {
+  if (segment->flags & GST_SEEK_FLAG_ACCURATE
+      || segment->flags & GST_SEEK_FLAG_KEY_UNIT) {
+    payload_parse_keyframe keyframe_seek = NULL;
     MpegTSBaseProgram *program = demux->program;
 
     if (program->streams[pid]) {
       switch (program->streams[pid]->stream_type) {
         case ST_VIDEO_MPEG1:
         case ST_VIDEO_MPEG2:
+          keyframe_seek = gst_tsdemux_has_mpeg2_keyframe;
+          break;
         case ST_VIDEO_MPEG4:
         case ST_VIDEO_H264:
         case ST_VIDEO_DIRAC:
@@ -760,8 +804,8 @@ gst_ts_demux_perform_seek (MpegTSBase * base, GstSegment * segment, guint16 pid)
     seekpcroffset.pcr = pcr_start.pcr;
     seekpcroffset.offset = pcr_start.offset;
     res =
-        gst_ts_demux_perform_accurate_seek (base, seektime, &seekpcroffset,
-        pcr_stop.offset - pcr_start.offset, pid);
+        gst_ts_demux_perform_auxiliary_seek (base, seektime, &seekpcroffset,
+        pcr_stop.offset - pcr_start.offset, pid, segment->flags, keyframe_seek);
   }
 
   segment->last_stop = seekpcroffset.gsttime;
@@ -810,8 +854,7 @@ gst_ts_demux_do_seek (MpegTSBase * base, GstEvent * event, guint16 pid)
   accurate = flags & GST_SEEK_FLAG_ACCURATE;
   flush = flags & GST_SEEK_FLAG_FLUSH;
 
-  if (flags & (GST_SEEK_FLAG_KEY_UNIT | GST_SEEK_FLAG_SEGMENT |
-          GST_SEEK_FLAG_SKIP)) {
+  if (flags & (GST_SEEK_FLAG_SEGMENT | GST_SEEK_FLAG_SKIP)) {
     GST_WARNING ("seek flags 0x%x are not supported", (int) flags);
     goto done;
   }