!7

// argument: mkv file to read

// output: png files

// from https://github.com/vzhn/ffmpeg-java-samples/blob/master/src/main/java/DemuxAndDecodeH264.java

lib 1400558 // javacpp ffmpeg examples for Win x64

p { DemuxAndDecodeH264.main(args); }

import org.bytedeco.javacpp.*;

import javax.imageio.ImageIO;

import java.awt.image.BufferedImage;

import java.awt.image.DataBufferByte;

import java.io.File;

import java.io.IOException;

import java.time.Duration;

import java.time.temporal.ChronoUnit;

import static org.bytedeco.javacpp.avcodec.*;

import static org.bytedeco.javacpp.avformat.*;

import static org.bytedeco.javacpp.avutil.*;

import org.bytedeco.javacpp.presets.avutil;

/**

* Read and decode h264 video from matroska (MKV) container

*/

sclass DemuxAndDecodeH264 {

    /** Matroska format context */

    private AVFormatContext avfmtCtx;

    /** Matroska video stream information  */

    private AVStream videoStream;

    /** matroska packet */

    private AVPacket avpacket;

    /** H264 Decoder ID */

    private AVCodec codec;

    /** H264 Decoder context */

    private AVCodecContext codecContext;

    /** yuv420 frame */

    private AVFrame yuv420Frame;

    /** RGB frame */

    private AVFrame rgbFrame;

    /** java RGB frame */

    private BufferedImage img;

    /** yuv420 to rgb converter */

    private swscale.SwsContext sws_ctx;

    /** number of frame */

    private int nframe;

    int maxFrames = 10;

    /* 1/1000 of second */

    private AVRational tb1000;

    private DemuxAndDecodeH264() {

        tb1000 = new AVRational();

        tb1000.num(1);

        tb1000.den(1000);

    }

    public static void main(String... argv) throws IOException {

        new DemuxAndDecodeH264().start(argv);

    }

    private void start(String[] argv) throws IOException {

        av_log_set_level(AV_LOG_VERBOSE);

        openInput(argv[0]);

        findVideoStream();

        initDecoder();

        initRgbFrame();

        initYuv420Frame();

        getSwsContext();

        avpacket = new avcodec.AVPacket();

        while ((av_read_frame(avfmtCtx, avpacket)) >= 0 && !done()) {

            if (avpacket.stream_index() == videoStream.index()) {

                processAVPacket(avpacket);

            }

            av_packet_unref(avpacket);

        }

        // now process delayed frames

        if (!done()) processAVPacket(null);

        free();

    }

    private AVFormatContext openInput(String file) throws IOException {

        avfmtCtx = new AVFormatContext(null);

        BytePointer filePointer = new BytePointer(file);

        int r = avformat.avformat_open_input(avfmtCtx, filePointer, null, null);

        filePointer.deallocate();

        if (r < 0) {

            avfmtCtx.close();

            throw new IOException("avformat_open_input error: " + r);

        }

        return avfmtCtx;

    }

    private void findVideoStream() throws IOException {

        int r = avformat_find_stream_info(avfmtCtx, (PointerPointer) null);

        if (r < 0) {

            avformat_close_input(avfmtCtx);

            avfmtCtx.close();

            throw new IOException("error: " + r);

        }

        PointerPointer<AVCodec> decoderRet = new PointerPointer<>(1);

        int videoStreamNumber = av_find_best_stream(avfmtCtx, AVMEDIA_TYPE_VIDEO, -1, -1, decoderRet, 0);

        if (videoStreamNumber < 0) {

            throw new IOException("failed to find video stream");

        }

        if (decoderRet.get(AVCodec.class).id() != AV_CODEC_ID_H264) {

            throw new IOException("failed to find h264 stream");

        }

        decoderRet.deallocate();

        videoStream =  avfmtCtx.streams(videoStreamNumber);

    }

    private void initDecoder() {

        codec = avcodec_find_decoder(AV_CODEC_ID_H264);

        codecContext = avcodec_alloc_context3(codec);

        if((codec.capabilities() & avcodec.AV_CODEC_CAP_TRUNCATED) != 0) {

            codecContext.flags(codecContext.flags() | avcodec.AV_CODEC_CAP_TRUNCATED);

        }

        avcodec_parameters_to_context(codecContext, videoStream.codecpar());

        if(avcodec_open2(codecContext, codec, (PointerPointer) null) < 0) {

            throw new RuntimeException("Error: could not open codec.\n");

        }

    }

    private void initYuv420Frame() {

        yuv420Frame = av_frame_alloc();

        if (yuv420Frame == null) {

            throw new RuntimeException("Could not allocate video frame\n");

        }

    }

    private void initRgbFrame() {

        rgbFrame = av_frame_alloc();

        rgbFrame.format(AV_PIX_FMT_BGR24);

        rgbFrame.width(codecContext.width());

        rgbFrame.height(codecContext.height());

        int ret = av_image_alloc(rgbFrame.data(),

                rgbFrame.linesize(),

                rgbFrame.width(),

                rgbFrame.height(),

                rgbFrame.format(),

                1);

        if (ret < 0) {

            throw new RuntimeException("could not allocate buffer!");

        }

        img = new BufferedImage(rgbFrame.width(), rgbFrame.height(), BufferedImage.TYPE_3BYTE_BGR);

    }

    private void getSwsContext() {

        sws_ctx = swscale.sws_getContext(

                codecContext.width(), codecContext.height(), codecContext.pix_fmt(),

                rgbFrame.width(), rgbFrame.height(), rgbFrame.format(),

                0, null, null, (DoublePointer) null);

    }

    private void processAVPacket(AVPacket avpacket) throws IOException {

        int ret = avcodec.avcodec_send_packet(codecContext, avpacket);

        if (ret < 0) {

            throw new RuntimeException("Error sending a packet for decoding\n");

        }

        receiveFrames();

    }

    bool done() { ret nframe >= maxFrames; }

    private void receiveFrames() throws IOException {

        int ret = 0;

        while (ret >= 0 && !done()) {

            ret = avcodec.avcodec_receive_frame(codecContext, yuv420Frame);

            if (ret == avutil.AVERROR_EAGAIN()

              || ret == org.bytedeco.javacpp.avutil.AVERROR_EOF()) {

                continue;

            } else

            if (ret < 0) {

                throw new RuntimeException("error during decoding");

            }

            swscale.sws_scale(sws_ctx, yuv420Frame.data(), yuv420Frame.linesize(), 0,

                    yuv420Frame.height(), rgbFrame.data(), rgbFrame.linesize());

            rgbFrame.best_effort_timestamp(yuv420Frame.best_effort_timestamp());

            processFrame(rgbFrame);

        }

    }

    private void processFrame(AVFrame rgbFrame) throws IOException {

        DataBufferByte buffer = (DataBufferByte) img.getRaster().getDataBuffer();

        rgbFrame.data(0).get(buffer.getData());

        long ptsMillis = av_rescale_q(rgbFrame.best_effort_timestamp(), videoStream.time_base(), tb1000);

        Duration d = Duration.of(ptsMillis, ChronoUnit.MILLIS);

        String name = String.format("img_%05d_%02d-%02d-%02d-%03d.png", ++nframe,

                d.toHoursPart(),

                d.toMinutesPart(),

                d.toSecondsPart(),

                d.toMillisPart());

        ImageIO.write(img, "png", new File(name));

    }

    private void free() {

        av_packet_unref(avpacket);

        avcodec.avcodec_close(codecContext);

        avcodec.avcodec_free_context(codecContext);

        swscale.sws_freeContext(sws_ctx);

        av_frame_free(rgbFrame);

        av_frame_free(yuv420Frame);

        avformat.avformat_close_input(avfmtCtx);

        avformat.avformat_free_context(avfmtCtx);

    }

}