mirror of
https://github.com/geekwenjie/SmartJavaAI.git
synced 2026-09-15 06:19:08 +00:00
【通用视觉】集成 OpenAI CLIP 模型，支持以图搜图、以文搜图、以图搜文等功能
【通用视觉】新增 YOLO 图像分类模型支持 【ASR/TTS】集成 Sherpa TTS（语音合成）与 ASR（语音识别）模块，支持中文、粤语、方言、英文等多种语言 【目标检测】优化视频目标检测功能
This commit is contained in:
dengwenjie
@@ -0,0 +1,62 @@
|
||||
package cn.smartjavaai.cls.config;
|
||||
|
||||
import cn.smartjavaai.action.enums.ActionRecModelEnum;
|
||||
import cn.smartjavaai.cls.enums.ClsModelEnum;
|
||||
import cn.smartjavaai.common.config.ModelConfig;
|
||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||
import lombok.Data;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 分类模型参数配置
|
||||
*
|
||||
* @author dwj
|
||||
*/
|
||||
@Data
|
||||
public class ClsModelConfig extends ModelConfig {
|
||||
|
||||
/**
|
||||
* 模型
|
||||
*/
|
||||
private ClsModelEnum modelEnum;
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 模型路径
|
||||
*/
|
||||
private String modelPath;
|
||||
|
||||
|
||||
/**
|
||||
* 允许的分类列表
|
||||
*/
|
||||
private List<String> allowedClasses;
|
||||
|
||||
/**
|
||||
* 置信度阈值
|
||||
*/
|
||||
private float threshold = 0.3f;
|
||||
|
||||
/**
|
||||
* 检测结果数量
|
||||
*/
|
||||
private int topK;
|
||||
|
||||
|
||||
|
||||
public ClsModelConfig() {
|
||||
}
|
||||
|
||||
public ClsModelConfig(ClsModelEnum modelEnum, DeviceEnum device) {
|
||||
this.modelEnum = modelEnum;
|
||||
setDevice(device);
|
||||
}
|
||||
|
||||
public ClsModelConfig(ClsModelEnum modelEnum) {
|
||||
this.modelEnum = modelEnum;
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package cn.smartjavaai.cls.criteria;
|
||||
|
||||
import ai.djl.Device;
|
||||
import ai.djl.modality.Classifications;
|
||||
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.action.config.ActionRecModelConfig;
|
||||
import cn.smartjavaai.action.enums.ActionRecModelEnum;
|
||||
import cn.smartjavaai.action.exception.ActionException;
|
||||
import cn.smartjavaai.action.model.CommonActionTranslator;
|
||||
import cn.smartjavaai.cls.config.ClsModelConfig;
|
||||
import cn.smartjavaai.cls.enums.ClsModelEnum;
|
||||
import cn.smartjavaai.cls.translator.YoloClsTranslator;
|
||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||
import cn.smartjavaai.common.utils.DJLCommonUtils;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
|
||||
import java.nio.file.Paths;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* 分类模型Criteria工厂
|
||||
* @author dwj
|
||||
*/
|
||||
public class ClsCriteriaFactory {
|
||||
|
||||
|
||||
/**
|
||||
* 创建动作识别Criteria
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public static Criteria<Image, Classifications> createCriteria(ClsModelConfig config) {
|
||||
Device device = null;
|
||||
if(!Objects.isNull(config.getDevice())){
|
||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu(config.getGpuId());
|
||||
}
|
||||
Translator<Image, Classifications> translator = getTranslator(config);
|
||||
//检查模型路径
|
||||
if (StringUtils.isBlank(config.getModelPath())){
|
||||
throw new ActionException("请指定模型路径");
|
||||
}
|
||||
boolean isUrl = DJLCommonUtils.hasSupportedProtocol(config.getModelPath());
|
||||
Criteria<Image, Classifications> criteria =
|
||||
Criteria.builder()
|
||||
.setTypes(Image.class, Classifications.class)
|
||||
.optModelUrls(isUrl ? config.getModelPath() : null)
|
||||
.optModelPath(isUrl ? null : 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, Classifications> getTranslator(ClsModelConfig config) {
|
||||
Translator<Image, Classifications> translator = null;
|
||||
if(config.getModelEnum() == ClsModelEnum.YOLOV11
|
||||
|| config.getModelEnum() == ClsModelEnum.YOLOV8){
|
||||
YoloClsTranslator.Builder builder = YoloClsTranslator.builder().optSynsetArtifactName("synset.txt");
|
||||
translator = builder.build();
|
||||
}
|
||||
return translator;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package cn.smartjavaai.cls.enums;
|
||||
|
||||
/**
|
||||
* 分类模型枚举
|
||||
* @author dwj
|
||||
*/
|
||||
public enum ClsModelEnum {
|
||||
|
||||
YOLOV8("OnnxRuntime",224,224),
|
||||
YOLOV11("OnnxRuntime",224,224);
|
||||
|
||||
/**
|
||||
* 模型输入尺寸：宽
|
||||
*/
|
||||
private final int inputWidth;
|
||||
|
||||
/**
|
||||
* 模型输入尺寸：高
|
||||
*/
|
||||
private final int inputHeight;
|
||||
|
||||
/**
|
||||
* 模型引擎
|
||||
*/
|
||||
private final String engine;
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static ClsModelEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (ClsModelEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
|
||||
|
||||
ClsModelEnum(String engine, int inputWidth, int inputHeight) {
|
||||
this.inputWidth = inputWidth;
|
||||
this.inputHeight = inputHeight;
|
||||
this.engine = engine;
|
||||
}
|
||||
|
||||
public int getInputWidth() {
|
||||
return inputWidth;
|
||||
}
|
||||
|
||||
public int getInputHeight() {
|
||||
return inputHeight;
|
||||
}
|
||||
|
||||
public String getEngine() {
|
||||
return engine;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package cn.smartjavaai.cls.exception;
|
||||
|
||||
/**
|
||||
* 分类模型异常
|
||||
* @author dwj
|
||||
*/
|
||||
public class ClsException extends RuntimeException{
|
||||
|
||||
public ClsException() {
|
||||
super();
|
||||
}
|
||||
|
||||
public ClsException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
|
||||
super(message, cause, enableSuppression, writableStackTrace);
|
||||
}
|
||||
|
||||
public ClsException(String message, Throwable cause) {
|
||||
super(message, cause);
|
||||
}
|
||||
|
||||
public ClsException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
public ClsException(Throwable cause) {
|
||||
super(cause);
|
||||
}
|
||||
|
||||
}
|
||||
45
vision/src/main/java/cn/smartjavaai/cls/model/ClsModel.java
Normal file
45
vision/src/main/java/cn/smartjavaai/cls/model/ClsModel.java
Normal file
@@ -0,0 +1,45 @@
|
||||
package cn.smartjavaai.cls.model;
|
||||
|
||||
import ai.djl.modality.Classifications;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import cn.smartjavaai.action.config.ActionRecModelConfig;
|
||||
import cn.smartjavaai.cls.config.ClsModelConfig;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
|
||||
/**
|
||||
* 图像分类模型
|
||||
* @author dwj
|
||||
*/
|
||||
public interface ClsModel extends AutoCloseable{
|
||||
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(ClsModelConfig config);
|
||||
|
||||
/**
|
||||
* 分类
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
default R<Classifications> detect(Image image){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
/**
|
||||
* 分类
|
||||
* @param imagePath
|
||||
* @return
|
||||
*/
|
||||
default R<Classifications> detect(String imagePath){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
default void setFromFactory(boolean fromFactory){
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package cn.smartjavaai.cls.model;
|
||||
|
||||
import cn.smartjavaai.action.model.CommonActionRecModel;
|
||||
import cn.smartjavaai.cls.config.ClsModelConfig;
|
||||
import cn.smartjavaai.cls.enums.ClsModelEnum;
|
||||
import cn.smartjavaai.common.config.Config;
|
||||
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 ClsModelFactory {
|
||||
|
||||
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
||||
private static volatile ClsModelFactory instance;
|
||||
|
||||
private static final ConcurrentHashMap<ClsModelEnum, ClsModel> modelMap = new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 模型注册表
|
||||
*/
|
||||
private static final Map<ClsModelEnum, Class<? extends ClsModel>> registry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
// 私有构造函数，防止外部创建实例
|
||||
private ClsModelFactory() {}
|
||||
|
||||
// 双重检查锁定的单例方法
|
||||
public static ClsModelFactory getInstance() {
|
||||
if (instance == null) {
|
||||
synchronized (ClsModelFactory.class) {
|
||||
if (instance == null) {
|
||||
instance = new ClsModelFactory();
|
||||
}
|
||||
}
|
||||
}
|
||||
return instance;
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取模型（通过配置）
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public ClsModel getModel(ClsModelConfig 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 ClsModel createFaceDetModel(ClsModelConfig config) {
|
||||
Class<?> clazz = registry.get(config.getModelEnum());
|
||||
if(clazz == null){
|
||||
throw new DetectionException("Unsupported model");
|
||||
}
|
||||
ClsModel model = null;
|
||||
try {
|
||||
model = (ClsModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new DetectionException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
model.setFromFactory(true);
|
||||
return model;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 注册模型
|
||||
* @param modelEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerAlgorithm(ClsModelEnum modelEnum, Class<? extends ClsModel> clazz) {
|
||||
registry.put(modelEnum, clazz);
|
||||
}
|
||||
|
||||
|
||||
|
||||
// 初始化默认算法
|
||||
static {
|
||||
registerAlgorithm(ClsModelEnum.YOLOV8, CommonClsModel.class);
|
||||
registerAlgorithm(ClsModelEnum.YOLOV11, CommonClsModel.class);
|
||||
log.debug("缓存目录：{}", Config.getCachePath());
|
||||
}
|
||||
|
||||
/**
|
||||
* 关闭所有已加载的模型
|
||||
*/
|
||||
public void closeAll() {
|
||||
modelMap.values().forEach(model -> {
|
||||
try {
|
||||
model.close();
|
||||
} catch (Exception e) {
|
||||
e.printStackTrace();
|
||||
}
|
||||
});
|
||||
modelMap.clear();
|
||||
}
|
||||
|
||||
/**
|
||||
* 移除缓存的模型
|
||||
* @param modelEnum
|
||||
*/
|
||||
public static void removeFromCache(ClsModelEnum modelEnum) {
|
||||
modelMap.remove(modelEnum);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
package cn.smartjavaai.cls.model;
|
||||
|
||||
import ai.djl.MalformedModelException;
|
||||
import ai.djl.engine.Engine;
|
||||
import ai.djl.inference.Predictor;
|
||||
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.repository.zoo.ModelNotFoundException;
|
||||
import ai.djl.repository.zoo.ZooModel;
|
||||
import cn.smartjavaai.action.criteria.ActionRecCriteriaFactory;
|
||||
import cn.smartjavaai.action.model.ActionRecModel;
|
||||
import cn.smartjavaai.action.model.ActionRecModelFactory;
|
||||
import cn.smartjavaai.cls.config.ClsModelConfig;
|
||||
import cn.smartjavaai.cls.criteria.ClsCriteriaFactory;
|
||||
import cn.smartjavaai.cls.exception.ClsException;
|
||||
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.FileUtils;
|
||||
import cn.smartjavaai.common.utils.ImageUtils;
|
||||
import cn.smartjavaai.common.utils.OpenCVUtils;
|
||||
import cn.smartjavaai.objectdetection.exception.DetectionException;
|
||||
import cn.smartjavaai.vision.utils.ClassificationFilter;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
import org.opencv.core.Mat;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.Objects;
|
||||
|
||||
/**
|
||||
* 通用图像分类模型
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class CommonClsModel implements ClsModel {
|
||||
|
||||
|
||||
private ClsModelConfig config;
|
||||
|
||||
private ZooModel<Image, Classifications> model;
|
||||
|
||||
private GenericObjectPool<Predictor<Image, Classifications>> predictorPool;
|
||||
|
||||
@Override
|
||||
public void loadModel(ClsModelConfig config) {
|
||||
if(Objects.isNull(config.getModelEnum())){
|
||||
throw new DetectionException("未配置模型枚举");
|
||||
}
|
||||
Criteria<Image, Classifications> criteria = ClsCriteriaFactory.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<Classifications> detect(Image image) {
|
||||
Classifications classifications = detectCore(image);
|
||||
// 过滤
|
||||
if(Objects.nonNull(classifications) && !classifications.items().isEmpty()){
|
||||
classifications = new ClassificationFilter(config.getAllowedClasses(), config.getThreshold()).filter(classifications);
|
||||
}
|
||||
return R.ok(classifications);
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public R<Classifications> detect(String imagePath) {
|
||||
if(!FileUtils.isFileExists(imagePath)){
|
||||
return R.fail(R.Status.FILE_NOT_FOUND);
|
||||
}
|
||||
Image img = null;
|
||||
try {
|
||||
img = SmartImageFactory.getInstance().fromFile(Paths.get(imagePath));
|
||||
return detect(img);
|
||||
} catch (IOException e) {
|
||||
throw new ClsException("无效的图片", e);
|
||||
} finally {
|
||||
ImageUtils.releaseOpenCVMat(img);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 模型核心推理方法
|
||||
* @param image
|
||||
* @return
|
||||
*/
|
||||
public Classifications detectCore(Image image) {
|
||||
Predictor<Image, Classifications> 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 {
|
||||
if (fromFactory) {
|
||||
// ActionRecModelFactory.removeFromCache(config.getModelEnum());
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
private boolean fromFactory = false;
|
||||
|
||||
@Override
|
||||
public void setFromFactory(boolean fromFactory) {
|
||||
this.fromFactory = fromFactory;
|
||||
}
|
||||
public boolean isFromFactory() {
|
||||
return fromFactory;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,386 @@
|
||||
package cn.smartjavaai.cls.translator;
|
||||
|
||||
import ai.djl.Model;
|
||||
import ai.djl.modality.Classifications;
|
||||
import ai.djl.modality.cv.Image;
|
||||
import ai.djl.modality.cv.transform.*;
|
||||
import ai.djl.modality.cv.util.NDImageUtils;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.translate.*;
|
||||
import ai.djl.util.Utils;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.net.MalformedURLException;
|
||||
import java.net.URL;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
|
||||
|
||||
public class YoloClsTranslator implements Translator<Image, Classifications> {
|
||||
|
||||
|
||||
protected float threshold;
|
||||
protected List<String> classes;
|
||||
protected boolean applyRatio;
|
||||
protected Pipeline pipeline;
|
||||
private Image.Flag flag;
|
||||
private Batchifier batchifier;
|
||||
protected int width;
|
||||
protected int height;
|
||||
protected int topk;
|
||||
|
||||
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 YoloClsTranslator(Builder builder) {
|
||||
this.threshold = builder.threshold;
|
||||
this.synsetLoader = builder.synsetLoader;
|
||||
this.applyRatio = builder.applyRatio;
|
||||
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) throws Exception {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
|
||||
//中心裁剪
|
||||
array = NDImageUtils.centerCrop(array);
|
||||
array = NDImageUtils.resize(array, width, height);
|
||||
// 转为 float32 且归一化到 0~1
|
||||
array = array.toType(DataType.FLOAT32, false).div(255f); // HWC
|
||||
// HWC -> CHW
|
||||
array = array.transpose(2, 0, 1); // CHW
|
||||
return new NDList(array);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Classifications processOutput(TranslatorContext ctx, NDList list) throws Exception {
|
||||
NDArray probabilitiesNd = list.singletonOrThrow();
|
||||
// probabilitiesNd = probabilitiesNd.softmax(0);
|
||||
return new Classifications(classes, probabilitiesNd, 5);
|
||||
}
|
||||
|
||||
|
||||
|
||||
public static class Builder {
|
||||
|
||||
|
||||
protected float threshold = 0.2F;
|
||||
protected boolean applyRatio;
|
||||
protected boolean removePadding;
|
||||
|
||||
protected int width = 224;
|
||||
protected int height = 224;
|
||||
protected Image.Flag flag;
|
||||
protected Pipeline pipeline;
|
||||
protected Batchifier batchifier;
|
||||
protected int topk = 5;
|
||||
|
||||
protected SynsetLoader synsetLoader;
|
||||
|
||||
public Builder() {
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds the translator.
|
||||
*
|
||||
* @return the new translator
|
||||
*/
|
||||
public YoloClsTranslator build() {
|
||||
if (pipeline == null) {
|
||||
addTransform(
|
||||
array -> array.transpose(2, 0, 1).toType(DataType.FLOAT32, false).div(255));
|
||||
}
|
||||
// validate();
|
||||
return new YoloClsTranslator(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();
|
||||
}
|
||||
|
||||
public Builder optTopk(int topk) {
|
||||
this.topk = topk;
|
||||
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.2F);
|
||||
String centerFit = ArgumentsUtil.stringValue(arguments, "centerFit", "false");
|
||||
this.removePadding = "true".equals(centerFit);
|
||||
String type = ArgumentsUtil.stringValue(arguments, "outputType", "AUTO");
|
||||
}
|
||||
|
||||
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"));
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user