临时提交

2026-09-13 13:18:58 +00:00 · 2025-08-29 10:30:35 +08:00
parent 8bf620a330
commit 86ea7eb03e
364 changed files with 8572 additions and 540 deletions
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/config/InstanceSegModelConfig.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/config/InstanceSegModelConfig.java
@@ -0,0 +1,55 @@
+package cn.smartjavaai.instanceseg.config;
+
+import cn.smartjavaai.common.config.ModelConfig;
+import cn.smartjavaai.common.enums.DeviceEnum;
+import cn.smartjavaai.instanceseg.enums.InstanceSegModelEnum;
+import cn.smartjavaai.objectdetection.constant.DetectorConstant;
+import lombok.Data;
+
+import java.util.List;
+
+/**
+ * 实例分割模型参数配置
+ *
+ * @author dwj
+ */
+@Data
+public class InstanceSegModelConfig extends ModelConfig {
+
+    /**
+     * 模型
+     */
+    private InstanceSegModelEnum modelEnum;
+
+
+    /**
+     * 模型路径
+     */
+    private String modelPath;
+
+
+    /**
+     * 允许的分类列表
+     */
+    private List<String> allowedClasses;
+
+    /**
+     * 置信度阈值
+     */
+    private float threshold = 0.3f;
+
+
+    public InstanceSegModelConfig() {
+    }
+
+    public InstanceSegModelConfig(InstanceSegModelEnum modelEnum, DeviceEnum device) {
+        this.modelEnum = modelEnum;
+        setDevice(device);
+    }
+
+    public InstanceSegModelConfig(InstanceSegModelEnum modelEnum) {
+        this.modelEnum = modelEnum;
+    }
+
+
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/criteria/InstanceSegCriteriaFactory.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/criteria/InstanceSegCriteriaFactory.java
@@ -0,0 +1,47 @@
+package cn.smartjavaai.instanceseg.criteria;
+
+import ai.djl.Device;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.training.util.ProgressBar;
+import cn.smartjavaai.common.enums.DeviceEnum;
+import cn.smartjavaai.instanceseg.config.InstanceSegModelConfig;
+import cn.smartjavaai.instanceseg.translator.YoloSegmentationTranslatorFactory2;
+import org.apache.commons.lang3.StringUtils;
+
+import java.nio.file.Paths;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * 实例分割Criteria工厂
+ * @author dwj
+ */
+public class InstanceSegCriteriaFactory {
+
+
+    public static Criteria<Image, DetectedObjects> createCriteria(InstanceSegModelConfig config) {
+        Device device = null;
+        if(!Objects.isNull(config.getDevice())){
+            device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
+        }
+        Criteria<Image, DetectedObjects> criteria = null;
+        ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
+        params.putAll(config.getCustomParams());
+//        YoloV5Translator.Builder builder = new YoloV5Translator.Builder()
+//                .optSynsetArtifactName("synset.txt").setPipeline()
+        criteria =
+                Criteria.builder()
+                        .setTypes(Image.class, DetectedObjects.class)
+                        .optModelUrls(StringUtils.isNotBlank(config.getModelPath()) ? null :
+                                config.getModelEnum().getModelUri())
+                        .optModelPath(StringUtils.isNotBlank(config.getModelPath()) ? Paths.get(config.getModelPath()) : null)
+                        .optDevice(device)
+                        .optEngine("PyTorch")
+                        .optTranslatorFactory(new YoloSegmentationTranslatorFactory2())
+                        .optProgress(new ProgressBar())
+                        .build();
+        return criteria;
+    }
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/entity/DetectParams.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/entity/DetectParams.java
@@ -0,0 +1,17 @@
+package cn.smartjavaai.instanceseg.entity;
+
+import lombok.Data;
+
+/**
+ * 检测参数
+ * @author dwj
+ */
+@Data
+public class DetectParams {
+
+    /**
+     * 置信度阈值
+     */
+    private float threshold = 0.3f;
+
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/enums/InstanceSegModelEnum.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/enums/InstanceSegModelEnum.java
@@ -0,0 +1,43 @@
+package cn.smartjavaai.instanceseg.enums;
+
+/**
+ * 实例分割模型枚举
+ * @author dwj
+ */
+public enum InstanceSegModelEnum {
+
+    SEG_YOLO11N_PYTORCH("djl://ai.djl.pytorch/yolo11n-seg"),
+
+    SEG_YOLOV8N_PYTORCH("djl://ai.djl.pytorch/yolo11n-seg"),
+
+    SEG_YOLO11N_ONNX("djl://ai.djl.onnxruntime/yolo11n-seg"),
+
+    SEG_YOLOV8N_ONNX("djl://ai.djl.onnxruntime/yolov8n-seg"),
+
+    SEG_MASK_RCNN("djl://ai.djl.mxnet/mask_rcnn");
+
+
+    /**
+     * 根据名称获取枚举 (忽略大小写和下划线变体)
+     */
+    public static InstanceSegModelEnum fromName(String name) {
+        String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
+        for (InstanceSegModelEnum model : values()) {
+            if (model.name().replaceAll("_", "").equals(formatted)) {
+                return model;
+            }
+        }
+        throw new IllegalArgumentException("未知模型名称: " + name);
+    }
+
+    private final String modelUri;
+
+    InstanceSegModelEnum(String modelUri) {
+        this.modelUri = modelUri;
+    }
+
+    public String getModelUri() {
+        return modelUri;
+    }
+
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/exception/InstanceSegException.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/exception/InstanceSegException.java
@@ -0,0 +1,29 @@
+package cn.smartjavaai.instanceseg.exception;
+
+/**
+ * 实例分割异常
+ * @author dwj
+ */
+public class InstanceSegException extends RuntimeException{
+
+    public InstanceSegException() {
+        super();
+    }
+
+    public InstanceSegException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+
+    public InstanceSegException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public InstanceSegException(String message) {
+        super(message);
+    }
+
+    public InstanceSegException(Throwable cause) {
+        super(cause);
+    }
+
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/model/CommonInstanceSegModel.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/model/CommonInstanceSegModel.java
@@ -0,0 +1,153 @@
+package cn.smartjavaai.instanceseg.model;
+
+import ai.djl.MalformedModelException;
+import ai.djl.engine.Engine;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ModelNotFoundException;
+import ai.djl.repository.zoo.ZooModel;
+import cn.smartjavaai.common.cv.SmartImageFactory;
+import cn.smartjavaai.common.entity.DetectionResponse;
+import cn.smartjavaai.common.entity.R;
+import cn.smartjavaai.common.pool.PredictorFactory;
+import cn.smartjavaai.instanceseg.config.InstanceSegModelConfig;
+import cn.smartjavaai.instanceseg.criteria.InstanceSegCriteriaFactory;
+import cn.smartjavaai.instanceseg.exception.InstanceSegException;
+import cn.smartjavaai.objectdetection.exception.DetectionException;
+import cn.smartjavaai.vision.utils.DetectedObjectsFilter;
+import cn.smartjavaai.vision.utils.DetectorUtils;
+import cn.smartjavaai.vision.utils.CategoryMaskFilter;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.pool2.impl.GenericObjectPool;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Objects;
+
+/**
+ * 实例分割模型
+ * @author dwj
+ */
+@Slf4j
+public class CommonInstanceSegModel implements InstanceSegModel {
+
+
+    private InstanceSegModelConfig config;
+
+    private ZooModel<Image, DetectedObjects> model;
+
+    private GenericObjectPool<Predictor<Image, DetectedObjects>> predictorPool;
+
+    @Override
+    public void loadModel(InstanceSegModelConfig config) {
+        if(Objects.isNull(config.getModelEnum())){
+            throw new DetectionException("未配置模型枚举");
+        }
+        Criteria<Image, DetectedObjects> criteria = InstanceSegCriteriaFactory.createCriteria(config);
+        this.config = config;
+        try {
+            model = criteria.loadModel();
+            // 创建池子：每个线程独享 Predictor
+            this.predictorPool = new GenericObjectPool<>(new PredictorFactory<>(model));
+            int predictorPoolSize = config.getPredictorPoolSize();
+            if(config.getPredictorPoolSize() <= 0){
+                predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
+            }
+            predictorPool.setMaxTotal(predictorPoolSize);
+            log.debug("当前设备: " + model.getNDManager().getDevice());
+            log.debug("当前引擎: " + Engine.getInstance().getEngineName());
+            log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
+        } catch (IOException | ModelNotFoundException | MalformedModelException e) {
+            throw new DetectionException("模型加载失败", e);
+        }
+    }
+
+    @Override
+    public R<DetectionResponse> detect(Image image) {
+        DetectedObjects detectedObjects = detectCore(image);
+        DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, image);
+        return R.ok(detectionResponse);
+    }
+
+    /**
+     * 模型核心推理方法
+     * @param image
+     * @return
+     */
+    @Override
+    public DetectedObjects detectCore(Image image) {
+        Predictor<Image, DetectedObjects> predictor = null;
+        try {
+            predictor = predictorPool.borrowObject();
+            DetectedObjects detectedObjects = predictor.predict(image);
+            //过滤
+            if(Objects.nonNull(detectedObjects) && detectedObjects.getNumberOfObjects() > 0){
+                DetectedObjectsFilter detectedObjectsFilter = new DetectedObjectsFilter(config.getAllowedClasses(), config.getThreshold());
+                detectedObjects = detectedObjectsFilter.filter(detectedObjects);
+            }
+            return detectedObjects;
+        } catch (Exception e) {
+            throw new DetectionException("实例分割错误", e);
+        }finally {
+            if (predictor != null) {
+                try {
+                    predictorPool.returnObject(predictor); //归还
+                    log.debug("释放资源");
+                } catch (Exception e) {
+                    log.warn("归还Predictor失败", e);
+                    try {
+                        predictor.close(); // 归还失败才销毁
+                    } catch (Exception ex) {
+                        log.error("关闭Predictor失败", ex);
+                    }
+                }
+            }
+        }
+    }
+
+    @Override
+    public R<DetectionResponse> detectAndDraw(Image image) {
+        DetectedObjects detectedObjects = detectCore(image);
+        image.drawBoundingBoxes(detectedObjects);
+        DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, image);
+        detectionResponse.setDrawnImage(image);
+        return R.ok(detectionResponse);
+    }
+
+    @Override
+    public R<DetectionResponse> detectAndDraw(String imagePath, String outputPath) {
+        try {
+            Image img = SmartImageFactory.getInstance().fromFile(Paths.get(imagePath));
+            DetectedObjects detectedObjects = detectCore(img);
+            img.drawBoundingBoxes(detectedObjects);
+            img.save(Files.newOutputStream(Paths.get(outputPath)), "png");
+            DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, img);
+            return R.ok(detectionResponse);
+        } catch (IOException e) {
+            throw new InstanceSegException(e);
+        }
+    }
+
+    @Override
+    public void close() throws Exception {
+        try {
+            if (predictorPool != null) {
+                predictorPool.close();
+            }
+        } catch (Exception e) {
+            log.warn("关闭 predictorPool 失败", e);
+        }
+        try {
+            if (model != null) {
+                model.close();
+            }
+        } catch (Exception e) {
+            log.warn("关闭 model 失败", e);
+        }
+    }
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/model/InstanceSegModel.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/model/InstanceSegModel.java
@@ -0,0 +1,47 @@
+package cn.smartjavaai.instanceseg.model;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.CategoryMask;
+import ai.djl.modality.cv.output.DetectedObjects;
+import cn.smartjavaai.common.entity.DetectionResponse;
+import cn.smartjavaai.common.entity.R;
+import cn.smartjavaai.instanceseg.config.InstanceSegModelConfig;
+
+import java.awt.image.BufferedImage;
+
+/**
+ * 实例分割模型
+ * @author dwj
+ */
+public interface InstanceSegModel extends AutoCloseable{
+
+
+    /**
+     * 加载模型
+     * @param config
+     */
+    void loadModel(InstanceSegModelConfig config);
+
+    /**
+     * 实例分割
+     * @param image
+     * @return
+     */
+    default R<DetectionResponse> detect(Image image){
+        throw new UnsupportedOperationException("默认不支持该功能");
+    }
+
+    default DetectedObjects detectCore(Image image){
+        throw new UnsupportedOperationException("默认不支持该功能");
+    }
+
+    default R<DetectionResponse> detectAndDraw(Image image){
+        throw new UnsupportedOperationException("默认不支持该功能");
+    }
+
+    default R<DetectionResponse> detectAndDraw(String imagePath, String outputPath){
+        throw new UnsupportedOperationException("默认不支持该功能");
+    }
+
+
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/translator/YoloSegmentationTranslator2.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/translator/YoloSegmentationTranslator2.java
@@ -0,0 +1,182 @@
+/*
+ * Copyright 2024 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.instanceseg.translator;
+
+import ai.djl.modality.cv.output.BoundingBox;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.Mask;
+import ai.djl.modality.cv.output.Rectangle;
+import ai.djl.modality.cv.transform.Resize;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.modality.cv.translator.YoloV5Translator;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Pipeline;
+import ai.djl.translate.TranslatorContext;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/** A translator for Yolov8 instance segmentation models. */
+public class YoloSegmentationTranslator2 extends YoloV5Translator {
+
+    private static final int[] AXIS_0 = {0};
+    private static final int[] AXIS_1 = {1};
+
+    private float threshold;
+    private float nmsThreshold;
+
+    /**
+     * Creates the instance segmentation translator from the given builder.
+     *
+     * @param builder the builder for the translator
+     */
+    public YoloSegmentationTranslator2(Builder builder) {
+        super(builder);
+        this.threshold = 0.25f;
+        this.nmsThreshold = 0.4F;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
+        NDArray pred = list.get(0);
+        NDArray protos = list.get(1);
+        int maskIndex = classes.size() + 4;
+        NDArray candidates = pred.get("4:" + maskIndex).max(AXIS_0).gt(threshold);
+        pred = pred.transpose();
+        NDArray sub = pred.get("..., :4");
+        sub = xywh2xyxy(sub);
+        pred = sub.concat(pred.get("..., 4:"), -1);
+        pred = pred.get(candidates);
+
+        NDList split = pred.split(new long[] {4, maskIndex}, 1);
+        NDArray box = split.get(0);
+
+        int numBox = Math.toIntExact(box.getShape().get(0));
+
+        float[] buf = box.toFloatArray();
+        float[] confidences = split.get(1).max(AXIS_1).toFloatArray();
+        long[] ids = split.get(1).argMax(1).toLongArray();
+
+        List<Rectangle> boxes = new ArrayList<>(numBox);
+        List<Double> scores = new ArrayList<>(numBox);
+        for (int i = 0; i < numBox; ++i) {
+            float xPos = buf[i * 4];
+            float yPos = buf[i * 4 + 1];
+            float w = buf[i * 4 + 2] - xPos;
+            float h = buf[i * 4 + 3] - yPos;
+            Rectangle rect = new Rectangle(xPos, yPos, w, h);
+            boxes.add(rect);
+            scores.add((double) confidences[i]);
+        }
+        List<Integer> nms = Rectangle.nms(boxes, scores, nmsThreshold);
+        long[] idx = nms.stream().mapToLong(Integer::longValue).toArray();
+        NDArray selected = box.getManager().create(idx);
+        NDArray masks = split.get(2).get(selected);
+
+        int maskW = Math.toIntExact(protos.getShape().get(2));
+        int maskH = Math.toIntExact(protos.getShape().get(1));
+
+        protos = protos.reshape(32, (long) maskH * maskW);
+        masks =
+                masks.matMul(protos)
+                        .reshape(nms.size(), maskH, maskW)
+                        .gt(0f)
+                        .toType(DataType.FLOAT32, true);
+
+        float[] maskArray = masks.toFloatArray();
+        box = box.get(selected);
+        buf = box.toFloatArray();
+
+        List<String> retClasses = new ArrayList<>();
+        List<Double> retProbs = new ArrayList<>();
+        List<BoundingBox> retBB = new ArrayList<>();
+        for (int i = 0; i < idx.length; ++i) {
+            float x = buf[i * 4] / width;
+            float y = buf[i * 4 + 1] / height;
+            float w = buf[i * 4 + 2] / width - x;
+            float h = buf[i * 4 + 3] / width - y;
+            int id = nms.get(i);
+            retClasses.add(classes.get((int) ids[id]));
+            retProbs.add((double) confidences[id]);
+
+            float[][] maskFloat = new float[maskH][maskW];
+            int pos = i * maskH * maskW;
+            for (int j = 0; j < maskH; j++) {
+                System.arraycopy(maskArray, pos + j * maskW, maskFloat[j], 0, maskW);
+            }
+            Mask bb = new Mask(x, y, w, h, maskFloat, true);
+            retBB.add(bb);
+        }
+        return new DetectedObjects(retClasses, retProbs, retBB);
+    }
+
+    private NDArray xywh2xyxy(NDArray array) {
+        NDArray xy = array.get("..., :2", new Object[0]);
+        NDArray wh = array.get("..., 2:", new Object[0]).div(2);
+        return xy.sub(wh).concat(xy.add(wh), -1);
+    }
+
+    /**
+     * Creates a builder to build a {@code YoloSegmentationTranslator}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Creates a builder to build a {@code YoloSegmentationTranslator} 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.optSynsetArtifactName("synset.txt");
+        builder.setImageSize(640, 640);
+        return builder;
+    }
+
+    /** The builder for instance segmentation translator. */
+    public static class Builder extends YoloV5Translator.Builder {
+
+        Builder() {}
+
+        /** {@inheritDoc} */
+        @Override
+        protected Builder self() {
+            return this;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public YoloSegmentationTranslator2 build() {
+            pipeline = new Pipeline();
+            pipeline.add(new Resize(640, 640));
+            pipeline.add(new ToTensor());
+//            validate();
+            return new YoloSegmentationTranslator2(this);
+        }
+    }
+
+
+
+
+
+}
--- a/vision/src/main/java/cn/smartjavaai/instanceseg/translator/YoloSegmentationTranslatorFactory2.java
+++ b/vision/src/main/java/cn/smartjavaai/instanceseg/translator/YoloSegmentationTranslatorFactory2.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2024 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.instanceseg.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.translator.ObjectDetectionTranslatorFactory;
+import ai.djl.translate.Translator;
+
+import java.io.Serializable;
+import java.util.Map;
+
+/** A translatorFactory that creates a {@link ai.djl.modality.cv.translator.YoloSegmentationTranslator} instance. */
+public class YoloSegmentationTranslatorFactory2 extends ObjectDetectionTranslatorFactory
+        implements Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    /** {@inheritDoc} */
+    @Override
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        Translator<Image, DetectedObjects> translator = YoloSegmentationTranslator2.builder(arguments).build();
+        return translator;
+    }
+}