darknoah99/DocLayout-YOLO-DocStructBench-onnx

Converted from wybxc/DocLayout-YOLO-DocStructBench-onnx

import {
  AutoProcessor,
  RawImage,
  AutoModel,
  env,
} from '@huggingface/transformers';


const model = await AutoModel.from_pretrained('darknoah99/DocLayout-YOLO-DocStructBench-onnx');
const processor = await AutoProcessor.from_pretrained('darknoah99/DocLayout-YOLO-DocStructBench-onnx', {});
let image;
if (isUrl(fileOrUrl)) image = await RawImage.fromURL(fileOrUrl);
else if (fs.statSync(fileOrUrl).isFile()) {
  const data = fs.readFileSync(fileOrUrl);
  const blob = new Blob([data]);
  image = await RawImage.fromBlob(blob);
}
const { pixel_values } = await processor(image);
const output = await model({ images: pixel_values });
const permuted = output.output0[0];
const result = [];
const threshold = 0.3;
const [scaledHeight, scaledWidth] = pixel_values.dims.slice(-2);
for (const [xc, yc, w, h, ...scores] of permuted.tolist()) {
  const x1 = (xc / scaledWidth) * image.width;
  const y1 = (yc / scaledHeight) * image.height;

  const x2 = (w / scaledWidth) * image.width;
  const y2 = (h / scaledHeight) * image.height;
  const score = scores[0];
  if (score > threshold) {
    const label = model.config.id2label[scores[1]];
    result.push({
      x1,
      x2,
      y1,
      y2,
      score,
      label,
      index: scores[1],
    });
  }
}
console.log(result)