From 67c9ea1983bd5f405b8fff3ff3581f486b3335af Mon Sep 17 00:00:00 2001 From: jeffqjiangNew <142832361+jeffqjiangNew@users.noreply.github.com> Date: Tue, 21 Nov 2023 14:33:15 -0500 Subject: [PATCH] Re-enabled SEI message processing. (#79) * * rocDecode/HEVC: Re-enabled SEI message processing. - Added "-sei" command option to extract SEI messages from stream. - Added support for SEI info parsing for multiple SEI NAL units with multiple payloads. - Fixed a SEI syntax parsing issue and several procssing issues. - Fixed a double free memory issue in SEI data buffer handling in HandlePictureDisplay(). * * rocDecode/HEVC: Coding format changes. * * rocDecode/HEVC: Minor format changes. [ROCm/rocdecode commit: ad186a90dad85f7a7f8c5ade8e839d552ad60b76] --- .../samples/videoDecode/videodecode.cpp | 11 +- projects/rocdecode/src/parser/hevc_parser.cpp | 128 +++++++++++------- projects/rocdecode/src/parser/hevc_parser.h | 22 +-- .../utils/rocvideodecode/roc_video_dec.cpp | 2 + 4 files changed, 104 insertions(+), 59 deletions(-) diff --git a/projects/rocdecode/samples/videoDecode/videodecode.cpp b/projects/rocdecode/samples/videoDecode/videodecode.cpp index e2fad9169c..4f4b71463d 100644 --- a/projects/rocdecode/samples/videoDecode/videodecode.cpp +++ b/projects/rocdecode/samples/videoDecode/videodecode.cpp @@ -42,6 +42,7 @@ void ShowHelpAndExit(const char *option = NULL) { << "-o Output File Path - dumps output if requested; optional" << std::endl << "-d GPU device ID (0 for the first device, 1 for the second, etc.); optional; default: 0" << std::endl << "-z force_zero_latency (force_zero_latency, Decoded frames will be flushed out for display immediately); optional;" << std::endl + << "-sei extract SEI messages; optional;" << std::endl << "-crop crop rectangle for output (not used when using interopped decoded frame); optional; default: 0" << std::endl; exit(0); } @@ -52,6 +53,7 @@ int main(int argc, char **argv) { int dump_output_frames = 0; int device_id = 0; bool b_force_zero_latency = false; // false by default: enabling this option might affect decoding performance + bool b_extract_sei_messages = false; Rect crop_rect = {}; Rect *p_crop_rect = nullptr; OutputSurfaceMemoryType mem_type = OUT_SURFACE_MEM_DEV_INTERNAL; // set to internal @@ -92,6 +94,13 @@ int main(int argc, char **argv) { b_force_zero_latency = true; continue; } + if (!strcmp(argv[i], "-sei")) { + if (i == argc) { + ShowHelpAndExit("-sei"); + } + b_extract_sei_messages = true; + continue; + } if (!strcmp(argv[i], "-crop")) { if (++i == argc || 4 != sscanf(argv[i], "%d,%d,%d,%d", &crop_rect.l, &crop_rect.t, &crop_rect.r, &crop_rect.b)) { ShowHelpAndExit("-crop"); @@ -108,7 +117,7 @@ int main(int argc, char **argv) { try { VideoDemuxer demuxer(input_file_path.c_str()); rocDecVideoCodec rocdec_codec_id = AVCodec2RocDecVideoCodec(demuxer.GetCodecID()); - RocVideoDecoder viddec(device_id, mem_type, rocdec_codec_id, false, b_force_zero_latency, p_crop_rect); + RocVideoDecoder viddec(device_id, mem_type, rocdec_codec_id, false, b_force_zero_latency, p_crop_rect, b_extract_sei_messages); std::string device_name, gcn_arch_name; int pci_bus_id, pci_domain_id, pci_device_id; diff --git a/projects/rocdecode/src/parser/hevc_parser.cpp b/projects/rocdecode/src/parser/hevc_parser.cpp index 9d9f0ba7e6..972d554d59 100644 --- a/projects/rocdecode/src/parser/hevc_parser.cpp +++ b/projects/rocdecode/src/parser/hevc_parser.cpp @@ -48,7 +48,12 @@ HEVCVideoParser::HEVCVideoParser() { m_pps_ = AllocStruct(MAX_PPS_COUNT); m_sh_ = AllocStruct(1); m_sh_copy_ = AllocStruct(1); - m_sei_message_ = AllocStruct(1); + + sei_rbsp_buf_ = nullptr; + sei_rbsp_buf_size_ = 0; + sei_payload_buf_ = nullptr; + sei_payload_buf_size_ = 0; + sei_message_list_.assign(INIT_SEI_MESSAGE_COUNT, {0}); memset(&curr_pic_info_, 0, sizeof(HevcPicInfo)); memset(&dpb_buffer_, 0, sizeof(DecodedPictureBuffer)); @@ -91,7 +96,7 @@ rocDecStatus HEVCVideoParser::ParseVideoData(RocdecSourceDataPacket *p_data) { // Whenever new sei message found if (pfn_get_sei_message_cb_ && sei_message_count_ > 0) { - FillSeiMessageCallbackFn(m_sei_message_); + FillSeiMessageCallbackFn(); } // Decode the picture @@ -121,8 +126,11 @@ HEVCVideoParser::~HEVCVideoParser() { if (m_sh_copy_) { delete m_sh_copy_; } - if (m_sei_message_) { - delete m_sei_message_; + if (sei_rbsp_buf_) { + delete [] sei_rbsp_buf_; + } + if (sei_payload_buf_) { + delete [] sei_payload_buf_; } } @@ -211,19 +219,11 @@ void HEVCVideoParser::FillSeqCallbackFn(SpsData* sps_data) { pfn_sequece_cb_(parser_params_.pUserData, &video_format_params_); } -void HEVCVideoParser::FillSeiMessageCallbackFn(SeiMessageData* sei_message_data) { +void HEVCVideoParser::FillSeiMessageCallbackFn() { sei_message_info_params_.sei_message_count = sei_message_count_; - sei_message_info_params_.pSEIMessage = nullptr; - sei_message_info_params_.pSEIMessage->sei_message_type = sei_message_data->payload_type; - sei_message_info_params_.pSEIMessage->sei_message_size = sei_message_data->payload_size; - // TODO: check reserve[3] values - for (int i = 0; i < 3; i++) { - sei_message_info_params_.pSEIMessage->reserved[i] = 0; - } - sei_message_info_params_.pSEIData = &sei_message_data; - - // TODO: check picIdx value - sei_message_info_params_.picIdx = 0; + sei_message_info_params_.pSEIMessage = sei_message_list_.data(); + sei_message_info_params_.pSEIData = (void*)sei_payload_buf_; + sei_message_info_params_.picIdx = curr_pic_info_.pic_idx; // callback function with RocdecSeiMessageInfo params filled out if (pfn_get_sei_message_cb_) pfn_get_sei_message_cb_(parser_params_.pUserData, &sei_message_info_params_); @@ -517,6 +517,7 @@ bool HEVCVideoParser::ParseFrameData(const uint8_t* p_stream, uint32_t frame_dat slice_num_ = 0; sei_message_count_ = 0; + sei_payload_size_ = 0; do { ret = GetNalUnit(); @@ -617,13 +618,24 @@ bool HEVCVideoParser::ParseFrameData(const uint8_t* p_stream, uint32_t frame_dat break; } - case NAL_UNIT_PREFIX_SEI: { - memcpy(m_rbsp_buf_, (frame_data_buffer_ptr_ + curr_start_code_offset_ + 5), ebsp_size); - m_rbsp_size_ = EBSPtoRBSP(m_rbsp_buf_, 0, ebsp_size); - // todo:: ParseSeiMessage is causing Segfault: disabling until we fix it + case NAL_UNIT_PREFIX_SEI: + case NAL_UNIT_SUFFIX_SEI: { if (pfn_get_sei_message_cb_) { - //ParseSeiMessage(m_rbsp_buf_, m_rbsp_size_); - //sei_message_count_++; + int sei_ebsp_size = nal_unit_size_ - 5; // copy the entire NAL unit + if (sei_rbsp_buf_) { + if (sei_ebsp_size > sei_rbsp_buf_size_) { + delete [] sei_rbsp_buf_; + sei_rbsp_buf_ = new uint8_t [sei_ebsp_size]; + sei_rbsp_buf_size_ = sei_ebsp_size; + } + } + else { + sei_rbsp_buf_size_ = sei_ebsp_size > INIT_SEI_PAYLOAD_BUF_SIZE ? sei_ebsp_size : INIT_SEI_PAYLOAD_BUF_SIZE; + sei_rbsp_buf_ = new uint8_t [sei_rbsp_buf_size_]; + } + memcpy(sei_rbsp_buf_, (frame_data_buffer_ptr_ + curr_start_code_offset_ + 5), sei_ebsp_size); + m_rbsp_size_ = EBSPtoRBSP(sei_rbsp_buf_, 0, sei_ebsp_size); + ParseSeiMessage(sei_rbsp_buf_, m_rbsp_size_); } break; } @@ -1788,33 +1800,55 @@ bool HEVCVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size) { } void HEVCVideoParser::ParseSeiMessage(uint8_t *nalu, size_t size) { - size_t offset = 0; - uint32_t sei_message_id = Parser::ExpGolomb::ReadUe(nalu, offset); - SeiMessageData *sei_message_ptr = &m_sei_message_[sei_message_id]; - memset(sei_message_ptr, 0, sizeof(SeiMessageData)); + int offset = 0; // byte offset + int payload_type; + int payload_size; - sei_message_ptr->payload_type = 0; - sei_message_ptr->payload_size = 0; - uint8_t temp_byte; - temp_byte = Parser::GetBit(nalu, offset); - while (temp_byte == 0XFF) { - sei_message_ptr->payload_type += 255; - temp_byte = Parser::GetBit(nalu, offset); - } - sei_message_ptr->payload_type += temp_byte; - sei_message_ptr->payload_size = 0; - temp_byte = Parser::GetBit(nalu, offset); - while (temp_byte == 0XFF) { - sei_message_ptr->payload_size += 255; - temp_byte = Parser::GetBit(nalu, offset); - } - sei_message_ptr->payload_size += temp_byte; + do { + payload_type = 0; + while (nalu[offset] == 0xFF) { + payload_type += 255; // ff_byte + offset++; + } + payload_type += nalu[offset]; // last_payload_type_byte + offset++; - // copy the payload to buffer - if (sei_message_ptr->payload_size > sizeof(m_sei_data_)) { - THROW("sei payload size is too big for sei buffer!"); - } - memcpy(m_sei_data_, sei_message_ptr, sei_message_ptr->payload_size); + payload_size = 0; + while (nalu[offset] == 0xFF) { + payload_size += 255; // ff_byte + offset++; + } + payload_size += nalu[offset]; // last_payload_size_byte + offset++; + + // We start with INIT_SEI_MESSAGE_COUNT. Should be enough for normal use cases. If not, resize. + if((sei_message_count_ + 1) > sei_message_list_.size()) { + sei_message_list_.resize((sei_message_count_ + 1)); + } + sei_message_list_[sei_message_count_].sei_message_type = payload_type; + sei_message_list_[sei_message_count_].sei_message_size = payload_size; + + if (sei_payload_buf_) { + if ((payload_size + sei_payload_size_) > sei_payload_buf_size_) { + uint8_t *tmp_ptr = new uint8_t [payload_size + sei_payload_size_]; + memcpy(tmp_ptr, sei_payload_buf_, sei_payload_size_); // save the existing payload + delete [] sei_payload_buf_; + sei_payload_buf_ = tmp_ptr; + } + } + else { + // First payload, sei_payload_size_ is 0. + sei_payload_buf_size_ = payload_size > INIT_SEI_PAYLOAD_BUF_SIZE ? payload_size : INIT_SEI_PAYLOAD_BUF_SIZE; + sei_payload_buf_ = new uint8_t [sei_payload_buf_size_]; + } + // Append the current payload to sei_payload_buf_ + memcpy(sei_payload_buf_ + sei_payload_size_, nalu + offset, payload_size); + + sei_payload_size_ += payload_size; + sei_message_count_++; + + offset += payload_size; + } while (offset < size && nalu[offset] != 0x80); } bool HEVCVideoParser::IsIdrPic(NalUnitHeader *nal_header_ptr) { diff --git a/projects/rocdecode/src/parser/hevc_parser.h b/projects/rocdecode/src/parser/hevc_parser.h index 6e39c70c85..a5ac3e30b9 100644 --- a/projects/rocdecode/src/parser/hevc_parser.h +++ b/projects/rocdecode/src/parser/hevc_parser.h @@ -41,11 +41,12 @@ extern int scaling_list_default_3[1][2][64]; #define MAX_SPS_COUNT 16 // 7.3.2.2.1 #define MAX_PPS_COUNT 64 // 7.4.3.3.1 #define RBSP_BUF_SIZE 1024 // enough to parse any parameter sets or slice headers -#define SEI_BUF_SIZE 1024 #define HVC_MAX_DPB_FRAMES 16 // (A-2) #define HEVC_MAX_NUM_REF_PICS 16 // 7.4.7.1. (num_tile_columns_minus1 + 1) * PicHeightInCtbsY − 1. Max tile columns = 20 (A.4.2). Pic height in 16x16 CTB of 8K = 270. #define MAX_ENTRY_POINT_OFFSETS 20 * 270 +#define INIT_SEI_MESSAGE_COUNT 16 // initial SEI message count +#define INIT_SEI_PAYLOAD_BUF_SIZE 1024 * 1024 // initial SEI payload buffer size, 1 MB class HEVCVideoParser : public RocVideoParser { @@ -603,12 +604,6 @@ protected: uint8_t slice_segment_header_extension_data_byte[256]; //u(8) } SliceHeaderData; - typedef struct { - uint8_t payload_type; - uint8_t reserved[3]; - uint32_t payload_size; - } SeiMessageData; - /*! \brief Inline function to Parse the NAL Unit Header * * \param [in] nal_unit A pointer of uint8_t containing the Demuxed output stream @@ -677,17 +672,22 @@ protected: PpsData* m_pps_ = nullptr; SliceHeaderData* m_sh_ = nullptr; SliceHeaderData* m_sh_copy_ = nullptr; - SeiMessageData* m_sei_message_ = nullptr; NalUnitHeader slice_nal_unit_header_; - uint8_t m_sei_data_[SEI_BUF_SIZE]; // to store SEI payload HevcPicInfo curr_pic_info_; bool b_new_picture_; int m_packet_count_; - int sei_message_count_; int m_rbsp_size_; uint8_t m_rbsp_buf_[RBSP_BUF_SIZE]; // to store parameter set or slice header RBSP + uint8_t *sei_rbsp_buf_; // buffer to store SEI RBSP. Allocated at run time. + uint32_t sei_rbsp_buf_size_; + std::vector sei_message_list_; + int sei_message_count_; // total SEI playload message count of the current frame. + uint8_t *sei_payload_buf_; // buffer to store SEI playload. Allocated at run time. + uint32_t sei_payload_buf_size_; + uint32_t sei_payload_size_; // total SEI payload size of the current frame + int slice_num_; uint8_t* pic_stream_data_ptr_; int pic_stream_data_size_; @@ -897,7 +897,7 @@ private: // functions to fill structures for callback functions void FillSeqCallbackFn(SpsData* sps_data); - void FillSeiMessageCallbackFn(SeiMessageData* sei_message_data); + void FillSeiMessageCallbackFn(); /*! \brief Function to fill the decode parameters and call back decoder to decode a picture * \return Return code in ParserResult form diff --git a/projects/rocdecode/utils/rocvideodecode/roc_video_dec.cpp b/projects/rocdecode/utils/rocvideodecode/roc_video_dec.cpp index 3dc58d3be2..b675cd6479 100644 --- a/projects/rocdecode/utils/rocvideodecode/roc_video_dec.cpp +++ b/projects/rocdecode/utils/rocvideodecode/roc_video_dec.cpp @@ -465,7 +465,9 @@ int RocVideoDecoder::HandlePictureDisplay(RocdecParserDispInfo *pDispInfo) { } } free(sei_message_display_q_[pDispInfo->picture_index].pSEIData); + sei_message_display_q_[pDispInfo->picture_index].pSEIData = NULL; // to avoid double free free(sei_message_display_q_[pDispInfo->picture_index].pSEIMessage); + sei_message_display_q_[pDispInfo->picture_index].pSEIMessage = NULL; // to avoid double free } }