临时提交

This commit is contained in:
dengwenjie
2025-08-31 18:41:27 +08:00
parent 86ea7eb03e
commit 2b044fda29
25 changed files with 617 additions and 1005 deletions

View File

@@ -9,24 +9,12 @@ import lombok.Data;
public enum FaceDetModelEnum {
RETINA_FACE("PyTorch",0,0, "https://resources.djl.ai/test-models/pytorch/retinaface.zip"),
RETINA_FACE_ONNX("OnnxRuntime",0,0, null),
RETINA_FACE_1080_720_ONNX("OnnxRuntime",1080,720, null),
RETINA_FACE_640_ONNX("OnnxRuntime",640,640, null),
RETINA_FACE_320_ONNX("OnnxRuntime",320,320, null),
RETINA_FACE_720_1280_ONNX("OnnxRuntime",720,1280, null),
RETINA_FACE_MOBILE_ONNX("OnnxRuntime",0,0, null),
RETINA_FACE_MOBILE_320_ONNX("OnnxRuntime",320,320, null),
RETINA_FACE_MOBILE_640_ONNX("OnnxRuntime",640,640, null),
RETINA_FACE_MOBILE_720_1280_ONNX("OnnxRuntime",720,1080, null),
ULTRA_LIGHT_FAST_GENERIC_FACE("PyTorch",0,0, "https://resources.djl.ai/test-models/pytorch/ultranet.zip"),
SEETA_FACE6_MODEL(null,0,0, null),
YOLOV8_FACE("OnnxRuntime",640,640, null),
YOLOV5_FACE_640("OnnxRuntime", 640,640, null),
YOLOV5_FACE_320("OnnxRuntime",320,320, null),
SCRFD_160("OnnxRuntime",160,160, null),
SCRFD_320("OnnxRuntime",320,320, null),
SCRFD_640("OnnxRuntime",640,640, null),
SCRFD_1280("OnnxRuntime",1280,1280, null),
MTCNN("OnnxRuntime",1280,1280, null);

View File

@@ -7,6 +7,7 @@ import cn.smartjavaai.face.enums.FaceDetModelEnum;
import cn.smartjavaai.face.exception.FaceException;
import cn.smartjavaai.face.model.facedect.CommonFaceDetModel;
import cn.smartjavaai.face.model.facedect.FaceDetModel;
import cn.smartjavaai.face.model.facedect.MtcnnFaceDetModel;
import cn.smartjavaai.face.model.facedect.SeetaFace6FaceDetModel;
import cn.smartjavaai.face.model.facerec.*;
import lombok.extern.slf4j.Slf4j;
@@ -122,11 +123,13 @@ public class FaceDetModelFactory {
// 初始化默认算法
static {
registerAlgorithm(FaceDetModelEnum.RETINA_FACE, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.RETINA_FACE_640_ONNX, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.RETINA_FACE_1080_720_ONNX, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.ULTRA_LIGHT_FAST_GENERIC_FACE, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.SEETA_FACE6_MODEL, SeetaFace6FaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.YOLOV8_FACE, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.YOLOV5_FACE_640, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.YOLOV5_FACE_320, CommonFaceDetModel.class);
registerAlgorithm(FaceDetModelEnum.MTCNN, MtcnnFaceDetModel.class);
log.debug("缓存目录:{}", Config.getCachePath());
}

View File

@@ -19,10 +19,7 @@ import cn.smartjavaai.common.entity.*;
import cn.smartjavaai.common.entity.Point;
import cn.smartjavaai.common.entity.face.FaceInfo;
import cn.smartjavaai.common.pool.PredictorFactory;
import cn.smartjavaai.common.utils.Base64ImageUtils;
import cn.smartjavaai.common.utils.FileUtils;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import cn.smartjavaai.common.utils.*;
import cn.smartjavaai.face.config.FaceDetConfig;
import cn.smartjavaai.face.exception.FaceException;
import cn.smartjavaai.face.model.facedect.criterial.FaceDetCriteriaFactory;
@@ -85,8 +82,8 @@ public class MtcnnFaceDetModel implements FaceDetModel{
Path rnetPath = modelPath.resolve("rnet_script.pt");
Path onetPath = modelPath.resolve("onet_script.pt");
pNetModel = getModel(pnetPath);
rNetModel = getModel(pnetPath);
oNetModel = getModel(pnetPath);
rNetModel = getModel(rnetPath);
oNetModel = getModel(onetPath);
this.pnetPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(pNetModel));
this.rnetPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(rNetModel));
@@ -266,22 +263,84 @@ public class MtcnnFaceDetModel implements FaceDetModel{
* @return
*/
public R<DetectionResponse> detect(Image image){
try (NDManager manager = NDManager.newBaseManager(pNetModel.getNDManager().getDevice())){
List<Double> scales = MtcnnProcess.generateScales(image);
NDArray imgs = MtcnnProcess.processInput(manager, image);
Predictor<NDList, NDList> pNetPredictor = null;
Predictor<NDList, NDList> rNetPredictor = null;
Predictor<NDList, NDList> oNetPredictor = null;
try (NDManager manager = pNetModel.getNDManager().newSubManager();){
pNetPredictor = pnetPredictorPool.borrowObject();
rNetPredictor = rnetPredictorPool.borrowObject();
oNetPredictor = onetPredictorPool.borrowObject();
int h = image.getHeight();
int w = image.getWidth();
NDList outputPnet = PNetModel.firstStage(manager, pnetPredictorPool.borrowObject(), imgs, scales, w, h);
//第一阶段
NDList outputPnet = PNetModel.firstStage(manager, pNetPredictor, image);
if(CollectionUtils.isEmpty(outputPnet)){
return R.fail(R.Status.NO_FACE_DETECTED);
}
NDArray boxes = outputPnet.get(0);
NDArray image_inds = outputPnet.get(1);
NDArray imgs = outputPnet.get(2);
if(DJLCommonUtils.isNDArrayEmpty(boxes) || DJLCommonUtils.isNDArrayEmpty(image_inds) || DJLCommonUtils.isNDArrayEmpty(imgs)){
return R.fail(R.Status.NO_FACE_DETECTED);
}
NDList pad = MtcnnUtils.pad(boxes, w, h);
NDList outputRnet = RNetModel.secondStage(manager, rnetPredictorPool.borrowObject(), imgs,boxes,pad, image_inds);
//第二阶段
NDList outputRnet = RNetModel.secondStage(manager, rNetPredictor, imgs,boxes,pad, image_inds);
if(CollectionUtils.isEmpty(outputRnet)){
return R.fail(R.Status.NO_FACE_DETECTED);
}
NDArray image_indsFiltered = outputRnet.get(0);
NDArray scoresFiltered = outputRnet.get(1);
MtcnnBatchResult oNetResult = ONetModel.thirdStage(manager, onetPredictorPool.borrowObject(), imgs,boxes, w, h, scoresFiltered, image_indsFiltered);
return R.ok(convertToDetectionResponse(oNetResult));
boxes = outputRnet.get(2);
if(DJLCommonUtils.isNDArrayEmpty(boxes) || DJLCommonUtils.isNDArrayEmpty(image_indsFiltered) || DJLCommonUtils.isNDArrayEmpty(scoresFiltered)){
return R.fail(R.Status.NO_FACE_DETECTED);
}
//第三阶段
MtcnnBatchResult oNetResult = ONetModel.thirdStage(manager, oNetPredictor, imgs,boxes, w, h, scoresFiltered, image_indsFiltered);
DetectionResponse detectionResponse = convertToDetectionResponse(oNetResult);
if(Objects.isNull(detectionResponse)){
return R.fail(R.Status.NO_FACE_DETECTED);
}
return R.ok(detectionResponse);
} catch (Exception e) {
throw new RuntimeException(e);
} finally {
if (pNetPredictor != null) {
try {
pnetPredictorPool.returnObject(pNetPredictor); //归还
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
pNetPredictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
if (rNetPredictor != null) {
try {
rnetPredictorPool.returnObject(rNetPredictor); //归还
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
rNetPredictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
if (oNetPredictor != null) {
try {
onetPredictorPool.returnObject(oNetPredictor); //归还
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
oNetPredictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
}
}
@@ -306,6 +365,9 @@ public class MtcnnFaceDetModel implements FaceDetModel{
NDArray probs = mtcnnBatchResult.probs.get(0);
NDArray points = mtcnnBatchResult.points.get(0);
if (DJLCommonUtils.isNDArrayEmpty(boxes) || DJLCommonUtils.isNDArrayEmpty(probs) || DJLCommonUtils.isNDArrayEmpty(points)){
return null;
}
long numBoxes = boxes.getShape().get(0);
for (int i = 0; i < numBoxes; i++) {
float[] boxCoords = boxes.get(i).toFloatArray(); // [x1, y1, x2, y2]

View File

@@ -8,6 +8,7 @@ import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.Translator;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import cn.smartjavaai.face.config.FaceDetConfig;
import cn.smartjavaai.face.config.FaceExpressionConfig;
import cn.smartjavaai.face.constant.FaceDetectConstant;
@@ -19,9 +20,7 @@ import cn.smartjavaai.face.exception.FaceException;
import cn.smartjavaai.face.model.expression.translator.DenseNetEmotionTranslator;
import cn.smartjavaai.face.model.expression.translator.FrEmotionTranslator;
import cn.smartjavaai.face.translator.FaceDetectionTranslator;
import cn.smartjavaai.face.translator.SCRFDFaceTranslator;
import cn.smartjavaai.face.translator.YoloV5FaceTranslator;
import cn.smartjavaai.face.translator.YoloV8FaceTranslator;
import org.apache.commons.lang3.StringUtils;
import java.nio.file.Paths;
@@ -73,33 +72,19 @@ public class FaceDetCriteriaFactory {
*/
public static Translator<Image, DetectedObjects> getTranslator(FaceDetConfig config) {
Translator<Image, DetectedObjects> translator = null;
if(config.getModelEnum() == FaceDetModelEnum.RETINA_FACE){
if(config.getModelEnum() == FaceDetModelEnum.RETINA_FACE || config.getModelEnum() == FaceDetModelEnum.RETINA_FACE_640_ONNX
|| config.getModelEnum() == FaceDetModelEnum.RETINA_FACE_1080_720_ONNX){
translator =
new FaceDetectionTranslator(config.getConfidenceThreshold(), config.getNmsThresh(), RetinaFaceConstant.variance, FaceDetectConstant.MAX_FACE_LIMIT, RetinaFaceConstant.scales, RetinaFaceConstant.steps);
new FaceDetectionTranslator(config.getConfidenceThreshold(), config.getNmsThresh(), RetinaFaceConstant.variance, FaceDetectConstant.MAX_FACE_LIMIT,
RetinaFaceConstant.scales, RetinaFaceConstant.steps, config.getModelEnum().getInputWidth(), config.getModelEnum().getInputHeight());
}else if (config.getModelEnum() == FaceDetModelEnum.ULTRA_LIGHT_FAST_GENERIC_FACE){
translator =
new FaceDetectionTranslator(config.getConfidenceThreshold(), config.getNmsThresh(), UltraLightFastGenericFaceConstant.variance, FaceDetectConstant.MAX_FACE_LIMIT, UltraLightFastGenericFaceConstant.scales, UltraLightFastGenericFaceConstant.steps);
}else if (config.getModelEnum() == FaceDetModelEnum.YOLOV8_FACE){
Map<String, Object> arguments = new HashMap<>();
arguments.put("width", 640);
arguments.put("height", 640);
arguments.put("resizeShort", 640);
arguments.put("centerFit", true);
translator = YoloV8FaceTranslator.builder(arguments).build();
new FaceDetectionTranslator(config.getConfidenceThreshold(), config.getNmsThresh(), UltraLightFastGenericFaceConstant.variance,
FaceDetectConstant.MAX_FACE_LIMIT, UltraLightFastGenericFaceConstant.scales, UltraLightFastGenericFaceConstant.steps, config.getModelEnum().getInputWidth(), config.getModelEnum().getInputHeight());
}else if (config.getModelEnum() == FaceDetModelEnum.YOLOV5_FACE_640
|| config.getModelEnum() == FaceDetModelEnum.YOLOV5_FACE_320){
Map<String, Object> arguments = new HashMap<>();
arguments.put("width", config.getModelEnum().getInputWidth());
arguments.put("height", config.getModelEnum().getInputHeight());
arguments.put("resizeShort", true);
arguments.put("centerFit", true);
translator = YoloV5FaceTranslator.builder(arguments).build();
}else if (config.getModelEnum() == FaceDetModelEnum.SCRFD_160
|| config.getModelEnum() == FaceDetModelEnum.SCRFD_320
|| config.getModelEnum() == FaceDetModelEnum.SCRFD_640
|| config.getModelEnum() == FaceDetModelEnum.SCRFD_1280){
translator =
new SCRFDFaceTranslator(config.getConfidenceThreshold(), config.getNmsThresh(), 100, new int[]{8, 16, 32});
translator = YoloV5FaceTranslator.builder()
.setImageSize(config.getModelEnum().getInputWidth(), config.getModelEnum().getInputWidth()).build();
}
return translator;
}

View File

@@ -65,7 +65,7 @@ public class ONetModel {
NDArray out0 = out.get(0).transpose(1, 0); // permute(1,0)
NDArray out1 = out.get(1).transpose(1, 0);
NDArray out2 = out.get(2).transpose(1, 0);
NDArray score = out1.get(1); // out1[1, :]
NDArray score = out2.get(1); // out1[1, :]
points = out1.duplicate();
NDArray ipass = score.gt(0.7); // score > threshold[1]
// ipass 为布尔/0-1张量长度应等于 points 的第 1 维(这里是 7
@@ -84,7 +84,7 @@ public class ONetModel {
NDArray boxesSelected = boxes.get(manager.create(validIndices)); // 行筛选
// 取前 4 列
boxesSelected = boxesSelected.get(":, 0:4"); // 只保留前 4 列
scoresFiltered = scoresFiltered.get(ipass).reshape(-1, 1); // score[ipass].unsqueeze(1)
scoresFiltered = score.get(ipass).reshape(-1, 1); // score[ipass].unsqueeze(1)
boxes = NDArrays.concat(new NDList(boxesSelected, scoresFiltered), 1); // 拼接成 (N,5)
// 筛选 image_inds
@@ -92,7 +92,6 @@ public class ONetModel {
NDArray mv = out0.transpose() // (N, 4)
.get(ipass); // 1-D 花式索引在第 0 维,得到 (k, 4)
System.out.println("----------");
// w_i = boxes[:, 2] - boxes[:, 0] + 1
NDArray w_i = boxes.get(":,2").sub(boxes.get(":,0")).add(1);

View File

@@ -23,53 +23,80 @@ import java.util.List;
public class PNetModel {
/**
* 输入图片
* @param imgs
* @param w
* @param h
* @param scale
* 生成金字塔缩放比例列表
* @param image
* @return
*/
public static NDList processInput(NDArray imgs, int w, int h, double scale){
int newH = (int) (h * scale + 1);
int newW = (int) (w * scale + 1);
// (N, H, W, C)
NDArray transposed = imgs.transpose(0, 2, 3, 1);
transposed = NDImageUtils.resize(transposed, newW, newH, Image.Interpolation.AREA);
// (N, C, H, W)
transposed = transposed.transpose(0, 3, 1, 2);
// 归一化
transposed = transposed.sub(127.5).mul(0.0078125f);
System.out.println("inputPnet: " + transposed.getShape());
return new NDList(transposed);
}
public static List<Double> generateScales(Image image){
long h = image.getHeight();
long w = image.getWidth();
// 计算最小缩放比例
double minsize = 20;
double m = 12.0 / minsize;
double minl = Math.min(h, w) * m;
public static NDList processOutput(NDList outputPnet, double scale, NDManager manager){
NDArray reg = outputPnet.get(0); // [B, 4, H, W]
NDArray probs = outputPnet.get(1); // [B, 2, H, W]
List<NDArray> boundingBox = generateBoundingBox(reg, probs.get(":, 1"), (float)scale, 0.6f);
NDArray boxes_scale = boundingBox.get(0); // [N,9]
NDArray imgIndND = boundingBox.get(1); // [N]
NDArray pick = NMSUtils.batchedNms(boxes_scale.get(":,:4"), boxes_scale.get(":,4"), imgIndND, 0.5f, manager);
return new NDList(boxes_scale, imgIndND, pick);
// 创建金字塔缩放比例列表
double factor = 0.709; // 你原代码的 factor
List<Double> scales = new ArrayList<>();
double scale_i = m;
while (minl >= 12) {
scales.add(scale_i);
scale_i *= factor;
minl *= factor;
}
return scales;
}
public static NDList firstStage(NDManager manager, Predictor<NDList, NDList> pnetPredictor, NDArray imgs, List<Double> scales, int width, int height) throws TranslateException {
public static NDArray pNetPre(Image input,NDManager manager){
// Image -> NDArray (H, W, C)
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
// 增加 batch 维度 (1, H, W, C) ----
array = array.expandDims(0);
// 交换维度 (N, C, H, W)
array = array.transpose(0, 3, 1, 2);
// 转成模型的数据类型
if (!array.getDataType().equals(DataType.FLOAT32)) {
array = array.toType(DataType.FLOAT32, false);
}
return array;
}
public static NDList firstStage(NDManager manager, Predictor<NDList, NDList> pnetPredictor, Image image) throws TranslateException {
List<Double> scales = MtcnnProcess.generateScales(image);
NDArray imgs = pNetPre(image,manager);
int h = image.getHeight();
int w = image.getWidth();
// 第一阶段
NDList boxes_list = new NDList();
NDList image_inds_list = new NDList();
NDList scale_picks_list = new NDList();
int offset = 0;
for (double scale : scales) {
NDList inputPnet = processInput(imgs, width, height, scale);
NDList outputPnet = pnetPredictor.predict(inputPnet);
NDList output = processOutput(outputPnet, scale, manager);
NDArray boxes_scale = output.get(0);
NDArray imgIndND = output.get(1);
NDArray pick = output.get(2);
int newH = (int) (h * scale + 1);
int newW = (int) (w * scale + 1);
// (N, H, W, C)
NDArray transposed = imgs.transpose(0, 2, 3, 1);
transposed = NDImageUtils.resize(transposed, newW, newH, Image.Interpolation.AREA);
// (N, C, H, W)
transposed = transposed.transpose(0, 3, 1, 2);
// 归一化
transposed = transposed.sub(127.5).mul(0.0078125f);
NDList outputPnet = pnetPredictor.predict(new NDList(transposed));
NDArray reg = outputPnet.get(0); // [B, 4, H, W]
NDArray probs = outputPnet.get(1); // [B, 2, H, W]
List<NDArray> boundingBox = generateBoundingBox(reg, probs.get(":, 1"), (float)scale, 0.6f);
NDArray boxes_scale = boundingBox.get(0); // [N,9]
NDArray imgIndND = boundingBox.get(1); // [N]
NDArray pick = NMSUtils.batchedNms(boxes_scale.get(":,:4"), boxes_scale.get(":,4"), imgIndND, 0.5f,manager);
boxes_list.add(boxes_scale);
image_inds_list.add(imgIndND);
scale_picks_list.add(pick.add(offset));
@@ -96,7 +123,6 @@ public class PNetModel {
boxes = boxes.get(pick);
image_inds = image_inds.get(pick);
System.out.println(Arrays.toString(boxes.get(0).toFloatArray()));
NDArray regw = boxes.get(":, 2").sub(boxes.get(":, 0"));
NDArray regh = boxes.get(":, 3").sub(boxes.get(":, 1"));
@@ -108,7 +134,7 @@ public class PNetModel {
boxes = NDArrays.stack(new NDList(qq1, qq2, qq3, qq4, boxes.get(":, 4")), 1);
boxes = MtcnnUtils.rerec(boxes);
return new NDList(boxes, image_inds);
return new NDList(boxes, image_inds, imgs);
}
/**
@@ -131,13 +157,8 @@ public class PNetModel {
// mask = probs >= thresh -> [B,H,W]
NDArray mask = probs.gte(threshold);
System.out.println("reg: " + reg.getShape());
System.out.println("probs shape: " + probs.getShape());
System.out.println("scale: " + scale);
System.out.println("mask: " + mask.getShape());
// mask_inds = mask.nonzero() -> [N,3] 每行: (batch, y, x)
NDArray maskInds = mask.nonzero();
System.out.println("maskInds: " + maskInds.getShape());
// image_inds = mask_inds[:, 0]
NDArray imageInds = maskInds.get(":,0");

View File

@@ -101,7 +101,7 @@ public class RNetModel {
log.debug("No face detected.");
return null;
}
return new NDList(image_indsFiltered, scoresFiltered);
return new NDList(image_indsFiltered, scoresFiltered,boxes);
}
}

View File

@@ -13,6 +13,7 @@
package cn.smartjavaai.face.translator;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.*;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDArrays;
@@ -22,7 +23,11 @@ import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
@@ -40,29 +45,50 @@ public class FaceDetectionTranslator implements Translator<Image, DetectedObject
private int[][] scales;
private int[] steps;
private int inputWidth = 0;
private int inputHeight = 0;
public FaceDetectionTranslator(
double confThresh,
double nmsThresh,
double[] variance,
int topK,
int[][] scales,
int[] steps) {
int[] steps,
int inputWidth,
int inputHeight) {
this.confThresh = confThresh;
this.nmsThresh = nmsThresh;
this.variance = variance;
this.topK = topK;
this.scales = scales;
this.steps = steps;
this.inputWidth = inputWidth;
this.inputHeight = inputHeight;
}
/** {@inheritDoc} */
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
ctx.setAttachment("sourceWidth", input.getWidth());
ctx.setAttachment("sourceHeight", input.getHeight());
ctx.setAttachment("width", inputWidth);
ctx.setAttachment("height", inputHeight);
NDArray array = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
if(inputWidth > 0 && inputHeight > 0){
//Letter box resize 640x640 with padding (保持比例,补边缘)
LetterBoxUtils.ResizeResult letterBoxResult = LetterBoxUtils.letterbox(ctx.getNDManager(), array, inputWidth, inputHeight, 114f, LetterBoxUtils.PaddingPosition.CENTER);
array = letterBoxResult.image;
ctx.setAttachment("needRecover", "1");//需要还原
ctx.setAttachment("scale", letterBoxResult.r);
}else{
ctx.setAttachment("needRecover", "0");//不需要还原
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
}
array = array.transpose(2, 0, 1).flip(0); // HWC -> CHW RGB -> BGR
// The network by default takes float32
if (!array.getDataType().equals(DataType.FLOAT32)) {
@@ -80,6 +106,13 @@ public class FaceDetectionTranslator implements Translator<Image, DetectedObject
int width = (int) ctx.getAttachment("width");
int height = (int) ctx.getAttachment("height");
int sourceWidth = (int) ctx.getAttachment("sourceWidth");
int sourceHeight = (int) ctx.getAttachment("sourceHeight");
String needRecover = (String) ctx.getAttachment("needRecover");
float scale = 0;
if("1".equals(needRecover)){
scale = (float) ctx.getAttachment("scale");
}
NDManager manager = ctx.getNDManager();
double scaleXY = variance[0];
@@ -91,7 +124,6 @@ public class FaceDetectionTranslator implements Translator<Image, DetectedObject
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:"));
@@ -141,6 +173,7 @@ public class FaceDetectionTranslator implements Translator<Image, DetectedObject
}
}
if (belowIoU) {
List<Point> keyPoints = new ArrayList<>();
for (int j = 0; j < 5; j++) { // 5 face landmarks
double x = landmsArr[j * 2];
@@ -150,9 +183,14 @@ public class FaceDetectionTranslator implements Translator<Image, DetectedObject
Landmark landmark =
new Landmark(boxArr[0], boxArr[1], boxArr[2], boxArr[3], keyPoints);
if(needRecover.equals("1")){
landmark = LetterBoxUtils.restoreBox(landmark, scale, sourceWidth, sourceHeight, width, height, true);
}
boxes.add(landmark);
recorder.put(classId, boxes);
String className = "Face"; // classes.get(classId)
int percent = (int) Math.round(probability * 100);
String className = "face " + percent + "%"; // classes.get(classId)
retNames.add(className);
retProbs.add(probability);
retBB.add(landmark);

View File

@@ -1,242 +0,0 @@
package cn.smartjavaai.face.translator;
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.Landmark;
import ai.djl.modality.cv.output.Point;
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.translate.Translator;
import ai.djl.translate.TranslatorContext;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import cn.smartjavaai.common.utils.NMSUtils;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* SCRFD Translator
*/
public class SCRFDFaceTranslator implements Translator<Image, DetectedObjects> {
private double confThresh;
private double nmsThresh;
private int topK;
private int[] steps;
private int inputWidth = 640;
private int inputHeight = 640;
public SCRFDFaceTranslator(
double confThresh,
double nmsThresh,
int topK,
int[] steps) {
this.confThresh = confThresh;
this.nmsThresh = nmsThresh;
this.topK = topK;
this.steps = steps;
}
/** {@inheritDoc} */
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
NDManager manager = ctx.getNDManager();
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
//Letter box resize 640x640 with padding (保持比例,补边缘)
LetterBoxUtils.ResizeResult letterBoxResult = LetterBoxUtils.letterbox(manager, array, inputWidth, inputHeight, 0f, LetterBoxUtils.PaddingPosition.LEFT_TOP);
ctx.setAttachment("scale", letterBoxResult.r);
array = letterBoxResult.image;
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);
}
// 归一化
array = array.sub(127.5).mul(0.0078125);
return new NDList(array);
}
/** {@inheritDoc} */
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
NDManager manager =
NDManager.newBaseManager(ctx.getNDManager().getDevice(), "PyTorch");
int sourceWidth = (int) ctx.getAttachment("width");
int sourceHeight = (int) ctx.getAttachment("height");
float detScale = (float) ctx.getAttachment("scale");
Map<String, NDArray> centerCache = new HashMap<>();
List<NDArray> scores_list = new ArrayList<>();
List<NDArray> bboxes_list = new ArrayList<>();
List<NDArray> kpss_list = new ArrayList<>();
// 步长长度
int fmc = steps.length;
int numAnchors = 2;
// 多尺度后处理
for (int idx = 0; idx < steps.length; idx++) {
int stride = steps[idx];
NDArray scores, bboxPreds, kpsPreds = null;
scores = list.get(idx);
bboxPreds = list.get(idx + fmc).mul(stride);
kpsPreds = list.get(idx + fmc * 2).mul(stride);
int height = inputHeight / stride;
int width = inputWidth / stride;
String key = height + "_" + width + "_" + stride;
// anchor centers cache
NDArray anchorCenters;
if (centerCache.containsKey(key)) {
anchorCenters = centerCache.get(key);
} else {
NDArray yv = manager.arange((float) height).reshape(height, 1).repeat(1, width);
NDArray xv = manager.arange((float) width).reshape(1, width).repeat(0, height);
// stack x, y 到最后一维
anchorCenters = xv.stack(yv, -1); // shape [height, width, 2]
// 乘 stride
anchorCenters = anchorCenters.mul(stride);
// 拉平成 [-1, 2],等价于 NumPy 的 reshape((-1, 2))
long total = anchorCenters.getShape().get(0) * anchorCenters.getShape().get(1);
// anchorCenters 现在是 [height*width, 2]
anchorCenters = anchorCenters.reshape(total, 2);
// 在第一维重复 numAnchors 次,直接拉平成最终形状
anchorCenters = anchorCenters.repeat(0, numAnchors); // shape [N*numAnchors, 2]
if (centerCache.size() < 100) {
centerCache.put(key, anchorCenters);
}
}
NDArray pos_mask = scores.gte(confThresh); // scores >= thresh
NDArray pos_inds = pos_mask.nonzero(); // shape: [N, 2]
pos_inds = pos_inds.get(":, 0"); // 取第一列的索引
// System.out.println(Arrays.toString(pos_inds.toLongArray()));
// 计算 bbox
NDArray bboxes = distance2bbox(anchorCenters, bboxPreds); // [num_anchors, 4]
// 取出符合阈值的
NDArray pos_scores = scores.get(pos_inds);
NDArray pos_bboxes = bboxes.get(pos_inds);
scores_list.add(pos_scores);
bboxes_list.add(pos_bboxes);
NDArray kpss = distance2kps(anchorCenters, kpsPreds); // [num_anchors, num_kps*2]
kpss = kpss.reshape(kpss.getShape().get(0), -1, 2); // reshape (N, -1, 2)
NDArray pos_kpss = kpss.get(pos_inds);
kpss_list.add(pos_kpss);
}
// 1. 合并 scores
NDArray scores = NDArrays.concat(new NDList(scores_list), 0);
NDArray scoresRavel = scores.reshape(-1);
// 2. 得到排序索引
long[] orderLong = scoresRavel.argSort().flip(0).get(":" + topK).toLongArray();
// 3. 合并 bboxes
NDArray bboxes = NDArrays.concat(new NDList(bboxes_list), 0).div(detScale);
NDArray kpss = NDArrays.concat(new NDList(kpss_list), 0).div(detScale);
// 4. 拼接 [x1,y1,x2,y2,score]
NDArray preDet = bboxes.concat(scores.reshape(-1,1), 1);
// 5. 按 order 排序
preDet = preDet.get(manager.create(orderLong));
// 6. NMS
int[] keep = NMSUtils.nms(preDet.get(":,0:4"), preDet.get(":,4"), (float)nmsThresh);
NDArray det = preDet.get(manager.create(keep));
// System.out.println(Arrays.toString(det.toFloatArray()));
if (kpss != null) {
kpss = kpss.get(manager.create(orderLong));
kpss = kpss.get(manager.create(keep));
}
List<String> retNames = new ArrayList<>();
List<Double> retProbs = new ArrayList<>();
List<BoundingBox> retBB = new ArrayList<>();
long numDet = det.getShape().get(0); // N
long numCols = det.getShape().get(1); // 应该是 5: x1,y1,x2,y2,score
float[] flat = det.toFloatArray(); // 一维
for (int i = 0; i < numDet; i++) {
int base = (int) (i * numCols);
float x1 = flat[base] / sourceWidth;
float y1 = flat[base + 1] / sourceHeight;
float x2 = flat[base + 2] / sourceWidth;
float y2 = flat[base + 3] / sourceHeight;
float score = flat[base + 4];
retNames.add("face"); // 类别
retProbs.add((double) score);
float width = x2 - x1;
float height = y2 - y1;
Landmark rect = new Landmark(x1, y1, width, height, decodeKps(kpss.get(i)));
retBB.add(rect);
}
return new DetectedObjects(retNames, retProbs, retBB);
}
public NDArray distance2bbox(NDArray points, NDArray distance) {
// points: [N, 2], distance: [N, 4]
NDArray x1 = points.get(":, 0").sub(distance.get(":, 0")); // x - left
NDArray y1 = points.get(":, 1").sub(distance.get(":, 1")); // y - top
NDArray x2 = points.get(":, 0").add(distance.get(":, 2")); // x + right
NDArray y2 = points.get(":, 1").add(distance.get(":, 3")); // y + bottom
// stack([x1, y1, x2, y2], axis=-1)
NDList list = new NDList(x1.expandDims(1), y1.expandDims(1), x2.expandDims(1), y2.expandDims(1));
return NDArrays.concat(list, 1); // axis=1 表示最后一维
}
public NDArray distance2kps(NDArray points, NDArray distance) {
// points: [N, 2], distance: [N, 2*num_kps]
int numKps = (int) distance.getShape().get(1) / 2;
List<NDArray> preds = new ArrayList<>();
for (int i = 0; i < numKps * 2; i += 2) {
NDArray px = points.get(":, " + (i % 2)).add(distance.get(":, " + i));
NDArray py = points.get(":, " + ((i % 2) + 1)).add(distance.get(":, " + (i + 1)));
preds.add(px);
preds.add(py);
}
// stack(preds, axis=-1)
NDList stackList = new NDList();
for (NDArray arr : preds) {
stackList.add(arr.expandDims(1));
}
return NDArrays.concat(stackList, 1); // shape [N, num_kps*2]
}
public List<Point> decodeKps(NDArray kpss) {
// 转成一维 float 数组
float[] flat = kpss.toFloatArray();
// reshape 成二维 [5][2]
int numPoints = (int) kpss.getShape().get(0); // 5
int dim = (int) kpss.getShape().get(1); // 2
float[][] kpsArray = new float[numPoints][dim];
for (int i = 0; i < numPoints; i++) {
for (int j = 0; j < dim; j++) {
kpsArray[i][j] = flat[i * dim + j];
}
}
// 转成 Point 数组
List<Point> points = new ArrayList<>();
for (int i = 0; i < numPoints; i++) {
points.add(new Point(Math.round(kpsArray[i][0]), Math.round(kpsArray[i][1])));
}
return points;
}
}

View File

@@ -8,6 +8,7 @@ import ai.djl.ndarray.NDList;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.*;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import java.util.*;
@@ -79,8 +80,59 @@ public class YoloV5FaceTranslator implements Translator<Image, DetectedObjects>
return builder;
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) throws Exception {
NDArray array = input.toNDArray(ctx.getNDManager(), flag);
// NDList list = pipeline.transform(new NDList(array));
// Shape shape = list.get(0).getShape();
// int processedWidth;
// int processedHeight;
// long[] dim = shape.getShape();
// if (NDImageUtils.isCHW(shape)) {
// processedWidth = (int) dim[dim.length - 1];
// processedHeight = (int) dim[dim.length - 2];
// } else {
// processedWidth = (int) dim[dim.length - 2];
// processedHeight = (int) dim[dim.length - 3];
// }
LetterBoxUtils.ResizeResult letterBoxResult = LetterBoxUtils.letterbox(ctx.getNDManager(), array, width, height, 114f, LetterBoxUtils.PaddingPosition.CENTER);
array = letterBoxResult.image;
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
ctx.setAttachment("processedWidth", width);
ctx.setAttachment("processedHeight", height);
ctx.setAttachment("scale", letterBoxResult.r);
// 转为 float32 且归一化到 0~1
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
// HWC -> CHW
array = array.transpose(2, 0, 1); // CHW
// return new NDList(array.expandDims(0));
return new NDList(array);
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) throws Exception {
int imageWidth = (Integer) ctx.getAttachment("width");
int imageHeight = (Integer) ctx.getAttachment("height");
float scale = (Float) ctx.getAttachment("scale");
switch (yoloOutputLayerType) {
case DETECT:
return processFromDetectOutput();
case AUTO:
if (list.get(0).getShape().dimension() > 2) {
return processFromDetectOutput();
} else {
return processFromBoxOutput(imageWidth, imageHeight, list, scale);
}
case BOX:
default:
return processFromBoxOutput(imageWidth, imageHeight, list, scale);
}
}
/** {@inheritDoc} */
protected DetectedObjects processFromBoxOutput(int imageWidth, int imageHeight, NDList list) {
protected DetectedObjects processFromBoxOutput(int imageWidth, int imageHeight, NDList list, float scale) {
float[] flattened = list.get(0).toFloatArray();
int sizeClasses = classes.size();
int stride = 15 + sizeClasses;
@@ -119,7 +171,7 @@ public class YoloV5FaceTranslator implements Translator<Image, DetectedObjects>
classIds.add(maxIndex);
}
}
return nms(imageWidth, imageHeight, boxes, classIds, scores);
return nms(imageWidth, imageHeight, boxes, classIds, scores, scale);
}
private DetectedObjects processFromDetectOutput() {
@@ -127,53 +179,16 @@ public class YoloV5FaceTranslator implements Translator<Image, DetectedObjects>
"detect layer output is not supported yet, check correct YoloV5 export format");
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) throws Exception {
int imageWidth = (Integer) ctx.getAttachment("width");
int imageHeight = (Integer) ctx.getAttachment("height");
switch (yoloOutputLayerType) {
case DETECT:
return processFromDetectOutput();
case AUTO:
if (list.get(0).getShape().dimension() > 2) {
return processFromDetectOutput();
} else {
return processFromBoxOutput(imageWidth, imageHeight, list);
}
case BOX:
default:
return processFromBoxOutput(imageWidth, imageHeight, list);
}
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) throws Exception {
NDArray array = input.toNDArray(ctx.getNDManager(), flag);
NDList list = pipeline.transform(new NDList(array));
Shape shape = list.get(0).getShape();
int processedWidth;
int processedHeight;
long[] dim = shape.getShape();
if (NDImageUtils.isCHW(shape)) {
processedWidth = (int) dim[dim.length - 1];
processedHeight = (int) dim[dim.length - 2];
} else {
processedWidth = (int) dim[dim.length - 2];
processedHeight = (int) dim[dim.length - 3];
}
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
ctx.setAttachment("processedWidth", processedWidth);
ctx.setAttachment("processedHeight", processedHeight);
return list;
}
protected DetectedObjects nms(
int imageWidth,
int imageHeight,
List<Landmark> boxes,
List<Integer> classIds,
List<Float> scores) {
List<Float> scores, float scale) {
List<String> retClasses = new ArrayList<>();
List<Double> retProbs = new ArrayList<>();
List<BoundingBox> retBB = new ArrayList<>();
@@ -196,34 +211,41 @@ public class YoloV5FaceTranslator implements Translator<Image, DetectedObjects>
for (int index : nms) {
int pos = map.get(index);
int id = classIds.get(pos);
retClasses.add(classes.get(id));
int percent = (int) Math.round(scores.get(pos).doubleValue() * 100);
String className = "face " + percent + "%"; // classes.get(classId)
retClasses.add(className);
retProbs.add(scores.get(pos).doubleValue());
// Rectangle rect = boxes.get(pos);
Landmark rect = boxes.get(pos);
List<Point> keypoints = new ArrayList<>();
if (removePadding) {
int padW = (width - imageWidth) / 2;
int padH = (height - imageHeight) / 2;
rect.getPath().forEach(point -> {
keypoints.add(new Point(point.getX() - padW, point.getY() - padH));
});
rect =
new Landmark(
(rect.getX() - padW) / imageWidth,
(rect.getY() - padH) / imageHeight,
rect.getWidth() / imageWidth,
rect.getHeight() / imageHeight,keypoints);
} else if (applyRatio) {
rect.getPath().forEach(point -> {
keypoints.add(new Point(point.getX() / width, point.getY() / height));
});
rect =
new Landmark(
rect.getX() / width,
rect.getY() / height,
rect.getWidth() / width,
rect.getHeight() / height,keypoints);
}
//恢复原图坐标(除回比例,减掉 padding
rect = LetterBoxUtils.restoreBox(rect, scale, imageWidth, imageHeight, width, height, false);
// if (removePadding) {
// int padW = (width - imageWidth) / 2;
// int padH = (height - imageHeight) / 2;
// rect.getPath().forEach(point -> {
// keypoints.add(new Point(point.getX() - padW, point.getY() - padH));
// });
// rect =
// new Landmark(
// (rect.getX() - padW) / imageWidth,
// (rect.getY() - padH) / imageHeight,
// rect.getWidth() / imageWidth,
// rect.getHeight() / imageHeight,keypoints);
// } else if (applyRatio) {
// rect.getPath().forEach(point -> {
// keypoints.add(new Point(point.getX() / width, point.getY() / height));
// });
// rect =
// new Landmark(
// rect.getX() / width,
// rect.getY() / height,
// rect.getWidth() / width,
// rect.getHeight() / height,keypoints);
// }
retBB.add(rect);
}
}

View File

@@ -1,476 +0,0 @@
package cn.smartjavaai.face.translator;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.*;
import ai.djl.modality.cv.transform.*;
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.*;
import java.util.*;
/**
* YoloV8Translator
* @author dwj
*/
public class YoloV8FaceTranslator implements Translator<Image, DetectedObjects> {
private int maxBoxes;
private YoloOutputType yoloOutputLayerType;
private float nmsThreshold;
protected float threshold;
// private BaseImageTranslator.SynsetLoader synsetLoader;
protected List<String> classes;
protected boolean applyRatio;
protected boolean removePadding;
protected Pipeline pipeline;
private Image.Flag flag;
private Batchifier batchifier;
protected int width;
protected int height;
/**
* Constructs an ImageTranslator with the provided builder.
*
* @param builder the data to build with
*/
protected YoloV8FaceTranslator(Builder builder) {
this.yoloOutputLayerType = builder.outputType;
this.nmsThreshold = builder.nmsThreshold;
maxBoxes = builder.maxBox;
this.threshold = builder.threshold;
// this.synsetLoader = builder.synsetLoader;
this.applyRatio = builder.applyRatio;
this.removePadding = builder.removePadding;
this.flag = builder.flag;
this.pipeline = builder.pipeline;
this.batchifier = builder.batchifier;
this.width = builder.width;
this.height = builder.height;
classes = Arrays.asList("face");
}
/**
* Creates a builder to build a {@code YoloV8Translator} with specified arguments.
*
* @return a new builder
*/
public static Builder builder() {
return new Builder();
}
/**
* Creates a builder to build a {@code YoloV8Translator} with specified arguments.
*
* @param arguments arguments to specify builder options
* @return a new builder
*/
public static Builder builder(Map<String, ?> arguments) {
Builder builder = new Builder();
builder.configPreProcess(arguments);
builder.configPostProcess(arguments);
return builder;
}
/** {@inheritDoc} */
protected DetectedObjects processFromBoxOutput(int imageWidth, int imageHeight, NDList list) {
NDArray rawResult = list.get(0);
NDArray reshapedResult = rawResult.transpose();
Shape shape = reshapedResult.getShape();
float[] buf = reshapedResult.toFloatArray();
int numberRows = Math.toIntExact(shape.get(0));
int nClasses = Math.toIntExact(shape.get(1));
int padding = nClasses - classes.size();
System.out.println(Arrays.toString(reshapedResult.get(0).toFloatArray()));
// if (padding != 0 && padding != 4) {
// throw new IllegalStateException(
// "Expected classes: " + (nClasses - 4) + ", got " + classes.size());
// }
ArrayList<Landmark> boxes = new ArrayList<>();
ArrayList<Float> scores = new ArrayList<>();
ArrayList<Integer> classIds = new ArrayList<>();
// reverse order search in heap; searches through #maxBoxes for optimization when set
for (int i = numberRows - 1; i > numberRows - maxBoxes; --i) {
int index = i * nClasses;
float maxClassProb = buf[index + 4];
// int maxIndex = -1;
// for (int c = 4; c < nClasses; c++) {
// float classProb = buf[index + c];
// if (classProb > maxClassProb) {
// maxClassProb = classProb;
// maxIndex = c;
// }
// }
// maxIndex -= padding;
if (maxClassProb > threshold) {
float xPos = buf[index]; // center x
float yPos = buf[index + 1]; // center y
float w = buf[index + 2];
float h = buf[index + 3];
Rectangle rect =
new Rectangle(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h);
// boxes.add(rect);
scores.add(maxClassProb);
classIds.add(0);
List<Point> keypoints = new ArrayList<>();
keypoints.add(new Point(buf[index + 5], buf[index + 6]));
keypoints.add(new Point(buf[index + 8], buf[index + 9]));
keypoints.add(new Point(buf[index + 11], buf[index + 12]));
keypoints.add(new Point(buf[index + 14], buf[index + 15]));
keypoints.add(new Point(buf[index + 17], buf[index + 18]));
Landmark kps = new Landmark(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h, keypoints);
boxes.add(kps);
}
}
return nms(imageWidth, imageHeight, boxes, classIds, scores);
}
private DetectedObjects processFromDetectOutput() {
throw new UnsupportedOperationException(
"detect layer output is not supported yet, check correct YoloV5 export format");
}
@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) throws Exception {
int imageWidth = (Integer) ctx.getAttachment("width");
int imageHeight = (Integer) ctx.getAttachment("height");
switch (yoloOutputLayerType) {
case DETECT:
return processFromDetectOutput();
case AUTO:
if (list.get(0).getShape().dimension() > 2) {
return processFromDetectOutput();
} else {
return processFromBoxOutput(imageWidth, imageHeight, list);
}
case BOX:
default:
return processFromBoxOutput(imageWidth, imageHeight, list);
}
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) throws Exception {
NDArray array = input.toNDArray(ctx.getNDManager(), flag);
NDList list = pipeline.transform(new NDList(array));
Shape shape = list.get(0).getShape();
int processedWidth;
int processedHeight;
long[] dim = shape.getShape();
if (NDImageUtils.isCHW(shape)) {
processedWidth = (int) dim[dim.length - 1];
processedHeight = (int) dim[dim.length - 2];
} else {
processedWidth = (int) dim[dim.length - 2];
processedHeight = (int) dim[dim.length - 3];
}
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
ctx.setAttachment("processedWidth", processedWidth);
ctx.setAttachment("processedHeight", processedHeight);
return list;
}
protected DetectedObjects nms(
int imageWidth,
int imageHeight,
List<Landmark> boxes,
List<Integer> classIds,
List<Float> scores) {
List<String> retClasses = new ArrayList<>();
List<Double> retProbs = new ArrayList<>();
List<BoundingBox> retBB = new ArrayList<>();
for (int classId = 0; classId < classes.size(); classId++) {
List<Rectangle> r = new ArrayList<>();
List<Double> s = new ArrayList<>();
List<Integer> map = new ArrayList<>();
for (int j = 0; j < classIds.size(); ++j) {
if (classIds.get(j) == classId) {
r.add(boxes.get(j));
s.add(scores.get(j).doubleValue());
map.add(j);
}
}
if (r.isEmpty()) {
continue;
}
List<Integer> nms = Rectangle.nms(r, s, nmsThreshold);
for (int index : nms) {
int pos = map.get(index);
int id = classIds.get(pos);
retClasses.add(classes.get(id));
retProbs.add(scores.get(pos).doubleValue());
// Rectangle rect = boxes.get(pos);
Landmark rect = boxes.get(pos);
List<Point> keypoints = new ArrayList<>();
if (removePadding) {
int padW = (width - imageWidth) / 2;
int padH = (height - imageHeight) / 2;
rect.getPath().forEach(point -> {
keypoints.add(new Point(point.getX() - padW, point.getY() - padH));
});
rect =
new Landmark(
(rect.getX() - padW) / imageWidth,
(rect.getY() - padH) / imageHeight,
rect.getWidth() / imageWidth,
rect.getHeight() / imageHeight,keypoints);
} else if (applyRatio) {
rect.getPath().forEach(point -> {
keypoints.add(new Point(point.getX() / width, point.getY() / height));
});
rect =
new Landmark(
rect.getX() / width,
rect.getY() / height,
rect.getWidth() / width,
rect.getHeight() / height,keypoints);
}
retBB.add(rect);
}
}
return new DetectedObjects(retClasses, retProbs, retBB);
}
public static class Builder {
private int maxBox = 8400;
YoloOutputType outputType;
float nmsThreshold;
protected float threshold = 0.2F;
protected boolean applyRatio;
protected boolean removePadding;
protected int width = 224;
protected int height = 224;
protected Image.Flag flag;
protected Pipeline pipeline;
protected Batchifier batchifier;
public Builder() {
this.outputType = YoloOutputType.AUTO;
this.nmsThreshold = 0.4F;
}
public Builder optOutputType(YoloOutputType outputType) {
this.outputType = outputType;
return this;
}
public Builder optNmsThreshold(float nmsThreshold) {
this.nmsThreshold = nmsThreshold;
return this;
}
/**
* Builds the translator.
*
* @return the new translator
*/
public YoloV8FaceTranslator build() {
if (pipeline == null) {
addTransform(
array -> array.transpose(2, 0, 1).toType(DataType.FLOAT32, false).div(255));
}
// validate();
return new YoloV8FaceTranslator(this);
}
protected Builder self() {
return this;
}
public Builder addTransform(Transform transform) {
if (this.pipeline == null) {
this.pipeline = new Pipeline();
}
this.pipeline.add(transform);
return this.self();
}
public Builder optApplyRatio(boolean value) {
this.applyRatio = value;
return this.self();
}
public Builder optFlag(Image.Flag flag) {
this.flag = flag;
return this.self();
}
public Builder setPipeline(Pipeline pipeline) {
this.pipeline = pipeline;
return this.self();
}
public Builder setImageSize(int width, int height) {
this.width = width;
this.height = height;
return this.self();
}
public Builder optBatchifier(Batchifier batchifier) {
this.batchifier = batchifier;
return this.self();
}
public Builder optThreshold(float threshold) {
this.threshold = threshold;
return this.self();
}
/** {@inheritDoc} */
protected void configPostProcess(Map<String, ?> arguments) {
if (ArgumentsUtil.booleanValue(arguments, "optApplyRatio") || ArgumentsUtil.booleanValue(arguments, "applyRatio")) {
this.optApplyRatio(true);
}
this.threshold = ArgumentsUtil.floatValue(arguments, "threshold", 0.2F);
String centerFit = ArgumentsUtil.stringValue(arguments, "centerFit", "false");
this.removePadding = "true".equals(centerFit);
String type = ArgumentsUtil.stringValue(arguments, "outputType", "AUTO");
this.outputType = YoloOutputType.valueOf(type.toUpperCase(Locale.ENGLISH));
this.nmsThreshold = ArgumentsUtil.floatValue(arguments, "nmsThreshold", 0.4F);
maxBox = ArgumentsUtil.intValue(arguments, "maxBox", 8400);
}
protected void configPreProcess(Map<String, ?> arguments) {
if (this.pipeline == null) {
this.pipeline = new Pipeline();
}
this.width = ArgumentsUtil.intValue(arguments, "width", 224);
this.height = ArgumentsUtil.intValue(arguments, "height", 224);
if (arguments.containsKey("flag")) {
this.flag = Image.Flag.valueOf(arguments.get("flag").toString());
}
String pad = ArgumentsUtil.stringValue(arguments, "pad", "false");
if ("true".equals(pad)) {
this.addTransform(new Pad(0.0));
} else if (!"false".equals(pad)) {
double padding = Double.parseDouble(pad);
this.addTransform(new Pad(padding));
}
String resize = ArgumentsUtil.stringValue(arguments, "resize", "false");
int w;
int shortEdge;
if ("true".equals(resize)) {
this.addTransform(new Resize(this.width, this.height));
} else if (!"false".equals(resize)) {
String[] tokens = resize.split("\\s*,\\s*");
w = (int)Double.parseDouble(tokens[0]);
if (tokens.length > 1) {
shortEdge = (int)Double.parseDouble(tokens[1]);
} else {
shortEdge = w;
}
Image.Interpolation interpolation;
if (tokens.length > 2) {
interpolation = Image.Interpolation.valueOf(tokens[2]);
} else {
interpolation = Image.Interpolation.BILINEAR;
}
this.addTransform(new Resize(w, shortEdge, interpolation));
}
String resizeShort = ArgumentsUtil.stringValue(arguments, "resizeShort", "false");
if ("true".equals(resizeShort)) {
w = Math.max(this.width, this.height);
this.addTransform(new ResizeShort(w));
} else if (!"false".equals(resizeShort)) {
String[] tokens = resizeShort.split("\\s*,\\s*");
shortEdge = (int)Double.parseDouble(tokens[0]);
int longEdge;
if (tokens.length > 1) {
longEdge = (int)Double.parseDouble(tokens[1]);
} else {
longEdge = -1;
}
Image.Interpolation interpolation;
if (tokens.length > 2) {
interpolation = Image.Interpolation.valueOf(tokens[2]);
} else {
interpolation = Image.Interpolation.BILINEAR;
}
this.addTransform(new ResizeShort(shortEdge, longEdge, interpolation));
}
if (ArgumentsUtil.booleanValue(arguments, "centerCrop", false)) {
this.addTransform(new CenterCrop(this.width, this.height));
}
if (ArgumentsUtil.booleanValue(arguments, "centerFit")) {
this.addTransform(new CenterFit(this.width, this.height));
}
if (ArgumentsUtil.booleanValue(arguments, "toTensor", true)) {
this.addTransform(new ToTensor());
}
String normalize = ArgumentsUtil.stringValue(arguments, "normalize", "false");
if ("true".equals(normalize)) {
float[] MEAN = new float[]{0.485F, 0.456F, 0.406F};
float[] STD = new float[]{0.229F, 0.224F, 0.225F};
this.addTransform(new Normalize(MEAN, STD));
} else if (!"false".equals(normalize)) {
String[] tokens = normalize.split("\\s*,\\s*");
if (tokens.length != 6) {
throw new IllegalArgumentException("Invalid normalize value: " + normalize);
}
float[] mean = new float[]{Float.parseFloat(tokens[0]), Float.parseFloat(tokens[1]), Float.parseFloat(tokens[2])};
float[] std = new float[]{Float.parseFloat(tokens[3]), Float.parseFloat(tokens[4]), Float.parseFloat(tokens[5])};
this.addTransform(new Normalize(mean, std));
}
String range = (String)arguments.get("range");
if ("0,1".equals(range)) {
this.addTransform((a) -> {
return a.div(255.0F);
});
} else if ("-1,1".equals(range)) {
this.addTransform((a) -> {
return a.div(128.0F).sub(1);
});
}
if (arguments.containsKey("batchifier")) {
this.batchifier = Batchifier.fromString((String)arguments.get("batchifier"));
}
}
}
public static enum YoloOutputType {
BOX,
DETECT,
AUTO;
private YoloOutputType() {
}
}
}

View File

@@ -182,7 +182,12 @@ public class FaceUtils {
DetectionRectangle rectangle = detectionInfo.getDetectionRectangle();
graphics.setColor(Color.RED);// 边框颜色
graphics.drawRect(rectangle.getX(), rectangle.getY(), rectangle.getWidth(), rectangle.getHeight());
drawText(graphics, "face", rectangle.getX(), rectangle.getY(), stroke, 4);
String className = "face";
if (detectionInfo.getScore() > 0){
int percent = (int) Math.round(detectionInfo.getScore() * 100);
className = "face " + percent + "%";
}
drawText(graphics, className , rectangle.getX(), rectangle.getY(), stroke, 4);
//绘制人脸关键点
if(detectionInfo.getFaceInfo() != null && detectionInfo.getFaceInfo().getKeyPoints() != null &&
!detectionInfo.getFaceInfo().getKeyPoints().isEmpty()){

View File

@@ -36,8 +36,10 @@ public class Test {
*/
public static FaceDetModel getFaceDetModel(){
FaceDetConfig config = new FaceDetConfig();
config.setModelEnum(FaceDetModelEnum.RETINA_FACE);//人脸检测模型
config.setConfidenceThreshold(FaceDetectConstant.DEFAULT_CONFIDENCE_THRESHOLD);//只返回相似度大于该值的人脸
// config.setModelEnum(FaceDetModelEnum.YOLOV8_FACE);//人脸检测模型
config.setModelPath("/Users/wenjie/Documents/develop/model/yolo-face/yolov8s-face-lindevs.onnx");
// config.setModelPath("/Users/wenjie/Documents/develop/face_model");
config.setConfidenceThreshold(0.2);//只返回相似度大于该值的人脸
config.setNmsThresh(FaceDetectConstant.NMS_THRESHOLD);//用于去除重复的人脸框,当两个框的重叠度超过该值时,只保留一个
return FaceDetModelFactory.getInstance().getModel(config);
}
@@ -81,6 +83,11 @@ public class Test {
// }
// }
// }
FaceDetModel faceDetModel = getFaceDetModel();
R<Void> result = faceDetModel.detectAndDraw("/Users/wenjie/Downloads/facetest/00974.png", "/Users/wenjie/Downloads/xx333.png");
log.info("result:{}", result.isSuccess() + " msg:" + result.getMessage());
}