Allow adjusting, adding, and deleting bounding boxes.

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<y2, x1<x2.
+function normalizeBox(b: { x1: number; y1: number; x2: number; y2: number }): BoxCoords {
+  const cl = (v: number) => 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<HTMLDivElement>(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 (
+    <div
+      ref={rootRef}
+      onPointerDown={onRootPointerDown}
+      className={classNames('absolute inset-0', drawing ? 'cursor-crosshair' : '')}
+    >
+      {boxes.map(box => {
+        const pos = draft && draft.index === box.elementIndex ? draft.box : box;
+        const isObj = box.type === 'obj';
+        const selected = selectedIndex === box.elementIndex;
+        return (
+          <div
+            key={box.elementIndex}
+            onPointerDown={drawing ? undefined : e => 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 && (
+              <span
+                title={box.label}
+                className={classNames(
+                  'absolute top-0 left-0 max-w-full px-1 py-0.5 text-[9px] leading-tight font-medium whitespace-pre-line break-words line-clamp-2 text-gray-900 pointer-events-none',
+                  { 'bg-cyan-400/90': isObj, 'bg-amber-400/90': !isObj },
+                )}
+              >
+                {box.label}
+              </span>
+            )}
+            {selected &&
+              !drawing &&
+              handles.map(h => (
+                <div
+                  key={h.mode}
+                  onPointerDown={e => 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,
+                  )}
+                />
+              ))}
+          </div>
+        );
+      })}
+      {drawDraft && (
+        <div
+          className="absolute border-2 border-dashed border-white bg-white/10 pointer-events-none"
+          style={{
+            left: `${Math.min(drawDraft.x1, drawDraft.x2) / 10}%`,
+            top: `${Math.min(drawDraft.y1, drawDraft.y2) / 10}%`,
+            width: `${Math.abs(drawDraft.x2 - drawDraft.x1) / 10}%`,
+            height: `${Math.abs(drawDraft.y2 - drawDraft.y1) / 10}%`,
+          }}
+        />
+      )}
+    </div>
+  );
+}
+
 // 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

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<string>('');
   const [isCaptionLoaded, setIsCaptionLoaded] = useState<boolean>(false);
   const [showBoxes, setShowBoxes] = useState<boolean>(false);
+  const [isEditingBoxes, setIsEditingBoxes] = useState<boolean>(false);
+  const [selectedBoxIndex, setSelectedBoxIndex] = useState<number | null>(null);
+  const [isDrawing, setIsDrawing] = useState<boolean>(false);
   const captionRef = useRef<string>('');
   const savedCaptionRef = useRef<string>('');
   const currentImgPathRef = useRef<string | null>(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 }}
+                    doubleClick={{ mode: 'toggle', step: 2, disabled: isEditingBoxes }}
-                    wheel={{ step: 0.2 }}
+                    wheel={{ step: 0.2, disabled: isEditingBoxes }}
-                    panning={{ disabled: false, allowRightClickPan: false }}
+                    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 && <BoundingBoxOverlay boxes={boundingBoxes} />}
+                        {isEditingBoxes ? (
+                          <BoundingBoxEditor
+                            boxes={editBoxes}
+                            selectedIndex={selectedBoxIndex}
+                            drawing={isDrawing}
+                            onSelect={setSelectedBoxIndex}
+                            onChangeBox={handleBoxChange}
+                            onCreateBox={handleCreateBox}
+                          />
+                        ) : (
+                          showBoxes && boundingBoxes && <BoundingBoxOverlay boxes={boundingBoxes} />
+                        )}
                       </div>
                     </TransformComponent>
                   </TransformWrapper>
@@ -371,12 +496,106 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
                   disabled={!isCaptionLoaded}
                 />
               </div>
+              {isEditingBoxes && (
+                <div className="rounded border border-gray-700 bg-gray-900 p-2 flex flex-col gap-2 text-xs">
+                  <div className="flex items-center gap-2">
+                    <button
+                      type="button"
+                      onClick={() => setIsDrawing(d => !d)}
+                      className={classNames('px-2 py-1 rounded border', {
+                        'bg-blue-600 border-blue-500 text-white': isDrawing,
+                        'border-gray-600 text-gray-300 hover:bg-gray-800': !isDrawing,
+                      })}
+                    >
+                      {isDrawing ? 'Cancel' : '+ Add Box'}
+                    </button>
+                    <span className="text-gray-500">
+                      {isDrawing
+                        ? 'Drag on the image to draw a new box'
+                        : 'Click a box to select; drag to move, handles to resize'}
+                    </span>
+                    <button
+                      type="button"
+                      onClick={saveCaption}
+                      disabled={isCaptionCurrent}
+                      className={classNames('ml-auto px-3 py-1 rounded border', {
+                        'bg-green-600 border-green-500 text-white hover:bg-green-500': !isCaptionCurrent,
+                        'border-gray-700 text-gray-500 cursor-default': isCaptionCurrent,
+                      })}
+                    >
+                      {isCaptionCurrent ? 'Saved' : 'Save'}
+                    </button>
+                  </div>
+                  {selectedElement && (
+                    <div className="flex flex-col gap-2 border-t border-gray-700 pt-2">
+                      <div className="flex items-center gap-2">
+                        <span className="text-gray-400">Type:</span>
+                        <button
+                          type="button"
+                          onClick={() => handleTypeChange('obj')}
+                          className={classNames('px-2 py-0.5 rounded border', {
+                            'bg-cyan-600 border-cyan-500 text-white': selectedElement.type !== 'text',
+                            'border-gray-600 text-gray-300 hover:bg-gray-800': selectedElement.type === 'text',
+                          })}
+                        >
+                          Object
+                        </button>
+                        <button
+                          type="button"
+                          onClick={() => handleTypeChange('text')}
+                          className={classNames('px-2 py-0.5 rounded border', {
+                            'bg-amber-600 border-amber-500 text-white': selectedElement.type === 'text',
+                            'border-gray-600 text-gray-300 hover:bg-gray-800': selectedElement.type !== 'text',
+                          })}
+                        >
+                          Text
+                        </button>
+                        <button
+                          type="button"
+                          onClick={() => handleDeleteBox(selectedBoxIndex!)}
+                          className="ml-auto px-2 py-0.5 rounded border border-red-700 text-red-400 hover:bg-red-900/40"
+                        >
+                          Delete
+                        </button>
+                      </div>
+                      {selectedElement.type === 'text' && (
+                        <label className="flex flex-col gap-1">
+                          <span className="text-gray-400">Text (shown in image)</span>
+                          <textarea
+                            className="w-full bg-gray-950 text-gray-100 rounded border border-gray-700 p-1 resize-none outline-none focus:border-blue-500"
+                            rows={2}
+                            value={selectedElement.text ?? ''}
+                            onChange={e => handleFieldChange('text', e.target.value)}
+                          />
+                        </label>
+                      )}
+                      <label className="flex flex-col gap-1">
+                        <span className="text-gray-400">Description</span>
+                        <textarea
+                          className="w-full bg-gray-950 text-gray-100 rounded border border-gray-700 p-1 resize-none outline-none focus:border-blue-500"
+                          rows={2}
+                          value={selectedElement.desc ?? ''}
+                          onChange={e => handleFieldChange('desc', e.target.value)}
+                        />
+                      </label>
+                    </div>
+                  )}
+                </div>
+              )}
             </div>
             <div className="absolute top-2 right-2 flex items-center gap-2 z-20">
               {canShowBoxes && (
                 <button
                   type="button"
-                  onClick={() => setShowBoxes(v => !v)}
+                  onClick={() => {
+                    const next = !showBoxes;
+                    setShowBoxes(next);
+                    if (!next) {
+                      setIsEditingBoxes(false);
+                      setSelectedBoxIndex(null);
+                      setIsDrawing(false);
+                    }
+                  }}
                   title={showBoxes ? 'Hide bounding boxes' : 'Show bounding boxes'}
                   className={classNames('bg-gray-900 rounded-full p-1 leading-[0px] hover:opacity-100', {
                     'opacity-100 text-blue-400': showBoxes,
@@ -386,6 +605,26 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
                   <SquareDashed />
                 </button>
               )}
+              {((canShowBoxes && showBoxes) || isEditingBoxes) && (
+                <button
+                  type="button"
+                  onClick={() => {
+                    const next = !isEditingBoxes;
+                    setIsEditingBoxes(next);
+                    if (!next) {
+                      setSelectedBoxIndex(null);
+                      setIsDrawing(false);
+                    }
+                  }}
+                  title={isEditingBoxes ? 'Done editing boxes' : 'Edit bounding boxes'}
+                  className={classNames('bg-gray-900 rounded-full p-1 leading-[0px] hover:opacity-100', {
+                    'opacity-100 text-blue-400': isEditingBoxes,
+                    'opacity-50': !isEditingBoxes,
+                  })}
+                >
+                  <Pencil />
+                </button>
+              )}
               <div className="bg-gray-900 rounded-full p-1 leading-[0px] opacity-50 hover:opacity-100">
                 <Menu>
                   <MenuButton>