Refactor state training grid to use classification card

web/src/components/card/ClassificationCard.tsx

@@ -16,6 +16,7 @@ type ClassificationCardProps = {
   threshold?: ClassificationThreshold;
   selected: boolean;
   i18nLibrary: string;
+  showArea?: boolean;
   onClick: (data: ClassificationItemData, meta: boolean) => void;
   children?: React.ReactNode;
 };
@@ -26,6 +27,7 @@ export function ClassificationCard({
   threshold,
   selected,
   i18nLibrary,
+  showArea = true,
   onClick,
   children,
 }: ClassificationCardProps) {
@@ -55,12 +57,12 @@ export function ClassificationCard({
   });
 
   const imageArea = useMemo(() => {
-    if (imgRef.current == null || !imageLoaded) {
+    if (!showArea || imgRef.current == null || !imageLoaded) {
       return undefined;
     }
 
     return imgRef.current.naturalWidth * imgRef.current.naturalHeight;
-  }, [imageLoaded]);
+  }, [showArea, imageLoaded]);
 
   return (
     <>

web/src/views/classification/ModelTrainingView.tsx

@@ -60,7 +60,7 @@ import { IoMdArrowRoundBack } from "react-icons/io";
 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 { ClassificationItemData, TrainFilter } from "@/types/classification";
 import { ClassificationCard } from "@/components/card/ClassificationCard";
 
 type ModelTrainingViewProps = {
@@ -682,20 +682,18 @@ function TrainGrid({
   onRefresh,
   onDelete,
 }: TrainGridProps) {
-  const { t } = useTranslation(["views/classificationModel"]);
+  const trainData = useMemo<ClassificationItemData[]>(
-
-  const trainData = useMemo(
     () =>
       trainImages
         .map((raw) => {
           const parts = raw.replaceAll(".webp", "").split("-");
           const rawScore = Number.parseFloat(parts[2]);
           return {
-            raw,
+            filename: raw,
-            timestamp: parts[0],
+            filepath: `clips/${model.name}/train/${raw}`,
-            label: parts[1],
+            timestamp: Number.parseFloat(parts[0]),
-            score: rawScore * 100,
+            name: parts[1],
-            truePositive: rawScore >= model.threshold,
+            score: rawScore,
           };
         })
         .filter((data) => {
@@ -703,10 +701,7 @@ function TrainGrid({
             return true;
           }
 
-          if (
+          if (trainFilter.classes && !trainFilter.classes.includes(data.name)) {
-            trainFilter.classes &&
-            !trainFilter.classes.includes(data.label)
-          ) {
             return false;
           }
 
@@ -726,10 +721,57 @@ function TrainGrid({
 
           return true;
         })
-        .sort((a, b) => b.timestamp.localeCompare(a.timestamp)),
+        .sort((a, b) => b.timestamp - a.timestamp),
     [model, trainImages, trainFilter],
   );
 
+  if (model.state_config) {
+    return (
+      <StateTrainGrid
+        model={model}
+        contentRef={contentRef}
+        classes={classes}
+        trainData={trainData}
+        selectedImages={selectedImages}
+        onClickImages={onClickImages}
+        onRefresh={onRefresh}
+        onDelete={onDelete}
+      />
+    );
+  }
+
+  return <div />;
+}
+
+type StateTrainGridProps = {
+  model: CustomClassificationModelConfig;
+  contentRef: MutableRefObject<HTMLDivElement | null>;
+  classes: string[];
+  trainData?: ClassificationItemData[];
+  selectedImages: string[];
+  onClickImages: (images: string[], ctrl: boolean) => void;
+  onRefresh: () => void;
+  onDelete: (ids: string[]) => void;
+};
+function StateTrainGrid({
+  model,
+  contentRef,
+  classes,
+  trainData,
+  selectedImages,
+  onClickImages,
+  onRefresh,
+  onDelete,
+}: StateTrainGridProps) {
+  const { t } = useTranslation(["views/classificationModel"]);
+
+  const threshold = useMemo(() => {
+    return {
+      recognition: model.threshold,
+      unknown: model.threshold,
+    };
+  }, [model]);
+
   return (
     <div
       ref={contentRef}
@@ -739,51 +781,20 @@ function TrainGrid({
       )}
     >
       {trainData?.map((data) => (
-        <div
+        <ClassificationCard
-          key={data.timestamp}
+          className="w-60 gap-2 rounded-lg bg-card p-2"
-          className={cn(
+          imgClassName="size-auto"
-            "flex w-56 cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]",
+          data={data}
-            selectedImages.includes(data.raw)
+          threshold={threshold}
-              ? "shadow-selected outline-selected"
+          selected={selectedImages.includes(data.filename)}
-              : "outline-transparent duration-500",
+          i18nLibrary="views/classificationModel"
-            isMobile && "w-[48%]",
+          showArea={false}
-          )}
+          onClick={(data, meta) => onClickImages([data.filename], meta)}
-          onClick={(e) => {
-            e.stopPropagation();
-            onClickImages([data.raw], e.ctrlKey || e.metaKey);
-          }}
         >
-          <div
-            className={cn(
-              "w-full overflow-hidden p-2 *:text-card-foreground",
-              isMobile && "flex justify-center",
-            )}
-          >
-            <img
-              className="w-56 rounded-lg"
-              src={`${baseUrl}clips/${model.name}/train/${data.raw}`}
-            />
-          </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 flex-col items-start text-xs text-primary-variant">
-                <div className="smart-capitalize">
-                  {data.label.replaceAll("_", " ")}
-                </div>
-                <div
-                  className={cn(
-                    "",
-                    data.truePositive ? "text-success" : "text-danger",
-                  )}
-                >
-                  {data.score}%
-                </div>
-              </div>
-              <div className="flex flex-row items-start justify-end gap-5 md:gap-4">
           <ClassificationSelectionDialog
             classes={classes}
             modelName={model.name}
-                  image={data.raw}
+            image={data.filename}
             onRefresh={onRefresh}
           >
             <TbCategoryPlus className="size-5 cursor-pointer text-primary-variant hover:text-primary" />
@@ -794,7 +805,7 @@ function TrainGrid({
                 className="size-5 cursor-pointer text-primary-variant hover:text-primary"
                 onClick={(e) => {
                   e.stopPropagation();
-                        onDelete([data.raw]);
+                  onDelete([data.filename]);
                 }}
               />
             </TooltipTrigger>
@@ -802,10 +813,7 @@ function TrainGrid({
               {t("button.deleteClassificationAttempts")}
             </TooltipContent>
           </Tooltip>
-              </div>
+        </ClassificationCard>
-            </div>
-          </div>
-        </div>
       ))}
     </div>
   );