import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { Image as KonvaImage, Layer, Rect, Stage, Text } from 'react-konva' import { motion, useDragControls } from 'motion/react' import { useStudio } from './StudioContext' import { clamp, confidenceStroke, getCellKey, getPages, getTablesForPage, heatmapColor, isSameCell, } from './studioUtils' import API from '../../api' function useLoadedImage(src, onError) { const [image, setImage] = useState(null) useEffect(() => { const nextImage = new window.Image() nextImage.crossOrigin = 'anonymous' nextImage.onload = () => setImage(nextImage) nextImage.onerror = () => { setImage(null) onError?.() } nextImage.src = src return () => { nextImage.onload = null nextImage.onerror = null } }, [onError, src]) return image } function TableLabel({ table, bbox, scale }) { const label = `Table ${table.table_id + 1}` const x = bbox[0] * scale + 8 const y = Math.max(12, bbox[1] * scale - 18) return ( ) } function intersectionArea(a, b) { const x1 = Math.max(a[0], b[0]) const y1 = Math.max(a[1], b[1]) const x2 = Math.min(a[2], b[2]) const y2 = Math.min(a[3], b[3]) return Math.max(0, x2 - x1) * Math.max(0, y2 - y1) } function centerOfBox(box) { return [(box[0] + box[2]) / 2, (box[1] + box[3]) / 2] } function resolveOcrItemCell(table, item) { const itemBox = Array.isArray(item?.bbox) && item.bbox.length === 4 ? item.bbox : null const cells = Array.isArray(table?.cells) ? table.cells : [] if (!itemBox || !cells.length) { return null } // Prefer explicit assignment if present in payload. if (Number.isFinite(item?.row) && Number.isFinite(item?.col)) { return { row: Number(item.row), col: Number(item.col) } } let bestOverlap = null let bestArea = 0 for (const cell of cells) { const cellBox = Array.isArray(cell?.bbox) && cell.bbox.length === 4 ? cell.bbox : null if (!cellBox) { continue } const area = intersectionArea(itemBox, cellBox) if (area > bestArea) { bestArea = area bestOverlap = cell } } if (bestOverlap && bestArea > 0) { return { row: bestOverlap.row, col: bestOverlap.col } } // Fallback: nearest cell center when overlap is zero due to OCR/table bbox drift. const [cx, cy] = centerOfBox(itemBox) let nearest = null let nearestDistSq = Number.POSITIVE_INFINITY for (const cell of cells) { const cellBox = Array.isArray(cell?.bbox) && cell.bbox.length === 4 ? cell.bbox : null if (!cellBox) { continue } const [ccx, ccy] = centerOfBox(cellBox) const dx = ccx - cx const dy = ccy - cy const distSq = dx * dx + dy * dy if (distSq < nearestDistSq) { nearestDistSq = distSq nearest = cell } } return nearest ? { row: nearest.row, col: nearest.col } : null } function normalizeBbox(bbox, maxWidth, maxHeight, minSize = 12) { let [x1, y1, x2, y2] = bbox.map((value) => Number(value) || 0) x1 = clamp(x1, 0, maxWidth) x2 = clamp(x2, 0, maxWidth) y1 = clamp(y1, 0, maxHeight) y2 = clamp(y2, 0, maxHeight) let left = Math.min(x1, x2) let right = Math.max(x1, x2) let top = Math.min(y1, y2) let bottom = Math.max(y1, y2) if ((right - left) < minSize) { const expand = (minSize - (right - left)) / 2 left = clamp(left - expand, 0, maxWidth) right = clamp(right + expand, 0, maxWidth) if ((right - left) < minSize) { if (left <= 0) right = clamp(minSize, 0, maxWidth) else left = clamp(maxWidth - minSize, 0, maxWidth) } } if ((bottom - top) < minSize) { const expand = (minSize - (bottom - top)) / 2 top = clamp(top - expand, 0, maxHeight) bottom = clamp(bottom + expand, 0, maxHeight) if ((bottom - top) < minSize) { if (top <= 0) bottom = clamp(minSize, 0, maxHeight) else top = clamp(maxHeight - minSize, 0, maxHeight) } } return [left, top, right, bottom] } function expandBbox(bbox, padding, maxWidth, maxHeight) { const pad = Math.max(0, Number(padding) || 0) if (!pad) { return normalizeBbox(bbox, maxWidth, maxHeight) } const [x1, y1, x2, y2] = bbox return normalizeBbox([x1 - pad, y1 - pad, x2 + pad, y2 + pad], maxWidth, maxHeight) } function resizeBbox(startBbox, startPoint, currentPoint, handle, maxWidth, maxHeight) { const [sx1, sy1, sx2, sy2] = startBbox const dx = currentPoint.x - startPoint.x const dy = currentPoint.y - startPoint.y let x1 = sx1 let y1 = sy1 let x2 = sx2 let y2 = sy2 if (handle.includes('w')) x1 = sx1 + dx if (handle.includes('e')) x2 = sx2 + dx if (handle.includes('n')) y1 = sy1 + dy if (handle.includes('s')) y2 = sy2 + dy return normalizeBbox([x1, y1, x2, y2], maxWidth, maxHeight) } function StudioPageCanvas({ page, tables, bboxOverrides, stageWidth, onImageError, onCopyText, hoveredCell, onHoverCellChange, onTableBBoxChange, onSelectionEnd, }) { const { state, dispatch } = useStudio() const [hoveredOcrKey, setHoveredOcrKey] = useState(null) // Use a ref (not state) for resizeDrag so that Konva mouse-move/up handlers // always read the *current* value — state updates are async and would lag. const resizeDragRef = useRef(null) const [, forceUpdate] = useState(0) const setResizeDrag = useCallback((val) => { resizeDragRef.current = typeof val === 'function' ? val(resizeDragRef.current) : val forceUpdate((n) => n + 1) }, []) const imageUrl = `${API}/api/image/${page.jobId}?page=${page.page_index}` const image = useLoadedImage(imageUrl, onImageError) const [pageWidth, pageHeight] = page.image_size || [0, 0] const safeWidth = Math.max(1, pageWidth || image?.naturalWidth || stageWidth || 1) const safeHeight = Math.max(1, pageHeight || image?.naturalHeight || 1) const drawnWidth = stageWidth ? Math.min(safeWidth, Math.max(240, stageWidth)) : safeWidth const scale = drawnWidth / safeWidth const stageHeight = safeHeight * scale const tableBboxFor = useCallback((table) => { const override = bboxOverrides?.[table.table_id] if (Array.isArray(override) && override.length === 4) { return override } return table.bbox }, [bboxOverrides]) const visibleTables = tables.filter((table) => (table.td_score ?? 0) >= state.tdThreshold) const isSelectionTool = state.activeTool === 'marquee' || state.activeTool === 'tableSelect' const selectionRect = state.selection && state.selection.pageIndex === page.page_index ? { left: Math.min(state.selection.start.x, state.selection.current.x), top: Math.min(state.selection.start.y, state.selection.current.y), right: Math.max(state.selection.start.x, state.selection.current.x), bottom: Math.max(state.selection.start.y, state.selection.current.y), } : null const handleMouseDown = (e) => { if (!isSelectionTool) return const stage = e.target.getStage() const pos = stage?.getPointerPosition() if (!pos) { return } dispatch({ type: 'setSelection', selection: { pageIndex: page.page_index, start: { x: pos.x / scale, y: pos.y / scale }, current: { x: pos.x / scale, y: pos.y / scale }, } }) } const handleMouseMove = (e) => { const drag = resizeDragRef.current if (drag) { const stage = e.target.getStage() const pos = stage?.getPointerPosition() if (!pos) { return } const nextBBox = resizeBbox( drag.startBBox, drag.startPoint, { x: pos.x / scale, y: pos.y / scale }, drag.handle, safeWidth, safeHeight, ) onTableBBoxChange?.(drag.tableId, nextBBox) return } if (!isSelectionTool || !state.selection || state.selection.pageIndex !== page.page_index) return const stage = e.target.getStage() const pos = stage?.getPointerPosition() if (!pos) { return } dispatch({ type: 'setSelection', selection: { ...state.selection, current: { x: pos.x / scale, y: pos.y / scale }, } }) } const handleMouseUp = () => { if (resizeDragRef.current) { setResizeDrag(null) return } if (!isSelectionTool || !state.selection || state.selection.pageIndex !== page.page_index) return const { start, current } = state.selection const rawBbox = [ Math.min(start.x, current.x), Math.min(start.y, current.y), Math.max(start.x, current.x), Math.max(start.y, current.y), ] // Ignore tiny / accidental drags if (Math.abs(rawBbox[2] - rawBbox[0]) <= 5 || Math.abs(rawBbox[3] - rawBbox[1]) <= 5) { dispatch({ type: 'setSelection', selection: null }) return } const bbox = expandBbox(rawBbox, state.selectionPadding, safeWidth, safeHeight) // For marquee OCR, fire immediately on mouse release — no button needed. if (state.activeTool === 'marquee') { onSelectionEnd?.({ bbox, clientX: ((bbox[0] + bbox[2]) / 2) * scale, clientY: (bbox[1] * scale) - 16, }) } // For tableSelect we leave the selection visible so the user can confirm. } useEffect(() => { const handleGlobalPointerUp = () => { if (resizeDragRef.current) { setResizeDrag(null) } } window.addEventListener('mouseup', handleGlobalPointerUp) return () => window.removeEventListener('mouseup', handleGlobalPointerUp) }, [setResizeDrag]) return (
{ onHoverCellChange?.(null) setResizeDrag(null) }} > {image && ( )} {visibleTables.map((table) => { const isActiveTable = state.activeTable === table.table_id const [x1, y1, x2, y2] = tableBboxFor(table) const width = (x2 - x1) * scale const height = (y2 - y1) * scale const x = x1 * scale const y = y1 * scale return ( dispatch({ type: 'setActiveTable', tableId: table.table_id })} /> ) })} {visibleTables.map((table) => ( ))} {visibleTables.flatMap((table) => (table.cells || []).map((cell) => { const key = getCellKey(table.table_id, cell.row, cell.col) const selected = isSameCell(state.activeCell, { tableId: table.table_id, row: cell.row, col: cell.col, }) const hovered = isSameCell(hoveredCell, { tableId: table.table_id, row: cell.row, col: cell.col, }) const [x1, y1, x2, y2] = cell.bbox const x = x1 * scale const y = y1 * scale const width = (x2 - x1) * scale const height = (y2 - y1) * scale const visibleStroke = state.showTableGrid ? '#E5E7EB' : 'transparent' const visibleFill = hovered ? 'rgba(59, 130, 246, 0.2)' : selected ? 'rgba(221, 107, 32, 0.18)' : state.showHeatmap && state.showTableGrid ? heatmapColor(cell.ocr_score, 0.14) : 'transparent' return ( { onHoverCellChange?.({ tableId: table.table_id, row: cell.row, col: cell.col, pageIndex: page.page_index, }) dispatch({ type: 'setActiveTable', tableId: table.table_id }) dispatch({ type: 'setActiveCell', cell: { tableId: table.table_id, row: cell.row, col: cell.col } }) }} onMouseLeave={() => onHoverCellChange?.(null)} onClick={() => { dispatch({ type: 'setActiveTable', tableId: table.table_id }) dispatch({ type: 'setActiveCell', cell: { tableId: table.table_id, row: cell.row, col: cell.col } }) }} listening perfectDrawEnabled={false} /> ) }))} {state.showTextBoxes && visibleTables.flatMap((table) => (table.ocr_items || []).map((item, index) => { const key = `${table.table_id}:${index}` const [x1, y1, x2, y2] = item.bbox const x = x1 * scale const y = y1 * scale const width = Math.max(1, (x2 - x1) * scale) const height = Math.max(1, (y2 - y1) * scale) const hovered = hoveredOcrKey === key return ( { setHoveredOcrKey(key) const targetCell = resolveOcrItemCell(table, item) if (targetCell) { dispatch({ type: 'setActiveTable', tableId: table.table_id }) onHoverCellChange?.({ tableId: table.table_id, row: targetCell.row, col: targetCell.col, pageIndex: page.page_index, }) } else { onHoverCellChange?.(null) } }} onMouseLeave={() => { setHoveredOcrKey((current) => (current === key ? null : current)) onHoverCellChange?.(null) }} onClick={(event) => onCopyText(item.text || '', event.evt.clientX, event.evt.clientY)} /> ) }))} {/* Corners rendered LAST so they stay interactive on top of cells */} {state.activeTool === 'pointer' && visibleTables.map((table) => { if (state.activeTable !== table.table_id) { return null } const [x1, y1, x2, y2] = tableBboxFor(table) const corners = [ { key: 'nw', x: x1, y: y1, cursor: 'nwse-resize' }, { key: 'ne', x: x2, y: y1, cursor: 'nesw-resize' }, { key: 'sw', x: x1, y: y2, cursor: 'nesw-resize' }, { key: 'se', x: x2, y: y2, cursor: 'nwse-resize' }, ] return corners.map((corner) => ( { event.cancelBubble = true const stage = event.target.getStage() const pos = stage?.getPointerPosition() if (!pos) { return } resizeDragRef.current = { tableId: table.table_id, handle: corner.key, startPoint: { x: pos.x / scale, y: pos.y / scale }, startBBox: tableBboxFor(table), } }} onMouseEnter={(event) => { const container = event.target.getStage()?.container() if (container) { container.style.cursor = corner.cursor } }} onMouseLeave={(event) => { const container = event.target.getStage()?.container() if (container) { container.style.cursor = 'default' } }} /> )) })} {selectionRect && ( )} {selectionRect && state.activeTool === 'tableSelect' && (
)}
) } export default function StudioDocumentPane({ annotation, jobId, onCopyText, hoveredCell, onHoverCellChange, onReprocessTable, onDeleteTable, onAddCell, onDeleteCell, onSelectionEnd, }) { const { state, dispatch } = useStudio() const scrollRef = useRef(null) const measureRef = useRef(null) const pageRefs = useRef([]) const [stageWidth, setStageWidth] = useState(720) const [tableBboxOverrides, setTableBboxOverrides] = useState({}) const [isReprocessingTable, setIsReprocessingTable] = useState(false) const [isDeletingTable, setIsDeletingTable] = useState(false) const [isAddingCell, setIsAddingCell] = useState(false) const [isDeletingCell, setIsDeletingCell] = useState(false) const [pendingDelete, setPendingDelete] = useState(false) // Palette State: Draggable & Resizable const [paletteSize, setPaletteSize] = useState({ width: 320, height: 420 }) const paletteRef = useRef(null) const dragControls = useDragControls() const handleResizeStart = (e) => { e.preventDefault() e.stopPropagation() const startX = e.pageX const startY = e.pageY const startWidth = paletteSize.width const startHeight = paletteSize.height const onMove = (moveEvent) => { setPaletteSize({ width: Math.max(240, startWidth + (moveEvent.pageX - startX)), height: Math.max(200, startHeight + (moveEvent.pageY - startY)) }) } const onUp = () => { window.removeEventListener('mousemove', onMove) window.removeEventListener('mouseup', onUp) } window.addEventListener('mousemove', onMove) window.addEventListener('mouseup', onUp) } const pages = useMemo( () => getPages(annotation).map((page) => ({ ...page, jobId })), [annotation, jobId], ) const activePageTables = useMemo( () => getTablesForPage(annotation, state.activePage), [annotation, state.activePage], ) const visibleActivePageTables = useMemo( () => activePageTables.filter((table) => (table.td_score ?? 0) >= state.tdThreshold), [activePageTables, state.tdThreshold], ) const selectedVisibleTable = useMemo( () => visibleActivePageTables.find((table) => table.table_id === state.activeTable) || null, [visibleActivePageTables, state.activeTable], ) const activeCellInSelectedTable = useMemo(() => { if (!selectedVisibleTable || !state.activeCell) { return null } if (state.activeCell.tableId !== selectedVisibleTable.table_id) { return null } return state.activeCell }, [selectedVisibleTable, state.activeCell]) const getTableBbox = useCallback((table) => { const override = tableBboxOverrides[table.table_id] return Array.isArray(override) && override.length === 4 ? override : table.bbox }, [tableBboxOverrides]) const clearTableBboxOverride = useCallback((tableId) => { setTableBboxOverrides((current) => { if (!(tableId in current)) { return current } const next = { ...current } delete next[tableId] return next }) }, []) useEffect(() => { const node = measureRef.current if (!node) { return undefined } const observer = new ResizeObserver((entries) => { const entry = entries[0] const width = entry?.contentRect?.width || node.clientWidth setStageWidth(clamp(width - 48, 280, 1200)) }) observer.observe(node) return () => observer.disconnect() }, []) useEffect(() => { const root = scrollRef.current const targets = pageRefs.current.filter(Boolean) if (!root || targets.length === 0) { return undefined } const observer = new IntersectionObserver( (entries) => { const visibleEntries = entries.filter((entry) => entry.isIntersecting && entry.intersectionRatio >= 0.5) if (!visibleEntries.length) { return } const mostVisible = visibleEntries.reduce((best, entry) => ( entry.intersectionRatio > best.intersectionRatio ? entry : best )) const pageIndex = Number(mostVisible.target.getAttribute('data-page-index')) if (!Number.isNaN(pageIndex) && pageIndex !== state.activePage) { dispatch({ type: 'setActivePage', page: pageIndex }) } }, { root, threshold: [0.5, 0.75], }, ) for (const target of targets) { observer.observe(target) } return () => observer.disconnect() }, [dispatch, state.activePage, pages.length]) useEffect(() => { const validIds = new Set((annotation?.tables || []).map((table) => table.table_id)) setTableBboxOverrides((current) => { const next = {} let changed = false for (const [tableId, bbox] of Object.entries(current)) { const numericId = Number(tableId) if (validIds.has(numericId)) { next[tableId] = bbox } else { changed = true } } return changed ? next : current }) }, [annotation]) const handleTableBBoxChange = useCallback((tableId, bbox) => { setTableBboxOverrides((current) => ({ ...current, [tableId]: bbox, })) }, []) const handleReprocessSelectedTable = useCallback(async () => { if (!selectedVisibleTable || !onReprocessTable || isReprocessingTable) { return } setIsReprocessingTable(true) try { const ok = await onReprocessTable({ tableId: selectedVisibleTable.table_id, pageIndex: state.activePage, bbox: getTableBbox(selectedVisibleTable), }) if (ok) { clearTableBboxOverride(selectedVisibleTable.table_id) } } finally { setIsReprocessingTable(false) } }, [ clearTableBboxOverride, getTableBbox, isReprocessingTable, onReprocessTable, selectedVisibleTable, state.activePage, ]) const handleDeleteSelectedTable = useCallback(async () => { if (!selectedVisibleTable || !onDeleteTable || isDeletingTable) { return } // First click arms the delete — second click (Confirm) fires it. if (!pendingDelete) { setPendingDelete(true) return } setPendingDelete(false) setIsDeletingTable(true) try { const ok = await onDeleteTable({ tableId: selectedVisibleTable.table_id, }) if (ok) { clearTableBboxOverride(selectedVisibleTable.table_id) } } finally { setIsDeletingTable(false) } }, [clearTableBboxOverride, isDeletingTable, onDeleteTable, pendingDelete, selectedVisibleTable]) const handleAddCellToSelectedTable = useCallback(async () => { if (!selectedVisibleTable || !onAddCell || isAddingCell || isDeletingCell || isDeletingTable || isReprocessingTable) { return } setIsAddingCell(true) try { await onAddCell({ tableId: selectedVisibleTable.table_id }) } finally { setIsAddingCell(false) } }, [isAddingCell, isDeletingCell, isDeletingTable, isReprocessingTable, onAddCell, selectedVisibleTable]) const handleDeleteSelectedCell = useCallback(async () => { if (!activeCellInSelectedTable || !onDeleteCell || isDeletingCell || isAddingCell || isDeletingTable || isReprocessingTable) { return } setIsDeletingCell(true) try { await onDeleteCell({ tableId: activeCellInSelectedTable.tableId, row: activeCellInSelectedTable.row, col: activeCellInSelectedTable.col, }) } finally { setIsDeletingCell(false) } }, [activeCellInSelectedTable, isAddingCell, isDeletingCell, isDeletingTable, isReprocessingTable, onDeleteCell]) return (
dragControls.start(e)} >
{pages.map((page, index) => (
{ pageRefs.current[index] = node }} className="document-page" data-page-index={page.page_index} >
Page {page.page_index + 1}
dispatch({ type: 'incrementErrors' })} onCopyText={onCopyText} hoveredCell={hoveredCell?.pageIndex === page.page_index ? hoveredCell : null} onHoverCellChange={onHoverCellChange} onTableBBoxChange={handleTableBBoxChange} onSelectionEnd={onSelectionEnd} />
))}
) }