package com.xindao.ocr.smartjavaai.model.common.direction.translator;
|
|
import ai.djl.modality.cv.Image;
|
import ai.djl.modality.cv.util.NDImageUtils;
|
import ai.djl.ndarray.NDArray;
|
import ai.djl.ndarray.NDList;
|
import ai.djl.ndarray.index.NDIndex;
|
import ai.djl.ndarray.types.Shape;
|
import ai.djl.translate.Batchifier;
|
import ai.djl.translate.Translator;
|
import ai.djl.translate.TranslatorContext;
|
import com.xindao.ocr.smartjavaai.entity.DirectionInfo;
|
|
import java.util.Arrays;
|
import java.util.List;
|
import java.util.Map;
|
|
/**
|
* 方向检测
|
*
|
* @author Calvin
|
* @mail 179209347@qq.com
|
* @website www.aias.top
|
*/
|
public class PpWordRotateTranslator implements Translator<Image, DirectionInfo> {
|
List<String> classes = Arrays.asList("No Rotate", "Rotate");
|
|
private String batchifier;
|
|
private int resizeHeight;
|
|
private int resizeWidth;
|
|
public PpWordRotateTranslator(Map<String, ?> arguments) {
|
batchifier = arguments.containsKey("batchifier")
|
? arguments.get("batchifier").toString()
|
: "padding";
|
|
resizeWidth = arguments.containsKey("resizeWidth")
|
? (Integer) arguments.get("resizeWidth")
|
: 192;
|
|
resizeHeight = arguments.containsKey("resizeHeight")
|
? (Integer) arguments.get("resizeHeight")
|
: 48;
|
}
|
|
@Override
|
public DirectionInfo processOutput(TranslatorContext ctx, NDList list) {
|
NDArray prob = list.singletonOrThrow();
|
float[] res = prob.toFloatArray();
|
int maxIndex = 0;
|
if (res[1] > res[0]) {
|
maxIndex = 1;
|
}
|
|
return new DirectionInfo(classes.get(maxIndex), Double.valueOf(res[maxIndex]));
|
}
|
|
// public NDList processInput2(TranslatorContext ctx, Image input){
|
// NDArray img = input.toNDArray(ctx.getNDManager());
|
// img = NDImageUtils.resize(img, 192, 48);
|
// img = NDImageUtils.toTensor(img).sub(0.5F).div(0.5F);
|
// img = img.expandDims(0);
|
// return new NDList(new NDArray[]{img});
|
// }
|
|
@Override
|
public NDList processInput(TranslatorContext ctx, Image input) {
|
NDArray img = input.toNDArray(ctx.getNDManager());
|
int imgC = 3;
|
int imgH = resizeHeight;
|
int imgW = resizeWidth;
|
|
NDArray array = ctx.getNDManager().zeros(new Shape(imgC, imgH, imgW));
|
|
int h = input.getHeight();
|
int w = input.getWidth();
|
int resized_w = 0;
|
|
float ratio = (float) w / (float) h;
|
if (Math.ceil(imgH * ratio) > imgW) {
|
resized_w = imgW;
|
} else {
|
resized_w = (int) (Math.ceil(imgH * ratio));
|
}
|
|
img = NDImageUtils.resize(img, resized_w, imgH);
|
|
img = NDImageUtils.toTensor(img).sub(0.5F).div(0.5F);
|
// img = img.transpose(2, 0, 1);
|
|
array.set(new NDIndex(":,:,0:" + resized_w), img);
|
|
// array = array.expandDims(0);
|
|
return new NDList(new NDArray[]{array});
|
}
|
|
@Override
|
public Batchifier getBatchifier() {
|
return Batchifier.fromString(batchifier);
|
}
|
|
}
|
