onnx seems to work...

Makefile

@@ -23,7 +23,7 @@ BEHAVE_VERSION := $(shell node determine_version_number.mjs)
 
 .PHONY=all public/app/bundled/libavjs lint public/app/tsc libavjs test build
 
-build: public/app/tsc public/app/bundled/libavjs-$(LIBAVJS_COMMIT)/version.txt $(STATIC_TARGET_MARKDOWN_FILES) public/app/bundled/tfjs-wasm
+build: public/app/tsc public/app/bundled/libavjs-$(LIBAVJS_COMMIT)/version.txt $(STATIC_TARGET_MARKDOWN_FILES) public/app/bundled/ort-wasm
 
 all: build test
 
@@ -76,6 +76,6 @@ clean:
 	@if [ -e public ]; then rm -r public; fi
 
 
-public/app/bundled/tfjs-wasm: $(wildcard node_modules/@tensorflow/tfjs-backend-wasm/dist/*.wasm)
+public/app/bundled/ort-wasm: $(wildcard node_modules/onnxruntime-web/dist/*.wasm)
 	@mkdir -p $@
-	@cp node_modules/@tensorflow/tfjs-backend-wasm/dist/*.wasm $@
+	@cp node_modules/onnxruntime-web/dist/*.wasm $@

package.json

@@ -3,9 +3,6 @@
   "dependencies": {
     "@material-symbols/font-400": "^0.15.0",
     "@reduxjs/toolkit": "^2.0.1",
-    "@tensorflow/tfjs": "^4.12.0",
-    "@tensorflow/tfjs-backend-wasm": "^4.12.0",
-    "@tensorflow/tfjs-backend-webgpu": "^4.12.0",
     "@types/dom-screen-wake-lock": "^1.0.3",
     "@types/dom-webcodecs": "^0.1.11",
     "@types/react": "^17.0.69",
@@ -24,7 +21,9 @@
     "libavjs-webcodecs-bridge": "github:Yahweasel/libavjs-webcodecs-bridge#98fb137f2b06029e376dc252c31d2f0b540a5919",
     "markdown-it": "^14.1.0",
     "markdown-it-attrs": "^4.2.0",
+    "onnxruntime-web": "^1.19.2",
     "preact": "^10.19.2",
+    "protobufjs": "^7.4.0",
     "react": "npm:@preact/compat@^17.1.2",
     "react-dom": "npm:@preact/compat@^17.1.2",
     "react-redux": "^8.1.3",

serve.py

@@ -10,6 +10,7 @@ def end_headers(self):
     def send_my_headers(self):
         self.send_header("Cross-Origin-Opener-Policy", "same-origin")
         self.send_header("Cross-Origin-Embedder-Policy", "require-corp")
+        self.send_header("Cross-Origin-Embedder-Policy", "credentialless")
 
 
 if __name__ == '__main__':

src/infer/Inferrer.tsx

@@ -29,7 +29,6 @@ function fileFilterForInfer(file: File): boolean | string {
 export function Inferrer(): JSX.Element {
   const [files, setFiles] = useState<FileTreeBranch>(new Map())
   const [concurrency, setConcurrency] = useState(1)
-  const [modelName, setModelName] = useState<string | null>(null)
   const [state, setState] = useState<"uploading" | "selectmodel" | "converting" | "done">("uploading")
   const [yoloSettings, setYoloSettings] = useState<YoloSettings | null>(null)
   const [destination, setDestination] = useState<FileSystemDirectoryHandle>()
@@ -85,14 +84,6 @@ export function Inferrer(): JSX.Element {
     void(loadCachedSettings().then(settings => setYoloSettings(settings)))
   }, [])
 
-  useEffect(() => {
-    if (!yoloSettings) {
-      setModelName(null)
-      return;
-    }
-    void(API.checkValidModel(yoloSettings.backend, yoloSettings.modelDirectory).then(response => setModelName(response.name)))
-  }, [yoloSettings])
-
   const [wakeLock, setWakeLock] = useState<WakeLockSentinel | null>(null)
   const [preventSleep, setPreventSleep] = useState(false)
 
@@ -147,7 +138,7 @@ export function Inferrer(): JSX.Element {
           Check the <a href="../help/infer.html">help page</a> or the <a href="../help/quickstart.html">quick start guide</a> for more information.
         </div>
         {yoloSettings ? <div className={css.explanation}>
-          Loaded model: {modelName !== null ? modelName : "<loading>"} ({yoloSettings.yoloVersion} / {yoloSettings.backend}) <button disabled={state!=="uploading"}
+          Loaded model: {yoloSettings.modelFilename !== null ? yoloSettings.modelFilename : "<loading>"} ({yoloSettings.yoloVersion} / {yoloSettings.backend}) <button disabled={state!=="uploading"}
             onClick={() => setState("selectmodel")}
           >change</button>
         </div> : <div className={css.explanation}>

src/infer/YoloSettings.tsx

@@ -1,33 +1,37 @@
 import { JSX } from "preact"
 import {useState, useEffect} from 'preact/hooks'
-import {YoloBackend, YoloSettings, YOLO_MODEL_NAME_FILE} from "../lib/tfjs-shared"
+import {YoloBackend, YoloSettings, YOLO_MODEL_NAME_FILE, YOLO_MODEL_DIRECTORY} from "../lib/tfjs-shared"
 import * as infercss from "./inferrer.module.css"
 import {getEntry, cp_r} from "../lib/fileutil"
-import { Checker, LiteralChecker, TypeChecker, getCheckerFromObject } from "../lib/typeCheck"
+import { Checker, LiteralChecker, StringChecker, getCheckerFromObject } from "../lib/typeCheck"
 import { API } from "../worker/Api"
+import { valueOrErrorAsync2 } from "src/lib/util"
 
 export const YOLO_SETTINGS_STORAGE_KEY = "YoloSettingsStorageKey"
-const YOLO_MODEL_DIRECTORY = "YoloModelDir"
 
 const YoloSettingsChecker: Checker<YoloSettings> = getCheckerFromObject({
   version: new LiteralChecker(1),
   yoloVersion: new LiteralChecker(["v5", "v8"]),
   backend: new LiteralChecker(["wasm", "webgl", "webgpu"]),
-  modelDirectory: new TypeChecker(FileSystemDirectoryHandle),
+  modelFilename: new StringChecker(),
 })
 
 export async function loadCachedSettings(): Promise<YoloSettings | null> {
   try {
     const opfsRoot = await navigator.storage.getDirectory()
     const opfsModelDir = await opfsRoot.getDirectoryHandle(YOLO_MODEL_DIRECTORY)
 
-    const settings = {
-      ...JSON.parse(localStorage.getItem(YOLO_SETTINGS_STORAGE_KEY)!),
-      modelDirectory: opfsModelDir
+    const settingsJSON = localStorage.getItem(YOLO_SETTINGS_STORAGE_KEY)
+    if (settingsJSON === null) {
+      console.log("No yolo settings found")
+      return null
     }
-    YoloSettingsChecker.assertInstance(settings)
-    await API.checkValidModel(settings.backend, opfsModelDir)
-    return settings
+    const yoloSettings = {
+      ...JSON.parse(settingsJSON)
+    }
+    YoloSettingsChecker.assertInstance(yoloSettings)
+    await API.checkValidModel(yoloSettings)
+    return yoloSettings
   } catch (e) {
     console.error("Problem retrieving settings:")
     console.error(e)
@@ -48,93 +52,81 @@ export function YoloSettingsDialog({
 }: Props): JSX.Element {
   const [yoloVersion, setYoloVersion] = useState<YoloSettings["yoloVersion"]>("v8")
   const [backend, setBackend] = useState<YoloSettings["backend"]>("webgl")
-  const [modelDir, setModelDir] = useState<FileSystemDirectoryHandle>()
+  const [newModelFile, setNewModelFile] = useState<FileSystemFileHandle>()
+  const modelFileName = newModelFile ? newModelFile.name
+    : yoloSettings ? yoloSettings.modelFilename : null
 
   useEffect(() => {
     if (yoloSettings === null) {
       setYoloVersion("v8")
       setBackend("webgl")
-      setModelDir(undefined)
     } else {
-      // basically we assume that the model is still in opfs
-      // Technically someone can have removed it since the last load of the site
-      // If so, bad luck. This will generate an error, and that's it
-      // Reload and you'll start with an empty model
       void(navigator.storage.getDirectory()
         .then(opfsRoot => opfsRoot.getDirectoryHandle(YOLO_MODEL_DIRECTORY))
-        .then(opfsModelDir => {
-          setYoloVersion(yoloSettings.yoloVersion)
-          setBackend(yoloSettings.backend)
-          setModelDir(opfsModelDir)
+        .then(opfsModelDir => opfsModelDir.getFileHandle(yoloSettings.modelFilename))
+        .then(fileHandle => fileHandle.getFile())
+        .then(file => file.arrayBuffer())
+        .catch(error => {
+          console.error("Error reading the model", error)
+          setYoloSettings(null)
         }))
     }
   }, [yoloSettings])
 
   async function save() {
-    localStorage.removeItem(YOLO_SETTINGS_STORAGE_KEY)
-    const newYoloSettingsWithoutModel: Omit<YoloSettings, "modelDirectory"> = {
+    const newYoloSettings = {
       version: 1,
       yoloVersion,
       backend,
-    }
-    const opfsRoot = await navigator.storage.getDirectory()
-    const modelDirRelativeToOPFSRoot = modelDir && await opfsRoot.resolve(modelDir)
-    const modelDirNeedsCopying /*as opposed to in OPFS*/ =
-      (modelDirRelativeToOPFSRoot ?? []).join("/") !== YOLO_MODEL_DIRECTORY
-    if (!modelDir) {
-      if (await getEntry(opfsRoot, [YOLO_MODEL_DIRECTORY])) {
-        await opfsRoot.removeEntry(YOLO_MODEL_DIRECTORY, {recursive: true})
-      } 
-      localStorage.removeItem(YOLO_SETTINGS_STORAGE_KEY)
-      setYoloSettings(null)
+      modelFilename: yoloSettings ? yoloSettings.modelFilename : null
+    } as Omit<YoloSettings, "modelFilename"> & {modelFilename: string | null}
+    localStorage.removeItem(YOLO_SETTINGS_STORAGE_KEY)
+    if (newModelFile) {
+      const opfsRoot = await navigator.storage.getDirectory()
+      const opfsModelDir = await opfsRoot.getDirectoryHandle(
+        YOLO_MODEL_DIRECTORY, {create: true})
+      const newModelName = newModelFile.name
+      const file = await opfsModelDir.getFileHandle(newModelName, {create: true})
+      const stream = await file.createWritable()
+      const data = await (await newModelFile.getFile()).arrayBuffer()
+      await stream.truncate(0)
+      await stream.write(data)
+      await stream.close()
+      newYoloSettings.modelFilename = newModelName
     } else {
-      if (modelDirNeedsCopying) {
-        await API.checkValidModel(backend, modelDir)
-        if (await getEntry(opfsRoot, [YOLO_MODEL_DIRECTORY])) {
-          await opfsRoot.removeEntry(YOLO_MODEL_DIRECTORY, {recursive: true})
-        } 
-        const opfsModelDir = await opfsRoot.getDirectoryHandle(
-          YOLO_MODEL_DIRECTORY, {create: true})
-        await cp_r(modelDir, opfsModelDir)
-        if (!await getEntry(opfsModelDir, [YOLO_MODEL_NAME_FILE])) {
-          const modelNameHandle = await opfsModelDir.getFileHandle(
-            YOLO_MODEL_NAME_FILE, {create: true})
-          const os = await modelNameHandle.createWritable()
-          await os.write(modelDir.name)
-          await os.close()
-        }
+      if (newYoloSettings.modelFilename === null) {
+        throw new Error("You need a model")
       }
+    }
+    YoloSettingsChecker.assertInstance(newYoloSettings)
       localStorage.setItem(
-        YOLO_SETTINGS_STORAGE_KEY, JSON.stringify(newYoloSettingsWithoutModel))
-      const opfsModelDir = await opfsRoot.getDirectoryHandle(YOLO_MODEL_DIRECTORY)
-      await API.checkValidModel(backend, opfsModelDir)
-
-      const newYoloSettings = {
-        ...newYoloSettingsWithoutModel, modelDirectory: opfsModelDir}
+        YOLO_SETTINGS_STORAGE_KEY, JSON.stringify(newYoloSettings))
       setYoloSettings(newYoloSettings)
-    }
     closeSettingsDialog()
   }
 
-  async function unloadModelDir() {
-    setModelDir(undefined)
-  }
-
-  async function setNewModelDir() {
-    const modelDirectory = await window.showDirectoryPicker({id: "modelpicker"})
-
-    try {
-      await API.checkValidModel(backend, modelDirectory)
-    } catch (e) {
-      console.log("opening of model failed", e)
-      window.alert("Opening of model failed; either the directory pointed to does not contain a valid model, or backend is not supported")
-      return
+  async function selectNewModelFile() {
+    const result = await valueOrErrorAsync2(() => window.showOpenFilePicker({
+      id: "modelpicker",
+      types: [{description: "ONNX model", accept: {"application/onnx": [".onnx"]}}]
+    }))
+    if ("error" in result) {
+      if ((result.error as DOMException).name === "AbortError") {
+        console.log("Cancel clicked, do nothing")
+        return
+      } else {
+        throw result.error
+      }
+    }
+    const modelFiles = result.value
+    if (modelFiles.length !== 1) {
+      throw new Error("There should (per spec) always be exactly one file")
     }
-    setModelDir(modelDirectory)
+    setNewModelFile(modelFiles[0])
   }
 
   if (yoloSettings === undefined) {
-    return  <></>
+    return  <div>Loading yolo settings....</div>
   }
 
   return <>
@@ -164,20 +156,19 @@ export function YoloSettingsDialog({
       </dd>
       <dt>Model</dt>
       <dd>
-        {modelDir !== undefined
+        {modelFileName !== null
           ? <div>
-            Model loaded
-            <button onClick={setNewModelDir}>Change model</button>
-            <button onClick={unloadModelDir}>Unload model</button>
+            Model {modelFileName} loaded
+            <button onClick={selectNewModelFile}>Change model</button>
             </div>
           : <>
-            <div>No model selected, load one here (or continue without a model)</div>
-            <button onClick={setNewModelDir}>Select model</button>
+            <div>No model selected, load one here</div>
+            <button onClick={selectNewModelFile}>Select model</button>
           </>
         }
       </dd>
     </dl>
-    <button onClick={save}>Save</button>
+    <button disabled={modelFileName === null} onClick={save}>Save</button>
     <button onClick={closeSettingsDialog}>Cancel</button>
   </>
 }

src/lib/tfjs-shared.ts

@@ -1,8 +1,9 @@
 export const YOLO_MODEL_NAME_FILE = "modelname.txt"
+export const YOLO_MODEL_DIRECTORY = "YoloModelDir"
 export type YoloSettings = {
   version: 1,
   yoloVersion: YoloVersion,
-  modelDirectory: FileSystemDirectoryHandle,
+  modelFilename: string,
   backend: YoloBackend,
 }
 

src/lib/util.ts

@@ -375,3 +375,20 @@ export async function debugImage(imageData: ImageData) {
   }
 }
 
+type ArgMinMaxReturn = {maxIndex: number, maxValue: number} | undefined
+export function argMax(items: ReadonlyArray<number>): ArgMinMaxReturn {
+  return items.reduce((acc, cur, index) => (
+    acc === undefined || cur > acc.maxValue
+      ? {maxValue: cur, maxIndex: index}
+      : acc),
+    undefined as ArgMinMaxReturn
+  )
+}
+export function argMin(items: ReadonlyArray<number>): ArgMinMaxReturn {
+  return items.reduce((acc, cur, index) => (
+    acc === undefined || cur < acc.maxValue
+      ? {maxValue: cur, maxIndex: index}
+      : acc),
+    undefined as ArgMinMaxReturn
+  )
+}

src/worker/Api.ts

@@ -2,7 +2,7 @@ declare const WORKER_URL: string;
 
 import { exhausted, promiseWithResolve} from "../lib/util"
 import {FileTreeLeaf} from "../lib/FileTree"
-import { YoloBackend, YoloSettings } from "../lib/tfjs-shared"
+import { YoloSettings } from "../lib/tfjs-shared"
 import { VideoMetadata } from "../lib/video-shared"
 import { tic } from "../lib/insight";
 
@@ -36,8 +36,7 @@ export type WorkerInferMethod = {
 export type WorkerCheckValidModel = {
   call: {
     method: "check_valid_model",
-    backend: YoloBackend
-    directory: FileSystemDirectoryHandle
+    yoloSettings: YoloSettings
   }
   message: {type: "done", result: {name: string}}
   | {type: "error", error: Error}
@@ -69,7 +68,7 @@ export class API {
     onProgress: (progress: FileTreeLeaf["progress"]) => void
   ): Promise<void> {
     const {promise, resolve, reject} = promiseWithResolve<void>()
-    const worker = new Worker(WORKER_URL, {name: "convertor"}) as ConvertWorker
+    const worker = new Worker(WORKER_URL, {name: "convertor", type: "module"}) as ConvertWorker
     worker.addEventListener("message", e => {
       const data = e.data as WorkerConvertMethod["message"]
       switch (data.type) {
@@ -99,7 +98,7 @@ export class API {
     onProgress: (progress: FileTreeLeaf["progress"]) => void,
   ): Promise<void> {
     const {promise, resolve, reject} = promiseWithResolve<void>()
-    const worker = new Worker(WORKER_URL, {name: "inferrer"}) as InferWorker
+    const worker = new Worker(WORKER_URL, {name: "inferrer", type: "module"}) as InferWorker
     worker.addEventListener("message", e => {
       const data = e.data as WorkerConvertMethod["message"]
       switch (data.type) {
@@ -122,11 +121,10 @@ export class API {
   }
 
   static checkValidModel(
-    yoloBackend: YoloBackend,
-    directory: FileSystemDirectoryHandle,
+    yoloSettings: YoloSettings,
   ): Promise<{name: string}> {
     const {promise, resolve, reject} = promiseWithResolve<{name: string}>()
-    const worker = new Worker(WORKER_URL, {name: "checkValidModel"}) as ValidModelWorker
+    const worker = new Worker(WORKER_URL, {name: "checkValidModel", type: "module"}) as ValidModelWorker
     worker.addEventListener("message", e => {
       const data = e.data as WorkerCheckValidModel["message"]
       switch (data.type) {
@@ -140,15 +138,15 @@ export class API {
           exhausted(data)
       }
     })
-    worker.postMessage({method: "check_valid_model", backend: yoloBackend, directory})
+    worker.postMessage({method: "check_valid_model", yoloSettings})
     return promise
   }
 
   static extractMetadata(
     file: File,
   ): Promise<VideoMetadata> {
     const {promise, resolve, reject} = promiseWithResolve<VideoMetadata>()
-    const worker = new Worker(WORKER_URL, {name: "extractMetadata"}) as ExtractMetadataWorker
+    const worker = new Worker(WORKER_URL, {name: "extractMetadata", type: "module"}) as ExtractMetadataWorker
     worker.addEventListener("message", e => {
       const data = e.data as WorkerExtractMetadata["message"]
       switch (data.type) {