diff --git a/applications/ftl2mkv/src/main.cpp b/applications/ftl2mkv/src/main.cpp index 4ee909feefc277a9134e19c432268452d8c67ebc..b555dcbf7f3494dbd66f87415f51cbb0cec3c49e 100644 --- a/applications/ftl2mkv/src/main.cpp +++ b/applications/ftl2mkv/src/main.cpp @@ -3,6 +3,7 @@ #include <ftl/codecs/reader.hpp> #include <ftl/codecs/packet.hpp> #include <ftl/rgbd/camera.hpp> +#include <ftl/codecs/hevc.hpp> #include <fstream> @@ -188,10 +189,7 @@ int main(int argc, char **argv) { bool keyframe = false; if (pkt.codec == codec_t::HEVC) { - // Obtain NAL unit type - int nal_type = (pkt.data[4] >> 1) & 0x3F; - // A type of 32 = VPS unit (so in this case a key frame) - if (nal_type == 32) { + if (ftl::codecs::hevc::isIFrame(pkt.data)) { seen_key[spkt.streamID] = true; keyframe = true; } diff --git a/components/codecs/include/ftl/codecs/hevc.hpp b/components/codecs/include/ftl/codecs/hevc.hpp new file mode 100644 index 0000000000000000000000000000000000000000..f658635d6f239b4aa7a21331f60f6936c517ba93 --- /dev/null +++ b/components/codecs/include/ftl/codecs/hevc.hpp @@ -0,0 +1,112 @@ +#ifndef _FTL_CODECS_HEVC_HPP_ +#define _FTL_CODECS_HEVC_HPP_ + +namespace ftl { +namespace codecs { + +/** + * H.265 / HEVC codec utility functions. + */ +namespace hevc { + +/** + * HEVC Network Abstraction Layer Unit types. + */ +enum class NALType : int { + CODED_SLICE_TRAIL_N = 0, + CODED_SLICE_TRAIL_R = 1, + + CODED_SLICE_TSA_N = 2, + CODED_SLICE_TSA_R = 3, + + CODED_SLICE_STSA_N = 4, + CODED_SLICE_STSA_R = 5, + + CODED_SLICE_RADL_N = 6, + CODED_SLICE_RADL_R = 7, + + CODED_SLICE_RASL_N = 8, + CODED_SLICE_RASL_R = 9, + + RESERVED_VCL_N10 = 10, + RESERVED_VCL_R11 = 11, + RESERVED_VCL_N12 = 12, + RESERVED_VCL_R13 = 13, + RESERVED_VCL_N14 = 14, + RESERVED_VCL_R15 = 15, + + CODED_SLICE_BLA_W_LP = 16, + CODED_SLICE_BLA_W_RADL = 17, + CODED_SLICE_BLA_N_LP = 18, + CODED_SLICE_IDR_W_RADL = 19, + CODED_SLICE_IDR_N_LP = 20, + CODED_SLICE_CRA = 21, + RESERVED_IRAP_VCL22 = 22, + RESERVED_IRAP_VCL23 = 23, + + RESERVED_VCL24 = 24, + RESERVED_VCL25 = 25, + RESERVED_VCL26 = 26, + RESERVED_VCL27 = 27, + RESERVED_VCL28 = 28, + RESERVED_VCL29 = 29, + RESERVED_VCL30 = 30, + RESERVED_VCL31 = 31, + + VPS = 32, + SPS = 33, + PPS = 34, + ACCESS_UNIT_DELIMITER = 35, + EOS = 36, + EOB = 37, + FILLER_DATA = 38, + PREFIX_SEI = 39, + SUFFIX_SEI = 40, + + RESERVED_NVCL41 = 41, + RESERVED_NVCL42 = 42, + RESERVED_NVCL43 = 43, + RESERVED_NVCL44 = 44, + RESERVED_NVCL45 = 45, + RESERVED_NVCL46 = 46, + RESERVED_NVCL47 = 47, + UNSPECIFIED_48 = 48, + UNSPECIFIED_49 = 49, + UNSPECIFIED_50 = 50, + UNSPECIFIED_51 = 51, + UNSPECIFIED_52 = 52, + UNSPECIFIED_53 = 53, + UNSPECIFIED_54 = 54, + UNSPECIFIED_55 = 55, + UNSPECIFIED_56 = 56, + UNSPECIFIED_57 = 57, + UNSPECIFIED_58 = 58, + UNSPECIFIED_59 = 59, + UNSPECIFIED_60 = 60, + UNSPECIFIED_61 = 61, + UNSPECIFIED_62 = 62, + UNSPECIFIED_63 = 63, + INVALID = 64 +}; + +/** + * Extract the NAL unit type from the first NAL header. + * With NvPipe, the 5th byte contains the NAL Unit header. + */ +inline NALType getNALType(const std::vector<uint8_t> &data) { + return static_cast<NALType>((data[4] >> 1) & 0x3F); +} + +/** + * Check the HEVC bitstream for an I-Frame. With NvPipe, all I-Frames start + * with a VPS NAL unit so just check for this. + */ +inline bool isIFrame(const std::vector<uint8_t> &data) { + return getNALType(data) == NALType::VPS; +} + +} +} +} + +#endif // _FTL_CODECS_HEVC_HPP_ diff --git a/components/codecs/src/nvpipe_decoder.cpp b/components/codecs/src/nvpipe_decoder.cpp index 39500ee2b248d5e887679633a49d865b722f13ab..b5e358388f74ac9dfb7df4082a56ab8426cf9f4f 100644 --- a/components/codecs/src/nvpipe_decoder.cpp +++ b/components/codecs/src/nvpipe_decoder.cpp @@ -3,6 +3,7 @@ #include <loguru.hpp> #include <ftl/cuda_util.hpp> +#include <ftl/codecs/hevc.hpp> //#include <cuda_runtime.h> #include <opencv2/core/cuda/common.hpp> @@ -58,8 +59,7 @@ bool NvPipeDecoder::decode(const ftl::codecs::Packet &pkt, cv::Mat &out) { if (pkt.codec == ftl::codecs::codec_t::HEVC) { // Obtain NAL unit type - int nal_type = (pkt.data[4] >> 1) & 0x3F; - if (nal_type == 32) seen_iframe_ = true; + if (ftl::codecs::hevc::isIFrame(pkt.data)) seen_iframe_ = true; } if (!seen_iframe_) return false; diff --git a/components/rgbd-sources/src/file_source.cpp b/components/rgbd-sources/src/file_source.cpp index cab5e084844921772d9e00d5a13d6cdaec8b0ac9..38acf1a7c47c35c3a2c88f505a801b3756168ffd 100644 --- a/components/rgbd-sources/src/file_source.cpp +++ b/components/rgbd-sources/src/file_source.cpp @@ -1,5 +1,6 @@ #include "file_source.hpp" +#include <ftl/codecs/hevc.hpp> #include <ftl/timer.hpp> using ftl::rgbd::detail::FileSource; @@ -42,10 +43,7 @@ FileSource::FileSource(ftl::rgbd::Source *s, ftl::codecs::Reader *r, int sid) : has_calibration_ = true; } else { if (pkt.codec == codec_t::HEVC) { - // Obtain NAL unit type - int nal_type = (pkt.data[4] >> 1) & 0x3F; - // A type of 32 = VPS unit, hence I-Frame in this case so skip past packets - if (nal_type == 32) _removeChannel(spkt.channel); + if (ftl::codecs::hevc::isIFrame(pkt.data)) _removeChannel(spkt.channel); } cache_[cache_write_].emplace_back(); auto &c = cache_[cache_write_].back();