OpenCVUtils.java

package io.lumify.opencvObjectDetector;

import org.opencv.core.Core;
import org.opencv.core.Mat;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.io.ByteArrayInputStream;
import java.io.IOException;

import static org.opencv.core.CvType.*;

public class OpenCVUtils {

    public static Mat bufferedImageToMat(BufferedImage image) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);

        if (image != null) {
            Mat mat;
            int numComponents = image.getColorModel().getNumComponents();
            switch (numComponents) {
                case 1:
                    mat = new Mat(image.getHeight(), image.getWidth(), CV_8UC1);
                    break;
                case 2:
                    mat = new Mat(image.getHeight(), image.getWidth(), CV_8UC2);
                    break;
                case 3:
                    mat = new Mat(image.getHeight(), image.getWidth(), CV_8UC3);
                    break;
                case 4:
                    mat = new Mat(image.getHeight(), image.getWidth(), CV_8UC4);
                    break;
                default:
                    throw new RuntimeException("Image has an unsupportable number of channels: " + numComponents);
            }

            if (image.getType() == BufferedImage.TYPE_BYTE_INDEXED) {
                image = convertByteIndexedImage(image);
            }

            byte[] pixelData = ((DataBufferByte) image.getRaster().getDataBuffer()).getData();
            mat.put(0, 0, pixelData);

            Mat mat3;
            if (numComponents == 3) {
                mat3 = mat;
            } else {
                mat3 = new Mat(image.getHeight(), image.getWidth(), CV_8UC3);
                mat.convertTo(mat3, CV_8UC3);
            }

            return mat3;
        }
        return null;
    }

    public static BufferedImage matToBufferedImage(Mat mat) throws IOException {
        byte[] pixelData = new byte[(int) (mat.total() * mat.channels())];
        mat.get(0, 0, pixelData);

        return ImageIO.read(new ByteArrayInputStream(pixelData));
    }

    private static BufferedImage convertByteIndexedImage(BufferedImage image) {
        BufferedImage convertedBImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_3BYTE_BGR);
        Graphics2D g2 = convertedBImage.createGraphics();
        g2.drawImage(image, 0, 0, null);
        g2.dispose();
        return convertedBImage;
    }
}