support skip hevc b frame

duiniuluantanqin · Feb 14, 2025 · ef578c9 · ef578c9
1 parent 5bf9eca
commit ef578c9
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 9 deletions.
diff --git a/trunk/src/app/srs_app_rtc_source.cpp b/trunk/src/app/srs_app_rtc_source.cpp
@@ -1154,10 +1154,16 @@ srs_error_t SrsRtcRtpBuilder::filter(SrsSharedPtrMessage* msg, SrsFormat* format
 
         // Because RTC does not support B-frame, so we will drop them.
         // TODO: Drop B-frame in better way, which not cause picture corruption.
-        if (!keep_bframe && format->vcodec->id == SrsVideoCodecIdAVC) {
+        if (!keep_bframe) {
             bool is_b_frame;
-            if ((err = SrsVideoFrame::parse_avc_b_frame(sample, is_b_frame)) != srs_success) {
-                return srs_error_wrap(err, "parse bframe");
+            if (format->vcodec->id == SrsVideoCodecIdAVC) {
+                if ((err = SrsVideoFrame::parse_avc_b_frame(sample, is_b_frame)) != srs_success) {
+                    return srs_error_wrap(err, "parse bframe");
+                }
+            } else if (format->vcodec->id == SrsVideoCodecIdHEVC) {
+                if ((err = SrsVideoFrame::parse_hevc_b_frame(sample, format, is_b_frame)) != srs_success) {
+                    return srs_error_wrap(err, "parse bframe");
+                }
             }
             if (is_b_frame) {
                 continue;

diff --git a/trunk/src/kernel/srs_kernel_codec.cpp b/trunk/src/kernel/srs_kernel_codec.cpp
@@ -703,7 +703,7 @@ srs_error_t SrsVideoFrame::parse_avc_nalu_type(const SrsSample* sample, SrsAvcNa
     srs_error_t err = srs_success;
 
     if (sample == NULL || sample->size < 1) {
-        return srs_error_new(ERROR_AVC_NALU_EMPTY, "empty nalu");
+        return srs_error_new(ERROR_NALU_EMPTY, "empty nalu");
     }
 
     uint8_t header = sample->bytes[0];
@@ -716,10 +716,6 @@ srs_error_t SrsVideoFrame::parse_avc_b_frame(const SrsSample* sample, bool& is_b
 {
     srs_error_t err = srs_success;
 
-    if (sample == NULL || sample->size < 1) {
-        return srs_error_new(ERROR_AVC_NALU_EMPTY, "empty nalu");
-    }
-
     SrsAvcNaluType nalu_type;
     if ((err = parse_avc_nalu_type(sample, nalu_type)) != srs_success) {
         return srs_error_wrap(err, "parse avc nalu type error");
@@ -755,6 +751,84 @@ srs_error_t SrsVideoFrame::parse_avc_b_frame(const SrsSample* sample, bool& is_b
     return err;
 }
 
+srs_error_t SrsVideoFrame::parse_hevc_nalu_type(const SrsSample *sample, SrsHevcNaluType &hevc_nalu_type)
+{
+    srs_error_t err = srs_success;
+
+    if (sample == NULL || sample->size < 1) {
+        return srs_error_new(ERROR_NALU_EMPTY, "empty nalu");
+    }
+
+    uint8_t header = sample->bytes[0];
+    hevc_nalu_type = (SrsHevcNaluType)((header >> 1) & 0x3f);
+
+    return err;
+}
+
+srs_error_t SrsVideoFrame::parse_hevc_b_frame(const SrsSample *sample, SrsFormat *format, bool &is_b_frame)
+{
+    srs_error_t err = srs_success;
+
+    SrsHevcNaluType nalu_type;
+    if ((err = parse_hevc_nalu_type(sample, nalu_type)) != srs_success) {
+        return srs_error_wrap(err, "parse hevc nalu type error");
+    }
+
+    SrsBuffer stream(sample->bytes, sample->size);
+    stream.skip(2);
+
+    // @see 7.3.6.1 General slice segment header syntax
+    // @doc ITU-T-H.265-2021.pdf, page 66.
+    SrsBitBuffer bs(&stream);   
+    uint8_t first_slice_segment_in_pic_flag = bs.read_bit();
+    if (nalu_type > SrsHevcNaluType_CODED_SLICE_BLA && nalu_type < SrsHevcNaluType_RESERVED_23) {
+        bs.skip_bits(1);
+        is_b_frame = false;
+        return err;
+    }
+
+    uint32_t slice_pic_parameter_set_id;
+    if ((err = bs.read_bits_ue(slice_pic_parameter_set_id)) != srs_success) {
+        return srs_error_wrap(err, "read slice pic parameter set id");
+    }
+
+    if (slice_pic_parameter_set_id >= SrsHevcMax_PPS_COUNT) {
+        return srs_error_new(ERROR_HEVC_DECODE_ERROR, "slice pic parameter set id out of range: %d", slice_pic_parameter_set_id);
+    }
+
+    SrsHevcRbspPps *pps = &(format->vcodec->hevc_dec_conf_record_.pps_table[slice_pic_parameter_set_id]);
+
+    uint8_t dependent_slice_segment_flag;
+    if (!first_slice_segment_in_pic_flag) {
+        if (pps->dependent_slice_segments_enabled_flag) {
+            dependent_slice_segment_flag = bs.read_bit();
+        } else {
+            dependent_slice_segment_flag = 0;
+        }
+    } else {
+        dependent_slice_segment_flag = 0;
+    }
+
+    if (dependent_slice_segment_flag) {
+        return srs_error_new(ERROR_HEVC_DECODE_ERROR, "dependent slice segment flag is not supported");
+    }
+
+    for (int i = 0; i < pps->num_extra_slice_header_bits; i++) {
+        bs.skip_bits(1);
+    }
+
+    uint32_t slice_type;
+    if ((err = bs.read_bits_ue(slice_type)) != srs_success) {
+        return srs_error_wrap(err, "read slice type");
+    }
+
+    is_b_frame = (slice_type == SrsHevcSliceTypeB) ? true : false;
+
+    // no need to evaluate the rest
+
+    return err;
+}
+
 SrsFormat::SrsFormat()
 {
     acodec = NULL;

diff --git a/trunk/src/kernel/srs_kernel_codec.hpp b/trunk/src/kernel/srs_kernel_codec.hpp
@@ -14,6 +14,7 @@
 
 class SrsBuffer;
 class SrsBitBuffer;
+class SrsFormat;
 
 /**
  * The video codec id.
@@ -498,6 +499,12 @@ enum SrsHevcNaluType {
 };
 #define SrsHevcNaluTypeParse(code) (SrsHevcNaluType)((code & 0x7E) >> 1)
 
+enum SrsHevcSliceType {
+    SrsHevcSliceTypeB = 0,
+    SrsHevcSliceTypeP = 1,
+    SrsHevcSliceTypeI = 2,
+};
+
 struct SrsHevcNalData {
     uint16_t nal_unit_length;
     std::vector<uint8_t> nal_unit_data;
@@ -1321,6 +1328,9 @@ class SrsVideoFrame : public SrsFrame
 public:
     static srs_error_t parse_avc_nalu_type(const SrsSample* sample, SrsAvcNaluType& avc_nalu_type);
     static srs_error_t parse_avc_b_frame(const SrsSample* sample, bool& is_b_frame);
+
+    static srs_error_t parse_hevc_nalu_type(const SrsSample* sample, SrsHevcNaluType& hevc_nalu_type);
+    static srs_error_t parse_hevc_b_frame(const SrsSample* sample, SrsFormat* format, bool& is_b_frame);
 };
 
 /**

diff --git a/trunk/src/kernel/srs_kernel_error.hpp b/trunk/src/kernel/srs_kernel_error.hpp
@@ -279,7 +279,7 @@
     XX(ERROR_HEVC_DECODE_ERROR             , 3099, "HevcDecode", "HEVC decode av stream failed")  \
     XX(ERROR_MP4_HVCC_CHANGE               , 3100, "Mp4HvcCChange", "MP4 does not support video HvcC change") \
     XX(ERROR_HEVC_API_NO_PREFIXED          , 3101, "HevcAnnexbPrefix", "No annexb prefix for HEVC decoder") \
-    XX(ERROR_AVC_NALU_EMPTY                , 3102, "AvcNaluEmpty", "AVC NALU is empty")
+    XX(ERROR_NALU_EMPTY                    , 3102, "NaluEmpty", "NALU is empty")
 
 /**************************************************/
 /* HTTP/StreamConverter protocol error. */