Move segment-anything tool (#4735)

Author: jpggvilaca

ui/package-lock.json

@@ -29,7 +29,6 @@ export default defineConfig({
             },
         }),
     ],
-
     source: {
         define: {
             ...publicVars,
@@ -54,4 +53,17 @@ export default defineConfig({
             },
         },
     },
+    server: {
+        headers: {
+            'Cross-Origin-Embedder-Policy': 'credentialless',
+            'Cross-Origin-Opener-Policy': 'same-origin',
+            'Content-Security-Policy':
+                "default-src 'self'; " +
+                "script-src 'self' 'unsafe-eval' blob:; " +
+                "worker-src 'self' blob:; " +
+                "connect-src 'self' http://localhost:7860 data:; " +
+                "img-src 'self' http://localhost:7860 data: blob:; " +
+                "style-src 'self' 'unsafe-inline';",
+        },
+    },
 });

ui/rsbuild.config.ts

@@ -1,6 +1,10 @@
 // Copyright (C) 2025 Intel Corporation
 // SPDX-License-Identifier: Apache-2.0
 
+import { CSSProperties } from 'react';
+
+import { isEmpty } from 'lodash-es';
+
 import { useAnnotator } from '../annotator-provider.component';
 import { useSelectedAnnotations } from '../select-annotation-provider.component';
 import { Annotation } from './annotation.component';
@@ -12,6 +16,17 @@ type AnnotationsProps = {
     isFocussed: boolean;
 };
 
+const DEFAULT_ANNOTATION_STYLES = {
+    fillOpacity: 0.4,
+    fill: 'var(--annotation-fill)',
+    stroke: 'var(--annotation-stroke)',
+    strokeLinecap: 'round',
+    strokeWidth: 'calc(1px / var(--zoom-scale))',
+    strokeDashoffset: 0,
+    strokeDasharray: 0,
+    strokeOpacity: 'var(--annotation-border-opacity, 1)',
+} satisfies CSSProperties;
+
 export const Annotations = ({ width, height, isFocussed }: AnnotationsProps) => {
     const { annotations } = useAnnotator();
     const { selectedAnnotations } = useSelectedAnnotations();
@@ -22,11 +37,17 @@ export const Annotations = ({ width, height, isFocussed }: AnnotationsProps) =>
         ...annotations.filter((a) => selectedAnnotations.has(a.id)),
     ];
 
+    if (isEmpty(annotations)) {
+        return <></>;
+    }
+
     return (
-        <MaskAnnotations annotations={orderedAnnotations} width={width} height={height} isEnabled={isFocussed}>
-            {orderedAnnotations.map((annotation) => (
-                <Annotation annotation={annotation} key={annotation.id} />
-            ))}
-        </MaskAnnotations>
+        <svg width={width} height={height} style={DEFAULT_ANNOTATION_STYLES}>
+            <MaskAnnotations annotations={orderedAnnotations} width={width} height={height} isEnabled={isFocussed}>
+                {orderedAnnotations.map((annotation) => (
+                    <Annotation annotation={annotation} key={annotation.id} />
+                ))}
+            </MaskAnnotations>
+        </svg>
     );
 };

ui/src/features/annotator/annotations/annotations.component.tsx

@@ -1,10 +1,9 @@
 // Copyright (C) 2025 Intel Corporation
 // SPDX-License-Identifier: Apache-2.0
 
-import { CSSProperties, MouseEvent } from 'react';
+import { PointerEvent } from 'react';
 
 import { Grid, View } from '@geti/ui';
-import { isEmpty } from 'lodash-es';
 
 import { ZoomProvider } from '../../components/zoom/zoom';
 import { ZoomTransform } from '../../components/zoom/zoom-transform';
@@ -15,29 +14,20 @@ import { useSelectedAnnotations } from './select-annotation-provider.component';
 import { ToolManager } from './tools/tool-manager.component';
 import { Annotation, DatasetItem } from './types';
 
-const DEFAULT_ANNOTATION_STYLES = {
-    fillOpacity: 0.4,
-    fill: 'var(--annotation-fill)',
-    stroke: 'var(--annotation-stroke)',
-    strokeLinecap: 'round',
-    strokeWidth: 'calc(1px / var(--zoom-scale))',
-    strokeDashoffset: 0,
-    strokeDasharray: 0,
-    strokeOpacity: 'var(--annotation-border-opacity, 1)',
-} satisfies CSSProperties;
-
 type AnnotatorCanvasProps = {
     mediaItem: DatasetItem;
     isFocussed: boolean;
 };
 export const AnnotatorCanvas = ({ mediaItem, isFocussed }: AnnotatorCanvasProps) => {
-    const { setSelectedAnnotations } = useSelectedAnnotations();
     const project_id = useProjectIdentifier();
+    const { setSelectedAnnotations } = useSelectedAnnotations();
+
     const size = { width: mediaItem.width, height: mediaItem.height };
-    // todo: pass media annotations
+    // TODO: pass media annotations
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
     const annotations: Annotation[] = [];
 
-    const handleClickOutside = (e: MouseEvent<SVGSVGElement>): void => {
+    const handleClickOutside = (e: PointerEvent<HTMLDivElement>): void => {
         if (e.target === e.currentTarget) {
             setSelectedAnnotations(new Set());
         }
@@ -51,20 +41,14 @@ export const AnnotatorCanvas = ({ mediaItem, isFocussed }: AnnotatorCanvasProps)
                         <img src={getImageUrl(project_id, String(mediaItem.id))} alt='Collected data' />
                     </View>
 
-                    {!isEmpty(annotations) && (
-                        <View gridArea={'innercanvas'}>
-                            <svg
-                                width={size.width}
-                                height={size.height}
-                                style={DEFAULT_ANNOTATION_STYLES}
-                                onClick={handleClickOutside}
-                            >
+                    <View gridArea={'innercanvas'}>
+                        <>
+                            <div onPointerDown={handleClickOutside}>
                                 <Annotations width={size.width} height={size.height} isFocussed={isFocussed} />
-
-                                <ToolManager />
-                            </svg>
-                        </View>
-                    )}
+                            </div>
+                            <ToolManager />
+                        </>
+                    </View>
                 </Grid>
             </ZoomTransform>
         </ZoomProvider>

ui/src/features/annotator/annotator-canvas.tsx

@@ -6,17 +6,23 @@ import { createContext, Dispatch, ReactNode, SetStateAction, useContext, useStat
 import { v4 as uuid } from 'uuid';
 
 import { ToolType } from '../../components/tool-selection-bar/tools/interface';
-import { Annotation, DatasetItem, Shape } from './types';
+import { useLoadImageQuery } from './hooks/use-load-image-query.hook';
+import { Annotation, DatasetItem, RegionOfInterest, Shape } from './types';
 
 type AnnotatorContext = {
+    // Tools
     activeTool: ToolType | null;
     setActiveTool: Dispatch<SetStateAction<ToolType>>;
 
+    // Annotations
+    annotations: Annotation[];
     addAnnotation: (shape: Shape) => void;
     updateAnnotation: (updatedAnnotation: Annotation) => void;
 
+    // Media item
     mediaItem: DatasetItem;
-    annotations: Annotation[];
+    image: ImageData;
+    roi: RegionOfInterest;
 };
 
 export const AnnotatorProviderContext = createContext<AnnotatorContext | null>(null);
@@ -26,6 +32,8 @@ export const AnnotatorProvider = ({ mediaItem, children }: { mediaItem: DatasetI
     // todo: pass media annotations
     const [annotations, setAnnotations] = useState<Annotation[]>([]);
 
+    const imageQuery = useLoadImageQuery(mediaItem);
+
     const updateAnnotation = (updatedAnnotation: Annotation) => {
         const { id } = updatedAnnotation;
 
@@ -57,6 +65,8 @@ export const AnnotatorProvider = ({ mediaItem, children }: { mediaItem: DatasetI
                 annotations,
 
                 mediaItem,
+                image: imageQuery.data,
+                roi: { x: 0, y: 0, width: mediaItem.width, height: mediaItem.height },
             }}
         >
             {children}

ui/src/features/annotator/annotator-provider.component.tsx

@@ -0,0 +1,30 @@
+// Copyright (C) 2025 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+import { useSuspenseQuery, UseSuspenseQueryResult } from '@tanstack/react-query';
+
+import { API_BASE_URL } from '../../../api/client';
+import { useProjectIdentifier } from '../../../hooks/use-project-identifier.hook';
+import { getImageData, loadImage } from '../tools/utils';
+import { DatasetItem } from '../types';
+
+export const useLoadImageQuery = (mediaItem: DatasetItem | undefined): UseSuspenseQueryResult<ImageData, unknown> => {
+    const projectId = useProjectIdentifier();
+
+    return useSuspenseQuery({
+        queryKey: ['mediaItem', mediaItem?.id, projectId],
+        queryFn: async () => {
+            if (mediaItem === undefined) {
+                throw new Error("Can't fetch undefined media item");
+            }
+
+            const imageUrl = `${API_BASE_URL}/api/projects/${projectId}/dataset/items/${mediaItem.id}/binary`;
+            const image = await loadImage(imageUrl);
+
+            return getImageData(image);
+        },
+        // The image of a media item never changes so we don't want to refetch stale data
+        staleTime: Infinity,
+        retry: 0,
+    });
+};

ui/src/features/annotator/hooks/use-load-image-query.hook.ts

@@ -0,0 +1,44 @@
+// Copyright (C) 2025 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+import { Flex, Heading, View } from '@geti/ui';
+
+import IntelBrandedLoadingGif from '../../assets/intel-loading.webp';
+
+export const AnnotatorLoading = ({ isLoading }: { isLoading: boolean }) => {
+    return (
+        <View
+            position={'absolute'}
+            left={0}
+            top={0}
+            right={0}
+            bottom={0}
+            UNSAFE_style={{
+                backgroundColor: 'var(--spectrum-alias-background-color-modal-overlay)',
+                zIndex: 10,
+            }}
+        >
+            <Flex direction={'column'} alignItems={'center'} justifyContent={'center'} height='100%' gap='size-100'>
+                {/*  eslint-disable-next-line jsx-a11y/img-redundant-alt */}
+                <img
+                    src={IntelBrandedLoadingGif}
+                    // eslint-disable-next-line jsx-a11y/no-noninteractive-element-to-interactive-role
+                    role='progressbar'
+                    alt='Processing image'
+                    style={{
+                        width: 300,
+                        height: 300,
+                    }}
+                />
+                <Heading
+                    level={1}
+                    UNSAFE_style={{
+                        textShadow: '1px 1px 2px black, 1px 1px 2px white',
+                    }}
+                >
+                    {isLoading && 'Processing image, please wait...'}
+                </Heading>
+            </Flex>
+        </View>
+    );
+};

ui/src/features/annotator/hooks/use-segment-anything.hook.ts

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { Divider } from '@geti/ui';
-import { BoundingBox, Selector } from '@geti/ui/icons';
+import { BoundingBox, SegmentAnythingIcon, Selector } from '@geti/ui/icons';
 
 import { ToolConfig } from '../../../components/tool-selection-bar/tools/interface';
 import { Tools } from '../../../components/tool-selection-bar/tools/tools.component';
@@ -17,7 +17,8 @@ const TASK_TOOL_CONFIG: Record<string, ToolConfig[]> = {
     ],
     segmentation: [
         { type: 'selection', icon: Selector },
-        // TODO: Add 'polygon' and 'sam' tools later
+        { type: 'sam', icon: SegmentAnythingIcon },
+        // TODO: Add 'polygon' tool later
     ],
 };
 

ui/src/features/annotator/loading.component.tsx

@@ -6,13 +6,13 @@ import { useAnnotator } from '../../annotator-provider.component';
 import { DrawingBox } from '../drawing-box-tool/drawing-box.component';
 
 export const BoundingBoxTool = () => {
-    const { mediaItem, addAnnotation } = useAnnotator();
+    const { mediaItem, addAnnotation, image } = useAnnotator();
     const { scale: zoom } = useZoom();
 
     return (
         <DrawingBox
             roi={{ x: 0, y: 0, width: mediaItem.width, height: mediaItem.height }}
-            image={new ImageData(mediaItem.width, mediaItem.height)}
+            image={image}
             zoom={zoom}
             onComplete={addAnnotation}
         />