Classification Model UI (#18571)

* Setup basic training structure

* Build out route

* Handle model configs

* Add image fetch APIs

* Implement model training screen with dataset selection

* Implement viewing of training images

* Adjust directories

* Implement viewing of images

* Add support for deleting images

* Implement full deletion

* Implement classification model training

* Improve naming

* More renaming

* Improve layout

* Reduce logging

* Cleanup

frigate/api/classification.py

@@ -21,7 +21,7 @@ from frigate.api.defs.request.classification_body import (
 from frigate.api.defs.tags import Tags
 from frigate.config import FrigateConfig
 from frigate.config.camera import DetectConfig
-from frigate.const import FACE_DIR, MODEL_CACHE_DIR
+from frigate.const import CLIPS_DIR, FACE_DIR
 from frigate.embeddings import EmbeddingsContext
 from frigate.models import Event
 from frigate.util.classification import train_classification_model
@@ -449,6 +449,50 @@ def transcribe_audio(request: Request, body: AudioTranscriptionBody):
 # custom classification training
 
 
+@router.get("/classification/{name}/dataset")
+def get_classification_dataset(name: str):
+    dataset_dict: dict[str, list[str]] = {}
+
+    dataset_dir = os.path.join(CLIPS_DIR, sanitize_filename(name), "dataset")
+
+    if not os.path.exists(dataset_dir):
+        return JSONResponse(status_code=200, content={})
+
+    for name in os.listdir(dataset_dir):
+        category_dir = os.path.join(dataset_dir, name)
+
+        if not os.path.isdir(category_dir):
+            continue
+
+        dataset_dict[name] = []
+
+        for file in filter(
+            lambda f: (f.lower().endswith((".webp", ".png", ".jpg", ".jpeg"))),
+            os.listdir(category_dir),
+        ):
+            dataset_dict[name].append(file)
+
+    return JSONResponse(status_code=200, content=dataset_dict)
+
+
+@router.get("/classification/{name}/train")
+def get_classification_images(name: str):
+    train_dir = os.path.join(CLIPS_DIR, sanitize_filename(name), "train")
+
+    if not os.path.exists(train_dir):
+        return JSONResponse(status_code=200, content=[])
+
+    return JSONResponse(
+        status_code=200,
+        content=list(
+            filter(
+                lambda f: (f.lower().endswith((".webp", ".png", ".jpg", ".jpeg"))),
+                os.listdir(train_dir),
+            )
+        ),
+    )
+
+
 @router.post("/classification/{name}/train")
 async def train_configured_model(
     request: Request, name: str, background_tasks: BackgroundTasks
@@ -466,10 +510,131 @@ 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,
     )
+
+
+@router.post(
+    "/classification/{name}/dataset/{category}/delete",
+    dependencies=[Depends(require_role(["admin"]))],
+)
+def delete_classification_dataset_images(
+    request: Request, name: str, category: str, body: dict = None
+):
+    config: FrigateConfig = request.app.frigate_config
+
+    if name not in config.classification.custom:
+        return JSONResponse(
+            content=(
+                {
+                    "success": False,
+                    "message": f"{name} is not a known classification model.",
+                }
+            ),
+            status_code=404,
+        )
+
+    json: dict[str, Any] = body or {}
+    list_of_ids = json.get("ids", "")
+    folder = os.path.join(
+        CLIPS_DIR, sanitize_filename(name), "dataset", sanitize_filename(category)
+    )
+
+    for id in list_of_ids:
+        file_path = os.path.join(folder, id)
+
+        if os.path.isfile(file_path):
+            os.unlink(file_path)
+
+    return JSONResponse(
+        content=({"success": True, "message": "Successfully deleted faces."}),
+        status_code=200,
+    )
+
+
+@router.post(
+    "/classification/{name}/dataset/categorize",
+    dependencies=[Depends(require_role(["admin"]))],
+)
+def categorize_classification_image(request: Request, name: str, body: dict = None):
+    config: FrigateConfig = request.app.frigate_config
+
+    if name not in config.classification.custom:
+        return JSONResponse(
+            content=(
+                {
+                    "success": False,
+                    "message": f"{name} is not a known classification model.",
+                }
+            ),
+            status_code=404,
+        )
+
+    json: dict[str, Any] = body or {}
+    category = sanitize_filename(json.get("category", ""))
+    training_file_name = sanitize_filename(json.get("training_file", ""))
+    training_file = os.path.join(CLIPS_DIR, name, "train", training_file_name)
+
+    if training_file_name and not os.path.isfile(training_file):
+        return JSONResponse(
+            content=(
+                {
+                    "success": False,
+                    "message": f"Invalid filename or no file exists: {training_file_name}",
+                }
+            ),
+            status_code=404,
+        )
+
+    new_name = f"{category}-{datetime.datetime.now().timestamp()}.png"
+    new_file_folder = os.path.join(CLIPS_DIR, name, "dataset", category)
+
+    if not os.path.exists(new_file_folder):
+        os.mkdir(new_file_folder)
+
+    # use opencv because webp images can not be used to train
+    img = cv2.imread(training_file)
+    cv2.imwrite(os.path.join(new_file_folder, new_name), img)
+    os.unlink(training_file)
+
+    return JSONResponse(
+        content=({"success": True, "message": "Successfully deleted faces."}),
+        status_code=200,
+    )
+
+
+@router.post(
+    "/classification/{name}/train/delete",
+    dependencies=[Depends(require_role(["admin"]))],
+)
+def delete_classification_train_images(request: Request, name: str, body: dict = None):
+    config: FrigateConfig = request.app.frigate_config
+
+    if name not in config.classification.custom:
+        return JSONResponse(
+            content=(
+                {
+                    "success": False,
+                    "message": f"{name} is not a known classification model.",
+                }
+            ),
+            status_code=404,
+        )
+
+    json: dict[str, Any] = body or {}
+    list_of_ids = json.get("ids", "")
+    folder = os.path.join(CLIPS_DIR, sanitize_filename(name), "train")
+
+    for id in list_of_ids:
+        file_path = os.path.join(folder, id)
+
+        if os.path.isfile(file_path):
+            os.unlink(file_path)
+
+    return JSONResponse(
+        content=({"success": True, "message": "Successfully deleted faces."}),
+        status_code=200,
+    )