diff --git a/frigate/api/classification.py b/frigate/api/classification.py index 9b116be10..6a738409e 100644 --- a/frigate/api/classification.py +++ b/frigate/api/classification.py @@ -870,6 +870,46 @@ def categorize_classification_image(request: Request, name: str, body: dict = No ) +@router.post( + "/classification/{name}/dataset/{category}/create", + response_model=GenericResponse, + dependencies=[Depends(require_role(["admin"]))], + summary="Create an empty classification category folder", + description="""Creates an empty folder for a classification category. + This is used to create folders for categories that don't have images yet. + Returns a success message or an error if the name is invalid.""", +) +def create_classification_category(request: Request, name: str, category: str): + 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, + ) + + category_folder = os.path.join( + CLIPS_DIR, sanitize_filename(name), "dataset", sanitize_filename(category) + ) + + os.makedirs(category_folder, exist_ok=True) + + return JSONResponse( + content=( + { + "success": True, + "message": f"Successfully created category folder: {category}", + } + ), + status_code=200, + ) + + @router.post( "/classification/{name}/train/delete", response_model=GenericResponse, diff --git a/web/public/locales/en/views/classificationModel.json b/web/public/locales/en/views/classificationModel.json index 6641b607e..7852bb550 100644 --- a/web/public/locales/en/views/classificationModel.json +++ b/web/public/locales/en/views/classificationModel.json @@ -166,6 +166,7 @@ "noImages": "No sample images generated", "classifying": "Classifying & Training...", "trainingStarted": "Training started successfully", + "modelCreated": "Model created successfully. Use the Recent Classifications view to add images for missing states, then train the model.", "errors": { "noCameras": "No cameras configured", "noObjectLabel": "No object label selected", diff --git a/web/src/components/classification/wizard/Step3ChooseExamples.tsx b/web/src/components/classification/wizard/Step3ChooseExamples.tsx index bae697f6e..47c990071 100644 --- a/web/src/components/classification/wizard/Step3ChooseExamples.tsx +++ b/web/src/components/classification/wizard/Step3ChooseExamples.tsx @@ -141,15 +141,49 @@ export default function Step3ChooseExamples({ ); await Promise.all(categorizePromises); - // Step 3: Kick off training - await axios.post(`/classification/${step1Data.modelName}/train`); + // Step 2.5: Create empty folders for classes that don't have any images + // This ensures all classes are available in the dataset view later + const classesWithImages = new Set( + Object.values(classifications).filter((c) => c && c !== "none"), + ); + const emptyFolderPromises = step1Data.classes + .filter((className) => !classesWithImages.has(className)) + .map((className) => + axios.post( + `/classification/${step1Data.modelName}/dataset/${className}/create`, + ), + ); + await Promise.all(emptyFolderPromises); - toast.success(t("wizard.step3.trainingStarted"), { - closeButton: true, - }); - setIsTraining(true); + // Step 3: Determine if we should train + // For state models, we need ALL states to have examples + // For object models, we need at least 2 classes with images + const allStatesHaveExamplesForTraining = + step1Data.modelType !== "state" || + step1Data.classes.every((className) => + classesWithImages.has(className), + ); + const shouldTrain = + allStatesHaveExamplesForTraining && classesWithImages.size >= 2; + + // Step 4: Kick off training only if we have enough classes with images + if (shouldTrain) { + await axios.post(`/classification/${step1Data.modelName}/train`); + + toast.success(t("wizard.step3.trainingStarted"), { + closeButton: true, + }); + setIsTraining(true); + } else { + // Don't train - not all states have examples + toast.success(t("wizard.step3.modelCreated"), { + closeButton: true, + }); + setIsTraining(false); + onClose(); + } }, - [step1Data, step2Data, t], + [step1Data, step2Data, t, onClose], ); const handleContinueClassification = useCallback(async () => {