set pts and dts of pes packets exactly according to specs

Originally committed as revision 16603 to svn://svn.ffmpeg.org/ffmpeg/trunk

2 files changed, 30 insertions, 4 deletions

diff --git a/libavformat/mpegtsenc.c b/libavformat/mpegtsenc.c
index dc8916baa6..0170a5acb8 100644
--- a/libavformat/mpegtsenc.c
+++ b/libavformat/mpegtsenc.c
@@ -21,6 +21,7 @@
 
 #include "libavutil/bswap.h"
 #include "libavutil/crc.h"
+#include "libavcodec/mpegvideo.h"
 #include "avformat.h"
 #include "mpegts.h"
 
@@ -670,6 +671,7 @@ static int mpegts_write_packet(AVFormatContext *s, AVPacket *pkt)
     uint8_t *buf= pkt->data;
     MpegTSWriteStream *ts_st = st->priv_data;
     int len, max_payload_size;
+    const uint8_t *access_unit_index = NULL;
 
     if (st->codec->codec_type == CODEC_TYPE_SUBTITLE) {
         /* for subtitle, a single PES packet must be generated */
@@ -683,6 +685,27 @@ static int mpegts_write_packet(AVFormatContext *s, AVPacket *pkt)
         return 0;
     }
     max_payload_size = DEFAULT_PES_PAYLOAD_SIZE;
+    if (st->codec->codec_id == CODEC_ID_MPEG2VIDEO ||
+        st->codec->codec_id == CODEC_ID_MPEG1VIDEO) {
+        const uint8_t *p = pkt->data;
+        const uint8_t *end = pkt->data+pkt->size;
+        uint32_t state = -1;
+        while (p < end) {
+            p = ff_find_start_code(p, end, &state);
+            if (state == PICTURE_START_CODE) {
+                access_unit_index = p - 4;
+                break;
+            }
+        }
+    } else if (st->codec->codec_type == CODEC_TYPE_AUDIO) {
+        access_unit_index = pkt->data;
+    }
+
+    if (!access_unit_index) {
+        av_log(s, AV_LOG_ERROR, "error, could not find access unit start\n");
+        return -1;
+    }
+
     while (size > 0) {
         len = max_payload_size - ts_st->payload_index;
         if (len > size)
@@ -691,16 +714,19 @@ static int mpegts_write_packet(AVFormatContext *s, AVPacket *pkt)
         buf += len;
         size -= len;
         ts_st->payload_index += len;
-        if (ts_st->payload_pts == AV_NOPTS_VALUE)
-            ts_st->payload_pts = pkt->pts;
-        if (ts_st->payload_dts == AV_NOPTS_VALUE)
+        if (access_unit_index && access_unit_index < buf &&
+            ts_st->payload_pts == AV_NOPTS_VALUE &&
+            ts_st->payload_dts == AV_NOPTS_VALUE) {
             ts_st->payload_dts = pkt->dts;
+            ts_st->payload_pts = pkt->pts;
+        }
         if (ts_st->payload_index >= max_payload_size) {
             mpegts_write_pes(s, st, ts_st->payload, ts_st->payload_index,
                              ts_st->payload_pts, ts_st->payload_dts);
             ts_st->payload_pts = AV_NOPTS_VALUE;
             ts_st->payload_dts = AV_NOPTS_VALUE;
             ts_st->payload_index = 0;
+            access_unit_index = NULL; // unset access unit to avoid setting pts/dts again
         }
     }
     return 0;
diff --git a/tests/libav.regression.ref b/tests/libav.regression.ref
index 22c1c40f92..412c652929 100644
--- a/tests/libav.regression.ref
+++ b/tests/libav.regression.ref
@@ -9,7 +9,7 @@ c351132527ccb1e8cab06cc0822fde23 *./tests/data/b-libav.rm
 bdb7484c68db722f66ba1630cf79844c *./tests/data/b-libav.mpg
 378880 ./tests/data/b-libav.mpg
 ./tests/data/b-libav.mpg CRC=0x2b71a386
-447b005e527cf495ec13092e788f028d *./tests/data/b-libav.ts
+d1ab4041e32fb802bb164844d91cc5fe *./tests/data/b-libav.ts
 471692 ./tests/data/b-libav.ts
 ./tests/data/b-libav.ts CRC=0xcc4948e1
 1b28a16652bb8ac528b33f7478ca18b6 *./tests/data/b-libav.swf