mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-04-05 14:47:40 +03:00
add ability to reclassify images
This commit is contained in:
parent
573a5ede62
commit
1ca138ff55
@ -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,
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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({
|
||||
<DrawerDescription>Details</DrawerDescription>
|
||||
</DrawerHeader>
|
||||
)}
|
||||
<DropdownMenuLabel>{t("categorizeImageAs")}</DropdownMenuLabel>
|
||||
<DropdownMenuLabel>
|
||||
{dialogLabel ?? t("categorizeImageAs")}
|
||||
</DropdownMenuLabel>
|
||||
<div
|
||||
className={cn(
|
||||
"flex max-h-[40dvh] flex-col overflow-y-auto",
|
||||
isMobile && "gap-2 pb-4",
|
||||
)}
|
||||
>
|
||||
{classes
|
||||
{filteredClasses
|
||||
.sort((a, b) => {
|
||||
if (a === "none") return 1;
|
||||
if (b === "none") return -1;
|
||||
@ -152,7 +173,7 @@ export default function ClassificationSelectionDialog({
|
||||
</div>
|
||||
</SelectorContent>
|
||||
</Selector>
|
||||
<TooltipContent>{t("categorizeImage")}</TooltipContent>
|
||||
<TooltipContent>{tooltipLabel ?? t("categorizeImage")}</TooltipContent>
|
||||
</Tooltip>
|
||||
</div>
|
||||
);
|
||||
|
||||
@ -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<HTMLDivElement | null>(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}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
@ -776,19 +809,23 @@ type DatasetGridProps = {
|
||||
contentRef: MutableRefObject<HTMLDivElement | null>;
|
||||
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)}
|
||||
>
|
||||
<ClassificationSelectionDialog
|
||||
classes={classes}
|
||||
modelName={modelName}
|
||||
image={image}
|
||||
excludeCategory={categoryName}
|
||||
dialogLabel={t("reclassifyImageAs")}
|
||||
tooltipLabel={t("reclassifyImage")}
|
||||
onCategorize={(newCat) => onReclassify(image, newCat)}
|
||||
>
|
||||
<BlurredIconButton>
|
||||
<TbCategoryPlus className="size-5" />
|
||||
</BlurredIconButton>
|
||||
</ClassificationSelectionDialog>
|
||||
<Tooltip>
|
||||
<TooltipTrigger>
|
||||
<LuTrash2
|
||||
|
||||
Loading…
Reference in New Issue
Block a user