ZF/SmartJavaAI
1
0
Fork 0
mirror of https://github.com/geekwenjie/SmartJavaAI.git synced 2026-09-20 01:29:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

临时提交

Browse Source
This commit is contained in:
dengwenjie
2025-08-29 10:30:35 +08:00
parent 8bf620a330
commit 86ea7eb03e
364 changed files with 8572 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
vision/src/main/java/cn/smartjavaai/semseg/config/SemSegModelConfig.java Normal file
View File

@@ -0,0 +1,49 @@
package cn.smartjavaai.semseg.config;
import cn.smartjavaai.common.config.ModelConfig;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.semseg.enums.SemSegModelEnum;
import lombok.Data;
import java.util.List;
/**
* 语义分割模型参数配置
*
* @author dwj
*/
@Data
public class SemSegModelConfig extends ModelConfig {
/**
* 模型
*/
private SemSegModelEnum modelEnum;
/**
* 模型路径
*/
private String modelPath;
/**
* 允许的分类列表
*/
private List<String> allowedClasses;
public SemSegModelConfig() {
}
public SemSegModelConfig(SemSegModelEnum modelEnum, DeviceEnum device) {
this.modelEnum = modelEnum;
setDevice(device);
}
public SemSegModelConfig(SemSegModelEnum modelEnum) {
this.modelEnum = modelEnum;
}
}

48
vision/src/main/java/cn/smartjavaai/semseg/criteria/SemSegCriteriaFactory.java Normal file
View File

@@ -0,0 +1,48 @@
package cn.smartjavaai.semseg.criteria;
import ai.djl.Device;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.CategoryMask;
import ai.djl.modality.cv.translator.SemanticSegmentationTranslatorFactory;
import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.semseg.config.SemSegModelConfig;
import cn.smartjavaai.semseg.enums.SemSegModelEnum;
import org.apache.commons.lang3.StringUtils;
import java.nio.file.Paths;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* 语义分割Criteria工厂
* @author dwj
*/
public class SemSegCriteriaFactory {
public static Criteria<Image, CategoryMask> createCriteria(SemSegModelConfig config) {
Device device = null;
if(!Objects.isNull(config.getDevice())){
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
}
Criteria<Image, CategoryMask> criteria = null;
ConcurrentHashMap params = new ConcurrentHashMap<String, String>();
params.putAll(config.getCustomParams());
if(config.getModelEnum() == SemSegModelEnum.DEEPLABV3){
criteria =
Criteria.builder()
.setTypes(Image.class, CategoryMask.class)
.optModelUrls(StringUtils.isNotBlank(config.getModelPath()) ? null :
config.getModelEnum().getModelUri())
.optModelPath(StringUtils.isNotBlank(config.getModelPath()) ? Paths.get(config.getModelPath()) : null)
.optTranslatorFactory(new SemanticSegmentationTranslatorFactory())
.optEngine("PyTorch")
.optDevice(device)
.optProgress(new ProgressBar())
.build();
}
return criteria;
}
}

17
vision/src/main/java/cn/smartjavaai/semseg/entity/DetectParams.java Normal file
View File

@@ -0,0 +1,17 @@
package cn.smartjavaai.semseg.entity;
import lombok.Data;
/**
* 检测参数
* @author dwj
*/
@Data
public class DetectParams {
/**
* 置信度阈值
*/
private float threshold = 0.3f;
}

34
vision/src/main/java/cn/smartjavaai/semseg/enums/SemSegModelEnum.java Normal file
View File

@@ -0,0 +1,34 @@
package cn.smartjavaai.semseg.enums;
/**
* 语义分割模型枚举
* @author dwj
*/
public enum SemSegModelEnum {
DEEPLABV3("djl://ai.djl.pytorch/deeplabv3");
/**
* 根据名称获取枚举 (忽略大小写和下划线变体)
*/
public static SemSegModelEnum fromName(String name) {
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
for (SemSegModelEnum model : values()) {
if (model.name().replaceAll("_", "").equals(formatted)) {
return model;
}
}
throw new IllegalArgumentException("未知模型名称: " + name);
}
private final String modelUri;
SemSegModelEnum(String modelUri) {
this.modelUri = modelUri;
}
public String getModelUri() {
return modelUri;
}
}

29
vision/src/main/java/cn/smartjavaai/semseg/exception/SemSegException.java Normal file
View File

@@ -0,0 +1,29 @@
package cn.smartjavaai.semseg.exception;
/**
* 语义分割异常
* @author dwj
*/
public class SemSegException extends RuntimeException{
public SemSegException() {
super();
}
public SemSegException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
super(message, cause, enableSuppression, writableStackTrace);
}
public SemSegException(String message, Throwable cause) {
super(message, cause);
}
public SemSegException(String message) {
super(message);
}
public SemSegException(Throwable cause) {
super(cause);
}
}

137
vision/src/main/java/cn/smartjavaai/semseg/model/CommonSemSegModel.java Normal file
View File

@@ -0,0 +1,137 @@
package cn.smartjavaai.semseg.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.ImageFactory;
import ai.djl.modality.cv.output.CategoryMask;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ZooModel;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.pool.PredictorFactory;
import cn.smartjavaai.common.utils.Base64ImageUtils;
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 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.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.Objects;
/**
* 语义分割模型
* @author dwj
*/
@Slf4j
public class CommonSemSegModel implements SemSegModel {
private SemSegModelConfig config;
private ZooModel<Image, CategoryMask> model;
private GenericObjectPool<Predictor<Image, CategoryMask>> predictorPool;
@Override
public void loadModel(SemSegModelConfig config) {
if(Objects.isNull(config.getModelEnum())){
throw new DetectionException("未配置模型枚举");
}
Criteria<Image, CategoryMask> criteria = SemSegCriteriaFactory.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<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()){
categoryMask = new CategoryMaskFilter(config.getAllowedClasses()).filter(categoryMask);
}
return R.ok(categoryMask);
}
/**
* 模型核心推理方法
* @param image
* @return
*/
public CategoryMask detectCore(Image image) {
Predictor<Image, CategoryMask> predictor = null;
try {
predictor = predictorPool.borrowObject();
return predictor.predict(image);
} 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 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);
}
}
}

42
vision/src/main/java/cn/smartjavaai/semseg/model/SemSegModel.java Normal file
View File

@@ -0,0 +1,42 @@
package cn.smartjavaai.semseg.model;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.CategoryMask;
import cn.smartjavaai.common.entity.DetectionResponse;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.semseg.config.SemSegModelConfig;
/**
* 语义分割模型
* @author dwj
*/
public interface SemSegModel extends AutoCloseable{
/**
* 加载模型
* @param config
*/
void loadModel(SemSegModelConfig config);
/**
* 语义分割
* @param base64Image
* @return
*/
default R<CategoryMask> detectBase64(String base64Image){
throw new UnsupportedOperationException("默认不支持该功能");
}
/**
* 语义分割
* @param image
* @return
*/
default R<CategoryMask> detect(Image image){
throw new UnsupportedOperationException("默认不支持该功能");
}
}