Files
labelmain/components/label/store.ts
2026-02-27 13:57:42 +08:00

555 lines
14 KiB
TypeScript

import { Comment } from "./api/label/typing"
import { del, get, set } from "idb-keyval"
import superjson from "superjson"
import { create } from "zustand"
import { persist, PersistStorage, StorageValue } from "zustand/middleware"
type LabelData = Map<string, Map<string, [number, any[], any, any[]][]>>
interface LabelState {
// pathId pathPoints detail hollowPathPoints
label: LabelData
labelTime: {
label: number
review1: number
review2: number
}
labelDefaultComments: Comment[]
setLabel: (label: LabelData) => void
stateStack: LabelData[]
setStateStack: (data: LabelData[]) => void
popStateStack: () => LabelData | null
pushStateStack: (newData: LabelData) => void
getLabelDetailDataByKeys: (
imageId: string,
categoryId: string,
id: number
) => any
updateLabelDetailDataByKeys: (
imageId: string,
categoryId: string,
id: number,
newDetail: any
) => void
addDetail: (
imageId: string,
categoryId: string,
id: number,
obj: {
[keys: string]: any
}
) => void
setLabelTime: (labelTime: {
label: number
review1: number
review2: number
}) => void
setLabelDefaultComments: (arr: Comment[]) => void
setImage: (
imageId: string,
operation: Map<string, [number, any[], any, any[]][]>
) => void
setOperation: (
imageId: string,
operationId: string,
path: [number, any[], any, any[]]
) => void
deleteOperation: (
imageId: string,
operationId: string,
pathId: number
) => void
resetLabel: () => void
}
interface ObjectState {
pathStatus: {
// key = imageId + pathId
[key: string]: {
["HIDE"]?: boolean
}
}
selectedPath: {
// key = imageId
[key: string]: number[]
}
operationStatus: {
// key = imageId + operationId
[key: string]: {
["HIDE"]?: boolean
["INPUT"]?: string
["COLLAPSE"]?: boolean
}
}
groupStatus: {
// key = imageId + groupId
[key: string]: {
["COLLAPSE"]?: boolean
}
}
imageStatus: {
// key = imageId
[key: string]: {
["INPUT"]?: string
["COLLAPSE"]?: boolean
["FILTERCOLOR"]?: {
[key: string]: boolean
}
}
}
updatePathStatus: (
id: string,
statusKey: "HIDE",
statusValue: boolean
) => void
updateSelectedPath: (
id: string,
type: "ADD" | "DELETE" | "ADDSOME",
pathValue?: number
) => void
resetPathAndOperationStatus: () => void
updateOperationStatus: (
id: string,
statusKey: "HIDE" | "COLLAPSE" | "INPUT",
statusValue: string | boolean
) => void
updateGroupStatus: (
id: string,
statusKey: "COLLAPSE",
statusValue: string | boolean
) => void
updateImageStatus: (
id: string,
statusKey: "INPUT" | "COLLAPSE" | "FILTERCOLOR",
statusValue:
| string
| boolean
| {
[key: string]: boolean
}
) => void
}
interface KeyEventState {
shift: boolean
focusInput: boolean
setShift: (value: boolean) => void
setFocusInput: (value: boolean) => void
}
const initialLabelState = {
label: new Map(),
labelTime: {
label: 0,
review1: 0,
review2: 0,
},
}
export const initialDetail = {
comment: [],
create_timestamp: 0,
first_modified_timestamp: 0,
first_modified_uid: 0,
last_modified_timestamp: 0,
last_modified_uid: 0,
sub_attributes: {},
prelabel: false,
}
// 自定义 storage 对象
export const storage: PersistStorage<any> = {
getItem: async (name: string) => {
// console.log(name, "has been retrieved");
let str = await get(name)
if (!str) return null
return superjson.parse(str) || null
},
setItem: async (name: string, value: StorageValue<any>): Promise<void> => {
// console.log(name, "with value", value, "has been saved");
await set(name, superjson.stringify(value))
},
removeItem: async (name: string): Promise<void> => {
console.log(name, "has been deleted")
await del(name)
},
}
// store path data
export const useLabelStore = create(
persist<LabelState>(
(storeSet) => ({
label: initialLabelState.label,
labelTime: initialLabelState.labelTime,
labelDefaultComments: [],
stateStack: [],
setLabel: (label) =>
storeSet((state) => {
return {
...state,
label,
}
}),
setStateStack: (data) =>
storeSet((state) => {
return {
...state,
stateStack: data,
}
}),
popStateStack: () => {
const arr: LabelData[] = useLabelStore.getState().stateStack
console.log("弹出栈", arr)
if (!arr.length) return null
if (arr.length === 1) {
return arr[0]
} else {
arr.pop()
storeSet((state) => {
return {
...state,
stateStack: arr,
}
})
return arr[arr.length - 1]
}
},
pushStateStack: (newData) => {
const arr: LabelData[] = useLabelStore.getState().stateStack
if (arr.length === 10) {
arr.shift()
}
arr.push(newData)
console.log(arr)
return storeSet((state) => {
return {
...state,
stateStack: arr,
}
})
},
getLabelDetailDataByKeys: (
imageId: string,
categoryId: string,
id: number
) => {
const data = useLabelStore.getState().label
if (!data.has(imageId)) return null
const image = data.get(imageId)
if (!image?.has(categoryId)) return null
const category = image.get(categoryId)
let detail = null
category?.map((item) => {
if (item[0] === id) detail = item[2]
})
return detail
},
updateLabelDetailDataByKeys: (imageId, categoryId, id, newData) => {
const labelData = useLabelStore.getState().label
if (!labelData.has(imageId)) return
const image = labelData.get(imageId)
if (!image?.has(categoryId)) return
const category = image.get(categoryId)
category?.map((item) => {
if (item[0] === id) item[2] = newData
})
},
addDetail: (imageId, categoryId, id, obj) => {
const labelData = useLabelStore.getState().label
if (!labelData.has(imageId)) return
const image = labelData.get(imageId)
if (!image?.has(categoryId)) return
const category = image.get(categoryId)
category?.map((item) => {
if (item[0] === id) item[2] = Object.assign(item[2], obj)
})
},
setLabelTime: (labelTime) =>
storeSet((state) => {
return {
...state,
labelTime,
}
}),
setLabelDefaultComments: (labelDefaultComments) =>
storeSet((state) => {
return {
...state,
labelDefaultComments,
}
}),
setImage: (imageId, operation) =>
storeSet((state) => {
state.label.set(imageId, operation)
return state
}),
setOperation: (imageId, operationId, path) =>
storeSet((state) => {
const labelData = useLabelStore.getState().label
if (!labelData.has(imageId)) {
labelData.set(imageId, new Map())
}
if (labelData.get(imageId)?.get(operationId)?.length) {
let result = [...labelData.get(imageId)?.get(operationId)!, path]
.reduce((acc, curr) => {
const key = curr[0]
acc.set(key, curr)
return acc
}, new Map())
.values()
labelData.get(imageId)?.set(operationId, Array.from(result))
} else {
labelData.get(imageId)?.set(operationId, [path])
}
console.log("绘制时", labelData)
useLabelStore.getState().pushStateStack(structuredClone(labelData))
return { ...state, label: structuredClone(labelData) }
}),
// 删除path
deleteOperation: (imageId, operationId, pathId) =>
storeSet((state) => {
if (state.label.get(imageId)?.get(operationId)?.length) {
let result =
state.label
.get(imageId)
?.get(operationId)
?.filter((item) => item[0] !== pathId) || []
state.label.get(imageId)?.set(operationId, result)
}
return state
}),
resetLabel: () =>
storeSet((state) => {
return {
...state,
label: initialLabelState.label,
labelTime: initialLabelState.labelTime,
}
}),
}),
{
name: "label-data-storage",
storage: storage,
}
)
)
const initialObjectState = {
pathStatus: {},
selectedPath: {},
operationStatus: {},
groupStatus: {},
imageStatus: {},
}
// store object list status
export const useObjectStore = create(
persist<ObjectState>(
(storeSet) => ({
pathStatus: initialObjectState.pathStatus,
selectedPath: initialObjectState.selectedPath,
operationStatus: initialObjectState.operationStatus,
groupStatus: initialObjectState.groupStatus,
imageStatus: initialObjectState.imageStatus,
updatePathStatus: (id: string, statusKey: "HIDE", statusValue: boolean) =>
storeSet((state) => {
let idStatus = state.pathStatus[id] || {}
idStatus = Object.assign(idStatus, {
[statusKey]: statusValue,
})
return {
...state,
pathStatus: {
...state.pathStatus,
[id]: idStatus,
},
}
}),
updateSelectedPath: (
id: string,
type: "ADD" | "DELETE" | "ADDSOME",
pathValue?: number
) =>
storeSet((state) => {
let idSelectedPath = state.selectedPath[id] || []
if (type === "ADD") {
if (useKeyEventStore.getState().shift) {
idSelectedPath = pathValue
? [...new Set([...idSelectedPath, pathValue])]
: idSelectedPath
} else {
idSelectedPath = pathValue ? [pathValue] : []
}
} else if (type === "ADDSOME") {
idSelectedPath = pathValue
? [...new Set([...idSelectedPath, pathValue])]
: idSelectedPath
} else if (type === "DELETE") {
idSelectedPath =
idSelectedPath?.filter((id) => id !== pathValue) || []
}
return {
...state,
selectedPath: {
...state.selectedPath,
[id]: idSelectedPath,
},
}
}),
resetPathAndOperationStatus: () =>
storeSet((state) => {
return {
...state,
selectedPath: {},
pathStatus: {},
operationStatus: {},
}
}),
updateOperationStatus: (
id: string,
statusKey: "HIDE" | "COLLAPSE" | "INPUT",
statusValue: string | boolean
) =>
storeSet((state) => {
let idStatus = state.operationStatus[id] || {}
idStatus = Object.assign(idStatus, {
[statusKey]: statusValue,
})
return {
...state,
operationStatus: {
...state.operationStatus,
[id]: idStatus,
},
}
}),
updateGroupStatus: (
id: string,
statusKey: "COLLAPSE",
statusValue: string | boolean
) =>
storeSet((state) => {
let idStatus = state.groupStatus[id] || {}
idStatus = Object.assign(idStatus, {
[statusKey]: statusValue,
})
return {
...state,
groupStatus: {
...state.groupStatus,
[id]: idStatus,
},
}
}),
updateImageStatus: (
id: string,
statusKey: "INPUT" | "COLLAPSE" | "FILTERCOLOR",
statusValue:
| string
| boolean
| {
[key: string]: boolean
}
) =>
storeSet((state) => {
let idStatus = state.imageStatus[id] || {}
idStatus = Object.assign(idStatus, {
[statusKey]: statusValue,
})
return {
...state,
imageStatus: {
...state.imageStatus,
[id]: idStatus,
},
}
}),
}),
{
name: "object-data-storage",
storage: storage,
partialize: (state: ObjectState) =>
Object.fromEntries(
Object.entries(state).filter(
([key]) => !["selectedPath"].includes(key)
)
) as any,
}
)
)
const initialKeyEventState = {
shift: false,
focusInput: false,
}
// store keyboard and mouse status
export const useKeyEventStore = create<KeyEventState>((set) => ({
shift: initialKeyEventState.shift,
setShift: (value: boolean) =>
set((state: KeyEventState) => ({ ...state, shift: value })),
focusInput: initialKeyEventState.focusInput,
setFocusInput: (value: boolean) =>
set((state: KeyEventState) => ({ ...state, focusInput: value })),
}))
interface ImageState {
images: Map<string, string>
setImages: (value: any) => void
}
interface VideoFrameState {
videoFrames: Map<string, string[]>
setVideoFrames: (key: string, value: string[]) => void
removeVideoFrames: (key: string) => void
}
export const useImagesStore = create(
persist<ImageState>(
(storeSet) => ({
images: new Map(),
setImages: (value: any) =>
storeSet((state) => {
return {
...state,
images: value,
}
}),
}),
{
name: "image-data-storage",
storage: storage,
}
)
)
export const useVideoFrameStore = create(
persist<VideoFrameState>(
(storeSet) => ({
videoFrames: new Map(),
setVideoFrames: (key: string, value: string[]) =>
storeSet((state) => {
const nextMap = new Map(state.videoFrames)
nextMap.set(key, value)
return {
...state,
videoFrames: nextMap,
}
}),
removeVideoFrames: (key: string) =>
storeSet((state) => {
const nextMap = new Map(state.videoFrames)
nextMap.delete(key)
return {
...state,
videoFrames: nextMap,
}
}),
}),
{
name: "video-frame-storage",
storage: storage,
}
)
)