From 220f3e1e7e4638833b32ebfccdd8919031c62f6d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 19 Feb 2026 08:11:51 +0000
Subject: [PATCH] Updated understanding: events support multiple classification
 models

Co-authored-by: Teagan42 <2989925+Teagan42@users.noreply.github.com>
---
 .../overlay/ClassificationSelectionDialog.tsx | 11 ++-
 .../classification/ModelTrainingView.tsx      | 83 +++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/web/src/components/overlay/ClassificationSelectionDialog.tsx b/web/src/components/overlay/ClassificationSelectionDialog.tsx
index 6398348a4..8e2037f18 100644
--- a/web/src/components/overlay/ClassificationSelectionDialog.tsx
+++ b/web/src/components/overlay/ClassificationSelectionDialog.tsx
@@ -35,6 +35,7 @@ type ClassificationSelectionDialogProps = {
   modelName: string;
   image: string;
   onRefresh: () => void;
+  onCategorize?: (category: string) => void; // Optional custom categorize handler
   children: ReactNode;
 };
 export default function ClassificationSelectionDialog({
@@ -43,12 +44,20 @@ export default function ClassificationSelectionDialog({
   modelName,
   image,
   onRefresh,
+  onCategorize,
   children,
 }: ClassificationSelectionDialogProps) {
   const { t } = useTranslation(["views/classificationModel"]);
 
   const onCategorizeImage = useCallback(
     (category: string) => {
+      // If custom categorize handler is provided, use it instead
+      if (onCategorize) {
+        onCategorize(category);
+        return;
+      }
+
+      // Default behavior: categorize single image
       axios
         .post(`/classification/${modelName}/dataset/categorize`, {
           category,
@@ -72,7 +81,7 @@ export default function ClassificationSelectionDialog({
           });
         });
     },
-    [modelName, image, onRefresh, t],
+    [modelName, image, onRefresh, onCategorize, t],
   );
 
   const isChildButton = useMemo(
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index 00e811d09..87e5407dd 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -458,6 +458,89 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
                 </>
               )}
             </div>
+            {pageToggle === "train" && (
+              <ClassificationSelectionDialog
+                classes={classes}
+                modelName={model.name}
+                image={selectedImages[0]}
+                onRefresh={(category: string) => {
+                  // Batch categorize all selected images
+                  let successCount = 0;
+                  let failCount = 0;
+                  const totalCount = selectedImages.length;
+
+                  selectedImages.forEach((filename, index) => {
+                    axios
+                      .post(`/classification/${model.name}/dataset/categorize`, {
+                        category,
+                        training_file: filename,
+                      })
+                      .then((resp) => {
+                        if (resp.status == 200) {
+                          successCount++;
+                        } else {
+                          failCount++;
+                        }
+
+                        // Show final toast after all requests complete
+                        if (index === totalCount - 1) {
+                          if (successCount === totalCount) {
+                            toast.success(
+                              t("toast.success.batchCategorized", {
+                                count: successCount,
+                              }),
+                              {
+                                position: "top-center",
+                              },
+                            );
+                          } else if (successCount > 0) {
+                            toast.warning(
+                              t("toast.warning.partialBatchCategorized", {
+                                success: successCount,
+                                total: totalCount,
+                              }),
+                              {
+                                position: "top-center",
+                              },
+                            );
+                          } else {
+                            toast.error(
+                              t("toast.error.batchCategorizeFailed", {
+                                count: totalCount,
+                              }),
+                              {
+                                position: "top-center",
+                              },
+                            );
+                          }
+                          setSelectedImages([]);
+                          refreshAll();
+                        }
+                      })
+                      .catch(() => {
+                        failCount++;
+                        if (index === totalCount - 1) {
+                          toast.error(
+                            t("toast.error.batchCategorizeFailed", {
+                              count: totalCount,
+                            }),
+                            {
+                              position: "top-center",
+                            },
+                          );
+                          setSelectedImages([]);
+                          refreshAll();
+                        }
+                      });
+                  });
+                }}
+              >
+                <Button className="flex gap-2">
+                  <TbCategoryPlus className="size-7 rounded-md p-1 text-secondary-foreground" />
+                  {isDesktop && t("button.categorizeImages")}
+                </Button>
+              </ClassificationSelectionDialog>
+            )}
             <Button
               className="flex gap-2"
               onClick={() => setDeleteDialogOpen(selectedImages)}