diff --git a/source/codecs/hevc.cpp b/source/codecs/hevc.cpp index d715562..cc51dbd 100644 --- a/source/codecs/hevc.cpp +++ b/source/codecs/hevc.cpp @@ -20,3 +20,217 @@ // SOFTWARE. #include "hevc.hpp" +#include "utility.hpp" + +enum class nal_unit_type : uint8_t { // 6 bits + TRAIL_N = 0, + TRAIL_R = 1, + TSA_N = 2, + TSA_R = 3, + STSA_N = 4, + STSA_R = 5, + RADL_N = 6, + RADL_R = 7, + RASL_N = 8, + RASL_R = 9, + RSV_VCL_N10 = 10, + RSV_VCL_R11 = 11, + RSV_VCL_N12 = 12, + RSV_VCL_R13 = 13, + RSV_VCL_N14 = 14, + RSV_VCL_R15 = 15, + BLA_W_LP = 16, + BLA_W_RADL = 17, + BLA_N_LP = 18, + IDR_W_RADL = 19, + IDR_N_LP = 20, + CRA = 21, + RSV_IRAP_VCL22 = 22, + RSV_IRAP_VCL23 = 23, + RSV_VCL24 = 24, + RSV_VCL25 = 25, + RSV_VCL26 = 26, + RSV_VCL27 = 27, + RSV_VCL28 = 28, + RSV_VCL29 = 29, + RSV_VCL30 = 30, + RSV_VCL31 = 31, + VPS = 32, + SPS = 33, + PPS = 34, + AUD = 35, + EOS = 36, + EOB = 37, + FD = 38, + PREFIX_SEI = 39, + SUFFIX_SEI = 40, + RSV_NVCL41 = 41, + RSV_NVCL42 = 42, + RSV_NVCL43 = 43, + RSV_NVCL44 = 44, + RSV_NVCL45 = 45, + RSV_NVCL46 = 46, + RSV_NVCL47 = 47, + UNSPEC48 = 48, + UNSPEC49 = 49, + UNSPEC50 = 50, + UNSPEC51 = 51, + UNSPEC52 = 52, + UNSPEC53 = 53, + UNSPEC54 = 54, + UNSPEC55 = 55, + UNSPEC56 = 56, + UNSPEC57 = 57, + UNSPEC58 = 58, + UNSPEC59 = 59, + UNSPEC60 = 60, + UNSPEC61 = 61, + UNSPEC62 = 62, + UNSPEC63 = 63, +}; + +struct hevc_nal_unit_header { + bool zero_bit : 1; + nal_unit_type nut : 6; + uint8_t layer_id : 6; + uint8_t temporal_id_plus1 : 3; +}; + +struct hevc_nal { + hevc_nal_unit_header* header; + size_t size = 0; + uint8_t* data = nullptr; +}; + +bool is_nal(uint8_t* data, uint8_t* end) +{ + size_t s = end - data; + if (s < 4) + return false; + + if (*data != 0x0) + return false; + if (*(data + 1) != 0x0) + return false; + if (*(data + 2) != 0x0) + return false; + if (*(data + 3) != 0x1) + return false; + + return true; +} + +bool seek_to_nal(uint8_t*& data, uint8_t* end) +{ + if (data > end) + return false; + + for (; data <= end; data++) { + if (is_nal(data, end)) { + return true; + } + } + + return false; +} + +size_t get_nal_size(uint8_t* data, uint8_t* end) +{ + uint8_t* ptr = data + 4; + if (!seek_to_nal(ptr, end)) { + return end - data; + } + return ptr - data; +} + +bool is_discard_marker(uint8_t* data, uint8_t* end) +{ + size_t s = end - data; + if (s < 4) + return false; + + if (*data != 0x0) + return false; + if (*(data + 1) != 0x0) + return false; + + if (*(data + 2) == 0x3) { + // Discard marker only if the next byte is not 0x0, 0x1, 0x2 or 0x3. + if (*(data + 3) != 0x0) + return false; + if (*(data + 3) != 0x1) + return false; + if (*(data + 3) != 0x2) + return false; + if (*(data + 3) != 0x3) + return false; + + return true; + } else { + if (*(data + 2) == 0x0) + return true; + if (*(data + 2) == 0x1) + return true; + if (*(data + 2) == 0x2) + return true; + + return false; + } +} + +bool should_discard_nal(uint8_t* data, uint8_t* end) +{ + if (data > end) + return true; + + for (; data <= end; data++) { + if (is_discard_marker(data, end)) + return true; + } + + return false; +} + +void progress_parse(uint8_t*& ptr, uint8_t* end, size_t& sz) +{ + ptr += sz; + sz = get_nal_size(ptr, end); +} + +void obsffmpeg::codecs::hevc::extract_header_sei(uint8_t* data, size_t sz_data, + std::vector& header, std::vector& sei) +{ + uint8_t* ptr = data; + uint8_t* end = data + sz_data; + + // Reserve enough memory to store the entire packet data if necessary. + header.reserve(sz_data); + sei.reserve(sz_data); + + if (!seek_to_nal(ptr, end)) { + return; + } + + for (size_t nal_sz = get_nal_size(ptr, end); nal_sz > 0; progress_parse(ptr, end, nal_sz)) { + if (should_discard_nal(ptr + 4, ptr + nal_sz)) { + continue; + } + + hevc_nal nal; + nal.header = reinterpret_cast(ptr + 4); + nal.size = nal_sz - 4 - 2; + nal.data = ptr + 4 + 2; + + switch (nal.header->nut) { + case nal_unit_type::VPS: + case nal_unit_type::SPS: + case nal_unit_type::PPS: + header.insert(header.end(), ptr, ptr + nal_sz); + break; + case nal_unit_type::PREFIX_SEI: + case nal_unit_type::SUFFIX_SEI: + sei.insert(sei.end(), ptr, ptr + nal_sz); + break; + } + } +} diff --git a/source/codecs/hevc.hpp b/source/codecs/hevc.hpp index bdc74cb..036a7e6 100644 --- a/source/codecs/hevc.hpp +++ b/source/codecs/hevc.hpp @@ -20,7 +20,7 @@ // SOFTWARE. #pragma once -#include +#include // Codec: HEVC #define P_HEVC "Codec.HEVC" @@ -60,6 +60,10 @@ namespace obsffmpeg { L6_2 = 186, UNKNOWN = -1, }; + + void extract_header_sei(uint8_t* data, size_t sz_data, + std::vector& header, std::vector& sei); + } // namespace hevc } // namespace codecs } // namespace obsffmpeg diff --git a/source/encoder.cpp b/source/encoder.cpp index 37efdf2..c755bd2 100644 --- a/source/encoder.cpp +++ b/source/encoder.cpp @@ -29,6 +29,7 @@ #include "plugin.hpp" #include "strings.hpp" #include "utility.hpp" +#include "codecs/hevc.hpp" extern "C" { #include @@ -1029,6 +1030,9 @@ int obsffmpeg::encoder::receive_packet(bool* received_packet, struct encoder_pac bfree(tmp_packet); bfree(tmp_header); bfree(tmp_sei); + } else if (_codec->id == AV_CODEC_ID_HEVC) { + obsffmpeg::codecs::hevc::extract_header_sei(_current_packet.data, _current_packet.size, + _extra_data, _sei_data); } _have_first_frame = true; }