Encode PCM stream into MP4 file stream

PCMEncoder.java

package me.tankery.mediabox.encoder;

import android.annotation.SuppressLint;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;

import timber.log.Timber;

/**
 * PCMEncoder
 *
 * Encode PCM stream into MP4 file stream.
 *
 * Created by tankery on 11/11/2017.
 */
public class PCMEncoder {

    private static final String COMPRESSED_AUDIO_FILE_MIME_TYPE = "audio/mp4a-latm";
    private static final int COMPRESSED_AUDIO_PROFILE = MediaCodecInfo.CodecProfileLevel.AACObjectLC;

    private static final int CODEC_TIMEOUT = 5000;

    private static final int[] SUPPORTED_SAMPLE_RATES = new int[] {
            96000,  // 0: 96000 Hz
            88200,  // 1: 88200 Hz
            64000,  // 2: 64000 Hz
            48000,  // 3: 48000 Hz
            44100,  // 4: 44100 Hz
            32000,  // 5: 32000 Hz
            24000,  // 6: 24000 Hz
            22050,  // 7: 22050 Hz
            16000,  // 8: 16000 Hz
            12000,  // 9: 12000 Hz
            11025,  // 10: 11025 Hz
            8000,   // 11: 8000 Hz
            7350,   // 12: 7350 Hz
                    // 13: Reserved
                    // 14: Reserved
                    // 15: frequency is written explictly
    };

    private final int bitrate;
    private final int sampleRate;
    private final int channelCount;
    private int freqIdx = -1;

    private MediaFormat mediaFormat;
    private MediaCodec mediaCodec;
    private ByteBuffer[] codecInputBuffers;
    private ByteBuffer[] codecOutputBuffers;
    private MediaCodec.BufferInfo bufferInfo;
    private int totalBytesRead;
    private long presentationTimeUs;
    private byte[] tempBuffer;

    /**
     * Creates encoder with given params for output file
     * @param bitrate Bitrate of the output file, higher bitrate brings better voice, but
     *                lager file. eg. 64k is enough for speech, its about 28MB/hour after encode.
     * @param sampleRate sampling rate of pcm.
     * @param channelCount channel count of pcm.
     */
    public PCMEncoder(final int bitrate, final int sampleRate, int channelCount) {
        this.bitrate = bitrate;
        this.sampleRate = sampleRate;
        this.channelCount = channelCount;
        for (int i = 0; i < SUPPORTED_SAMPLE_RATES.length; i++) {
            if (sampleRate == SUPPORTED_SAMPLE_RATES[i]) {
                freqIdx = i;
                break;
            }
        }
        if (freqIdx == -1) {
            throw new IllegalArgumentException("Not support sample rate " + sampleRate);
        }
    }

    /**
     * Encodes input stream
     *
     * @throws IOException
     */
    public void encode(InputStream inputStream, OutputStream outputStream) throws IOException {
        prepare();

        boolean hasMoreData = true;
        while (hasMoreData) {
            hasMoreData = readInputs(inputStream);
            writeOutputs(outputStream);
        }

        inputStream.close();
        outputStream.close();

        stop();
    }

    private void prepare() {
        try {
            Timber.d("Preparing PCMEncoder");

            mediaFormat = MediaFormat.createAudioFormat(COMPRESSED_AUDIO_FILE_MIME_TYPE, sampleRate, channelCount);
            mediaFormat.setInteger(MediaFormat.KEY_AAC_PROFILE, COMPRESSED_AUDIO_PROFILE);
            mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);

            mediaCodec = MediaCodec.createEncoderByType(COMPRESSED_AUDIO_FILE_MIME_TYPE);
            mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
            mediaCodec.start();

            codecInputBuffers = mediaCodec.getInputBuffers();
            codecOutputBuffers = mediaCodec.getOutputBuffers();

            bufferInfo = new MediaCodec.BufferInfo();

            totalBytesRead = 0;
            presentationTimeUs = 0;
            tempBuffer = new byte[2 * sampleRate];
        } catch (IOException e) {
            Timber.e(e, "Exception while initializing PCMEncoder");
        }
    }

    private void stop() {
        Timber.d("Stopping PCMEncoder");
        mediaCodec.stop();
        mediaCodec.release();
    }

    private boolean readInputs(InputStream inputStream) throws IOException {
        boolean hasMoreData = true;
        int inputBufferIndex = 0;
        int currentBatchRead = 0;
        while (inputBufferIndex != -1 && hasMoreData && currentBatchRead <= 50 * sampleRate) {
            inputBufferIndex = mediaCodec.dequeueInputBuffer(CODEC_TIMEOUT);

            if (inputBufferIndex >= 0) {
                ByteBuffer buffer = codecInputBuffers[inputBufferIndex];
                buffer.clear();

                int bytesRead = inputStream.read(tempBuffer, 0, buffer.limit());
                if (bytesRead == -1) {
                    mediaCodec.queueInputBuffer(inputBufferIndex, 0, 0, presentationTimeUs, MediaCodec.BUFFER_FLAG_END_OF_STREAM);
                    hasMoreData = false;
                } else {
                    totalBytesRead += bytesRead;
                    currentBatchRead += bytesRead;
                    buffer.put(tempBuffer, 0, bytesRead);
                    mediaCodec.queueInputBuffer(inputBufferIndex, 0, bytesRead, presentationTimeUs, 0);
                    presentationTimeUs = 1000000L * (totalBytesRead / 2) / sampleRate;
                }
            }
        }
        return hasMoreData;
    }

    @SuppressLint("WrongConstant")
    private void writeOutputs(OutputStream outputStream) throws IOException {
        int outputBufferIndex = 0;
        while (outputBufferIndex != MediaCodec.INFO_TRY_AGAIN_LATER) {
            outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, CODEC_TIMEOUT);
            if (outputBufferIndex >= 0) {
                if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0 && bufferInfo.size != 0) {
                    mediaCodec.releaseOutputBuffer(outputBufferIndex, false);
                } else {
                    // Write ADTS header and AAC data to frame.
                    int outPacketSize = bufferInfo.size + 7;    // 7 is ADTS size
                    byte[] data = new byte[outPacketSize];  //space for ADTS header included
                    addADTStoPacket(data, outPacketSize);

                    ByteBuffer encodedData = codecOutputBuffers[outputBufferIndex];
                    encodedData.position(bufferInfo.offset);
                    encodedData.limit(bufferInfo.offset + bufferInfo.size);
                    encodedData.get(data, 7, bufferInfo.size);
                    encodedData.clear();
                    outputStream.write(data, 0, outPacketSize);  //open FileOutputStream beforehand

                    mediaCodec.releaseOutputBuffer(outputBufferIndex, false);
                }
            } else if (outputBufferIndex == MediaCodec.INFO_OUTPUT_BUFFERS_CHANGED) {
                codecOutputBuffers = mediaCodec.getOutputBuffers();
            } else if (outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                mediaFormat = mediaCodec.getOutputFormat();
            }
        }
    }

    /**
     *  Add ADTS header at the beginning of each and every AAC packet.
     *  This is needed as MediaCodec encoder generates a packet of raw
     *  AAC data.
     *
     *  Note the packetLen must count in the ADTS header itself.
     **/
    private void addADTStoPacket(byte[] packet, int packetLen) {
        int profile = COMPRESSED_AUDIO_PROFILE;
        int freqIdx = this.freqIdx;
        int chanCfg = channelCount;

        // fill in ADTS data
        packet[0] = (byte)0xFF;
        packet[1] = (byte)0xF9;
        packet[2] = (byte)(((profile-1)<<6) + (freqIdx<<2) +(chanCfg>>2));
        packet[3] = (byte)(((chanCfg&3)<<6) + (packetLen>>11));
        packet[4] = (byte)((packetLen&0x7FF) >> 3);
        packet[5] = (byte)(((packetLen&7)<<5) + 0x1F);
        packet[6] = (byte)0xFC;
    }
}