From e33f2a91596539ff361eac167aade565bd8d3e29 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Mon, 3 Nov 2025 05:58:02 -0700
Subject: [PATCH] Fix classification back step

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

diff --git a/web/src/components/classification/wizard/Step3ChooseExamples.tsx b/web/src/components/classification/wizard/Step3ChooseExamples.tsx
index 06bb2bbad..68da03eaf 100644
--- a/web/src/components/classification/wizard/Step3ChooseExamples.tsx
+++ b/web/src/components/classification/wizard/Step3ChooseExamples.tsx
@@ -317,6 +317,21 @@ export default function Step3ChooseExamples({
     return unclassifiedImages.length === 0;
   }, [unclassifiedImages]);
 
+  const handleBack = useCallback(() => {
+    if (currentClassIndex > 0) {
+      const previousClass = allClasses[currentClassIndex - 1];
+      setCurrentClassIndex((prev) => prev - 1);
+
+      // Restore selections for the previous class
+      const previousSelections = Object.entries(imageClassifications)
+        .filter(([_, className]) => className === previousClass)
+        .map(([imageName, _]) => imageName);
+      setSelectedImages(new Set(previousSelections));
+    } else {
+      onBack();
+    }
+  }, [currentClassIndex, allClasses, imageClassifications, onBack]);
+
   return (
     <div className="flex flex-col gap-6">
       {isTraining ? (
@@ -420,7 +435,7 @@ export default function Step3ChooseExamples({
 
       {!isTraining && (
         <div className="flex flex-col gap-3 pt-3 sm:flex-row sm:justify-end sm:gap-4">
-          <Button type="button" onClick={onBack} className="sm:flex-1">
+          <Button type="button" onClick={handleBack} className="sm:flex-1">
             {t("button.back", { ns: "common" })}
           </Button>
           <Button