package com.xindao.ocr.smartjavaai.utils;
|
|
import ai.djl.modality.cv.Image;
|
import ai.djl.modality.cv.ImageFactory;
|
import ai.djl.modality.cv.output.BoundingBox;
|
import ai.djl.modality.cv.output.DetectedObjects;
|
import ai.djl.modality.cv.output.Landmark;
|
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.opencv.OpenCVImageFactory;
|
import cn.smartjavaai.common.entity.DetectionRectangle;
|
import cn.smartjavaai.common.entity.Point;
|
import cn.smartjavaai.common.utils.ImageUtils;
|
import cn.smartjavaai.common.utils.OpenCVUtils;
|
import cn.smartjavaai.common.utils.PointUtils;
|
import com.xindao.ocr.smartjavaai.entity.*;
|
import com.xindao.ocr.smartjavaai.enums.AngleEnum;
|
import com.xindao.ocr.smartjavaai.enums.PlateType;
|
import com.xindao.ocr.smartjavaai.opencv.OcrNDArrayUtils;
|
import lombok.extern.slf4j.Slf4j;
|
import org.apache.commons.collections.CollectionUtils;
|
import org.opencv.core.Mat;
|
import org.opencv.core.Scalar;
|
import org.opencv.imgproc.Imgproc;
|
|
import java.awt.*;
|
import java.awt.image.BufferedImage;
|
import java.util.*;
|
import java.util.List;
|
|
/**
|
* @author dwj
|
* @date 2025/4/22
|
*/
|
@Slf4j
|
public class OcrUtils {
|
|
|
/**
|
* 转换为OcrBox
|
* @param dt_boxes
|
* @return
|
*/
|
public static List<OcrBox> convertToOcrBox(NDList dt_boxes) {
|
List<OcrBox> boxList = new ArrayList<>();
|
for (NDArray box : dt_boxes) {
|
float[] pointsArr = box.toFloatArray();
|
OcrBox ocrBox = new OcrBox(
|
new Point(pointsArr[0], pointsArr[1]),
|
new Point(pointsArr[2], pointsArr[3]),
|
new Point(pointsArr[4], pointsArr[5]),
|
new Point(pointsArr[6], pointsArr[7])
|
);
|
boxList.add(ocrBox);
|
}
|
return boxList;
|
}
|
|
/**
|
* 转换为OcrBox
|
* @param ndLists
|
* @return
|
*/
|
public static List<List<OcrBox>> convertToOcrBox(List<NDList> ndLists) {
|
if (ndLists == null || ndLists.isEmpty()) {
|
return Collections.emptyList();
|
}
|
List<List<OcrBox>> boxLists = new ArrayList<>();
|
for (NDList dt_boxes : ndLists) {
|
boxLists.add(convertToOcrBox(dt_boxes));
|
}
|
return boxLists;
|
}
|
|
|
|
|
/**
|
* 图片旋转
|
*
|
* @param manager
|
* @param image
|
* @return
|
*/
|
public static Image rotateImg(NDManager manager, Image image) {
|
NDArray rotated = NDImageUtils.rotate90(image.toNDArray(manager), 1);
|
return ImageFactory.getInstance().fromNDArray(rotated);
|
}
|
|
/**
|
* 逆时针旋转图片
|
*
|
* @param image
|
* @param times
|
* @return
|
*/
|
public static Image rotateImg(Image image, int times) {
|
try (NDManager manager = NDManager.newBaseManager()) {
|
NDArray rotated = NDImageUtils.rotate90(image.toNDArray(manager), times);
|
return OpenCVImageFactory.getInstance().fromNDArray(rotated);
|
}
|
}
|
|
|
/**
|
* 逆时针旋转图片
|
*
|
* @param image
|
* @param angleEnum
|
* @return
|
*/
|
public static Image rotateImg(Image image, AngleEnum angleEnum) {
|
try (NDManager manager = NDManager.newBaseManager()) {
|
int times = 0;
|
switch (angleEnum) {
|
case ANGLE_90:
|
times = 1;
|
break;
|
case ANGLE_180:
|
times = 2;
|
break;
|
case ANGLE_270:
|
times = 3;
|
break;
|
}
|
NDArray rotated = NDImageUtils.rotate90(image.toNDArray(manager), times);
|
return OpenCVImageFactory.getInstance().fromNDArray(rotated);
|
}
|
}
|
|
|
/**
|
* 转换为OcrInfo
|
* @param lines
|
* @return
|
*/
|
public static OcrInfo convertToOcrInfo(List<ArrayList<RotatedBoxCompX>> lines){
|
if(Objects.isNull(lines) || lines.size() == 0){
|
return null;
|
}
|
List<List<OcrItem>> lineList = new ArrayList<List<OcrItem>>();
|
String fullText = "";
|
for(ArrayList<RotatedBoxCompX> boxList : lines){
|
List<OcrItem> line = new ArrayList<OcrItem>();
|
for(RotatedBoxCompX box : boxList){
|
float[] pointsArr = box.getBox().toFloatArray();
|
float[] lt = Arrays.copyOfRange(pointsArr, 0, 2);
|
float[] rt = Arrays.copyOfRange(pointsArr, 2, 4);
|
float[] rb = Arrays.copyOfRange(pointsArr, 4, 6);
|
float[] lb = Arrays.copyOfRange(pointsArr, 6, 8);
|
OcrBox ocrBox = new OcrBox(new Point(lt[0], lt[1]), new Point(rt[0], rt[1]), new Point(rb[0], rb[1]), new Point(lb[0], lb[1]));
|
OcrItem ocrItem = new OcrItem(ocrBox, box.getText());
|
line.add(ocrItem);
|
String text = box.getText();
|
if(text.trim().equals(""))
|
continue;
|
fullText += text + " ";
|
}
|
lineList.add(line);
|
fullText += '\n';
|
}
|
return new OcrInfo(lineList, fullText);
|
}
|
|
public static OcrInfo convertRotatedBoxesToOcrItems(List<RotatedBox> rotatedBoxes) {
|
OcrInfo ocrInfo = new OcrInfo();
|
List<OcrItem> ocrItems = new ArrayList<>();
|
StringBuilder fullText = new StringBuilder();
|
for (RotatedBox rotatedBox : rotatedBoxes) {
|
NDArray box = rotatedBox.getBox();
|
float[] points = box.toFloatArray();
|
Point topLeft = new Point(points[0], points[1]);
|
Point topRight = new Point(points[2], points[3]);
|
Point bottomRight = new Point(points[4], points[5]);
|
Point bottomLeft = new Point(points[6], points[7]);
|
|
OcrBox ocrBox = new OcrBox(topLeft, topRight, bottomRight, bottomLeft);
|
String text = rotatedBox.getText();
|
|
OcrItem item = new OcrItem();
|
item.setOcrBox(ocrBox);
|
item.setText(text);
|
ocrItems.add(item);
|
fullText.append(text + " ");
|
}
|
if (fullText.length() > 0) {
|
fullText.deleteCharAt(fullText.length() - 1);
|
}
|
ocrInfo.setOcrItemList(ocrItems);
|
ocrInfo.setFullText(fullText.toString());
|
return ocrInfo;
|
}
|
|
|
/**
|
* 透视变换 + 裁剪
|
* @param srcMat
|
* @param landMarks
|
* @return
|
*/
|
public static Image transformAndCrop(Mat srcMat, List<ai.djl.modality.cv.output.Point> landMarks){
|
if (landMarks == null || landMarks.size() != 4) {
|
throw new IllegalArgumentException("必须提供4个关键点");
|
}
|
|
// 步骤 1：排序为 左上、右上、右下、左下
|
List<ai.djl.modality.cv.output.Point> ordered = PointUtils.orderPoints(landMarks);
|
|
ai.djl.modality.cv.output.Point lt = ordered.get(0);
|
ai.djl.modality.cv.output.Point rt = ordered.get(1);
|
ai.djl.modality.cv.output.Point rb = ordered.get(2);
|
ai.djl.modality.cv.output.Point lb = ordered.get(3);
|
|
// 步骤 2：计算目标图像尺寸（宽、高）
|
int img_crop_width = (int) Math.max(
|
PointUtils.distance(lt, rt),
|
PointUtils.distance(rb, lb)
|
);
|
int img_crop_height = (int) Math.max(
|
PointUtils.distance(lt, lb),
|
PointUtils.distance(rt, rb)
|
);
|
|
// 步骤 3：构造目标坐标点
|
List<ai.djl.modality.cv.output.Point> dstPoints = Arrays.asList(
|
new ai.djl.modality.cv.output.Point(0, 0),
|
new ai.djl.modality.cv.output.Point(img_crop_width, 0),
|
new ai.djl.modality.cv.output.Point(img_crop_width, img_crop_height),
|
new ai.djl.modality.cv.output.Point(0, img_crop_height)
|
);
|
|
// 步骤 4：透视变换
|
Mat srcPoint2f = OcrNDArrayUtils.toMat(ordered);
|
Mat dstPoint2f = OcrNDArrayUtils.toMat(dstPoints);
|
Mat cvMat = OpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
|
|
// 步骤 5：转为 DJL Image + 裁剪
|
Image subImg = OpenCVImageFactory.getInstance().fromImage(cvMat);
|
subImg = subImg.getSubImage(0, 0, img_crop_width, img_crop_height);
|
|
// 释放资源
|
cvMat.release();
|
srcPoint2f.release();
|
dstPoint2f.release();
|
|
return subImg;
|
}
|
|
|
|
/**
|
* 透视变换+裁剪
|
* @param srcMat
|
* @param box
|
* @return
|
*/
|
public static Image transformAndCrop(Mat srcMat, OcrBox box){
|
float[] pointsArr = box.toFloatArray();
|
float[] lt = Arrays.copyOfRange(pointsArr, 0, 2);
|
float[] rt = Arrays.copyOfRange(pointsArr, 2, 4);
|
float[] rb = Arrays.copyOfRange(pointsArr, 4, 6);
|
float[] lb = Arrays.copyOfRange(pointsArr, 6, 8);
|
int img_crop_width = (int) Math.max(PointUtils.distance(lt, rt), PointUtils.distance(rb, lb));
|
int img_crop_height = (int) Math.max(PointUtils.distance(lt, lb), PointUtils.distance(rt, rb));
|
List<ai.djl.modality.cv.output.Point> srcPoints = new ArrayList<>();
|
srcPoints.add(new ai.djl.modality.cv.output.Point(lt[0], lt[1]));
|
srcPoints.add(new ai.djl.modality.cv.output.Point(rt[0], rt[1]));
|
srcPoints.add(new ai.djl.modality.cv.output.Point(rb[0], rb[1]));
|
srcPoints.add(new ai.djl.modality.cv.output.Point(lb[0], lb[1]));
|
List<ai.djl.modality.cv.output.Point> dstPoints = new ArrayList<>();
|
dstPoints.add(new ai.djl.modality.cv.output.Point(0, 0));
|
dstPoints.add(new ai.djl.modality.cv.output.Point(img_crop_width, 0));
|
dstPoints.add(new ai.djl.modality.cv.output.Point(img_crop_width, img_crop_height));
|
dstPoints.add(new ai.djl.modality.cv.output.Point(0, img_crop_height));
|
Mat srcPoint2f = OcrNDArrayUtils.toMat(srcPoints);
|
Mat dstPoint2f = OcrNDArrayUtils.toMat(dstPoints);
|
//透视变换
|
Mat cvMat = OpenCVUtils.perspectiveTransform(srcMat, srcPoint2f, dstPoint2f);
|
Image subImg = OpenCVImageFactory.getInstance().fromImage(cvMat);
|
//ImageUtils.saveImage(subImg, i + ".png", "build/output");
|
//变换后裁剪
|
subImg = subImg.getSubImage(0, 0, img_crop_width, img_crop_height);
|
cvMat.release();
|
srcPoint2f.release();
|
dstPoint2f.release();
|
return subImg;
|
}
|
|
/**
|
* 绘制文本框
|
*
|
* @param mat
|
* @param boxList
|
*/
|
public static void drawRect(Mat mat, List<OcrBox> boxList) {
|
for(OcrBox ocrBox : boxList){
|
Imgproc.line(mat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 255, 0), 1);
|
Imgproc.line(mat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 255, 0),1);
|
Imgproc.line(mat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 255, 0),1);
|
Imgproc.line(mat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 255, 0), 1);
|
}
|
}
|
|
|
/**
|
* 绘制文本框及文本
|
* @param image
|
* @param ocrInfo
|
*/
|
public static void drawRectWithText(BufferedImage image, OcrInfo ocrInfo, int fontSize) {
|
// 将绘制图像转换为Graphics2D
|
Graphics2D g = (Graphics2D) image.getGraphics();
|
try {
|
Font font = new Font("楷体", Font.PLAIN, fontSize);
|
g.setFont(font);
|
g.setColor(new Color(0, 0, 255));
|
// 声明画笔属性 ：粗 细（单位像素）末端无修饰 折线处呈尖角
|
BasicStroke bStroke = new BasicStroke(2, BasicStroke.CAP_BUTT, BasicStroke.JOIN_MITER);
|
g.setStroke(bStroke);
|
List<OcrItem> ocrItemList = ocrInfo.getOcrItemList();
|
if(CollectionUtils.isNotEmpty(ocrInfo.getLineList())){
|
ocrItemList = ocrInfo.flattenLines();
|
}
|
for(OcrItem item : ocrItemList){
|
OcrBox box = item.getOcrBox();
|
int[] xPoints = {
|
(int)box.getTopLeft().getX(),
|
(int)box.getTopRight().getX(),
|
(int)box.getBottomRight().getX(),
|
(int)box.getBottomLeft().getX(),
|
(int)box.getTopLeft().getX()
|
};
|
int[] yPoints = {
|
(int)box.getTopLeft().getY(),
|
(int)box.getTopRight().getY(),
|
(int)box.getBottomRight().getY(),
|
(int)box.getBottomLeft().getY(),
|
(int)box.getTopLeft().getY()
|
};
|
g.drawPolyline(xPoints, yPoints, 5);
|
g.drawString(item.getText(), xPoints[0], yPoints[0]);
|
}
|
} finally {
|
g.dispose();
|
}
|
}
|
|
|
/**
|
* 绘制文本框及文本
|
* @param srcMat
|
* @param itemList
|
*/
|
public static void drawRectWithText(Mat srcMat, List<OcrItem> itemList) {
|
for(OcrItem item : itemList){
|
OcrBox ocrBox = item.getOcrBox();
|
Imgproc.line(srcMat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 255, 0), 1);
|
Imgproc.line(srcMat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 255, 0),1);
|
Imgproc.line(srcMat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 255, 0),1);
|
Imgproc.line(srcMat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 255, 0), 1);
|
// 中文乱码
|
Imgproc.putText(srcMat, item.getAngle().getValue(), ocrBox.getTopLeft().toCvPoint(), Imgproc.FONT_HERSHEY_SCRIPT_SIMPLEX, 1.0, new Scalar(0, 255, 0), 1);
|
}
|
}
|
|
public static List<PlateInfo> convertToPlateInfo(DetectedObjects detectedObjects, Image image) {
|
List<PlateInfo> plateInfoList = new ArrayList<>();
|
Iterator iterator = detectedObjects.items().iterator();
|
int index = 0;
|
while(iterator.hasNext()) {
|
DetectedObjects.DetectedObject result = (DetectedObjects.DetectedObject)iterator.next();
|
BoundingBox box = result.getBoundingBox();
|
List<Point> keyPoints = new ArrayList<Point>();
|
if(box instanceof Landmark){
|
box.getBounds().getPath().forEach(point -> {
|
keyPoints.add(new Point(point.getX(), point.getY()));
|
});
|
}
|
int x = (int)(box.getBounds().getX() * image.getWidth());
|
int y = (int)(box.getBounds().getY() * image.getHeight());
|
int width = (int)(box.getBounds().getWidth() * image.getWidth());
|
int height = (int)(box.getBounds().getHeight() * image.getHeight());
|
// 修正边界，防止越界
|
if (x < 0) x = 0;
|
if (y < 0) y = 0;
|
if (x + width > image.getWidth()) width = image.getWidth() - x;
|
if (y + height > image.getHeight()) height = image.getHeight() - y;
|
PlateInfo plateInfo = new PlateInfo();
|
plateInfo.setPlateType(PlateType.fromClassName(detectedObjects.getClassNames().get(index)));
|
plateInfo.setScore(detectedObjects.getProbabilities().get(index).floatValue());
|
plateInfo.setDetectionRectangle(new DetectionRectangle(x, y, width, height));
|
OcrBox ocrBox = new OcrBox(keyPoints.get(0), keyPoints.get(1), keyPoints.get(2), keyPoints.get(3));
|
plateInfo.setBox(ocrBox);
|
plateInfoList.add(plateInfo);
|
index++;
|
}
|
return plateInfoList;
|
}
|
|
/**
|
* 绘制车牌信息
|
* @param srcMat
|
* @param plateInfoList
|
*/
|
public static void drawPlateInfo(Mat srcMat, List<PlateInfo> plateInfoList) {
|
for(PlateInfo plateInfo : plateInfoList){
|
OcrBox ocrBox = plateInfo.getBox();
|
Imgproc.line(srcMat, ocrBox.getTopLeft().toCvPoint(), ocrBox.getTopRight().toCvPoint(), new Scalar(0, 0, 255), 1);
|
Imgproc.line(srcMat, ocrBox.getTopRight().toCvPoint(), ocrBox.getBottomRight().toCvPoint(), new Scalar(0, 0, 255),1);
|
Imgproc.line(srcMat, ocrBox.getBottomRight().toCvPoint(), ocrBox.getBottomLeft().toCvPoint(), new Scalar(0, 0, 255),1);
|
Imgproc.line(srcMat, ocrBox.getBottomLeft().toCvPoint(), ocrBox.getTopLeft().toCvPoint(), new Scalar(0, 0, 255), 1);
|
// 中文乱码
|
ImageUtils.putTextWithBackground(srcMat, plateInfo.getPlateNumber() + " " + plateInfo.getPlateColor(), ocrBox.getTopLeft().toCvPoint(), new Scalar(255, 255, 255), new Scalar(0, 0, 0), 1);
|
}
|
}
|
|
/**
|
* 在图像上绘制带白色背景、黑色文字的文本
|
*/
|
public static void drawPlateInfo(BufferedImage image, List<PlateInfo> plateInfoList) {
|
// 将绘制图像转换为Graphics2D
|
Graphics2D graphics = (Graphics2D) image.getGraphics();
|
try {
|
graphics.setColor(Color.RED);// 边框颜色
|
graphics.setStroke(new BasicStroke(2)); // 线宽2像素
|
graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
|
RenderingHints.VALUE_ANTIALIAS_ON); // 抗锯齿
|
int stroke = 2;
|
for(PlateInfo plateInfo : plateInfoList){
|
DetectionRectangle rectangle = plateInfo.getDetectionRectangle();
|
graphics.setColor(Color.RED);// 边框颜色
|
//绘制车牌框
|
graphics.drawRect(rectangle.getX(), rectangle.getY(), rectangle.getWidth(), rectangle.getHeight());
|
graphics.setColor(Color.BLACK);// 字体颜色
|
ImageUtils.drawText(graphics, plateInfo.getPlateNumber() + " " + plateInfo.getPlateColor(), rectangle.getX(), rectangle.getY(), stroke, 4);
|
OcrBox ocrBox = plateInfo.getBox();
|
//绘制关键点
|
graphics.setColor(Color.BLUE);
|
graphics.drawRect((int)ocrBox.getTopLeft().getX(), (int)ocrBox.getTopLeft().getY(), 2, 2);
|
graphics.setColor(Color.GREEN);
|
graphics.drawRect((int)ocrBox.getTopRight().getX(), (int)ocrBox.getTopRight().getY(), 2, 2);
|
graphics.setColor(Color.RED);
|
graphics.drawRect((int)ocrBox.getBottomLeft().getX(), (int)ocrBox.getBottomLeft().getY(), 2, 2);
|
graphics.setColor(Color.CYAN);
|
graphics.drawRect((int)ocrBox.getBottomRight().getX(), (int)ocrBox.getBottomRight().getY(), 2, 2);
|
}
|
} finally {
|
graphics.dispose();
|
}
|
}
|
|
|
}
|
