From 3378ad07b198d2332ce51241c3a1f2253fb269d4 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Tue, 7 Oct 2025 06:55:41 -0600
Subject: [PATCH] Update classification data card to use classification card

---
 web/public/locales/en/common.json             |  5 +-
 .../components/card/ClassificationCard.tsx    |  6 +-
 .../classification/ModelTrainingView.tsx      | 72 +++++++------------
 3 files changed, 35 insertions(+), 48 deletions(-)

diff --git a/web/public/locales/en/common.json b/web/public/locales/en/common.json
index c22a9227d..501f2d4bc 100644
--- a/web/public/locales/en/common.json
+++ b/web/public/locales/en/common.json
@@ -263,5 +263,8 @@
     "desc": "Page not found"
   },
   "selectItem": "Select {{item}}",
-  "readTheDocumentation": "Read the documentation"
+  "readTheDocumentation": "Read the documentation",
+  "information": {
+    "pixels": "{{area}}px"
+  }
 }
diff --git a/web/src/components/card/ClassificationCard.tsx b/web/src/components/card/ClassificationCard.tsx
index 7e694c595..235224db6 100644
--- a/web/src/components/card/ClassificationCard.tsx
+++ b/web/src/components/card/ClassificationCard.tsx
@@ -11,6 +11,7 @@ import { useTranslation } from "react-i18next";
 
 type ClassificationCardProps = {
   className?: string;
+  imgClassName?: string;
   data: ClassificationItemData;
   threshold?: ClassificationThreshold;
   selected: boolean;
@@ -20,6 +21,7 @@ type ClassificationCardProps = {
 };
 export function ClassificationCard({
   className,
+  imgClassName,
   data,
   threshold,
   selected,
@@ -75,7 +77,7 @@ export function ClassificationCard({
           <img
             ref={imgRef}
             onLoad={() => setImageLoaded(true)}
-            className={cn("size-44", isMobile && "w-full")}
+            className={cn("size-44", imgClassName, isMobile && "w-full")}
             src={`${baseUrl}${data.filepath}`}
             onClick={(e) => {
               e.stopPropagation();
@@ -84,7 +86,7 @@ export function ClassificationCard({
           />
           {imageArea != undefined && (
             <div className="absolute bottom-1 right-1 z-10 rounded-lg bg-black/50 px-2 py-1 text-xs text-white">
-              {t("pixels", { area: imageArea })}
+              {t("information.pixels", { ns: "common", area: imageArea })}
             </div>
           )}
         </div>
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index 7fe241496..68a2bfbbc 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -61,6 +61,7 @@ import { MdAutoFixHigh } from "react-icons/md";
 import TrainFilterDialog from "@/components/overlay/dialog/TrainFilterDialog";
 import useApiFilter from "@/hooks/use-api-filter";
 import { TrainFilter } from "@/types/classification";
+import { ClassificationCard } from "@/components/card/ClassificationCard";
 
 type ModelTrainingViewProps = {
   model: CustomClassificationModelConfig;
@@ -626,53 +627,34 @@ function DatasetGrid({
       className="scrollbar-container flex flex-wrap gap-2 overflow-y-auto p-2"
     >
       {classData.map((image) => (
-        <div
-          className={cn(
-            "flex w-60 cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]",
-            selectedImages.includes(image)
-              ? "shadow-selected outline-selected"
-              : "outline-transparent duration-500",
-          )}
-          onClick={(e) => {
-            e.stopPropagation();
-
-            if (e.ctrlKey || e.metaKey) {
-              onClickImages([image], true);
-            }
+        <ClassificationCard
+          key={image}
+          className="w-60 gap-4 rounded-lg bg-card p-2"
+          imgClassName="size-auto"
+          data={{
+            filename: image,
+            filepath: `clips/${modelName}/dataset/${categoryName}/${image}`,
+            name: "",
           }}
+          selected={selectedImages.includes(image)}
+          i18nLibrary="views/classificationModel"
+          onClick={(data, _) => onClickImages([data.filename], true)}
         >
-          <div
-            className={cn(
-              "w-full overflow-hidden p-2 *:text-card-foreground",
-              isMobile && "flex justify-center",
-            )}
-          >
-            <img
-              className="rounded-lg"
-              src={`${baseUrl}clips/${modelName}/dataset/${categoryName}/${image}`}
-            />
-          </div>
-          <div className="rounded-b-lg bg-card p-3">
-            <div className="flex w-full flex-row items-center justify-between gap-2">
-              <div className="flex w-full flex-row items-start justify-end gap-5 md:gap-4">
-                <Tooltip>
-                  <TooltipTrigger>
-                    <LuTrash2
-                      className="size-5 cursor-pointer text-primary-variant hover:text-primary"
-                      onClick={(e) => {
-                        e.stopPropagation();
-                        onDelete([image]);
-                      }}
-                    />
-                  </TooltipTrigger>
-                  <TooltipContent>
-                    {t("button.deleteClassificationAttempts")}
-                  </TooltipContent>
-                </Tooltip>
-              </div>
-            </div>
-          </div>
-        </div>
+          <Tooltip>
+            <TooltipTrigger>
+              <LuTrash2
+                className="size-5 cursor-pointer text-primary-variant hover:text-primary"
+                onClick={(e) => {
+                  e.stopPropagation();
+                  onDelete([image]);
+                }}
+              />
+            </TooltipTrigger>
+            <TooltipContent>
+              {t("button.deleteClassificationAttempts")}
+            </TooltipContent>
+          </Tooltip>
+        </ClassificationCard>
       ))}
     </div>
   );