danvics/openreader
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(pdf): add strict validation for layout model output shapes and extend test coverage

Browse source 
Add explicit error handling for invalid or inconsistent pred_boxes and logits array lengths in runLayoutModel to prevent silent failures. Expand test suite to verify correct region filtering and error scenarios, ensuring only labeled regions are returned and malformed outputs are handled robustly.
This commit is contained in:
Richard R 2026-06-03 04:35:27 -06:00
parent 6024ed2244
commit bc5388af6b
2 changed files with 61 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
compute/core/src/pdf/runLayoutModel.ts
View file

@ -286,11 +286,29 @@ export async function runLayoutModel(input: RunLayoutInput): Promise<LayoutRegio
const logits = output.logits?.data as Float32Array | undefined;
const predBoxes = output.pred_boxes?.data as Float32Array | undefined;
if (!logits || !predBoxes) return [];
if (predBoxes.length === 0 && logits.length === 0) return [];
if (predBoxes.length === 0 || logits.length === 0) {
throw new Error(
`layout-model-invalid-output-shape: pred_boxes length=${predBoxes.length}, logits length=${logits.length}`,
);
}
const numQueries = Math.floor(predBoxes.length / 4);
if (numQueries <= 0) return [];
const classCount = Math.floor(logits.length / numQueries);
if (classCount <= 0) return [];
if (predBoxes.length % 4 !== 0) {
throw new Error(`layout-model-invalid-pred-box-shape: length ${predBoxes.length} is not divisible by 4`);
}
const numQueries = predBoxes.length / 4;
if (numQueries <= 0) {
throw new Error(`layout-model-invalid-pred-box-shape: expected positive query count, got ${numQueries}`);
}
if (logits.length % numQueries !== 0) {
throw new Error(
`layout-model-invalid-logit-shape: length ${logits.length} is not divisible by query count ${numQueries}`,
);
}
const classCount = logits.length / numQueries;
if (classCount <= 0) {
throw new Error(`layout-model-invalid-logit-shape: expected positive class count, got ${classCount}`);
}
const regions: LayoutRegion[] = [];
for (let queryIdx = 0; queryIdx < numQueries; queryIdx += 1) {

39
compute/core/tests/control-plane/run-layout-model.test.ts
View file

@ -133,4 +133,43 @@ describe('runLayoutModel', () => {
expect.closeTo(40, 5),
]);
});
test('drops unlabeled query winners and keeps only labeled regions', async () => {
mockState.runOutput = {
logits: {
data: new Float32Array([
0.1,
0.2,
5,
4,
0.1,
0.1,
]),
},
pred_boxes: {
data: new Float32Array([
0.25, 0.25, 0.3, 0.3,
0.75, 0.75, 0.3, 0.3,
]),
},
};
const { runLayoutModel } = await import('../../src/pdf/runLayoutModel');
const regions = await runLayoutModel({
pageWidth: 100,
pageHeight: 100,
textItems: [{} as never],
pageImage: Buffer.from([1]),
});
expect(regions).toHaveLength(1);
expect(regions[0]?.label).toBe('text');
expect(regions[0]?.confidence).toBeCloseTo(0.96109135, 5);
expect(regions[0]?.bbox).toEqual([
expect.closeTo(60, 5),
expect.closeTo(60, 5),
expect.closeTo(90, 5),
expect.closeTo(90, 5),
]);
});
});