Drastically improved the performance of the dataset viewer on large datasets by switching to virtualization. Fixed issue with images loading when they have two periods in a row ..

ui/src/components/DatasetImageCard.tsx

@@ -1,10 +1,11 @@
-import React, { useRef, useEffect, useState, ReactNode, KeyboardEvent } from 'react';
-import { FaTrashAlt, FaEye, FaEyeSlash } from 'react-icons/fa';
+import React, { useEffect, useState, ReactNode, KeyboardEvent, useRef } from 'react';
+import { FaTrashAlt } from 'react-icons/fa';
 import { openConfirm } from './ConfirmModal';
 import classNames from 'classnames';
 import { apiClient } from '@/utils/api';
 import AudioPlayer from './AudioPlayer';
 import { isVideo, isAudio } from '@/utils/basic';
+import useCaptionBatch, { setCachedCaption } from '@/hooks/useCaptionBatch';
 
 interface DatasetImageCardProps {
   imageUrl: string;
@@ -27,131 +28,85 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
   onImageClick,
   captionRefreshKey = 0,
 }) => {
-  const cardRef = useRef<HTMLDivElement>(null);
-  const [isVisible, setIsVisible] = useState<boolean>(false);
-  const [inViewport, setInViewport] = useState<boolean>(false);
   const [loaded, setLoaded] = useState<boolean>(false);
-  const [isCaptionLoaded, setIsCaptionLoaded] = useState<boolean>(false);
+  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
+  const [pollTick, setPollTick] = useState(0);
+
+  const combinedRefreshKey = captionRefreshKey + pollTick;
+  const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(imageUrl, combinedRefreshKey);
+
   const [caption, setCaption] = useState<string>('');
   const [savedCaption, setSavedCaption] = useState<string>('');
-  const abortControllerRef = useRef<AbortController | null>(null);
+  const dirtyRef = useRef<boolean>(false);
 
-  const fetchCaption = async () => {
-    if (isCaptionLoaded) return;
-    abortControllerRef.current?.abort();
-    const controller = new AbortController();
-    abortControllerRef.current = controller;
-    apiClient
-      .post(`/api/caption/get`, { imgPath: imageUrl }, { signal: controller.signal })
-      .then(res => res.data)
-      .then(data => {
-        console.log('Caption fetched:', data);
-        if (data) {
-          data = `${data}`;
-        }
-        setCaption(data || '');
-        setSavedCaption((data || '').trim());
-        setIsCaptionLoaded(true);
-      })
-      .catch(error => {
-        if (controller.signal.aborted) return;
-        console.error('Error fetching caption:', error);
-      })
-      .finally(() => {
-        if (abortControllerRef.current === controller) {
-          abortControllerRef.current = null;
-        }
-      });
-  };
+  // Sync from the fetched caption, but don't clobber unsaved local edits.
+  useEffect(() => {
+    if (!isCaptionLoaded) return;
+    if (dirtyRef.current) return;
+    setCaption(fetchedCaption);
+    setSavedCaption(fetchedCaption.trim());
+  }, [fetchedCaption, isCaptionLoaded]);
+
+  // Poll while auto-captioning so backend-written captions show up.
+  useEffect(() => {
+    if (!isAutoCaptioning) return;
+    const interval = setInterval(() => setPollTick(t => t + 1), 5000);
+    return () => clearInterval(interval);
+  }, [isAutoCaptioning]);
 
   const saveCaption = () => {
     const trimmedCaption = caption.trim();
-    if (trimmedCaption === savedCaption) return;
+    if (trimmedCaption === savedCaption) {
+      dirtyRef.current = false;
+      return;
+    }
     apiClient
       .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption })
-      .then(res => res.data)
-      .then(data => {
-        console.log('Caption saved:', data);
+      .then(() => {
         setSavedCaption(trimmedCaption);
+        setCachedCaption(imageUrl, trimmedCaption);
+        dirtyRef.current = false;
       })
       .catch(error => {
         console.error('Error saving caption:', error);
       });
   };
 
-  // Only fetch caption when the component is both in viewport and visible
+  // Save any pending edit if the card unmounts (e.g. scrolled out of the virtualized window).
+  const latestRef = useRef({ caption, savedCaption, imageUrl });
   useEffect(() => {
-    if (inViewport && isVisible) {
-      fetchCaption();
-    }
-  }, [inViewport, isVisible, isCaptionLoaded]);
-
-  // Poll for caption updates every 5 seconds while auto-captioning
+    latestRef.current = { caption, savedCaption, imageUrl };
+  });
   useEffect(() => {
-    if (!isAutoCaptioning || !inViewport || !isVisible) return;
-    const interval = setInterval(() => {
-      // Reset so fetchCaption will re-fetch
-      setIsCaptionLoaded(false);
-    }, 5000);
-    return () => clearInterval(interval);
-  }, [isAutoCaptioning, inViewport, isVisible]);
-
-  // External trigger (e.g. caption edited in the full-screen viewer) — re-fetch
-  useEffect(() => {
-    if (captionRefreshKey === 0) return;
-    setIsCaptionLoaded(false);
-  }, [captionRefreshKey]);
-
-  useEffect(() => {
-    // Create intersection observer to check viewport visibility
-    const observer = new IntersectionObserver(
-      entries => {
-        if (entries[0].isIntersecting) {
-          setInViewport(true);
-          // Initialize isVisible to true when first coming into view
-          if (!isVisible) {
-            setIsVisible(true);
-          }
-        } else {
-          setInViewport(false);
-          // Cancel any in-flight caption fetch when scrolling away
-          abortControllerRef.current?.abort();
-        }
-      },
-      { threshold: 0.1 },
-    );
-
-    if (cardRef.current) {
-      observer.observe(cardRef.current);
-    }
-
     return () => {
-      observer.disconnect();
+      if (!dirtyRef.current) return;
+      const { caption: c, savedCaption: s, imageUrl: url } = latestRef.current;
+      const trimmed = c.trim();
+      if (trimmed === s) return;
+      apiClient
+        .post('/api/img/caption', { imgPath: url, caption: trimmed })
+        .then(() => setCachedCaption(url, trimmed))
+        .catch(err => console.error('Error saving caption on unmount:', err));
     };
   }, []);
 
-  const toggleVisibility = (): void => {
-    setIsVisible(prev => !prev);
-    if (!isVisible && !isCaptionLoaded) {
-      fetchCaption();
-    }
-  };
-
   const handleLoad = (): void => {
     setLoaded(true);
   };
 
   const handleKeyDown = (e: KeyboardEvent<HTMLTextAreaElement>): void => {
-    // If Enter is pressed without Shift, prevent default behavior and save
     if (e.key === 'Enter' && !e.shiftKey) {
       e.preventDefault();
       saveCaption();
     }
   };
 
-  const isCaptionCurrent = caption.trim() === savedCaption;
+  const handleCaptionChange = (value: string) => {
+    dirtyRef.current = value.trim() !== savedCaption;
+    setCaption(value);
+  };
 
-  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
+  const isCaptionCurrent = caption.trim() === savedCaption;
 
   const isItAVideo = isVideo(imageUrl);
   const isItAudio = isAudio(imageUrl);
@@ -159,66 +114,49 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
 
   return (
     <div className={`flex flex-col ${className}`}>
-      {/* Square image container */}
-      <div
-        ref={cardRef}
-        className="relative w-full"
-        style={{ paddingBottom: '100%' }} // Make it square
-      >
+      <div className="relative w-full" style={{ paddingBottom: '100%' }}>
         <div className="absolute inset-0 rounded-t-lg shadow-md">
-          {inViewport && isVisible && (
-            <>
-              {isItAVideo && (
-                <video
-                  src={`/api/img/${encodeURIComponent(imageUrl)}`}
-                  className={`w-full h-full object-contain`}
-                  autoPlay={false}
-                  loop
-                  muted
-                  controls
-                />
-              )}
-              {isItAudio && !showAudioPlayer && (
-                <div
-                  className="w-full h-full cursor-pointer flex items-center justify-center bg-gray-900"
-                  onClick={() => setShowAudioPlayer(true)}
-                >
-                  <img
-                    src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
-                    alt={alt}
-                    className="w-full h-full object-contain"
-                    onError={e => {
-                      (e.target as HTMLImageElement).style.display = 'none';
-                    }}
-                  />
-                </div>
-              )}
-              {isItAudio && showAudioPlayer && (
-                <AudioPlayer
-                  src={`/api/img/${encodeURIComponent(imageUrl)}`}
-                  title={imageUrl.replace(/^.*[\\/]/, '')}
-                />
-              )}
-              {isItImage && (
-                <img
-                  src={`/api/img/${encodeURIComponent(imageUrl)}`}
-                  alt={alt}
-                  onLoad={handleLoad}
-                  onClick={onImageClick}
-                  className={classNames('w-full h-full object-contain transition-opacity duration-300', {
-                    'opacity-100': loaded,
-                    'opacity-0': !loaded,
-                    'cursor-zoom-in': !!onImageClick,
-                  })}
-                />
-              )}
-            </>
+          {isItAVideo && (
+            <video
+              src={`/api/img/${encodeURIComponent(imageUrl)}`}
+              className={`w-full h-full object-contain`}
+              autoPlay={false}
+              loop
+              muted
+              controls
+            />
           )}
-          {!isVisible && (
-            <div className="absolute inset-0 flex items-center justify-center bg-gray-800 bg-opacity-75 rounded-t-lg">
-              <span className="text-white text-lg"></span>
+          {isItAudio && !showAudioPlayer && (
+            <div
+              className="w-full h-full cursor-pointer flex items-center justify-center bg-gray-900"
+              onClick={() => setShowAudioPlayer(true)}
+            >
+              <img
+                src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
+                alt={alt}
+                className="w-full h-full object-contain"
+                onError={e => {
+                  (e.target as HTMLImageElement).style.display = 'none';
+                }}
+              />
             </div>
           )}
+          {isItAudio && showAudioPlayer && (
+            <AudioPlayer src={`/api/img/${encodeURIComponent(imageUrl)}`} title={imageUrl.replace(/^.*[\\/]/, '')} />
+          )}
+          {isItImage && (
+            <img
+              src={`/api/img/${encodeURIComponent(imageUrl)}`}
+              alt={alt}
+              onLoad={handleLoad}
+              onClick={onImageClick}
+              className={classNames('w-full h-full object-contain transition-opacity duration-300', {
+                'opacity-100': loaded,
+                'opacity-0': !loaded,
+                'cursor-zoom-in': !!onImageClick,
+              })}
+            />
+          )}
           {children && <div className="absolute inset-0 flex items-center justify-center">{children}</div>}
           <div className="absolute top-1 right-1 flex space-x-2 z-10">
             <button
@@ -254,7 +192,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
           'border-transparent border-2': isCaptionCurrent,
         })}
       >
-        {inViewport && isVisible && (isCaptionLoaded || caption) && (
+        {isCaptionLoaded ? (
           <form
             onSubmit={e => {
               e.preventDefault();
@@ -269,17 +207,11 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
               value={caption}
               rows={3}
               readOnly={isAutoCaptioning}
-              onChange={e => setCaption(e.target.value)}
+              onChange={e => handleCaptionChange(e.target.value)}
               onKeyDown={handleKeyDown}
             />
           </form>
-        )}
-        {(!inViewport || !isVisible) && isCaptionLoaded && (
-          <div className="w-full h-full flex items-center justify-center text-gray-400">
-            {isVisible ? 'Scroll into view to edit caption' : 'Show content to edit caption'}
-          </div>
-        )}
-        {!isCaptionLoaded && !caption && (
+        ) : (
           <div className="w-full h-full flex items-center justify-center text-gray-400">Loading caption...</div>
         )}
       </div>

ui/src/components/layout.tsx

@@ -1,4 +1,5 @@
 'use client';
+import React from 'react';
 import classNames from 'classnames';
 import ThemeLogo from './ThemeLogo';
 import { mobileSidebarState } from './Sidebar';
@@ -35,10 +36,14 @@ export const TopBar: React.FC<Props> = ({ children, className }) => {
   );
 };
 
-export const MainContent: React.FC<Props> = ({ children, className }) => {
+export const MainContent = React.forwardRef<HTMLDivElement, Props>(({ children, className }, ref) => {
   return (
-    <div className={classNames('pt-14 px-2 sm:px-4 absolute top-0 left-0 w-full h-full overflow-auto', className)}>
+    <div
+      ref={ref}
+      className={classNames('pt-14 px-2 sm:px-4 absolute top-0 left-0 w-full h-full overflow-auto', className)}
+    >
       {children ? children : null}
     </div>
   );
-};
+});
+MainContent.displayName = 'MainContent';