package com.xindao.ocr.smartjavaai.factory;
|
|
import cn.smartjavaai.common.config.Config;
|
import com.xindao.ocr.smartjavaai.config.PlateDetModelConfig;
|
import com.xindao.ocr.smartjavaai.config.PlateRecModelConfig;
|
import com.xindao.ocr.smartjavaai.enums.PlateDetModelEnum;
|
import com.xindao.ocr.smartjavaai.enums.PlateRecModelEnum;
|
import com.xindao.ocr.smartjavaai.exception.OcrException;
|
import com.xindao.ocr.smartjavaai.model.plate.CRNNPlateRecModel;
|
import com.xindao.ocr.smartjavaai.model.plate.PlateDetModel;
|
import com.xindao.ocr.smartjavaai.model.plate.PlateRecModel;
|
import com.xindao.ocr.smartjavaai.model.plate.Yolov5PlateDetModel;
|
import lombok.extern.slf4j.Slf4j;
|
|
import java.util.Map;
|
import java.util.Objects;
|
import java.util.concurrent.ConcurrentHashMap;
|
|
/**
|
* 车牌识别模型工厂
|
* @author dwj
|
*/
|
@Slf4j
|
public class PlateModelFactory {
|
|
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
private static volatile PlateModelFactory instance;
|
|
/**
|
* 模型缓存
|
*/
|
private static final ConcurrentHashMap<PlateDetModelEnum, PlateDetModel> detModelMap = new ConcurrentHashMap<>();
|
|
/**
|
* 模型缓存
|
*/
|
private static final ConcurrentHashMap<PlateRecModelEnum, PlateRecModel> recModelMap = new ConcurrentHashMap<>();
|
|
|
/**
|
* 模型注册表
|
*/
|
private static final Map<PlateDetModelEnum, Class<? extends PlateDetModel>> detModelRegistry =
|
new ConcurrentHashMap<>();
|
|
/**
|
* 模型注册表
|
*/
|
private static final Map<PlateRecModelEnum, Class<? extends PlateRecModel>> recModelRegistry =
|
new ConcurrentHashMap<>();
|
|
|
public static PlateModelFactory getInstance() {
|
if (instance == null) {
|
synchronized (PlateModelFactory.class) {
|
if (instance == null) {
|
instance = new PlateModelFactory();
|
}
|
}
|
}
|
return instance;
|
}
|
|
|
|
/**
|
* 注册模型
|
* @param plateDetModelEnum
|
* @param clazz
|
*/
|
private static void registerDetModel(PlateDetModelEnum plateDetModelEnum, Class<? extends PlateDetModel> clazz) {
|
detModelRegistry.put(plateDetModelEnum, clazz);
|
}
|
|
/**
|
* 注册模型
|
* @param plateRecModelEnum
|
* @param clazz
|
*/
|
private static void registerRecModel(PlateRecModelEnum plateRecModelEnum, Class<? extends PlateRecModel> clazz) {
|
recModelRegistry.put(plateRecModelEnum, clazz);
|
}
|
|
|
/**
|
* 获取模型
|
* @param config
|
* @return
|
*/
|
public PlateDetModel getDetModel(PlateDetModelConfig config) {
|
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
throw new OcrException("未配置OCR模型");
|
}
|
return detModelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
return createDetModel(config);
|
});
|
}
|
|
/**
|
* 获取模型
|
* @param config
|
* @return
|
*/
|
public PlateRecModel getRecModel(PlateRecModelConfig config) {
|
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
throw new OcrException("未配置OCR模型");
|
}
|
return recModelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
return createRecModel(config);
|
});
|
}
|
|
|
|
/**
|
* 创建检测模型
|
* @param config
|
* @return
|
*/
|
private PlateDetModel createDetModel(PlateDetModelConfig config) {
|
Class<?> clazz = detModelRegistry.get(config.getModelEnum());
|
if(clazz == null){
|
throw new OcrException("Unsupported model");
|
}
|
PlateDetModel model = null;
|
try {
|
model = (PlateDetModel) clazz.newInstance();
|
} catch (InstantiationException | IllegalAccessException e) {
|
throw new OcrException(e);
|
}
|
model.loadModel(config);
|
return model;
|
}
|
|
/**
|
* 创建识别模型
|
* @param config
|
* @return
|
*/
|
private PlateRecModel createRecModel(PlateRecModelConfig config) {
|
Class<?> clazz = recModelRegistry.get(config.getModelEnum());
|
if(clazz == null){
|
throw new OcrException("Unsupported model");
|
}
|
PlateRecModel model = null;
|
try {
|
model = (PlateRecModel) clazz.newInstance();
|
} catch (InstantiationException | IllegalAccessException e) {
|
throw new OcrException(e);
|
}
|
model.loadModel(config);
|
return model;
|
}
|
|
|
// 初始化默认算法
|
static {
|
registerDetModel(PlateDetModelEnum.YOLOV5, Yolov5PlateDetModel.class);
|
registerDetModel(PlateDetModelEnum.YOLOV7, Yolov5PlateDetModel.class);
|
registerRecModel(PlateRecModelEnum.PLATE_REC_CRNN, CRNNPlateRecModel.class);
|
log.debug("缓存目录：{}", Config.getCachePath());
|
}
|
|
}
|
