Review fix

Signed-off-by: Janusz Lisiecki <[email protected]>
NVIDIA · Feb 5, 2025 · 8aad640 · 8aad640
1 parent 522b1dd
commit 8aad640
Showing 1 changed file with 16 additions and 14 deletions.
diff --git a/dali/operators/reader/loader/video/frames_decoder.cc b/dali/operators/reader/loader/video/frames_decoder.cc
@@ -377,25 +377,27 @@ void FramesDecoder::BuildIndex() {
       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
-          // NAL unit type is in the fifth byte for H.264
-          uint8_t nal_unit_type = packet->data[4] & 0x1F;
-          if (nal_unit_type == 5) {  // IDR frame for H.264
-            entry.is_keyframe = true;
+        for (int i = 0; i < packet->size - 4; i++) {
+          // find start of the NAL unit and then check it type
+          if (packet->data[i] == 0x00 && packet->data[i+1] == 0x00 &&
+              packet->data[i+2] == 0x01) {
+            uint8_t nal_unit_type = packet->data[i+3] & 0x1F;
+            if (nal_unit_type == 5) {
+              entry.is_keyframe = true;
+            }
           }
         }
       } else if (codec_id == AV_CODEC_ID_HEVC) {
-      // Further check if it's an IDR frame for H.265
-        if (packet->size >= 6) {  // Ensure there's enough data to inspect
-          // NAL unit type is in the sixth byte for H.265
-          uint8_t nal_unit_type = (packet->data[4] >> 1) & 0x3F;
-          if (nal_unit_type >= 16 && nal_unit_type <= 21) {  // IDR frame types for H.265
-            entry.is_keyframe = true;
+        for (int i = 0; i < packet->size - 4; i++) {
+          // find start of the NAL unit and then check it type
+          if (packet->data[i] == 0x00 && packet->data[i+1] == 0x00 &&
+              packet->data[i+2] == 0x01) {
+            uint8_t nal_unit_type = (packet->data[i+3] >> 1) & 0x3F;
+            if (nal_unit_type >= 16 && nal_unit_type <= 21) {
+              entry.is_keyframe = true;
+            }
           }
         }
       } else {