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 (
{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}
/>
))}
)
}