ZF/SmartJavaAI
1
0
Fork 0
mirror of https://github.com/geekwenjie/SmartJavaAI.git synced 2026-09-10 03:28:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

- 【人脸检测】新增6个模型(MTCNN、YOLOV5、RetinaFace小尺寸版)，大幅提升性能

Browse Source 
- 【人脸识别】新增Seetaface6轻量模型
- 【目标检测】支持视频流目标检测（rtsp、视频文件等）
- 【目标检测】支持tensorflow2目标检测模型
- 【目标检测】新增行人检测模型(yolo-person)
- 【通用视觉】新增4个动作识别模型
- 【通用视觉】新增语义分割模型
- 【通用视觉】新增5个实例分割模型(含yolov8-seg、yolov11-seg)
- 【通用视觉】新增yolo-obb11旋转框检测(含yolov11-obb)
- 【通用视觉】新增5个姿态估计模型(含yolov8-pose、yolov11-pose)
This commit is contained in:
dengwenjie
2025-09-07 17:19:19 +08:00
parent 2b044fda29
commit a8e7ce6c4e
102 changed files with 4473 additions and 1368 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
vision/pom.xml
View File

@@ -36,32 +36,9 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.bytedeco</groupId>
<artifactId>javacpp</artifactId>
<version>1.5.10</version>
<classifier>macosx-arm64</classifier>
</dependency>
<dependency>
<groupId>org.bytedeco</groupId>
<artifactId>ffmpeg</artifactId>
<version>6.1.1-1.5.10</version>
<classifier>macosx-arm64</classifier>
</dependency>
<dependency>
<groupId>org.bytedeco</groupId>
<artifactId>openblas</artifactId>
<version>0.3.26-1.5.10</version>
<classifier>macosx-arm64</classifier>
</dependency>
<dependency>
<groupId>org.bytedeco</groupId>
<artifactId>opencv</artifactId>
<version>4.9.0-1.5.10</version>
<classifier>macosx-arm64</classifier>
</dependency>
</dependencies>

65
vision/src/main/java/cn/smartjavaai/action/criteria/ActionRecCriteriaFactory.java
View File

@@ -3,8 +3,10 @@ package cn.smartjavaai.action.criteria;
import ai.djl.Device;
import ai.djl.modality.Classifications;
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 ai.djl.translate.Translator;
import cn.smartjavaai.action.config.ActionRecModelConfig;
import cn.smartjavaai.action.enums.ActionRecModelEnum;
import cn.smartjavaai.action.model.CommonActionTranslator;
@@ -23,45 +25,48 @@ import java.util.concurrent.ConcurrentHashMap;
public class ActionRecCriteriaFactory {
/**
* 创建动作识别Criteria
* @param config
* @return
*/
public static Criteria<Image, Classifications> createCriteria(ActionRecModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, Classifications> criteria = null;
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
params.putAll(config.getCustomParams());
if(config.getModelEnum() == ActionRecModelEnum.VIT_BASE_PATCH16_224){
criteria =
Criteria.builder()
.setTypes(Image.class, Classifications.class)
.optModelUrls(StringUtils.isNotBlank(config.getModelPath()) ? null :
config.getModelEnum().getModelUri())
.optModelPath(StringUtils.isNotBlank(config.getModelPath()) ? Paths.get(config.getModelPath()) : null)
.optEngine("PyTorch")
.optDevice(device)
.optProgress(new ProgressBar())
.build();
}else {
Translator<Image, Classifications> translator = getTranslator(config);
if(StringUtils.isBlank(config.getModelEnum().getModelUrl())){
//检查模型路径
if (StringUtils.isBlank(config.getModelPath())){
throw new ActionException("请指定模型路径");
}
int width = 224;
int height = 224;
if(config.getModelEnum() == ActionRecModelEnum.INCEPTIONV3_KINETICS400){
width = 299;
height = 299;
}
criteria =
Criteria.builder()
.setTypes(Image.class, Classifications.class)
.optTranslator(new CommonActionTranslator(width, height))
.optEngine("OnnxRuntime")
.optModelPath(Paths.get(config.getModelPath()))
.optDevice(device)
.optProgress(new ProgressBar())
.build();
}
Criteria<Image, Classifications> criteria =
Criteria.builder()
.setTypes(Image.class, Classifications.class)
.optModelUrls(StringUtils.isNotBlank(config.getModelPath()) ? null : config.getModelEnum().getModelUrl())
.optModelPath(StringUtils.isNotBlank(config.getModelPath()) ? Paths.get(config.getModelPath()) : null)
.optTranslator(translator)
.optDevice(device)
.optProgress(new ProgressBar())
.optEngine(config.getModelEnum().getEngine())
.build();
return criteria;
}
/**
* 获取动作识别Translator
* @param config
* @return
*/
public static Translator<Image, Classifications> getTranslator(ActionRecModelConfig config) {
Translator<Image, Classifications> translator = null;
if(config.getModelEnum() == ActionRecModelEnum.INCEPTIONV1_KINETICS400_ONNX
|| config.getModelEnum() == ActionRecModelEnum.INCEPTIONV3_KINETICS400_ONNX
|| config.getModelEnum() == ActionRecModelEnum.INCEPTIONV3_KINETICS400_ONNX){
translator =new CommonActionTranslator(config.getModelEnum().getInputWidth(), config.getModelEnum().getInputHeight());
}
return translator;
}
}

51
vision/src/main/java/cn/smartjavaai/action/enums/ActionRecModelEnum.java
View File

@@ -6,21 +6,33 @@ package cn.smartjavaai.action.enums;
*/
public enum ActionRecModelEnum {
VIT_BASE_PATCH16_224("djl://ai.djl.pytorch/Human-Action-Recognition-VIT-Base-patch16-224"),
VIT_BASE_PATCH16_224_DJL("PyTorch",0,0,"djl://ai.djl.pytorch/Human-Action-Recognition-VIT-Base-patch16-224"),
INCEPTIONV3_KINETICS400(""),
INCEPTIONV3_KINETICS400_ONNX("OnnxRuntime",299,299,""),
INCEPTIONV1_KINETICS400(""),
INCEPTIONV1_KINETICS400_ONNX("OnnxRuntime",224,224,""),
RESNET18_V1B_KINETICS400(""),
RESNET_V1B_KINETICS400_ONNX("OnnxRuntime",224,224,"");
RESNET34_V1B_KINETICS400(""),
/**
* 模型输入尺寸：宽
*/
private final int inputWidth;
RESNET50_V1B_KINETICS400(""),
/**
* 模型输入尺寸：高
*/
private final int inputHeight;
RESNET101_V1B_KINETICS400(""),
/**
* 模型地址
*/
private final String modelUrl;
RESNET152_V1B_KINETICS400("");
/**
* 模型引擎
*/
private final String engine;
/**
* 根据名称获取枚举 (忽略大小写和下划线变体)
@@ -35,14 +47,27 @@ public enum ActionRecModelEnum {
throw new IllegalArgumentException("未知模型名称: " + name);
}
private final String modelUri;
ActionRecModelEnum(String modelUri) {
this.modelUri = modelUri;
ActionRecModelEnum(String engine, int inputWidth, int inputHeight, String modelUrl) {
this.inputWidth = inputWidth;
this.inputHeight = inputHeight;
this.modelUrl = modelUrl;
this.engine = engine;
}
public String getModelUri() {
return modelUri;
public int getInputWidth() {
return inputWidth;
}
public int getInputHeight() {
return inputHeight;
}
public String getModelUrl() {
return modelUrl;
}
public String getEngine() {
return engine;
}
}

9
vision/src/main/java/cn/smartjavaai/action/model/ActionRecModel.java
View File

@@ -24,15 +24,6 @@ public interface ActionRecModel extends AutoCloseable{
*/
void loadModel(ActionRecModelConfig config);
/**
* 动作检测
* @param base64Image
* @return
*/
default R<Classifications> detectBase64(String base64Image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 动作检测
* @param image

111
vision/src/main/java/cn/smartjavaai/action/model/ActionRecModelFactory.java Normal file
View File

@@ -0,0 +1,111 @@
package cn.smartjavaai.action.model;
import cn.smartjavaai.action.config.ActionRecModelConfig;
import cn.smartjavaai.action.enums.ActionRecModelEnum;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import cn.smartjavaai.objectdetection.model.person.CommonPersonDetModel;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 动作识别 模型工厂
* @author dwj
*/
@Slf4j
public class ActionRecModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile ActionRecModelFactory instance;
private static final ConcurrentHashMap<ActionRecModelEnum, ActionRecModel> modelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<ActionRecModelEnum, Class<? extends ActionRecModel>> registry =
new ConcurrentHashMap<>();
// 私有构造函数，防止外部创建实例
private ActionRecModelFactory() {}
// 双重检查锁定的单例方法
public static ActionRecModelFactory getInstance() {
if (instance == null) {
synchronized (ActionRecModelFactory.class) {
if (instance == null) {
instance = new ActionRecModelFactory();
}
}
}
return instance;
}
/**
* 获取模型（通过配置）
* @param config
* @return
*/
public ActionRecModel getModel(ActionRecModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型");
}
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createFaceDetModel(config);
});
}
/**
* 使用ModelConfig创建模型
* @param config
* @return
*/
private ActionRecModel createFaceDetModel(ActionRecModelConfig config) {
Class<?> clazz = registry.get(config.getModelEnum());
if(clazz == null){
throw new DetectionException("Unsupported model");
}
ActionRecModel model = null;
try {
model = (ActionRecModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new DetectionException(e);
}
model.loadModel(config);
return model;
}
/**
* 注册模型
* @param modelEnum
* @param clazz
*/
private static void registerAlgorithm(ActionRecModelEnum modelEnum, Class<? extends ActionRecModel> clazz) {
registry.put(modelEnum, clazz);
}
/**
* 移除缓存的模型
* @param modelEnum
*/
public static void removeFromCache(ActionRecModelEnum modelEnum) {
modelMap.remove(modelEnum);
}
// 初始化默认算法
static {
registerAlgorithm(ActionRecModelEnum.INCEPTIONV1_KINETICS400_ONNX, CommonActionRecModel.class);
registerAlgorithm(ActionRecModelEnum.INCEPTIONV3_KINETICS400_ONNX, CommonActionRecModel.class);
registerAlgorithm(ActionRecModelEnum.RESNET_V1B_KINETICS400_ONNX, CommonActionRecModel.class);
registerAlgorithm(ActionRecModelEnum.VIT_BASE_PATCH16_224_DJL, CommonActionRecModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

16
vision/src/main/java/cn/smartjavaai/action/model/CommonActionRecModel.java
View File

@@ -68,26 +68,12 @@ public class CommonActionRecModel implements ActionRecModel{
}
}
@Override
public R<Classifications> detectBase64(String base64Image) {
if(StringUtils.isBlank(base64Image)){
return R.fail(R.Status.INVALID_IMAGE);
}
try {
byte[] imageData = Base64ImageUtils.base64ToImage(base64Image);
Image image = ImageFactory.getInstance().fromInputStream(new ByteArrayInputStream(imageData));
return detect(image);
} catch (IOException e) {
throw new DetectionException("读取图片异常", e);
}
}
@Override
public R<Classifications> detect(Image image) {
Classifications classifications = detectCore(image);
// 过滤
if(config.getThreshold() > 0 && CollectionUtils.isNotEmpty(config.getAllowedClasses())
&& Objects.nonNull(classifications) && !classifications.items().isEmpty()){
if(Objects.nonNull(classifications) && !classifications.items().isEmpty()){
classifications = new ClassificationFilter(config.getAllowedClasses(), config.getThreshold()).filter(classifications);
}
return R.ok(classifications);

1
vision/src/main/java/cn/smartjavaai/action/model/CommonActionTranslator.java
View File

@@ -116,7 +116,6 @@ public class CommonActionTranslator implements Translator<Image, Classifications
float[] std = {0.229f * 255, 0.224f * 255, 0.225f * 255};
// 增加 batch 维度，变成 (1, H, W, C)
array = array.expandDims(0);
System.out.println(Arrays.toString(array.getShape().getShape()));
return new NDList(array);
}

2
vision/src/main/java/cn/smartjavaai/instanceseg/config/InstanceSegModelConfig.java
View File

@@ -36,7 +36,7 @@ public class InstanceSegModelConfig extends ModelConfig {
/**
* 置信度阈值
*/
private float threshold = 0.3f;
private float threshold = 0.25f;
public InstanceSegModelConfig() {

49
vision/src/main/java/cn/smartjavaai/instanceseg/criteria/InstanceSegCriteriaFactory.java
View File

@@ -3,10 +3,12 @@ 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.modality.cv.translator.InstanceSegmentationTranslatorFactory;
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.enums.InstanceSegModelEnum;
import cn.smartjavaai.instanceseg.translator.YoloSegmentationTranslatorFactory2;
import org.apache.commons.lang3.StringUtils;
@@ -27,21 +29,42 @@ public class InstanceSegCriteriaFactory {
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());
// 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();
if(config.getModelEnum() == InstanceSegModelEnum.SEG_MASK_RCNN){
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(config.getModelEnum().getEngine())
.optArgument("normalize","true")
.optArgument("synsetFileName","classes.txt")
.optTranslatorFactory(new InstanceSegmentationTranslatorFactory())
.optProgress(new ProgressBar())
.build();
}else{
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)
.optArgument("width", config.getModelEnum().getInputWidth())
.optArgument("height", config.getModelEnum().getInputHeight())
.optArgument("resize", "true")
.optArgument("threshold", config.getThreshold())
.optEngine(config.getModelEnum().getEngine())
.optTranslatorFactory(new YoloSegmentationTranslatorFactory2())
.optProgress(new ProgressBar())
.build();
}
return criteria;
}
}

41
vision/src/main/java/cn/smartjavaai/instanceseg/enums/InstanceSegModelEnum.java
View File

@@ -6,15 +6,15 @@ package cn.smartjavaai.instanceseg.enums;
*/
public enum InstanceSegModelEnum {
SEG_YOLO11N_PYTORCH("djl://ai.djl.pytorch/yolo11n-seg"),
SEG_YOLO11N_PYTORCH("PyTorch", 640, 640, "djl://ai.djl.pytorch/yolo11n-seg"),
SEG_YOLOV8N_PYTORCH("djl://ai.djl.pytorch/yolo11n-seg"),
SEG_YOLOV8N_PYTORCH("PyTorch", 640, 640, "djl://ai.djl.pytorch/yolov8n-seg"),
SEG_YOLO11N_ONNX("djl://ai.djl.onnxruntime/yolo11n-seg"),
SEG_YOLO11N_ONNX("OnnxRuntime", 640, 640, "djl://ai.djl.onnxruntime/yolo11n-seg"),
SEG_YOLOV8N_ONNX("djl://ai.djl.onnxruntime/yolov8n-seg"),
SEG_YOLOV8N_ONNX("OnnxRuntime", 640, 640, "djl://ai.djl.onnxruntime/yolov8n-seg"),
SEG_MASK_RCNN("djl://ai.djl.mxnet/mask_rcnn");
SEG_MASK_RCNN("MXNet", 0,0, "djl://ai.djl.mxnet/mask_rcnn");
/**
@@ -30,14 +30,43 @@ public enum InstanceSegModelEnum {
throw new IllegalArgumentException("未知模型名称: " + name);
}
/**
* 模型输入尺寸：宽
*/
private final int inputWidth;
/**
* 模型输入尺寸：高
*/
private final int inputHeight;
private final String modelUri;
InstanceSegModelEnum(String modelUri) {
/**
* 模型引擎
*/
private final String engine;
InstanceSegModelEnum(String engine, int inputWidth, int inputHeight, String modelUri) {
this.inputWidth = inputWidth;
this.inputHeight = inputHeight;
this.modelUri = modelUri;
this.engine = engine;
}
public String getModelUri() {
return modelUri;
}
public String getEngine() {
return engine;
}
public int getInputWidth() {
return inputWidth;
}
public int getInputHeight() {
return inputHeight;
}
}

6
vision/src/main/java/cn/smartjavaai/instanceseg/model/CommonInstanceSegModel.java
View File

@@ -113,6 +113,9 @@ public class CommonInstanceSegModel implements InstanceSegModel {
@Override
public R<DetectionResponse> detectAndDraw(Image image) {
DetectedObjects detectedObjects = detectCore(image);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
image.drawBoundingBoxes(detectedObjects);
DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, image);
detectionResponse.setDrawnImage(image);
@@ -124,6 +127,9 @@ public class CommonInstanceSegModel implements InstanceSegModel {
try {
Image img = SmartImageFactory.getInstance().fromFile(Paths.get(imagePath));
DetectedObjects detectedObjects = detectCore(img);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
img.drawBoundingBoxes(detectedObjects);
img.save(Files.newOutputStream(Paths.get(outputPath)), "png");
DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, img);

111
vision/src/main/java/cn/smartjavaai/instanceseg/model/InstanceSegModelFactory.java Normal file
View File

@@ -0,0 +1,111 @@
package cn.smartjavaai.instanceseg.model;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.instanceseg.config.InstanceSegModelConfig;
import cn.smartjavaai.instanceseg.enums.InstanceSegModelEnum;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 实例分割 模型工厂
* @author dwj
*/
@Slf4j
public class InstanceSegModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile InstanceSegModelFactory instance;
private static final ConcurrentHashMap<InstanceSegModelEnum, InstanceSegModel> modelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<InstanceSegModelEnum, Class<? extends InstanceSegModel>> registry =
new ConcurrentHashMap<>();
// 私有构造函数，防止外部创建实例
private InstanceSegModelFactory() {}
// 双重检查锁定的单例方法
public static InstanceSegModelFactory getInstance() {
if (instance == null) {
synchronized (InstanceSegModelFactory.class) {
if (instance == null) {
instance = new InstanceSegModelFactory();
}
}
}
return instance;
}
/**
* 获取模型（通过配置）
* @param config
* @return
*/
public InstanceSegModel getModel(InstanceSegModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型");
}
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createFaceDetModel(config);
});
}
/**
* 使用ModelConfig创建模型
* @param config
* @return
*/
private InstanceSegModel createFaceDetModel(InstanceSegModelConfig config) {
Class<?> clazz = registry.get(config.getModelEnum());
if(clazz == null){
throw new DetectionException("Unsupported model");
}
InstanceSegModel model = null;
try {
model = (InstanceSegModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new DetectionException(e);
}
model.loadModel(config);
return model;
}
/**
* 注册模型
* @param modelEnum
* @param clazz
*/
private static void registerAlgorithm(InstanceSegModelEnum modelEnum, Class<? extends InstanceSegModel> clazz) {
registry.put(modelEnum, clazz);
}
/**
* 移除缓存的模型
* @param modelEnum
*/
public static void removeFromCache(InstanceSegModelEnum modelEnum) {
modelMap.remove(modelEnum);
}
// 初始化默认算法
static {
registerAlgorithm(InstanceSegModelEnum.SEG_YOLOV8N_ONNX, CommonInstanceSegModel.class);
registerAlgorithm(InstanceSegModelEnum.SEG_YOLOV8N_PYTORCH, CommonInstanceSegModel.class);
registerAlgorithm(InstanceSegModelEnum.SEG_YOLO11N_PYTORCH, CommonInstanceSegModel.class);
registerAlgorithm(InstanceSegModelEnum.SEG_YOLO11N_ONNX, CommonInstanceSegModel.class);
registerAlgorithm(InstanceSegModelEnum.SEG_MASK_RCNN, CommonInstanceSegModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

59
vision/src/main/java/cn/smartjavaai/obb/config/ObbDetModelConfig.java Normal file
View File

@@ -0,0 +1,59 @@
package cn.smartjavaai.obb.config;
import cn.smartjavaai.common.config.ModelConfig;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.obb.enums.ObbDetModelEnum;
import lombok.Data;
import java.util.List;
/**
* 旋转框模型参数配置
*
* @author dwj
*/
@Data
public class ObbDetModelConfig extends ModelConfig {
/**
* 模型
*/
private ObbDetModelEnum modelEnum;
/**
* 模型路径
*/
private String modelPath;
/**
* 允许的分类列表
*/
private List<String> allowedClasses;
/**
* 置信度阈值
*/
private float threshold = 0.25f;
/**
* 按置信度分数排序后，最多保留的检测框数量
*/
private int topK;
public ObbDetModelConfig() {
}
public ObbDetModelConfig(ObbDetModelEnum modelEnum, DeviceEnum device) {
this.modelEnum = modelEnum;
setDevice(device);
}
public ObbDetModelConfig(ObbDetModelEnum modelEnum) {
this.modelEnum = modelEnum;
}
}

70
vision/src/main/java/cn/smartjavaai/obb/criteria/ObbDetCriteriaFactory.java Normal file
View File

@@ -0,0 +1,70 @@
package cn.smartjavaai.obb.criteria;
import ai.djl.Device;
import ai.djl.modality.cv.Image;
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.obb.config.ObbDetModelConfig;
import cn.smartjavaai.obb.entity.ObbResult;
import cn.smartjavaai.obb.enums.ObbDetModelEnum;
import cn.smartjavaai.obb.exception.ObbDetException;
import cn.smartjavaai.obb.translator.YoloV11OddTranslator;
import org.apache.commons.lang3.StringUtils;
import java.nio.file.Paths;
import java.util.Objects;
/**
* 旋转框Criteria工厂
* @author dwj
*/
public class ObbDetCriteriaFactory {
/**
* 创建旋转框检测Criteria
* @param config
* @return
*/
public static Criteria<Image, ObbResult> createCriteria(ObbDetModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Translator<Image, ObbResult> translator = getTranslator(config);
//检查模型路径
if (StringUtils.isBlank(config.getModelPath())){
throw new ObbDetException("请指定模型路径");
}
Criteria<Image, ObbResult> criteria =
Criteria.builder()
.setTypes(Image.class, ObbResult.class)
.optModelPath(Paths.get(config.getModelPath()))
.optTranslator(translator)
.optDevice(device)
.optProgress(new ProgressBar())
.optEngine(config.getModelEnum().getEngine())
.build();
return criteria;
}
/**
* 获取旋转框检测Translator
* @param config
* @return
*/
public static Translator<Image, ObbResult> getTranslator(ObbDetModelConfig config) {
Translator<Image, ObbResult> translator = null;
if (config.getModelEnum() == ObbDetModelEnum.YOLOV11){
translator = YoloV11OddTranslator.builder()
.setImageSize(config.getModelEnum().getInputWidth(), config.getModelEnum().getInputHeight())
.optThreshold(config.getThreshold() > 0 ? config.getThreshold() : 0.25f)
.optNmsThreshold(0.45f)
.optSynsetArtifactName("synset.txt").build();
}
return translator;
}
}

20
vision/src/main/java/cn/smartjavaai/obb/entity/ObbResult.java Normal file
View File

@@ -0,0 +1,20 @@
package cn.smartjavaai.obb.entity;
import lombok.Data;
import java.util.List;
/**
* odd检测结果
* @author dwj
*/
@Data
public class ObbResult {
private List<YoloRotatedBox> rotatedBoxeList;
public ObbResult(List<YoloRotatedBox> rotatedBoxeList) {
this.rotatedBoxeList = rotatedBoxeList;
}
}

98
vision/src/main/java/cn/smartjavaai/obb/entity/YoloRotatedBox.java Normal file
View File

@@ -0,0 +1,98 @@
package cn.smartjavaai.obb.entity;
import cn.smartjavaai.common.entity.Point;
import java.util.ArrayList;
import java.util.List;
/**
* 旋转框
* @author dwj
*/
public class YoloRotatedBox {
public float cx, cy, w, h, angle;
public float score;
public String className;
public YoloRotatedBox(float cx, float cy, float w, float h, float angle, String className, float score) {
this.cx = cx;
this.cy = cy;
this.w = w;
this.h = h;
this.angle = angle;
this.className = className;
this.score = score;
}
public static double probiou(YoloRotatedBox b1, YoloRotatedBox b2, double eps) {
double[] c1 = covarianceMatrix(b1.w, b1.h, b1.angle);
double[] c2 = covarianceMatrix(b2.w, b2.h, b2.angle);
double a1 = c1[0], b1v = c1[1], c1v = c1[2];
double a2 = c2[0], b2v = c2[1], c2v = c2[2];
double x1 = b1.cx, y1 = b1.cy;
double x2 = b2.cx, y2 = b2.cy;
double t1 = ((a1 + a2) * Math.pow(y1 - y2, 2) + (b1v + b2v) * Math.pow(x1 - x2, 2))
/ ((a1 + a2) * (b1v + b2v) - Math.pow(c1v + c2v, 2) + eps);
double t2 = ((c1v + c2v) * (x2 - x1) * (y1 - y2))
/ ((a1 + a2) * (b1v + b2v) - Math.pow(c1v + c2v, 2) + eps);
double t3 = Math.log(((a1 + a2) * (b1v + b2v) - Math.pow(c1v + c2v, 2))
/ (4 * Math.sqrt(a1 * b1v - Math.pow(c1v, 2)) * Math.sqrt(a2 * b2v - Math.pow(c2v, 2)) + eps) + eps);
// 2. probiou 内部 clamp 严格对应 Python
double bd = 0.25 * t1 + 0.5 * t2 + 0.5 * t3;
bd = Math.max(Math.min(bd, 100.0), eps);
double hd = Math.sqrt(1.0 - Math.exp(-bd) + eps);
return 1 - hd;
}
private static double[] covarianceMatrix(double w, double h, double r) {
double a = Math.pow(w, 2) / 12.0;
double b = Math.pow(h, 2) / 12.0;
double cos = Math.cos(r);
double sin = Math.sin(r);
double aVal = a * cos * cos + b * sin * sin;
double bVal = a * sin * sin + b * cos * cos;
double cVal = (a - b) * sin * cos;
return new double[]{aVal, bVal, cVal};
}
/**
* 转换为4点坐标
* @return
*/
public List<Point> toPoints() {
double cos = Math.cos(angle);
double sin = Math.sin(angle);
// vec1 = [w/2*cos, w/2*sin]
double vec1x = w / 2.0 * cos;
double vec1y = w / 2.0 * sin;
// vec2 = [-h/2*sin, h/2*cos]
double vec2x = -h / 2.0 * sin;
double vec2y = h / 2.0 * cos;
List<Point> points = new ArrayList<>(4);
// pt1 = ctr + vec1 + vec2
points.add(new Point((int) Math.round(cx + vec1x + vec2x), (int) Math.round(cy + vec1y + vec2y)));
// pt2 = ctr + vec1 - vec2
points.add(new Point((int) Math.round(cx + vec1x - vec2x), (int) Math.round(cy + vec1y - vec2y)));
// pt3 = ctr - vec1 - vec2
points.add(new Point((int) Math.round(cx - vec1x - vec2x), (int) Math.round(cy - vec1y - vec2y)));
// pt4 = ctr - vec1 + vec2
points.add(new Point((int) Math.round(cx - vec1x + vec2x), (int) Math.round(cy - vec1y + vec2y)));
return points;
}
}

59
vision/src/main/java/cn/smartjavaai/obb/enums/ObbDetModelEnum.java Normal file
View File

@@ -0,0 +1,59 @@
package cn.smartjavaai.obb.enums;
/**
* 旋转框检测模型枚举
* @author dwj
*/
public enum ObbDetModelEnum {
YOLOV11("OnnxRuntime", 1024, 1024);
/**
* 模型引擎
*/
private final String engine;
/**
* 模型输入尺寸：宽
*/
private final int inputWidth;
/**
* 模型输入尺寸：高
*/
private final int inputHeight;
ObbDetModelEnum(String engine, int inputWidth, int inputHeight) {
this.engine = engine;
this.inputWidth = inputWidth;
this.inputHeight = inputHeight;
}
public String getEngine() {
return engine;
}
public int getInputWidth() {
return inputWidth;
}
public int getInputHeight() {
return inputHeight;
}
/**
* 根据名称获取枚举 (忽略大小写和下划线变体)
*/
public static ObbDetModelEnum fromName(String name) {
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
for (ObbDetModelEnum model : values()) {
if (model.name().replaceAll("_", "").equals(formatted)) {
return model;
}
}
throw new IllegalArgumentException("未知模型名称: " + name);
}
}

29
vision/src/main/java/cn/smartjavaai/obb/exception/ObbDetException.java Normal file
View File

@@ -0,0 +1,29 @@
package cn.smartjavaai.obb.exception;
/**
* 旋转框检测异常
* @author dwj
*/
public class ObbDetException extends RuntimeException{
public ObbDetException() {
super();
}
public ObbDetException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
super(message, cause, enableSuppression, writableStackTrace);
}
public ObbDetException(String message, Throwable cause) {
super(message, cause);
}
public ObbDetException(String message) {
super(message);
}
public ObbDetException(Throwable cause) {
super(cause);
}
}

160
vision/src/main/java/cn/smartjavaai/obb/model/CommonObbDetModel.java Normal file
View File

@@ -0,0 +1,160 @@
package cn.smartjavaai.obb.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.obb.config.ObbDetModelConfig;
import cn.smartjavaai.obb.criteria.ObbDetCriteriaFactory;
import cn.smartjavaai.obb.entity.ObbResult;
import cn.smartjavaai.obb.exception.ObbDetException;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import cn.smartjavaai.vision.utils.DetectedObjectsFilter;
import cn.smartjavaai.vision.utils.DetectorUtils;
import cn.smartjavaai.vision.utils.ObbResultFilter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.opencv.core.Mat;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Objects;
/**
* 旋转框模型
* @author dwj
*/
@Slf4j
public class CommonObbDetModel implements ObbDetModel {
private ObbDetModelConfig config;
private ZooModel<Image, ObbResult> model;
private GenericObjectPool<Predictor<Image, ObbResult>> predictorPool;
@Override
public void loadModel(ObbDetModelConfig config) {
if(Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型枚举");
}
Criteria<Image, ObbResult> criteria = ObbDetCriteriaFactory.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) {
ObbResult obbResult = detectCore(image);
DetectionResponse detectionResponse = DetectorUtils.obbToToDetectionResponse(obbResult);
return R.ok(detectionResponse);
}
/**
* 模型核心推理方法
* @param image
* @return
*/
@Override
public ObbResult detectCore(Image image) {
Predictor<Image, ObbResult> predictor = null;
try {
predictor = predictorPool.borrowObject();
ObbResult obbResult = predictor.predict(image);
//过滤
if(Objects.nonNull(obbResult) && CollectionUtils.isNotEmpty(obbResult.getRotatedBoxeList())){
ObbResultFilter obbResultFilter = new ObbResultFilter(config.getAllowedClasses(), config.getTopK());
obbResult = obbResultFilter.filter(obbResult);
}
return obbResult;
} 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) {
ObbResult obbResult = detectCore(image);
if(Objects.isNull(obbResult) || CollectionUtils.isEmpty(obbResult.getRotatedBoxeList())){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
DetectorUtils.drawRectWithText(image, obbResult.getRotatedBoxeList());
DetectionResponse detectionResponse = DetectorUtils.obbToToDetectionResponse(obbResult);
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));
ObbResult obbResult = detectCore(img);
if(Objects.isNull(obbResult) || CollectionUtils.isEmpty(obbResult.getRotatedBoxeList())){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
DetectorUtils.drawRectWithText(img, obbResult.getRotatedBoxeList());
img.save(Files.newOutputStream(Paths.get(outputPath)), "png");
DetectionResponse detectionResponse = DetectorUtils.obbToToDetectionResponse(obbResult);
return R.ok(detectionResponse);
} catch (IOException e) {
throw new ObbDetException(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);
}
}
}

60
vision/src/main/java/cn/smartjavaai/obb/model/ObbDetModel.java Normal file
View File

@@ -0,0 +1,60 @@
package cn.smartjavaai.obb.model;
import ai.djl.modality.cv.Image;
import cn.smartjavaai.common.entity.DetectionResponse;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.obb.config.ObbDetModelConfig;
import cn.smartjavaai.obb.entity.ObbResult;
/**
* 旋转框检测模型
* @author dwj
*/
public interface ObbDetModel extends AutoCloseable{
/**
* 加载模型
* @param config
*/
void loadModel(ObbDetModelConfig config);
/**
* 旋转框
* @param image
* @return
*/
default R<DetectionResponse> detect(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 旋转框检测 核心方法
* @param image
* @return
*/
default ObbResult detectCore(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 旋转框检测并绘制
* @param image
* @return
*/
default R<DetectionResponse> detectAndDraw(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 旋转框检测并绘制
* @param imagePath
* @param outputPath
* @return
*/
default R<DetectionResponse> detectAndDraw(String imagePath, String outputPath){
throw new UnsupportedOperationException("默认不支持该功能");
}
}

107
vision/src/main/java/cn/smartjavaai/obb/model/ObbDetModelFactory.java Normal file
View File

@@ -0,0 +1,107 @@
package cn.smartjavaai.obb.model;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.obb.config.ObbDetModelConfig;
import cn.smartjavaai.obb.enums.ObbDetModelEnum;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 旋转框检测 模型工厂
* @author dwj
*/
@Slf4j
public class ObbDetModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile ObbDetModelFactory instance;
private static final ConcurrentHashMap<ObbDetModelEnum, ObbDetModel> modelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<ObbDetModelEnum, Class<? extends ObbDetModel>> registry =
new ConcurrentHashMap<>();
// 私有构造函数，防止外部创建实例
private ObbDetModelFactory() {}
// 双重检查锁定的单例方法
public static ObbDetModelFactory getInstance() {
if (instance == null) {
synchronized (ObbDetModelFactory.class) {
if (instance == null) {
instance = new ObbDetModelFactory();
}
}
}
return instance;
}
/**
* 获取模型（通过配置）
* @param config
* @return
*/
public ObbDetModel getModel(ObbDetModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型");
}
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createFaceDetModel(config);
});
}
/**
* 使用ModelConfig创建模型
* @param config
* @return
*/
private ObbDetModel createFaceDetModel(ObbDetModelConfig config) {
Class<?> clazz = registry.get(config.getModelEnum());
if(clazz == null){
throw new DetectionException("Unsupported model");
}
ObbDetModel model = null;
try {
model = (ObbDetModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new DetectionException(e);
}
model.loadModel(config);
return model;
}
/**
* 注册模型
* @param modelEnum
* @param clazz
*/
private static void registerAlgorithm(ObbDetModelEnum modelEnum, Class<? extends ObbDetModel> clazz) {
registry.put(modelEnum, clazz);
}
/**
* 移除缓存的模型
* @param modelEnum
*/
public static void removeFromCache(ObbDetModelEnum modelEnum) {
modelMap.remove(modelEnum);
}
// 初始化默认算法
static {
registerAlgorithm(ObbDetModelEnum.YOLOV11, CommonObbDetModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

532
vision/src/main/java/cn/smartjavaai/obb/translator/YoloV11OddTranslator.java Normal file
View File

@@ -0,0 +1,532 @@
package cn.smartjavaai.obb.translator;
import ai.djl.Model;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.transform.*;
import ai.djl.ndarray.NDArray;
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.*;
import ai.djl.util.Utils;
import cn.smartjavaai.common.utils.LetterBoxUtils;
import cn.smartjavaai.obb.entity.ObbResult;
import cn.smartjavaai.obb.entity.YoloRotatedBox;
import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.*;
/**
* YoloV11OddTranslator
* @author dwj
*/
public class YoloV11OddTranslator implements Translator<Image, ObbResult> {
private int maxBoxes;
private YoloOutputType yoloOutputLayerType;
private float nmsThreshold;
protected float threshold;
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;
private SynsetLoader synsetLoader;
@Override
public void prepare(TranslatorContext ctx) throws IOException {
if (this.classes == null) {
this.classes = this.synsetLoader.load(ctx.getModel());
}
}
/**
* Constructs an ImageTranslator with the provided builder.
*
* @param builder the data to build with
*/
protected YoloV11OddTranslator(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;
}
/**
* 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;
}
@Override
public NDList processInput(TranslatorContext ctx, Image input) {
NDManager manager = ctx.getNDManager();
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
int imageWidth = (int) array.getShape().get(1);
int imageHeight = (int) array.getShape().get(0);
//Letter box resize 640x640 with padding (保持比例，补边缘)
LetterBoxUtils.ResizeResult letterBoxResult = LetterBoxUtils.letterbox(manager, array, width, height, 114f, LetterBoxUtils.PaddingPosition.CENTER);
array = letterBoxResult.image;
// 转为 float32 且归一化到 0~1
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
// HWC -> CHW
array = array.transpose(2, 0, 1); // CHW
ctx.setAttachment("width", input.getWidth());
ctx.setAttachment("height", input.getHeight());
ctx.setAttachment("processedWidth", width);
ctx.setAttachment("processedHeight", height);
return new NDList(array);
}
/** {@inheritDoc} */
protected ObbResult processFromBoxOutput(int imageWidth, int imageHeight, int processedWidth, int processedHeight, NDList list) {
float scale = Math.min((float) processedWidth / imageWidth, (float) processedHeight / imageHeight);
float padW = (processedWidth - imageWidth * scale) / 2;
float padH = (processedHeight - imageHeight * scale) / 2;
//[cx,cy,w,h,class*15,rotated]
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));
// reverse order search in heap; searches through #maxBoxes for optimization when set
List<YoloRotatedBox> rotatedBoxes = new ArrayList<>();
for (int i = numberRows - 1; i > numberRows - maxBoxes; --i) {
int index = i * nClasses;
// 找最大类别
float maxClassProb = -1f;
int maxIndex = -1;
for (int c = 4; c < nClasses - 1; c++) { // 类别从4开始，-1是rotated
float classProb = buf[index + c];
if (classProb > maxClassProb) {
maxClassProb = classProb;
maxIndex = c - 4;
}
}
if (maxClassProb > threshold) {
float cx = buf[index];
float cy = buf[index + 1];
float w = buf[index + 2];
float h = buf[index + 3];
cx = (cx - padW) / scale;
cy = (cy - padH) / scale;
w = w / scale;
h = h / scale;
float angle = buf[index + nClasses - 1]; // 最后一个是旋转角度
YoloRotatedBox rotatedBox = new YoloRotatedBox(cx, cy, w, h, angle, classes.get(maxIndex), maxClassProb);
rotatedBoxes.add(rotatedBox);
}
}
List<YoloRotatedBox> rotatedBoxeList = rotatedNMS(rotatedBoxes, nmsThreshold);
return new ObbResult(rotatedBoxeList);
}
public static List<YoloRotatedBox> rotatedNMS(List<YoloRotatedBox> boxes, double iouThreshold) {
List<YoloRotatedBox> keep = new ArrayList<>();
boolean[] removed = new boolean[boxes.size()];
// 按 score 降序
// boxes.sort((b1, b2) -> Float.compare(b2.score, b1.score));
for (int i = 0; i < boxes.size(); i++) {
if (removed[i]) continue;
YoloRotatedBox ibox = boxes.get(i);
keep.add(ibox);
for (int j = i + 1; j < boxes.size(); j++) {
if (removed[j]) continue;
YoloRotatedBox jbox = boxes.get(j);
if (!ibox.className.equals(jbox.className)) continue;
double iou = YoloRotatedBox.probiou(ibox, jbox, 1e-7);
if (iou > iouThreshold - 1e-7) {
removed[j] = true;
}
}
}
return keep;
}
private ObbResult processFromDetectOutput() {
throw new UnsupportedOperationException(
"detect layer output is not supported yet, check correct YoloV5 export format");
}
@Override
public ObbResult processOutput(TranslatorContext ctx, NDList list) throws Exception {
int imageWidth = (Integer) ctx.getAttachment("width");
int imageHeight = (Integer) ctx.getAttachment("height");
int processedWidth = (Integer) ctx.getAttachment("processedWidth");
int processedHeight = (Integer) ctx.getAttachment("processedHeight");
switch (yoloOutputLayerType) {
case DETECT:
return processFromDetectOutput();
case AUTO:
if (list.get(0).getShape().dimension() > 2) {
return processFromDetectOutput();
} else {
return processFromBoxOutput(imageWidth, imageHeight, processedWidth, processedHeight, list);
}
case BOX:
default:
return processFromBoxOutput(imageWidth, imageHeight, processedWidth, processedHeight, list);
}
}
public static class Builder {
private int maxBox = 8400;
YoloOutputType outputType;
float nmsThreshold;
protected float threshold = 0.25F;
protected boolean applyRatio;
protected boolean removePadding;
protected int width = 224;
protected int height = 224;
protected Image.Flag flag;
protected Pipeline pipeline;
protected Batchifier batchifier;
protected SynsetLoader synsetLoader;
public Builder() {
this.outputType = YoloOutputType.AUTO;
this.nmsThreshold = 0.45F;
}
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 YoloV11OddTranslator build() {
if (pipeline == null) {
addTransform(
array -> array.transpose(2, 0, 1).toType(DataType.FLOAT32, false).div(255));
}
// validate();
return new YoloV11OddTranslator(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();
}
/**
* Sets the name of the synset file listing the potential classes for an image.
*
* @param synsetArtifactName a file listing the potential classes for an image
* @return the builder
*/
public Builder optSynsetArtifactName(String synsetArtifactName) {
synsetLoader = new SynsetLoader(synsetArtifactName);
return self();
}
/**
* Sets the URL of the synset file.
*
* @param synsetUrl the URL of the synset file
* @return the builder
*/
public Builder optSynsetUrl(String synsetUrl) {
try {
this.synsetLoader = new SynsetLoader(new URL(synsetUrl));
} catch (MalformedURLException e) {
throw new IllegalArgumentException("Invalid synsetUrl: " + synsetUrl, e);
}
return self();
}
/**
* Sets the potential classes for an image.
*
* @param synset the potential classes for an image
* @return the builder
*/
public Builder optSynset(List<String> synset) {
synsetLoader = new SynsetLoader(synset);
return 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.25F);
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.45F);
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() {
}
}
protected static final class SynsetLoader {
private String synsetFileName;
private URL synsetUrl;
private List<String> synset;
public SynsetLoader(List<String> synset) {
this.synset = synset;
}
public SynsetLoader(URL synsetUrl) {
this.synsetUrl = synsetUrl;
}
public SynsetLoader(String synsetFileName) {
this.synsetFileName = synsetFileName;
}
public List<String> load(Model model) throws IOException {
if (synset != null) {
return synset;
} else if (synsetUrl != null) {
try (InputStream is = synsetUrl.openStream()) {
return Utils.readLines(is);
}
}
return model.getArtifact(synsetFileName, Utils::readLines);
}
}
}

6
vision/src/main/java/cn/smartjavaai/objectdetection/config/PersonDetModelConfig.java
View File

@@ -34,12 +34,6 @@ public class PersonDetModelConfig extends ModelConfig {
*/
private String modelPath;
/**
* 允许的分类列表
*/
private List<String> allowedClasses;
/**
* 按置信度分数排序后，最多保留的检测框数量
*/

16
vision/src/main/java/cn/smartjavaai/objectdetection/criteria/CriteriaBuilderFactory.java
View File

@@ -17,23 +17,23 @@ public class CriteriaBuilderFactory {
public static Criteria<Image, DetectedObjects> createCriteria(DetectorModelConfig config) {
//以下模型modelPath不允许为空
if(config.getModelEnum() == DetectorModelEnum.YOLOV8_OFFICIAL ||
config.getModelEnum() == DetectorModelEnum.YOLOV12_OFFICIAL ||
config.getModelEnum() == DetectorModelEnum.YOLOV8_CUSTOM ||
config.getModelEnum() == DetectorModelEnum.YOLOV12_CUSTOM ||
if(config.getModelEnum() == DetectorModelEnum.YOLOV8_OFFICIAL_ONNX ||
config.getModelEnum() == DetectorModelEnum.YOLOV12_OFFICIAL_ONNX ||
config.getModelEnum() == DetectorModelEnum.YOLOV8_CUSTOM_ONNX ||
config.getModelEnum() == DetectorModelEnum.YOLOV12_CUSTOM_ONNX ||
config.getModelEnum() == DetectorModelEnum.TENSORFLOW2_OFFICIAL){
if(StringUtils.isBlank(config.getModelPath())){
throw new DetectionException("modelPath is null");
}
}
switch (config.getModelEnum()) {
case YOLOV8_OFFICIAL:
case YOLOV8_OFFICIAL_ONNX:
return new YoloCriteriaBuilder().buildCriteria(config);
case YOLOV12_OFFICIAL:
case YOLOV12_OFFICIAL_ONNX:
return new YoloCriteriaBuilder().buildCriteria(config);
case YOLOV8_CUSTOM:
case YOLOV8_CUSTOM_ONNX:
return new YoloCriteriaBuilder().buildCriteria(config);
case YOLOV12_CUSTOM:
case YOLOV12_CUSTOM_ONNX:
return new YoloCriteriaBuilder().buildCriteria(config);
case TENSORFLOW2_OFFICIAL:
return new Tensorflow2CriteriaBuilder().buildCriteria(config);

49
vision/src/main/java/cn/smartjavaai/objectdetection/enums/DetectorModelEnum.java
View File

@@ -8,40 +8,43 @@ package cn.smartjavaai.objectdetection.enums;
public enum DetectorModelEnum {
// resnet50 系列
SSD_300_RESNET50("ai.djl.pytorch/ssd/0.0.1/ssd_300_resnet50"),
SSD_512_RESNET50_V1_VOC("ai.djl./ssd/0.0.1/ssd_512_resnet50_v1_voc"),
SSD_300_RESNET50_DJL("ai.djl.pytorch/ssd/0.0.1/ssd_300_resnet50"),
SSD_512_RESNET50_V1_VOC_DJL("ai.djl./ssd/0.0.1/ssd_512_resnet50_v1_voc"),
// vgg16 系列
SSD_512_VGG16_ATROUS_COCO("ai.djl.mxnet/ssd/0.0.1/ssd_512_vgg16_atrous_coco"),
SSD_300_VGG16_ATROUS_VOC("ai.djl.mxnet/ssd/0.0.1/ssd_300_vgg16_atrous_voc"),
SSD_512_VGG16_ATROUS_COCO_DJL("ai.djl.mxnet/ssd/0.0.1/ssd_512_vgg16_atrous_coco"),
SSD_300_VGG16_ATROUS_VOC_DJL("ai.djl.mxnet/ssd/0.0.1/ssd_300_vgg16_atrous_voc"),
// mobilenet 系列
SSD_512_MOBILENET1_VOC("ai.djl.mxnet/ssd/0.0.1/ssd_512_mobilenet1.0_voc"),
SSD_512_MOBILENET1_VOC_DJL("ai.djl.mxnet/ssd/0.0.1/ssd_512_mobilenet1.0_voc"),
// YOLO 系列
YOLOV8N("ai.djl.pytorch/yolov8n/0.0.1/yolov8n"),
YOLO11N("ai.djl.pytorch/yolo11n/0.0.1/yolo11n"),
YOLOV5S("ai.djl.pytorch/yolo5s/0.0.1/yolov5s"),
YOLOV5S_ONNXRUNTIME("ai.djl.onnxruntime/yolo5s/0.0.1/yolo5s"),
YOLO("ai.djl.mxnet/yolo/0.0.1/yolo"),
// YOLOV8N("ai.djl.pytorch/yolov8n/0.0.1/yolov8n"),
// YOLO11N("ai.djl.pytorch/yolo11n/0.0.1/yolo11n"),
YOLOV5S_DJL("ai.djl.pytorch/yolo5s/0.0.1/yolov5s"),
YOLOV5S_ONNX_DJL("ai.djl.onnxruntime/yolo5s/0.0.1/yolo5s"),
YOLO_DJL("ai.djl.mxnet/yolo/0.0.1/yolo"),
// YOLOv3 变体
YOLO3_DARKNET_VOC_416("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_voc_416"),
YOLO3_MOBILENET_VOC_320("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_voc_320"),
YOLO3_MOBILENET_VOC_416("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_voc_41"),
YOLO3_DARKNET_COCO_320("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_coco_320"),
YOLO3_DARKNET_COCO_416("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_coco_416"),
YOLO3_DARKNET_COCO_608("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_coco_608"),
YOLO3_MOBILENET_COCO_320("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_coco_320"),
YOLO3_MOBILENET_COCO_416("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_coco_416"),
YOLO3_MOBILENET_COCO_608("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_coco_608"),
YOLO3_DARKNET_VOC_416_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_voc_416"),
YOLO3_MOBILENET_VOC_320_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_voc_320"),
YOLO3_MOBILENET_VOC_416_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_voc_41"),
YOLO3_DARKNET_COCO_320_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_coco_320"),
YOLO3_DARKNET_COCO_416_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_coco_416"),
YOLO3_DARKNET_COCO_608_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_darknet_coco_608"),
YOLO3_MOBILENET_COCO_320_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_coco_320"),
YOLO3_MOBILENET_COCO_416_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_coco_416"),
YOLO3_MOBILENET_COCO_608_DJL("ai.djl.mxnet/yolo/0.0.1/yolo3_mobilenet_coco_608"),
YOLOV12_OFFICIAL(""),
YOLOV8_OFFICIAL(""),
YOLOV8_CUSTOM(""),
YOLOV8_OFFICIAL_ONNX(""),
YOLOV11_OFFICIAL_ONNX(""),
YOLOV12_OFFICIAL_ONNX(""),
YOLOV12_CUSTOM(""),
YOLOV8_CUSTOM_ONNX(""),
YOLOV11_CUSTOM_ONNX(""),
YOLOV12_CUSTOM_ONNX(""),
// TensorFlow 2.x 官方模型
TENSORFLOW2_OFFICIAL("");

9
vision/src/main/java/cn/smartjavaai/objectdetection/model/DetectorModel.java
View File

@@ -11,9 +11,9 @@ import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ZooModel;
import cn.smartjavaai.common.entity.DetectionResponse;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.pool.PredictorFactory;
import cn.smartjavaai.common.utils.FileUtils;
import cn.smartjavaai.common.utils.FrameConverterUtil;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import cn.smartjavaai.objectdetection.config.DetectorModelConfig;
@@ -119,6 +119,9 @@ public class DetectorModel implements AutoCloseable{
try {
img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
DetectedObjects detectedObjects = detect(img);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
throw new DetectionException("未检测到图片中的物体");
}
img.drawBoundingBoxes(detectedObjects);
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
// 调用 save 方法将 Image 写入字节流
@@ -186,6 +189,9 @@ public class DetectorModel implements AutoCloseable{
}
Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
DetectedObjects detectedObjects = detect(img);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
throw new DetectionException("未检测到图片中的物体");
}
img.drawBoundingBoxes(detectedObjects);
try {
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
@@ -225,7 +231,6 @@ public class DetectorModel implements AutoCloseable{
if (predictor != null) {
try {
predictorPool.returnObject(predictor); //归还
log.debug("释放资源");
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {

12
vision/src/main/java/cn/smartjavaai/objectdetection/model/ObjectDetectionModelFactory.java
View File

@@ -61,12 +61,12 @@ public class ObjectDetectionModelFactory {
* 获取默认模型
* @return
*/
public DetectorModel getModel() {
// 初始化默认配置
DetectorModelConfig config = new DetectorModelConfig();
config.setModelEnum(DetectorModelEnum.YOLO11N);
return getModel(config);
}
// public DetectorModel getModel() {
// // 初始化默认配置
// DetectorModelConfig config = new DetectorModelConfig();
// config.setModelEnum(DetectorModelEnum.YOLO11N);
// return getModel(config);
// }
/**
* 关闭所有已加载的模型

18
vision/src/main/java/cn/smartjavaai/objectdetection/model/person/CommonPersonDetModel.java
View File

@@ -8,10 +8,12 @@ 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.hutool.core.img.ImgUtil;
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.common.utils.ImageUtils;
import cn.smartjavaai.obb.config.ObbDetModelConfig;
import cn.smartjavaai.obb.criteria.ObbDetCriteriaFactory;
import cn.smartjavaai.obb.entity.ObbResult;
@@ -27,6 +29,7 @@ import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.opencv.core.Mat;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
@@ -92,7 +95,7 @@ public class CommonPersonDetModel implements PersonDetModel {
DetectedObjects detectedObjects = predictor.predict(image);
//过滤
if(Objects.nonNull(detectedObjects) && detectedObjects.getNumberOfObjects() > 0){
DetectedObjectsFilter detectedObjectsFilter = new DetectedObjectsFilter(config.getAllowedClasses(), config.getTopK());
DetectedObjectsFilter detectedObjectsFilter = new DetectedObjectsFilter(null, config.getTopK());
detectedObjects = detectedObjectsFilter.filter(detectedObjects);
}
return detectedObjects;
@@ -118,9 +121,13 @@ public class CommonPersonDetModel implements PersonDetModel {
@Override
public R<DetectionResponse> detectAndDraw(Image image) {
DetectedObjects detectedObjects = detectCore(image);
image.drawBoundingBoxes(detectedObjects);
DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, image);
detectionResponse.setDrawnImage(image);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
Image drawnImage = ImageUtils.copy(image);
drawnImage.drawBoundingBoxes(detectedObjects);
DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, drawnImage);
detectionResponse.setDrawnImage(drawnImage);
return R.ok(detectionResponse);
}
@@ -129,6 +136,9 @@ public class CommonPersonDetModel implements PersonDetModel {
try {
Image img = SmartImageFactory.getInstance().fromFile(Paths.get(imagePath));
DetectedObjects detectedObjects = detectCore(img);
if(Objects.isNull(detectedObjects) || detectedObjects.getNumberOfObjects() == 0){
return R.fail(R.Status.NO_OBJECT_DETECTED);
}
img.drawBoundingBoxes(detectedObjects);
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
// 调用 save 方法将 Image 写入字节流

109
vision/src/main/java/cn/smartjavaai/objectdetection/model/person/PersonDetModelFactory.java Normal file
View File

@@ -0,0 +1,109 @@
package cn.smartjavaai.objectdetection.model.person;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.objectdetection.config.PersonDetModelConfig;
import cn.smartjavaai.objectdetection.config.PersonDetModelConfig;
import cn.smartjavaai.objectdetection.enums.PersonDetectorModelEnum;
import cn.smartjavaai.objectdetection.enums.PersonDetectorModelEnum;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 行人检测 模型工厂
* @author dwj
*/
@Slf4j
public class PersonDetModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile PersonDetModelFactory instance;
private static final ConcurrentHashMap<PersonDetectorModelEnum, PersonDetModel> modelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<PersonDetectorModelEnum, Class<? extends PersonDetModel>> registry =
new ConcurrentHashMap<>();
// 私有构造函数，防止外部创建实例
private PersonDetModelFactory() {}
// 双重检查锁定的单例方法
public static PersonDetModelFactory getInstance() {
if (instance == null) {
synchronized (PersonDetModelFactory.class) {
if (instance == null) {
instance = new PersonDetModelFactory();
}
}
}
return instance;
}
/**
* 获取模型（通过配置）
* @param config
* @return
*/
public PersonDetModel getModel(PersonDetModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型");
}
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createFaceDetModel(config);
});
}
/**
* 使用ModelConfig创建模型
* @param config
* @return
*/
private PersonDetModel createFaceDetModel(PersonDetModelConfig config) {
Class<?> clazz = registry.get(config.getModelEnum());
if(clazz == null){
throw new DetectionException("Unsupported model");
}
PersonDetModel model = null;
try {
model = (PersonDetModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new DetectionException(e);
}
model.loadModel(config);
return model;
}
/**
* 注册模型
* @param modelEnum
* @param clazz
*/
private static void registerAlgorithm(PersonDetectorModelEnum modelEnum, Class<? extends PersonDetModel> clazz) {
registry.put(modelEnum, clazz);
}
/**
* 移除缓存的模型
* @param modelEnum
*/
public static void removeFromCache(PersonDetectorModelEnum modelEnum) {
modelMap.remove(modelEnum);
}
// 初始化默认算法
static {
registerAlgorithm(PersonDetectorModelEnum.YOLOV8_PERSON, CommonPersonDetModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

16
vision/src/main/java/cn/smartjavaai/objectdetection/stream/StreamDetectionListener.java
View File

@@ -12,5 +12,21 @@ import java.util.List;
public interface StreamDetectionListener {
/**
* 当检测到目标时回调
* @param detectionInfoList 目标信息列表
* @param image 检测到的图片
*/
void onObjectDetected(List<DetectionInfo> detectionInfoList, Image image);
/**
* 当视频文件读取完毕时回调
*/
void onStreamEnded();
/**
* 当视频流断开连接时回调
*/
void onStreamDisconnected();
}

226
vision/src/main/java/cn/smartjavaai/objectdetection/stream/StreamDetector.java
View File

@@ -1,15 +1,16 @@
package cn.smartjavaai.objectdetection.stream;
import ai.djl.inference.Predictor;
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.Mask;
import ai.djl.modality.cv.output.Rectangle;
import cn.hutool.core.date.LocalDateTimeUtil;
import cn.hutool.core.lang.UUID;
import cn.smartjavaai.common.entity.*;
import cn.smartjavaai.common.enums.VideoSourceType;
import cn.smartjavaai.common.utils.FrameConverterUtil;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import cn.smartjavaai.objectdetection.exception.DetectionException;
@@ -18,11 +19,10 @@ import cn.smartjavaai.vision.utils.DetectorUtils;
import lombok.extern.slf4j.Slf4j;
import nu.pattern.OpenCV;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.bytedeco.ffmpeg.global.avutil;
import org.bytedeco.javacv.*;
import org.bytedeco.opencv.global.opencv_imgcodecs;
import org.opencv.core.Core;
import org.opencv.core.Mat;
import org.opencv.imgcodecs.Imgcodecs;
import java.util.*;
import java.util.concurrent.*;
@@ -41,10 +41,14 @@ public class StreamDetector implements AutoCloseable{
private DetectorModel detectorModel;
private String streamUrl;
private ExecutorService grabberExecutor; // 专门抓帧的线程
private ExecutorService processorExecutor; // 专门处理帧的线程
//专门抓帧的线程
private ExecutorService grabberExecutor;
//专门处理帧的线程
private ExecutorService processorExecutor;
//回调线程池
ExecutorService callbackExecutor;
private int frameDetectionInterval = 1;
private int repeatGap = 5; // 秒
private long repeatGap = 5; // 秒
private volatile boolean isRunning;
private FrameGrabber grabber;
private StreamDetectionListener listener;
@@ -55,6 +59,18 @@ public class StreamDetector implements AutoCloseable{
private Map<String, Long> lastDetectTime = new ConcurrentHashMap<>();
private BlockingQueue<Frame> frameQueue = new LinkedBlockingQueue<>(100);
private GenericObjectPool<Predictor<Image, DetectedObjects>> predictorPool;
private Predictor<Image, DetectedObjects> predictor;
boolean grabberFinished = false; // 标记结束
//空帧数量
private int nullFrameCount = 0;
// 连续多少次空帧认为断联
private static final int MAX_NULL_FRAMES = 5;
public static Builder builder() { return new Builder(); }
private StreamDetector(Builder builder) {
@@ -64,9 +80,8 @@ public class StreamDetector implements AutoCloseable{
this.listener = builder.listener;
this.sourceType = builder.sourceType;
this.cameraIndex = builder.cameraIndex;
this.repeatGap = builder.repeatGap;
this.converterToMat = new OpenCVFrameConverter.ToOrgOpenCvCoreMat();
}
private void initializeGrabber() throws FrameGrabber.Exception {
@@ -78,34 +93,61 @@ public class StreamDetector implements AutoCloseable{
if (sourceType == VideoSourceType.STREAM) {
grabber.setOption("rtsp_transport", "tcp");
grabber.setOption("buffer_size", "1024000");
grabber.setOption("stimeout", "20000000");
grabber.setOption("max_delay", "500000");
grabber.setOption("stimeout", "2000000"); // 超时：单位微秒，这里是2秒
grabber.setOption("rw_timeout", "2000000"); // 读超时
grabber.setOption("max_delay", "5000000");
grabber.setOption("timeout", "2000000"); // 总超时
}
}
//日志级别
avutil.av_log_set_level(avutil.AV_LOG_ERROR);
grabber.start();
if(sourceType == VideoSourceType.FILE){
// 总帧数
int totalFrames = grabber.getLengthInFrames();
log.info("视频帧数：{}", totalFrames);
}
}
public void startDetection() {
if (isRunning) return;
isRunning = true;
if (isRunning){
throw new RuntimeException("当前正在运行中");
}
grabberFinished = false;
//获取模型Predictor
predictorPool = detectorModel.getPool();
try {
predictor = predictorPool.borrowObject();
} catch (Exception e) {
throw new DetectionException(e);
}
// 初始化抓帧线程池
if (grabberExecutor == null) grabberExecutor = Executors.newSingleThreadExecutor();
if (grabberExecutor == null || grabberExecutor.isShutdown())
grabberExecutor = Executors.newFixedThreadPool(2);
// 初始化帧处理线程池
if (processorExecutor == null) processorExecutor = Executors.newSingleThreadExecutor();
if (processorExecutor == null || processorExecutor.isShutdown())
processorExecutor = Executors.newFixedThreadPool(2);
if (callbackExecutor == null || callbackExecutor.isShutdown())
callbackExecutor = Executors.newFixedThreadPool(4);
try {
initializeGrabber();
} catch (FrameGrabber.Exception e) {
throw new DetectionException("视频流检测启动失败", e);
}
isRunning = true;
log.debug("视频流处理已启动");
// 初始化抓帧线程
grabberExecutor.submit(() -> {
try {
initializeGrabber();
processFrames();
} catch (Exception e) {
log.error("视频流处理异常", e);
} finally {
release();
//标识抓取已结束
grabberFinished = true;
}
});
log.info("视频流处理已启动");
// 初始化队列处理线程：解决回调比较耗时，导致线程池爆满
startFrameProcessor();
}
@@ -114,22 +156,42 @@ public class StreamDetector implements AutoCloseable{
* 负责抓取视频帧到队列
*/
private void processFrames() {
int frameCount = 0;
while (isRunning) {
long frameCount = 0;
while (!grabberFinished && isRunning) {
try {
Frame frame = grabber.grab();
if (frame == null || frame.image == null) continue;
Frame frame = grabber.grabFrame();
if (frame == null || frame.image == null) {
if(sourceType == VideoSourceType.FILE){
log.debug("视频检测结束");
grabberFinished = true;
break;
}else{
log.debug("未检测到视频帧");
nullFrameCount++;
if (nullFrameCount > MAX_NULL_FRAMES) {
log.warn("检测到视频断开，已超过最大空帧次数");
if(isRunning){
stopDetection();
}
if (listener != null) {
listener.onStreamDisconnected();
}
break;
}
continue;
}
}
nullFrameCount = 0; // 只要拿到正常帧就清零
frameCount++;
if (frameCount % frameDetectionInterval != 0) continue;
Frame currentFrame = frame.clone();
frameQueue.offer(currentFrame); // 队列满则丢弃，可改为 put 阻塞
frameQueue.offer(currentFrame);
// log.debug("正在抓取第{}帧，当前帧数：{}", frameCount, frameQueue.size());
} catch (Exception e) {
log.error("抓取视频帧异常", e);
if (e instanceof FFmpegFrameGrabber.Exception) reconnect();
}
}
log.debug("抓取帧线程退出");
}
/**
@@ -137,18 +199,29 @@ public class StreamDetector implements AutoCloseable{
*/
private void startFrameProcessor() {
processorExecutor.submit(() -> {
log.info("帧处理线程已启动");
while (isRunning || !frameQueue.isEmpty()) {
log.debug("帧处理线程已启动");
while ((!grabberFinished || !frameQueue.isEmpty()) && isRunning) {
try {
Frame frame = frameQueue.poll(100, TimeUnit.MILLISECONDS);
if (frame != null) processFrame(frame);
if (frame != null) {
processFrame(frame);
}
} catch (InterruptedException e) {
e.printStackTrace();
log.debug("帧处理线程被中断，准备退出");
break;
} catch (Exception e) {
log.error("帧处理异常", e);
}
}
if(isRunning){
stopDetection();
}
if (listener != null) {
listener.onStreamEnded();
}
log.debug("帧处理线程退出");
});
}
private void processFrame(Frame frame) {
@@ -158,12 +231,16 @@ public class StreamDetector implements AutoCloseable{
if (mat == null) return;
Image image = ImageFactory.getInstance().fromImage(mat);
DetectedObjects detectedObjects = detectorModel.detect(image);
// log.debug("检测结果：{}", detectedObjects.toString());
DetectedObjects detectedObjects = predictor.predict(image);
// log.info("内部检测结果：{}", detectedObjects.toString());
DetectionResponse detectionResponse = DetectorUtils.convertToDetectionResponse(detectedObjects, image);
if(Objects.isNull(detectionResponse)){
return;
}
List<DetectionInfo> filtered = filterRepeatedObjects(detectionResponse);
if (!filtered.isEmpty() && listener != null) {
listener.onObjectDetected(filtered, image); // 同帧多物体一次回调
Image copyImage = image.duplicate();
callbackExecutor.submit(() -> listener.onObjectDetected(filtered, copyImage));
}
} catch (Throwable e) {
e.printStackTrace();
@@ -187,47 +264,86 @@ public class StreamDetector implements AutoCloseable{
return result;
}
private void reconnect() {
log.info("尝试重新连接视频流");
try {
release();
Thread.sleep(5000);
initializeGrabber();
} catch (Exception e) {
log.error("重新连接RTSP流失败", e);
/**
* 开始检测下一个视频文件
*/
public void startNextVideo(String videoPath) {
if(!grabberFinished){
throw new DetectionException("当前视频未检测结束，请先关闭当前检测，再切换下一个视频");
}
if(sourceType != VideoSourceType.FILE){
throw new DetectionException("sourceType不是文件");
}
this.streamUrl = videoPath;
this.grabberFinished = false;
startDetection();
}
public void stopDetection() { isRunning = false; }
private void release() {
public void stopDetection() {
log.debug("停止检测中...");
isRunning = false;
grabberFinished = true;
if (predictor != null && predictorPool != null) {
try {
predictorPool.returnObject(predictor); //归还
predictor = null;
predictorPool = null;
} catch (Exception e) {
log.warn("归还Predictor失败", e);
try {
predictor.close(); // 归还失败才销毁
} catch (Exception ex) {
log.error("关闭Predictor失败", ex);
}
}
}
if (grabber != null) {
try { grabber.stop(); grabber.release(); }
catch (FrameGrabber.Exception e) { log.error("释放Grabber失败", e); }
try {
grabber.stop(); grabber.release();
grabber = null;
}catch (FrameGrabber.Exception e) {
log.error("释放Grabber失败", e);
}
}
log.debug("停止检测完毕");
}
@Override
public void close() {
stopDetection();
if (grabberExecutor != null) grabberExecutor.shutdownNow();
if (processorExecutor != null) processorExecutor.shutdownNow();
release();
// if(isRunning){
// System.out.println("--isRunning：" + isRunning);
// stopDetection();
// }
if (grabberExecutor != null){
grabberExecutor.shutdownNow();
}
if (processorExecutor != null) {
processorExecutor.shutdownNow();
}
if (callbackExecutor != null) {
callbackExecutor.shutdownNow();
}
}
public static class Builder {
private DetectorModel detectorModel;
private String streamUrl;
private ExecutorService executorService;
private int frameDetectionInterval = 1;
private StreamDetectionListener listener;
private VideoSourceType sourceType = VideoSourceType.STREAM; // 默认流
private int cameraIndex = 0; // 默认第一个摄像头
private long repeatGap = 5;//同物体重复检测间隔
public Builder detectorModel(DetectorModel m) { this.detectorModel = m; return this; }
public Builder streamUrl(String url) { this.streamUrl = url; return this; }
public Builder executorService(ExecutorService es) { this.executorService = es; return this; }
public Builder listener(StreamDetectionListener listener) { this.listener = listener; return this; }
public Builder repeatGap(long repeatGap) {
this.repeatGap = repeatGap;
return this;
}
public Builder sourceType(VideoSourceType sourceType) {
this.sourceType = sourceType;
return this;
@@ -268,10 +384,6 @@ public class StreamDetector implements AutoCloseable{
throw new DetectionException("不支持的视频源类型: " + sourceType);
}
if (executorService == null) {
executorService = Executors.newFixedThreadPool(2); // 至少2个线程
}
return new StreamDetector(this);
}

13
vision/src/main/java/cn/smartjavaai/objectdetection/translator/YoloV8PersonDetTranslator.java
View File

@@ -136,7 +136,6 @@ public class YoloV8PersonDetTranslator implements Translator<Image, DetectedObje
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());
@@ -208,17 +207,7 @@ public class YoloV8PersonDetTranslator implements Translator<Image, DetectedObje
retProbs.add(scores.get(pos).doubleValue());
// Rectangle rect = boxes.get(pos);
Rectangle rect = boxes.get(pos);
if (removePadding) {
rect =
LetterBoxUtils.restoreBox(rect, scale, origImageWidth, origImageHeight, width, height);
} else if (applyRatio) {
rect =
new Rectangle(
rect.getX() / width,
rect.getY() / height,
rect.getWidth() / width,
rect.getHeight() / height);
}
rect = LetterBoxUtils.restoreBox(rect, scale, origImageWidth, origImageHeight, width, height);
retBB.add(rect);
}
}

52
vision/src/main/java/cn/smartjavaai/pose/config/PoseModelConfig.java Normal file
View File

@@ -0,0 +1,52 @@
package cn.smartjavaai.pose.config;
import cn.smartjavaai.action.enums.ActionRecModelEnum;
import cn.smartjavaai.common.config.ModelConfig;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.pose.enums.PoseModelEnum;
import lombok.Data;
import java.util.List;
/**
* 姿态估计模型参数配置
*
* @author dwj
*/
@Data
public class PoseModelConfig extends ModelConfig {
/**
* 模型
*/
private PoseModelEnum modelEnum;
/**
* 模型路径
*/
private String modelPath;
/**
* 置信度阈值
*/
private float threshold;
public PoseModelConfig() {
}
public PoseModelConfig(PoseModelEnum modelEnum, DeviceEnum device) {
this.modelEnum = modelEnum;
setDevice(device);
}
public PoseModelConfig(PoseModelEnum modelEnum) {
this.modelEnum = modelEnum;
}
}

78
vision/src/main/java/cn/smartjavaai/pose/criteria/PoseCriteriaFactory.java Normal file
View File

@@ -0,0 +1,78 @@
package cn.smartjavaai.pose.criteria;
import ai.djl.Device;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.Joints;
import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.Translator;
import cn.smartjavaai.action.config.ActionRecModelConfig;
import cn.smartjavaai.action.enums.ActionRecModelEnum;
import cn.smartjavaai.action.model.CommonActionTranslator;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.action.exception.ActionException;
import cn.smartjavaai.common.utils.DJLCommonUtils;
import cn.smartjavaai.obb.config.ObbDetModelConfig;
import cn.smartjavaai.obb.entity.ObbResult;
import cn.smartjavaai.obb.enums.ObbDetModelEnum;
import cn.smartjavaai.obb.exception.ObbDetException;
import cn.smartjavaai.obb.translator.YoloV11OddTranslator;
import cn.smartjavaai.objectdetection.constant.DetectorConstant;
import cn.smartjavaai.pose.config.PoseModelConfig;
import cn.smartjavaai.pose.exception.PoseException;
import org.apache.commons.lang3.StringUtils;
import java.nio.file.Paths;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 姿势估计模型工厂
* @author dwj
*/
public class PoseCriteriaFactory {
/**
* 创建姿势估计Criteria
* @param config
* @return
*/
public static Criteria<Image, Joints[]> createCriteria(PoseModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, Joints[]> criteria = null;
//DJL官方提供模型
if(StringUtils.isNotBlank(config.getModelEnum().getModelUri())){
criteria = createDJLCriteria(config, device);
}
return criteria;
}
/**
* 创建DJL官方模型Criteria
* 需要模型同目录下存在：serving.properties
* @param config
* @return
*/
public static Criteria<Image, Joints[]> createDJLCriteria(PoseModelConfig config, Device device) {
if(StringUtils.isNotBlank(config.getModelPath())
&& DJLCommonUtils.isServingPropertiesExists(Paths.get(config.getModelPath()))){
throw new PoseException("模型所在目录未找到 serving.properties 文件");
}
Criteria<Image, Joints[]> criteria =
Criteria.builder()
.setTypes(Image.class, Joints[].class)
.optModelUrls(StringUtils.isNotBlank(config.getModelPath()) ? null : config.getModelEnum().getModelUri())
.optModelPath(StringUtils.isNotBlank(config.getModelPath()) ? Paths.get(config.getModelPath()) : null)
.optDevice(device)
.optArgument("threshold", config.getThreshold() > 0 ? config.getThreshold() : null)
.optProgress(new ProgressBar())
.build();
return criteria;
}
}

42
vision/src/main/java/cn/smartjavaai/pose/enums/PoseModelEnum.java Normal file
View File

@@ -0,0 +1,42 @@
package cn.smartjavaai.pose.enums;
/**
* 姿态估计模型枚举
* @author dwj
*/
public enum PoseModelEnum {
YOLO11N_POSE_PT("djl://ai.djl.pytorch/yolo11n-pose"),
YOLOV8N_POSE_PT("djl://ai.djl.pytorch/yolov8n-pose"),
YOLO11N_POSE_ONNX("djl://ai.djl.onnxruntime/yolo11n-pose"),
YOLOV8N_POSE_ONNX("djl://ai.djl.onnxruntime/yolov8n-pose");
// SIMPLE_POSE_MXNET("djl://ai.djl.mxnet/simple_pose");
/**
* 根据名称获取枚举 (忽略大小写和下划线变体)
*/
public static PoseModelEnum fromName(String name) {
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
for (PoseModelEnum model : values()) {
if (model.name().replaceAll("_", "").equals(formatted)) {
return model;
}
}
throw new IllegalArgumentException("未知模型名称: " + name);
}
private final String modelUri;
PoseModelEnum(String modelUri) {
this.modelUri = modelUri;
}
public String getModelUri() {
return modelUri;
}
}

29
vision/src/main/java/cn/smartjavaai/pose/exception/PoseException.java Normal file
View File

@@ -0,0 +1,29 @@
package cn.smartjavaai.pose.exception;
/**
* 动作检测异常
* @author dwj
*/
public class PoseException extends RuntimeException{
public PoseException() {
super();
}
public PoseException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
super(message, cause, enableSuppression, writableStackTrace);
}
public PoseException(String message, Throwable cause) {
super(message, cause);
}
public PoseException(String message) {
super(message);
}
public PoseException(Throwable cause) {
super(cause);
}
}

147
vision/src/main/java/cn/smartjavaai/pose/model/CommonPoseModel.java Normal file
View File

@@ -0,0 +1,147 @@
package cn.smartjavaai.pose.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.Joints;
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.common.utils.ImageUtils;
import cn.smartjavaai.obb.exception.ObbDetException;
import cn.smartjavaai.objectdetection.config.PersonDetModelConfig;
import cn.smartjavaai.objectdetection.criteria.PersonDetCriteriaFactory;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import cn.smartjavaai.pose.config.PoseModelConfig;
import cn.smartjavaai.pose.criteria.PoseCriteriaFactory;
import cn.smartjavaai.vision.utils.DetectedObjectsFilter;
import cn.smartjavaai.vision.utils.DetectorUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.io.ByteArrayOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Objects;
/**
* 姿态估计模型
* @author dwj
*/
@Slf4j
public class CommonPoseModel implements PoseModel {
private PoseModelConfig config;
private ZooModel<Image, Joints[]> model;
private GenericObjectPool<Predictor<Image, Joints[]>> predictorPool;
@Override
public void loadModel(PoseModelConfig config) {
if(Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型枚举");
}
Criteria<Image, Joints[]> criteria = PoseCriteriaFactory.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);
}
}
/**
* 模型核心推理方法
* @param image
* @return
*/
@Override
public R<Joints[]> detect(Image image) {
Predictor<Image, Joints[]> predictor = null;
try {
predictor = predictorPool.borrowObject();
Joints[] joints = predictor.predict(image);
return R.ok(joints);
} 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 Image detectAndDraw(Image image) {
Image drawnImage = ImageUtils.copy(image);
R<Joints[]> allJoints = detect(drawnImage);
for (Joints joints : allJoints.getData()) {
drawnImage.drawJoints(joints);
}
return drawnImage;
}
@Override
public R<Joints[]> detectAndDraw(String imagePath, String outputPath) {
try {
Image img = SmartImageFactory.getInstance().fromFile(Paths.get(imagePath));
R<Joints[]> allJoints = detect(img);
for (Joints joints : allJoints.getData()) {
img.drawJoints(joints);
}
// 调用 save 方法将 Image 写入字节流
img.save(Files.newOutputStream(Paths.get(outputPath)), "png");
return allJoints;
} catch (IOException e) {
throw new ObbDetException(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);
}
}
}

112
vision/src/main/java/cn/smartjavaai/pose/model/PoseDetModelFactory.java Normal file
View File

@@ -0,0 +1,112 @@
package cn.smartjavaai.pose.model;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.obb.model.CommonObbDetModel;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import cn.smartjavaai.pose.config.PoseModelConfig;
import cn.smartjavaai.pose.enums.PoseModelEnum;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 姿态估计 模型工厂
* @author dwj
*/
@Slf4j
public class PoseDetModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile PoseDetModelFactory instance;
private static final ConcurrentHashMap<PoseModelEnum, PoseModel> modelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<PoseModelEnum, Class<? extends PoseModel>> registry =
new ConcurrentHashMap<>();
// 私有构造函数，防止外部创建实例
private PoseDetModelFactory() {}
// 双重检查锁定的单例方法
public static PoseDetModelFactory getInstance() {
if (instance == null) {
synchronized (PoseDetModelFactory.class) {
if (instance == null) {
instance = new PoseDetModelFactory();
}
}
}
return instance;
}
/**
* 获取模型（通过配置）
* @param config
* @return
*/
public PoseModel getModel(PoseModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型");
}
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createFaceDetModel(config);
});
}
/**
* 使用ModelConfig创建模型
* @param config
* @return
*/
private PoseModel createFaceDetModel(PoseModelConfig config) {
Class<?> clazz = registry.get(config.getModelEnum());
if(clazz == null){
throw new DetectionException("Unsupported model");
}
PoseModel model = null;
try {
model = (PoseModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new DetectionException(e);
}
model.loadModel(config);
return model;
}
/**
* 注册模型
* @param modelEnum
* @param clazz
*/
private static void registerAlgorithm(PoseModelEnum modelEnum, Class<? extends PoseModel> clazz) {
registry.put(modelEnum, clazz);
}
/**
* 移除缓存的模型
* @param modelEnum
*/
public static void removeFromCache(PoseModelEnum modelEnum) {
modelMap.remove(modelEnum);
}
// 初始化默认算法
static {
registerAlgorithm(PoseModelEnum.YOLOV8N_POSE_ONNX, CommonPoseModel.class);
registerAlgorithm(PoseModelEnum.YOLO11N_POSE_ONNX, CommonPoseModel.class);
registerAlgorithm(PoseModelEnum.YOLOV8N_POSE_PT, CommonPoseModel.class);
registerAlgorithm(PoseModelEnum.YOLO11N_POSE_PT, CommonPoseModel.class);
// registerAlgorithm(PoseModelEnum.SIMPLE_POSE_MXNET, CommonPoseModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

53
vision/src/main/java/cn/smartjavaai/pose/model/PoseModel.java Normal file
View File

@@ -0,0 +1,53 @@
package cn.smartjavaai.pose.model;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Joints;
import cn.smartjavaai.common.entity.DetectionResponse;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.objectdetection.config.PersonDetModelConfig;
import cn.smartjavaai.pose.config.PoseModelConfig;
/**
* 姿态估计模型
* @author dwj
*/
public interface PoseModel extends AutoCloseable{
/**
* 加载模型
* @param config
*/
void loadModel(PoseModelConfig config);
/**
* 姿态估计
* @param image
* @return
*/
default R<Joints[]> detect(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 姿态估计并绘制
* @param image
* @return
*/
default Image detectAndDraw(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 姿态估计并绘制
* @param imagePath
* @param outputPath
* @return
*/
default R<Joints[]> detectAndDraw(String imagePath, String outputPath){
throw new UnsupportedOperationException("默认不支持该功能");
}
}

52
vision/src/main/java/cn/smartjavaai/semseg/model/CommonSemSegModel.java
View File

@@ -6,23 +6,32 @@ import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.CategoryMask;
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.common.utils.Base64ImageUtils;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.instanceseg.exception.InstanceSegException;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import cn.smartjavaai.semseg.config.SemSegModelConfig;
import cn.smartjavaai.semseg.criteria.SemSegCriteriaFactory;
import cn.smartjavaai.vision.utils.CategoryMaskFilter;
import cn.smartjavaai.vision.utils.DetectorUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.impl.GenericObjectPool;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Objects;
/**
@@ -63,26 +72,12 @@ public class CommonSemSegModel implements SemSegModel {
}
}
@Override
public R<CategoryMask> detectBase64(String base64Image) {
if(StringUtils.isBlank(base64Image)){
return R.fail(R.Status.INVALID_IMAGE);
}
try {
byte[] imageData = Base64ImageUtils.base64ToImage(base64Image);
Image image = ImageFactory.getInstance().fromInputStream(new ByteArrayInputStream(imageData));
return detect(image);
} catch (IOException e) {
throw new DetectionException("读取图片异常", e);
}
}
@Override
public R<CategoryMask> detect(Image image) {
CategoryMask categoryMask = detectCore(image);
// 过滤
if(CollectionUtils.isNotEmpty(config.getAllowedClasses())
&& Objects.nonNull(categoryMask) && !categoryMask.getClasses().isEmpty()){
&& Objects.nonNull(categoryMask) && CollectionUtils.isNotEmpty(categoryMask.getClasses())){
categoryMask = new CategoryMaskFilter(config.getAllowedClasses()).filter(categoryMask);
}
return R.ok(categoryMask);
@@ -117,6 +112,33 @@ public class CommonSemSegModel implements SemSegModel {
}
}
@Override
public R<CategoryMask> detectAndDraw(String imagePath, String outputPath) {
try {
Image img = SmartImageFactory.getInstance().fromFile(Paths.get(imagePath));
CategoryMask categoryMask = detectCore(img);
if(Objects.isNull(categoryMask) || CollectionUtils.isEmpty(categoryMask.getClasses())){
throw new InstanceSegException("未检测到实例");
}
ImageUtils.drawMask(categoryMask, img, 180, 0);
img.save(Files.newOutputStream(Paths.get(outputPath)), "png");
return R.ok(categoryMask);
} catch (IOException e) {
throw new InstanceSegException(e);
}
}
@Override
public Image detectAndDraw(Image image) {
CategoryMask categoryMask = detectCore(image);
if(Objects.isNull(categoryMask) || CollectionUtils.isEmpty(categoryMask.getClasses())){
throw new InstanceSegException("未检测到实例");
}
Image drawnImage = ImageUtils.copy(image);
ImageUtils.drawMask(categoryMask, drawnImage, 180, 0);
return drawnImage;
}
@Override
public void close() throws Exception {
try {

18
vision/src/main/java/cn/smartjavaai/semseg/model/SemSegModel.java
View File

@@ -20,15 +20,6 @@ public interface SemSegModel extends AutoCloseable{
*/
void loadModel(SemSegModelConfig config);
/**
* 语义分割
* @param base64Image
* @return
*/
default R<CategoryMask> detectBase64(String base64Image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 语义分割
* @param image
@@ -38,5 +29,14 @@ public interface SemSegModel extends AutoCloseable{
throw new UnsupportedOperationException("默认不支持该功能");
}
default R<CategoryMask> detectAndDraw(String imagePath, String outputPath){
throw new UnsupportedOperationException("默认不支持该功能");
}
default Image detectAndDraw(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
}

107
vision/src/main/java/cn/smartjavaai/semseg/model/SemSegModelFactory.java Normal file
View File

@@ -0,0 +1,107 @@
package cn.smartjavaai.semseg.model;
import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.objectdetection.exception.DetectionException;
import cn.smartjavaai.semseg.config.SemSegModelConfig;
import cn.smartjavaai.semseg.enums.SemSegModelEnum;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 语义分割 模型工厂
* @author dwj
*/
@Slf4j
public class SemSegModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
private static volatile SemSegModelFactory instance;
private static final ConcurrentHashMap<SemSegModelEnum, SemSegModel> modelMap = new ConcurrentHashMap<>();
/**
* 模型注册表
*/
private static final Map<SemSegModelEnum, Class<? extends SemSegModel>> registry =
new ConcurrentHashMap<>();
// 私有构造函数，防止外部创建实例
private SemSegModelFactory() {}
// 双重检查锁定的单例方法
public static SemSegModelFactory getInstance() {
if (instance == null) {
synchronized (SemSegModelFactory.class) {
if (instance == null) {
instance = new SemSegModelFactory();
}
}
}
return instance;
}
/**
* 获取模型（通过配置）
* @param config
* @return
*/
public SemSegModel getModel(SemSegModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型");
}
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
return createFaceDetModel(config);
});
}
/**
* 使用ModelConfig创建模型
* @param config
* @return
*/
private SemSegModel createFaceDetModel(SemSegModelConfig config) {
Class<?> clazz = registry.get(config.getModelEnum());
if(clazz == null){
throw new DetectionException("Unsupported model");
}
SemSegModel model = null;
try {
model = (SemSegModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new DetectionException(e);
}
model.loadModel(config);
return model;
}
/**
* 注册模型
* @param modelEnum
* @param clazz
*/
private static void registerAlgorithm(SemSegModelEnum modelEnum, Class<? extends SemSegModel> clazz) {
registry.put(modelEnum, clazz);
}
/**
* 移除缓存的模型
* @param modelEnum
*/
public static void removeFromCache(SemSegModelEnum modelEnum) {
modelMap.remove(modelEnum);
}
// 初始化默认算法
static {
registerAlgorithm(SemSegModelEnum.DEEPLABV3, CommonSemSegModel.class);
log.debug("缓存目录：{}", Config.getCachePath());
}
}

60
vision/src/main/java/cn/smartjavaai/vision/utils/DetectorUtils.java
View File

@@ -6,6 +6,7 @@ import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Mask;
import ai.djl.modality.cv.output.Rectangle;
import cn.smartjavaai.common.entity.*;
import cn.smartjavaai.common.entity.Point;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.obb.entity.ObbResult;
import cn.smartjavaai.obb.entity.YoloRotatedBox;
@@ -15,6 +16,7 @@ import org.opencv.core.Scalar;
import org.opencv.imgproc.Imgproc;
import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
@@ -94,7 +96,7 @@ public class DetectorUtils {
/**
* 绘制文本框及文本
* 绘制旋转框及文本（opencv）
* @param srcMat
* @param rotatedBoxeList
*/
@@ -111,4 +113,60 @@ public class DetectorUtils {
}
/**
* 绘制旋转框及文本（BufferedImage）
* @param image
* @param rotatedBoxeList
*/
public static void drawRectWithText(BufferedImage image, List<YoloRotatedBox> rotatedBoxeList) {
Graphics2D g2d = image.createGraphics();
// 抗锯齿，让线条和文字更平滑
g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
g2d.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_ON);
// 在左上角画文字，带背景
FontMetrics fm = g2d.getFontMetrics();
for (YoloRotatedBox box : rotatedBoxeList) {
// 画笔颜色和字体
g2d.setColor(Color.GREEN);
g2d.setStroke(new BasicStroke(1.5f)); // 线条粗细
g2d.setFont(new Font("SansSerif", Font.PLAIN, 16)); // 设置中文字体，避免乱码
List<Point> points = box.toPoints();
// 画矩形的4条边
g2d.drawLine((int) points.get(0).getX(), (int) points.get(0).getY(), (int) points.get(1).getX(), (int) points.get(1).getY());
g2d.drawLine((int) points.get(1).getX(), (int) points.get(1).getY(), (int) points.get(2).getX(), (int) points.get(2).getY());
g2d.drawLine((int) points.get(2).getX(), (int) points.get(2).getY(), (int) points.get(3).getX(), (int) points.get(3).getY());
g2d.drawLine((int) points.get(3).getX(), (int) points.get(3).getY(), (int) points.get(0).getX(), (int) points.get(0).getY());
int percent = (int) Math.round(box.score * 100);
String className = box.className + " " + percent + "%";
int textWidth = fm.stringWidth(box.className + " " + box.score);
int textHeight = fm.getHeight();
// 文字位置
int textX = (int) points.get(0).getX();
int textY = (int) points.get(0).getY() - 5;
// 画背景矩形（半透明）
g2d.setColor(new Color(0, 0, 0, 128));
g2d.fillRect(textX, textY - textHeight, textWidth, textHeight);
g2d.setColor(Color.WHITE);
// 在左上角画文字
g2d.drawString(className, (int) points.get(0).getX(), (int) points.get(0).getY() - 5);
}
g2d.dispose();
}
/**
* 绘制旋转框及文本
* @param image
* @param rotatedBoxeList
*/
public static void drawRectWithText(Image image, List<YoloRotatedBox> rotatedBoxeList){
if(image.getWrappedImage() instanceof BufferedImage){
drawRectWithText((BufferedImage) image.getWrappedImage(), rotatedBoxeList);
}else{
drawRectWithText((Mat)image.getWrappedImage(), rotatedBoxeList);
}
}
}

12
vision/src/main/java/cn/smartjavaai/vision/utils/ObbResultFilter.java
View File

@@ -2,6 +2,7 @@ package cn.smartjavaai.vision.utils;
import cn.smartjavaai.obb.entity.ObbResult;
import cn.smartjavaai.obb.entity.YoloRotatedBox;
import org.apache.commons.collections.CollectionUtils;
import java.util.Collections;
import java.util.List;
@@ -26,11 +27,12 @@ public class ObbResultFilter {
return new ObbResult(Collections.emptyList());
}
// 1. 按类别过滤
List<YoloRotatedBox> filtered = result.getRotatedBoxeList().stream()
.filter(box -> allowedClasses.contains(box.className))
.collect(Collectors.toList());
List<YoloRotatedBox> filtered = result.getRotatedBoxeList();
if(CollectionUtils.isNotEmpty(allowedClasses)){
filtered = result.getRotatedBoxeList().stream()
.filter(box -> allowedClasses.contains(box.className))
.collect(Collectors.toList());
}
// 2. 按 score 排序
// filtered.sort((a, b) -> Float.compare(b.score, a.score));