From 7aaef29eb7be81d78a2b1a8f22188111dc868abb Mon Sep 17 00:00:00 2001 From: Rajy Rawther Date: Thu, 11 Apr 2024 05:21:50 -0700 Subject: [PATCH] videodemuxer: added seek functionality (#311) * demux add seek functionality * addressed review comments [ROCm/rocdecode commit: 77e4c3150a781d2279147206cca8f8368ae47c9b] --- projects/rocdecode/utils/video_demuxer.h | 233 ++++++++++++++++++++++- 1 file changed, 230 insertions(+), 3 deletions(-) diff --git a/projects/rocdecode/utils/video_demuxer.h b/projects/rocdecode/utils/video_demuxer.h index 58a14a675a..e4a2eebc66 100644 --- a/projects/rocdecode/utils/video_demuxer.h +++ b/projects/rocdecode/utils/video_demuxer.h @@ -40,6 +40,83 @@ extern "C" { * \brief AMD The rocDecode video demuxer API. */ +/** + * @brief Enum for Seek mode + * + */ +typedef enum SeekModeEnum { + SEEK_MODE_EXACT_FRAME = 0, + SEEK_MODE_PREV_KEY_FRAME = 1, + SEEK_MODE_NUM, +} SeekMode; + +/** + * @brief Enum for Seek Criteria + * + */ +typedef enum SeekCriteriaEnum { + SEEK_CRITERIA_FRAME_NUM = 0, + SEEK_CRITERIA_TIME_STAMP = 1, + SEEK_CRITERIA_NUM, +} SeekCriteria; + +struct PacketData { + int32_t key; + int64_t pts; + int64_t dts; + uint64_t pos; + uintptr_t bsl_data; + uint64_t bsl; + uint64_t duration; +}; + +class VideoSeekContext { +public: + VideoSeekContext() + : use_seek_(false), seek_frame_(0), seek_mode_(SEEK_MODE_PREV_KEY_FRAME), seek_crit_(SEEK_CRITERIA_FRAME_NUM), + out_frame_pts_(0), out_frame_duration_(0), num_frames_decoded_(0U) {} + + VideoSeekContext(uint64_t frame_id) + : use_seek_(true), seek_frame_(frame_id), seek_mode_(SEEK_MODE_PREV_KEY_FRAME), + seek_crit_(SEEK_CRITERIA_FRAME_NUM), out_frame_pts_(0), out_frame_duration_(0), num_frames_decoded_(0U) {} + + VideoSeekContext& operator=(const VideoSeekContext& other) { + use_seek_ = other.use_seek_; + seek_frame_ = other.seek_frame_; + seek_mode_ = other.seek_mode_; + seek_crit_ = other.seek_crit_; + out_frame_pts_ = other.out_frame_pts_; + out_frame_duration_ = other.out_frame_duration_; + num_frames_decoded_ = other.num_frames_decoded_; + return *this; + } + + /* Will be set to false when not seeking, true otherwise; + */ + bool use_seek_; + + /* Frame we want to get. Set by user. + * Shall be set to frame timestamp in case seek is done by time. + */ + uint64_t seek_frame_; + + /* Mode in which we seek. */ + SeekMode seek_mode_; + + /* Criteria by which we seek. */ + SeekCriteria seek_crit_; + + /* PTS of frame found after seek. */ + int64_t out_frame_pts_; + + /* Duration of frame found after seek. */ + int64_t out_frame_duration_; + + /* Number of frames that were decoded during seek. */ + uint64_t num_frames_decoded_; + +}; + // Video Demuxer Interface class class VideoDemuxer { @@ -54,13 +131,29 @@ class VideoDemuxer { VideoDemuxer(StreamProvider *stream_provider) : VideoDemuxer(CreateFmtContextUtil(stream_provider)) {av_io_ctx_ = av_fmt_input_ctx_->pb;} ~VideoDemuxer(); bool Demux(uint8_t **video, int *video_size, int64_t *pts = nullptr); + bool Seek(VideoSeekContext& seek_ctx, uint8_t** pp_video, int* video_size); const uint32_t GetWidth() const { return width_;} const uint32_t GetHeight() const { return height_;} const uint32_t GetChromaHeight() const { return chroma_height_;} const uint32_t GetBitDepth() const { return bit_depth_;} const uint32_t GetBytePerPixel() const { return byte_per_pixel_;} const uint32_t GetBitRate() const { return bit_rate_;} - const double GetFrameRate() const {return frame_rate_.den != 0 ? frame_rate_.num / frame_rate_.den : 0;}; + const double GetFrameRate() const {return frame_rate_;}; + bool IsVFR() const { return frame_rate_ != avg_frame_rate_; }; + int64_t TsFromTime(double ts_sec) { + // Convert integer timestamp representation to AV_TIME_BASE and switch to fixed_point + auto const ts_tbu = llround(ts_sec * AV_TIME_BASE); + + // Rescale the timestamp to value represented in stream base units; + AVRational time_factor = {1, AV_TIME_BASE}; + return av_rescale_q(ts_tbu, time_factor, av_fmt_input_ctx_->streams[av_stream_]->time_base); + } + + int64_t TsFromFrameNumber(int64_t frame_num) { + auto const ts_sec = static_cast(frame_num) / frame_rate_; + return TsFromTime(ts_sec); + } + private: VideoDemuxer(AVFormatContext *av_fmt_input_ctx); AVFormatContext *CreateFmtContextUtil(StreamProvider *stream_provider); @@ -73,12 +166,14 @@ class VideoDemuxer { AVBSFContext *av_bsf_ctx_ = nullptr; AVCodecID av_video_codec_id_; AVPixelFormat chroma_format_; - AVRational frame_rate_ = {}; + double frame_rate_ = 0.0; + double avg_frame_rate_ = 0.0; uint8_t *data_with_header_ = nullptr; int av_stream_ = 0; bool is_h264_ = false; bool is_hevc_ = false; bool is_mpeg4_ = false; + bool is_seekable_ = false; int64_t default_time_scale_ = 1000; double time_base_ = 0.0; uint32_t frame_count_ = 0; @@ -197,7 +292,10 @@ VideoDemuxer::VideoDemuxer(AVFormatContext *av_fmt_input_ctx) : av_fmt_input_ctx height_ = av_fmt_input_ctx_->streams[av_stream_]->codecpar->height; chroma_format_ = (AVPixelFormat)av_fmt_input_ctx_->streams[av_stream_]->codecpar->format; bit_rate_ = av_fmt_input_ctx_->streams[av_stream_]->codecpar->bit_rate; - frame_rate_ = av_fmt_input_ctx_->streams[av_stream_]->r_frame_rate; + if (av_fmt_input_ctx_->streams[av_stream_]->r_frame_rate.den != 0) + frame_rate_ = static_cast(av_fmt_input_ctx_->streams[av_stream_]->r_frame_rate.num) / static_cast(av_fmt_input_ctx_->streams[av_stream_]->r_frame_rate.den); + if (av_fmt_input_ctx_->streams[av_stream_]->avg_frame_rate.den != 0) + avg_frame_rate_ = static_cast(av_fmt_input_ctx_->streams[av_stream_]->avg_frame_rate.num) / static_cast(av_fmt_input_ctx_->streams[av_stream_]->avg_frame_rate.den); switch (chroma_format_) { case AV_PIX_FMT_YUV420P10LE: @@ -255,6 +353,9 @@ VideoDemuxer::VideoDemuxer(AVFormatContext *av_fmt_input_ctx) : av_fmt_input_ctx || !strcmp(av_fmt_input_ctx_->iformat->long_name, "FLV (Flash Video)") || !strcmp(av_fmt_input_ctx_->iformat->long_name, "Matroska / WebM")); + // Check if the input file allow seek functionality. + is_seekable_ = av_fmt_input_ctx_->iformat->read_seek || av_fmt_input_ctx_->iformat->read_seek2; + if (is_h264_) { const AVBitStreamFilter *bsf = av_bsf_get_by_name("h264_mp4toannexb"); if (!bsf) { @@ -293,6 +394,132 @@ VideoDemuxer::VideoDemuxer(AVFormatContext *av_fmt_input_ctx) : av_fmt_input_ctx } } +bool VideoDemuxer::Seek(VideoSeekContext& seek_ctx, uint8_t** pp_video, int* video_size) { + /* !!! IMPORTANT !!! + * Across this function, packet decode timestamp (DTS) values are used to + * compare given timestamp against. This is done because DTS values shall + * monotonically increase during the course of decoding unlike PTS values + * which may be affected by frame reordering due to B frames. + */ + + if (!is_seekable_) { + std::cerr << "ERROR: Seek isn't supported for this input." << std::endl; + return false; + } + + if (IsVFR() && (SEEK_CRITERIA_FRAME_NUM == seek_ctx.seek_crit_)) { + std::cerr << "ERROR: Can't seek by frame number in VFR sequences. Seek by timestamp instead." << std::endl; + return false; + } + + // Seek for single frame; + auto seek_frame = [&](VideoSeekContext const& seek_ctx, int flags) { + bool seek_backward = true; + int64_t timestamp = 0; + int ret = 0; + + switch (seek_ctx.seek_crit_) { + case SEEK_CRITERIA_FRAME_NUM: + timestamp = TsFromFrameNumber(seek_ctx.seek_frame_); + ret = av_seek_frame(av_fmt_input_ctx_, av_stream_, timestamp, seek_backward ? AVSEEK_FLAG_BACKWARD | flags : flags); + break; + case SEEK_CRITERIA_TIME_STAMP: + timestamp = TsFromTime(seek_ctx.seek_frame_); + ret = av_seek_frame(av_fmt_input_ctx_, av_stream_, timestamp, seek_backward ? AVSEEK_FLAG_BACKWARD | flags : flags); + break; + default: + std::cerr << "ERROR: Invalid seek mode" << std::endl; + ret = -1; + } + + if (ret < 0) { + throw std::runtime_error("ERROR: seeking for frame: "); + } + }; + + // Check if frame satisfies seek conditions; + auto is_seek_done = [&](PacketData& pkt_data, VideoSeekContext const& seek_ctx) { + int64_t target_ts = 0; + + switch (seek_ctx.seek_crit_) { + case SEEK_CRITERIA_FRAME_NUM: + target_ts = TsFromFrameNumber(seek_ctx.seek_frame_); + break; + case SEEK_CRITERIA_TIME_STAMP: + target_ts = TsFromTime(seek_ctx.seek_frame_); + break; + default: + std::cerr << "ERROR::Invalid seek criteria" << std::endl; + return -1; + } + + if (pkt_data.dts == target_ts) { + return 0; + } + else if (pkt_data.dts > target_ts) { + return 1; + } + else { + return -1; + }; + }; + + /* This will seek for exact frame number; + * Note that decoder may not be able to decode such frame; */ + auto seek_for_exact_frame = [&](PacketData& pkt_data, VideoSeekContext& seek_ctx) { + // Repetititive seek until seek condition is satisfied; + VideoSeekContext tmp_ctx(seek_ctx.seek_frame_); + seek_frame(tmp_ctx, AVSEEK_FLAG_ANY); + + int seek_done = 0; + do { + if (!Demux(pp_video, video_size)) { + break; + } + seek_done = is_seek_done(pkt_data, seek_ctx); + + // We've gone too far and need to seek backwards; + if (seek_done > 0) { + tmp_ctx.seek_frame_--; + seek_frame(tmp_ctx, AVSEEK_FLAG_ANY); + } + // Need to read more frames until we reach requested number; + else if (seek_done < 0) { + continue; + } + } while (seek_done != 0); + + seek_ctx.out_frame_pts_ = pkt_data.pts; + seek_ctx.out_frame_duration_ = pkt_data.duration; + }; + + // Seek for closest key frame in the past; + auto seek_for_prev_key_frame = [&](PacketData& pkt_data, VideoSeekContext& seek_ctx) { + seek_frame(seek_ctx.seek_crit_, AVSEEK_FLAG_BACKWARD); + Demux(pp_video, video_size); + seek_ctx.out_frame_pts_ = pkt_data.pts; + seek_ctx.out_frame_duration_ = pkt_data.duration; + }; + + PacketData pktData; + pktData.bsl_data = size_t(*pp_video); + pktData.bsl = *video_size; + + switch (seek_ctx.seek_mode_) { + case SEEK_MODE_EXACT_FRAME: + seek_for_exact_frame(pktData, seek_ctx); + break; + case SEEK_MODE_PREV_KEY_FRAME: + seek_for_prev_key_frame(pktData, seek_ctx); + break; + default: + throw std::runtime_error("ERROR::Unsupported seek mode"); + break; + } + + return true; +} + AVFormatContext *VideoDemuxer::CreateFmtContextUtil(StreamProvider *stream_provider) { AVFormatContext *ctx = nullptr; if (!(ctx = avformat_alloc_context())) {