113 lines
3.5 KiB
TypeScript
113 lines
3.5 KiB
TypeScript
import { useMemo } from 'react';
|
|
import { DetectionData } from '@/types';
|
|
import {
|
|
getEnabledDetectionTypes,
|
|
DETECTION_TYPES,
|
|
} from '@/constants/detectionModeConfig';
|
|
|
|
export const useFrameDetectionMap = (
|
|
data: DetectionData,
|
|
detectionType: string,
|
|
) => {
|
|
const frameDetectionMap = useMemo(() => {
|
|
const map = new Map<number, any[]>();
|
|
const detectionMode = data.detection_mode || detectionType;
|
|
const enabledTypes = getEnabledDetectionTypes(detectionMode);
|
|
const enabledKeys = new Set(enabledTypes.map((t) => t.id));
|
|
|
|
if (data.frames && Array.isArray(data.frames)) {
|
|
data.frames.forEach((frameData) => {
|
|
const frameId = frameData.frame_id;
|
|
const flatDetections = (frameData as any).detections;
|
|
|
|
if (flatDetections && Array.isArray(flatDetections)) {
|
|
const filteredDetections = flatDetections
|
|
.filter((d: any) => {
|
|
const type = (d.type || '').toLowerCase();
|
|
return enabledKeys.has(type);
|
|
})
|
|
.map((d: any) => ({
|
|
...d,
|
|
_detType: (d.type || '').toLowerCase(),
|
|
// Map old ID fields and new ID fields dynamically for backward compatibility
|
|
[`${(d.type || '').toLowerCase()}_id`]: d.detection_id,
|
|
}));
|
|
|
|
if (filteredDetections.length > 0) {
|
|
map.set(frameId, filteredDetections);
|
|
}
|
|
} else {
|
|
// Legacy format: separate arrays (potholes, signboards, etc.)
|
|
let detections: any[] = [];
|
|
|
|
enabledTypes.forEach((type) => {
|
|
// Try common plural naming conventions for legacy support
|
|
const possibleKeys = [
|
|
`${type.id}s`,
|
|
type.id.endsWith('y')
|
|
? `${type.id.slice(0, -1)}ies`
|
|
: `${type.id}s`,
|
|
type.listKey.replace('_list', 's'),
|
|
type.listKey,
|
|
];
|
|
|
|
for (const listKey of possibleKeys) {
|
|
if ((frameData as any)[listKey]) {
|
|
detections = [
|
|
...detections,
|
|
...(frameData as any)[listKey].map((item: any) => ({
|
|
...item,
|
|
_detType: type.id,
|
|
type: type.id,
|
|
[`${type.id.toLowerCase()}_id`]:
|
|
(item as any).detection_id ??
|
|
(item as any)[`${type.id.toLowerCase()}_id`] ??
|
|
item.id,
|
|
})),
|
|
];
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
if (detections.length > 0) map.set(frameId, detections);
|
|
}
|
|
});
|
|
}
|
|
return map;
|
|
}, [data, detectionType]);
|
|
|
|
const sortedDetectionIndices = useMemo(() => {
|
|
return Array.from(frameDetectionMap.keys()).sort((a, b) => a - b);
|
|
}, [frameDetectionMap]);
|
|
|
|
const getNearestDetections = (
|
|
frame: number,
|
|
sortedIndices: number[],
|
|
maxSkip = 3,
|
|
) => {
|
|
const exact = frameDetectionMap.get(frame);
|
|
if (exact) return exact;
|
|
|
|
let low = 0,
|
|
high = sortedIndices.length - 1,
|
|
targetIndex = -1;
|
|
|
|
while (low <= high) {
|
|
const mid = Math.floor((low + high) / 2);
|
|
if (sortedIndices[mid] <= frame) {
|
|
targetIndex = sortedIndices[mid];
|
|
low = mid + 1;
|
|
} else {
|
|
high = mid - 1;
|
|
}
|
|
}
|
|
|
|
return targetIndex !== -1 && frame - targetIndex <= maxSkip
|
|
? frameDetectionMap.get(targetIndex)
|
|
: undefined;
|
|
};
|
|
|
|
return { frameDetectionMap, getNearestDetections, sortedDetectionIndices };
|
|
};
|