diff --git a/frigate/api/classification.py b/frigate/api/classification.py
index 5e5320813..ec53a47af 100644
--- a/frigate/api/classification.py
+++ b/frigate/api/classification.py
@@ -435,7 +435,7 @@ def transcribe_audio(request: Request, body: AudioTranscriptionBody):
 def get_classification_dataset(name: str):
     dataset_dict: dict[str, list[str]] = {}
 
-    dataset_dir = os.path.join(MODEL_CACHE_DIR, f"{sanitize_filename(name)}/dataset")
+    dataset_dir = os.path.join(CLIPS_DIR, sanitize_filename(name), "dataset")
 
     if not os.path.exists(dataset_dir):
         return JSONResponse(status_code=200, content={})
@@ -459,7 +459,7 @@ def get_classification_dataset(name: str):
 
 @router.get("/classification/{name}/train")
 def get_classification_images(name: str):
-    train_dir = os.path.join(CLIPS_DIR, sanitize_filename(name))
+    train_dir = os.path.join(CLIPS_DIR, sanitize_filename(name), "train")
 
     if not os.path.exists(train_dir):
         return JSONResponse(status_code=200, content=[])
@@ -492,9 +492,7 @@ async def train_configured_model(
             status_code=404,
         )
 
-    background_tasks.add_task(
-        train_classification_model, os.path.join(MODEL_CACHE_DIR, name)
-    )
+    background_tasks.add_task(train_classification_model, name)
     return JSONResponse(
         content={"success": True, "message": "Started classification model training."},
         status_code=200,
diff --git a/web/src/views/classification/ModelSelectionView.tsx b/web/src/views/classification/ModelSelectionView.tsx
index 502a53251..fdc83dc82 100644
--- a/web/src/views/classification/ModelSelectionView.tsx
+++ b/web/src/views/classification/ModelSelectionView.tsx
@@ -38,12 +38,10 @@ export default function ModelSelectionView({
     <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]",
             isMobile && "w-full",
-            false
-              ? "shadow-selected outline-selected"
-              : "outline-transparent duration-500",
           )}
           onClick={() => onClick(config)}
           onContextMenu={() => {
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index c7bd0e18e..1a76b5c5e 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -40,8 +40,12 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
 
   // dataset
 
-  const { data: trainImages } = useSWR(`classification/${model.name}/train`);
-  const { data: dataset } = useSWR(`classification/${model.name}/dataset`);
+  const { data: trainImages } = useSWR<string[]>(
+    `classification/${model.name}/train`,
+  );
+  const { data: dataset } = useSWR<{ [id: string]: string[] }>(
+    `classification/${model.name}/dataset`,
+  );
 
   // actions
 
@@ -54,8 +58,8 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
       <div className="flex flex-row justify-between gap-2 align-middle">
         <LibrarySelector
           pageToggle={pageToggle}
-          dataset={dataset}
-          trainImages={trainImages}
+          dataset={dataset || {}}
+          trainImages={trainImages || []}
           setPageToggle={setPageToggle}
           onDelete={() => {}}
           onRename={() => {}}
@@ -65,13 +69,18 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
       {pageToggle == "train" ? (
         <TrainGrid
           model={model}
-          trainImages={trainImages}
+          trainImages={trainImages || []}
           selected={false}
           onClickImages={() => {}}
           onDelete={() => {}}
         />
       ) : (
-        <DatasetGrid />
+        <DatasetGrid
+          modelName={model.name}
+          categoryName={pageToggle}
+          images={dataset?.[pageToggle] || []}
+          onDelete={() => {}}
+        />
       )}
     </div>
   );
@@ -259,51 +268,65 @@ function LibrarySelector({
 }
 
 type DatasetGridProps = {
-  name: string;
+  modelName: string;
+  categoryName: string;
   images: string[];
-  onDelete: (name: string, ids: string[]) => void;
+  onDelete: (modelName: string, categoryName: string, ids: string[]) => void;
 };
-function DatasetGrid({ name, images, onDelete }: DatasetGridProps) {
+function DatasetGrid({
+  modelName,
+  categoryName,
+  images,
+  onDelete,
+}: DatasetGridProps) {
   const { t } = useTranslation(["views/classificationModel"]);
 
   return (
-    <div
-      className={cn(
-        "flex cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]",
-      )}
-    >
-      <div
-        className={cn(
-          "w-full overflow-hidden p-2 *:text-card-foreground",
-          isMobile && "flex justify-center",
-        )}
-      >
-        <img
-          className="h-40 rounded-lg"
-          src={`${baseUrl}clips/faces/${name}/${images}`}
-        />
-      </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">{name}</div>
+    <div className="grid grid-cols-10 gap-2 overflow-y-auto">
+      {images.map((image) => (
+        <div
+          className={cn(
+            "flex h-60 cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]",
+            "outline-transparent duration-500",
+          )}
+          onClick={() => {
+            //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="rounded-lg"
+              src={`${baseUrl}clips/${modelName}/dataset/${categoryName}/${image}`}
+            />
           </div>
-          <div className="flex 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(name, images);
-                  }}
-                />
-              </TooltipTrigger>
-              <TooltipContent>{t("button.deleteFaceAttempts")}</TooltipContent>
-            </Tooltip>
+          <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(modelName, categoryName, [image]);
+                      }}
+                    />
+                  </TooltipTrigger>
+                  <TooltipContent>
+                    {t("button.deleteClassificationAttempts")}
+                  </TooltipContent>
+                </Tooltip>
+              </div>
+            </div>
           </div>
         </div>
-      </div>
+      ))}
     </div>
   );
 }
@@ -339,9 +362,10 @@ function TrainGrid({
   );
 
   return (
-    <div className="grid size-full grid-cols-10 gap-2 overflow-y-auto">
-      {trainData.map((data) => (
+    <div className="grid grid-cols-10 gap-2 overflow-y-auto">
+      {trainData?.map((data) => (
         <div
+          key={data.timestamp}
           className={cn(
             "flex cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]",
             selected
@@ -361,7 +385,7 @@ function TrainGrid({
           >
             <img
               className="h-48 rounded-lg"
-              src={`${baseUrl}clips/${model.name}/${data.raw}`}
+              src={`${baseUrl}clips/${model.name}/train/${data.raw}`}
             />
           </div>
           <div className="rounded-b-lg bg-card p-3">