Add VFR detection (NVIDIA#3921)

Add VFR check to FramesDecoder. Signed-off-by: Albert Wolant <[email protected]>
cyyever · Jun 7, 2022 · b9fc8b1 · b9fc8b1
1 parent b7d98cf
commit b9fc8b1
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 0 deletions.
diff --git a/dali/operators/reader/loader/video/frames_decoder.cc b/dali/operators/reader/loader/video/frames_decoder.cc
@@ -103,6 +103,7 @@ FramesDecoder::FramesDecoder(const std::string &filename)
       " Supported codecs: h264, HEVC."));
   InitAvState();
   BuildIndex();
+  DetectVfr();
 }
 
 void FramesDecoder::BuildIndex() {
@@ -134,6 +135,23 @@ void FramesDecoder::BuildIndex() {
   Reset();
 }
 
+void FramesDecoder::DetectVfr() {
+  if (NumFrames() < 3) {
+    is_vfr_ = false;
+    return;
+  }
+
+  int pts_step = index_[1].pts - index_[0].pts;
+  for (int frame_id = 2; frame_id < NumFrames(); ++frame_id) {
+    if ((index_[frame_id].pts - index_[frame_id - 1].pts) != pts_step) {
+      is_vfr_ = true;
+      return;
+    }
+  }
+
+  is_vfr_ = false;
+}
+
 void FramesDecoder::CopyToOutput(uint8_t *data) {
   LazyInitSwContext();
 

diff --git a/dali/operators/reader/loader/video/frames_decoder.h b/dali/operators/reader/loader/video/frames_decoder.h
@@ -125,6 +125,14 @@ class DLL_PUBLIC FramesDecoder {
     return Channels() * Width() * Height();
   }
 
+    /**
+   * @brief Is video variable frame rate
+   * 
+   */
+  bool IsVfr() const {
+    return is_vfr_;
+  }
+
   /**
    * @brief Reads next frame of the video and copies it to the provided buffer, if copy_to_output is True.
    * 
@@ -203,9 +211,12 @@ class DLL_PUBLIC FramesDecoder {
 
   bool CheckCodecSupport();
 
+  void DetectVfr();
+
   int channels_ = 3;
   bool flush_state_ = false;
   std::string filename_;
+  bool is_vfr_ = false;
 };
 }  // namespace dali
 

diff --git a/dali/operators/reader/loader/video/frames_decoder_test.cc b/dali/operators/reader/loader/video/frames_decoder_test.cc
@@ -36,6 +36,7 @@ class FramesDecoderTestBase : public VideoTestBase {
     ASSERT_EQ(decoder.Width(), ground_truth.Width());
     ASSERT_EQ(decoder.Channels(), ground_truth.NumChannels());
     ASSERT_EQ(decoder.NumFrames(), ground_truth.NumFrames());
+    ASSERT_EQ(decoder.IsVfr(), ground_truth.IsVfr());
 
     // Iterate through the whole video in order
     for (int i = 0; i < decoder.NumFrames(); ++i) {

diff --git a/dali/operators/reader/loader/video/video_test_base.h b/dali/operators/reader/loader/video/video_test_base.h
@@ -60,6 +60,8 @@ class TestVideo {
 
   void CompareFrameAvgError(int frame_id, const uint8_t *frame, double eps = 1.0);
 
+  bool IsVfr() { return is_vfr_; }
+
   std::vector<cv::Mat> frames_;
   bool is_vfr_ = false;
 };