diff --git a/web/src/views/classification/ModelSelectionView.tsx b/web/src/views/classification/ModelSelectionView.tsx index ef415efb4..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 { @@ -48,22 +49,45 @@ type ModelCardProps = { 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 (