|
--- |
|
license: apache-2.0 |
|
library_name: transformers.js |
|
tags: |
|
- DocLayout-yolo |
|
- transformers.js |
|
- onnx |
|
--- |
|
Converted from [wybxc/DocLayout-YOLO-DocStructBench-onnx](https://huggingface.co/wybxc/DocLayout-YOLO-DocStructBench-onnx) |
|
|
|
```ts |
|
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) |
|
``` |