mirror of
https://github.com/geekwenjie/SmartJavaAI.git
synced 2026-09-17 15:39:03 +00:00
初始提交
This commit is contained in:
dengwenjie
@@ -0,0 +1,212 @@
|
||||
/*
|
||||
* Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
|
||||
* with the License. A copy of the License is located at
|
||||
*
|
||||
* http://aws.amazon.com/apache2.0/
|
||||
*
|
||||
* or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
|
||||
* OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
|
||||
* and limitations under the License.
|
||||
*/
|
||||
package cn.smartjavaai.face;
|
||||
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.*;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.ndarray.types.Shape;
|
||||
import ai.djl.translate.Translator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* Face Detection Translator
|
||||
*/
|
||||
public class FaceDetectionTranslator implements Translator<Image, DetectedObjects> {
|
||||
|
||||
private double confThresh;
|
||||
private double nmsThresh;
|
||||
private int topK;
|
||||
private double[] variance;
|
||||
private int[][] scales;
|
||||
private int[] steps;
|
||||
|
||||
public FaceDetectionTranslator(
|
||||
double confThresh,
|
||||
double nmsThresh,
|
||||
double[] variance,
|
||||
int topK,
|
||||
int[][] scales,
|
||||
int[] steps) {
|
||||
this.confThresh = confThresh;
|
||||
this.nmsThresh = nmsThresh;
|
||||
this.variance = variance;
|
||||
this.topK = topK;
|
||||
this.scales = scales;
|
||||
this.steps = steps;
|
||||
}
|
||||
|
||||
/** {@inheritDoc} */
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, Image input) {
|
||||
|
||||
ctx.setAttachment("width", input.getWidth());
|
||||
ctx.setAttachment("height", input.getHeight());
|
||||
|
||||
NDArray array = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
|
||||
array = array.transpose(2, 0, 1).flip(0); // HWC -> CHW RGB -> BGR
|
||||
// The network by default takes float32
|
||||
if (!array.getDataType().equals(DataType.FLOAT32)) {
|
||||
array = array.toType(DataType.FLOAT32, false);
|
||||
}
|
||||
NDArray mean =
|
||||
ctx.getNDManager().create(new float[] {104f, 117f, 123f}, new Shape(3, 1, 1));
|
||||
array = array.sub(mean);
|
||||
return new NDList(array);
|
||||
}
|
||||
|
||||
/** {@inheritDoc} */
|
||||
@Override
|
||||
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
|
||||
|
||||
int width = (int) ctx.getAttachment("width");
|
||||
int height = (int) ctx.getAttachment("height");
|
||||
|
||||
NDManager manager = ctx.getNDManager();
|
||||
double scaleXY = variance[0];
|
||||
double scaleWH = variance[1];
|
||||
|
||||
NDArray prob = list.get(1).get(":, 1:");
|
||||
prob =
|
||||
NDArrays.stack(
|
||||
new NDList(
|
||||
prob.argMax(1).toType(DataType.FLOAT32, false),
|
||||
prob.max(new int[] {1})));
|
||||
|
||||
NDArray boxRecover = boxRecover(manager, width, height, scales, steps);
|
||||
NDArray boundingBoxes = list.get(0);
|
||||
NDArray bbWH = boundingBoxes.get(":, 2:").mul(scaleWH).exp().mul(boxRecover.get(":, 2:"));
|
||||
NDArray bbXY =
|
||||
boundingBoxes
|
||||
.get(":, :2")
|
||||
.mul(scaleXY)
|
||||
.mul(boxRecover.get(":, 2:"))
|
||||
.add(boxRecover.get(":, :2"))
|
||||
.sub(bbWH.mul(0.5f));
|
||||
|
||||
boundingBoxes = NDArrays.concat(new NDList(bbXY, bbWH), 1);
|
||||
|
||||
NDArray landms = list.get(2);
|
||||
landms = decodeLandm(landms, boxRecover, scaleXY);
|
||||
|
||||
// filter the result below the threshold
|
||||
NDArray cutOff = prob.get(1).gt(confThresh);
|
||||
boundingBoxes = boundingBoxes.transpose().booleanMask(cutOff, 1).transpose();
|
||||
landms = landms.transpose().booleanMask(cutOff, 1).transpose();
|
||||
prob = prob.booleanMask(cutOff, 1);
|
||||
|
||||
// start categorical filtering
|
||||
long[] order = prob.get(1).argSort().get(":" + topK).toLongArray();
|
||||
prob = prob.transpose();
|
||||
List<String> retNames = new ArrayList<>();
|
||||
List<Double> retProbs = new ArrayList<>();
|
||||
List<BoundingBox> retBB = new ArrayList<>();
|
||||
|
||||
Map<Integer, List<BoundingBox>> recorder = new ConcurrentHashMap<>();
|
||||
|
||||
for (int i = order.length - 1; i >= 0; i--) {
|
||||
long currMaxLoc = order[i];
|
||||
float[] classProb = prob.get(currMaxLoc).toFloatArray();
|
||||
int classId = (int) classProb[0];
|
||||
double probability = classProb[1];
|
||||
|
||||
double[] boxArr = boundingBoxes.get(currMaxLoc).toDoubleArray();
|
||||
double[] landmsArr = landms.get(currMaxLoc).toDoubleArray();
|
||||
Rectangle rect = new Rectangle(boxArr[0], boxArr[1], boxArr[2], boxArr[3]);
|
||||
List<BoundingBox> boxes = recorder.getOrDefault(classId, new ArrayList<>());
|
||||
boolean belowIoU = true;
|
||||
for (BoundingBox box : boxes) {
|
||||
if (box.getIoU(rect) > nmsThresh) {
|
||||
belowIoU = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (belowIoU) {
|
||||
List<Point> keyPoints = new ArrayList<>();
|
||||
for (int j = 0; j < 5; j++) { // 5 face landmarks
|
||||
double x = landmsArr[j * 2];
|
||||
double y = landmsArr[j * 2 + 1];
|
||||
keyPoints.add(new Point(x * width, y * height));
|
||||
}
|
||||
Landmark landmark =
|
||||
new Landmark(boxArr[0], boxArr[1], boxArr[2], boxArr[3], keyPoints);
|
||||
|
||||
boxes.add(landmark);
|
||||
recorder.put(classId, boxes);
|
||||
String className = "Face"; // classes.get(classId)
|
||||
retNames.add(className);
|
||||
retProbs.add(probability);
|
||||
retBB.add(landmark);
|
||||
}
|
||||
}
|
||||
|
||||
return new DetectedObjects(retNames, retProbs, retBB);
|
||||
}
|
||||
|
||||
private NDArray boxRecover(
|
||||
NDManager manager, int width, int height, int[][] scales, int[] steps) {
|
||||
int[][] aspectRatio = new int[steps.length][2];
|
||||
for (int i = 0; i < steps.length; i++) {
|
||||
int wRatio = (int) Math.ceil((float) width / steps[i]);
|
||||
int hRatio = (int) Math.ceil((float) height / steps[i]);
|
||||
aspectRatio[i] = new int[] {hRatio, wRatio};
|
||||
}
|
||||
|
||||
List<double[]> defaultBoxes = new ArrayList<>();
|
||||
|
||||
for (int idx = 0; idx < steps.length; idx++) {
|
||||
int[] scale = scales[idx];
|
||||
for (int h = 0; h < aspectRatio[idx][0]; h++) {
|
||||
for (int w = 0; w < aspectRatio[idx][1]; w++) {
|
||||
for (int i : scale) {
|
||||
double skx = i * 1.0 / width;
|
||||
double sky = i * 1.0 / height;
|
||||
double cx = (w + 0.5) * steps[idx] / width;
|
||||
double cy = (h + 0.5) * steps[idx] / height;
|
||||
defaultBoxes.add(new double[] {cx, cy, skx, sky});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
double[][] boxes = new double[defaultBoxes.size()][defaultBoxes.get(0).length];
|
||||
for (int i = 0; i < defaultBoxes.size(); i++) {
|
||||
boxes[i] = defaultBoxes.get(i);
|
||||
}
|
||||
return manager.create(boxes).clip(0.0, 1.0);
|
||||
}
|
||||
|
||||
// decode face landmarks, 5 points per face
|
||||
private NDArray decodeLandm(NDArray pre, NDArray priors, double scaleXY) {
|
||||
NDArray point1 =
|
||||
pre.get(":, :2").mul(scaleXY).mul(priors.get(":, 2:")).add(priors.get(":, :2"));
|
||||
NDArray point2 =
|
||||
pre.get(":, 2:4").mul(scaleXY).mul(priors.get(":, 2:")).add(priors.get(":, :2"));
|
||||
NDArray point3 =
|
||||
pre.get(":, 4:6").mul(scaleXY).mul(priors.get(":, 2:")).add(priors.get(":, :2"));
|
||||
NDArray point4 =
|
||||
pre.get(":, 6:8").mul(scaleXY).mul(priors.get(":, 2:")).add(priors.get(":, :2"));
|
||||
NDArray point5 =
|
||||
pre.get(":, 8:10").mul(scaleXY).mul(priors.get(":, 2:")).add(priors.get(":, :2"));
|
||||
return NDArrays.concat(new NDList(point1, point2, point3, point4, point5), 1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user