All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.jcodec.codecs.h264.H264Encoder Maven / Gradle / Ivy

There is a newer version: 0.2.5
Show newest version
package org.jcodec.codecs.h264;
import static java.lang.System.arraycopy;
import static org.jcodec.codecs.h264.H264Utils.escapeNAL;

import org.jcodec.codecs.h264.encode.DumbRateControl;
import org.jcodec.codecs.h264.encode.EncodedMB;
import org.jcodec.codecs.h264.encode.MBEncoderHelper;
import org.jcodec.codecs.h264.encode.MBEncoderI16x16;
import org.jcodec.codecs.h264.encode.MBEncoderP16x16;
import org.jcodec.codecs.h264.encode.MotionEstimator;
import org.jcodec.codecs.h264.encode.RateControl;
import org.jcodec.codecs.h264.io.CAVLC;
import org.jcodec.codecs.h264.io.model.MBType;
import org.jcodec.codecs.h264.io.model.NALUnit;
import org.jcodec.codecs.h264.io.model.NALUnitType;
import org.jcodec.codecs.h264.io.model.PictureParameterSet;
import org.jcodec.codecs.h264.io.model.RefPicMarkingIDR;
import org.jcodec.codecs.h264.io.model.SeqParameterSet;
import org.jcodec.codecs.h264.io.model.SliceHeader;
import org.jcodec.codecs.h264.io.model.SliceType;
import org.jcodec.codecs.h264.io.write.CAVLCWriter;
import org.jcodec.codecs.h264.io.write.SliceHeaderWriter;
import org.jcodec.common.VideoEncoder;
import org.jcodec.common.io.BitWriter;
import org.jcodec.common.logging.Logger;
import org.jcodec.common.model.ColorSpace;
import org.jcodec.common.model.Picture;
import org.jcodec.common.model.Size;
import org.jcodec.common.tools.MathUtil;

import java.nio.ByteBuffer;

/**
 * This class is part of JCodec ( www.jcodec.org ) This software is distributed
 * under FreeBSD License
 * 
 * MPEG 4 AVC ( H.264 ) Encoder
 * 
 * Conforms to H.264 ( ISO/IEC 14496-10 ) specifications
 * 
 * @author The JCodec project
 * 
 */
public class H264Encoder extends VideoEncoder {

    // private static final int QP = 20;
    private static final int KEY_INTERVAL_DEFAULT = 25;

    public static H264Encoder createH264Encoder() {
        return new H264Encoder(new DumbRateControl());
    }

    private CAVLC[] cavlc;
    private byte[][] leftRow;
    private byte[][] topLine;
    private RateControl rc;
    private int frameNumber;
    private int keyInterval;

    private int maxPOC;

    private int maxFrameNumber;

    private SeqParameterSet sps;

    private PictureParameterSet pps;

    private MBEncoderI16x16 mbEncoderI16x16;

    private MBEncoderP16x16 mbEncoderP16x16;

    private Picture ref;
    private Picture picOut;
    private EncodedMB[] topEncoded;

    private EncodedMB outMB;

    public H264Encoder(RateControl rc) {
        this.rc = rc;
        this.keyInterval = KEY_INTERVAL_DEFAULT;
    }

    public int getKeyInterval() {
        return keyInterval;
    }

    public void setKeyInterval(int keyInterval) {
        this.keyInterval = keyInterval;
    }
    
    /**
     * Encode this picture into h.264 frame. Frame type will be selected by
     * encoder.
     */
    public EncodedFrame encodeFrame(Picture pic, ByteBuffer _out) {
        if (pic.getColor() != ColorSpace.YUV420J)
            throw new IllegalArgumentException("Input picture color is not supported: " + pic.getColor());
        
        if (frameNumber >= keyInterval) {
            frameNumber = 0;
        }

        SliceType sliceType = frameNumber == 0 ? SliceType.I : SliceType.P;
        boolean idr = frameNumber == 0;

        ByteBuffer data = doEncodeFrame(pic, _out, idr, frameNumber++, sliceType);
        return new EncodedFrame(data, idr);
    }

    /**
     * Encode this picture as an IDR frame. IDR frame starts a new independently
     * decodeable video sequence
     * 
     * @param pic
     * @param _out
     * @return
     */
    public ByteBuffer encodeIDRFrame(Picture pic, ByteBuffer _out) {
        frameNumber = 0;
        return doEncodeFrame(pic, _out, true, frameNumber, SliceType.I);
    }

    /**
     * Encode this picture as a P-frame. P-frame is an frame predicted from one
     * or more of the previosly decoded frame and is usually 10x less in size
     * then the IDR frame.
     * 
     * @param pic
     * @param _out
     * @return
     */
    public ByteBuffer encodePFrame(Picture pic, ByteBuffer _out) {
        frameNumber++;
        return doEncodeFrame(pic, _out, true, frameNumber, SliceType.P);
    }

    public ByteBuffer doEncodeFrame(Picture pic, ByteBuffer _out, boolean idr, int frameNumber, SliceType frameType) {
        ByteBuffer dup = _out.duplicate();

        if (idr) {
            sps = initSPS(new Size(pic.getCroppedWidth(), pic.getCroppedHeight()));
            pps = initPPS();

            maxPOC = 1 << (sps.log2MaxPicOrderCntLsbMinus4 + 4);
            maxFrameNumber = 1 << (sps.log2MaxFrameNumMinus4 + 4);
        }

        if (idr) {
            dup.putInt(0x1);
            new NALUnit(NALUnitType.SPS, 3).write(dup);
            writeSPS(dup, sps);

            dup.putInt(0x1);
            new NALUnit(NALUnitType.PPS, 3).write(dup);
            writePPS(dup, pps);
        }

        int mbWidth = sps.picWidthInMbsMinus1 + 1;
        int mbHeight = sps.picHeightInMapUnitsMinus1 + 1;

        leftRow = new byte[][] { new byte[16], new byte[8], new byte[8] };
        topLine = new byte[][] { new byte[mbWidth << 4], new byte[mbWidth << 3], new byte[mbWidth << 3] };
        picOut = Picture.create(mbWidth << 4, mbHeight << 4, ColorSpace.YUV420J);

        outMB = new EncodedMB();
        topEncoded = new EncodedMB[mbWidth];
        for (int i = 0; i < mbWidth; i++)
            topEncoded[i] = new EncodedMB();

        encodeSlice(sps, pps, pic, dup, idr, frameNumber, frameType);

        putLastMBLine();

        ref = picOut;

        dup.flip();
        return dup;
    }

    private void writePPS(ByteBuffer dup, PictureParameterSet pps) {
        ByteBuffer tmp = ByteBuffer.allocate(1024);
        pps.write(tmp);
        tmp.flip();
        escapeNAL(tmp, dup);
    }

    private void writeSPS(ByteBuffer dup, SeqParameterSet sps) {
        ByteBuffer tmp = ByteBuffer.allocate(1024);
        sps.write(tmp);
        tmp.flip();
        escapeNAL(tmp, dup);
    }

    public PictureParameterSet initPPS() {
        PictureParameterSet pps = new PictureParameterSet();
        pps.picInitQpMinus26 = rc.getInitQp(SliceType.I) - 26;
        return pps;
    }

    public SeqParameterSet initSPS(Size sz) {
        SeqParameterSet sps = new SeqParameterSet();
        sps.picWidthInMbsMinus1 = ((sz.getWidth() + 15) >> 4) - 1;
        sps.picHeightInMapUnitsMinus1 = ((sz.getHeight() + 15) >> 4) - 1;
        sps.chromaFormatIdc = ColorSpace.YUV420J;
        sps.profileIdc = 66;
        sps.levelIdc = 40;
        sps.frameMbsOnlyFlag = true;
        sps.log2MaxFrameNumMinus4 = Math.max(0, MathUtil.log2(keyInterval) - 3);

        int codedWidth = (sps.picWidthInMbsMinus1 + 1) << 4;
        int codedHeight = (sps.picHeightInMapUnitsMinus1 + 1) << 4;
        sps.frameCroppingFlag = codedWidth != sz.getWidth() || codedHeight != sz.getHeight();
        sps.frameCropRightOffset = (codedWidth - sz.getWidth() + 1) >> 1;
        sps.frameCropBottomOffset = (codedHeight - sz.getHeight() + 1) >> 1;

        return sps;
    }

    private void encodeSlice(SeqParameterSet sps, PictureParameterSet pps, Picture pic, ByteBuffer dup, boolean idr,
            int frameNum, SliceType sliceType) {
        if (idr && sliceType != SliceType.I) {
            idr = false;
            Logger.warn("Illegal value of idr = true when sliceType != I");
        }
        cavlc = new CAVLC[] { new CAVLC(sps, pps, 2, 2), new CAVLC(sps, pps, 1, 1), new CAVLC(sps, pps, 1, 1) };
        mbEncoderI16x16 = new MBEncoderI16x16(cavlc, leftRow, topLine);
        mbEncoderP16x16 = new MBEncoderP16x16(sps, ref, cavlc, new MotionEstimator(16));

        rc.reset();
        int qp = rc.getInitQp(sliceType);

        dup.putInt(0x1);
        new NALUnit(idr ? NALUnitType.IDR_SLICE : NALUnitType.NON_IDR_SLICE, 3).write(dup);
        SliceHeader sh = new SliceHeader();
        sh.sliceType = sliceType;
        if (idr)
            sh.refPicMarkingIDR = new RefPicMarkingIDR(false, false);
        sh.pps = pps;
        sh.sps = sps;
        sh.picOrderCntLsb = (frameNum << 1) % maxPOC;
        sh.frameNum = frameNum % maxFrameNumber;
        sh.sliceQpDelta = qp - (pps.picInitQpMinus26 + 26);

        ByteBuffer buf = ByteBuffer.allocate(pic.getWidth() * pic.getHeight());
        BitWriter sliceData = new BitWriter(buf);
        new SliceHeaderWriter().write(sh, idr, 2, sliceData);

        for (int mbY = 0; mbY < sps.picHeightInMapUnitsMinus1 + 1; mbY++) {
            for (int mbX = 0; mbX < sps.picWidthInMbsMinus1 + 1; mbX++) {
                if (sliceType == SliceType.P) {
                    CAVLCWriter.writeUE(sliceData, 0); // number of skipped mbs
                }

                MBType mbType = selectMBType(sliceType);

                if (mbType == MBType.I_16x16) {
                    // I16x16 carries part of layout information in the
                    // macroblock type
                    // itself for this reason we'll have to decide it now to
                    // embed into
                    // macroblock type
                    int predMode = mbEncoderI16x16.getPredMode(pic, mbX, mbY);
                    int cbpChroma = mbEncoderI16x16.getCbpChroma(pic, mbX, mbY);
                    int cbpLuma = mbEncoderI16x16.getCbpLuma(pic, mbX, mbY);

                    int i16x16TypeOffset = (cbpLuma / 15) * 12 + cbpChroma * 4 + predMode;
                    int mbTypeOffset = sliceType == SliceType.P ? 5 : 0;

                    CAVLCWriter.writeUE(sliceData, mbTypeOffset + mbType.code() + i16x16TypeOffset);
                } else {
                    CAVLCWriter.writeUE(sliceData, mbType.code());
                }

                BitWriter candidate;
                int qpDelta;
                do {
                    candidate = sliceData.fork();
                    qpDelta = rc.getQpDelta();
                    encodeMacroblock(mbType, pic, mbX, mbY, candidate, qp, qpDelta);
                } while (!rc.accept(candidate.position() - sliceData.position()));
                sliceData = candidate;
                qp += qpDelta;

                collectPredictors(outMB.getPixels(), mbX);
                addToReference(mbX, mbY);
            }
        }
        sliceData.write1Bit(1);
        sliceData.flush();
        buf = sliceData.getBuffer();
        buf.flip();

        escapeNAL(buf, dup);
    }

    private void encodeMacroblock(MBType mbType, Picture pic, int mbX, int mbY, BitWriter candidate, int qp, int qpDelta) {
        if (mbType == MBType.I_16x16)
            mbEncoderI16x16.encodeMacroblock(pic, mbX, mbY, candidate, outMB, mbX > 0 ? topEncoded[mbX - 1] : null,
                    mbY > 0 ? topEncoded[mbX] : null, qp + qpDelta, qpDelta);
        else if (mbType == MBType.P_16x16)
            mbEncoderP16x16.encodeMacroblock(pic, mbX, mbY, candidate, outMB, mbX > 0 ? topEncoded[mbX - 1] : null,
                    mbY > 0 ? topEncoded[mbX] : null, qp + qpDelta, qpDelta);
        else
            throw new RuntimeException("Macroblock of type " + mbType + " is not supported.");
    }

    private MBType selectMBType(SliceType sliceType) {
        if (sliceType == SliceType.I)
            return MBType.I_16x16;
        else if (sliceType == SliceType.P)
            return MBType.P_16x16;
        else
            throw new RuntimeException("Unsupported slice type");
    }

    private void addToReference(int mbX, int mbY) {
        if (mbY > 0)
            MBEncoderHelper.putBlkPic(picOut, topEncoded[mbX].getPixels(), mbX << 4, (mbY - 1) << 4);
        EncodedMB tmp = topEncoded[mbX];
        topEncoded[mbX] = outMB;
        outMB = tmp;
    }

    private void putLastMBLine() {
        int mbWidth = sps.picWidthInMbsMinus1 + 1;
        int mbHeight = sps.picHeightInMapUnitsMinus1 + 1;
        for (int mbX = 0; mbX < mbWidth; mbX++)
            MBEncoderHelper.putBlkPic(picOut, topEncoded[mbX].getPixels(), mbX << 4, (mbHeight - 1) << 4);
    }

    private void collectPredictors(Picture outMB, int mbX) {
        arraycopy(outMB.getPlaneData(0), 240, topLine[0], mbX << 4, 16);
        arraycopy(outMB.getPlaneData(1), 56, topLine[1], mbX << 3, 8);
        arraycopy(outMB.getPlaneData(2), 56, topLine[2], mbX << 3, 8);

        copyCol(outMB.getPlaneData(0), 15, 16, leftRow[0]);
        copyCol(outMB.getPlaneData(1), 7, 8, leftRow[1]);
        copyCol(outMB.getPlaneData(2), 7, 8, leftRow[2]);
    }

    private void copyCol(byte[] planeData, int off, int stride, byte[] out) {
        for (int i = 0; i < out.length; i++) {
            out[i] = planeData[off];
            off += stride;
        }
    }

    @Override
    public ColorSpace[] getSupportedColorSpaces() {
        return new ColorSpace[] { ColorSpace.YUV420J };
    }

    @Override
    public int estimateBufferSize(Picture frame) {
        return Math.max(1 << 16, frame.getWidth() * frame.getHeight() / 2);
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy