From 90a2084f70da5629cfdd159fe445ab3e0eb296b1 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 4 Jun 2026 20:29:31 -0600 Subject: [PATCH] Allow adjusting, adding, and deleting bounding boxes. --- ui/src/components/BoundingBoxOverlay.tsx | 252 ++++++++++++++++++++++ ui/src/components/DatasetImageViewer.tsx | 261 ++++++++++++++++++++++- 2 files changed, 502 insertions(+), 11 deletions(-) diff --git a/ui/src/components/BoundingBoxOverlay.tsx b/ui/src/components/BoundingBoxOverlay.tsx index 669541bd..44ee961e 100644 --- a/ui/src/components/BoundingBoxOverlay.tsx +++ b/ui/src/components/BoundingBoxOverlay.tsx @@ -1,3 +1,4 @@ +import { useRef, useState } from 'react'; import classNames from 'classnames'; // A single bbox parsed from an Ideogram-style caption/prompt. Stored coords are @@ -11,6 +12,12 @@ export interface OverlayBox { type: 'obj' | 'text'; } +// Same as OverlayBox but carries the element's index in the caption's element +// array, so edits can be written back to the right element. +export interface EditableBox extends OverlayBox { + elementIndex: number; +} + // Returns the list of boxes if the text is an Ideogram bbox-JSON caption/prompt // with at least one bbox, otherwise null (normal captions/prompts get no overlay). export function parseBoundingBoxes(text: string): OverlayBox[] | null { @@ -36,6 +43,251 @@ export function parseBoundingBoxes(text: string): OverlayBox[] | null { return boxes.length > 0 ? boxes : null; } +// Build editable boxes (tagged with their element index) from an already-parsed +// caption object. Returns [] if the object has no usable elements/boxes. +export function extractBoxes(data: any): EditableBox[] { + const elements = data?.compositional_deconstruction?.elements; + if (!Array.isArray(elements)) return []; + const boxes: EditableBox[] = []; + elements.forEach((el: any, idx: number) => { + const bb = el?.bbox; + if (Array.isArray(bb) && bb.length === 4 && bb.every((n: any) => typeof n === 'number')) { + const isText = el.type === 'text'; + const label = (isText ? el.text : el.desc) ?? ''; + boxes.push({ + y1: bb[0], + x1: bb[1], + y2: bb[2], + x2: bb[3], + label: `${label}`, + type: isText ? 'text' : 'obj', + elementIndex: idx, + }); + } + }); + return boxes; +} + +// Parse the caption for EDITING: returns the parsed object plus its boxes. null if +// the text is not an Ideogram bbox-JSON with at least one box. +export function parseCaptionForEditing(text: string): { data: any; boxes: EditableBox[] } | null { + const trimmed = text.trim(); + if (!trimmed.startsWith('{')) return null; + let data: any; + try { + data = JSON.parse(trimmed); + } catch { + return null; + } + const boxes = extractBoxes(data); + return boxes.length > 0 ? { data, boxes } : null; +} + +export interface BoxCoords { + y1: number; + x1: number; + y2: number; + x2: number; +} + +// Clamp/round/normalize coords to integer 0-1000 with y1 Math.max(0, Math.min(1000, Math.round(v))); + const x1 = cl(Math.min(b.x1, b.x2)); + const x2 = cl(Math.max(b.x1, b.x2)); + const y1 = cl(Math.min(b.y1, b.y2)); + const y2 = cl(Math.max(b.y1, b.y2)); + return { y1, x1, y2, x2 }; +} + +type DragMode = 'move' | 'nw' | 'ne' | 'sw' | 'se'; +const MIN_SIZE = 8; // minimum normalized box span so boxes can't collapse + +// Editable overlay: click to select, drag a box body to move, drag a corner +// handle to resize, ✕ to delete. With `drawing` on, drag on empty space to +// create a new box. All geometry commits fire on pointer-up. Place inside the +// same `relative` image-wrapping container as BoundingBoxOverlay; coordinate +// math uses this layer's own rect, so it's correct at any zoom. +export function BoundingBoxEditor({ + boxes, + selectedIndex, + drawing, + onSelect, + onChangeBox, + onCreateBox, +}: { + boxes: EditableBox[]; + selectedIndex: number | null; + drawing: boolean; + onSelect: (elementIndex: number | null) => void; + onChangeBox: (elementIndex: number, box: BoxCoords) => void; + onCreateBox: (box: BoxCoords) => void; +}) { + const rootRef = useRef(null); + const [draft, setDraft] = useState<{ index: number; box: OverlayBox } | null>(null); + const [drawDraft, setDrawDraft] = useState<{ x1: number; y1: number; x2: number; y2: number } | null>(null); + + const toNorm = (clientX: number, clientY: number) => { + const rect = rootRef.current!.getBoundingClientRect(); + return { + x: Math.max(0, Math.min(1000, ((clientX - rect.left) / rect.width) * 1000)), + y: Math.max(0, Math.min(1000, ((clientY - rect.top) / rect.height) * 1000)), + }; + }; + + const startDrag = (e: React.PointerEvent, box: EditableBox, mode: DragMode) => { + // Stop the event from reaching the zoom/pan wrapper so dragging never pans. + e.preventDefault(); + e.stopPropagation(); + onSelect(box.elementIndex); + const root = rootRef.current; + if (!root) return; + const startX = e.clientX; + const startY = e.clientY; + const start: OverlayBox = { ...box }; + let current: OverlayBox = { ...box }; + + const move = (ev: PointerEvent) => { + const rect = root.getBoundingClientRect(); + if (!rect.width || !rect.height) return; + const dx = ((ev.clientX - startX) / rect.width) * 1000; + const dy = ((ev.clientY - startY) / rect.height) * 1000; + let { x1, y1, x2, y2 } = start; + if (mode === 'move') { + const w = x2 - x1; + const h = y2 - y1; + const nx1 = Math.max(0, Math.min(1000 - w, x1 + dx)); + const ny1 = Math.max(0, Math.min(1000 - h, y1 + dy)); + x1 = nx1; + y1 = ny1; + x2 = nx1 + w; + y2 = ny1 + h; + } else { + if (mode.includes('w')) x1 = Math.max(0, Math.min(x2 - MIN_SIZE, x1 + dx)); + if (mode.includes('e')) x2 = Math.min(1000, Math.max(x1 + MIN_SIZE, x2 + dx)); + if (mode.includes('n')) y1 = Math.max(0, Math.min(y2 - MIN_SIZE, y1 + dy)); + if (mode.includes('s')) y2 = Math.min(1000, Math.max(y1 + MIN_SIZE, y2 + dy)); + } + current = { ...start, x1, y1, x2, y2 }; + setDraft({ index: box.elementIndex, box: current }); + }; + + const up = () => { + window.removeEventListener('pointermove', move); + window.removeEventListener('pointerup', up); + setDraft(null); + onChangeBox(box.elementIndex, normalizeBox(current)); + }; + + window.addEventListener('pointermove', move); + window.addEventListener('pointerup', up); + }; + + // Pointer-down on the empty root: draw a new box (drawing mode) or deselect. + const onRootPointerDown = (e: React.PointerEvent) => { + if (!drawing) { + onSelect(null); + return; + } + e.preventDefault(); + e.stopPropagation(); + if (!rootRef.current) return; + const p0 = toNorm(e.clientX, e.clientY); + let cur = { x1: p0.x, y1: p0.y, x2: p0.x, y2: p0.y }; + setDrawDraft(cur); + + const move = (ev: PointerEvent) => { + const p = toNorm(ev.clientX, ev.clientY); + cur = { x1: p0.x, y1: p0.y, x2: p.x, y2: p.y }; + setDrawDraft({ ...cur }); + }; + const up = () => { + window.removeEventListener('pointermove', move); + window.removeEventListener('pointerup', up); + setDrawDraft(null); + const nb = normalizeBox(cur); + if (nb.x2 - nb.x1 >= MIN_SIZE && nb.y2 - nb.y1 >= MIN_SIZE) onCreateBox(nb); + }; + window.addEventListener('pointermove', move); + window.addEventListener('pointerup', up); + }; + + const handles: { mode: DragMode; cls: string }[] = [ + { mode: 'nw', cls: 'top-0 left-0 -translate-x-1/2 -translate-y-1/2 cursor-nwse-resize' }, + { mode: 'ne', cls: 'top-0 right-0 translate-x-1/2 -translate-y-1/2 cursor-nesw-resize' }, + { mode: 'sw', cls: 'bottom-0 left-0 -translate-x-1/2 translate-y-1/2 cursor-nesw-resize' }, + { mode: 'se', cls: 'bottom-0 right-0 translate-x-1/2 translate-y-1/2 cursor-nwse-resize' }, + ]; + + return ( +
+ {boxes.map(box => { + const pos = draft && draft.index === box.elementIndex ? draft.box : box; + const isObj = box.type === 'obj'; + const selected = selectedIndex === box.elementIndex; + return ( +
startDrag(e, box, 'move')} + className={classNames('absolute border-2', { + 'pointer-events-none': drawing, + 'cursor-move touch-none': !drawing, + 'border-white ring-2 ring-blue-400': selected && !drawing, + 'border-cyan-400': isObj && !(selected && !drawing), + 'border-amber-400': !isObj && !(selected && !drawing), + })} + style={{ + left: `${pos.x1 / 10}%`, + top: `${pos.y1 / 10}%`, + width: `${(pos.x2 - pos.x1) / 10}%`, + height: `${(pos.y2 - pos.y1) / 10}%`, + }} + > + {box.label && ( + + {box.label} + + )} + {selected && + !drawing && + handles.map(h => ( +
startDrag(e, { ...box, ...pos }, h.mode)} + className={classNames( + 'absolute w-3 h-3 rounded-sm border border-gray-900 bg-white touch-none', + h.cls, + )} + /> + ))} +
+ ); + })} + {drawDraft && ( +
+ )} +
+ ); +} + // Absolute overlay layer of boxes. Place inside a `relative` container that wraps // the image (e.g. inside react-zoom-pan-pinch's TransformComponent) so boxes // track the image during zoom/pan. Coords are percentages of the frame, so this diff --git a/ui/src/components/DatasetImageViewer.tsx b/ui/src/components/DatasetImageViewer.tsx index 5619564b..f973a00a 100644 --- a/ui/src/components/DatasetImageViewer.tsx +++ b/ui/src/components/DatasetImageViewer.tsx @@ -2,7 +2,7 @@ import { useState, useEffect, useMemo, useCallback, useRef } from 'react'; import { createPortal } from 'react-dom'; import { Dialog, DialogBackdrop, DialogPanel } from '@headlessui/react'; -import { Cog, SquareDashed } from 'lucide-react'; +import { Cog, SquareDashed, Pencil } from 'lucide-react'; import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/react'; import classNames from 'classnames'; import { openConfirm } from './ConfirmModal'; @@ -10,7 +10,15 @@ import { apiClient } from '@/utils/api'; import { isVideo, isAudio } from '@/utils/basic'; import AudioPlayer from './AudioPlayer'; import { TransformWrapper, TransformComponent } from 'react-zoom-pan-pinch'; -import BoundingBoxOverlay, { parseBoundingBoxes } from './BoundingBoxOverlay'; +import BoundingBoxOverlay, { BoundingBoxEditor, parseBoundingBoxes, extractBoxes } from './BoundingBoxOverlay'; + +function safeParse(text: string): any { + try { + return JSON.parse(text); + } catch { + return null; + } +} interface Props { imgPath: string | null; // current image path @@ -27,6 +35,9 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre const [savedCaption, setSavedCaption] = useState(''); const [isCaptionLoaded, setIsCaptionLoaded] = useState(false); const [showBoxes, setShowBoxes] = useState(false); + const [isEditingBoxes, setIsEditingBoxes] = useState(false); + const [selectedBoxIndex, setSelectedBoxIndex] = useState(null); + const [isDrawing, setIsDrawing] = useState(false); const captionRef = useRef(''); const savedCaptionRef = useRef(''); const currentImgPathRef = useRef(null); @@ -34,6 +45,13 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre useEffect(() => setMounted(true), []); + // Leave box-edit mode whenever the image changes, to avoid accidental edits. + useEffect(() => { + setIsEditingBoxes(false); + setSelectedBoxIndex(null); + setIsDrawing(false); + }, [imgPath]); + // open/close based on external value useEffect(() => { setIsOpen(Boolean(imgPath)); @@ -88,6 +106,14 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre [onCaptionSaved], ); + // Stable handle to the latest saveCaptionForPath so the fetch effect doesn't + // re-run (and re-fetch, blanking the caption) every time a save changes its + // identity via the parent's onCaptionSaved. + const saveCaptionForPathRef = useRef(saveCaptionForPath); + useEffect(() => { + saveCaptionForPathRef.current = saveCaptionForPath; + }, [saveCaptionForPath]); + const saveCaption = useCallback(() => { if (!imgPath) return; saveCaptionForPath(imgPath, caption, savedCaption); @@ -97,7 +123,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre useEffect(() => { const previousPath = currentImgPathRef.current; if (previousPath && previousPath !== imgPath) { - saveCaptionForPath(previousPath, captionRef.current, savedCaptionRef.current); + saveCaptionForPathRef.current(previousPath, captionRef.current, savedCaptionRef.current); } currentImgPathRef.current = imgPath; @@ -137,7 +163,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre return () => { controller.abort(); }; - }, [imgPath, saveCaptionForPath]); + }, [imgPath]); // Save any pending caption when the viewer fully unmounts useEffect(() => { @@ -194,6 +220,82 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre }); }, [imgPath, onChange, refreshImages]); + // Mutate the caption JSON's element array, updating the textarea ONLY (no + // network). Persisting happens via the Save button / nav auto-save. Returns + // whatever the mutator returns. + const editCaption = useCallback( + (fn: (elements: any[], data: any) => any): any => { + let data: any; + try { + data = JSON.parse(caption); + } catch { + return undefined; + } + const elements = data?.compositional_deconstruction?.elements; + if (!Array.isArray(elements)) return undefined; + const result = fn(elements, data); + setCaption(JSON.stringify(data, null, 2)); + return result; + }, + [caption], + ); + + const handleBoxChange = useCallback( + (elementIndex: number, box: { y1: number; x1: number; y2: number; x2: number }) => { + editCaption(els => { + if (els[elementIndex]) els[elementIndex].bbox = [box.y1, box.x1, box.y2, box.x2]; + }); + }, + [editCaption], + ); + + const handleDeleteBox = useCallback( + (elementIndex: number) => { + editCaption(els => { + els.splice(elementIndex, 1); + }); + setSelectedBoxIndex(null); + }, + [editCaption], + ); + + const handleCreateBox = useCallback( + (box: { y1: number; x1: number; y2: number; x2: number }) => { + const newIndex = editCaption(els => { + els.push({ type: 'obj', bbox: [box.y1, box.x1, box.y2, box.x2], desc: '' }); + return els.length - 1; + }); + setSelectedBoxIndex(typeof newIndex === 'number' ? newIndex : null); + setIsDrawing(false); + }, + [editCaption], + ); + + const handleFieldChange = useCallback( + (field: 'desc' | 'text', value: string) => { + editCaption(els => { + if (selectedBoxIndex != null && els[selectedBoxIndex]) els[selectedBoxIndex][field] = value; + }); + }, + [editCaption, selectedBoxIndex], + ); + + const handleTypeChange = useCallback( + (type: 'obj' | 'text') => { + editCaption(els => { + const el = selectedBoxIndex != null ? els[selectedBoxIndex] : null; + if (!el) return; + el.type = type; + if (type === 'text') { + if (el.text == null) el.text = ''; + } else { + delete el.text; + } + }); + }, + [editCaption, selectedBoxIndex], + ); + // keyboard events while open — skip nav while caption textarea is focused useEffect(() => { const handleKeyDown = (event: KeyboardEvent) => { @@ -219,7 +321,12 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre break; case 'Delete': case 'Backspace': - handleDelete(); + // While editing boxes, Delete removes the selected box (never the image). + if (isEditingBoxes) { + if (selectedBoxIndex != null) handleDeleteBox(selectedBoxIndex); + } else { + handleDelete(); + } break; default: break; @@ -227,7 +334,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre }; window.addEventListener('keydown', handleKeyDown); return () => window.removeEventListener('keydown', handleKeyDown); - }, [isOpen, onCancel, handlePrev, handleNext, handleDelete]); + }, [isOpen, onCancel, handlePrev, handleNext, handleDelete, isEditingBoxes, selectedBoxIndex, handleDeleteBox]); // Touch swipe navigation const touchStartRef = useRef<{ x: number; y: number } | null>(null); @@ -285,6 +392,13 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre const boundingBoxes = useMemo(() => parseBoundingBoxes(caption), [caption]); const canShowBoxes = Boolean(boundingBoxes && imgPath && !isAudio(imgPath) && !isVideo(imgPath)); + // Boxes and the selected element are derived from the (locally edited) caption. + const editBoxes = useMemo(() => extractBoxes(safeParse(caption)), [caption]); + const selectedElement = useMemo(() => { + if (selectedBoxIndex == null) return null; + return safeParse(caption)?.compositional_deconstruction?.elements?.[selectedBoxIndex] ?? null; + }, [caption, selectedBoxIndex]); + if (!mounted) return null; return createPortal( @@ -323,9 +437,9 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre initialScale={1} minScale={1} maxScale={6} - doubleClick={{ mode: 'toggle', step: 2 }} - wheel={{ step: 0.2 }} - panning={{ disabled: false, allowRightClickPan: false }} + doubleClick={{ mode: 'toggle', step: 2, disabled: isEditingBoxes }} + wheel={{ step: 0.2, disabled: isEditingBoxes }} + panning={{ disabled: isEditingBoxes, allowRightClickPan: false }} onTransform={(_ref, state) => { zoomedRef.current = state.scale > 1.01; }} @@ -338,7 +452,18 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre draggable={false} className="w-auto h-auto max-w-full sm:max-w-[95vw] max-h-[70vh] object-contain select-none !pointer-events-auto" /> - {showBoxes && boundingBoxes && } + {isEditingBoxes ? ( + + ) : ( + showBoxes && boundingBoxes && + )}
@@ -371,12 +496,106 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre disabled={!isCaptionLoaded} />
+ {isEditingBoxes && ( +
+
+ + + {isDrawing + ? 'Drag on the image to draw a new box' + : 'Click a box to select; drag to move, handles to resize'} + + +
+ {selectedElement && ( +
+
+ Type: + + + +
+ {selectedElement.type === 'text' && ( +