Remove images when wizard is cancelled

This commit is contained in:
Nicolas Mowen 2025-11-10 06:21:15 -07:00
parent 8aad89a83a
commit e0389382f8
2 changed files with 28 additions and 18 deletions

View File

@ -953,31 +953,29 @@ async def generate_object_examples(request: Request, body: GenerateObjectExample
dependencies=[Depends(require_role(["admin"]))], dependencies=[Depends(require_role(["admin"]))],
summary="Delete a classification model", summary="Delete a classification model",
description="""Deletes a specific classification model and all its associated data. description="""Deletes a specific classification model and all its associated data.
The name must exist in the classification models. Returns a success message or an error if the name is invalid.""", Works even if the model is not in the config (e.g., partially created during wizard).
Returns a success message.""",
) )
def delete_classification_model(request: Request, name: str): def delete_classification_model(request: Request, name: str):
config: FrigateConfig = request.app.frigate_config sanitized_name = sanitize_filename(name)
if name not in config.classification.custom:
return JSONResponse(
content=(
{
"success": False,
"message": f"{name} is not a known classification model.",
}
),
status_code=404,
)
# Delete the classification model's data directory in clips # Delete the classification model's data directory in clips
data_dir = os.path.join(CLIPS_DIR, sanitize_filename(name)) data_dir = os.path.join(CLIPS_DIR, sanitized_name)
if os.path.exists(data_dir): if os.path.exists(data_dir):
shutil.rmtree(data_dir) try:
shutil.rmtree(data_dir)
logger.info(f"Deleted classification data directory for {name}")
except Exception as e:
logger.debug(f"Failed to delete data directory for {name}: {e}")
# Delete the classification model's files in model_cache # Delete the classification model's files in model_cache
model_dir = os.path.join(MODEL_CACHE_DIR, sanitize_filename(name)) model_dir = os.path.join(MODEL_CACHE_DIR, sanitized_name)
if os.path.exists(model_dir): if os.path.exists(model_dir):
shutil.rmtree(model_dir) try:
shutil.rmtree(model_dir)
logger.info(f"Deleted classification model directory for {name}")
except Exception as e:
logger.debug(f"Failed to delete model directory for {name}: {e}")
return JSONResponse( return JSONResponse(
content=( content=(

View File

@ -15,6 +15,7 @@ import Step3ChooseExamples, {
} from "./wizard/Step3ChooseExamples"; } from "./wizard/Step3ChooseExamples";
import { cn } from "@/lib/utils"; import { cn } from "@/lib/utils";
import { isDesktop } from "react-device-detect"; import { isDesktop } from "react-device-detect";
import axios from "axios";
const OBJECT_STEPS = [ const OBJECT_STEPS = [
"wizard.steps.nameAndDefine", "wizard.steps.nameAndDefine",
@ -120,7 +121,18 @@ export default function ClassificationModelWizardDialog({
dispatch({ type: "PREVIOUS_STEP" }); dispatch({ type: "PREVIOUS_STEP" });
}; };
const handleCancel = () => { const handleCancel = async () => {
// Clean up any generated training images if we're cancelling from Step 3
if (wizardState.step1Data && wizardState.step3Data?.examplesGenerated) {
try {
await axios.delete(
`/classification/${wizardState.step1Data.modelName}`,
);
} catch (error) {
// Silently fail - user is already cancelling
}
}
dispatch({ type: "RESET" }); dispatch({ type: "RESET" });
onClose(); onClose();
}; };