diff --git a/frigate/api/classification.py b/frigate/api/classification.py index 1b91afeea..60f1899d5 100644 --- a/frigate/api/classification.py +++ b/frigate/api/classification.py @@ -31,9 +31,10 @@ from frigate.api.defs.response.generic_response import GenericResponse from frigate.api.defs.tags import Tags from frigate.config import FrigateConfig from frigate.config.camera import DetectConfig -from frigate.const import CLIPS_DIR, FACE_DIR +from frigate.const import CLIPS_DIR, FACE_DIR, MODEL_CACHE_DIR from frigate.embeddings import EmbeddingsContext from frigate.models import Event +from frigate.util.builtin import find_config_file, update_yaml_file_bulk from frigate.util.classification import ( collect_object_classification_examples, collect_state_classification_examples, @@ -828,12 +829,34 @@ def delete_classification_model(request: Request, name: str): status_code=404, ) - # Delete the classification model's data directory - model_dir = os.path.join(CLIPS_DIR, sanitize_filename(name)) + # Delete the classification model's data directory in clips + data_dir = os.path.join(CLIPS_DIR, sanitize_filename(name)) + if os.path.exists(data_dir): + shutil.rmtree(data_dir) + # Delete the classification model's files in model_cache + model_dir = os.path.join(MODEL_CACHE_DIR, sanitize_filename(name)) if os.path.exists(model_dir): shutil.rmtree(model_dir) + # Remove the model from the config file + config_file = find_config_file() + try: + # Setting value to empty string deletes the key + updates = {f"classification.custom.{name}": None} + update_yaml_file_bulk(config_file, updates) + except Exception as e: + logger.error(f"Error updating config file: {e}") + return JSONResponse( + content=( + { + "success": False, + "message": f"Failed to update config file: {str(e)}", + } + ), + status_code=500, + ) + return JSONResponse( content=( {