Treat non-IDR I-pictures as key frames (issue #1)

Refer to hls.js/avc-video-parser.ts "isKey" variables of mp4fragmenter.cpp and tsmemseg.cpp have different roles.
xtne6f · Mar 13, 2024 · eca62f8 · eca62f8
1 parent dcc47db
commit eca62f8
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 4 deletions.
diff --git a/mp4fragmenter.cpp b/mp4fragmenter.cpp
@@ -391,6 +391,25 @@ void CMp4Fragmenter::AddVideoPes(const std::vector<uint8_t> &pes, bool h265)
                             // 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) &&
+                                (nal[2] != 0 || nal[3] != 0 || nal[4] != 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
@@ -540,7 +540,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;
                                 }
@@ -554,7 +554,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);