From af5f3e5f84909e7d3161a5a7d902e7b814d3124d Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Wed, 16 Jul 2025 15:48:09 -0600
Subject: [PATCH] Improve model train button

---
 web/public/locales/en/views/classificationModel.json | 3 ++-
 web/src/views/classification/ModelTrainingView.tsx   | 9 +++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/web/public/locales/en/views/classificationModel.json b/web/public/locales/en/views/classificationModel.json
index 0af0179b9..47b2b13bf 100644
--- a/web/public/locales/en/views/classificationModel.json
+++ b/web/public/locales/en/views/classificationModel.json
@@ -3,7 +3,8 @@
     "deleteClassificationAttempts": "Delete Classification Images",
     "renameCategory": "Rename Class",
     "deleteCategory": "Delete Class",
-    "deleteImages": "Delete Images"
+    "deleteImages": "Delete Images",
+    "trainModel": "Train Model"
   },
   "toast": {
     "success": {
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index c8c2266d5..c489bc699 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -50,6 +50,7 @@ import { ModelState } from "@/types/ws";
 import ActivityIndicator from "@/components/indicators/activity-indicator";
 import { useNavigate } from "react-router-dom";
 import { IoMdArrowRoundBack } from "react-icons/io";
+import { MdAutoFixHigh } from "react-icons/md";
 
 type ModelTrainingViewProps = {
   model: CustomClassificationModelConfig;
@@ -345,8 +346,12 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
             onClick={trainModel}
             disabled={modelState != "complete"}
           >
-            Train Model
-            {modelState == "training" && <ActivityIndicator size={20} />}
+            {modelState == "training" ? (
+              <ActivityIndicator size={20} />
+            ) : (
+              <MdAutoFixHigh />
+            )}
+            {t("button.trainModel")}
           </Button>
         )}
       </div>