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>(new Set()); + const [cacheKey, setCacheKey] = useState(Date.now()); + const [loadedImages, setLoadedImages] = useState>(new Set()); + + const handleImageLoad = useCallback((imageName: string) => { + setLoadedImages((prev) => new Set(prev).add(imageName)); + }, []); const { data: trainImages, mutate: refreshTrainImages } = useSWR( 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) && ( +
+ +
+ )} {`Example handleImageLoad(imageName)} /> );