From b0429431023b7412df1122d9b7e25faff38be80f Mon Sep 17 00:00:00 2001 From: sheenachhabra Date: Tue, 12 Sep 2023 08:24:12 -0700 Subject: [PATCH] Add different modes in MuxerWrapper For pause and resume feature we need to use same `MuxerWrapper` for `remuxing processed video` and then to `process remaining video`. In order to use same `MuxerWrapper` across `different Exports` we need to preserve its state. PiperOrigin-RevId: 564728396 --- .../media3/transformer/MuxerWrapper.java | 133 +++++++++-- .../media3/transformer/Transformer.java | 4 +- .../media3/transformer/MuxerWrapperTest.java | 210 ++++++++++++++++++ 3 files changed, 333 insertions(+), 14 deletions(-) create mode 100644 libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java b/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java index 0db0bccac9..ba42f77f0e 100644 --- a/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java +++ b/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java @@ -21,9 +21,11 @@ import static androidx.media3.common.util.Assertions.checkNotNull; import static androidx.media3.common.util.Assertions.checkState; import static androidx.media3.common.util.Util.contains; import static java.lang.Math.max; +import static java.lang.annotation.ElementType.TYPE_USE; import static java.util.concurrent.TimeUnit.MILLISECONDS; import android.util.SparseArray; +import androidx.annotation.IntDef; import androidx.annotation.IntRange; import androidx.annotation.Nullable; import androidx.media3.common.C; @@ -34,6 +36,10 @@ import androidx.media3.common.util.Util; import androidx.media3.effect.DebugTraceUtil; import com.google.common.collect.ImmutableList; import java.io.File; +import java.lang.annotation.Documented; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; import java.nio.ByteBuffer; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -47,6 +53,26 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; *

This wrapper can contain at most one video track and one audio track. */ /* package */ final class MuxerWrapper { + /** Different modes for muxing. */ + @Documented + @Retention(RetentionPolicy.SOURCE) + @Target(TYPE_USE) + @IntDef({MUXER_MODE_DEFAULT, MUXER_MODE_MUX_PARTIAL_VIDEO, MUXER_MODE_APPEND_VIDEO}) + public @interface MuxerMode {} + + /** The default muxer mode. */ + public static final int MUXER_MODE_DEFAULT = 0; + + /** + * Used for muxing a partial video. The video {@link TrackInfo} is kept the same when {@linkplain + * #changeToAppendVideoMode() transitioning} to {@link #MUXER_MODE_APPEND_VIDEO} after finishing + * muxing partial video. Only one video track can be {@link #addTrackFormat(Format) added} in this + * mode. + */ + public static final int MUXER_MODE_MUX_PARTIAL_VIDEO = 1; + + /** Used for appending the remaining video samples with the previously muxed partial video. */ + public static final int MUXER_MODE_APPEND_VIDEO = 2; private static final String TIMER_THREAD_NAME = "Muxer:Timer"; private static final String MUXER_TIMEOUT_ERROR_FORMAT_STRING = @@ -83,21 +109,45 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; private @MonotonicNonNull ScheduledFuture abortScheduledFuture; private boolean isAborted; private @MonotonicNonNull Muxer muxer; + private @MuxerMode int muxerMode; + private volatile boolean muxedPartialVideo; private volatile int additionalRotationDegrees; private volatile int trackCount; - /** Creates an instance. */ - public MuxerWrapper(String outputPath, Muxer.Factory muxerFactory, Listener listener) { + /** + * Creates an instance. + * + * @param outputPath The output file path to write the media data to. + * @param muxerFactory A {@link Muxer.Factory} to create a {@link Muxer}. + * @param listener A {@link MuxerWrapper.Listener}. + * @param muxerMode The {@link MuxerMode}. The initial mode must be {@link #MUXER_MODE_DEFAULT} or + * {@link #MUXER_MODE_MUX_PARTIAL_VIDEO}. + */ + public MuxerWrapper( + String outputPath, Muxer.Factory muxerFactory, Listener listener, @MuxerMode int muxerMode) { this.outputPath = outputPath; this.muxerFactory = muxerFactory; this.listener = listener; - + checkArgument(muxerMode == MUXER_MODE_DEFAULT || muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO); + this.muxerMode = muxerMode; trackTypeToInfo = new SparseArray<>(); previousTrackType = C.TRACK_TYPE_NONE; abortScheduledExecutorService = Util.newSingleThreadScheduledExecutor(TIMER_THREAD_NAME); } + /** + * Changes {@link MuxerMode} to {@link #MUXER_MODE_APPEND_VIDEO}. + * + *

This method must be called only after partial video is muxed using {@link + * #MUXER_MODE_MUX_PARTIAL_VIDEO}. + */ + public void changeToAppendVideoMode() { + checkState(muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO); + + muxerMode = MUXER_MODE_APPEND_VIDEO; + } + /** * Sets the clockwise rotation to add to the {@linkplain #addTrackFormat(Format) video track's} * rotation, in degrees. @@ -123,12 +173,23 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; *

The track count must be set before any track format is {@linkplain #addTrackFormat(Format) * added}. * + *

When using muxer mode other than {@link #MUXER_MODE_DEFAULT}, the track count must be 1. + * *

Can be called from any thread. * * @throws IllegalStateException If a track format was {@linkplain #addTrackFormat(Format) added} * before calling this method. */ public void setTrackCount(@IntRange(from = 1) int trackCount) { + if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO || muxerMode == MUXER_MODE_APPEND_VIDEO) { + checkArgument( + trackCount == 1, + "Only one video track can be added in MUXER_MODE_MUX_PARTIAL_VIDEO and" + + " MUXER_MODE_APPEND_VIDEO"); + if (muxerMode == MUXER_MODE_APPEND_VIDEO) { + return; + } + } checkState( trackTypeToInfo.size() == 0, "The track count cannot be changed after adding track formats."); @@ -158,8 +219,11 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; * *

{@link Muxer#addMetadata(Metadata)} is called if the {@link Format#metadata} is present. * - * @param format The {@link Format} to be added. - * @throws IllegalArgumentException If the format is unsupported. + * @param format The {@link Format} to be added. In {@link #MUXER_MODE_APPEND_VIDEO} mode, the + * added {@link Format} must match the existing {@link Format} set when the muxer was in + * {@link #MUXER_MODE_MUX_PARTIAL_VIDEO} mode. + * @throws IllegalArgumentException If the format is unsupported or if it does not match the + * existing format in {@link #MUXER_MODE_APPEND_VIDEO} mode. * @throws IllegalStateException If the number of formats added exceeds the {@linkplain * #setTrackCount track count}, if {@link #setTrackCount(int)} has not been called or if there * is already a track of that {@link C.TrackType}. @@ -167,6 +231,27 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; * the track. */ public void addTrackFormat(Format format) throws Muxer.MuxerException { + if (muxerMode == MUXER_MODE_APPEND_VIDEO) { + checkState(contains(trackTypeToInfo, C.TRACK_TYPE_VIDEO)); + TrackInfo videoTrackInfo = trackTypeToInfo.get(C.TRACK_TYPE_VIDEO); + + // Ensure that video formats are the same. Some fields like codecs, averageBitrate, framerate, + // etc, don't match exactly in the Extractor output format and the Encoder output + // format but these fields can be ignored. + Format existingFormat = videoTrackInfo.format; + checkArgument(Util.areEqual(existingFormat.sampleMimeType, format.sampleMimeType)); + checkArgument(existingFormat.width == format.width); + checkArgument(existingFormat.height == format.height); + checkArgument(existingFormat.initializationDataEquals(format)); + checkArgument(Util.areEqual(existingFormat.colorInfo, format.colorInfo)); + + checkNotNull(muxer); + resetAbortTimer(); + return; + } else if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO) { + checkArgument(MimeTypes.isVideo(format.sampleMimeType)); + } + int trackCount = this.trackCount; checkState(trackCount > 0, "The track count should be set before the formats are added."); checkState(trackTypeToInfo.size() < trackCount, "All track formats have already been added."); @@ -275,25 +360,44 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; DebugTraceUtil.logEvent(DebugTraceUtil.EVENT_MUXER_TRACK_ENDED_AUDIO, trackInfo.timeUs); } - trackTypeToInfo.delete(trackType); - if (trackTypeToInfo.size() == 0) { - abortScheduledExecutorService.shutdownNow(); - if (!isEnded) { + if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO) { + muxedPartialVideo = true; + } else { + trackTypeToInfo.delete(trackType); + if (trackTypeToInfo.size() == 0) { isEnded = true; - listener.onEnded(Util.usToMs(maxEndedTrackTimeUs), getCurrentOutputSizeBytes()); } } + + if (muxedPartialVideo) { + listener.onEnded(Util.usToMs(maxEndedTrackTimeUs), getCurrentOutputSizeBytes()); + if (abortScheduledFuture != null) { + abortScheduledFuture.cancel(/* mayInterruptIfRunning= */ false); + } + return; + } + + if (isEnded) { + listener.onEnded(Util.usToMs(maxEndedTrackTimeUs), getCurrentOutputSizeBytes()); + abortScheduledExecutorService.shutdownNow(); + } } - /** Returns whether all the tracks are {@linkplain #endTrack(int) ended}. */ + /** + * Returns whether all the tracks are {@linkplain #endTrack(int) ended} or a partial video is + * completely muxed using {@link #MUXER_MODE_MUX_PARTIAL_VIDEO}. + */ public boolean isEnded() { - return isEnded; + return isEnded || (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO && muxedPartialVideo); } /** * Finishes writing the output and releases any resources associated with muxing. * - *

The muxer cannot be used anymore once this method has been called. + *

When this method is called in {@link #MUXER_MODE_MUX_PARTIAL_VIDEO} mode, the resources are + * not released and the {@link MuxerWrapper} can be reused after {@link #changeToAppendVideoMode() + * changing mode} to {@link #MUXER_MODE_APPEND_VIDEO} mode. In all other modes the {@link + * MuxerWrapper} cannot be used anymore once this method has been called. * * @param forCancellation Whether the reason for releasing the resources is the transformation * cancellation. @@ -301,6 +405,9 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull; * and {@code forCancellation} is false. */ public void release(boolean forCancellation) throws Muxer.MuxerException { + if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO && !forCancellation) { + return; + } isReady = false; abortScheduledExecutorService.shutdownNow(); if (muxer != null) { diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/Transformer.java b/libraries/transformer/src/main/java/androidx/media3/transformer/Transformer.java index 816eca5350..0bee9c8058 100644 --- a/libraries/transformer/src/main/java/androidx/media3/transformer/Transformer.java +++ b/libraries/transformer/src/main/java/androidx/media3/transformer/Transformer.java @@ -832,7 +832,9 @@ public final class Transformer { public void start(Composition composition, String path) { ComponentListener componentListener = new ComponentListener(composition); startInternal( - composition, new MuxerWrapper(path, muxerFactory, componentListener), componentListener); + composition, + new MuxerWrapper(path, muxerFactory, componentListener, MuxerWrapper.MUXER_MODE_DEFAULT), + componentListener); } /** diff --git a/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java b/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java new file mode 100644 index 0000000000..af7f174373 --- /dev/null +++ b/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java @@ -0,0 +1,210 @@ +/* + * Copyright 2023 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package androidx.media3.transformer; + +import static androidx.media3.common.MimeTypes.VIDEO_H264; +import static androidx.media3.transformer.MuxerWrapper.MUXER_MODE_MUX_PARTIAL_VIDEO; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import androidx.media3.common.C; +import androidx.media3.common.ColorInfo; +import androidx.media3.common.Format; +import androidx.test.ext.junit.runners.AndroidJUnit4; +import com.google.common.collect.ImmutableList; +import java.io.IOException; +import java.nio.ByteBuffer; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; + +/** Unit tests for {@link MuxerWrapper}. */ +@RunWith(AndroidJUnit4.class) +public class MuxerWrapperTest { + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private static final Format FAKE_VIDEO_TRACK_FORMAT = + new Format.Builder() + .setSampleMimeType(VIDEO_H264) + .setWidth(1080) + .setHeight(720) + .setInitializationData(ImmutableList.of(new byte[] {1, 2, 3, 4})) + .setColorInfo(ColorInfo.SDR_BT709_LIMITED) + .build(); + private static final Format FAKE_AUDIO_TRACK_FORMAT = + new Format.Builder() + .setSampleMimeType("audio/mp4a-latm") + .setSampleRate(40000) + .setChannelCount(2) + .build(); + + private static final ByteBuffer FAKE_SAMPLE = ByteBuffer.wrap(new byte[] {1, 2, 3, 4}); + + private String outputFilePath; + + @Before + public void setUp() throws IOException { + outputFilePath = temporaryFolder.newFile("output.mp4").getPath(); + } + + @Test + public void changeToAppendVideoMode_afterDefaultMode_throws() { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MuxerWrapper.MUXER_MODE_DEFAULT); + + assertThrows(IllegalStateException.class, muxerWrapper::changeToAppendVideoMode); + } + + @Test + public void setTrackCount_toTwoInMuxPartialVideoMode_throws() { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + + assertThrows(IllegalArgumentException.class, () -> muxerWrapper.setTrackCount(2)); + } + + @Test + public void setTrackCount_toTwoInAppendVideoMode_throws() throws Exception { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + muxerWrapper.setTrackCount(1); + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + muxerWrapper.writeSample( + C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0); + muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO); + muxerWrapper.changeToAppendVideoMode(); + + assertThrows(IllegalArgumentException.class, () -> muxerWrapper.setTrackCount(2)); + } + + @Test + public void addTrackFormat_withAudioFormatInMuxPartialVideoMode_throws() { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + muxerWrapper.setTrackCount(1); + + assertThrows( + IllegalArgumentException.class, () -> muxerWrapper.addTrackFormat(FAKE_AUDIO_TRACK_FORMAT)); + } + + @Test + public void addTrackFormat_withSameVideoFormatInAppendVideoMode_doesNotThrow() throws Exception { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + muxerWrapper.setTrackCount(1); + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + muxerWrapper.writeSample( + C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0); + muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO); + muxerWrapper.changeToAppendVideoMode(); + muxerWrapper.setTrackCount(1); + + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + } + + @Test + public void addTrackFormat_withDifferentVideoFormatInAppendVideoMode_throws() throws Exception { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + muxerWrapper.setTrackCount(1); + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + muxerWrapper.writeSample( + C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0); + muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO); + muxerWrapper.changeToAppendVideoMode(); + muxerWrapper.setTrackCount(1); + Format differentVideoFormat = FAKE_VIDEO_TRACK_FORMAT.buildUpon().setHeight(5000).build(); + + assertThrows( + IllegalArgumentException.class, () -> muxerWrapper.addTrackFormat(differentVideoFormat)); + } + + @Test + public void isEnded_afterPartialVideoMuxed_returnsTrue() throws Exception { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + muxerWrapper.setTrackCount(1); + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + muxerWrapper.writeSample( + C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0); + muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO); + + assertThat(muxerWrapper.isEnded()).isTrue(); + } + + @Test + public void isEnded_afterStartingAppendVideo_returnsFalse() throws Exception { + MuxerWrapper muxerWrapper = + new MuxerWrapper( + outputFilePath, + new DefaultMuxer.Factory(), + new NoOpMuxerListenerImpl(), + MUXER_MODE_MUX_PARTIAL_VIDEO); + muxerWrapper.setTrackCount(1); + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + muxerWrapper.writeSample( + C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0); + muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO); + muxerWrapper.changeToAppendVideoMode(); + muxerWrapper.setTrackCount(1); + muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT); + + assertThat(muxerWrapper.isEnded()).isFalse(); + } + + private static final class NoOpMuxerListenerImpl implements MuxerWrapper.Listener { + + @Override + public void onTrackEnded( + @C.TrackType int trackType, Format format, int averageBitrate, int sampleCount) {} + + @Override + public void onEnded(long durationMs, long fileSizeBytes) {} + + @Override + public void onError(ExportException exportException) {} + } +}