@cornerstonejs/ai 3.0.0-beta.2 → 3.0.0-beta.4

package/dist/esm/ONNXSegmentationController.d.ts

@@ -0,0 +1,130 @@
+import type { Types } from '@cornerstonejs/core';
+import * as cornerstoneTools from '@cornerstonejs/tools';
+export type ModelType = {
+    name: string;
+    key: string;
+    url: string;
+    size: number;
+    opt?: Record<string, unknown>;
+};
+export declare enum Loggers {
+    Log = "status",
+    Encoder = "encoder",
+    Decoder = "decoder"
+}
+export default class ONNXSegmentationController {
+    static MarkerInclude: string;
+    static MarkerExclude: string;
+    static BoxPrompt: string;
+    static viewportOptions: {
+        displayArea: Types.DisplayArea;
+        background: Types.Point3;
+    };
+    maxWidth: number;
+    maxHeight: number;
+    modelWidth: number;
+    modelHeight: number;
+    tool: any;
+    static MODELS: {
+        sam_l: ({
+            name: string;
+            url: string;
+            size: number;
+            key: string;
+            feedType: string;
+        } | {
+            name: string;
+            url: string;
+            size: number;
+            key: string;
+            feedType?: undefined;
+        })[];
+        sam_h: ({
+            name: string;
+            url: string;
+            size: number;
+            key: string;
+            feedType: string;
+        } | {
+            name: string;
+            url: string;
+            size: number;
+            key: string;
+            feedType?: undefined;
+        })[];
+    };
+    canvas: HTMLCanvasElement;
+    canvasMask: HTMLCanvasElement;
+    private sessions;
+    private config;
+    private points;
+    private labels;
+    private worldPoints;
+    private loadingAI;
+    protected viewport: any;
+    protected excludeTool: string;
+    protected currentImage: any;
+    private listeners;
+    protected desiredImage: {
+        imageId: any;
+        sampleImageId: any;
+        imageIndex: number;
+        decoder: any;
+        encoder: any;
+    };
+    protected imageEncodings: Map<any, any>;
+    protected sharedImageEncoding: any;
+    protected boxRadius: number;
+    protected imageImageData: any;
+    protected isGpuInUse: boolean;
+    protected annotationsNeedUpdating: boolean;
+    protected maskImageData: any;
+    protected promptAnnotationTypes: string[];
+    protected islandFillOptions: {
+        maxInternalRemove: number;
+        fillInternalEdge: boolean;
+    };
+    protected pCutoff: number;
+    constructor(options?: {
+        listeners: any;
+        getPromptAnnotations: any;
+        promptAnnotationTypes: any;
+        models: any;
+        modelName: any;
+        islandFillOptions: any;
+    });
+    initModel(): Promise<unknown>;
+    setPCutoff(cutoff: number): void;
+    initViewport(viewport: any): void;
+    acceptPreview(element: any): void;
+    rejectPreview(element: any): void;
+    interpolateScroll(viewport?: any, dir?: number): Promise<void>;
+    protected log(logger: Loggers, ...args: any[]): void;
+    protected getPromptAnnotations: (viewport?: any) => cornerstoneTools.Types.Annotations;
+    protected viewportRenderedListener: (_event: any) => void;
+    protected annotationModifiedListener: (_event?: any) => void;
+    disconnectViewport(viewport: any): void;
+    protected load(): Promise<void>;
+    clear(viewport: any): void;
+    cacheImageEncodings(current?: any, offset?: number, length?: number): any;
+    protected handleImage({ imageId, sampleImageId }: {
+        imageId: any;
+        sampleImageId: any;
+    }, imageSession: any): Promise<void>;
+    protected runDecode(): Promise<void>;
+    tryLoad(options?: {
+        resetImage: boolean;
+    }): void;
+    mapAnnotationPoint(worldPoint: any): number[];
+    updateAnnotations(): void;
+    restoreImageEncoding(session: any, imageId: any): Promise<any>;
+    loadStorageImageEncoding(session: any, imageId: any, index?: any): Promise<any>;
+    storeImageEncoding(session: any, imageId: any, data: any): Promise<void>;
+    createLabelmap(mask: any, canvasPosition: any, _points: any, _labels: any): void;
+    decode(points: any, labels: any, useSession?: any): Promise<void>;
+    fetchAndCacheModel(url: any, name: any): Promise<ArrayBuffer>;
+    loadModels(models: any, imageSession?: any): Promise<void>;
+    getDirectoryForImageId(session: any, imageId: any): Promise<any>;
+    getFileNameForImageId(imageId: any, extension: any): any;
+    getConfig(modelName?: string): any;
+}

package/dist/esm/ONNXSegmentationController.js

@@ -0,0 +1,836 @@
+import { utilities, eventTarget, Enums } from '@cornerstonejs/core';
+import * as cornerstoneTools from '@cornerstonejs/tools';
+import { segmentation as cstSegmentation, LabelmapBaseTool, } from '@cornerstonejs/tools';
+const { strategies } = cstSegmentation;
+const { fillInsideCircle } = strategies;
+import ort from 'onnxruntime-web/webgpu';
+import { vec3 } from 'gl-matrix';
+const { annotation } = cornerstoneTools;
+const { state: annotationState } = annotation;
+const { Events } = Enums;
+const { Events: toolsEvents } = cornerstoneTools.Enums;
+const { segmentation } = cornerstoneTools;
+const { filterAnnotationsForDisplay } = cornerstoneTools.utilities.planar;
+const { IslandRemoval } = cornerstoneTools.utilities;
+const { triggerSegmentationDataModified } = segmentation.triggerSegmentationEvents;
+function cloneTensor(t) {
+    return new ort.Tensor(t.type, Float32Array.from(t.data), t.dims);
+}
+function feedForSam(emb, points, labels, modelSize = [1024, 1024]) {
+    const maskInput = new ort.Tensor(new Float32Array(256 * 256), [1, 1, 256, 256]);
+    const hasMask = new ort.Tensor(new Float32Array([0]), [1]);
+    const originalImageSize = new ort.Tensor(new Float32Array(modelSize), [2]);
+    const pointCoords = new ort.Tensor(new Float32Array(points), [
+        1,
+        points.length / 2,
+        2,
+    ]);
+    const pointLabels = new ort.Tensor(new Float32Array(labels), [
+        1,
+        labels.length,
+    ]);
+    const key = (emb.image_embeddings && 'image_embeddings') || 'embeddings';
+    return {
+        image_embeddings: cloneTensor(emb[key]),
+        point_coords: pointCoords,
+        point_labels: pointLabels,
+        mask_input: maskInput,
+        has_mask_input: hasMask,
+        orig_im_size: originalImageSize,
+    };
+}
+function getBuffer(fileData) {
+    return new Promise((resolve) => {
+        const reader = new FileReader();
+        reader.readAsArrayBuffer(fileData);
+        reader.onload = function () {
+            const arrayBuffer = reader.result;
+            const bytes = new Float32Array(arrayBuffer);
+            resolve(bytes);
+        };
+    });
+}
+export var Loggers;
+(function (Loggers) {
+    Loggers["Log"] = "status";
+    Loggers["Encoder"] = "encoder";
+    Loggers["Decoder"] = "decoder";
+})(Loggers || (Loggers = {}));
+export default class ONNXSegmentationController {
+    static { this.MarkerInclude = 'MarkerInclude'; }
+    static { this.MarkerExclude = 'MarkerExclude'; }
+    static { this.BoxPrompt = 'BoxPrompt'; }
+    static { this.viewportOptions = {
+        displayArea: {
+            storeAsInitialCamera: true,
+            interpolationType: Enums.InterpolationType.NEAREST,
+            imageArea: [1, 1],
+            imageCanvasPoint: {
+                imagePoint: [0.5, 0.5],
+                canvasPoint: [0.5, 0.5],
+            },
+        },
+        background: [0, 0, 0.2],
+    }; }
+    static { this.MODELS = {
+        sam_l: [
+            {
+                name: 'sam-l-encoder',
+                url: '/sam_l/vit_l_encoder.onnx',
+                size: 1224,
+                key: 'encoder',
+                feedType: 'images',
+            },
+            {
+                name: 'sam-l-decoder',
+                url: '/sam_l/vit_l_decoder.onnx',
+                size: 17,
+                key: 'decoder',
+            },
+        ],
+        sam_h: [
+            {
+                name: 'sam-h-encoder',
+                url: '/sam_h/vit_h_encoder.onnx',
+                size: 18,
+                key: 'encoder',
+                feedType: 'images',
+            },
+            {
+                name: 'sam-h-decoder',
+                url: '/sam_h/vit_h_decoder.onnx',
+                size: 1,
+                key: 'decoder',
+            },
+        ],
+    }; }
+    constructor(options = {
+        listeners: null,
+        getPromptAnnotations: null,
+        promptAnnotationTypes: null,
+        models: null,
+        modelName: null,
+        islandFillOptions: undefined,
+    }) {
+        this.maxWidth = 1024;
+        this.maxHeight = 1024;
+        this.modelWidth = 1024;
+        this.modelHeight = 1024;
+        this.canvas = document.createElement('canvas');
+        this.canvasMask = document.createElement('canvas');
+        this.sessions = [];
+        this.points = [];
+        this.labels = [];
+        this.worldPoints = new Array();
+        this.excludeTool = ONNXSegmentationController.MarkerExclude;
+        this.listeners = [console.log];
+        this.desiredImage = {
+            imageId: null,
+            sampleImageId: null,
+            imageIndex: -1,
+            decoder: null,
+            encoder: null,
+        };
+        this.imageEncodings = new Map();
+        this.boxRadius = 5;
+        this.isGpuInUse = false;
+        this.annotationsNeedUpdating = false;
+        this.promptAnnotationTypes = [
+            ONNXSegmentationController.MarkerInclude,
+            ONNXSegmentationController.MarkerExclude,
+            ONNXSegmentationController.BoxPrompt,
+        ];
+        this.islandFillOptions = {
+            maxInternalRemove: 16,
+            fillInternalEdge: true,
+        };
+        this.pCutoff = 64;
+        this.getPromptAnnotations = (viewport = this.viewport) => {
+            const annotations = [];
+            const { element } = viewport;
+            for (const annotationName of this.promptAnnotationTypes) {
+                annotations.push(...annotationState.getAnnotations(annotationName, element));
+            }
+            const currentAnnotations = filterAnnotationsForDisplay(this.viewport, annotations);
+            return currentAnnotations;
+        };
+        this.viewportRenderedListener = (_event) => {
+            const { viewport, currentImage, desiredImage } = this;
+            desiredImage.imageId =
+                viewport.getCurrentImageId() || viewport.getViewReferenceId();
+            desiredImage.imageIndex = viewport.getCurrentImageIdIndex();
+            if (!desiredImage.imageId) {
+                return;
+            }
+            if (desiredImage.imageId.startsWith('volumeId:')) {
+                desiredImage.sampleImageId = viewport.getImageIds(viewport.getVolumeId())[0];
+            }
+            else {
+                desiredImage.sampleImageId = desiredImage.imageId;
+            }
+            if (desiredImage.imageId === currentImage?.imageId) {
+                return;
+            }
+            const { canvasMask } = this;
+            const ctxMask = canvasMask.getContext('2d');
+            ctxMask.clearRect(0, 0, canvasMask.width, canvasMask.height);
+            this.tryLoad({ resetImage: true });
+        };
+        this.annotationModifiedListener = (_event) => {
+            const currentAnnotations = this.getPromptAnnotations();
+            if (!currentAnnotations.length) {
+                return;
+            }
+            this.annotationsNeedUpdating = true;
+            this.tryLoad();
+        };
+        if (options.listeners) {
+            this.listeners = [...options.listeners];
+        }
+        if (options.getPromptAnnotations) {
+            this.getPromptAnnotations = options.getPromptAnnotations;
+        }
+        this.promptAnnotationTypes =
+            options.promptAnnotationTypes || this.promptAnnotationTypes;
+        if (options.models) {
+            Object.assign(ONNXSegmentationController.MODELS, options.models);
+        }
+        this.config = this.getConfig(options.modelName);
+        this.islandFillOptions =
+            options.islandFillOptions ?? this.islandFillOptions;
+    }
+    initModel() {
+        if (!this.loadingAI) {
+            this.loadingAI = this.load();
+        }
+        return this.loadingAI;
+    }
+    setPCutoff(cutoff) {
+        this.pCutoff = cutoff;
+        this.annotationsNeedUpdating = true;
+        this.tryLoad();
+    }
+    initViewport(viewport) {
+        const { desiredImage } = this;
+        if (this.viewport) {
+            this.disconnectViewport(this.viewport);
+        }
+        this.currentImage = null;
+        this.viewport = viewport;
+        const brushInstance = new LabelmapBaseTool({}, {
+            configuration: {
+                strategies: {
+                    FILL_INSIDE_CIRCLE: fillInsideCircle,
+                },
+                activeStrategy: 'FILL_INSIDE_CIRCLE',
+                preview: {
+                    enabled: true,
+                    previewColors: {
+                        0: [255, 255, 255, 128],
+                        1: [0, 255, 255, 192],
+                        2: [255, 0, 255, 255],
+                    },
+                },
+            },
+        });
+        this.tool = brushInstance;
+        desiredImage.imageId =
+            viewport.getCurrentImageId?.() || viewport.getViewReferenceId();
+        if (desiredImage.imageId.startsWith('volumeId:')) {
+            desiredImage.sampleImageId = viewport.getImageIds(viewport.getVolumeId())[0];
+        }
+        else {
+            desiredImage.sampleImageId = desiredImage.imageId;
+        }
+        viewport.element.addEventListener(Events.IMAGE_RENDERED, this.viewportRenderedListener);
+        const boundListener = this.annotationModifiedListener;
+        eventTarget.addEventListener(toolsEvents.ANNOTATION_ADDED, boundListener);
+        eventTarget.addEventListener(toolsEvents.ANNOTATION_MODIFIED, boundListener);
+        eventTarget.addEventListener(toolsEvents.ANNOTATION_COMPLETED, boundListener);
+        if (desiredImage.imageId) {
+            this.tryLoad();
+        }
+    }
+    acceptPreview(element) {
+        this.tool.acceptPreview(element);
+    }
+    rejectPreview(element) {
+        this.tool.rejectPreview(element);
+    }
+    async interpolateScroll(viewport = this.viewport, dir = 1) {
+        const { element } = viewport;
+        this.tool.acceptPreview(element);
+        const promptAnnotations = this.getPromptAnnotations(viewport);
+        if (!promptAnnotations.length) {
+            return;
+        }
+        const currentSliceIndex = viewport.getCurrentImageIdIndex();
+        const { focalPoint } = viewport.getCamera();
+        const viewRef = viewport.getViewReference({
+            sliceIndex: currentSliceIndex + dir,
+        });
+        if (!viewRef || viewRef.sliceIndex === currentSliceIndex) {
+            console.warn('No next image in direction', dir, currentSliceIndex);
+            return;
+        }
+        viewport.scroll(dir);
+        await new Promise((resolve) => window.setTimeout(resolve, 250));
+        const nextAnnotations = this.getPromptAnnotations(viewport);
+        if (nextAnnotations.length > 0) {
+            return;
+        }
+        const { focalPoint: newFocal } = viewport.getCamera();
+        const newDelta = vec3.sub(vec3.create(), newFocal, focalPoint);
+        for (const annotation of promptAnnotations) {
+            annotation.interpolationUID ||= crypto.randomUUID();
+            const newAnnotation = structuredClone(annotation);
+            newAnnotation.annotationUID = undefined;
+            Object.assign(newAnnotation.metadata, viewRef);
+            newAnnotation.cachedStats = {};
+            for (const handle of newAnnotation.data.handles.points) {
+                vec3.add(handle, handle, newDelta);
+            }
+            annotationState.addAnnotation(newAnnotation, viewport.element);
+        }
+        viewport.render();
+    }
+    log(logger, ...args) {
+        for (const listener of this.listeners) {
+            listener(logger, ...args);
+        }
+    }
+    disconnectViewport(viewport) {
+        viewport.element.removeEventListener(Events.IMAGE_RENDERED, this.viewportRenderedListener);
+        const boundListener = this.annotationModifiedListener;
+        eventTarget.removeEventListener(toolsEvents.ANNOTATION_MODIFIED, boundListener);
+        eventTarget.removeEventListener(toolsEvents.ANNOTATION_COMPLETED, boundListener);
+    }
+    async load() {
+        const { sessions } = this;
+        this.canvas.style.cursor = 'wait';
+        let loader;
+        for (let i = 0; i < 2; i++) {
+            sessions.push({
+                sessionIndex: i,
+                encoder: null,
+                decoder: null,
+                imageEmbeddings: null,
+                isLoading: false,
+                canvas: i === 0 ? this.canvas : document.createElement('canvas'),
+            });
+            if (i === 0) {
+                loader = this.loadModels(ONNXSegmentationController.MODELS[this.config.model], sessions[i]).catch((e) => {
+                    this.log(Loggers.Log, "Couldn't load models", e);
+                });
+                await loader;
+            }
+            else {
+                sessions[i].encoder = sessions[0].encoder;
+            }
+            sessions[i].loader = loader;
+        }
+    }
+    clear(viewport) {
+        this.points = [];
+        this.labels = [];
+        this.getPromptAnnotations(viewport).forEach((annotation) => annotationState.removeAnnotation(annotation.annotationUID));
+        this.tool.rejectPreview(this.viewport.element);
+    }
+    async cacheImageEncodings(current = this.viewport.getCurrentImageIdIndex(), offset = 0, length = 1000_000) {
+        const { viewport, imageEncodings } = this;
+        if (offset >= length) {
+            return;
+        }
+        const index = (offset + current) % length;
+        const view = viewport.getViewReference({ sliceIndex: index });
+        if (!view) {
+            length = index;
+            return this.cacheImageEncodings(current, offset, length);
+        }
+        const imageId = view.referencedImageId ||
+            viewport.getViewReferenceId({ sliceIndex: index });
+        if (!imageEncodings.has(imageId)) {
+            await this.loadStorageImageEncoding(current, imageId, index);
+        }
+        if (imageEncodings.has(imageId)) {
+            this.cacheImageEncodings(current, offset + 1, length);
+            return;
+        }
+        this.tryLoad();
+        if (this.isGpuInUse) {
+            setTimeout(() => this.cacheImageEncodings(current, offset), 500);
+            return;
+        }
+        this.log(Loggers.Log, 'Caching', index, imageId);
+        const sampleImageId = viewport.getImageIds()[0];
+        this.handleImage({ imageId, sampleImageId }, this.sessions[1]).then(() => {
+            this.cacheImageEncodings(current, offset + 1, length);
+        });
+    }
+    async handleImage({ imageId, sampleImageId }, imageSession) {
+        if (imageId === imageSession.imageId || this.isGpuInUse) {
+            return;
+        }
+        const { viewport, desiredImage } = this;
+        this.isGpuInUse = true;
+        imageSession.imageId = imageId;
+        imageSession.sampleImageId = sampleImageId;
+        try {
+            const isCurrent = desiredImage.imageId === imageId;
+            const { canvas } = imageSession;
+            if (isCurrent) {
+                this.log(Loggers.Encoder, `Loading image on ${imageSession.sessionIndex}`);
+                this.log(Loggers.Decoder, 'Awaiting image');
+                canvas.style.cursor = 'wait';
+            }
+            this.points = [];
+            this.labels = [];
+            const width = this.maxWidth;
+            const height = this.maxHeight;
+            canvas.width = width;
+            canvas.height = height;
+            imageSession.imageEmbeddings = undefined;
+            const size = canvas.style.width;
+            const ctx = canvas.getContext('2d', { willReadFrequently: true });
+            ctx.clearRect(0, 0, width, height);
+            const renderArguments = {
+                canvas,
+                imageId,
+                viewportOptions: {
+                    ...viewport.defaultOptions,
+                    ...ONNXSegmentationController.viewportOptions,
+                },
+                viewReference: null,
+                renderingEngineId: viewport.getRenderingEngine().id,
+            };
+            if (imageId.startsWith('volumeId:')) {
+                const viewRef = viewport.getViewReference();
+                renderArguments.viewReference = viewRef;
+                renderArguments.imageId = null;
+            }
+            imageSession.canvasPosition = await utilities.loadImageToCanvas(renderArguments);
+            canvas.style.width = size;
+            canvas.style.height = size;
+            if (isCurrent) {
+                this.log(Loggers.Encoder, `Rendered image on ${imageSession.sessionIndex}`);
+            }
+            this.imageImageData = ctx.getImageData(0, 0, width, height);
+            const data = await this.restoreImageEncoding(imageSession, imageId);
+            if (data) {
+                imageSession.imageEmbeddings = data;
+                if (desiredImage.imageId === imageId) {
+                    this.log(Loggers.Encoder, 'Cached Image');
+                    canvas.style.cursor = 'default';
+                }
+            }
+            else {
+                const t = await ort.Tensor.fromImage(this.imageImageData, {
+                    resizedWidth: this.modelWidth,
+                    resizedHeight: this.modelHeight,
+                });
+                const { feedType = 'input_image' } = this.config.encoder;
+                const feed = (feedType === 'images' && { images: t }) ||
+                    (feedType === 'pixelValues' && { pixel_values: t }) || {
+                    input_image: t,
+                };
+                await imageSession.loader;
+                const session = await imageSession.encoder;
+                if (!session) {
+                    this.log(Loggers.Log, '****** No session');
+                    return;
+                }
+                const start = performance.now();
+                imageSession.imageEmbeddings = session.run(feed);
+                const data = await imageSession.imageEmbeddings;
+                this.storeImageEncoding(imageSession, imageId, data);
+                if (desiredImage.imageId === imageId) {
+                    this.log(Loggers.Encoder, `Image Ready ${imageSession.sessionIndex} ${(performance.now() - start).toFixed(1)} ms`);
+                    canvas.style.cursor = 'default';
+                }
+            }
+        }
+        finally {
+            this.isGpuInUse = false;
+        }
+        this.tryLoad();
+    }
+    async runDecode() {
+        const { canvas } = this;
+        if (this.isGpuInUse || !this.currentImage?.imageEmbeddings) {
+            return;
+        }
+        this.isGpuInUse = true;
+        try {
+            this.canvas.style.cursor = 'wait';
+            await this.decode(this.points, this.labels);
+        }
+        finally {
+            canvas.style.cursor = 'default';
+            this.isGpuInUse = false;
+        }
+    }
+    tryLoad(options = { resetImage: false }) {
+        const { viewport, desiredImage } = this;
+        if (!desiredImage.imageId || options.resetImage) {
+            desiredImage.imageId =
+                viewport.getCurrentImageId() || viewport.getViewReferenceId();
+            this.currentImage = null;
+        }
+        const [session] = this.sessions;
+        if (session.imageId === desiredImage.imageId) {
+            if (this.currentImage !== session) {
+                this.currentImage = session;
+            }
+            this.updateAnnotations();
+            return;
+        }
+        this.handleImage(desiredImage, session);
+    }
+    mapAnnotationPoint(worldPoint) {
+        const { viewport } = this;
+        const canvasPoint = viewport.worldToCanvas(worldPoint);
+        const { width, height } = viewport.canvas;
+        const { width: destWidth, height: destHeight } = this.canvas;
+        const x = Math.trunc((canvasPoint[0] * destWidth * devicePixelRatio) / width);
+        const y = Math.trunc((canvasPoint[1] * destHeight * devicePixelRatio) / height);
+        return [x, y];
+    }
+    updateAnnotations() {
+        if (this.isGpuInUse ||
+            !this.annotationsNeedUpdating ||
+            !this.currentImage) {
+            return;
+        }
+        const promptAnnotations = this.getPromptAnnotations();
+        this.annotationsNeedUpdating = false;
+        this.points = [];
+        this.labels = [];
+        this.worldPoints = [];
+        if (!promptAnnotations?.length) {
+            return;
+        }
+        for (const annotation of promptAnnotations) {
+            const handle = annotation.data.handles.points[0];
+            const point = this.mapAnnotationPoint(handle);
+            this.points.push(...point);
+            if (annotation.metadata.toolName === ONNXSegmentationController.BoxPrompt) {
+                this.labels.push(2, 3);
+                this.points.push(...this.mapAnnotationPoint(annotation.data.handles.points[3]));
+            }
+            else {
+                const label = annotation.metadata.toolName === this.excludeTool ? 0 : 1;
+                if (label) {
+                    this.worldPoints.push(handle);
+                }
+                this.labels.push(label);
+            }
+        }
+        this.runDecode();
+    }
+    async restoreImageEncoding(session, imageId) {
+        if (!this.sharedImageEncoding) {
+            return;
+        }
+        if (!this.imageEncodings.has(imageId)) {
+            await this.loadStorageImageEncoding(session, imageId);
+        }
+        const floatData = this.imageEncodings.get(imageId);
+        if (floatData) {
+            const key = (this.sharedImageEncoding.image_embeddings && 'image_embeddings') ||
+                'embeddings';
+            this.sharedImageEncoding[key].cpuData.set(floatData);
+            return this.sharedImageEncoding;
+        }
+    }
+    async loadStorageImageEncoding(session, imageId, index = null) {
+        try {
+            const root = await this.getDirectoryForImageId(session, imageId);
+            const name = this.getFileNameForImageId(imageId, this.config.model);
+            if (!root || !name) {
+                return null;
+            }
+            const fileHandle = await findFileEntry(root, name);
+            if (!fileHandle) {
+                return null;
+            }
+            this.log(Loggers.Log, 'Loading from storage', index || imageId, name);
+            const file = await fileHandle.getFile();
+            if (file) {
+                const buffer = await getBuffer(file);
+                this.imageEncodings.set(imageId, buffer);
+            }
+        }
+        catch (e) {
+            this.log(Loggers.Log, 'Unable to fetch file', imageId, e);
+        }
+    }
+    async storeImageEncoding(session, imageId, data) {
+        if (!this.sharedImageEncoding) {
+            this.sharedImageEncoding = data;
+        }
+        const storeData = (data.image_embeddings || data.embeddings)?.cpuData;
+        if (!storeData) {
+            console.log('Unable to store data', data);
+            return;
+        }
+        const writeData = new Float32Array(storeData);
+        this.imageEncodings.set(imageId, writeData);
+        try {
+            const root = await this.getDirectoryForImageId(session, imageId);
+            const name = this.getFileNameForImageId(imageId, this.config.model);
+            if (!root || !name) {
+                return;
+            }
+            const fileHandle = await root.getFileHandle(name, { create: true });
+            const writable = await fileHandle.createWritable();
+            await writable.write(writeData);
+            await writable.close();
+        }
+        catch (e) {
+            this.log(Loggers.Log, 'Unable to write', imageId, e);
+        }
+    }
+    createLabelmap(mask, canvasPosition, _points, _labels) {
+        const { canvas, viewport } = this;
+        const preview = this.tool.addPreview(viewport.element);
+        const { previewSegmentIndex, memo, segmentationId, segmentIndex } = preview;
+        const previewVoxelManager = memo?.voxelManager || preview.previewVoxelManager;
+        const { dimensions } = previewVoxelManager;
+        const { data } = mask;
+        const { origin, topRight, bottomLeft } = canvasPosition;
+        const downVec = vec3.subtract(vec3.create(), bottomLeft, origin);
+        const rightVec = vec3.subtract(vec3.create(), topRight, origin);
+        vec3.scale(downVec, downVec, 1 / canvas.height);
+        vec3.scale(rightVec, rightVec, 1 / canvas.width);
+        const worldPointJ = vec3.create();
+        const worldPoint = vec3.create();
+        const imageData = viewport.getDefaultImageData();
+        for (let j = 0; j < canvas.height; j++) {
+            vec3.scaleAndAdd(worldPointJ, origin, downVec, j);
+            for (let i = 0; i < canvas.width; i++) {
+                vec3.scaleAndAdd(worldPoint, worldPointJ, rightVec, i);
+                const ijkPoint = imageData.worldToIndex(worldPoint).map(Math.round);
+                if (ijkPoint.findIndex((v, index) => v < 0 || v >= dimensions[index]) !==
+                    -1) {
+                    continue;
+                }
+                const maskIndex = 4 * (i + j * this.maxWidth);
+                const v = data[maskIndex];
+                if (v > this.pCutoff) {
+                    previewVoxelManager.setAtIJKPoint(ijkPoint, previewSegmentIndex);
+                }
+                else {
+                    previewVoxelManager.setAtIJKPoint(ijkPoint, null);
+                }
+            }
+        }
+        const voxelManager = previewVoxelManager.sourceVoxelManager || previewVoxelManager;
+        if (this.islandFillOptions) {
+            const islandRemoval = new IslandRemoval(this.islandFillOptions);
+            if (islandRemoval.initialize(viewport, voxelManager, {
+                previewSegmentIndex,
+                segmentIndex,
+                points: this.worldPoints.map((point) => imageData.worldToIndex(point).map(Math.round)),
+            })) {
+                islandRemoval.floodFillSegmentIsland();
+                islandRemoval.removeExternalIslands();
+                islandRemoval.removeInternalIslands();
+            }
+        }
+        triggerSegmentationDataModified(segmentationId);
+    }
+    async decode(points, labels, useSession = this.currentImage) {
+        const { canvas, canvasMask, imageImageData, desiredImage, boxRadius } = this;
+        const ctx = canvas.getContext('2d', { willReadFrequently: true });
+        ctx.clearRect(0, 0, canvas.width, canvas.height);
+        canvas.width = imageImageData.width;
+        canvas.height = imageImageData.height;
+        canvasMask.width = imageImageData.width;
+        canvasMask.height = imageImageData.height;
+        if (!useSession || useSession.imageId !== desiredImage.imageId) {
+            this.log(Loggers.Log, '***** Image not current, need to wait for current image');
+            return;
+        }
+        ctx.putImageData(imageImageData, 0, 0);
+        if (points.length) {
+            if (!useSession.imageEmbeddings) {
+                await useSession.encoder;
+            }
+            const emb = await useSession.imageEmbeddings;
+            const session = useSession.decoder;
+            const feed = feedForSam(emb, points, labels);
+            const res = await session.run(feed);
+            for (let i = 0; i < points.length; i += 2) {
+                const label = labels[i / 2];
+                ctx.fillStyle = label ? 'blue' : 'pink';
+                ctx.fillRect(points[i] - boxRadius, points[i + 1] - boxRadius, 2 * boxRadius, 2 * boxRadius);
+            }
+            const mask = res.masks;
+            this.maskImageData = mask.toImageData();
+            this.createLabelmap(this.maskImageData, useSession.canvasPosition, points, labels);
+            ctx.globalAlpha = 0.3;
+            const { data } = this.maskImageData;
+            const counts = [];
+            for (let i = 0; i < data.length; i += 4) {
+                const v = data[i];
+                if (v > 0) {
+                    if (v < 255) {
+                        data[i] = 0;
+                        if (v > 192) {
+                            data[i + 1] = 255;
+                        }
+                        else {
+                            data[i + 2] = v + 64;
+                        }
+                    }
+                    counts[v] = 1 + (counts[v] || 0);
+                }
+            }
+            const bitmap = await createImageBitmap(this.maskImageData);
+            ctx.drawImage(bitmap, 0, 0);
+            const ctxMask = canvasMask.getContext('2d');
+            ctxMask.globalAlpha = 0.9;
+            ctxMask.drawImage(bitmap, 0, 0);
+        }
+    }
+    async fetchAndCacheModel(url, name) {
+        try {
+            const cache = await caches.open('onnx');
+            let cachedResponse = await cache.match(url);
+            if (cachedResponse == undefined) {
+                await cache.add(url);
+                cachedResponse = await cache.match(url);
+                this.log(Loggers.Log, `${name} (network)`);
+            }
+            else {
+                this.log(Loggers.Log, `${name} (cached)`);
+            }
+            const data = await cachedResponse.arrayBuffer();
+            return data;
+        }
+        catch (error) {
+            this.log(Loggers.Log, `${name} (network)`);
+            return await fetch(url).then((response) => response.arrayBuffer());
+        }
+    }
+    async loadModels(models, imageSession = this.currentImage) {
+        const cache = await caches.open('onnx');
+        let missing = 0;
+        const urls = [];
+        for (const model of Object.values(models)) {
+            const cachedResponse = await cache.match(model.url);
+            if (cachedResponse === undefined) {
+                missing += model.size;
+            }
+            urls.push(model.url);
+        }
+        if (missing > 0) {
+            this.log(Loggers.Log, `downloading ${missing} MB from network ... it might take a while`);
+        }
+        else {
+            this.log(Loggers.Log, 'loading...');
+        }
+        const start = performance.now();
+        for (const model of Object.values(models)) {
+            const opt = {
+                executionProviders: [this.config.provider],
+                enableMemPattern: false,
+                enableCpuMemArena: false,
+                extra: {
+                    session: {
+                        disable_prepacking: '1',
+                        use_device_allocator_for_initializers: '1',
+                        use_ort_model_bytes_directly: '1',
+                        use_ort_model_bytes_for_initializers: '1',
+                    },
+                },
+                interOpNumThreads: 4,
+                intraOpNumThreads: 2,
+            };
+            const model_bytes = await this.fetchAndCacheModel(model.url, model.name);
+            const extra_opt = model.opt || {};
+            const sessionOptions = { ...opt, ...extra_opt };
+            this.config[model.key] = model;
+            imageSession[model.key] = await ort.InferenceSession.create(model_bytes, sessionOptions);
+        }
+        const stop = performance.now();
+        this.log(Loggers.Log, `ready, ${(stop - start).toFixed(1)}ms`, urls.join(', '));
+    }
+    async getDirectoryForImageId(session, imageId) {
+        if (imageId.indexOf('/studies/') === -1 ||
+            imageId.indexOf('/instances/') === -1) {
+            imageId = session.sampleImageId;
+            if (!imageId ||
+                imageId.indexOf('/studies/') === -1 ||
+                imageId.indexOf('/instances/') === -1) {
+                return null;
+            }
+        }
+        const studySeriesUids = imageId
+            .split('/studies/')[1]
+            .split('/instances/')[0]
+            .split('/');
+        const [studyUID, _series, seriesUID] = studySeriesUids;
+        const root = await window.navigator.storage.getDirectory();
+        const modelRoot = await getOrCreateDir(root, this.config.model);
+        const studyRoot = await getOrCreateDir(modelRoot, studyUID);
+        const seriesRoot = await getOrCreateDir(studyRoot, seriesUID);
+        return seriesRoot;
+    }
+    getFileNameForImageId(imageId, extension) {
+        if (imageId.startsWith('volumeId:')) {
+            const sliceIndex = imageId.indexOf('sliceIndex=');
+            const focalPoint = imageId.indexOf('&focalPoint=');
+            const name = imageId
+                .substring(sliceIndex, focalPoint)
+                .replace('&', '.')
+                .replace('sliceIndex=', 'volume.');
+            return name + extension;
+        }
+        const instancesLocation = imageId.indexOf('/instances/');
+        if (instancesLocation != -1) {
+            const sopLocation = instancesLocation + 11;
+            const nextSlash = imageId.indexOf('/', sopLocation);
+            return imageId.substring(sopLocation, nextSlash) + extension;
+        }
+    }
+    getConfig(modelName = 'sam_b') {
+        if (this.config) {
+            return this.config;
+        }
+        const query = window.location.search.substring(1);
+        const config = {
+            model: modelName,
+            provider: 'webgpu',
+            device: 'gpu',
+            threads: 4,
+            local: null,
+            isSlimSam: false,
+        };
+        const vars = query.split('&');
+        for (let i = 0; i < vars.length; i++) {
+            const pair = vars[i].split('=');
+            if (pair[0] in config) {
+                config[pair[0]] = decodeURIComponent(pair[1]);
+            }
+        }
+        config.threads = parseInt(String(config.threads));
+        config.local = parseInt(config.local);
+        ort.env.wasm.wasmPaths = 'dist/';
+        ort.env.wasm.numThreads = config.threads;
+        ort.env.wasm.proxy = config.provider == 'wasm';
+        this.config = config;
+        return config;
+    }
+}
+async function getOrCreateDir(dir, name) {
+    return ((await findFileEntry(dir, name)) ||
+        dir.getDirectoryHandle(name, { create: true }));
+}
+async function findFileEntry(dir, name) {
+    for await (const [key, value] of dir) {
+        if (key === name) {
+            return value;
+        }
+    }
+}

package/dist/esm/index.d.ts

@@ -0,0 +1,2 @@
+import ONNXSegmentationController from './ONNXSegmentationController';
+export { ONNXSegmentationController };

package/dist/esm/index.js

@@ -0,0 +1,2 @@
+import ONNXSegmentationController from './ONNXSegmentationController';
+export { ONNXSegmentationController };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@cornerstonejs/ai",
-  "version": "3.0.0-beta.2",
+  "version": "3.0.0-beta.4",
   "description": "AI and ML Interfaces for Cornerstone3D",
   "files": [
     "dist"
@@ -21,7 +21,8 @@
   },
   "scripts": {
     "test": "jest --testTimeout 60000",
-    "clean": "rimraf dist",
+    "clean": "rm -rf node_modules/.cache/storybook && shx rm -rf dist",
+    "clean:deep": "yarn run clean && shx rm -rf node_modules",
     "build": "yarn run build:esm",
     "build:esm": "tsc --project ./tsconfig.json",
     "build:esm:watch": "tsc --project ./tsconfig.json --watch",
@@ -30,11 +31,14 @@
     "build:update-api": "yarn run build:esm && api-extractor run --local",
     "start": "tsc --project ./tsconfig.json --watch",
     "format": "prettier --write 'src/**/*.js' 'test/**/*.js'",
-    "lint": "eslint --fix ."
+    "lint": "eslint --fix .",
+    "format-check": "npx eslint ./src --quiet",
+    "api-check": "api-extractor --debug run ",
+    "prepublishOnly": "yarn clean && yarn build"
   },
   "repository": {
     "type": "git",
-    "url": "git+https://github.com/dcmjs-org/dcmjs.git"
+    "url": "git+https://github.com/cornerstonejs/cornerstone3D.git"
   },
   "author": "@cornerstonejs",
   "license": "MIT",
@@ -53,8 +57,8 @@
     "onnxruntime-web": "1.17.1"
   },
   "peerDependencies": {
-    "@cornerstonejs/core": "^3.0.0-beta.2",
-    "@cornerstonejs/tools": "^3.0.0-beta.2"
+    "@cornerstonejs/core": "^3.0.0-beta.4",
+    "@cornerstonejs/tools": "^3.0.0-beta.4"
   },
-  "gitHead": "0e36ef232d9195cbecdf99982dedf6b73cd9d1ee"
+  "gitHead": "93464de8c8c6f01ca4ca35140c8f910b4a226582"
 }