基本实现Mac环境下没有大问题

Signed-off-by: JoeAllen <2494136784@qq.com>
This commit is contained in:
2022-07-28 13:34:08 +08:00
parent 7cdf89a036
commit 92166625c2
7 changed files with 737 additions and 0 deletions
+134
View File
@@ -0,0 +1,134 @@
package cn.heavenlybook;
import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;
import cn.heavenlybook.utils.AnimalsClassification;
import cn.heavenlybook.utils.FaceDetection;
import cn.heavenlybook.utils.OpenCVImageUtil;
import lombok.extern.java.Log;
import org.bytedeco.javacv.CanvasFrame;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.javacv.OpenCVFrameGrabber;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_core.Point;
import org.bytedeco.opencv.opencv_core.Rect;
import org.bytedeco.opencv.opencv_core.Scalar;
import javax.swing.*;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.util.List;
import static org.bytedeco.opencv.global.opencv_imgproc.*;
/**
* @author ht
*/
@Log
public class Main {
public static void main(String[] args) throws IOException, ModelException, TranslateException {
float shrink = 0.5f;
float threshold = 0.7f;
Criteria<Image, DetectedObjects> criteria = new FaceDetection().criteria(shrink, threshold);
// 开启摄像头,获取图像(得到的图像为frame类型,需要转换为mat类型进行检测和识别)
OpenCVFrameGrabber grabber = new OpenCVFrameGrabber(0);
grabber.start();
// Frame与Mat转换
OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();
// 新建一个预览窗口
CanvasFrame canvas = new CanvasFrame("人脸检测");
canvas.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
canvas.setVisible(true);
canvas.setFocusable(true);
// 窗口置顶
if (canvas.isAlwaysOnTopSupported()) {
canvas.setAlwaysOnTop(true);
}
Frame frame;
try (ZooModel<Image, DetectedObjects> model = ModelZoo.loadModel(criteria);
Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
// 获取图像帧
while (canvas.isVisible() && (frame = grabber.grab()) != null) {
// 将获取的frame转化成mat数据类型
Mat img = converter.convert(frame);
BufferedImage buffImg = OpenCVImageUtil.mat2BufferedImage(img);
Image image = ImageFactory.getInstance().fromImage(buffImg);
// region 动物识别
Classifications classifications = AnimalsClassification.predict(image);
Classifications.Classification bestItem = classifications.best();
System.out.println(bestItem.getClassName() + " : " + bestItem.getProbability());
// List<Classifications.Classification> items = classifications.items();
// List<String> names = new ArrayList<>();
// List<Double> probs = new ArrayList<>();
// for (int i = 0; i < items.size(); i++) {
// Classifications.Classification item = items.get(i);
// names.add(item.getClassName());
// probs.add(item.getProbability());
// }
log.info(classifications.toJson());
// endregion
int imageWidth = image.getWidth();
int imageHeight = image.getHeight();
DetectedObjects detections = predictor.predict(image);
List<DetectedObjects.DetectedObject> items = detections.items();
// 遍历人脸
for (DetectedObjects.DetectedObject item : items) {
BoundingBox box = item.getBoundingBox();
Rectangle rectangle = box.getBounds();
int x = (int) (rectangle.getX() * imageWidth);
int y = (int) (rectangle.getY() * imageHeight);
Rect face =
new Rect(
x,
y,
(int) (rectangle.getWidth() * imageWidth),
(int) (rectangle.getHeight() * imageHeight));
// 绘制人脸矩形区域,scalar色彩顺序:BGR(蓝绿红)
rectangle(img, face, new Scalar(0, 0, 255, 1));
int posX = Math.max(face.tl().x() - 10, 0);
int posY = Math.max(face.tl().y() - 10, 0);
// 在人脸矩形上面绘制文字
putText(
img,
"Face",
new Point(posX, posY),
FONT_HERSHEY_COMPLEX,
1.0,
new Scalar(0, 0, 255, 2.0));
}
// 显示视频图像
canvas.showImage(frame);
}
}
canvas.dispose();
grabber.close();
}
}
@@ -0,0 +1,80 @@
package cn.heavenlybook.utils;
import ai.djl.Model;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import ai.djl.util.Utils;
import java.io.IOException;
import java.io.InputStream;
import java.util.List;
/**
* @author ht
*/
public class AnimalTranslator implements Translator<Image, Classifications> {
List<String> classes;
public AnimalTranslator() {}
@Override
public void prepare(TranslatorContext ctx) throws IOException {
Model model = ctx.getModel();
try (InputStream is = model.getArtifact("label_list.txt").openStream()) {
classes = Utils.readLines(is, true);
// classes.add(0, "blank");
// classes.add("");
}
}
@Override
public Classifications processOutput(TranslatorContext ctx, NDList list) {
NDArray prob = list.singletonOrThrow();
return new Classifications(this.classes, prob);
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) throws Exception {
NDArray array = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
float percent = 256f / Math.min(input.getWidth(), input.getHeight());
int resizedWidth = Math.round(input.getWidth() * percent);
int resizedHeight = Math.round(input.getHeight() * percent);
// img = img.resize((resizedWidth, resizedHeight), Image.LANCZOS)
array = NDImageUtils.resize(array, resizedWidth, resizedHeight);
array = NDImageUtils.centerCrop(array, 224, 224);
// The network by default takes float32
if (!array.getDataType().equals(DataType.FLOAT32)) {
array = array.toType(DataType.FLOAT32, false);
}
array = array.transpose(2, 0, 1).div(255f); // HWC -> CHW RGB
NDArray mean =
ctx.getNDManager().create(new float[] {0.485f, 0.456f, 0.406f}, new Shape(3, 1, 1));
NDArray std =
ctx.getNDManager().create(new float[] {0.229f, 0.224f, 0.225f}, new Shape(3, 1, 1));
array = array.sub(mean);
array = array.div(std);
array = array.expandDims(0);
return new NDList(array);
}
@Override
public Batchifier getBatchifier() {
return null;
}
}
@@ -0,0 +1,83 @@
package cn.heavenlybook.utils;
import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
// https://www.paddlepaddle.org.cn/hubdetail?name=resnet50_vd_animals&en_category=ImageClassification
//
// Users/calvin/Desktop/Download/browser/PaddleHub-release-v2.1/modules/image/classification/resnet50_vd_animals
/**
* @author ht
*/
public final class AnimalsClassification {
private AnimalsClassification() {}
public static Classifications predict(Image img)
throws IOException, ModelException, TranslateException {
Classifications classifications = AnimalsClassification.classfier(img);
List<Classifications.Classification> items = classifications.items();
double sum = 0;
double max = 0;
double[] probArr = new double[items.size()];
List<String> names = new ArrayList<>();
List<Double> probs = new ArrayList<>();
for (int i = 0; i < items.size(); i++) {
Classifications.Classification item = items.get(i);
double prob = item.getProbability();
probArr[i] = prob;
if (prob > max) {
max = prob;
}
}
for (int i = 0; i < items.size(); i++) {
probArr[i] = Math.exp(probArr[i] - max);
sum = sum + probArr[i];
}
for (int i = 0; i < items.size(); i++) {
Classifications.Classification item = items.get(i);
names.add(item.getClassName());
probs.add(probArr[i]);
}
return new Classifications(names, probs);
}
public static Classifications classfier(Image img)
throws IOException, ModelException, TranslateException {
Criteria<Image, Classifications> criteria =
Criteria.builder()
.optEngine("PaddlePaddle")
.setTypes(Image.class, Classifications.class)
.optModelUrls("https://aias-home.oss-cn-beijing.aliyuncs.com/models/animals.zip")
// .optModelUrls("/Users/calvin/model/animals/")
.optModelName("inference")
.optTranslator(new AnimalTranslator())
.optProgress(new ProgressBar())
.build();
try (ZooModel<Image, Classifications> rotateModel = ModelZoo.loadModel(criteria)) {
try (Predictor<Image, Classifications> classifier = rotateModel.newPredictor()) {
Classifications classifications = classifier.predict(img);
return classifications;
}
}
}
}
@@ -0,0 +1,103 @@
package cn.heavenlybook.utils;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.Shape;
import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import lombok.extern.java.Log;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
/**
* @author ht
*/
@Log
public final class FaceDetection {
public FaceDetection() {}
public Criteria<Image, DetectedObjects> criteria(float shrink, float threshold) {
Criteria<Image, DetectedObjects> criteria =
Criteria.builder()
.optEngine("PaddlePaddle")
.setTypes(Image.class, DetectedObjects.class)
.optModelUrls(
"https://aias-home.oss-cn-beijing.aliyuncs.com/models/face_mask/face_detection.zip")
// .optModelUrls("/Users/calvin/model/face_mask/pyramidbox_lite/")
// .optModelName("inference")
.optProgress(new ProgressBar())
.optTranslator(new FaceTranslator(shrink, threshold))
.build();
return criteria;
}
private static final class FaceTranslator implements Translator<Image, DetectedObjects> {
private final float shrink;
private final float threshold;
private final List<String> className;
FaceTranslator(float shrink, float threshold) {
this.shrink = shrink;
this.threshold = threshold;
className = Arrays.asList("Not Face", "Face");
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
return processImageOutput(list, className, threshold);
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
return processImageInput(ctx.getNDManager(), input, shrink);
}
@Override
public Batchifier getBatchifier() {
return null;
}
NDList processImageInput(NDManager manager, Image input, float shrink) {
NDArray array = input.toNDArray(manager);
Shape shape = array.getShape();
array =
NDImageUtils.resize(array, (int) (shape.get(1) * shrink), (int) (shape.get(0) * shrink));
array = array.transpose(2, 0, 1).flip(0); // HWC -> CHW BGR -> RGB
NDArray mean = manager.create(new float[] {104f, 117f, 123f}, new Shape(3, 1, 1));
array = array.sub(mean).mul(0.007843f); // normalization
array = array.expandDims(0); // make batch dimension
return new NDList(array);
}
DetectedObjects processImageOutput(NDList list, List<String> className, float threshold) {
NDArray result = list.singletonOrThrow();
float[] probabilities = result.get(":,1").toFloatArray();
List<String> names = new ArrayList<>();
List<Double> prob = new ArrayList<>();
List<BoundingBox> boxes = new ArrayList<>();
for (int i = 0; i < probabilities.length; i++) {
if (probabilities[i] >= threshold) {
float[] array = result.get(i).toFloatArray();
names.add(className.get((int) array[0]));
prob.add((double) probabilities[i]);
boxes.add(new Rectangle(array[2], array[3], array[4] - array[2], array[5] - array[3]));
}
}
return new DetectedObjects(names, prob, boxes);
}
}
}
@@ -0,0 +1,79 @@
package cn.heavenlybook.utils;
import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
// https://www.paddlepaddle.org.cn/hubdetail?name=mobilenet_v2_animals&en_category=ImageClassification
//
// /Users/calvin/Desktop/Download/browser/PaddleHub-release-v2.1/modules/image/classification/mobilenet_v2_animals
public final class LightAnimalsClassification {
private LightAnimalsClassification() {}
public static Classifications predict(Image img)
throws IOException, ModelException, TranslateException {
Classifications classifications = LightAnimalsClassification.classfier(img);
List<Classifications.Classification> items = classifications.items();
double sum = 0;
double max = 0;
double[] probArr = new double[items.size()];
List<String> names = new ArrayList<>();
List<Double> probs = new ArrayList<>();
for (int i = 0; i < items.size(); i++) {
Classifications.Classification item = items.get(i);
double prob = item.getProbability();
probArr[i] = prob;
if (prob > max) max = prob;
}
for (int i = 0; i < items.size(); i++) {
probArr[i] = Math.exp(probArr[i] - max);
sum = sum + probArr[i];
}
for (int i = 0; i < items.size(); i++) {
Classifications.Classification item = items.get(i);
names.add(item.getClassName());
probs.add(probArr[i]);
}
return new Classifications(names, probs);
}
public static Classifications classfier(Image img)
throws IOException, ModelException, TranslateException {
Criteria<Image, Classifications> criteria =
Criteria.builder()
.optEngine("PaddlePaddle")
.setTypes(Image.class, Classifications.class)
.optModelUrls(
"https://aias-home.oss-cn-beijing.aliyuncs.com/models/mobilenet_animals.zip")
// .optModelUrls("/Users/calvin/model/mobilenet_animals/")
.optModelName("inference")
.optTranslator(new AnimalTranslator())
.optProgress(new ProgressBar())
.build();
try (ZooModel rotateModel = ModelZoo.loadModel(criteria)) {
try (Predictor<Image, Classifications> classifier = rotateModel.newPredictor()) {
Classifications classifications = classifier.predict(img);
return classifications;
}
}
}
}
@@ -0,0 +1,158 @@
package cn.heavenlybook.utils;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.opencv.opencv_core.Mat;
import java.awt.*;
import java.awt.image.*;
import java.nio.ByteBuffer;
import java.nio.IntBuffer;
/**
* 图片类型转换
*
* @author Calvin
*/
public class OpenCVImageUtil {
/**
* 将 BufferedImage 转 Mat
*
* @param original
*/
public static Mat bufferedImage2Mat(BufferedImage original) {
OpenCVFrameConverter.ToMat cv = new OpenCVFrameConverter.ToMat();
return cv.convertToMat(new Java2DFrameConverter().convert(original));
}
/**
* 将mat转BufferedImage
*
* @param matrix
*/
public static BufferedImage mat2BufferedImage(Mat matrix) {
int cols = matrix.cols();
int rows = matrix.rows();
int elemSize = (int) matrix.elemSize();
byte[] data = new byte[cols * rows * elemSize];
matrix.data().get(data);
int type = 0;
switch (matrix.channels()) {
case 1:
type = BufferedImage.TYPE_BYTE_GRAY;
break;
case 3:
type = BufferedImage.TYPE_3BYTE_BGR;
byte b;
for (int i = 0; i < data.length; i = i + 3) {
b = data[i];
data[i] = data[i + 2];
data[i + 2] = b;
}
break;
default:
return null;
}
BufferedImage image = new BufferedImage(cols, rows, type);
// BufferedImage对象中最重要的两个组件为Raster和ColorModel,分别用于存储图像的像素数据与颜色数据。
// 表示像素矩形数组的类。Raster 封装存储样本值的 DataBuffer,以及描述如何在 DataBuffer 中定位给定样本值的 SampleModel。
// 由于Raster对象是BufferedImage对象中的像素数据存储对象,因此,BufferedImage支持从Raster对象中获取任意位置(x,y)点的像素值p(x,y)。
image.getRaster().setDataElements(0, 0, cols, rows, data);
return image;
}
/**
* 将bufferImage转Mat
*
* @param original
* @param matType
* @param msg
* @param x
* @param y
*/
public static Mat bufferedImage2Mat(
BufferedImage original, int matType, String msg, int x, int y) {
Graphics2D g = original.createGraphics();
try {
g.setComposite(AlphaComposite.Src);
g.drawImage(original, 0, 0, null);
g.drawString(msg, x, y);
} finally {
g.dispose();
}
Mat mat = new Mat(original.getHeight(), original.getWidth(), matType);
mat.data().put(((DataBufferByte) original.getRaster().getDataBuffer()).getData());
return mat;
}
/**
* 24位BGR数组转BufferedImage
*
* @param src -bgr排列的24位图像像素数据数组
* @param width -宽度
* @param height-高度
* @return
*/
public static BufferedImage BGR2BufferedImage(byte[] src, int width, int height) {
BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR);
Raster ra = image.getRaster();
DataBuffer out = ra.getDataBuffer();
DataBufferByte db = (DataBufferByte) out;
ByteBuffer.wrap(db.getData()).put(src, 0, src.length);
return image;
}
/**
* 24位BGR字节缓冲转BufferedImage
*
* @param src -bgr排列的24位图像像素数据字节缓冲
* @param width -宽度
* @param height-高度
* @return
*/
public static BufferedImage BGR2BufferedImage(ByteBuffer src, int width, int height) {
BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR);
Raster ra = image.getRaster();
DataBuffer out = ra.getDataBuffer();
DataBufferByte db = (DataBufferByte) out;
ByteBuffer.wrap(db.getData()).put(src);
return image;
}
/**
* 24位整型BGR字节缓冲转BufferedImage
*
* @param src -bgr排列的24位图像像素整型缓冲(int由3个byte组成)
* @param width -宽度
* @param height-高度
* @return
*/
public static BufferedImage BGR2BufferedImage(IntBuffer src, int width, int height) {
BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_BGR);
Raster ra = image.getRaster();
DataBuffer out = ra.getDataBuffer();
DataBufferInt db = (DataBufferInt) out;
IntBuffer.wrap(db.getData()).put(src);
return image;
}
/**
* 24位整型RGB字节缓冲转BufferedImage
*
* @param src --rgb排列的24位图像像素整型缓冲(int由3个byte组成)
* @param width -宽度
* @param height-高度
* @return
*/
public static BufferedImage RGB2BufferedImage(IntBuffer src, int width, int height) {
BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
Raster ra = image.getRaster();
DataBuffer out = ra.getDataBuffer();
DataBufferInt db = (DataBufferInt) out;
IntBuffer.wrap(db.getData()).put(src);
return image;
}
}