From 73eceaf24260f90e39c980237fa3c01a7b6904cd Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Fri, 19 Dec 2025 08:39:35 -0600
Subject: [PATCH] add cache key and activity indicator for loading
 classification wizard images

---
 .../wizard/Step3ChooseExamples.tsx               | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/web/src/components/classification/wizard/Step3ChooseExamples.tsx b/web/src/components/classification/wizard/Step3ChooseExamples.tsx
index ec45f4b4d..e3dd04afc 100644
--- a/web/src/components/classification/wizard/Step3ChooseExamples.tsx
+++ b/web/src/components/classification/wizard/Step3ChooseExamples.tsx
@@ -45,6 +45,12 @@ export default function Step3ChooseExamples({
   const [isProcessing, setIsProcessing] = useState(false);
   const [currentClassIndex, setCurrentClassIndex] = useState(0);
   const [selectedImages, setSelectedImages] = useState<Set<string>>(new Set());
+  const [cacheKey, setCacheKey] = useState<number>(Date.now());
+  const [loadedImages, setLoadedImages] = useState<Set<string>>(new Set());
+
+  const handleImageLoad = useCallback((imageName: string) => {
+    setLoadedImages((prev) => new Set(prev).add(imageName));
+  }, []);
 
   const { data: trainImages, mutate: refreshTrainImages } = useSWR<string[]>(
     hasGenerated ? `classification/${step1Data.modelName}/train` : null,
@@ -332,6 +338,8 @@ export default function Step3ChooseExamples({
       setHasGenerated(true);
       toast.success(t("wizard.step3.generateSuccess"));
 
+      // Update cache key to force image reload
+      setCacheKey(Date.now());
       await refreshTrainImages();
     } catch (error) {
       const axiosError = error as {
@@ -565,10 +573,16 @@ export default function Step3ChooseExamples({
                       )}
                       onClick={() => toggleImageSelection(imageName)}
                     >
+                      {!loadedImages.has(imageName) && (
+                        <div className="flex h-full items-center justify-center">
+                          <ActivityIndicator className="size-6" />
+                        </div>
+                      )}
                       <img
-                        src={`${baseUrl}clips/${step1Data.modelName}/train/${imageName}`}
+                        src={`${baseUrl}clips/${step1Data.modelName}/train/${imageName}?t=${cacheKey}`}
                         alt={`Example ${index + 1}`}
                         className="h-full w-full object-cover"
+                        onLoad={() => handleImageLoad(imageName)}
                       />
                     </div>
                   );