Optimize AvcChunkHandler to use normal ChunkClock if no B Frames.

ba0b991d · Dustin · 5b952294 · ba0b991d · ba0b991d · ba0b991d
Commit ba0b991d authored Feb 04, 2022 by Dustin
Showing with 54 additions and 25 deletions
library/extractor/src/main/java/com/google/android/exoplayer2/extractor/avi/AvcChunkHandler.java
library/extractor/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java
library/extractor/src/test/java/com/google/android/exoplayer2/extractor/avi/AvcChunkPeekerTest.java
--- a/library/extractor/src/main/java/com/google/android/exoplayer2/extractor/avi/AvcChunkHandler.java
+++ b/library/extractor/src/main/java/com/google/android/exoplayer2/extractor/avi/AvcChunkHandler.java
@@ -16,6 +16,7 @@
 package com.google.android.exoplayer2.extractor.avi;
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.extractor.ExtractorInput;
@@ -43,28 +44,36 @@ public class AvcChunkHandler extends NalChunkPeeker {
  public AvcChunkHandler(int id, @NonNull TrackOutput trackOutput,
      @NonNull ChunkClock clock, Format.Builder formatBuilder) {
-    super(id, trackOutput, new PicCountClock(clock.durationUs, clock.chunks), 16);
+    super(id, trackOutput, clock, 16);
    this.formatBuilder = formatBuilder;
  }
-  @NonNull
+  @Nullable
-  @Override
+  @VisibleForTesting
-  public PicCountClock getClock() {
+  PicCountClock getPicCountClock() {
-    return (PicCountClock) clock;
+    if (clock instanceof PicCountClock) {
+      return (PicCountClock)clock;
+    } else {
+      return null;
+    }
  }
  @Override
  boolean skip(byte nalType) {
-    return false;
+    if (clock instanceof PicCountClock) {
+      return false;
+    } else {
+      //If the clock is regular clock, skip "normal" frames
+      return nalType >= 0 && nalType <= NAL_TYPE_IDR;
+    }
  }
  /**
   * Greatly simplified way to calculate the picOrder
   * Full logic is here
   * https://chromium.googlesource.com/chromium/src/media/+/refs/heads/main/video/h264_poc.cc
-   * @param nalTypeOffset
   */
-  void updatePicCountClock(final int nalTypeOffset) {
+  void updatePicCountClock(final int nalTypeOffset, final PicCountClock picCountClock) {
    final ParsableNalUnitBitArray in = new ParsableNalUnitBitArray(buffer, nalTypeOffset + 1, buffer.length);
    //slide_header()
    in.readUnsignedExpGolombCodedInt(); //first_mb_in_slice
@@ -84,10 +93,10 @@ public class AvcChunkHandler extends NalChunkPeeker {
    if (spsData.picOrderCountType == 0) {
      int picOrderCountLsb = in.readBits(spsData.picOrderCntLsbLength);
      //Log.d("Test", "FrameNum: " + frame + " cnt=" + picOrderCountLsb);
-      getClock().setPicCount(picOrderCountLsb);
+      picCountClock.setPicCount(picOrderCountLsb);
      return;
    } else if (spsData.picOrderCountType == 2) {
-      getClock().setPicCount(frameNum);
+      picCountClock.setPicCount(frameNum);
      return;
    }
    clock.setIndex(clock.getIndex());
@@ -98,11 +107,20 @@ public class AvcChunkHandler extends NalChunkPeeker {
    final int spsStart = nalTypeOffset + 1;
    nalTypeOffset = seekNextNal(input, spsStart);
    spsData = NalUnitUtil.parseSpsNalUnitPayload(buffer, spsStart, pos);
-    if (spsData.picOrderCountType == 0) {
+    //If we have B Frames, upgrade to PicCountClock
-      getClock().setMaxPicCount(1 << spsData.picOrderCntLsbLength, 2);
+    final PicCountClock picCountClock;
-    } else if (spsData.picOrderCountType == 2) {
+    if (spsData.maxNumRefFrames > 1 && !(clock instanceof PicCountClock)) {
-      //Plus one because we double the frame number
+      clock = picCountClock = new PicCountClock(clock.durationUs, clock.chunks);
-      getClock().setMaxPicCount(1 << spsData.frameNumLength, 1);
+    } else {
+      picCountClock = getPicCountClock();
+    }
+    if (picCountClock != null) {
+      if (spsData.picOrderCountType == 0) {
+        picCountClock.setMaxPicCount(1 << spsData.picOrderCntLsbLength, 2);
+      } else if (spsData.picOrderCountType == 2) {
+        //Plus one because we double the frame number
+        picCountClock.setMaxPicCount(1 << spsData.frameNumLength, 1);
+      }
    }
    if (spsData.pixelWidthHeightRatio != pixelWidthHeightRatio) {
      pixelWidthHeightRatio = spsData.pixelWidthHeightRatio;
@@ -121,11 +139,17 @@ public class AvcChunkHandler extends NalChunkPeeker {
        case 2:
        case 3:
        case 4:
-          updatePicCountClock(nalTypeOffset);
+          if (clock instanceof PicCountClock) {
+            updatePicCountClock(nalTypeOffset, (PicCountClock)clock);
+          }
          return;
-        case NAL_TYPE_IDR:
+        case NAL_TYPE_IDR: {
-          getClock().syncIndexes();
+          final PicCountClock picCountClock = getPicCountClock();
+          if (picCountClock != null) {
+            picCountClock.syncIndexes();
+          }
          return;
+        }
        case NAL_TYPE_AUD:
        case NAL_TYPE_SEI:
        case NAL_TYPE_PPS: {

--- a/library/extractor/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java
+++ b/library/extractor/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java
@@ -33,6 +33,7 @@ public final class NalUnitUtil {
    public final int constraintsFlagsAndReservedZero2Bits;
    public final int levelIdc;
    public final int seqParameterSetId;
+    public final int maxNumRefFrames;
    public final int width;
    public final int height;
    public final float pixelWidthHeightRatio;
@@ -48,6 +49,7 @@ public final class NalUnitUtil {
        int constraintsFlagsAndReservedZero2Bits,
        int levelIdc,
        int seqParameterSetId,
+        int maxNumRefFrames,
        int width,
        int height,
        float pixelWidthHeightRatio,
@@ -61,6 +63,7 @@ public final class NalUnitUtil {
      this.constraintsFlagsAndReservedZero2Bits = constraintsFlagsAndReservedZero2Bits;
      this.levelIdc = levelIdc;
      this.seqParameterSetId = seqParameterSetId;
+      this.maxNumRefFrames = maxNumRefFrames;
      this.width = width;
      this.height = height;
      this.pixelWidthHeightRatio = pixelWidthHeightRatio;
@@ -367,7 +370,7 @@ public final class NalUnitUtil {
        data.readUnsignedExpGolombCodedInt(); // offset_for_ref_frame[i]
      }
    }
-    data.readUnsignedExpGolombCodedInt(); // max_num_ref_frames
+    int maxNumRefFrames = data.readUnsignedExpGolombCodedInt(); // max_num_ref_frames
    data.skipBit(); // gaps_in_frame_num_value_allowed_flag
    int picWidthInMbs = data.readUnsignedExpGolombCodedInt() + 1;
@@ -427,6 +430,7 @@ public final class NalUnitUtil {
        constraintsFlagsAndReservedZero2Bits,
        levelIdc,
        seqParameterSetId,
+        maxNumRefFrames,
        frameWidth,
        frameHeight,
        pixelWidthHeightRatio,

--- a/library/extractor/src/test/java/com/google/android/exoplayer2/extractor/avi/AvcChunkPeekerTest.java
+++ b/library/extractor/src/test/java/com/google/android/exoplayer2/extractor/avi/AvcChunkPeekerTest.java
@@ -35,8 +35,8 @@ public class AvcChunkPeekerTest {
      setSampleMimeType(MimeTypes.VIDEO_H264).
      setWidth(1280).setHeight(720).setFrameRate(24000f/1001f);
-  private static final byte[] P_SLICE = {00,00,00,01,0x41,(byte)0x9A,0x13,0x36,0x21,0x3A,0x5F,
+  private static final byte[] P_SLICE = {0,0,0,1,0x41,(byte)0x9A,0x13,0x36,0x21,0x3A,0x5F,
-      (byte)0xFE,(byte)0x9E,0x10,00,00};
+      (byte)0xFE,(byte)0x9E,0x10,0,0};
  private FakeTrackOutput fakeTrackOutput;
  private AvcChunkHandler avcChunkPeeker;
@@ -61,19 +61,20 @@ public class AvcChunkPeekerTest {
  @Test
  public void peek_givenStreamHeader() throws IOException {
    peekStreamHeader();
-    final PicCountClock picCountClock = avcChunkPeeker.getClock();
+    final PicCountClock picCountClock = avcChunkPeeker.getPicCountClock();
+    Assert.assertNotNull(picCountClock);
    Assert.assertEquals(64, picCountClock.getMaxPicCount());
    Assert.assertEquals(0, avcChunkPeeker.getSpsData().picOrderCountType);
    Assert.assertEquals(1.18f, fakeTrackOutput.lastFormat.pixelWidthHeightRatio, 0.01f);
  }
  @Test
-  public void peek_givenStreamHeaderAndPSlice() throws IOException {
+  public void newChunk_givenStreamHeaderAndPSlice() throws IOException {
    peekStreamHeader();
-    final PicCountClock picCountClock = avcChunkPeeker.getClock();
+    final PicCountClock picCountClock = avcChunkPeeker.getPicCountClock();
    final FakeExtractorInput input = new FakeExtractorInput.Builder().setData(P_SLICE).build();
-    avcChunkPeeker.peek(input, P_SLICE.length);
+    avcChunkPeeker.newChunk(0, P_SLICE.length, input);
    Assert.assertEquals(12, picCountClock.getLastPicCount());
  }