Add very special handling for the MJPG video format

[kaka/cakelight.git] / src / kaka / cakelight / FrameGrabber.java

package kaka.cakelight;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.util.Optional;

import static kaka.cakelight.Main.log;

public class FrameGrabber implements Closeable {
    private Configuration config;
    private File file;
    private int bytesPerFrame;
    private InputStream fileStream;
    private final ByteArrayOutputStream bufferedBytes = new ByteArrayOutputStream();

    private FrameGrabber() {
    }

    public static FrameGrabber from(File videoDevice, Configuration config) {
        FrameGrabber fg = new FrameGrabber();
        fg.config = config;
        fg.file = videoDevice;
        fg.bytesPerFrame = config.video.width * config.video.height * config.video.bpp;
        fg.prepare();
        return fg;
    }

    private boolean prepare() {
        try {
            fileStream = new FileInputStream(file);
            return true;
        } catch (FileNotFoundException e) {
            // TODO: handle java.io.FileNotFoundException: /dev/video1 (Permission denied)
            e.printStackTrace();
            return false;
        }
    }

    /**
     * Must be run in the same thread as {@link #prepare}.
     */
    public Optional<VideoFrame> grabFrame() {
        try {
            byte[] data;
            if (config.video.mjpg) {
                byte[] jpgData = readStreamingJpgData();
                if (jpgData == null) {
                    return Optional.empty();
                }
                saveTemporaryJpgFile(jpgData);
                byte[] bmpData = convertJpgFileToByteArray();
                if (bmpData == null) {
                    return Optional.empty();
                }
                data = bmpData;
            } else {
                data = new byte[bytesPerFrame];
                int count = fileStream.read(data);
                if (count != bytesPerFrame) {
                    log("Expected to read " + bytesPerFrame + " bytes per frame but read " + count);
                }
            }

            return Optional.of(VideoFrame.of(data, config));
        } catch (IOException e) {
            e.printStackTrace();
        }

        return Optional.empty();
    }

    private byte[] readStreamingJpgData() throws IOException {
        byte[] data;
        byte[] batch = new byte[1024];
        boolean lastByteIsXX = false;
        loop:
        while (true) {
            int batchCount = fileStream.read(batch);
            if (batchCount == -1) {
                return null;
            }
            if (lastByteIsXX) {
                if (batch[0] == (byte) 0xd8) {
                    data = bufferedBytes.toByteArray();
                    bufferedBytes.reset();
                    bufferedBytes.write(0xff);
                    bufferedBytes.write(batch, 0, batchCount);
                    break;
                }
                bufferedBytes.write(0xff);
            }
            for (int i = 0; i < batchCount - 1; i++) {
                if (batch[i] == (byte) 0xff && batch[i + 1] == (byte) 0xd8) { // start of jpeg
                    if (i > 0) {
                        bufferedBytes.write(batch, 0, i);
                    }
                    data = bufferedBytes.toByteArray();
                    bufferedBytes.reset();
                    bufferedBytes.write(batch, i, batchCount - i);
                    break loop;
                }
            }
            lastByteIsXX = batch[batchCount - 1] == (byte) 0xff;
            bufferedBytes.write(batch, 0, batchCount - (lastByteIsXX ? 1 : 0));
        }
        return data;
    }

    private void saveTemporaryJpgFile(byte[] data) throws IOException {
        try (FileOutputStream fos = new FileOutputStream("/tmp/cakelight-video-stream.jpg")) {
            fos.write(data);
        }
    }

    private byte[] convertJpgFileToByteArray() throws IOException {
        BufferedImage image = ImageIO.read(new File("/tmp/cakelight-video-stream.jpg"));
        if (image != null) { // will almost always be null the first time
            try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
                ImageIO.write(image, "bmp", baos);
                baos.flush();
                return baos.toByteArray();
            }
        } else {
            return null;
        }
    }

    @Override
    public void close() throws IOException {
        fileStream.close();
    }
}