diff --git a/frigate/api/classification.py b/frigate/api/classification.py index 5e1087d17..6b11d3903 100644 --- a/frigate/api/classification.py +++ b/frigate/api/classification.py @@ -787,6 +787,101 @@ def delete_classification_dataset_images( ) +@router.post( + "/classification/{name}/dataset/{category}/reclassify", + response_model=GenericResponse, + dependencies=[Depends(require_role(["admin"]))], + summary="Reclassify a dataset image to a different category", + description="""Moves a single dataset image from one category to another. + The image is re-saved as PNG in the target category and removed from the source.""", +) +def reclassify_classification_image( + 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 {} + image_id = sanitize_filename(json.get("id", "")) + new_category = sanitize_filename(json.get("new_category", "")) + + if not image_id or not new_category: + return JSONResponse( + content=( + { + "success": False, + "message": "Both 'id' and 'new_category' are required.", + } + ), + status_code=400, + ) + + if new_category == category: + return JSONResponse( + content=( + { + "success": False, + "message": "New category must differ from the current category.", + } + ), + status_code=400, + ) + + sanitized_name = sanitize_filename(name) + source_folder = os.path.join( + CLIPS_DIR, sanitized_name, "dataset", sanitize_filename(category) + ) + source_file = os.path.join(source_folder, image_id) + + if not os.path.isfile(source_file): + return JSONResponse( + content=( + { + "success": False, + "message": f"Image not found: {image_id}", + } + ), + status_code=404, + ) + + random_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6)) + timestamp = datetime.datetime.now().timestamp() + new_name = f"{new_category}-{timestamp}-{random_id}.png" + target_folder = os.path.join(CLIPS_DIR, sanitized_name, "dataset", new_category) + + os.makedirs(target_folder, exist_ok=True) + + img = cv2.imread(source_file) + cv2.imwrite(os.path.join(target_folder, new_name), img) + os.unlink(source_file) + + # Clean up empty source folder (unless it is "none") + if ( + os.path.exists(source_folder) + and not os.listdir(source_folder) + and category.lower() != "none" + ): + os.rmdir(source_folder) + + # Mark dataset as changed so UI knows retraining is needed + write_training_metadata(sanitized_name, 0) + + return JSONResponse( + content=({"success": True, "message": "Successfully reclassified image."}), + status_code=200, + ) + + @router.put( "/classification/{name}/dataset/{old_category}/rename", response_model=GenericResponse, diff --git a/web/public/locales/en/views/classificationModel.json b/web/public/locales/en/views/classificationModel.json index 03704fb50..6192cea7f 100644 --- a/web/public/locales/en/views/classificationModel.json +++ b/web/public/locales/en/views/classificationModel.json @@ -26,6 +26,7 @@ "deletedModel_one": "Successfully deleted {{count}} model", "deletedModel_other": "Successfully deleted {{count}} models", "categorizedImage": "Successfully Classified Image", + "reclassifiedImage": "Successfully Reclassified Image", "trainedModel": "Successfully trained model.", "trainingModel": "Successfully started model training.", "updatedModel": "Successfully updated model configuration", @@ -43,7 +44,8 @@ "trainingFailed": "Model training failed. Check Frigate logs for details.", "trainingFailedToStart": "Failed to start model training: {{errorMessage}}", "updateModelFailed": "Failed to update model: {{errorMessage}}", - "renameCategoryFailed": "Failed to rename class: {{errorMessage}}" + "renameCategoryFailed": "Failed to rename class: {{errorMessage}}", + "reclassifyFailed": "Failed to reclassify image: {{errorMessage}}" } }, "deleteCategory": { @@ -92,6 +94,8 @@ }, "categorizeImageAs": "Classify Image As:", "categorizeImage": "Classify Image", + "reclassifyImageAs": "Reclassify Image As:", + "reclassifyImage": "Reclassify Image", "menu": { "objects": "Objects", "states": "States" diff --git a/web/src/components/overlay/ClassificationSelectionDialog.tsx b/web/src/components/overlay/ClassificationSelectionDialog.tsx index f99011423..e64710026 100644 --- a/web/src/components/overlay/ClassificationSelectionDialog.tsx +++ b/web/src/components/overlay/ClassificationSelectionDialog.tsx @@ -34,7 +34,11 @@ type ClassificationSelectionDialogProps = { classes: string[]; modelName: string; image: string; - onRefresh: () => void; + onRefresh?: () => void; + onCategorize?: (category: string) => void; + excludeCategory?: string; + dialogLabel?: string; + tooltipLabel?: string; children: ReactNode; }; export default function ClassificationSelectionDialog({ @@ -43,12 +47,21 @@ export default function ClassificationSelectionDialog({ modelName, image, onRefresh, + onCategorize, + excludeCategory, + dialogLabel, + tooltipLabel, children, }: ClassificationSelectionDialogProps) { const { t } = useTranslation(["views/classificationModel"]); const onCategorizeImage = useCallback( (category: string) => { + if (onCategorize) { + onCategorize(category); + return; + } + axios .post(`/classification/${modelName}/dataset/categorize`, { category, @@ -59,7 +72,7 @@ export default function ClassificationSelectionDialog({ toast.success(t("toast.success.categorizedImage"), { position: "top-center", }); - onRefresh(); + onRefresh?.(); } }) .catch((error) => { @@ -72,7 +85,13 @@ export default function ClassificationSelectionDialog({ }); }); }, - [modelName, image, onRefresh, t], + [modelName, image, onRefresh, onCategorize, t], + ); + + const filteredClasses = useMemo( + () => + excludeCategory ? classes.filter((c) => c !== excludeCategory) : classes, + [classes, excludeCategory], ); const isChildButton = useMemo( @@ -118,14 +137,16 @@ export default function ClassificationSelectionDialog({ Details )} - {t("categorizeImageAs")} + + {dialogLabel ?? t("categorizeImageAs")} +
- {classes + {filteredClasses .sort((a, b) => { if (a === "none") return 1; if (b === "none") return -1; @@ -152,7 +173,7 @@ export default function ClassificationSelectionDialog({
- {t("categorizeImage")} + {tooltipLabel ?? t("categorizeImage")} ); diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx index 25214b703..cb0edb61c 100644 --- a/web/src/views/classification/ModelTrainingView.tsx +++ b/web/src/views/classification/ModelTrainingView.tsx @@ -304,6 +304,37 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) { [pageToggle, model, refreshTrain, refreshDataset, t], ); + const onReclassify = useCallback( + (image: string, newCategory: string) => { + axios + .post( + `/classification/${model.name}/dataset/${pageToggle}/reclassify`, + { + id: image, + new_category: newCategory, + }, + ) + .then((resp) => { + if (resp.status == 200) { + toast.success(t("toast.success.reclassifiedImage"), { + position: "top-center", + }); + refreshDataset(); + } + }) + .catch((error) => { + const errorMessage = + error.response?.data?.message || + error.response?.data?.detail || + "Unknown error"; + toast.error(t("toast.error.reclassifyFailed", { errorMessage }), { + position: "top-center", + }); + }); + }, + [pageToggle, model, refreshDataset, t], + ); + // keyboard const contentRef = useRef(null); @@ -535,10 +566,12 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) { contentRef={contentRef} modelName={model.name} categoryName={pageToggle} + classes={Object.keys(dataset || {})} images={dataset?.[pageToggle] || []} selectedImages={selectedImages} onClickImages={onClickImages} onDelete={onDelete} + onReclassify={onReclassify} /> )} @@ -776,19 +809,23 @@ type DatasetGridProps = { contentRef: MutableRefObject; modelName: string; categoryName: string; + classes: string[]; images: string[]; selectedImages: string[]; onClickImages: (images: string[], ctrl: boolean) => void; onDelete: (ids: string[]) => void; + onReclassify: (image: string, newCategory: string) => void; }; function DatasetGrid({ contentRef, modelName, categoryName, + classes, images, selectedImages, onClickImages, onDelete, + onReclassify, }: DatasetGridProps) { const { t } = useTranslation(["views/classificationModel"]); @@ -816,6 +853,19 @@ function DatasetGrid({ i18nLibrary="views/classificationModel" onClick={(data, _) => onClickImages([data.filename], true)} > + onReclassify(image, newCat)} + > + + + +