mpegtsdemux: payload parsing for H.264

gst/mpegtsdemux/payload_parsers.c

@@ -27,6 +27,11 @@
 #define PICTURE_START_CODE 0x00000100
 #define GROUP_START_CODE   0x000001B8
 
+#define SLICE_NAL_UNIT_TYPE     0x01
+#define SLICE_IDR_NAL_UNIT_TYPE 0x05
+#define SEI_NAL_UNIT_TYPE       0x06
+
+#define SEI_TYPE_RECOVERY_POINT 0x06
 
 typedef struct Mpeg2PictureHeader
 {
@@ -43,6 +48,13 @@ typedef struct Mpeg2PictureHeader
   guint8 backward_f_code;
 } Mpeg2PictureHeader;
 
+/* shortened slice header */
+typedef struct H264SliceHeader
+{
+  guint32 first_mb_in_slice;
+  guint8 slice_type;
+} H264SliceHeader;
+
 
 static guint8 *
 find_start_code (guint32 * start_code, guint8 * buffer, guint8 * buffer_end)
@@ -93,8 +105,6 @@ gboolean
 gst_tsdemux_has_mpeg2_keyframe (guint32 * state,
     MpegTSPacketizerPacket * packet)
 {
-
-  //guint32 i = 0;
   guint8 *data = packet->payload;
   guint8 *data_end = packet->data_end;
 
@@ -116,13 +126,183 @@ gst_tsdemux_has_mpeg2_keyframe (guint32 * state,
     } 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);
+
+      *state = 0xffffffff;
       return success && hdr.picture_coding_type == 1;
     }
   }
 
   return FALSE;
 }
+
+/* variable length Exp-Golomb parsing according to H.264 spec 9.1*/
+static gboolean
+read_golomb (GstBitReader * br, guint32 * value)
+{
+  guint8 b, leading_zeros = -1;
+  *value = 1;
+
+  for (b = 0; !b; leading_zeros++) {
+    if (!gst_bit_reader_get_bits_uint8 (br, &b, 1))
+      return FALSE;
+    *value *= 2;
+  }
+
+  *value = (*value >> 1) - 1;
+  if (leading_zeros > 0) {
+    guint32 tmp = 0;
+    if (!gst_bit_reader_get_bits_uint32 (br, &tmp, leading_zeros))
+      return FALSE;
+    *value += tmp;
+  }
+
+  return TRUE;
+}
+
+/* just parse the requirred bits of the slice header */
+static gboolean
+parse_h264_slice_header (H264SliceHeader * hdr, guint8 * buffer,
+    guint8 * buffer_end)
+{
+  guint32 value;
+  GstBitReader br = GST_BIT_READER_INIT (buffer, buffer_end - buffer);
+
+  if (!read_golomb (&br, &value))
+    return FALSE;
+  hdr->first_mb_in_slice = value;
+
+  if (!read_golomb (&br, &value))
+    return FALSE;
+  hdr->slice_type = value;
+
+  return TRUE;
+}
+
+enum H264SliceTypes
+{
+  h264_p_slice = 0,
+  h264_b_slice,
+  h264_i_slice,
+  h264_sp_slice,
+  h264_si_slice,
+  h264_p_slice_a,
+  h264_b_slice_a,
+  h264_i_slice_a,
+  h264_sp_slice_a,
+  h264_si_slice_a,
+};
+
+static gboolean
+is_key_slice (guint8 slice_type)
+{
+  switch (slice_type) {
+    case h264_i_slice:
+    case h264_si_slice:
+    case h264_i_slice_a:
+    case h264_si_slice_a:
+      return TRUE;
+  }
+  return FALSE;
+}
+
+gboolean
+gst_tsdemux_has_h264_keyframe (guint32 * state, MpegTSPacketizerPacket * packet)
+{
+  guint8 *data = packet->payload;
+  guint8 *data_end = packet->data_end;
+
+  GST_LOG ("state: 0x%08x", *state);
+
+  while (data <= data_end) {
+    guint8 nal_unit_type;
+    guint8 *next_data = NULL;
+
+    data = find_start_code (state, data, data_end);
+
+    if (!data)
+      goto beach;
+
+    GST_LOG ("found start code: 0x%08x", *state);
+
+    /* determine length */
+    nal_unit_type = *state & 0x1f;
+    next_data = find_start_code (state, data, data_end);
+
+    if (nal_unit_type == SEI_NAL_UNIT_TYPE && !next_data) {
+      GST_WARNING ("NAL unit 0x%02x not completely in ts packet",
+          nal_unit_type);
+      goto beach;
+    }
+    next_data -= 4;
+
+    switch (nal_unit_type) {
+      case SLICE_IDR_NAL_UNIT_TYPE:
+        GST_DEBUG ("found SLICE_IDR NAL unit type");
+        *state = 0xffffffff;
+        return TRUE;
+      case SLICE_NAL_UNIT_TYPE:
+      {
+        H264SliceHeader hdr = { 0 };
+        gboolean success;
+
+        success = parse_h264_slice_header (&hdr, data, data_end);
+        GST_DEBUG ("found SLICE NAL unit type with slice type %d",
+            hdr.slice_type);
+
+        *state = 0xffffffff;
+        return success && is_key_slice (hdr.slice_type);
+      }
+      case SEI_NAL_UNIT_TYPE:
+      {
+        guint32 recovery_frame_count;
+        GstBitReader br = GST_BIT_READER_INIT (data, next_data - data);
+
+        break;
+
+        /* SEI message is at least 24 bit long */
+        while (gst_bit_reader_get_remaining (&br) >= 24) {
+          gint type = 0, size = 0;
+          guint8 tmp = 0;
+
+          do {
+            if (!gst_bit_reader_get_bits_uint8 (&br, &tmp, 8))
+              goto beach;
+            type += tmp;
+          } while (tmp == 255);
+
+          do {
+            if (!gst_bit_reader_get_bits_uint8 (&br, &tmp, 8))
+              goto beach;
+            size += tmp;
+          } while (tmp == 255);
+
+
+          GST_LOG ("found SEI msg type: %d, len: %d", type, size);
+
+          switch (type) {
+            case SEI_TYPE_RECOVERY_POINT:
+              if (!read_golomb (&br, &recovery_frame_count))
+                return FALSE;
+              gst_bit_reader_skip (&br, 1);     /* exact_match */
+              gst_bit_reader_skip (&br, 1);     /* broken_link_flag */
+              gst_bit_reader_skip (&br, 2);     /* changing_slice_group_idc */
+              GST_DEBUG ("found SEI with recovery point message, "
+                  "recovery_frame_count: %d", recovery_frame_count);
+              return TRUE;
+            default:
+              /* skip all other sei messages */
+              gst_bit_reader_skip (&br, size * 8);
+          }
+        }
+      }
+        data = next_data;
+        *state = 0xffffffff;
+    }
+  }
+beach:
+  return FALSE;
+}

gst/mpegtsdemux/payload_parsers.h

@@ -27,3 +27,6 @@ typedef gboolean (*payload_parse_keyframe) (guint32 *state, MpegTSPacketizerPack
 
 gboolean
 gst_tsdemux_has_mpeg2_keyframe (guint32 *state, MpegTSPacketizerPacket * packet);
+
+gboolean
+gst_tsdemux_has_h264_keyframe (guint32 *state, MpegTSPacketizerPacket * packet);

gst/mpegtsdemux/tsdemux.c

@@ -792,8 +792,10 @@ gst_ts_demux_perform_seek (MpegTSBase * base, GstSegment * segment, guint16 pid)
         case ST_VIDEO_MPEG2:
           keyframe_seek = gst_tsdemux_has_mpeg2_keyframe;
           break;
-        case ST_VIDEO_MPEG4:
         case ST_VIDEO_H264:
+          keyframe_seek = gst_tsdemux_has_h264_keyframe;
+          break;
+        case ST_VIDEO_MPEG4:
         case ST_VIDEO_DIRAC:
           GST_WARNING ("no payload parser for stream 0x%04x type: 0x%02x", pid,
               program->streams[pid]->stream_type);