Make frames decoder to build index without file decoding

dali/operators/reader/loader/video/frames_decoder.cc

@@ -365,33 +365,62 @@ bool FramesDecoder::IsFormatSeekable() {
 }
 
 void FramesDecoder::BuildIndex() {
-  // TODO(awolant): Optimize this function for:
-  //  - CFR
-  //  - index present in the header
-
   index_ = vector<IndexEntry>();
 
   int last_keyframe = -1;
-  while (ReadRegularFrame(nullptr, false)) {
+  while (av_read_frame(av_state_->ctx_, av_state_->packet_) >= 0) {
     IndexEntry entry;
-    entry.is_keyframe = av_state_->frame_->flags & AV_FRAME_FLAG_KEY;
-    entry.pts = av_state_->frame_->pts;
-    entry.is_flush_frame = false;
+    // We want to make sure that we call av_packet_unref in every iteration
+    auto packet = AVPacketScope(av_state_->packet_, av_packet_unref);
 
+    if (packet->stream_index != av_state_->stream_id_) {
+      continue;
+    }
+    if (packet->flags & AV_PKT_FLAG_KEY) {
+      if (index_->size() == 0) {
+        entry.is_keyframe = true;
+      }
+      auto codec_id = av_state_->ctx_->streams[packet->stream_index]->codecpar->codec_id;
+      if (codec_id == AV_CODEC_ID_H264) {
+        if (packet->size >= 5) {  // Ensure there's enough data to inspect
+          uint8_t nal_unit_type = packet->data[4] & 0x1F;
+          if (nal_unit_type == 5) {  // IDR frame for H.264
+            entry.is_keyframe = true;
+          }
+        }
+      } else if (codec_id == AV_CODEC_ID_HEVC) {
+        if (packet->size >= 6) {  // Ensure there's enough data to inspect
+          uint8_t nal_unit_type = (packet->data[4] >> 1) & 0x3F;
+          if (nal_unit_type >= 16 && nal_unit_type <= 21) {
+            entry.is_keyframe = true;
+          }
+        }
+      } else {
+        // for other formats it is enough to check (packet->flags & AV_PKT_FLAG_KEY)
+        entry.is_keyframe = true;
+      }
+    } else {
+      entry.is_keyframe = false;
+    }
     if (entry.is_keyframe) {
       last_keyframe = index_->size();
     }
+    if (packet->pts != AV_NOPTS_VALUE) {
+      entry.pts = packet->pts;
+    } else {
+      entry.pts = packet->dts;
+    }
+    entry.is_flush_frame = false;
     entry.last_keyframe_id = last_keyframe;
     index_->push_back(entry);
   }
-  while (ReadFlushFrame(nullptr, false)) {
-    IndexEntry entry;
-    entry.is_keyframe = false;
-    entry.pts = av_state_->frame_->pts;
-    entry.is_flush_frame = true;
-    entry.last_keyframe_id = last_keyframe;
-    index_->push_back(entry);
-  }
+  av_packet_unref(av_state_->packet_);
+  index_->back().is_flush_frame = true;
+  // Make sure that the index is sorted by pts as some frames may not be in the presentation
+  // order in the container
+  std::sort(index_->begin(), index_->end(), [](const IndexEntry &a, const IndexEntry &b) {
+    return a.pts < b.pts;
+  });
   Reset();
 }