docs-6.1.0/utils/video_demuxer.h Source File

docs-6.1.0/utils/video_demuxer.h Source File#

rocDecode: docs-6.1.0/utils/video_demuxer.h Source File
video_demuxer.h
Go to the documentation of this file.
1 /*
2 Copyright (c) 2023 - 2024 Advanced Micro Devices, Inc. All rights reserved.
3 
4 Permission is hereby granted, free of charge, to any person obtaining a copy
5 of this software and associated documentation files (the "Software"), to deal
6 in the Software without restriction, including without limitation the rights
7 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8 copies of the Software, and to permit persons to whom the Software is
9 furnished to do so, subject to the following conditions:
10 
11 The above copyright notice and this permission notice shall be included in
12 all copies or substantial portions of the Software.
13 
14 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
19 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
20 THE SOFTWARE.
21 */
22 
23 #pragma once
24 
25 #include <iostream>
26 extern "C" {
27  #include <libavcodec/avcodec.h>
28  #include <libavformat/avformat.h>
29 }
30 
31 #include "rocdecode.h"
41 // Video Demuxer Interface class
42 class VideoDemuxer {
43  public:
45  public:
46  virtual ~StreamProvider() {}
47  virtual int GetData(uint8_t *buf, int buf_size) = 0;
48  };
49  AVCodecID GetCodecID() { return av_video_codec_id_; };
50  VideoDemuxer(const char *input_file_path) : VideoDemuxer(CreateFmtContextUtil(input_file_path)) {}
51  VideoDemuxer(StreamProvider *stream_provider) : VideoDemuxer(CreateFmtContextUtil(stream_provider)) {av_io_ctx_ = av_fmt_input_ctx_->pb;}
52  ~VideoDemuxer();
53  bool Demux(uint8_t **video, int *video_size, int64_t *pts = nullptr);
54  private:
55  VideoDemuxer(AVFormatContext *av_fmt_input_ctx);
56  AVFormatContext *CreateFmtContextUtil(StreamProvider *stream_provider);
57  AVFormatContext *CreateFmtContextUtil(const char *input_file_path);
58  static int ReadPacket(void *data, uint8_t *buf, int buf_size);
59  AVFormatContext *av_fmt_input_ctx_ = nullptr;
60  AVIOContext *av_io_ctx_ = nullptr;
61  AVPacket* packet_ = nullptr;
62  AVPacket* packet_filtered_ = nullptr;
63  AVBSFContext *av_bsf_ctx_ = nullptr;
64  AVCodecID av_video_codec_id_;
65  uint8_t *data_with_header_ = nullptr;
66  int av_stream_ = 0;
67  bool is_h264_ = false;
68  bool is_hevc_ = false;
69  bool is_mpeg4_ = false;
70  int64_t default_time_scale_ = 1000;
71  double time_base_ = 0.0;
72  unsigned int frame_count_ = 0;
73 };
74 
75 VideoDemuxer::~VideoDemuxer() {
76  if (!av_fmt_input_ctx_) {
77  return;
78  }
79  if (packet_) {
80  av_packet_free(&packet_);
81  }
82  if (packet_filtered_) {
83  av_packet_free(&packet_filtered_);
84  }
85  if (av_bsf_ctx_) {
86  av_bsf_free(&av_bsf_ctx_);
87  }
88  avformat_close_input(&av_fmt_input_ctx_);
89  if (av_io_ctx_) {
90  av_freep(&av_io_ctx_->buffer);
91  av_freep(&av_io_ctx_);
92  }
93  if (data_with_header_) {
94  av_free(data_with_header_);
95  }
96 }
97 
98 bool VideoDemuxer::Demux(uint8_t **video, int *video_size, int64_t *pts) {
99  if (!av_fmt_input_ctx_) {
100  return false;
101  }
102  *video_size = 0;
103  if (packet_->data) {
104  av_packet_unref(packet_);
105  }
106  int ret = 0;
107  while ((ret = av_read_frame(av_fmt_input_ctx_, packet_)) >= 0 && packet_->stream_index != av_stream_) {
108  av_packet_unref(packet_);
109  }
110  if (ret < 0) {
111  return false;
112  }
113  if (is_h264_ || is_hevc_) {
114  if (packet_filtered_->data) {
115  av_packet_unref(packet_filtered_);
116  }
117  if (av_bsf_send_packet(av_bsf_ctx_, packet_) != 0) {
118  std::cerr << "ERROR: av_bsf_send_packet failed!" << std::endl;
119  return false;
120  }
121  if (av_bsf_receive_packet(av_bsf_ctx_, packet_filtered_) != 0) {
122  std::cerr << "ERROR: av_bsf_receive_packet failed!" << std::endl;
123  return false;
124  }
125  *video = packet_filtered_->data;
126  *video_size = packet_filtered_->size;
127  if (pts)
128  *pts = (int64_t) (packet_filtered_->pts * default_time_scale_ * time_base_);
129  } else {
130  if (is_mpeg4_ && (frame_count_ == 0)) {
131  int ext_data_size = av_fmt_input_ctx_->streams[av_stream_]->codecpar->extradata_size;
132  if (ext_data_size > 0) {
133  data_with_header_ = (uint8_t *)av_malloc(ext_data_size + packet_->size - 3 * sizeof(uint8_t));
134  if (!data_with_header_) {
135  std::cerr << "ERROR: av_malloc failed!" << std::endl;
136  return false;
137  }
138  memcpy(data_with_header_, av_fmt_input_ctx_->streams[av_stream_]->codecpar->extradata, ext_data_size);
139  memcpy(data_with_header_ + ext_data_size, packet_->data + 3, packet_->size - 3 * sizeof(uint8_t));
140  *video = data_with_header_;
141  *video_size = ext_data_size + packet_->size - 3 * sizeof(uint8_t);
142  }
143  } else {
144  *video = packet_->data;
145  *video_size = packet_->size;
146  }
147  if (pts)
148  *pts = (int64_t)(packet_->pts * default_time_scale_ * time_base_);
149  }
150  frame_count_++;
151  return true;
152 }
153 
154 VideoDemuxer::VideoDemuxer(AVFormatContext *av_fmt_input_ctx) : av_fmt_input_ctx_(av_fmt_input_ctx) {
155  av_log_set_level(AV_LOG_QUIET);
156  if (!av_fmt_input_ctx_) {
157  std::cerr << "ERROR: av_fmt_input_ctx_ is not vaild!" << std::endl;
158  return;
159  }
160  packet_ = av_packet_alloc();
161  packet_filtered_ = av_packet_alloc();
162  if (!packet_ || !packet_filtered_) {
163  std::cerr << "ERROR: av_packet_alloc failed!" << std::endl;
164  return;
165  }
166  if (avformat_find_stream_info(av_fmt_input_ctx_, nullptr) < 0) {
167  std::cerr << "ERROR: avformat_find_stream_info failed!" << std::endl;
168  return;
169  }
170  av_stream_ = av_find_best_stream(av_fmt_input_ctx_, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
171  if (av_stream_ < 0) {
172  std::cerr << "ERROR: av_find_best_stream failed!" << std::endl;
173  av_packet_free(&packet_);
174  av_packet_free(&packet_filtered_);
175  return;
176  }
177  av_video_codec_id_ = av_fmt_input_ctx_->streams[av_stream_]->codecpar->codec_id;
178  AVRational time_base = av_fmt_input_ctx_->streams[av_stream_]->time_base;
179  time_base_ = av_q2d(time_base);
180 
181  is_h264_ = av_video_codec_id_ == AV_CODEC_ID_H264 && (!strcmp(av_fmt_input_ctx_->iformat->long_name, "QuickTime / MOV")
182  || !strcmp(av_fmt_input_ctx_->iformat->long_name, "FLV (Flash Video)")
183  || !strcmp(av_fmt_input_ctx_->iformat->long_name, "Matroska / WebM"));
184  is_hevc_ = av_video_codec_id_ == AV_CODEC_ID_HEVC && (!strcmp(av_fmt_input_ctx_->iformat->long_name, "QuickTime / MOV")
185  || !strcmp(av_fmt_input_ctx_->iformat->long_name, "FLV (Flash Video)")
186  || !strcmp(av_fmt_input_ctx_->iformat->long_name, "Matroska / WebM"));
187  is_mpeg4_ = av_video_codec_id_ == AV_CODEC_ID_MPEG4 && (!strcmp(av_fmt_input_ctx_->iformat->long_name, "QuickTime / MOV")
188  || !strcmp(av_fmt_input_ctx_->iformat->long_name, "FLV (Flash Video)")
189  || !strcmp(av_fmt_input_ctx_->iformat->long_name, "Matroska / WebM"));
190 
191  if (is_h264_) {
192  const AVBitStreamFilter *bsf = av_bsf_get_by_name("h264_mp4toannexb");
193  if (!bsf) {
194  std::cerr << "ERROR: av_bsf_get_by_name() failed" << std::endl;
195  av_packet_free(&packet_);
196  av_packet_free(&packet_filtered_);
197  return;
198  }
199  if (av_bsf_alloc(bsf, &av_bsf_ctx_) != 0) {
200  std::cerr << "ERROR: av_bsf_alloc failed!" << std::endl;
201  return;
202  }
203  avcodec_parameters_copy(av_bsf_ctx_->par_in, av_fmt_input_ctx_->streams[av_stream_]->codecpar);
204  if (av_bsf_init(av_bsf_ctx_) < 0) {
205  std::cerr << "ERROR: av_bsf_init failed!" << std::endl;
206  return;
207  }
208  }
209  if (is_hevc_) {
210  const AVBitStreamFilter *bsf = av_bsf_get_by_name("hevc_mp4toannexb");
211  if (!bsf) {
212  std::cerr << "ERROR: av_bsf_get_by_name() failed" << std::endl;
213  av_packet_free(&packet_);
214  av_packet_free(&packet_filtered_);
215  return;
216  }
217  if (av_bsf_alloc(bsf, &av_bsf_ctx_) != 0 ) {
218  std::cerr << "ERROR: av_bsf_alloc failed!" << std::endl;
219  return;
220  }
221  avcodec_parameters_copy(av_bsf_ctx_->par_in, av_fmt_input_ctx_->streams[av_stream_]->codecpar);
222  if (av_bsf_init(av_bsf_ctx_) < 0) {
223  std::cerr << "ERROR: av_bsf_init failed!" << std::endl;
224  return;
225  }
226  }
227 }
228 
229 AVFormatContext *VideoDemuxer::CreateFmtContextUtil(StreamProvider *stream_provider) {
230  AVFormatContext *ctx = nullptr;
231  if (!(ctx = avformat_alloc_context())) {
232  std::cerr << "ERROR: avformat_alloc_context failed" << std::endl;
233  return nullptr;
234  }
235  uint8_t *avioc_buffer = nullptr;
236  int avioc_buffer_size = 100 * 1024 * 1024;
237  avioc_buffer = (uint8_t *)av_malloc(avioc_buffer_size);
238  if (!avioc_buffer) {
239  std::cerr << "ERROR: av_malloc failed!" << std::endl;
240  return nullptr;
241  }
242  av_io_ctx_ = avio_alloc_context(avioc_buffer, avioc_buffer_size,
243  0, stream_provider, &ReadPacket, nullptr, nullptr);
244  if (!av_io_ctx_) {
245  std::cerr << "ERROR: avio_alloc_context failed!" << std::endl;
246  return nullptr;
247  }
248  ctx->pb = av_io_ctx_;
249 
250  if (avformat_open_input(&ctx, nullptr, nullptr, nullptr) != 0) {
251  std::cerr << "ERROR: avformat_open_input failed!" << std::endl;
252  return nullptr;
253  }
254  return ctx;
255 }
256 
257 AVFormatContext *VideoDemuxer::CreateFmtContextUtil(const char *input_file_path) {
258  avformat_network_init();
259  AVFormatContext *ctx = nullptr;
260  if (avformat_open_input(&ctx, input_file_path, nullptr, nullptr) != 0 ) {
261  std::cerr << "ERROR: avformat_open_input failed!" << std::endl;
262  return nullptr;
263  }
264  return ctx;
265 }
266 
267 int VideoDemuxer::ReadPacket(void *data, uint8_t *buf, int buf_size) {
268  return ((StreamProvider *)data)->GetData(buf, buf_size);
269 }
270 
271 static inline rocDecVideoCodec AVCodec2RocDecVideoCodec(AVCodecID av_codec) {
272  switch (av_codec) {
273  case AV_CODEC_ID_MPEG1VIDEO : return rocDecVideoCodec_MPEG1;
274  case AV_CODEC_ID_MPEG2VIDEO : return rocDecVideoCodec_MPEG2;
275  case AV_CODEC_ID_MPEG4 : return rocDecVideoCodec_MPEG4;
276  case AV_CODEC_ID_H264 : return rocDecVideoCodec_AVC;
277  case AV_CODEC_ID_HEVC : return rocDecVideoCodec_HEVC;
278  case AV_CODEC_ID_VP8 : return rocDecVideoCodec_VP8;
279  case AV_CODEC_ID_VP9 : return rocDecVideoCodec_VP9;
280  case AV_CODEC_ID_MJPEG : return rocDecVideoCodec_JPEG;
281  case AV_CODEC_ID_AV1 : return rocDecVideoCodec_AV1;
282  default : return rocDecVideoCodec_NumCodecs;
283  }
284 }
Definition: video_demuxer.h:44
Definition: video_demuxer.h:42
The AMD rocDecode Library.
@ rocDecVideoCodec_AV1
Definition: rocdecode.h:84
@ rocDecVideoCodec_MPEG4
Definition: rocdecode.h:81
@ rocDecVideoCodec_HEVC
Definition: rocdecode.h:83
@ rocDecVideoCodec_MPEG1
Definition: rocdecode.h:79
@ rocDecVideoCodec_JPEG
Definition: rocdecode.h:87
@ rocDecVideoCodec_VP9
Definition: rocdecode.h:86
@ rocDecVideoCodec_AVC
Definition: rocdecode.h:82
@ rocDecVideoCodec_MPEG2
Definition: rocdecode.h:80
@ rocDecVideoCodec_VP8
Definition: rocdecode.h:85
@ rocDecVideoCodec_NumCodecs
Definition: rocdecode.h:88