layout.cpp

// layout.cpp - layout.mnn 的 MNN 推理实现 (预处理 + 推理 + 解析)
//
// 实测确认 (probe_mnn_run.py, WSL MNN 3.6.1):
//   fetch_name_0 [300,7] = (label_id, score, x1, y1, x2, y2, order)
//     - 坐标在 800x800 输入坐标系, 回原图需 x*=w/800, y*=h/800
//     - 图内不做阈值和 NMS: 300 行全量输出, 同一区域会以不同 label 重复出现
//   fetch_name_1 [1]     = 常量 300 (top-k), 无用
//   fetch_name_2 [300,200,200] int32 = 每 query 的 1/4 分辨率 mask (polygon 用, 此处忽略)
#include "layout.hpp"

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <memory>
#include <numeric>
#include <vector>

#include <opencv2/imgproc.hpp>

#include "MNN/Interpreter.hpp"
#include "log.hpp"

// 可视化中文渲染: 需 OpenCV freetype 模块 (libopencv-freetype-dev) + CJK 字体。
// 没有时自动退化为英文标签 (Hershey 字体), 编译不报错。
#if __has_include(<opencv2/freetype.hpp>)
#define HAVE_OPENCV_FREETYPE 1
#include <opencv2/freetype.hpp>
#endif

namespace {

constexpr int kInputSize = 800;

// 类别表, 顺序即 label_id (来源: PP-DocLayoutV3_onnx 的 inference.yml)
const char* kLabels[] = {
    "abstract",        "algorithm",   "aside_text",    "chart",
    "content",         "display_formula", "doc_title", "figure_title",
    "footer",          "footer_image", "footnote",     "formula_number",
    "header",          "header_image", "image",        "inline_formula",
    "number",          "paragraph_title", "reference", "reference_content",
    "seal",            "table",       "text",          "vertical_text",
    "vision_footnote",
};
constexpr int kNumLabels = static_cast<int>(sizeof(kLabels) / sizeof(kLabels[0]));

// 中文类名, 下标 = label_id (与 kLabels 一一对应)
const char* kLabelsZh[] = {
    "摘要",     "算法",     "侧边文本", "图表",     "目录",
    "行间公式", "文档标题", "图题",     "页脚",     "页脚图片",
    "脚注",     "公式编号", "页眉",     "页眉图片", "图片",
    "行内公式", "编号",     "段落标题", "参考文献", "参考文献内容",
    "印章",     "表格",     "正文",     "竖排文本", "图注",
};
static_assert(sizeof(kLabelsZh) / sizeof(kLabelsZh[0]) == kNumLabels,
              "中英文类名表数量不一致");

const char* label_zh(int label_id) {
  return (label_id >= 0 && label_id < kNumLabels) ? kLabelsZh[label_id] : "?";
}

// 按类别给框着色 (BGR): 同类业务一个色系
cv::Scalar label_color(int label_id) {
  switch (label_id) {
    case 4:    // content
    case 22:   // text
    case 23:   // vertical_text
    case 0:    // abstract
    case 1:    // algorithm
    case 2:    // aside_text
    case 10:   // footnote
    case 16:   // number
    case 18:   // reference
    case 19:   // reference_content
      return {255, 0, 0};    // 蓝
    case 6:    // doc_title
    case 17:   // paragraph_title
    case 7:    // figure_title
      return {0, 255, 255};  // 黄
    case 5:    // display_formula
    case 15:   // inline_formula
    case 11:   // formula_number
      return {0, 140, 255};  // 橙
    case 14:   // image
    case 3:    // chart
    case 24:   // vision_footnote
      return {0, 0, 255};    // 红
    case 21:                 // table
      return {0, 255, 0};    // 绿
    case 20:                 // seal
      return {255, 0, 255};  // 紫
    case 12:   // header
    case 13:   // header_image
    case 8:    // footer
    case 9:    // footer_image
      return {255, 255, 0};  // 青
    default:
      return {160, 160, 160};  // 未知类别 灰
  }
}

// 探测系统中文字体
std::string find_cjk_font() {
  const char* candidates[] = {
      "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
      "/usr/share/fonts/opentype/noto/NotoSerifCJK-Regular.ttc",
      "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
      "/usr/share/fonts/truetype/arphic/uming.ttc",
      "C:/Windows/Fonts/msyh.ttc",
      "C:/Windows/Fonts/simhei.ttf",
  };
  for (const char* p : candidates) {
    FILE* f = fopen(p, "rb");
    if (f != nullptr) {
      fclose(f);
      return p;
    }
  }
  return "";
}

// 模型张量名 (Paddle 导出签名, 与 netron / 实测一致)
constexpr const char* kInImage = "image";
constexpr const char* kInImShape = "im_shape";
constexpr const char* kInScale = "scale_factor";
constexpr const char* kOutBoxes = "fetch_name_0";

void fill_float_input(MNN::Interpreter* net, MNN::Session* session,
                      const char* name, const void* data, size_t nbytes) {
  MNN::Tensor* t = net->getSessionInput(session, name);
  if (t == nullptr) {
    fprintf(stderr, "[layout] 找不到输入张量: %s\n", name);
    return;
  }
  MNN::Tensor host(t, t->getDimensionType());
  memcpy(host.host<void>(), data, nbytes);
  t->copyFromHostTensor(&host);
}

std::shared_ptr<MNN::Tensor> fetch_output(MNN::Interpreter* net,
                                          MNN::Session* session,
                                          const char* name) {
  MNN::Tensor* t = net->getSessionOutput(session, name);
  if (t == nullptr) {
    fprintf(stderr, "[layout] 找不到输出张量: %s\n", name);
    return nullptr;
  }
  auto host = std::make_shared<MNN::Tensor>(t, t->getDimensionType());
  t->copyToHostTensor(host.get());
  return host;
}

void dump_tensor(const char* name, const MNN::Tensor* t, bool as_int) {
  printf("[dump] %s  shape=[", name);
  for (int i = 0; i < t->dimensions(); ++i) {
    printf("%d%s", t->length(i), (i + 1 < t->dimensions()) ? " x " : "");
  }
  printf("]  first values:");
  const int n = static_cast<int>(t->elementSize());
  const int show = n < 12 ? n : 12;
  for (int i = 0; i < show; ++i) {
    if (as_int)
      printf(" %d", t->host<int32_t>()[i]);
    else
      printf(" %.3f", t->host<float>()[i]);
  }
  printf("\n");
}

// ---- 以下后处理与 glmocr/utils/layout_postprocess_utils.py 对齐 ----

struct Box {
  int cls = 0;
  float score = 0.f;
  float x1 = 0, y1 = 0, x2 = 0, y2 = 0;
  int order = 0;
};

// IoU, 含 Paddle 遗留的 +1 (与 layout_postprocess_utils.iou 一致)
float legacy_iou(const Box& a, const Box& b) {
  const float ix1 = std::max(a.x1, b.x1);
  const float iy1 = std::max(a.y1, b.y1);
  const float ix2 = std::min(a.x2, b.x2);
  const float iy2 = std::min(a.y2, b.y2);
  const float inter =
      std::max(0.0f, ix2 - ix1 + 1) * std::max(0.0f, iy2 - iy1 + 1);
  const float area1 = (a.x2 - a.x1 + 1) * (a.y2 - a.y1 + 1);
  const float area2 = (b.x2 - b.x1 + 1) * (b.y2 - b.y1 + 1);
  return inter / (area1 + area2 - inter);
}

// 同类 IoU>=0.6 抑制 / 异类 IoU>=0.98 抑制 (与 layout_postprocess_utils.nms 一致)
std::vector<Box> nms(std::vector<Box> boxes) {
  std::stable_sort(boxes.begin(), boxes.end(),
                   [](const Box& a, const Box& b) { return a.score > b.score; });
  std::vector<bool> removed(boxes.size(), false);
  std::vector<Box> kept;
  for (size_t i = 0; i < boxes.size(); ++i) {
    if (removed[i]) continue;
    kept.push_back(boxes[i]);
    for (size_t j = i + 1; j < boxes.size(); ++j) {
      if (removed[j]) continue;
      const float thresh =
          (boxes[i].cls == boxes[j].cls) ? 0.6f : 0.98f;
      if (legacy_iou(boxes[i], boxes[j]) >= thresh) removed[j] = true;
    }
  }
  return kept;
}

// ---- 以下为补齐的后处理步骤 (与 layout_postprocess_utils.py 对齐) ----

// is_contained: box a 是否被 box b 包含。
// 非对称: 分母是 a 自身面积; 宽高【不带】+1 (与 legacy_iou 的 +1 约定不同, 勿混!)
// 对应 layout_postprocess_utils.is_contained
bool is_contained(const Box& a, const Box& b) {
  const float a_area = (a.x2 - a.x1) * (a.y2 - a.y1);
  if (a_area <= 0.f) return false;
  const float ix1 = std::max(a.x1, b.x1);
  const float iy1 = std::max(a.y1, b.y1);
  const float ix2 = std::min(a.x2, b.x2);
  const float iy2 = std::min(a.y2, b.y2);
  const float inter =
      std::max(0.0f, ix2 - ix1) * std::max(0.0f, iy2 - iy1);
  return inter / a_area >= 0.8f;
}

// 整页 image 大框过滤 (对应 apply_layout_postprocess 的 filter_large_image 段):
// 仅 cls==14(image), 横图阈值 0.82 / 竖图 0.93; 框面积先 clamp 再算;
// 全部滤光时保留原样不清空 (Python: if len(filtered)>0 才替换)
void filter_large_images(std::vector<Box>& boxes, int img_w, int img_h) {
  if (boxes.size() <= 1) return;
  const float area_thres = (img_w > img_h) ? 0.82f : 0.93f;
  const float img_area = static_cast<float>(img_w) * img_h;
  std::vector<Box> filtered;
  filtered.reserve(boxes.size());
  for (const Box& b : boxes) {
    if (b.cls != 14) {  // image
      filtered.push_back(b);
      continue;
    }
    const float x1 = std::max(0.f, b.x1), y1 = std::max(0.f, b.y1);
    const float x2 = std::min(static_cast<float>(img_w), b.x2);
    const float y2 = std::min(static_cast<float>(img_h), b.y2);
    const float box_area = (x2 - x1) * (y2 - y1);
    if (box_area <= area_thres * img_area) filtered.push_back(b);
  }
  if (!filtered.empty()) boxes = std::move(filtered);
}

// 单个类别的 containment 检查 (对应 check_containment 的 category+mode 分支):
// preserve 框跳过位置在外层 i 循环 (连 contains_other 都不标记);
// large 模式只考察 j.cls==category; small 模式只考察 i.cls==category
void check_containment(const std::vector<Box>& boxes, int category,
                       bool large_mode, std::vector<char>& contains,
                       std::vector<char>& contained) {
  const int n = static_cast<int>(boxes.size());
  for (int i = 0; i < n; ++i) {
    if (boxes[i].cls == 14 || boxes[i].cls == 20 || boxes[i].cls == 3)
      continue;  // preserve: image/seal/chart 永不判为被包含
    for (int j = 0; j < n; ++j) {
      if (i == j) continue;
      bool hit = false;
      if (large_mode)
        hit = (boxes[j].cls == category) && is_contained(boxes[i], boxes[j]);
      else
        hit = (boxes[i].cls == category) && is_contained(boxes[i], boxes[j]);
      if (hit) {
        contained[i] = 1;
        contains[j] = 1;
      }
    }
  }
}

// 嵌套框合并 (对应 dict 模式 merge_bboxes_mode, glmocr/config.yaml:221-246):
// 全类 large, 仅 18(reference)=small; preserve={image(14), seal(20), chart(3)}
// 逐类独立 check 后 & 到同一 keep_mask, 最后一次性过滤
void merge_nested_boxes(std::vector<Box>& boxes) {
  if (boxes.empty()) return;
  const int n = static_cast<int>(boxes.size());

  // 收集当前出现的类别
  std::vector<int> categories;
  std::vector<bool> seen(25, false);
  for (const Box& b : boxes)
    if (!seen[b.cls]) {
      seen[b.cls] = true;
      categories.push_back(b.cls);
    }

  std::vector<char> keep(n, 1);
  std::vector<char> contains(n), contained(n);
  for (int category : categories) {
    const bool large_mode = (category != 18);  // 18=reference -> small
    std::fill(contains.begin(), contains.end(), 0);
    std::fill(contained.begin(), contained.end(), 0);
    check_containment(boxes, category, large_mode, contains, contained);
    for (int i = 0; i < n; ++i) {
      if (large_mode)
        keep[i] &= (contained[i] == 0);
      else
        keep[i] &= (contains[i] == 0) || (contained[i] == 1);
    }
  }

  std::vector<Box> merged;
  merged.reserve(n);
  for (int i = 0; i < n; ++i)
    if (keep[i]) merged.push_back(boxes[i]);
  boxes = std::move(merged);
}

// 中心等比扩框 (对应 unclip_boxes 的 tuple 分支); ratio=1.0 时数值不变
void unclip(std::vector<Box>& boxes, float w_ratio, float h_ratio) {
  if (w_ratio == 1.0f && h_ratio == 1.0f) return;
  for (Box& b : boxes) {
    const float w = b.x2 - b.x1;
    const float h = b.y2 - b.y1;
    const float cx = b.x1 + w / 2.f;
    const float cy = b.y1 + h / 2.f;
    b.x1 = cx - w * w_ratio / 2.f;
    b.x2 = cx + w * w_ratio / 2.f;
    b.y1 = cy - h * h_ratio / 2.f;
    b.y2 = cy + h * h_ratio / 2.f;
  }
}

// 最终 clamp 到图界 + 剔除无效框 (对应输出阶段的 clamp + skip invalid)
// 此时的框即 OCR 裁剪框
void clamp_and_validate(std::vector<Box>& boxes, int img_w, int img_h) {
  std::vector<Box> valid;
  valid.reserve(boxes.size());
  for (Box b : boxes) {
    b.x1 = std::max(0.f, std::min(b.x1, static_cast<float>(img_w)));
    b.y1 = std::max(0.f, std::min(b.y1, static_cast<float>(img_h)));
    b.x2 = std::max(0.f, std::min(b.x2, static_cast<float>(img_w)));
    b.y2 = std::max(0.f, std::min(b.y2, static_cast<float>(img_h)));
    if (b.x1 < b.x2 && b.y1 < b.y2) valid.push_back(b);
  }
  boxes = std::move(valid);
}

}  // namespace

const char* layout_label_name(int label_id) {
  return (label_id >= 0 && label_id < kNumLabels) ? kLabels[label_id] : "?";
}

LayoutDetector::~LayoutDetector() {
#if HAVE_OPENCV_FREETYPE
  if (ft2_ != nullptr)
    delete static_cast<cv::Ptr<cv::freetype::FreeType2>*>(ft2_);
#endif
  if (net_ != nullptr) {
    if (session_ != nullptr) net_->releaseSession(session_);
    delete net_;
  }
}

bool LayoutDetector::loadModel(const std::string& mnn_path, int num_threads) {
  {
    LOG_SCOPE("createFromFile");
    net_ = MNN::Interpreter::createFromFile(mnn_path.c_str());
  }
  if (net_ == nullptr) return false;

  MNN::ScheduleConfig config;
  config.type = MNN_FORWARD_CPU;  // GPU: MNN_FORWARD_OPENCL / MNN_FORWARD_CUDA
  config.numThread = num_threads;
  {
    LOG_SCOPE("createSession+resizeSession");
    session_ = net_->createSession(config);
    if (session_ == nullptr) return false;

    // 输入 shape 全部固定, resize 一次即可
    const int s = kInputSize;
    if (MNN::Tensor* t = net_->getSessionInput(session_, kInImage))
      net_->resizeTensor(t, {1, 3, s, s});
    if (MNN::Tensor* t = net_->getSessionInput(session_, kInImShape))
      net_->resizeTensor(t, {1, 2});
    if (MNN::Tensor* t = net_->getSessionInput(session_, kInScale))
      net_->resizeTensor(t, {1, 2});
    net_->resizeSession(session_);
  }
  return true;
}

bool LayoutDetector::detect(const cv::Mat& bgr, std::vector<LayoutBox>* out) {
  out->clear();
  if (session_ == nullptr) return false;
  const int s = kInputSize;

  // ---- 1. 预处理: BGR->RGB / 拉伸 800x800 / ÷255 / HWC->CHW ----
  std::vector<float> image;
  float h = 0.f, w = 0.f;
  {
    LOG_SCOPE("预处理 (resize/CHW)");
    cv::Mat rgb, resized;
    cv::cvtColor(bgr, rgb, cv::COLOR_BGR2RGB);
    cv::resize(rgb, resized, cv::Size(s, s), 0, 0, cv::INTER_CUBIC);

    image.resize(3 * s * s);
    const float kInv255 = 1.0f / 255.0f;
    for (int y = 0; y < s; ++y) {
      const cv::Vec3b* row = resized.ptr<cv::Vec3b>(y);
      for (int x = 0; x < s; ++x) {
        image[0 * s * s + y * s + x] = row[x][0] * kInv255;  // R
        image[1 * s * s + y * s + x] = row[x][1] * kInv255;  // G
        image[2 * s * s + y * s + x] = row[x][2] * kInv255;  // B
      }
    }
    h = static_cast<float>(bgr.rows);
    w = static_cast<float>(bgr.cols);
  }

  // ---- 2. 辅助输入 (实测用这组值, 输出坐标即 800 输入坐标系) ----
  const float im_shape[2] = {h, w};
  const float scale_factor[2] = {h / s, w / s};

  fill_float_input(net_, session_, kInImage, image.data(),
                   image.size() * sizeof(float));
  fill_float_input(net_, session_, kInImShape, im_shape, sizeof(im_shape));
  fill_float_input(net_, session_, kInScale, scale_factor,
                   sizeof(scale_factor));

  // ---- 3. 推理 + 取输出 ----
  std::shared_ptr<MNN::Tensor> boxes_t;
  {
    LOG_SCOPE("MNN 推理");
    net_->runSession(session_);
    boxes_t = fetch_output(net_, session_, kOutBoxes);
  }
  if (boxes_t == nullptr) return false;

  if (debug_) {
    dump_tensor(kOutBoxes, boxes_t.get(), false);
    printf("[dump] 原图 %d x %d\n", bgr.cols, bgr.rows);
  }

  // ---- 4. 解析 fetch_name_0 [300,7] ----
  if (boxes_t->dimensions() != 2 || boxes_t->length(1) != 7) {
    fprintf(stderr, "[layout] fetch_name_0 形状非 [N,7] (dims=%d)\n",
            boxes_t->dimensions());
    return false;
  }
  const float kx = w / s;  // x: 800 系 -> 原图宽
  const float ky = h / s;  // y: 800 系 -> 原图高
  const int n = boxes_t->length(0);
  const float* p = boxes_t->host<float>();

  // 原始 [300,7] 全量 dump (供 verify_postprocess.py 隔离验证后处理链)
  if (debug_ && !raw_dump_path_.empty()) {
    FILE* f = fopen(raw_dump_path_.c_str(), "w");
    if (f != nullptr) {
      for (int i = 0; i < n; ++i) {
        const float* r = p + i * 7;
        fprintf(f, "%.6f %.6f %.6f %.6f %.6f %.6f %.6f\n", r[0], r[1], r[2],
                r[3], r[4], r[5], r[6]);
      }
      fclose(f);
      printf("[dump] 原始输出已保存: %s\n", raw_dump_path_.c_str());
    }
  }

  std::vector<Box> cands;
  cands.reserve(n);
  for (int i = 0; i < n; ++i) {
    const float* r = p + i * 7;
    if (r[1] < score_threshold_) continue;  // 图内无阈值, 这里过滤
    Box b;
    b.cls = static_cast<int>(r[0]);
    b.score = r[1];
    b.x1 = std::max(0.0f, r[2] * kx);
    b.y1 = std::max(0.0f, r[3] * ky);
    b.x2 = std::min(w, r[4] * kx);
    b.y2 = std::min(h, r[5] * ky);
    b.order = static_cast<int>(r[6]);
    if (b.x1 >= b.x2 || b.y1 >= b.y2) continue;
    cands.push_back(b);
  }

  // ---- 5. 后处理流水 (顺序与 layout_postprocess_utils.apply_layout_postprocess 一致) ----
  LOG_SCOPE("后处理 (NMS/大图过滤/嵌套合并/排序/unclip)");
  std::vector<Box> kept = nms(std::move(cands));
  // 5a. 整页 image 大框过滤 (NMS 之后)
  filter_large_images(kept, bgr.cols, bgr.rows);
  // 5b. 嵌套框合并 (per-class large/small, preserve=image/seal/chart)
  merge_nested_boxes(kept);
  // 5c. 按阅读顺序排序 (stable: order 并列时维持 NMS 后的 score 降序,
  //     对应 np.argsort 的前序保持行为)
  std::stable_sort(kept.begin(), kept.end(),
                   [](const Box& a, const Box& b) { return a.order < b.order; });
  // 5d. unclip 扩框 + 最终 clamp + 无效框剔除; 得到的框即 OCR 裁剪框
  unclip(kept, unclip_w_, unclip_h_);
  clamp_and_validate(kept, bgr.cols, bgr.rows);

  out->reserve(kept.size());
  for (const Box& b : kept) {
    LayoutBox d;
    d.label_id = b.cls;
    d.score = b.score;
    d.x1 = b.x1;
    d.y1 = b.y1;
    d.x2 = b.x2;
    d.y2 = b.y2;
    d.order = b.order > 0 ? b.order : -1;  // 对应 Python 的 order None
    out->push_back(d);
  }
  return true;
}

cv::Mat LayoutDetector::draw(const cv::Mat& bgr,
                             const std::vector<LayoutBox>& boxes) {
  cv::Mat vis = bgr.clone();
  if (boxes.empty()) return vis;

  // 字号/线宽随图分辨率缩放 (默认字号 = 图高/60)
  const int font_h = font_size_ > 0 ? font_size_ : std::max(24, vis.rows / 60);
  const int thick = std::max(2, vis.rows / 800);

#if HAVE_OPENCV_FREETYPE
  // 懒加载中文字体 (整个生命周期只加载一次)
  auto* ft = static_cast<cv::Ptr<cv::freetype::FreeType2>*>(ft2_);
  if (ft == nullptr) {
    std::string font = font_path_;
    if (font.empty()) font = find_cjk_font();
    if (!font.empty()) {
      ft = new cv::Ptr<cv::freetype::FreeType2>(cv::freetype::createFreeType2());
      (*ft)->loadFontData(font, 0);
      ft2_ = ft;
      printf("[layout] 中文字体: %s, 字号 %dpx\n", font.c_str(), font_h);
    } else {
      fprintf(stderr, "[layout] 未找到中文字体, 标签退化为英文\n");
    }
  }
  const bool use_ft = (ft != nullptr);
#endif

  for (const auto& b : boxes) {
    const cv::Scalar color = label_color(b.label_id);
    const cv::Point p1(static_cast<int>(b.x1), static_cast<int>(b.y1));
    const cv::Point p2(static_cast<int>(b.x2), static_cast<int>(b.y2));
    cv::rectangle(vis, p1, p2, color, thick);

    // 标签: 中文类名 + score
    char score_txt[16];
    snprintf(score_txt, sizeof(score_txt), " %.2f", b.score);
    std::string txt;
#if HAVE_OPENCV_FREETYPE
    if (use_ft)
      txt = std::string(label_zh(b.label_id)) + score_txt;
    else
#endif
      txt = std::string(layout_label_name(b.label_id)) + score_txt;

    // 文字底端默认在框外左上; 太靠近图顶时放框内
    cv::Point org(p1.x + 4, p1.y - 6);
    if (org.y - font_h < 0) org.y = p1.y + font_h + 6;

#if HAVE_OPENCV_FREETYPE
    if (use_ft) {
      (*ft)->putText(vis, txt, org, font_h, color, -1, cv::LINE_AA, true);
    } else
#endif
    {
      cv::putText(vis, txt, org, cv::FONT_HERSHEY_SIMPLEX, font_h / 30.0,
                  color, thick, cv::LINE_AA);
    }
  }
  return vis;
}