修复人脸识别算法facenet-pytorch实现方式

smartjavaai-face/src/main/java/cn/smartjavaai/face/algo/FeatureExtractionAlgo.java

@@ -1,5 +1,6 @@
 package cn.smartjavaai.face.algo;
 
+import ai.djl.Device;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
@@ -14,6 +15,7 @@ import cn.smartjavaai.face.AbstractFaceAlgorithm;
 import cn.smartjavaai.face.FaceDetectedResult;
 import cn.smartjavaai.face.FaceDetectionTranslator;
 import cn.smartjavaai.face.ModelConfig;
+import cn.smartjavaai.face.translator.FaceFeatureTranslator;
 import org.apache.commons.lang3.StringUtils;
 
 import java.io.InputStream;
@@ -50,32 +52,34 @@ public class FeatureExtractionAlgo extends AbstractFaceAlgorithm {
      * 加载人脸特征提取模型
      * @param config
      * @throws Exception
-     *//*
+     */
     @Override
     public void loadFaceFeatureModel(ModelConfig config) throws Exception {
         String normalize = mean.stream().map(Object::toString).collect(Collectors.joining(","));
-        faceFeatureCriteria = Criteria.builder()
+        faceFeatureCriteria =
+                Criteria.builder()
                         .setTypes(Image.class, float[].class)
+                        .optModelName("face_feature") // specify model file prefix
                         .optModelUrls(StringUtils.isNotBlank(config.getModelPath()) ? null :
                                 "https://resources.djl.ai/test-models/pytorch/face_feature.zip")
                         .optModelPath(StringUtils.isNotBlank(config.getModelPath()) ? Paths.get(config.getModelPath()) : null)
-                        .optModelName("face_feature") // specify model file prefix
+                        .optTranslator(new FaceFeatureTranslator())
                         .optArgument("normalize", normalize)
-                        .optTranslatorFactory(new ImageFeatureExtractorFactory())
-                        .optProgress(new ProgressBar())
                         .optEngine("PyTorch") // Use PyTorch engine
+                        .optProgress(new ProgressBar())
                         .build();
+
         model = faceFeatureCriteria.loadModel();
         predictor = model.newPredictor();
     }
 
 
-    *//**
+    /**
      * 特征提取
      * @param imagePath 图片路径
      * @return
      * @throws Exception
-     *//*
+     */
     @Override
     public float[] featureExtraction(String imagePath) throws Exception {
         Path imageFile = Paths.get(imagePath);
@@ -84,12 +88,12 @@ public class FeatureExtractionAlgo extends AbstractFaceAlgorithm {
         return predictor.predict(img);
     }
 
-    *//**
+    /**
      * 特征提取
      * @param inputStream 输入流
      * @return
      * @throws Exception
-     *//*
+     */
     @Override
     public float[] featureExtraction(InputStream inputStream) throws Exception {
         Image img = ImageFactory.getInstance().fromInputStream(inputStream);
@@ -97,13 +101,13 @@ public class FeatureExtractionAlgo extends AbstractFaceAlgorithm {
         return predictor.predict(img);
     }
 
-    *//**
+    /**
      * 计算相似度
      * @param feature1 图1特征
      * @param feature2 图2特征
      * @return
      * @throws Exception
-     *//*
+     */
     @Override
     public float calculSimilar(float[] feature1, float[] feature2) throws Exception {
         float ret = 0.0f;
@@ -118,13 +122,13 @@ public class FeatureExtractionAlgo extends AbstractFaceAlgorithm {
         return (float) ((ret / Math.sqrt(mod1) / Math.sqrt(mod2) + 1) / 2.0f);
     }
 
-    *//**
+    /**
      * 特征比较
      * @param imagePath1 图1路径
      * @param imagePath2 图2路径
      * @return
      * @throws Exception
-     *//*
+     */
     @Override
     public float featureComparison(String imagePath1, String imagePath2) throws Exception {
         float[] feature1 = featureExtraction(imagePath1);
@@ -132,19 +136,19 @@ public class FeatureExtractionAlgo extends AbstractFaceAlgorithm {
         return calculSimilar(feature1, feature2);
     }
 
-    *//**
+    /**
      * 特征比较
      * @param inputStream1 图1输入流
      * @param inputStream2 图2输入流
      * @return
      * @throws Exception
-     *//*
+     */
     @Override
     public float featureComparison(InputStream inputStream1, InputStream inputStream2) throws Exception {
         float[] feature1 = featureExtraction(inputStream1);
         float[] feature2 = featureExtraction(inputStream2);
         return calculSimilar(feature1, feature2);
-    }*/
+    }
 
     /*@Override
     public float[] recognize(FaceRegion region) {

smartjavaai-face/src/main/java/cn/smartjavaai/face/translator/FaceFeatureTranslator.java

@@ -0,0 +1,52 @@
+package cn.smartjavaai.face.translator;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.transform.Normalize;
+import ai.djl.modality.cv.transform.Resize;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Pipeline;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+/**
+ * @author 邓文杰
+ * @date 2025/3/31
+ */
+public final class FaceFeatureTranslator implements Translator<Image, float[]> {
+
+    public FaceFeatureTranslator() {
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Image input) {
+        NDArray array = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
+        Pipeline pipeline = new Pipeline();
+        pipeline
+                .add(new Resize(180))
+                .add(new ToTensor())
+                .add(new Normalize(
+                        new float[]{127.5f / 255.0f, 127.5f / 255.0f, 127.5f / 255.0f},
+                        new float[]{128.0f / 255.0f, 128.0f / 255.0f, 128.0f / 255.0f}));
+
+        return pipeline.transform(new NDList(array));
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public float[] processOutput(TranslatorContext ctx, NDList list) {
+        return list.singletonOrThrow().toFloatArray();
+    }
+
+    @Override
+    public Batchifier getBatchifier() {
+        return Batchifier.STACK;
+    }
+}