@jax-data-science/sci-image-visualizer 0.8.3 → 0.10.0

package/fesm2022/onnx-sam.worker.js

@@ -1,14 +1,41 @@
 // src/lib/toolbar/segmentation/onnx-sam.worker.ts
+import { LruCache } from "napari-js";
 import * as ort2 from "onnxruntime-web";
 
 // src/lib/toolbar/segmentation/sam-onnx-core.ts
 import * as ort from "onnxruntime-web";
 var PIXEL_MEAN = [123.675, 116.28, 103.53];
 var PIXEL_STD = [58.395, 57.12, 57.375];
-var MODEL_CACHE = "sam-onnx";
-async function fetchModel(url, onProgress) {
-  const cache = typeof caches !== "undefined" ? await caches.open(MODEL_CACHE).catch(() => null) : null;
-  const hit = cache ? await cache.match(url) : null;
+var MODEL_CACHE_PREFIX = "sam-onnx";
+var MODEL_CACHE = `${MODEL_CACHE_PREFIX}-v2`;
+var staleCachesPurged = null;
+function purgeStaleModelCaches() {
+  staleCachesPurged ?? (staleCachesPurged = caches.keys().then(
+    (names) => Promise.all(
+      names.filter((n) => n.startsWith(MODEL_CACHE_PREFIX) && n !== MODEL_CACHE).map((n) => caches.delete(n))
+    )
+  ).then(() => void 0).catch(() => void 0));
+  return staleCachesPurged;
+}
+async function openModelCache() {
+  if (typeof caches === "undefined")
+    return null;
+  await purgeStaleModelCaches();
+  return caches.open(MODEL_CACHE).catch(() => null);
+}
+function modelCacheKey(url, revision) {
+  if (!revision)
+    return url;
+  return `${url}${url.includes("?") ? "&" : "?"}siv-model-rev=${encodeURIComponent(revision)}`;
+}
+async function cacheModel(cache, url, key, body) {
+  await cache.delete(url, { ignoreSearch: true }).catch(() => void 0);
+  await cache.put(key, body).catch(() => void 0);
+}
+async function fetchModel(url, onProgress, revision) {
+  const key = modelCacheKey(url, revision);
+  const cache = await openModelCache();
+  const hit = cache ? await cache.match(key) : null;
   if (hit) {
     onProgress?.(1);
     return await hit.arrayBuffer();
@@ -21,7 +48,7 @@ async function fetchModel(url, onProgress) {
     const buf = await resp.arrayBuffer();
     onProgress?.(1);
     if (cache)
-      await cache.put(url, new Response(buf)).catch(() => void 0);
+      await cacheModel(cache, url, key, new Response(buf));
     return buf;
   }
   const reader = resp.body.getReader();
@@ -42,7 +69,7 @@ async function fetchModel(url, onProgress) {
     off += c.length;
   }
   if (cache) {
-    await cache.put(url, new Response(out.slice(), { headers: { "content-length": String(received) } })).catch(() => void 0);
+    await cacheModel(cache, url, key, new Response(out, { headers: { "content-length": String(received) } }));
   }
   onProgress?.(1);
   return out.buffer;
@@ -149,7 +176,7 @@ var encoder = null;
 var decoder = null;
 var inputSize = 1024;
 var nextToken = 1;
-var embeddings = /* @__PURE__ */ new Map();
+var embeddings = new LruCache(3);
 function post(msg, transfer = []) {
   self.postMessage(msg, transfer);
 }
@@ -158,8 +185,12 @@ async function onLoad(msg) {
   inputSize = msg.inputSize || 1024;
   const hasGpu = typeof navigator !== "undefined" && "gpu" in navigator;
   const eps = msg.encoderProviders && msg.encoderProviders.length ? msg.encoderProviders : hasGpu ? ["webgpu", "wasm"] : ["wasm"];
-  const encBuf = await fetchModel(msg.encoderUrl, (f) => post({ id: msg.id, type: "progress", fraction: f }));
-  const decBuf = await fetchModel(msg.decoderUrl);
+  const encBuf = await fetchModel(
+    msg.encoderUrl,
+    (f) => post({ id: msg.id, type: "progress", fraction: f }),
+    msg.revision
+  );
+  const decBuf = await fetchModel(msg.decoderUrl, void 0, msg.revision);
   encoder = await ort2.InferenceSession.create(encBuf, { executionProviders: eps });
   decoder = await ort2.InferenceSession.create(decBuf, { executionProviders: ["wasm"] });
   post({ id: msg.id, type: "loaded" });
@@ -171,9 +202,15 @@ async function onEmbed(msg) {
   const e = await runEncoder(encoder, data, msg.width, msg.height, inputSize);
   const token = nextToken++;
   embeddings.set(token, e);
-  if (embeddings.size > 3)
-    embeddings.delete(embeddings.keys().next().value);
-  post({ id: msg.id, type: "embedded", token, scale: e.scale, imageWidth: e.imageWidth, imageHeight: e.imageHeight, dims: e.dims });
+  post({
+    id: msg.id,
+    type: "embedded",
+    token,
+    scale: e.scale,
+    imageWidth: e.imageWidth,
+    imageHeight: e.imageHeight,
+    dims: e.dims
+  });
 }
 async function onDecode(msg) {
   if (!decoder)
@@ -182,10 +219,9 @@ async function onDecode(msg) {
   if (!e)
     throw new Error("SAM embedding expired; re-encode the image.");
   const r = await runDecoder(decoder, e, msg.prompt);
-  post(
-    { id: msg.id, type: "decoded", width: r.width, height: r.height, iou: r.iou, buffer: r.mask.buffer },
-    [r.mask.buffer]
-  );
+  post({ id: msg.id, type: "decoded", width: r.width, height: r.height, iou: r.iou, buffer: r.mask.buffer }, [
+    r.mask.buffer
+  ]);
 }
 self.onmessage = async (ev) => {
   const msg = ev.data;