package com.xindao.ocr.smartjavaai.model.common.detect.translator;
|
|
import ai.djl.modality.cv.Image;
|
import ai.djl.modality.cv.util.NDImageUtils;
|
import ai.djl.ndarray.NDArray;
|
import ai.djl.ndarray.NDArrays;
|
import ai.djl.ndarray.NDList;
|
import ai.djl.ndarray.NDManager;
|
import ai.djl.ndarray.index.NDIndex;
|
import ai.djl.ndarray.types.DataType;
|
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.opencv.OcrNDArrayUtils;
|
import org.opencv.core.*;
|
import org.opencv.imgproc.Imgproc;
|
|
import java.util.ArrayList;
|
import java.util.List;
|
import java.util.Map;
|
|
/**
|
* 文字检测前后处理
|
*
|
* @author Calvin
|
* @mail 179209347@qq.com
|
* @website www.aias.top
|
*/
|
public class PPOCRDetTranslator implements Translator<Image, NDList> {
|
// det_algorithm == "DB"
|
private final float thresh = 0.3f;
|
private final boolean use_dilation = false;
|
private final String score_mode = "fast";
|
private final String box_type = "quad";
|
|
//检测的图像边长限制
|
private final int limit_side_len;
|
//输出的最大文本框数量
|
private final int max_candidates;
|
//文本框最小尺寸阈值
|
private final int min_size;
|
//文本框的分数阈值
|
private final float box_thresh;
|
|
/**
|
* 这个参数是检测后处理时控制文本框大小的，默认1.6，可以尝试改成2.5或者更大，反之，如果觉得文本框不够紧凑，也可以把该参数调小。
|
* 检测框大小过于紧贴文字或检测框过大，可以调整db_unclip_ratio这个参数，加大参数可以扩大检测框，减小参数可以减小检测框大小；
|
*/
|
private final float unclip_ratio;
|
private float ratio_h;
|
private float ratio_w;
|
private int img_height;
|
private int img_width;
|
|
private String batchifier;
|
|
public PPOCRDetTranslator(Map<String, ?> arguments) {
|
limit_side_len =
|
arguments.containsKey("limit_side_len")
|
? Integer.parseInt(arguments.get("limit_side_len").toString())
|
: 960;
|
max_candidates =
|
arguments.containsKey("max_candidates")
|
? Integer.parseInt(arguments.get("max_candidates").toString())
|
: 1000;
|
min_size =
|
arguments.containsKey("min_size")
|
? Integer.parseInt(arguments.get("min_size").toString())
|
: 3;
|
box_thresh =
|
arguments.containsKey("box_thresh")
|
? Float.parseFloat(arguments.get("box_thresh").toString())
|
: 0.6f; // 0.5f
|
unclip_ratio =
|
arguments.containsKey("unclip_ratio")
|
? Float.parseFloat(arguments.get("unclip_ratio").toString())
|
: 1.6f;
|
|
batchifier = arguments.containsKey("batchifier")
|
? arguments.get("batchifier").toString()
|
: "stack";
|
}
|
|
@Override
|
public NDList processOutput(TranslatorContext ctx, NDList list) {
|
NDManager manager = ctx.getNDManager();
|
NDArray pred = list.get(0);
|
pred = pred.squeeze();
|
NDArray segmentation = pred.gt(thresh); // thresh=0.3 .mul(255f)
|
|
segmentation = segmentation.toType(DataType.UINT8, true);
|
Shape shape = segmentation.getShape();
|
int rows = (int) shape.get(0);
|
int cols = (int) shape.get(1);
|
|
Mat newMask = new Mat();
|
if (this.use_dilation) {
|
Mat mask = new Mat();
|
//convert from NDArray to Mat
|
Mat srcMat = OcrNDArrayUtils.uint8NDArrayToMat(segmentation);
|
// size 越小，腐蚀的单位越小，图片越接近原图
|
// Mat dilation_kernel = Imgproc.getStructuringElement(Imgproc.MORPH_RECT, new Size(2, 2));
|
Mat dilation_kernel = OcrNDArrayUtils.uint8ArrayToMat(new byte[][]{{1, 1}, {1, 1}});
|
/**
|
* 膨胀说明： 图像的一部分区域与指定的核进行卷积， 求核的最`大`值并赋值给指定区域。 膨胀可以理解为图像中`高亮区域`的'领域扩大'。
|
* 意思是高亮部分会侵蚀不是高亮的部分，使高亮部分越来越多。
|
*/
|
Imgproc.dilate(srcMat, mask, dilation_kernel);
|
//destination Matrix
|
Scalar scalar = new Scalar(255);
|
Core.multiply(mask, scalar, newMask);
|
// release Mat
|
mask.release();
|
srcMat.release();
|
dilation_kernel.release();
|
} else {
|
Mat srcMat = OcrNDArrayUtils.uint8NDArrayToMat(segmentation);
|
//destination Matrix
|
Scalar scalar = new Scalar(255);
|
Core.multiply(srcMat, scalar, newMask);
|
// release Mat
|
srcMat.release();
|
}
|
|
NDArray boxes = boxes_from_bitmap(manager, pred, newMask);
|
|
//boxes[:, :, 0] = boxes[:, :, 0] / ratio_w
|
NDArray boxes1 = boxes.get(":, :, 0").div(ratio_w);
|
boxes.set(new NDIndex(":, :, 0"), boxes1);
|
//boxes[:, :, 1] = boxes[:, :, 1] / ratio_h
|
NDArray boxes2 = boxes.get(":, :, 1").div(ratio_h);
|
boxes.set(new NDIndex(":, :, 1"), boxes2);
|
|
NDList dt_boxes = this.filter_tag_det_res(boxes);
|
|
dt_boxes.detach();
|
|
// release Mat
|
newMask.release();
|
|
return dt_boxes;
|
}
|
|
|
private NDList filter_tag_det_res(NDArray dt_boxes) {
|
NDList boxesList = new NDList();
|
|
int num = (int) dt_boxes.getShape().get(0);
|
for (int i = 0; i < num; i++) {
|
NDArray box = dt_boxes.get(i);
|
box = order_points_clockwise(box);
|
box = clip_det_res(box);
|
float[] box0 = box.get(0).toFloatArray();
|
float[] box1 = box.get(1).toFloatArray();
|
float[] box3 = box.get(3).toFloatArray();
|
int rect_width = (int) Math.sqrt(Math.pow(box1[0] - box0[0], 2) + Math.pow(box1[1] - box0[1], 2));
|
int rect_height = (int) Math.sqrt(Math.pow(box3[0] - box0[0], 2) + Math.pow(box3[1] - box0[1], 2));
|
if (rect_width <= 3 || rect_height <= 3)
|
continue;
|
boxesList.add(box);
|
}
|
|
return boxesList;
|
}
|
|
private NDArray clip_det_res(NDArray points) {
|
for (int i = 0; i < points.getShape().get(0); i++) {
|
int value = Math.max((int) points.get(i, 0).toFloatArray()[0], 0);
|
value = Math.min(value, img_width - 1);
|
points.set(new NDIndex(i + ",0"), value);
|
value = Math.max((int) points.get(i, 1).toFloatArray()[0], 0);
|
value = Math.min(value, img_height - 1);
|
points.set(new NDIndex(i + ",1"), value);
|
}
|
|
return points;
|
}
|
|
/**
|
* sort the points based on their x-coordinates
|
* 顺时针
|
*
|
* @param pts
|
* @return
|
*/
|
|
private NDArray order_points_clockwise(NDArray pts) {
|
NDList list = new NDList();
|
long[] indexes = pts.get(":, 0").argSort().toLongArray();
|
|
// grab the left-most and right-most points from the sorted
|
// x-roodinate points
|
Shape s1 = pts.getShape();
|
NDArray leftMost1 = pts.get(indexes[0] + ",:");
|
NDArray leftMost2 = pts.get(indexes[1] + ",:");
|
NDArray leftMost = leftMost1.concat(leftMost2).reshape(2, 2);
|
NDArray rightMost1 = pts.get(indexes[2] + ",:");
|
NDArray rightMost2 = pts.get(indexes[3] + ",:");
|
NDArray rightMost = rightMost1.concat(rightMost2).reshape(2, 2);
|
|
// now, sort the left-most coordinates according to their
|
// y-coordinates so we can grab the top-left and bottom-left
|
// points, respectively
|
indexes = leftMost.get(":, 1").argSort().toLongArray();
|
NDArray lt = leftMost.get(indexes[0] + ",:");
|
NDArray lb = leftMost.get(indexes[1] + ",:");
|
indexes = rightMost.get(":, 1").argSort().toLongArray();
|
NDArray rt = rightMost.get(indexes[0] + ",:");
|
NDArray rb = rightMost.get(indexes[1] + ",:");
|
|
list.add(lt);
|
list.add(rt);
|
list.add(rb);
|
list.add(lb);
|
|
NDArray rect = NDArrays.concat(list).reshape(4, 2);
|
return rect;
|
}
|
|
/**
|
* Get boxes from the binarized image predicted by DB
|
*
|
* @param manager
|
* @param pred the binarized image predicted by DB.
|
* @param bitmap new 'pred' after threshold filtering.
|
*/
|
private NDArray boxes_from_bitmap(NDManager manager, NDArray pred, Mat bitmap) {
|
int dest_height = (int) pred.getShape().get(0);
|
int dest_width = (int) pred.getShape().get(1);
|
int height = bitmap.rows();
|
int width = bitmap.cols();
|
|
List<MatOfPoint> contours = new ArrayList<>();
|
Mat hierarchy = new Mat();
|
// 寻找轮廓
|
Imgproc.findContours(
|
bitmap,
|
contours,
|
hierarchy,
|
Imgproc.RETR_LIST,
|
Imgproc.CHAIN_APPROX_SIMPLE);
|
|
int num_contours = Math.min(contours.size(), max_candidates);
|
NDList boxList = new NDList();
|
float[] scores = new float[num_contours];
|
|
for (int index = 0; index < num_contours; index++) {
|
MatOfPoint contour = contours.get(index);
|
MatOfPoint2f newContour = new MatOfPoint2f(contour.toArray());
|
float[][] pointsArr = new float[4][2];
|
int sside = get_mini_boxes(newContour, pointsArr);
|
if (sside < this.min_size)
|
continue;
|
NDArray points = manager.create(pointsArr);
|
float score = box_score_fast(manager, pred, points);
|
if (score < this.box_thresh)
|
continue;
|
|
NDArray box = unclip(manager, points); // TODO get_mini_boxes(box)
|
|
// box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
|
NDArray boxes1 = box.get(":,0").div(width).mul(dest_width).round().clip(0, dest_width);
|
box.set(new NDIndex(":, 0"), boxes1);
|
// box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
|
NDArray boxes2 = box.get(":,1").div(height).mul(dest_height).round().clip(0, dest_height);
|
box.set(new NDIndex(":, 1"), boxes2);
|
|
boxList.add(box);
|
scores[index] = score;
|
|
// release memory
|
contour.release();
|
newContour.release();
|
}
|
|
NDArray boxes = NDArrays.stack(boxList);
|
|
// release
|
hierarchy.release();
|
|
return boxes;
|
}
|
|
/**
|
* Shrink or expand the boxaccording to 'unclip_ratio'
|
*
|
* @param points The predicted box.
|
* @return uncliped box
|
*/
|
private NDArray unclip(NDManager manager, NDArray points) {
|
points = order_points_clockwise(points);
|
float[] pointsArr = points.toFloatArray();
|
float[] lt = java.util.Arrays.copyOfRange(pointsArr, 0, 2);
|
float[] lb = java.util.Arrays.copyOfRange(pointsArr, 6, 8);
|
|
float[] rt = java.util.Arrays.copyOfRange(pointsArr, 2, 4);
|
float[] rb = java.util.Arrays.copyOfRange(pointsArr, 4, 6);
|
|
float width = distance(lt, rt);
|
float height = distance(lt, lb);
|
|
if (width > height) {
|
float k = (lt[1] - rt[1]) / (lt[0] - rt[0]); // y = k * x + b
|
|
float delta_dis = height;
|
float delta_x = (float) Math.sqrt((delta_dis * delta_dis) / (k * k + 1));
|
float delta_y = Math.abs(k * delta_x);
|
|
if (k > 0) {
|
pointsArr[0] = lt[0] - delta_x + delta_y;
|
pointsArr[1] = lt[1] - delta_y - delta_x;
|
pointsArr[2] = rt[0] + delta_x + delta_y;
|
pointsArr[3] = rt[1] + delta_y - delta_x;
|
|
pointsArr[4] = rb[0] + delta_x - delta_y;
|
pointsArr[5] = rb[1] + delta_y + delta_x;
|
pointsArr[6] = lb[0] - delta_x - delta_y;
|
pointsArr[7] = lb[1] - delta_y + delta_x;
|
} else {
|
pointsArr[0] = lt[0] - delta_x - delta_y;
|
pointsArr[1] = lt[1] + delta_y - delta_x;
|
pointsArr[2] = rt[0] + delta_x - delta_y;
|
pointsArr[3] = rt[1] - delta_y - delta_x;
|
|
pointsArr[4] = rb[0] + delta_x + delta_y;
|
pointsArr[5] = rb[1] - delta_y + delta_x;
|
pointsArr[6] = lb[0] - delta_x + delta_y;
|
pointsArr[7] = lb[1] + delta_y + delta_x;
|
}
|
} else {
|
float k = (lt[1] - rt[1]) / (lt[0] - rt[0]); // y = k * x + b
|
|
float delta_dis = width;
|
float delta_y = (float) Math.sqrt((delta_dis * delta_dis) / (k * k + 1));
|
float delta_x = Math.abs(k * delta_y);
|
|
if (k > 0) {
|
pointsArr[0] = lt[0] + delta_x - delta_y;
|
pointsArr[1] = lt[1] - delta_y - delta_x;
|
pointsArr[2] = rt[0] + delta_x + delta_y;
|
pointsArr[3] = rt[1] - delta_y + delta_x;
|
|
pointsArr[4] = rb[0] - delta_x + delta_y;
|
pointsArr[5] = rb[1] + delta_y + delta_x;
|
pointsArr[6] = lb[0] - delta_x - delta_y;
|
pointsArr[7] = lb[1] + delta_y - delta_x;
|
} else {
|
pointsArr[0] = lt[0] - delta_x - delta_y;
|
pointsArr[1] = lt[1] - delta_y + delta_x;
|
pointsArr[2] = rt[0] - delta_x + delta_y;
|
pointsArr[3] = rt[1] - delta_y - delta_x;
|
|
pointsArr[4] = rb[0] + delta_x + delta_y;
|
pointsArr[5] = rb[1] + delta_y - delta_x;
|
pointsArr[6] = lb[0] + delta_x - delta_y;
|
pointsArr[7] = lb[1] + delta_y + delta_x;
|
}
|
}
|
points = manager.create(pointsArr).reshape(4, 2);
|
|
return points;
|
}
|
|
private float distance(float[] point1, float[] point2) {
|
float disX = point1[0] - point2[0];
|
float disY = point1[1] - point2[1];
|
float dis = (float) Math.sqrt(disX * disX + disY * disY);
|
return dis;
|
}
|
|
/**
|
* Get boxes from the contour or box.
|
*
|
* @param contour The predicted contour.
|
* @param pointsArr The predicted box.
|
* @return smaller side of box
|
*/
|
private int get_mini_boxes(MatOfPoint2f contour, float[][] pointsArr) {
|
// https://blog.csdn.net/qq_37385726/article/details/82313558
|
// bounding_box[1] - rect 返回矩形的长和宽
|
RotatedRect rect = Imgproc.minAreaRect(contour);
|
Mat points = new Mat();
|
Imgproc.boxPoints(rect, points);
|
|
float[][] fourPoints = new float[4][2];
|
for (int row = 0; row < 4; row++) {
|
fourPoints[row][0] = (float) points.get(row, 0)[0];
|
fourPoints[row][1] = (float) points.get(row, 1)[0];
|
}
|
|
float[] tmpPoint = new float[2];
|
for (int i = 0; i < 4; i++) {
|
for (int j = i + 1; j < 4; j++) {
|
if (fourPoints[j][0] < fourPoints[i][0]) {
|
tmpPoint[0] = fourPoints[i][0];
|
tmpPoint[1] = fourPoints[i][1];
|
fourPoints[i][0] = fourPoints[j][0];
|
fourPoints[i][1] = fourPoints[j][1];
|
fourPoints[j][0] = tmpPoint[0];
|
fourPoints[j][1] = tmpPoint[1];
|
}
|
}
|
}
|
|
int index_1 = 0;
|
int index_2 = 1;
|
int index_3 = 2;
|
int index_4 = 3;
|
|
if (fourPoints[1][1] > fourPoints[0][1]) {
|
index_1 = 0;
|
index_4 = 1;
|
} else {
|
index_1 = 1;
|
index_4 = 0;
|
}
|
|
if (fourPoints[3][1] > fourPoints[2][1]) {
|
index_2 = 2;
|
index_3 = 3;
|
} else {
|
index_2 = 3;
|
index_3 = 2;
|
}
|
|
pointsArr[0] = fourPoints[index_1];
|
pointsArr[1] = fourPoints[index_2];
|
pointsArr[2] = fourPoints[index_3];
|
pointsArr[3] = fourPoints[index_4];
|
|
int height = rect.boundingRect().height;
|
int width = rect.boundingRect().width;
|
int sside = Math.min(height, width);
|
|
// release
|
points.release();
|
|
return sside;
|
}
|
|
/**
|
* Calculate the score of box.
|
*
|
* @param bitmap The binarized image predicted by DB.
|
* @param points The predicted box
|
* @return
|
*/
|
private float box_score_fast(NDManager manager, NDArray bitmap, NDArray points) {
|
NDArray box = points.get(":");
|
long h = bitmap.getShape().get(0);
|
long w = bitmap.getShape().get(1);
|
// xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
|
int xmin = box.get(":, 0").min().floor().clip(0, w - 1).toType(DataType.INT32, true).toIntArray()[0];
|
int xmax = box.get(":, 0").max().ceil().clip(0, w - 1).toType(DataType.INT32, true).toIntArray()[0];
|
int ymin = box.get(":, 1").min().floor().clip(0, h - 1).toType(DataType.INT32, true).toIntArray()[0];
|
int ymax = box.get(":, 1").max().ceil().clip(0, h - 1).toType(DataType.INT32, true).toIntArray()[0];
|
|
NDArray mask = manager.zeros(new Shape(ymax - ymin + 1, xmax - xmin + 1), DataType.UINT8);
|
|
box.set(new NDIndex(":, 0"), box.get(":, 0").sub(xmin));
|
box.set(new NDIndex(":, 1"), box.get(":, 1").sub(ymin));
|
|
//mask - convert from NDArray to Mat
|
Mat maskMat = OcrNDArrayUtils.uint8NDArrayToMat(mask);
|
|
//mask - convert from NDArray to Mat - 4 rows, 2 cols
|
Mat boxMat = OcrNDArrayUtils.floatNDArrayToMat(box, CvType.CV_32S);
|
|
// boxMat.reshape(1, new int[]{1, 4, 2});
|
List<MatOfPoint> pts = new ArrayList<>();
|
MatOfPoint matOfPoint = OcrNDArrayUtils.matToMatOfPoint(boxMat); // new MatOfPoint(boxMat);
|
pts.add(matOfPoint);
|
Imgproc.fillPoly(maskMat, pts, new Scalar(1));
|
|
|
NDArray subBitMap = bitmap.get(ymin + ":" + (ymax + 1) + "," + xmin + ":" + (xmax + 1));
|
Mat bitMapMat = OcrNDArrayUtils.floatNDArrayToMat(subBitMap);
|
|
Scalar score = Core.mean(bitMapMat, maskMat);
|
float scoreValue = (float) score.val[0];
|
// release
|
maskMat.release();
|
boxMat.release();
|
bitMapMat.release();
|
|
return scoreValue;
|
}
|
|
@Override
|
public NDList processInput(TranslatorContext ctx, Image input) {
|
NDArray img = input.toNDArray(ctx.getNDManager());
|
int h = input.getHeight();
|
int w = input.getWidth();
|
img_height = h;
|
img_width = w;
|
|
// limit the max side
|
float ratio = 1.0f;
|
if (Math.max(h, w) > limit_side_len) {
|
if (h > w) {
|
ratio = (float) limit_side_len / (float) h;
|
} else {
|
ratio = (float) limit_side_len / (float) w;
|
}
|
}
|
|
int resize_h = (int) (h * ratio);
|
int resize_w = (int) (w * ratio);
|
|
resize_h = Math.round((float) resize_h / 32f) * 32;
|
resize_w = Math.round((float) resize_w / 32f) * 32;
|
|
ratio_h = resize_h / (float) h;
|
ratio_w = resize_w / (float) w;
|
|
img = NDImageUtils.resize(img, resize_w, resize_h);
|
|
img = NDImageUtils.toTensor(img);
|
|
img =
|
NDImageUtils.normalize(
|
img,
|
new float[]{0.485f, 0.456f, 0.406f},
|
new float[]{0.229f, 0.224f, 0.225f});
|
|
// img = img.expandDims(0);
|
|
return new NDList(img);
|
}
|
|
@Override
|
public Batchifier getBatchifier() {
|
return Batchifier.fromString(batchifier);
|
}
|
}
|
