Extract explanations to its own query

jpggvilaca · jpggvilaca · commit 03117ac33e82 · 2025-07-08T11:55:48.000+02:00
diff --git a/web_ui/packages/core/src/requests/query-keys.ts b/web_ui/packages/core/src/requests/query-keys.ts
@@ -118,6 +118,13 @@ const getSelectedMediaItemQueryKeys = () => {
             taskId?: string,
             roiId?: string
         ) => [...commonKey(mediaIdentifier), taskId, roiId, `${prefix}-predictions`, predictionCache],
+
+        EXPLANATIONS: (
+            datasetIdentifier: DatasetIdentifier,
+            mediaIdentifier: MediaIdentifier | undefined,
+            taskId?: string
+        ) => ['explanations', datasetIdentifier, mediaIdentifier, taskId],
+
         SELECTED: (mediaIdentifier: MediaIdentifier | undefined, taskId?: string) => [
             ...commonKey(mediaIdentifier),
             taskId,
diff --git a/web_ui/src/core/annotations/services/prediction-service.interface.ts b/web_ui/src/core/annotations/services/prediction-service.interface.ts
@@ -6,7 +6,10 @@ import { Explanation } from '../prediction.interface';
 
 export interface PredictionResult {
     annotations: ReadonlyArray<Annotation>;
-    maps: Explanation[];
+}
+
+export interface ExplanationResult {
+    maps: ReadonlyArray<Explanation>;
 }
 
 export interface InferenceServerStatusResult {
diff --git a/web_ui/src/pages/annotator/providers/annotator-provider/annotator-provider.component.tsx b/web_ui/src/pages/annotator/providers/annotator-provider/annotator-provider.component.tsx
@@ -69,7 +69,7 @@ export const AnnotatorProvider = ({ children }: AnnotatorProviderProps): JSX.Ele
 
     const [activeTool, setActiveTool] = useState<ToolType>(() => defaultToolForProject(activeDomains));
 
-    const { selectedMediaItem, predictionsQuery } = useSelectedMediaItem();
+    const { selectedMediaItem, predictionsQuery, explanationsQuery } = useSelectedMediaItem();
     const isTaskChainSelectedClassification = isTaskChainDomainSelected(DOMAIN.CLASSIFICATION);
 
     const initialPredictionAnnotations = predictionsQuery.data?.annotations;
@@ -79,6 +79,7 @@ export const AnnotatorProvider = ({ children }: AnnotatorProviderProps): JSX.Ele
         userProjectSettings,
         isTaskChainSelectedClassification
     );
+    const explanations = explanationsQuery.data || [];
 
     const { undoRedoActions, ...userAnnotationScene } = useAnnotationSceneState(
         initialAnnotations,
@@ -109,7 +110,7 @@ export const AnnotatorProvider = ({ children }: AnnotatorProviderProps): JSX.Ele
                                     settings={userProjectSettings}
                                     userAnnotationScene={userAnnotationScene}
                                     initPredictions={initialPredictionAnnotations}
-                                    explanations={selectedMediaItem?.predictions?.maps || EMPTY_EXPLANATION}
+                                    explanations={explanations || EMPTY_EXPLANATION}
                                 >
                                     <SubmitAnnotationsProvider
                                         settings={userProjectSettings}
diff --git a/web_ui/src/pages/annotator/providers/annotator-provider/use-initial-annotations.hook.test.tsx b/web_ui/src/pages/annotator/providers/annotator-provider/use-initial-annotations.hook.test.tsx
@@ -33,7 +33,8 @@ const geSelectedMediaItem = (predictions: Annotation[] = [], annotations: Annota
     ...getMockedImageMediaItem({}),
     image: getMockedImage(),
     annotations,
-    predictions: { annotations: predictions, maps: [] },
+    predictions: { annotations: predictions },
+    explanations: [],
 });
 const getInitialPredictionConfig = (isEnabled = true) => ({
     [FEATURES_KEYS.INITIAL_PREDICTION]: {
diff --git a/web_ui/src/pages/annotator/providers/prediction-provider/prediction-provider.component.tsx b/web_ui/src/pages/annotator/providers/prediction-provider/prediction-provider.component.tsx
@@ -326,7 +326,7 @@ export const PredictionProvider = ({
         selectedInput,
         taskId: String(selectedTask?.id),
         enabled: isPredictionsQueryEnabled,
-        onSuccess: runWhenNotDrawing(({ annotations: newRawPredictions, maps: newMaps }: PredictionResult) => {
+        onSuccess: runWhenNotDrawing(({ annotations: newRawPredictions }: PredictionResult) => {
             const selectedPredictions = selectAnnotations(newRawPredictions, userSceneSelectedInputs);
 
             if (isEmpty(userAnnotations)) {
@@ -348,7 +348,6 @@ export const PredictionProvider = ({
                 canUpdatePrediction && userAnnotationScene.updateAnnotation(newPrediction);
             }
 
-            setExplanations(newMaps);
             setRawPredictions(selectedPredictions);
 
             // Optionally update video timeline predictions
diff --git a/web_ui/src/pages/annotator/providers/prediction-provider/use-prediction-roi-query.hook.test.tsx b/web_ui/src/pages/annotator/providers/prediction-provider/use-prediction-roi-query.hook.test.tsx
@@ -135,7 +135,7 @@ describe('usePredictionsRoiQuery', () => {
         });
 
         await waitFor(() => {
-            expect(result.current.data).toEqual({ annotations: [], maps: [] });
+            expect(result.current.data).toEqual({ annotations: [] });
         });
     });
 
@@ -149,7 +149,6 @@ describe('usePredictionsRoiQuery', () => {
         await waitFor(() => {
             expect(result.current.data).toEqual({
                 annotations: [],
-                maps: [],
             });
         });
 
@@ -236,7 +235,7 @@ describe('usePredictionsRoiQuery', () => {
             });
 
             await waitFor(() => {
-                expect(result.current.data).toEqual({ annotations: [], maps: [] });
+                expect(result.current.data).toEqual({ annotations: [] });
             });
         });
     });
diff --git a/web_ui/src/pages/annotator/providers/prediction-provider/use-prediction-roi-query.hook.ts b/web_ui/src/pages/annotator/providers/prediction-provider/use-prediction-roi-query.hook.ts
@@ -9,11 +9,7 @@ import { useQuery } from '@tanstack/react-query';
 import { AxiosError } from 'axios';
 
 import { TaskChainInput } from '../../../../core/annotations/annotation.interface';
-import {
-    PredictionCache,
-    PredictionMode,
-    PredictionResult,
-} from '../../../../core/annotations/services/prediction-service.interface';
+import { PredictionMode, PredictionResult } from '../../../../core/annotations/services/prediction-service.interface';
 import { getPredictionCache } from '../../../../core/annotations/services/utils';
 import { usePrevious } from '../../../../hooks/use-previous/use-previous.hook';
 import { useProject } from '../../../project-details/providers/project-provider/project-provider.component';
@@ -45,21 +41,17 @@ export const usePredictionsRoiQuery = ({
     const { inferenceService } = useApplicationServices();
 
     const isValidRoi = roiId !== undefined && prevRoi !== roiId;
-
     const isQueryEnabled = enabled && isValidRoi;
 
     const predictionMode = isActiveLearningMode ? PredictionMode.AUTO : PredictionMode.ONLINE;
     const predictionCache = getPredictionCache(predictionMode);
-    const isPredictionCacheNever = predictionCache === PredictionCache.NEVER;
 
     const handleSuccessRef = useRef(onSuccess);
 
     useLayoutEffect(() => {
         handleSuccessRef.current = onSuccess;
     }, [onSuccess]);
 
-    // TODO: extract explanation, combine with other prediction query
-
     const query = useQuery<PredictionResult, AxiosError>({
         queryKey: QUERY_KEYS.SELECTED_MEDIA_ITEM.PREDICTIONS(
             selectedMediaItem?.identifier,
@@ -69,39 +61,25 @@ export const usePredictionsRoiQuery = ({
             roiId
         ),
         queryFn: async ({ signal }) => {
-            if (isQueryEnabled === false) {
-                return { maps: [], annotations: [] };
+            if (!isQueryEnabled || !selectedMediaItem) {
+                return { annotations: [] };
             }
 
-            if (selectedMediaItem === undefined) {
-                return { maps: [], annotations: [] };
-            }
-
-            const explainPromise = isPredictionCacheNever
-                ? inferenceService.getExplanations(datasetIdentifier, selectedMediaItem, taskId, selectedInput, signal)
-                : Promise.resolve([]);
-
-            return Promise.allSettled([
-                inferenceService.getPredictions(
-                    datasetIdentifier,
-                    project.labels,
-                    selectedMediaItem,
-                    predictionCache,
-                    taskId,
-                    selectedInput,
-                    signal
-                ),
-                explainPromise,
-            ]).then(([annotationsResponse, mapsResponse]) => {
-                const maps = mapsResponse.status === 'fulfilled' ? mapsResponse.value : [];
-                const annotations = annotationsResponse.status === 'fulfilled' ? annotationsResponse.value : [];
+            const annotations = await inferenceService.getPredictions(
+                datasetIdentifier,
+                project.labels,
+                selectedMediaItem,
+                predictionCache,
+                taskId,
+                selectedInput,
+                signal
+            );
 
-                handleSuccessRef.current !== undefined && handleSuccessRef.current({ annotations, maps });
+            handleSuccessRef.current?.({ annotations });
 
-                return { annotations, maps };
-            });
+            return { annotations };
         },
-        initialData: { maps: [], annotations: [] },
+        initialData: { annotations: [] },
         enabled: isQueryEnabled,
     });
 
diff --git a/web_ui/src/pages/annotator/providers/selected-media-item-provider/default-selected-media-item-provider.component.tsx b/web_ui/src/pages/annotator/providers/selected-media-item-provider/default-selected-media-item-provider.component.tsx
@@ -7,6 +7,7 @@ import QUERY_KEYS from '@geti/core/src/requests/query-keys';
 import { useQuery, UseQueryResult } from '@tanstack/react-query';
 import { noop } from 'lodash-es';
 
+import { ExplanationResult } from '../../../../core/annotations/services/inference-service.interface';
 import { PredictionResult } from '../../../../core/annotations/services/prediction-service.interface';
 import { SelectedMediaItemContext, SelectedMediaItemProps } from './selected-media-item-provider.component';
 import { SelectedMediaItem } from './selected-media-item.interface';
@@ -40,6 +41,7 @@ export const DefaultSelectedMediaItemProvider = ({
         selectedMediaItemQuery,
         setSelectedMediaItem: noop,
         predictionsQuery: { isLoading: false } as UseQueryResult<PredictionResult>,
+        explanationsQuery: { isLoading: false } as UseQueryResult<ExplanationResult>,
     };
 
     return <SelectedMediaItemContext.Provider value={value}>{children}</SelectedMediaItemContext.Provider>;
diff --git a/web_ui/src/pages/annotator/providers/selected-media-item-provider/selected-media-item-provider.component.tsx b/web_ui/src/pages/annotator/providers/selected-media-item-provider/selected-media-item-provider.component.tsx
@@ -12,6 +12,7 @@ import { AxiosError } from 'axios';
 import { isEmpty, isEqual } from 'lodash-es';
 
 import { Annotation } from '../../../../core/annotations/annotation.interface';
+import { ExplanationResult } from '../../../../core/annotations/services/inference-service.interface';
 import { PredictionMode, PredictionResult } from '../../../../core/annotations/services/prediction-service.interface';
 import { InferenceModel } from '../../../../core/annotations/services/visual-prompt-service';
 import { MediaItem } from '../../../../core/media/media.interface';
@@ -31,6 +32,7 @@ import { hasEmptyLabels } from '../../utils';
 import { useTask } from '../task-provider/task-provider.component';
 import { SelectedMediaItem } from './selected-media-item.interface';
 import { useAnnotationsQuery } from './use-annotations-query.hook';
+import { useExplanationsQuery } from './use-explanation-query.hook';
 import { useLoadImageQuery } from './use-load-image-query.hook';
 import { usePredictionsQuery } from './use-predictions-query.hook';
 import { useSelectedInferenceModel } from './use-selected-inference-model';
@@ -43,6 +45,7 @@ import {
 
 export interface SelectedMediaItemProps {
     predictionsQuery: UseQueryResult<PredictionResult>;
+    explanationsQuery: UseQueryResult<ExplanationResult>;
     selectedMediaItem: SelectedMediaItem | undefined;
     selectedMediaItemQuery: UseQueryResult<SelectedMediaItem>;
     setSelectedMediaItem: (media: MediaItem | undefined) => void;
@@ -87,7 +90,7 @@ const useIsSuggestPredictionEnabled = (projectIdentifier: ProjectIdentifier): bo
 };
 
 /**
- * Load either online or auto predicitons based on the annotator mode
+ * Load either online or auto predictions based on the annotator mode
  * When the user switches between both modes we don't want to refetch predictions,
  * so in both cases we will keep the queries mounted
  */
@@ -169,6 +172,10 @@ export const SelectedMediaItemProvider = ({ children }: SelectedMediaItemProvide
     });
 
     const predictionsQuery = usePredictionsQueryBasedOnAnnotatorMode(mediaItem);
+    const explanationsQuery = useExplanationsQuery({
+        datasetIdentifier,
+        mediaItem,
+    });
 
     const selectedMediaItemQueryKey = [
         ...QUERY_KEYS.SELECTED_MEDIA_ITEM.SELECTED(pendingMediaItem?.identifier, selectedTask?.id),
@@ -201,7 +208,7 @@ export const SelectedMediaItemProvider = ({ children }: SelectedMediaItemProvide
                             }
                         }
 
-                        const predictionsData = predictionsQuery.data ?? { maps: [], annotations: [] };
+                        const predictionsData = predictionsQuery.data ?? { annotations: [] };
 
                         resolve([imageQuery.data, annotationsQuery.data, predictionsData]);
                     }
@@ -266,6 +273,7 @@ export const SelectedMediaItemProvider = ({ children }: SelectedMediaItemProvide
 
     const value: SelectedMediaItemProps = {
         predictionsQuery,
+        explanationsQuery,
         selectedMediaItem,
         selectedMediaItemQuery,
         setSelectedMediaItem: setPendingMediaItem,
diff --git a/web_ui/src/pages/annotator/providers/selected-media-item-provider/selected-media-item.interface.ts b/web_ui/src/pages/annotator/providers/selected-media-item-provider/selected-media-item.interface.ts
@@ -2,11 +2,13 @@
 // LIMITED EDGE SOFTWARE DISTRIBUTION LICENSE
 
 import { Annotation } from '../../../../core/annotations/annotation.interface';
+import { ExplanationResult } from '../../../../core/annotations/services/inference-service.interface';
 import { PredictionResult } from '../../../../core/annotations/services/prediction-service.interface';
 import { MediaItem } from '../../../../core/media/media.interface';
 
 export type SelectedMediaItem = MediaItem & {
     image: ImageData;
     annotations: Annotation[];
     predictions?: PredictionResult;
+    explanations?: ExplanationResult;
 };
diff --git a/web_ui/src/pages/annotator/providers/selected-media-item-provider/use-explanation-query.hook.ts b/web_ui/src/pages/annotator/providers/selected-media-item-provider/use-explanation-query.hook.ts
@@ -0,0 +1,45 @@
+// Copyright (C) 2022-2025 Intel Corporation
+// LIMITED EDGE SOFTWARE DISTRIBUTION LICENSE
+
+import QUERY_KEYS from '@geti/core/src/requests/query-keys';
+import { useApplicationServices } from '@geti/core/src/services/application-services-provider.component';
+import { QueryKey, useQuery, UseQueryResult } from '@tanstack/react-query';
+import { AxiosError } from 'axios';
+
+import { Explanation } from '../../../../core/annotations/prediction.interface';
+import { MediaItem } from '../../../../core/media/media.interface';
+import { DatasetIdentifier } from '../../../../core/projects/dataset.interface';
+
+interface UseGetExplanations {
+    datasetIdentifier: DatasetIdentifier;
+    mediaItem: MediaItem | undefined;
+    enabled?: boolean;
+    taskId?: string;
+}
+
+export const useExplanationsQuery = ({
+    datasetIdentifier,
+    mediaItem,
+    enabled = true,
+    taskId,
+}: UseGetExplanations): UseQueryResult<Explanation[], AxiosError> => {
+    const { inferenceService } = useApplicationServices();
+
+    const queryKey: QueryKey = QUERY_KEYS.SELECTED_MEDIA_ITEM.EXPLANATIONS(
+        datasetIdentifier,
+        mediaItem?.identifier,
+        taskId
+    );
+
+    return useQuery<Explanation[], AxiosError>({
+        queryKey,
+        queryFn: async ({ signal }) => {
+            if (!mediaItem) throw new Error("Can't fetch undefined media item");
+
+            return inferenceService.getExplanations(datasetIdentifier, mediaItem, taskId, undefined, signal);
+        },
+        enabled: enabled && !!mediaItem,
+        staleTime: 5 * 60_000,
+        gcTime: 0,
+    });
+};
diff --git a/web_ui/src/pages/annotator/providers/selected-media-item-provider/use-predictions-query.hook.ts b/web_ui/src/pages/annotator/providers/selected-media-item-provider/use-predictions-query.hook.ts
diff --git a/web_ui/src/pages/project-details/components/project-model/training-dataset/training-dataset-details-preview/training-dataset-details-preview.component.tsx b/web_ui/src/pages/project-details/components/project-model/training-dataset/training-dataset-details-preview/training-dataset-details-preview.component.tsx
diff --git a/web_ui/src/pages/project-details/components/project-test/test-details-preview/test-details-preview.component.tsx b/web_ui/src/pages/project-details/components/project-test/test-details-preview/test-details-preview.component.tsx
diff --git a/web_ui/src/pages/project-details/components/project-test/test-details-preview/use-test-results-query.hook.ts b/web_ui/src/pages/project-details/components/project-test/test-details-preview/use-test-results-query.hook.ts
