From da2d52645fe328100ba31bb669a8858cff699c42 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Mon, 23 Jun 2025 08:40:28 -0600
Subject: [PATCH] Classification model cover images (#18843)

* Move to separate component

* Add cover images for clssification models
---
 .../classification/ModelSelectionView.tsx     | 74 ++++++++++++++-----
 1 file changed, 54 insertions(+), 20 deletions(-)

diff --git a/web/src/views/classification/ModelSelectionView.tsx b/web/src/views/classification/ModelSelectionView.tsx
index 63133842a..aa2f94c6a 100644
--- a/web/src/views/classification/ModelSelectionView.tsx
+++ b/web/src/views/classification/ModelSelectionView.tsx
@@ -1,3 +1,4 @@
+import { baseUrl } from "@/api/baseUrl";
 import ActivityIndicator from "@/components/indicators/activity-indicator";
 import { cn } from "@/lib/utils";
 import {
@@ -37,27 +38,60 @@ export default function ModelSelectionView({
   return (
     <div className="flex size-full gap-2 p-2">
       {classificationConfigs.map((config) => (
-        <div
-          key={config.name}
-          className={cn(
-            "flex h-52 cursor-pointer flex-col gap-2 rounded-lg bg-card p-2 outline outline-[3px]",
-            "outline-transparent duration-500",
-            isMobile && "w-full",
-          )}
-          onClick={() => onClick(config)}
-          onContextMenu={() => {
-            // e.stopPropagation();
-            // e.preventDefault();
-            // handleClickEvent(true);
-          }}
-        >
-          <div className="size-48"></div>
-          <div className="smart-capitalize">
-            {config.name} ({config.state_config != null ? "State" : "Object"}{" "}
-            Classification)
-          </div>
-        </div>
+        <ModelCard config={config} onClick={() => onClick(config)} />
       ))}
     </div>
   );
 }
+
+type ModelCardProps = {
+  config: CustomClassificationModelConfig;
+  onClick: () => void;
+};
+function ModelCard({ config, onClick }: ModelCardProps) {
+  const { data: dataset } = useSWR<{
+    [id: string]: string[];
+  }>(`classification/${config.name}/dataset`, { revalidateOnFocus: false });
+
+  const coverImages = useMemo(() => {
+    if (!dataset) {
+      return {};
+    }
+
+    const imageMap: { [key: string]: string } = {};
+
+    for (const [key, imageList] of Object.entries(dataset)) {
+      if (imageList.length > 0) {
+        imageMap[key] = imageList[0];
+      }
+    }
+
+    return imageMap;
+  }, [dataset]);
+
+  return (
+    <div
+      key={config.name}
+      className={cn(
+        "flex h-60 cursor-pointer flex-col items-center gap-2 rounded-lg bg-card p-2 outline outline-[3px]",
+        "outline-transparent duration-500",
+        isMobile && "w-full",
+      )}
+      onClick={() => onClick()}
+    >
+      <div className="grid size-48 grid-cols-2 gap-2">
+        {Object.entries(coverImages).map(([key, image]) => (
+          <img
+            key={key}
+            className=""
+            src={`${baseUrl}clips/${config.name}/dataset/${key}/${image}`}
+          />
+        ))}
+      </div>
+      <div className="smart-capitalize">
+        {config.name} ({config.state_config != null ? "State" : "Object"}{" "}
+        Classification)
+      </div>
+    </div>
+  );
+}