Paddle Structure

package nexus.io.cv.paddle.utils;

import org.yaml.snakeyaml.Yaml;

import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class ModelLabelUtils {

  public static List<String> loadLabelsFromYaml(String yamlPath) throws IOException {
    Yaml yaml = new Yaml();
    try (InputStream in = Files.newInputStream(Paths.get(yamlPath))) {
      Map<String, Object> root = yaml.load(in);
      Object labelsObj = root.get("label_list");
      if (labelsObj instanceof List<?>) {
        return ((List<?>) labelsObj).stream().map(String::valueOf).collect(Collectors.toList());
      }
      return new ArrayList<>();
    }
  }

  public static List<String> loadLabels(Path yamlPath) throws IOException {
    if (!Files.exists(yamlPath)) {
      throw new IllegalArgumentException("inference.yml 不存在: " + yamlPath);
    }

    Yaml yaml = new Yaml();
    try (InputStream in = Files.newInputStream(yamlPath)) {
      Map<String, Object> root = yaml.load(in);
      if (root == null) {
        return new ArrayList<>();
      }

      Object labelsObj = root.get("label_list");
      if (labelsObj instanceof List<?>) {
        List<String> labels = new ArrayList<>();
        for (Object o : (List<?>) labelsObj) {
          labels.add(String.valueOf(o));
        }
        return labels;
      }

      return new ArrayList<>();
    }
  }

  public static Map<Integer, String> loadLabelMap(Path yamlPath) throws IOException {
    List<String> labels = loadLabels(yamlPath);
    Map<Integer, String> map = new LinkedHashMap<>();
    for (int i = 0; i < labels.size(); i++) {
      map.put(i, labels.get(i));
    }
    return map;
  }
}

package nexus.io.cv.paddle.utils;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.Map;

import org.junit.Test;

import nexus.io.tio.utils.json.JsonUtils;

public class ModelLabelUtilsTest {

  @Test
  public void test() {
    String path = "F:\\code\\python\\python-study\\11_python_ai-study\\python-pp-structure-study\\PP-DocLayout_plus-L_infer-onnx\\inference.yml";
    try {
      Map<Integer, String> map = ModelLabelUtils.loadLabelMap(Paths.get(path));
      System.out.println(JsonUtils.toJson(map));
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

}

im_shape = [800, 800];           // resize 后尺寸
scale_factor = [scaleY, scaleX]; // 缩放比例

package nexus.io.cv.paddle.demo;

import java.io.IOException;
import java.nio.FloatBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Random;
import java.util.TreeMap;

import org.opencv.core.Mat;
import org.opencv.core.Point;
import org.opencv.core.Scalar;
import org.opencv.core.Size;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONWriter;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;
import nexus.io.cv.model.DetectionResult;
import nexus.io.cv.model.PreprocessResult;
import nexus.io.cv.model.ProcessResult;

public class PPDocLayoutOnnxDemo {

  private static final int INPUT_W = 800;
  private static final int INPUT_H = 800;

  /**
   * 该模型通常已经输出原图坐标
   */
  private static final boolean OUTPUT_ALREADY_ORIGINAL_COORDS = true;

  private static final boolean DEBUG_LOG = true;
  // ======================================

  // ================= 主流程：单张 =================

  public static ProcessResult processSingleImage(OrtEnvironment env, OrtSession session, Path imgPath,
      List<String> labels, float confThres) throws Exception {

    Mat bgr = Imgcodecs.imread(imgPath.toString());
    if (bgr.empty()) {
      throw new IllegalArgumentException("无法读取图片: " + imgPath);
    }

    Mat original = bgr.clone();
    PreprocessResult pre = null;

    try {
      pre = preprocess(env, bgr);
      List<DetectionResult> detections = infer(session, pre, labels, confThres);

      if (!OUTPUT_ALREADY_ORIGINAL_COORDS) {
        remapBoxesToOriginal(detections, pre.origW, pre.origH, pre.scaleX, pre.scaleY);
      } else {
        clampBoxes(detections, pre.origW, pre.origH);
      }

      ProcessResult result = new ProcessResult();
      result.imgPath = imgPath;
      result.originalImage = original;
      result.detections = detections;
      return result;

    } finally {
      if (pre != null) {
        closeQuietly(pre.imageTensor);
        closeQuietly(pre.imShapeTensor);
        closeQuietly(pre.scaleFactorTensor);
      }
      bgr.release();
    }
  }

  // ================= 前处理：单张 =================

  private static PreprocessResult preprocess(OrtEnvironment env, Mat bgr) throws OrtException {
    int origH = bgr.rows();
    int origW = bgr.cols();

    float scaleY = (float) INPUT_H / (float) origH;
    float scaleX = (float) INPUT_W / (float) origW;

    Mat resized = new Mat();
    Imgproc.resize(bgr, resized, new Size(INPUT_W, INPUT_H), 0, 0, Imgproc.INTER_LINEAR);

    Mat rgb = new Mat();
    Imgproc.cvtColor(resized, rgb, Imgproc.COLOR_BGR2RGB);

    float[] imageData = new float[3 * INPUT_H * INPUT_W];

    for (int y = 0; y < INPUT_H; y++) {
      for (int x = 0; x < INPUT_W; x++) {
        double[] px = rgb.get(y, x);
        imageData[0 * INPUT_H * INPUT_W + y * INPUT_W + x] = (float) px[0] / 255.0f;
        imageData[1 * INPUT_H * INPUT_W + y * INPUT_W + x] = (float) px[1] / 255.0f;
        imageData[2 * INPUT_H * INPUT_W + y * INPUT_W + x] = (float) px[2] / 255.0f;
      }
    }

    float[] imShapeData = new float[] { INPUT_H, INPUT_W };
    float[] scaleFactorData = new float[] { scaleY, scaleX };

    if (DEBUG_LOG) {
      System.out.printf(Locale.US, "preprocess file=%s orig=(%d,%d) scaleFactor=[%.6f, %.6f]%n", "single", origW, origH,
          scaleY, scaleX);
    }

    PreprocessResult r = new PreprocessResult();
    r.imageTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(imageData), new long[] { 1, 3, INPUT_H, INPUT_W });
    r.imShapeTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(imShapeData), new long[] { 1, 2 });
    r.scaleFactorTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(scaleFactorData), new long[] { 1, 2 });
    r.origW = origW;
    r.origH = origH;
    r.scaleX = scaleX;
    r.scaleY = scaleY;

    resized.release();
    rgb.release();

    return r;
  }

  // ================= 推理 + 后处理：单张 =================

  private static List<DetectionResult> infer(OrtSession session, PreprocessResult pre, List<String> labels,
      float confThres) throws OrtException {

    Map<String, OnnxTensor> inputs = new HashMap<>();
    inputs.put("image", pre.imageTensor);
    inputs.put("im_shape", pre.imShapeTensor);
    inputs.put("scale_factor", pre.scaleFactorTensor);

    try (OrtSession.Result outputs = session.run(inputs)) {
      Object out0 = outputs.get("fetch_name_0").orElseThrow().getValue();
      Object out1 = outputs.get("fetch_name_1").orElseThrow().getValue();

      if (DEBUG_LOG) {
        System.out.println("fetch_name_0 class = " + out0.getClass().getName());
        System.out.println("fetch_name_1 class = " + out1.getClass().getName());
      }

      float[][] boxes = as2DFloatArray(out0);
      int[] boxesNum = as1DIntArray(out1);

      if (boxes == null) {
        throw new IllegalStateException("fetch_name_0 不是预期二维数组");
      }

      if (DEBUG_LOG) {
        System.out
            .println("fetch_name_0 rows = " + boxes.length + ", cols = " + (boxes.length > 0 ? boxes[0].length : 0));
        System.out.println("fetch_name_1 = " + Arrays.toString(boxesNum));

        for (int i = 0; i < Math.min(10, boxes.length); i++) {
          System.out.println("row[" + i + "] = " + Arrays.toString(boxes[i]));
        }
      }

      List<DetectionResult> detections = new ArrayList<>();
      int validRows = (boxesNum != null && boxesNum.length > 0) ? boxesNum[0] : boxes.length;
      validRows = Math.min(validRows, boxes.length);

      for (int i = 0; i < validRows; i++) {
        DetectionResult d = parseDetectionRow(boxes[i], labels, confThres);
        if (d != null) {
          detections.add(d);
        }
      }

      return detections;
    }
  }

  private static DetectionResult parseDetectionRow(float[] row, List<String> labels, float confThres) {
    if (row == null || row.length < 6) {
      return null;
    }

    int clsId = Math.round(row[0]);
    float score = row[1];
    float x1 = row[2];
    float y1 = row[3];
    float x2 = row[4];
    float y2 = row[5];

    if (clsId < 0) {
      return null;
    }
    if (score < confThres) {
      return null;
    }
    if (x2 <= x1 || y2 <= y1) {
      return null;
    }

    String label = clsId < labels.size() ? labels.get(clsId) : ("class_" + clsId);

    DetectionResult d = new DetectionResult();
    d.clsId = clsId;
    d.label = label;
    d.confidence = score;
    d.bbox = new float[] { x1, y1, x2, y2 };
    return d;
  }

  // ================= 坐标处理 =================

  private static void remapBoxesToOriginal(List<DetectionResult> detections, int origW, int origH, float scaleX,
      float scaleY) {
    for (DetectionResult d : detections) {
      d.bbox[0] = clamp(d.bbox[0] / scaleX, 0, origW - 1);
      d.bbox[1] = clamp(d.bbox[1] / scaleY, 0, origH - 1);
      d.bbox[2] = clamp(d.bbox[2] / scaleX, 0, origW - 1);
      d.bbox[3] = clamp(d.bbox[3] / scaleY, 0, origH - 1);
    }
  }

  private static void clampBoxes(List<DetectionResult> detections, int origW, int origH) {
    for (DetectionResult d : detections) {
      d.bbox[0] = clamp(d.bbox[0], 0, origW - 1);
      d.bbox[1] = clamp(d.bbox[1], 0, origH - 1);
      d.bbox[2] = clamp(d.bbox[2], 0, origW - 1);
      d.bbox[3] = clamp(d.bbox[3], 0, origH - 1);
    }
  }

  // ================= 结果保存 =================
  public static void saveVisualized(Mat original, List<DetectionResult> detections, Path imgPath, String outDir) {
    Mat vis = original.clone();

    for (DetectionResult d : detections) {
      Scalar color = colorOfClass(d.clsId);
      Point p1 = new Point(d.bbox[0], d.bbox[1]);
      Point p2 = new Point(d.bbox[2], d.bbox[3]);

      Imgproc.rectangle(vis, p1, p2, color, 2);

      String text = d.label + " " + String.format(Locale.US, "%.2f", d.confidence);
      double textY = Math.max(15, d.bbox[1] - 5);
      Imgproc.putText(vis, text, new Point(d.bbox[0], textY), Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, color, 2);
    }

    Path out = Paths.get(outDir, "vis", imgPath.getFileName().toString());
    Imgcodecs.imwrite(out.toString(), vis);
    vis.release();
    original.release();
  }

  public static void saveJson(List<DetectionResult> detections, Path imgPath, String outDir) throws IOException {
    List<Map<String, Object>> boxes = new ArrayList<>();
    for (DetectionResult d : detections) {
      Map<String, Object> m = new LinkedHashMap<>();
      m.put("cls_id", d.clsId);
      m.put("label", d.label);
      m.put("score", d.confidence);
      m.put("coordinate", Arrays.asList(d.bbox[0], d.bbox[1], d.bbox[2], d.bbox[3]));
      boxes.add(m);
    }

    Map<String, Object> res = new LinkedHashMap<>();
    res.put("input_path", imgPath.toString());
    res.put("page_index", null);
    res.put("boxes", boxes);

    Map<String, Object> root = new LinkedHashMap<>();
    root.put("res", res);
    root.put("debug_detection_count", detections.size());
    root.put("debug_saved_at", System.currentTimeMillis());

    Path out = Paths.get(outDir, "json", getStem(imgPath) + ".json");
    String json = JSON.toJSONString(root, JSONWriter.Feature.PrettyFormat);
    Files.writeString(out, json);
  }

  public static void saveYoloTxt(List<DetectionResult> detections, Path imgPath, String outDir) throws IOException {
    Mat img = Imgcodecs.imread(imgPath.toString());
    int w = img.cols();
    int h = img.rows();
    img.release();

    List<String> lines = new ArrayList<>();
    for (DetectionResult d : detections) {
      float x1 = clamp(d.bbox[0], 0, w);
      float y1 = clamp(d.bbox[1], 0, h);
      float x2 = clamp(d.bbox[2], 0, w);
      float y2 = clamp(d.bbox[3], 0, h);

      if (x2 <= x1 || y2 <= y1) {
        continue;
      }

      float cx = ((x1 + x2) / 2.0f) / w;
      float cy = ((y1 + y2) / 2.0f) / h;
      float bw = (x2 - x1) / w;
      float bh = (y2 - y1) / h;

      lines.add(String.format(Locale.US, "%d %.6f %.6f %.6f %.6f", d.clsId, cx, cy, bw, bh));
    }

    Path out = Paths.get(outDir, "labels", getStem(imgPath) + ".txt");
    Files.write(out, lines);
  }

  public static void saveClassesTxt(Map<Integer, String> classMap, Path out) throws IOException {
    if (classMap.isEmpty()) {
      Files.write(out, Collections.emptyList());
      return;
    }

    int maxId = classMap.keySet().stream().max(Integer::compareTo).orElse(0);
    List<String> lines = new ArrayList<>();
    for (int i = 0; i <= maxId; i++) {
      lines.add(classMap.getOrDefault(i, "class_" + i));
    }
    Files.write(out, lines);
  }

  public static Map<Integer, String> toClassMap(List<DetectionResult> detections) {
    Map<Integer, String> map = new TreeMap<>();
    for (DetectionResult d : detections) {
      map.put(d.clsId, d.label);
    }
    return map;
  }

  private static String getStem(Path p) {
    String name = p.getFileName().toString();
    int idx = name.lastIndexOf('.');
    return idx >= 0 ? name.substring(0, idx) : name;
  }

  private static float[][] as2DFloatArray(Object val) {
    if (val instanceof float[][]) {
      return (float[][]) val;
    }
    if (val instanceof float[][][]) {
      float[][][] arr = (float[][][]) val;
      return arr.length > 0 ? arr[0] : null;
    }
    return null;
  }

  private static int[] as1DIntArray(Object val) {
    if (val instanceof long[]) {
      long[] src = (long[]) val;
      int[] dst = new int[src.length];
      for (int i = 0; i < src.length; i++) {
        dst[i] = (int) src[i];
      }
      return dst;
    }
    if (val instanceof int[]) {
      return (int[]) val;
    }
    if (val instanceof float[]) {
      float[] src = (float[]) val;
      int[] dst = new int[src.length];
      for (int i = 0; i < src.length; i++) {
        dst[i] = Math.round(src[i]);
      }
      return dst;
    }
    if (val instanceof long[][]) {
      long[][] src = (long[][]) val;
      if (src.length == 0) {
        return new int[0];
      }
      int[] dst = new int[src[0].length];
      for (int i = 0; i < src[0].length; i++) {
        dst[i] = (int) src[0][i];
      }
      return dst;
    }
    if (val instanceof int[][]) {
      int[][] src = (int[][]) val;
      if (src.length == 0) {
        return new int[0];
      }
      return src[0];
    }
    if (val instanceof float[][]) {
      float[][] src = (float[][]) val;
      if (src.length == 0) {
        return new int[0];
      }
      int[] dst = new int[src[0].length];
      for (int i = 0; i < src[0].length; i++) {
        dst[i] = Math.round(src[0][i]);
      }
      return dst;
    }
    return null;
  }

  private static Scalar colorOfClass(int clsId) {
    Random random = new Random(clsId * 2027L + 17);
    return new Scalar(50 + random.nextInt(206), 50 + random.nextInt(206), 50 + random.nextInt(206));
  }

  private static float clamp(float v, float min, float max) {
    return Math.max(min, Math.min(max, v));
  }

  private static void closeQuietly(AutoCloseable c) {
    if (c != null) {
      try {
        c.close();
      } catch (Exception ignored) {
      }
    }
  }
}

package nexus.io.cv.paddle.demo;

import java.io.IOException;
import java.nio.file.DirectoryStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

import org.junit.Test;

import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtSession;
import nexus.io.cv.model.DetectionResult;
import nexus.io.cv.model.ProcessResult;
import nexus.io.cv.paddle.utils.ModelLabelUtils;

public class PPDocLayoutOnnxDemoTest {

  Set<String> SUPPORTED_EXTS = new HashSet<>(Arrays.asList(".jpg", ".jpeg", ".png", ".bmp"));

  @Test
  public void test() throws Exception {

    nu.pattern.OpenCV.loadLocally();
    // ================= 配置 =================
    String MODEL_DIR = "F:\\code\\python\\python-study\\11_python_ai-study\\python-pp-structure-study\\PP-DocLayout_plus-L_infer-onnx";
    String MODEL_PATH = MODEL_DIR + "\\inference.onnx";
    String MODEL_YML = MODEL_DIR + "\\inference.yml";

    String INPUT_DIR = "F:\\code\\project\\project-videotutor\\mc-qa-cv\\mc-qa-cv-base\\upload";
    String OUTPUT_DIR = "F:\\code\\project\\project-videotutor\\mc-qa-cv\\mc-qa-cv-base\\output\\java\\single";

    float CONF_THRESHOLD = 0.35f;
    String DEBUG_FILE_NAME = "2027202380864421888.png";

    ensureDir(Paths.get(OUTPUT_DIR));
    ensureDir(Paths.get(OUTPUT_DIR, "vis"));
    ensureDir(Paths.get(OUTPUT_DIR, "json"));
    ensureDir(Paths.get(OUTPUT_DIR, "labels"));

    List<String> labels = ModelLabelUtils.loadLabelsFromYaml(MODEL_YML);
    System.out.println("labels count = " + labels.size());

    Path imgPath = resolveSingleImage(INPUT_DIR, DEBUG_FILE_NAME);
    if (imgPath == null) {
      System.out.println("没有可处理图片");
      return;
    }

    try (OrtEnvironment env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions options = new OrtSession.SessionOptions();
        OrtSession session = env.createSession(MODEL_PATH, options)) {

      System.out.println("load ok");
      System.out.println("inputs: " + session.getInputInfo().keySet());
      System.out.println("outputs: " + session.getOutputInfo().keySet());

      ProcessResult result = PPDocLayoutOnnxDemo.processSingleImage(env, session, imgPath, labels, CONF_THRESHOLD);

      System.out.println("final detections size = " + result.detections.size());
      for (DetectionResult d : result.detections) {
        System.out.printf(Locale.US, "det: cls=%d label=%s score=%.6f box=[%.2f, %.2f, %.2f, %.2f]%n", d.clsId, d.label,
            d.confidence, d.bbox[0], d.bbox[1], d.bbox[2], d.bbox[3]);
      }

      PPDocLayoutOnnxDemo.saveVisualized(result.originalImage, result.detections, imgPath, OUTPUT_DIR);
      PPDocLayoutOnnxDemo.saveJson(result.detections, imgPath, OUTPUT_DIR);
      PPDocLayoutOnnxDemo.saveYoloTxt(result.detections, imgPath, OUTPUT_DIR);
      Map<Integer, String> classMap = PPDocLayoutOnnxDemo.toClassMap(result.detections);
      PPDocLayoutOnnxDemo.saveClassesTxt(classMap, Paths.get(OUTPUT_DIR, "classes.txt"));

      System.out.println();
      System.out.println("========== 处理完成 ==========");
      System.out.println("JSON目录:   " + Paths.get(OUTPUT_DIR, "json"));
      System.out.println("YOLO目录:   " + Paths.get(OUTPUT_DIR, "labels"));
      System.out.println("可视化目录: " + Paths.get(OUTPUT_DIR, "vis"));
    }
  }

  private Path resolveSingleImage(String inputDir, String fileName) throws IOException {
    Path root = Paths.get(inputDir);
    if (!Files.exists(root)) {
      throw new IllegalArgumentException("输入目录不存在: " + inputDir);
    }

    if (fileName != null && !fileName.trim().isEmpty()) {
      Path p = root.resolve(fileName);
      if (!Files.exists(p) || !Files.isRegularFile(p)) {
        throw new IllegalArgumentException("指定文件不存在: " + p);
      }
      String ext = getExtLower(p);
      if (!SUPPORTED_EXTS.contains(ext)) {
        throw new IllegalArgumentException("不支持的图片后缀: " + ext);
      }
      return p;
    }

    try (DirectoryStream<Path> stream = Files.newDirectoryStream(root)) {
      for (Path p : stream) {
        if (!Files.isRegularFile(p)) {
          continue;
        }
        String ext = getExtLower(p);
        if (SUPPORTED_EXTS.contains(ext)) {
          return p;
        }
      }
    }
    return null;
  }

  private static String getExtLower(Path p) {
    String name = p.getFileName().toString();
    int idx = name.lastIndexOf('.');
    return idx >= 0 ? name.substring(idx).toLowerCase(Locale.ROOT) : "";
  }

  private static void ensureDir(Path path) throws IOException {
    if (!Files.exists(path)) {
      Files.createDirectories(path);
    }
  }

}