Refactor state training grid to use classification card

This commit is contained in:
Nicolas Mowen 2025-10-07 07:16:59 -06:00
parent 3378ad07b1
commit bc9a63d905
2 changed files with 91 additions and 81 deletions

View File

@ -16,6 +16,7 @@ type ClassificationCardProps = {
threshold?: ClassificationThreshold; threshold?: ClassificationThreshold;
selected: boolean; selected: boolean;
i18nLibrary: string; i18nLibrary: string;
showArea?: boolean;
onClick: (data: ClassificationItemData, meta: boolean) => void; onClick: (data: ClassificationItemData, meta: boolean) => void;
children?: React.ReactNode; children?: React.ReactNode;
}; };
@ -26,6 +27,7 @@ export function ClassificationCard({
threshold, threshold,
selected, selected,
i18nLibrary, i18nLibrary,
showArea = true,
onClick, onClick,
children, children,
}: ClassificationCardProps) { }: ClassificationCardProps) {
@ -55,12 +57,12 @@ export function ClassificationCard({
}); });
const imageArea = useMemo(() => { const imageArea = useMemo(() => {
if (imgRef.current == null || !imageLoaded) { if (!showArea || imgRef.current == null || !imageLoaded) {
return undefined; return undefined;
} }
return imgRef.current.naturalWidth * imgRef.current.naturalHeight; return imgRef.current.naturalWidth * imgRef.current.naturalHeight;
}, [imageLoaded]); }, [showArea, imageLoaded]);
return ( return (
<> <>

View File

@ -60,7 +60,7 @@ import { IoMdArrowRoundBack } from "react-icons/io";
import { MdAutoFixHigh } from "react-icons/md"; import { MdAutoFixHigh } from "react-icons/md";
import TrainFilterDialog from "@/components/overlay/dialog/TrainFilterDialog"; import TrainFilterDialog from "@/components/overlay/dialog/TrainFilterDialog";
import useApiFilter from "@/hooks/use-api-filter"; import useApiFilter from "@/hooks/use-api-filter";
import { TrainFilter } from "@/types/classification"; import { ClassificationItemData, TrainFilter } from "@/types/classification";
import { ClassificationCard } from "@/components/card/ClassificationCard"; import { ClassificationCard } from "@/components/card/ClassificationCard";
type ModelTrainingViewProps = { type ModelTrainingViewProps = {
@ -682,20 +682,18 @@ function TrainGrid({
onRefresh, onRefresh,
onDelete, onDelete,
}: TrainGridProps) { }: TrainGridProps) {
const { t } = useTranslation(["views/classificationModel"]); const trainData = useMemo<ClassificationItemData[]>(
const trainData = useMemo(
() => () =>
trainImages trainImages
.map((raw) => { .map((raw) => {
const parts = raw.replaceAll(".webp", "").split("-"); const parts = raw.replaceAll(".webp", "").split("-");
const rawScore = Number.parseFloat(parts[2]); const rawScore = Number.parseFloat(parts[2]);
return { return {
raw, filename: raw,
timestamp: parts[0], filepath: `clips/${model.name}/train/${raw}`,
label: parts[1], timestamp: Number.parseFloat(parts[0]),
score: rawScore * 100, name: parts[1],
truePositive: rawScore >= model.threshold, score: rawScore,
}; };
}) })
.filter((data) => { .filter((data) => {
@ -703,10 +701,7 @@ function TrainGrid({
return true; return true;
} }
if ( if (trainFilter.classes && !trainFilter.classes.includes(data.name)) {
trainFilter.classes &&
!trainFilter.classes.includes(data.label)
) {
return false; return false;
} }
@ -726,10 +721,57 @@ function TrainGrid({
return true; return true;
}) })
.sort((a, b) => b.timestamp.localeCompare(a.timestamp)), .sort((a, b) => b.timestamp - a.timestamp),
[model, trainImages, trainFilter], [model, trainImages, trainFilter],
); );
if (model.state_config) {
return (
<StateTrainGrid
model={model}
contentRef={contentRef}
classes={classes}
trainData={trainData}
selectedImages={selectedImages}
onClickImages={onClickImages}
onRefresh={onRefresh}
onDelete={onDelete}
/>
);
}
return <div />;
}
type StateTrainGridProps = {
model: CustomClassificationModelConfig;
contentRef: MutableRefObject<HTMLDivElement | null>;
classes: string[];
trainData?: ClassificationItemData[];
selectedImages: string[];
onClickImages: (images: string[], ctrl: boolean) => void;
onRefresh: () => void;
onDelete: (ids: string[]) => void;
};
function StateTrainGrid({
model,
contentRef,
classes,
trainData,
selectedImages,
onClickImages,
onRefresh,
onDelete,
}: StateTrainGridProps) {
const { t } = useTranslation(["views/classificationModel"]);
const threshold = useMemo(() => {
return {
recognition: model.threshold,
unknown: model.threshold,
};
}, [model]);
return ( return (
<div <div
ref={contentRef} ref={contentRef}
@ -739,73 +781,39 @@ function TrainGrid({
)} )}
> >
{trainData?.map((data) => ( {trainData?.map((data) => (
<div <ClassificationCard
key={data.timestamp} className="w-60 gap-2 rounded-lg bg-card p-2"
className={cn( imgClassName="size-auto"
"flex w-56 cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]", data={data}
selectedImages.includes(data.raw) threshold={threshold}
? "shadow-selected outline-selected" selected={selectedImages.includes(data.filename)}
: "outline-transparent duration-500", i18nLibrary="views/classificationModel"
isMobile && "w-[48%]", showArea={false}
)} onClick={(data, meta) => onClickImages([data.filename], meta)}
onClick={(e) => {
e.stopPropagation();
onClickImages([data.raw], e.ctrlKey || e.metaKey);
}}
> >
<div <ClassificationSelectionDialog
className={cn( classes={classes}
"w-full overflow-hidden p-2 *:text-card-foreground", modelName={model.name}
isMobile && "flex justify-center", image={data.filename}
)} onRefresh={onRefresh}
> >
<img <TbCategoryPlus className="size-5 cursor-pointer text-primary-variant hover:text-primary" />
className="w-56 rounded-lg" </ClassificationSelectionDialog>
src={`${baseUrl}clips/${model.name}/train/${data.raw}`} <Tooltip>
/> <TooltipTrigger>
</div> <LuTrash2
<div className="rounded-b-lg bg-card p-3"> className="size-5 cursor-pointer text-primary-variant hover:text-primary"
<div className="flex w-full flex-row items-center justify-between gap-2"> onClick={(e) => {
<div className="flex flex-col items-start text-xs text-primary-variant"> e.stopPropagation();
<div className="smart-capitalize"> onDelete([data.filename]);
{data.label.replaceAll("_", " ")} }}
</div> />
<div </TooltipTrigger>
className={cn( <TooltipContent>
"", {t("button.deleteClassificationAttempts")}
data.truePositive ? "text-success" : "text-danger", </TooltipContent>
)} </Tooltip>
> </ClassificationCard>
{data.score}%
</div>
</div>
<div className="flex flex-row items-start justify-end gap-5 md:gap-4">
<ClassificationSelectionDialog
classes={classes}
modelName={model.name}
image={data.raw}
onRefresh={onRefresh}
>
<TbCategoryPlus className="size-5 cursor-pointer text-primary-variant hover:text-primary" />
</ClassificationSelectionDialog>
<Tooltip>
<TooltipTrigger>
<LuTrash2
className="size-5 cursor-pointer text-primary-variant hover:text-primary"
onClick={(e) => {
e.stopPropagation();
onDelete([data.raw]);
}}
/>
</TooltipTrigger>
<TooltipContent>
{t("button.deleteClassificationAttempts")}
</TooltipContent>
</Tooltip>
</div>
</div>
</div>
</div>
))} ))}
</div> </div>
); );