Added ability to set the caption extention in dataset viewer, captioner, and trainer so one dataset can have multiple caption styles in different files with different extensions. Added dataset caption template for a blank ideogram 4 formatted template.

2026-06-06 08:32:24 -06:00
parent 10cdeb394e
commit 41157b460c
19 changed files with 219 additions and 270 deletions
--- a/ui/src/components/AutoCaptionButton.tsx
+++ b/ui/src/components/AutoCaptionButton.tsx
@@ -8,9 +8,10 @@ import { Loader2 } from 'lucide-react';
 type AutoCaptionButtonProps = {
  datasetPath: string;
  setIsAutoCaptioning?: (isAutoCaptioning: boolean) => void;
+  captionExt?: string;
 };

-export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning }: AutoCaptionButtonProps) {
+export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning, captionExt }: AutoCaptionButtonProps) {
  const { job, status, refreshJob } = useJobByRef(datasetPath, 5000);
  useEffect(() => {
    if (setIsAutoCaptioning) {
@@ -34,9 +35,13 @@ export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning }:
    <Button
      className="text-white bg-blue-600 px-2 sm:px-3 py-1 rounded-md mr-1 sm:mr-2 text-sm sm:text-base whitespace-nowrap"
      onClick={() =>
-        openCaptionDatasetModal(datasetPath, () => {
-          refreshJob();
-        })
+        openCaptionDatasetModal(
+          datasetPath,
+          () => {
+            refreshJob();
+          },
+          { defaultCaptionExt: captionExt },
+        )
      }
    >
      <span className="hidden sm:inline">Auto Caption</span>
--- a/ui/src/components/CaptionDatasetModal.tsx
+++ b/ui/src/components/CaptionDatasetModal.tsx
@@ -22,6 +22,7 @@ export interface CaptionDatasetModalState {
  datasetPath: string;
  jobId?: string | null;
  cloneId?: string | null;
+  defaultCaptionExt?: string;
  onClose?: () => void;
 }

@@ -30,13 +31,14 @@ export const captionDatasetModalState = createGlobalState<CaptionDatasetModalSta
 export const openCaptionDatasetModal = (
  datasetPath: string,
  onClose?: () => void,
-  options?: { jobId?: string | null; cloneId?: string | null },
+  options?: { jobId?: string | null; cloneId?: string | null; defaultCaptionExt?: string },
 ) => {
  captionDatasetModalState.set({
    datasetPath,
    onClose,
    jobId: options?.jobId ?? null,
    cloneId: options?.cloneId ?? null,
+    defaultCaptionExt: options?.defaultCaptionExt,
  });
 };

@@ -64,6 +66,11 @@ export const CaptionDatasetModal: React.FC = () => {
    if (modalInfo?.datasetPath) {
      setJobConfig(modalInfo.datasetPath, 'config.process[0].caption.path_to_caption');
    }
+    // default the caption extension to the current header selection. Editing it
+    // here only changes this job, not the header (separate state).
+    if (modalInfo?.defaultCaptionExt) {
+      setJobConfig(modalInfo.defaultCaptionExt, 'config.process[0].caption.caption_extension');
+    }
  }, [modalInfo]);

  // clone existing caption job
--- a/ui/src/components/CaptionSimpleJob.tsx
+++ b/ui/src/components/CaptionSimpleJob.tsx
@@ -121,6 +121,20 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
            }}
            options={quantizationOptions}
          />
+          <div className="mt-4">
+            <CreatableSelectInput
+              label="Caption Extension"
+              value={jobConfig.config.process[0].caption.caption_extension || 'txt'}
+              onChange={value => {
+                setJobConfig(value, 'config.process[0].caption.caption_extension');
+              }}
+              options={[
+                { value: 'txt', label: 'txt' },
+                { value: 'json', label: 'json' },
+                { value: 'caption', label: 'caption' },
+              ]}
+            />
+          </div>
          {additionalSections.includes('caption.max_res') && (
            <div className="mt-4">
              <SelectInput
--- a/ui/src/components/DatasetImageCard.tsx
+++ b/ui/src/components/DatasetImageCard.tsx
@@ -18,6 +18,7 @@ interface DatasetImageCardProps {
  captionRefreshKey?: number;
  observerRoot?: Element | null;
  rootMargin?: string;
+  captionExt?: string;
 }

 const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
@@ -31,6 +32,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  captionRefreshKey = 0,
  observerRoot = null,
  rootMargin = '200px 0px',
+  captionExt = 'txt',
 }) => {
  const [loaded, setLoaded] = useState<boolean>(false);
  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
@@ -110,6 +112,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(
    isVisible ? imageUrl : null,
    combinedRefreshKey,
+    captionExt,
  );

  const [caption, setCaption] = useState<string>('');
@@ -138,10 +141,10 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
      return;
    }
    apiClient
-      .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption })
+      .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption, ext: captionExt })
      .then(() => {
        setSavedCaption(trimmedCaption);
-        setCachedCaption(imageUrl, trimmedCaption);
+        setCachedCaption(imageUrl, trimmedCaption, captionExt);
        dirtyRef.current = false;
      })
      .catch(error => {
@@ -150,19 +153,19 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  };

  // Save any pending edit if the card unmounts (e.g. scrolled out of the virtualized window).
-  const latestRef = useRef({ caption, savedCaption, imageUrl });
+  const latestRef = useRef({ caption, savedCaption, imageUrl, captionExt });
  useEffect(() => {
-    latestRef.current = { caption, savedCaption, imageUrl };
+    latestRef.current = { caption, savedCaption, imageUrl, captionExt };
  });
  useEffect(() => {
    return () => {
      if (!dirtyRef.current) return;
-      const { caption: c, savedCaption: s, imageUrl: url } = latestRef.current;
+      const { caption: c, savedCaption: s, imageUrl: url, captionExt: ext } = latestRef.current;
      const trimmed = c.trim();
      if (trimmed === s) return;
      apiClient
-        .post('/api/img/caption', { imgPath: url, caption: trimmed })
-        .then(() => setCachedCaption(url, trimmed))
+        .post('/api/img/caption', { imgPath: url, caption: trimmed, ext })
+        .then(() => setCachedCaption(url, trimmed, ext))
        .catch(err => console.error('Error saving caption on unmount:', err));
    };
  }, []);
--- a/ui/src/components/DatasetImageViewer.tsx
+++ b/ui/src/components/DatasetImageViewer.tsx
@@ -12,6 +12,7 @@ import AudioPlayer from './AudioPlayer';
 import { TransformWrapper, TransformComponent } from 'react-zoom-pan-pinch';
 import { BoundingBoxEditor, parseBoundingBoxes, extractBoxes } from './BoundingBoxOverlay';
 import IdeogramCaptionSidebar, { isIdeogramCaption } from './IdeogramCaptionSidebar';
+import datasetTemplates from '@/helpers/datasetTemplates';

 function safeParse(text: string): any {
  try {
@@ -27,9 +28,17 @@ interface Props {
  onChange: (nextPath: string | null) => void; // parent setter
  refreshImages?: () => void;
  onCaptionSaved?: (imgPath: string, caption: string) => void;
+  captionExt?: string;
 }

-export default function DatasetImageViewer({ imgPath, imageList, onChange, refreshImages, onCaptionSaved }: Props) {
+export default function DatasetImageViewer({
+  imgPath,
+  imageList,
+  onChange,
+  refreshImages,
+  onCaptionSaved,
+  captionExt = 'txt',
+}: Props) {
  const [mounted, setMounted] = useState(false);
  const [isOpen, setIsOpen] = useState(Boolean(imgPath));
  const [caption, setCaption] = useState<string>('');
@@ -98,7 +107,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
      const trimmed = value.trim();
      if (trimmed === prevSaved.trim()) return;
      apiClient
-        .post('/api/img/caption', { imgPath: path, caption: trimmed })
+        .post('/api/img/caption', { imgPath: path, caption: trimmed, ext: captionExt })
        .then(() => {
          if (currentImgPathRef.current === path) {
            setSavedCaption(trimmed);
@@ -109,7 +118,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
          console.error('Error saving caption:', error);
        });
    },
-    [onCaptionSaved],
+    [onCaptionSaved, captionExt],
  );

  // Stable handle to the latest saveCaptionForPath so the fetch effect doesn't
@@ -151,7 +160,11 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
      // transformResponse identity: keep the caption as a raw string. Axios's
      // default parses any JSON-looking body into an object (our bbox captions
      // are JSON), which would render as "[object Object]".
-      .post('/api/caption/get', { imgPath }, { signal: controller.signal, transformResponse: [d => d] })
+      .post(
+        '/api/caption/get',
+        { imgPath, ext: captionExt },
+        { signal: controller.signal, transformResponse: [d => d] },
+      )
      .then(res => res.data)
      .then(data => {
        if (controller.signal.aborted) return;
@@ -169,7 +182,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
    return () => {
      controller.abort();
    };
-  }, [imgPath]);
+  }, [imgPath, captionExt]);

  // Save any pending caption when the viewer fully unmounts
  useEffect(() => {
@@ -509,6 +522,23 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
                  {currentIndex >= 0 ? `${currentIndex + 1} / ${imageList.length}` : ''}
                </div>
              </div>
+              {isCaptionLoaded && caption.trim() === '' && (
+                <select
+                  className="w-full bg-gray-900 border border-gray-700 text-gray-100 text-sm rounded p-2 outline-none focus:ring-0 focus:outline-none"
+                  value=""
+                  onChange={e => {
+                    const template = datasetTemplates[e.target.value];
+                    if (template) setCaption(template.trim());
+                  }}
+                >
+                  <option value="">Templates...</option>
+                  {Object.keys(datasetTemplates).map(key => (
+                    <option key={key} value={key}>
+                      {key}
+                    </option>
+                  ))}
+                </select>
+              )}
              {isIdeogram ? (
                <IdeogramCaptionSidebar
                  caption={caption}
--- a/ui/src/components/formInputs.tsx
+++ b/ui/src/components/formInputs.tsx
@@ -363,13 +363,15 @@ export const CreatableSelectInput = (props: CreatableSelectInputProps) => {
        </label>
      )}
      <div className="flex gap-2">
-        <div className={isCustom ? 'w-1/3' : 'w-full'}>
+        <div className={isCustom ? 'w-20 shrink-0' : 'w-full'}>
          <Select
            value={selectedOption}
            options={selectOptions}
            isDisabled={props.disabled}
            className="aitk-react-select-container"
            classNamePrefix="aitk-react-select"
+            menuPosition="fixed"
+            menuPlacement="auto"
            formatOptionLabel={(option: unknown) => {
              const opt = option as SelectOption;
              return opt.value === CUSTOM_SELECT_VALUE ? (
@@ -397,7 +399,7 @@ export const CreatableSelectInput = (props: CreatableSelectInputProps) => {
            type="text"
            value={value}
            onChange={e => onChange(e.target.value)}
-            className={`${inputClasses} w-2/3`}
+            className={`${inputClasses} flex-1 min-w-0`}
            placeholder={props.placeholder ?? 'Enter custom value'}
            disabled={props.disabled}
            autoFocus