Merge branch 'master' into master-with-d

xtne6f · Mar 19, 2024 · b6d5f11 · b6d5f11
2 parents 0b3eee7 + b8b4673
commit b6d5f11
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 11 deletions.
diff --git a/mp4fragmenter.cpp b/mp4fragmenter.cpp
@@ -62,12 +62,22 @@ template<class P>
 void ParseNals(const uint8_t *payload, size_t payloadSize, P onNalProc)
 {
     size_t nalPos = 0;
-    for (size_t i = 0; i + 1 < payloadSize; ++i) {
-        if (i + 2 == payloadSize || (payload[i] == 0 && payload[i + 1] == 0 && payload[i + 2] == 1)) {
+    for (size_t i = 2;;) {
+        if (i >= payloadSize || (payload[i] == 1 && payload[i - 1] == 0 && payload[i - 2] == 0)) {
             if (nalPos != 0) {
-                onNalProc(payload + nalPos, (i + 2 == payloadSize ? payloadSize : i - (payload[i - 1] == 0)) - nalPos);
+                onNalProc(payload + nalPos, (i >= payloadSize ? payloadSize : i - 2 - (payload[i - 3] == 0)) - nalPos);
             }
-            nalPos = i + 3;
+            if (i >= payloadSize) {
+                break;
+            }
+            nalPos = i + 1;
+            i += 3;
+        }
+        else if (payload[i] > 0) {
+            i += 3;
+        }
+        else {
+            ++i;
         }
     }
 }
@@ -377,10 +387,26 @@ void CMp4Fragmenter::AddVideoPes(const std::vector<uint8_t> &pes, bool h265)
                         // Drop SEI
                     }
                     else {
-                        if (h265 ? (nalUnitType == 19 || nalUnitType == 20 || nalUnitType == 21) : (nalUnitType == 5)) {
-                            // IRAP
+                        if (h265 ? (nalUnitType >= 16 && nalUnitType <= 21) : (nalUnitType == 5)) {
+                            // IRAP (BLA or CRA or IDR)
                             isKey = true;
                         }
+                        else if (!h265 && nalUnitType == 1) {
+                            // Non-IDR
+                            // Emulation prevention should not appear unless first_mb_in_slice value is huge
+                            if (len >= 5 && (nal[1] != 0 || nal[2] != 0 || nal[3] != 3)) {
+                                uint8_t sliceIntro[16] = {};
+                                std::copy(nal + 1, nal + 5, sliceIntro);
+                                size_t pos = 0;
+                                // first_mb_in_slice
+                                ReadUegBits(sliceIntro, pos);
+                                int sliceType = ReadUegBits(sliceIntro, pos);
+                                if (sliceType == 2 || sliceType == 4 || sliceType == 7 || sliceType == 9) {
+                                    // I or SI picture
+                                    isKey = true;
+                                }
+                            }
+                        }
                         sampleSize += 4 + len;
                         PushUint(m_videoMdat, static_cast<uint32_t>(len));
                         m_videoMdat.insert(m_videoMdat.end(), nal, nal + len);

diff --git a/tsmemseg.cpp b/tsmemseg.cpp
@@ -387,7 +387,7 @@ void WriteSegmentHeader(std::vector<uint8_t> &buf, uint32_t segCount, bool isMp4
 {
     // NULL TS header
     buf[0] = 0x47;
-    buf[1] = 0x01;
+    buf[1] = 0x1f;
     buf[2] = 0xff;
     buf[3] = 0x10;
     WriteUint32(&buf[4], segCount);
@@ -542,7 +542,7 @@ void ProcessSegmentation(FILE *fp, bool enableFragmentation, uint32_t targetDura
                         if (pid == pat.first_pmt.first_video_pid) {
                             nalState = 0;
                             if (9 + pesHeaderLength < payloadSize) {
-                                if (contains_nal_irap(&nalState, payload + 9 + pesHeaderLength, payloadSize - (9 + pesHeaderLength), h265)) {
+                                if (contains_nal_idr_or_cra(&nalState, payload + 9 + pesHeaderLength, payloadSize - (9 + pesHeaderLength), h265)) {
                                     isKey = !isFirstKey;
                                     isFirstKey = false;
                                 }
@@ -556,7 +556,7 @@ void ProcessSegmentation(FILE *fp, bool enableFragmentation, uint32_t targetDura
                     }
                 }
                 else if (pid == pat.first_pmt.first_video_pid) {
-                    if (contains_nal_irap(&nalState, payload, payloadSize, h265)) {
+                    if (contains_nal_idr_or_cra(&nalState, payload, payloadSize, h265)) {
                         isKey = !isFirstKey;
                         isFirstKey = false;
                     }

diff --git a/util.cpp b/util.cpp
@@ -149,7 +149,7 @@ void extract_pmt(PMT *pmt, const uint8_t *payload, int payload_size, int unit_st
     while (!done);
 }
 
-int contains_nal_irap(int *nal_state, const uint8_t *payload, int payload_size, bool h_265)
+int contains_nal_idr_or_cra(int *nal_state, const uint8_t *payload, int payload_size, bool h_265)
 {
     for (int i = 0; i < payload_size; ++i) {
         // 0,1,2: Searching for NAL start code

diff --git a/util.hpp b/util.hpp
@@ -45,7 +45,7 @@ uint32_t calc_crc32(const uint8_t *data, int data_size, uint32_t crc = 0xfffffff
 int extract_psi(PSI *psi, const uint8_t *payload, int payload_size, int unit_start, int counter);
 void extract_pat(PAT *pat, const uint8_t *payload, int payload_size, int unit_start, int counter);
 void extract_pmt(PMT *pmt, const uint8_t *payload, int payload_size, int unit_start, int counter);
-int contains_nal_irap(int *nal_state, const uint8_t *payload, int payload_size, bool h_265);
+int contains_nal_idr_or_cra(int *nal_state, const uint8_t *payload, int payload_size, bool h_265);
 int get_ts_payload_size(const uint8_t *packet);
 int64_t get_pes_timestamp(const uint8_t *data_5bytes);