diff --git a/compute/core/src/pdf/runLayoutModel.ts b/compute/core/src/pdf/runLayoutModel.ts index 93122ac..c97dd6b 100644 --- a/compute/core/src/pdf/runLayoutModel.ts +++ b/compute/core/src/pdf/runLayoutModel.ts @@ -234,40 +234,24 @@ function clampBox( return [x0, y0, x1, y1]; } -function sigmoid(value: number): number { - if (value >= 0) { - const z = Math.exp(-value); - return 1 / (1 + z); - } - const z = Math.exp(value); - return z / (1 + z); -} - -function buildOrderSequence(orderLogits: Float32Array, numQueries: number): number[] { - const orderVotes = new Float32Array(numQueries); - - for (let col = 0; col < numQueries; col += 1) { - let votes = 0; - - for (let row = 0; row < col; row += 1) { - votes += sigmoid(orderLogits[row * numQueries + col] ?? 0); +function softmaxMax(logits: Float32Array, offset: number, count: number): { index: number; score: number } { + let maxLogit = Number.NEGATIVE_INFINITY; + let maxIndex = 0; + for (let i = 0; i < count; i += 1) { + const value = logits[offset + i]; + if (value > maxLogit) { + maxLogit = value; + maxIndex = i; } - - for (let row = col + 1; row < numQueries; row += 1) { - votes += 1 - sigmoid(orderLogits[col * numQueries + row] ?? 0); - } - - orderVotes[col] = votes; } - const pointers = Array.from({ length: numQueries }, (_, index) => index) - .sort((a, b) => orderVotes[a]! - orderVotes[b]!); - const orderSeq = new Array(numQueries); - for (let rank = 0; rank < pointers.length; rank += 1) { - const queryIndex = pointers[rank]!; - orderSeq[queryIndex] = rank; + let sum = 0; + for (let i = 0; i < count; i += 1) { + sum += Math.exp(logits[offset + i] - maxLogit); } - return orderSeq; + + const score = sum > 0 ? (1 / sum) : 0; + return { index: maxIndex, score }; } function normalizeModelLabel(rawLabel: string): string { @@ -301,26 +285,24 @@ export async function runLayoutModel(input: RunLayoutInput): Promise= numQueries * numQueries - ? buildOrderSequence(orderLogits, numQueries) - : null; - - const detections: Array<{ - queryIdx: number; - classIdx: number; - score: number; - order: number; - rawBox: [number, number, number, number]; - }> = []; + const regions: LayoutRegion[] = []; for (let queryIdx = 0; queryIdx < numQueries; queryIdx += 1) { + const cls = softmaxMax(logits, queryIdx * classCount, classCount); + const rawLabel = idToLabel[cls.index]; + if (!rawLabel) continue; + const mapped = LABEL_MAP[normalizeModelLabel(rawLabel)]; + if (!mapped) continue; + + const minScore = CLASS_MIN_SCORE[mapped] ?? MIN_SCORE; + if (!Number.isFinite(cls.score) || cls.score < minScore) continue; + const cx = predBoxes[queryIdx * 4 + 0] * pageWidth; const cy = predBoxes[queryIdx * 4 + 1] * pageHeight; const w = predBoxes[queryIdx * 4 + 2] * pageWidth; @@ -332,33 +314,7 @@ export async function runLayoutModel(input: RunLayoutInput): Promise { - if (Math.abs(b.score - a.score) > 1e-6) return b.score - a.score; - return a.order - b.order; - }); - - for (const detection of detections.slice(0, numQueries)) { - const rawLabel = idToLabel[detection.classIdx]; - if (!rawLabel) continue; - const mapped = LABEL_MAP[normalizeModelLabel(rawLabel)]; - if (!mapped) continue; - - const minScore = CLASS_MIN_SCORE[mapped] ?? MIN_SCORE; - if (!Number.isFinite(detection.score) || detection.score < minScore) continue; - - const clamped = clampBox(detection.rawBox, pageWidth, pageHeight); + const clamped = clampBox(rawBox, pageWidth, pageHeight); if (!clamped) continue; const sizeRule = MIN_REGION_SIZE[mapped]; @@ -371,7 +327,7 @@ export async function runLayoutModel(input: RunLayoutInput): Promise ({ + runOutput: { + logits: { data: new Float32Array() }, + pred_boxes: { data: new Float32Array() }, + }, +})); + +vi.mock('onnxruntime-node', () => ({ + InferenceSession: { + create: vi.fn(async () => ({ + run: vi.fn(async () => mockState.runOutput), + })), + }, + Tensor: class Tensor { + constructor( + public type: string, + public data: Float32Array, + public dims: number[], + ) {} + }, +})); + +vi.mock('fs/promises', () => ({ + readFile: vi.fn(async (path: string) => { + if (path === '/tmp/model-config.json') { + return JSON.stringify({ + id2label: { + 0: 'text', + 1: 'table', + }, + }); + } + + if (path === '/tmp/model-preprocessor.json') { + return JSON.stringify({ + size: { width: 2, height: 2 }, + rescale_factor: 1 / 255, + image_mean: [0, 0, 0], + image_std: [1, 1, 1], + }); + } + + throw new Error(`unexpected readFile path: ${path}`); + }), +})); + +vi.mock('@napi-rs/canvas', () => { + const createCanvas = (width: number, height: number) => ({ + getContext: () => ({ + fillStyle: '#ffffff', + fillRect: () => {}, + drawImage: () => {}, + imageSmoothingEnabled: true, + getImageData: () => ({ + data: new Uint8ClampedArray(width * height * 4).fill(255), + }), + }), + }); + const loadImage = vi.fn(async () => ({ width: 2, height: 2 })); + + return { + createCanvas, + loadImage, + default: { + createCanvas, + loadImage, + }, + }; +}); + +vi.mock('../../src/pdf/model', () => ({ + ensureModel: vi.fn(async () => '/tmp/model.onnx'), + MODEL_CONFIG_PATH: '/tmp/model-config.json', + MODEL_PREPROCESSOR_PATH: '/tmp/model-preprocessor.json', +})); + +vi.mock('../../src/config/cpu-budget', () => ({ + getOnnxThreadsPerJob: vi.fn(() => 1), +})); + +describe('runLayoutModel', () => { + beforeEach(() => { + vi.resetModules(); + mockState.runOutput = { + logits: { data: new Float32Array() }, + pred_boxes: { data: new Float32Array() }, + }; + }); + + test('keeps one winner per query instead of dropping later queries behind duplicate class rows', async () => { + mockState.runOutput = { + logits: { + data: new Float32Array([ + 3, + 4, + 2.5, + 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(2); + expect(regions[0]?.label).toBe('text'); + expect(regions[0]?.confidence).toBeCloseTo(0.9168273, 6); + expect(regions[0]?.bbox).toEqual([ + expect.closeTo(60, 5), + expect.closeTo(60, 5), + expect.closeTo(90, 5), + expect.closeTo(90, 5), + ]); + expect(regions[1]?.label).toBe('table'); + expect(regions[1]?.confidence).toBeCloseTo(0.73105858, 6); + expect(regions[1]?.bbox).toEqual([ + expect.closeTo(10, 5), + expect.closeTo(10, 5), + expect.closeTo(40, 5), + expect.closeTo(40, 5), + ]); + }); +}); diff --git a/compute/worker/tests/api/routes.test.ts b/compute/worker/tests/api/routes.test.ts index 6ba641d..dad89d9 100644 --- a/compute/worker/tests/api/routes.test.ts +++ b/compute/worker/tests/api/routes.test.ts @@ -120,7 +120,7 @@ describe('compute worker API routes', () => { expect(stream.body).toContain('"status":"succeeded"'); }); - test('marks stale running whisper and pdf ops failed during startup reconciliation but leaves queued ops on the conservative path', async () => { + test('marks stale running whisper and pdf ops failed during request-time orphan recovery but leaves queued ops on the conservative path', async () => { const now = Date.now(); fake.seedState({ opId: 'op-stale-whisper-running', @@ -159,6 +159,8 @@ describe('compute worker API routes', () => { updatedAt: now - 310_000, }); + // GET /ops/:opId resolves via getOpState(), which first awaits the shared + // orphanRecoveryPromise path through ensureOrphanedOpRecovery(). const fetch = await runtime.app.inject({ method: 'GET', url: '/ops/op-stale-whisper-running', diff --git a/src/components/reader/PdfLayoutScan.module.css b/src/components/reader/PdfLayoutScan.module.css index 0c44b79..5ee5a31 100644 --- a/src/components/reader/PdfLayoutScan.module.css +++ b/src/components/reader/PdfLayoutScan.module.css @@ -51,9 +51,9 @@ color: var(--background); background: var(--accent); box-shadow: 0 2px 10px color-mix(in srgb, var(--accent) 50%, transparent); - animation: tagIn 0.34s var(--ease); + animation: tag-in 0.34s var(--ease); } -@keyframes tagIn { +@keyframes tag-in { from { opacity: 0; transform: translateY(-4px) scale(0.92); } to { opacity: 1; transform: translateY(0) scale(1); } }