From b34cce3e9d9342e1f43c6185aaa932a1d0643346 Mon Sep 17 00:00:00 2001 From: Joe Bowser Date: Wed, 7 Jan 2026 20:53:44 -0800 Subject: [PATCH] Add SAM3 (Segment Anything 3) interactive segmentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrates SAM3 as a toggle-able overlay tool for interactive image segmentation with ONNX Runtime backend (CUDA/CoreML/DirectML with CPU fallback). Backend: - sam3_engine.rs: ONNX inference with embedding caching - sam3_image_processor.rs: Image preprocessing and mask postprocessing - Tauri commands for load/unload/encode/segment operations - Model download from onnx-community/sam3-tracker-ONNX Frontend: - Sam3Overlay: Interactive mask overlay with point/box prompts - Sam3Controls: Mask quality selector, opacity, export, crop-to-mask - Toolbar toggle button with download indicator - Download dialog with progress for SAM3 model Usage: Click sparkle button to enable, left-click for foreground points, right-click for background, drag to draw bounding box. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- src-tauri/Cargo.lock | 515 +------------------------ src-tauri/Cargo.toml | 1 + src-tauri/src/lib.rs | 162 +++++++- src-tauri/src/model_manager.rs | 92 +++++ src-tauri/src/sam3_engine.rs | 391 +++++++++++++++++++ src-tauri/src/sam3_image_processor.rs | 198 ++++++++++ src-tauri/tests/sam3_test.rs | 206 ++++++++++ src/App.tsx | 449 ++++++++++++++++++++- src/components/DownloadModelDialog.tsx | 25 +- src/components/ImageViewer.tsx | 97 ++++- src/components/Sam3Controls.css | 217 +++++++++++ src/components/Sam3Controls.tsx | 140 +++++++ src/components/Sam3Overlay.css | 64 +++ src/components/Sam3Overlay.tsx | 227 +++++++++++ src/components/Toolbar.css | 41 ++ src/components/Toolbar.tsx | 45 ++- src/types/index.ts | 3 + src/types/sam3.ts | 94 +++++ 18 files changed, 2447 insertions(+), 520 deletions(-) create mode 100644 src-tauri/src/sam3_engine.rs create mode 100644 src-tauri/src/sam3_image_processor.rs create mode 100644 src-tauri/tests/sam3_test.rs create mode 100644 src/components/Sam3Controls.css create mode 100644 src/components/Sam3Controls.tsx create mode 100644 src/components/Sam3Overlay.css create mode 100644 src/components/Sam3Overlay.tsx create mode 100644 src/types/sam3.ts diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index e43e308..8b09958 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -31,24 +31,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "aligned" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "377e4c0ba83e4431b10df45c1d4666f178ea9c552cac93e60c3a88bf32785923" -dependencies = [ - "as-slice", -] - -[[package]] -name = "aligned-vec" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" -dependencies = [ - "equator", -] - [[package]] name = "alloc-no-stdlib" version = "2.0.4" @@ -79,38 +61,12 @@ version = "1.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" -[[package]] -name = "arbitrary" -version = "1.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" - -[[package]] -name = "arg_enum_proc_macro" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.110", -] - [[package]] name = "arrayvec" version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" -[[package]] -name = "as-slice" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516" -dependencies = [ - "stable_deref_trait", -] - [[package]] name = "ashpd" version = "0.11.0" @@ -298,49 +254,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "av-scenechange" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f321d77c20e19b92c39e7471cf986812cbb46659d2af674adc4331ef3f18394" -dependencies = [ - "aligned", - "anyhow", - "arg_enum_proc_macro", - "arrayvec", - "log", - "num-rational", - "num-traits", - "pastey", - "rayon", - "thiserror 2.0.17", - "v_frame", - "y4m", -] - -[[package]] -name = "av1-grain" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cfddb07216410377231960af4fcab838eaa12e013417781b78bd95ee22077f8" -dependencies = [ - "anyhow", - "arrayvec", - "log", - "nom 8.0.0", - "num-rational", - "v_frame", -] - -[[package]] -name = "avif-serialize" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47c8fbc0f831f4519fe8b810b6a7a91410ec83031b8233f730a0480029f6a23f" -dependencies = [ - "arrayvec", -] - [[package]] name = "base64" version = "0.13.1" @@ -370,10 +283,10 @@ name = "baseweightcanvas" version = "0.1.0" dependencies = [ "anyhow", + "base64 0.22.1", "directories", "futures-util", - "image 0.24.9", - "image 0.25.9", + "image", "libloading 0.8.9", "ndarray", "once_cell", @@ -413,15 +326,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "bitstream-io" -version = "4.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60d4bd9d1db2c6bdf285e223a7fa369d5ce98ec767dec949c6ca62863ce61757" -dependencies = [ - "core2", -] - [[package]] name = "block-buffer" version = "0.10.4" @@ -483,12 +387,6 @@ dependencies = [ "alloc-stdlib", ] -[[package]] -name = "built" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4ad8f11f288f48ca24471bbd51ac257aaeaaa07adae295591266b792902ae64" - [[package]] name = "bumpalo" version = "3.19.0" @@ -507,12 +405,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" -[[package]] -name = "byteorder-lite" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495" - [[package]] name = "bytes" version = "1.11.0" @@ -605,8 +497,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97463e1064cb1b1c1384ad0a0b9c8abd0988e2a91f52606c80ef14aadb63e36" dependencies = [ "find-msvc-tools", - "jobserver", - "libc", "shlex", ] @@ -780,15 +670,6 @@ dependencies = [ "libc", ] -[[package]] -name = "core2" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b49ba7ef1ad6107f8824dbe97de947cbaac53c44e7f9756a1fba0d37c1eec505" -dependencies = [ - "memchr", -] - [[package]] name = "cpufeatures" version = "0.2.17" @@ -1260,26 +1141,6 @@ dependencies = [ "syn 2.0.110", ] -[[package]] -name = "equator" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" -dependencies = [ - "equator-macro", -] - -[[package]] -name = "equator-macro" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.110", -] - [[package]] name = "equivalent" version = "1.0.2" @@ -1364,26 +1225,6 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" -[[package]] -name = "fax" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f05de7d48f37cd6730705cbca900770cab77a89f413d23e100ad7fad7795a0ab" -dependencies = [ - "fax_derive", -] - -[[package]] -name = "fax_derive" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0aca10fb742cb43f9e7bb8467c91aa9bcb8e3ffbc6a6f7389bb93ffc920577d" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.110", -] - [[package]] name = "fdeflate" version = "0.3.7" @@ -1745,16 +1586,6 @@ dependencies = [ "weezl", ] -[[package]] -name = "gif" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f954a9e9159ec994f73a30a12b96a702dde78f5547bcb561174597924f7d4162" -dependencies = [ - "color_quant", - "weezl", -] - [[package]] name = "gio" version = "0.18.4" @@ -2138,7 +1969,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cc50b891e4acf8fe0e71ef88ec43ad82ee07b3810ad09de10f1d01f072ed4b98" dependencies = [ "byteorder", - "png 0.17.16", + "png", ] [[package]] @@ -2259,54 +2090,14 @@ dependencies = [ "byteorder", "color_quant", "exr", - "gif 0.13.3", + "gif", "jpeg-decoder", "num-traits", - "png 0.17.16", + "png", "qoi", - "tiff 0.9.1", + "tiff", ] -[[package]] -name = "image" -version = "0.25.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6506c6c10786659413faa717ceebcb8f70731c0a60cbae39795fdf114519c1a" -dependencies = [ - "bytemuck", - "byteorder-lite", - "color_quant", - "exr", - "gif 0.14.0", - "image-webp", - "moxcms", - "num-traits", - "png 0.18.0", - "qoi", - "ravif", - "rayon", - "rgb", - "tiff 0.10.3", - "zune-core 0.5.0", - "zune-jpeg 0.5.5", -] - -[[package]] -name = "image-webp" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3" -dependencies = [ - "byteorder-lite", - "quick-error", -] - -[[package]] -name = "imgref" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7c5cedc30da3a610cac6b4ba17597bdf7152cf974e8aab3afb3d54455e371c8" - [[package]] name = "indexmap" version = "1.9.3" @@ -2352,17 +2143,6 @@ dependencies = [ "cfb", ] -[[package]] -name = "interpolate_name" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.110", -] - [[package]] name = "ipnet" version = "2.11.0" @@ -2458,16 +2238,6 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8eaf4bc02d17cbdd7ff4c7438cafcdf7fb9a4613313ad11b4f8fefe7d3fa0130" -[[package]] -name = "jobserver" -version = "0.1.34" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" -dependencies = [ - "getrandom 0.3.4", - "libc", -] - [[package]] name = "jpeg-decoder" version = "0.3.2" @@ -2574,16 +2344,6 @@ version = "0.2.177" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" -[[package]] -name = "libfuzzer-sys" -version = "0.4.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5037190e1f70cbeef565bd267599242926f724d3b8a9f510fd7e0b540cfa4404" -dependencies = [ - "arbitrary", - "cc", -] - [[package]] name = "libloading" version = "0.7.4" @@ -2642,15 +2402,6 @@ version = "0.4.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" -[[package]] -name = "loop9" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fae87c125b03c1d2c0150c90365d7d6bcc53fb73a9acaef207d2d065860f062" -dependencies = [ - "imgref", -] - [[package]] name = "mac" version = "0.1.1" @@ -2714,16 +2465,6 @@ dependencies = [ "rawpointer", ] -[[package]] -name = "maybe-rayon" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ea1f30cedd69f0a2954655f7188c6a834246d2bcf1e315e2ac40c4b24dc9519" -dependencies = [ - "cfg-if", - "rayon", -] - [[package]] name = "memchr" version = "2.7.6" @@ -2794,16 +2535,6 @@ dependencies = [ "syn 2.0.110", ] -[[package]] -name = "moxcms" -version = "0.7.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fbdd3d7436f8b5e892b8b7ea114271ff0fa00bc5acae845d53b07d498616ef6" -dependencies = [ - "num-traits", - "pxfm", -] - [[package]] name = "muda" version = "0.17.1" @@ -2819,7 +2550,7 @@ dependencies = [ "objc2-core-foundation", "objc2-foundation 0.3.2", "once_cell", - "png 0.17.16", + "png", "serde", "thiserror 2.0.17", "windows-sys 0.60.2", @@ -2922,31 +2653,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "nom" -version = "8.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" -dependencies = [ - "memchr", -] - -[[package]] -name = "noop_proc_macro" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" - -[[package]] -name = "num-bigint" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" -dependencies = [ - "num-integer", - "num-traits", -] - [[package]] name = "num-complex" version = "0.4.6" @@ -2962,17 +2668,6 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" -[[package]] -name = "num-derive" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.110", -] - [[package]] name = "num-integer" version = "0.1.46" @@ -2982,17 +2677,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "num-traits" version = "0.2.19" @@ -3473,12 +3157,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pastey" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" - [[package]] name = "pathdiff" version = "0.2.3" @@ -3689,19 +3367,6 @@ dependencies = [ "miniz_oxide", ] -[[package]] -name = "png" -version = "0.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97baced388464909d42d89643fe4361939af9b7ce7a31ee32a168f832a70f2a0" -dependencies = [ - "bitflags 2.10.0", - "crc32fast", - "fdeflate", - "flate2", - "miniz_oxide", -] - [[package]] name = "polling" version = "3.11.0" @@ -3829,34 +3494,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "profiling" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3eb8486b569e12e2c32ad3e204dbaba5e4b5b216e9367044f25f1dba42341773" -dependencies = [ - "profiling-procmacros", -] - -[[package]] -name = "profiling-procmacros" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52717f9a02b6965224f95ca2a81e2e0c5c43baacd28ca057577988930b6c3d5b" -dependencies = [ - "quote", - "syn 2.0.110", -] - -[[package]] -name = "pxfm" -version = "0.1.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3cbdf373972bf78df4d3b518d07003938e2c7d1fb5891e55f9cb6df57009d84" -dependencies = [ - "num-traits", -] - [[package]] name = "qoi" version = "0.4.1" @@ -3866,12 +3503,6 @@ dependencies = [ "bytemuck", ] -[[package]] -name = "quick-error" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" - [[package]] name = "quick-xml" version = "0.37.5" @@ -4015,56 +3646,6 @@ dependencies = [ "rand_core 0.5.1", ] -[[package]] -name = "rav1e" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43b6dd56e85d9483277cde964fd1bdb0428de4fec5ebba7540995639a21cb32b" -dependencies = [ - "aligned-vec", - "arbitrary", - "arg_enum_proc_macro", - "arrayvec", - "av-scenechange", - "av1-grain", - "bitstream-io", - "built", - "cfg-if", - "interpolate_name", - "itertools", - "libc", - "libfuzzer-sys", - "log", - "maybe-rayon", - "new_debug_unreachable", - "noop_proc_macro", - "num-derive", - "num-traits", - "paste", - "profiling", - "rand 0.9.2", - "rand_chacha 0.9.0", - "simd_helpers", - "thiserror 2.0.17", - "v_frame", - "wasm-bindgen", -] - -[[package]] -name = "ravif" -version = "0.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef69c1990ceef18a116855938e74793a5f7496ee907562bd0857b6ac734ab285" -dependencies = [ - "avif-serialize", - "imgref", - "loop9", - "quick-error", - "rav1e", - "rayon", - "rgb", -] - [[package]] name = "raw-window-handle" version = "0.6.2" @@ -4256,12 +3837,6 @@ dependencies = [ "windows-sys 0.59.0", ] -[[package]] -name = "rgb" -version = "0.8.52" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c6a884d2998352bb4daf0183589aec883f16a6da1f4dde84d8e2e9a5409a1ce" - [[package]] name = "ring" version = "0.17.14" @@ -4677,15 +4252,6 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" -[[package]] -name = "simd_helpers" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95890f873bec569a0362c235787f3aca6e1e887302ba4840839bcc6459c42da6" -dependencies = [ - "quote", -] - [[package]] name = "siphasher" version = "0.3.11" @@ -4792,7 +4358,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" dependencies = [ "base64 0.13.1", - "nom 7.1.3", + "nom", "serde", "unicode-segmentation", ] @@ -5232,7 +4798,7 @@ dependencies = [ "ico", "json-patch", "plist", - "png 0.17.16", + "png", "proc-macro2", "quote", "semver", @@ -5517,20 +5083,6 @@ dependencies = [ "weezl", ] -[[package]] -name = "tiff" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af9605de7fee8d9551863fd692cce7637f548dbd9db9180fcc07ccc6d26c336f" -dependencies = [ - "fax", - "flate2", - "half", - "quick-error", - "weezl", - "zune-jpeg 0.4.21", -] - [[package]] name = "time" version = "0.3.44" @@ -5855,7 +5407,7 @@ dependencies = [ "objc2-core-graphics", "objc2-foundation 0.3.2", "once_cell", - "png 0.17.16", + "png", "serde", "thiserror 2.0.17", "windows-sys 0.60.2", @@ -6054,17 +5606,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "v_frame" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "666b7727c8875d6ab5db9533418d7c764233ac9c0cff1d469aec8fa127597be2" -dependencies = [ - "aligned-vec", - "num-traits", - "wasm-bindgen", -] - [[package]] name = "vcpkg" version = "0.2.15" @@ -7016,12 +6557,6 @@ dependencies = [ "rustix", ] -[[package]] -name = "y4m" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" - [[package]] name = "yoke" version = "0.8.1" @@ -7187,18 +6722,6 @@ dependencies = [ "syn 2.0.110", ] -[[package]] -name = "zune-core" -version = "0.4.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f423a2c17029964870cfaabb1f13dfab7d092a62a29a89264f4d36990ca414a" - -[[package]] -name = "zune-core" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "111f7d9820f05fd715df3144e254d6fc02ee4088b0644c0ffd0efc9e6d9d2773" - [[package]] name = "zune-inflate" version = "0.2.54" @@ -7208,24 +6731,6 @@ dependencies = [ "simd-adler32", ] -[[package]] -name = "zune-jpeg" -version = "0.4.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29ce2c8a9384ad323cf564b67da86e21d3cfdff87908bc1223ed5c99bc792713" -dependencies = [ - "zune-core 0.4.12", -] - -[[package]] -name = "zune-jpeg" -version = "0.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc6fb7703e32e9a07fb3f757360338b3a567a5054f21b5f52a666752e333d58e" -dependencies = [ - "zune-core 0.5.0", -] - [[package]] name = "zvariant" version = "5.8.0" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 0fb42dd..00bcc85 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -35,6 +35,7 @@ image = "0.24" ndarray = "0.16.1" tokenizers = { version = "0.22.1" } thiserror = "1.0" +base64 = "0.22" # ONNX Runtime with platform-specific execution providers [target.'cfg(target_os = "windows")'.dependencies] diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e4f59d2..8aaf381 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -11,11 +11,14 @@ pub mod model_manager; pub mod inference_engine; mod audio_decoder; pub mod vlm_onnx; +pub mod sam3_engine; +pub mod sam3_image_processor; use model_manager::{ModelManager, DownloadProgress}; use inference_engine::{InferenceEngine, SharedInferenceEngine, create_shared_engine}; use audio_decoder::decode_audio_file; use vlm_onnx::VlmOnnx; +use sam3_engine::{Sam3Engine, Sam3SegmentResult}; // Download cancellation state pub type DownloadCancellation = Arc; @@ -23,6 +26,9 @@ pub type DownloadCancellation = Arc; // Shared ONNX engine pub type SharedOnnxEngine = Arc>>; +// Shared SAM3 engine +pub type SharedSam3Engine = Arc>>; + // Chat message for conversation history #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatMessage { @@ -400,6 +406,151 @@ async fn is_onnx_model_downloaded(model_id: String) -> Result { .map_err(|e| e.to_string()) } +// ============================================================================ +// SAM3 Commands +// ============================================================================ + +// Load SAM3 models +#[tauri::command] +async fn sam3_load( + sam3_engine: State<'_, SharedSam3Engine>, +) -> Result<(), String> { + println!("Loading SAM3 models..."); + + let manager = ModelManager::new().map_err(|e| e.to_string())?; + let (vision_path, decoder_path) = manager + .get_sam3_model_paths() + .await + .map_err(|e| e.to_string())?; + + println!("SAM3 Vision: {:?}", vision_path); + println!("SAM3 Decoder: {:?}", decoder_path); + + // Load the SAM3 model in a blocking task + let engine = tokio::task::spawn_blocking(move || { + Sam3Engine::load(&vision_path, &decoder_path) + }) + .await + .map_err(|e| format!("Task join error: {}", e))? + .map_err(|e| e.to_string())?; + + // Store the engine + let mut engine_lock = sam3_engine.lock().await; + *engine_lock = Some(engine); + + println!("SAM3 loaded successfully"); + Ok(()) +} + +// Unload SAM3 models to free GPU memory +#[tauri::command] +async fn sam3_unload( + sam3_engine: State<'_, SharedSam3Engine>, +) -> Result<(), String> { + println!("Unloading SAM3..."); + + let mut engine_lock = sam3_engine.lock().await; + *engine_lock = None; + + println!("SAM3 unloaded"); + Ok(()) +} + +// Check if SAM3 is loaded +#[tauri::command] +async fn sam3_is_loaded( + sam3_engine: State<'_, SharedSam3Engine>, +) -> Result { + let engine_lock = sam3_engine.lock().await; + Ok(engine_lock.is_some()) +} + +// Encode an image with SAM3 vision encoder +#[tauri::command] +async fn sam3_encode_image( + image_data: Vec, + sam3_engine: State<'_, SharedSam3Engine>, +) -> Result<(), String> { + println!("Encoding image with SAM3..."); + + let mut engine_lock = sam3_engine.lock().await; + let engine = engine_lock.as_mut() + .ok_or_else(|| "SAM3 not loaded".to_string())?; + + // Clone the image data for the blocking task + let image_data_clone = image_data.clone(); + + // We need to work around the borrow checker by taking ownership + let mut engine_taken = engine_lock.take() + .ok_or_else(|| "SAM3 not loaded".to_string())?; + drop(engine_lock); + + // Run encoding in a blocking task + let result = tokio::task::spawn_blocking(move || { + engine_taken.encode_image(&image_data_clone)?; + Ok::<_, anyhow::Error>(engine_taken) + }) + .await + .map_err(|e| format!("Task join error: {}", e))?; + + // Put the engine back + let engine_back = result.map_err(|e| e.to_string())?; + let mut engine_lock = sam3_engine.lock().await; + *engine_lock = Some(engine_back); + + println!("Image encoded"); + Ok(()) +} + +// Run SAM3 segmentation with prompts +#[tauri::command] +async fn sam3_segment( + points: Vec<(f32, f32)>, + labels: Vec, + bbox: Option<(f32, f32, f32, f32)>, + sam3_engine: State<'_, SharedSam3Engine>, +) -> Result { + let mut engine_lock = sam3_engine.lock().await; + let engine = engine_lock.as_mut() + .ok_or_else(|| "SAM3 not loaded".to_string())?; + + // Run segmentation (fast, can run in async context) + engine.segment(&points, &labels, bbox) + .map_err(|e| e.to_string()) +} + +// Check if SAM3 model is downloaded +#[tauri::command] +async fn sam3_is_model_downloaded() -> Result { + let manager = ModelManager::new().map_err(|e| e.to_string())?; + manager + .is_sam3_downloaded() + .await + .map_err(|e| e.to_string()) +} + +// Download SAM3 model +#[tauri::command] +async fn sam3_download_model( + app: tauri::AppHandle, + cancellation: State<'_, DownloadCancellation>, +) -> Result<(), String> { + cancellation.store(false, Ordering::SeqCst); + + let manager = ModelManager::new().map_err(|e| e.to_string())?; + let cancel_flag = cancellation.inner().clone(); + + manager + .download_sam3_model( + move |progress: DownloadProgress| { + let _ = app.emit("download-progress", &progress); + }, + cancel_flag, + ) + .await + .map_err(|e| e.to_string()) +} + #[cfg_attr(mobile, tauri::mobile_entry_point)] pub fn run() { tauri::Builder::default() @@ -407,6 +558,7 @@ pub fn run() { .plugin(tauri_plugin_dialog::init()) .manage(create_shared_engine()) .manage(Arc::new(tokio::sync::Mutex::new(None::))) // ONNX engine + .manage(Arc::new(tokio::sync::Mutex::new(None::))) // SAM3 engine .manage(Arc::new(AtomicBool::new(false))) // Download cancellation flag .invoke_handler(tauri::generate_handler![ greet, @@ -424,7 +576,15 @@ pub fn run() { download_onnx_model, load_onnx_model, generate_onnx_response, - is_onnx_model_downloaded + is_onnx_model_downloaded, + // SAM3 commands + sam3_load, + sam3_unload, + sam3_is_loaded, + sam3_encode_image, + sam3_segment, + sam3_is_model_downloaded, + sam3_download_model ]) .setup(|app| { // Create menu diff --git a/src-tauri/src/model_manager.rs b/src-tauri/src/model_manager.rs index 748e1f0..1532bf6 100644 --- a/src-tauri/src/model_manager.rs +++ b/src-tauri/src/model_manager.rs @@ -620,6 +620,98 @@ impl ModelManager { Err(_) => Ok(false), } } + + // ============================================================================ + // SAM3 Model Management + // ============================================================================ + + /// SAM3 model directory name + const SAM3_MODEL_ID: &'static str = "sam3-tracker"; + + /// Get SAM3 model paths (vision encoder and decoder) + pub async fn get_sam3_model_paths(&self) -> Result<(PathBuf, PathBuf)> { + let sam3_dir = self.models_dir.join(Self::SAM3_MODEL_ID); + + if !sam3_dir.exists() { + return Err(anyhow!("SAM3 model not found. Please download it first.")); + } + + // Look for ONNX files + let mut vision_encoder: Option = None; + let mut decoder: Option = None; + + let mut entries = tokio::fs::read_dir(&sam3_dir).await?; + while let Some(entry) = entries.next_entry().await? { + if let Some(filename) = entry.file_name().to_str() { + let path = entry.path(); + let filename_lower = filename.to_lowercase(); + + if filename_lower.contains("vision_encoder") && filename_lower.ends_with(".onnx") { + vision_encoder = Some(path); + } else if (filename_lower.contains("decoder") || filename_lower.contains("prompt_encoder")) + && filename_lower.ends_with(".onnx") + { + decoder = Some(path); + } + } + } + + match (vision_encoder, decoder) { + (Some(vision), Some(dec)) => Ok((vision, dec)), + _ => Err(anyhow!("Missing SAM3 model files (need vision_encoder*.onnx and *decoder*.onnx)")), + } + } + + /// Check if SAM3 model is downloaded + pub async fn is_sam3_downloaded(&self) -> Result { + match self.get_sam3_model_paths().await { + Ok(_) => Ok(true), + Err(_) => Ok(false), + } + } + + /// Download SAM3 model from HuggingFace + pub async fn download_sam3_model( + &self, + progress_callback: F, + cancel_flag: Arc, + ) -> Result<()> + where + F: Fn(DownloadProgress) + Send + Clone + 'static, + { + // SAM3 tracker ONNX models from HuggingFace + // Using the quantized Q4 models for efficiency + // Note: ONNX models use external data files (.onnx_data) for weights + let repo = "onnx-community/sam3-tracker-ONNX"; + let files = vec![ + ("onnx/vision_encoder_q4.onnx", "vision_encoder_q4.onnx"), + ("onnx/vision_encoder_q4.onnx_data", "vision_encoder_q4.onnx_data"), + ("onnx/prompt_encoder_mask_decoder_q4.onnx", "prompt_encoder_mask_decoder_q4.onnx"), + ("onnx/prompt_encoder_mask_decoder_q4.onnx_data", "prompt_encoder_mask_decoder_q4.onnx_data"), + ]; + + let sam3_dir = self.models_dir.join(Self::SAM3_MODEL_ID); + fs::create_dir_all(&sam3_dir).await?; + + for (remote_path, local_name) in files { + if cancel_flag.load(Ordering::SeqCst) { + return Err(anyhow!("Download cancelled")); + } + + let url = format!( + "https://huggingface.co/{}/resolve/main/{}", + repo, remote_path + ); + let local_path = sam3_dir.join(local_name); + + println!("Downloading {} to {:?}", url, local_path); + + self.download_file(&url, &local_path, local_name, progress_callback.clone(), cancel_flag.clone()) + .await?; + } + + Ok(()) + } } #[cfg(test)] diff --git a/src-tauri/src/sam3_engine.rs b/src-tauri/src/sam3_engine.rs new file mode 100644 index 0000000..4f6f7c3 --- /dev/null +++ b/src-tauri/src/sam3_engine.rs @@ -0,0 +1,391 @@ +// SAM3 ONNX Engine for Baseweight Canvas +// Segment Anything 3 - interactive image segmentation + +use anyhow::{Result, Context}; +use ndarray::{Array2, Array3, Array4, ArrayD}; +use ort::{ + session::{Session, builder::GraphOptimizationLevel}, + value::Value, +}; +use serde::Serialize; +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::collections::hash_map::DefaultHasher; +use std::path::Path; + +use crate::sam3_image_processor::Sam3ImageProcessor; + +#[cfg(target_os = "windows")] +use ort::execution_providers::DirectMLExecutionProvider; + +#[cfg(target_os = "linux")] +use ort::execution_providers::CUDAExecutionProvider; + +#[cfg(target_os = "macos")] +use ort::execution_providers::CoreMLExecutionProvider; + +/// Cached image embeddings from the vision encoder +struct CachedEmbeddings { + /// Hash of the input image data for cache invalidation + image_hash: u64, + /// Multi-scale embeddings from vision encoder + embeddings_0: ArrayD, + embeddings_1: ArrayD, + embeddings_2: ArrayD, + /// Original image dimensions + original_size: (u32, u32), +} + +/// Result from SAM3 segmentation +#[derive(Debug, Clone, Serialize)] +pub struct Sam3SegmentResult { + /// Three masks as base64-encoded PNGs (different granularities) + pub masks: Vec, + /// IoU confidence scores for each mask + pub iou_scores: Vec, + /// Object presence probability (sigmoid of object_score_logits) + pub object_score: f32, + /// Original image dimensions + pub width: u32, + pub height: u32, +} + +/// SAM3 ONNX inference engine +pub struct Sam3Engine { + vision_session: Session, + decoder_session: Session, + image_processor: Sam3ImageProcessor, + cached_embeddings: Option, +} + +impl Sam3Engine { + /// Load SAM3 models from disk + pub fn load(vision_model_path: &Path, decoder_model_path: &Path) -> Result { + println!("Loading SAM3 models..."); + println!(" Vision encoder: {:?}", vision_model_path); + println!(" Decoder: {:?}", decoder_model_path); + + // Create vision session with platform-specific execution providers + let vision_session = Self::create_session(vision_model_path, "vision encoder")?; + + // Create decoder session + let decoder_session = Self::create_session(decoder_model_path, "decoder")?; + + // Print model info + println!("Vision encoder inputs:"); + for input in &vision_session.inputs { + println!(" {}: {:?}", input.name, input.input_type); + } + println!("Vision encoder outputs:"); + for output in &vision_session.outputs { + println!(" {}: {:?}", output.name, output.output_type); + } + + println!("Decoder inputs:"); + for input in &decoder_session.inputs { + println!(" {}: {:?}", input.name, input.input_type); + } + println!("Decoder outputs:"); + for output in &decoder_session.outputs { + println!(" {}: {:?}", output.name, output.output_type); + } + + let image_processor = Sam3ImageProcessor::new(); + + Ok(Self { + vision_session, + decoder_session, + image_processor, + cached_embeddings: None, + }) + } + + /// Create an ONNX session with platform-specific execution providers + fn create_session(model_path: &Path, name: &str) -> Result { + #[cfg(target_os = "windows")] + { + // Try DirectML, fall back to CPU + let dml_result = Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_execution_providers([DirectMLExecutionProvider::default().build()])? + .commit_from_file(model_path); + + match dml_result { + Ok(session) => { + println!(" {}: using DirectML", name); + Ok(session) + } + Err(e) => { + println!(" DirectML failed for {} ({}), falling back to CPU", name, e); + Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .commit_from_file(model_path) + .context(format!("Failed to create {} session", name)) + } + } + } + + #[cfg(target_os = "linux")] + { + // Try CUDA first, fall back to CPU if it fails + let cuda_result = Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_execution_providers([CUDAExecutionProvider::default().build()])? + .commit_from_file(model_path); + + match cuda_result { + Ok(session) => { + println!(" {}: using CUDA", name); + Ok(session) + } + Err(e) => { + println!(" CUDA failed for {} ({}), falling back to CPU", name, e); + Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .commit_from_file(model_path) + .context(format!("Failed to create {} session", name)) + } + } + } + + #[cfg(target_os = "macos")] + { + // Try CoreML, fall back to CPU + let coreml_result = Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_execution_providers([CoreMLExecutionProvider::default().build()])? + .commit_from_file(model_path); + + match coreml_result { + Ok(session) => { + println!(" {}: using CoreML", name); + Ok(session) + } + Err(e) => { + println!(" CoreML failed for {} ({}), falling back to CPU", name, e); + Session::builder()? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .commit_from_file(model_path) + .context(format!("Failed to create {} session", name)) + } + } + } + } + + /// Compute a hash of image data for cache invalidation + fn compute_image_hash(image_data: &[u8]) -> u64 { + let mut hasher = DefaultHasher::new(); + image_data.hash(&mut hasher); + hasher.finish() + } + + /// Encode an image with the vision encoder + /// Results are cached for subsequent segmentation calls + pub fn encode_image(&mut self, image_data: &[u8]) -> Result<()> { + let image_hash = Self::compute_image_hash(image_data); + + // Check if we already have cached embeddings for this image + if let Some(ref cached) = self.cached_embeddings { + if cached.image_hash == image_hash { + println!("Using cached image embeddings"); + return Ok(()); + } + } + + println!("Encoding image with vision encoder..."); + + // Load and preprocess image + let image = image::load_from_memory(image_data) + .context("Failed to load image")?; + let (pixel_values, original_size, _resized_size) = self.image_processor.preprocess(image)?; + + println!(" Original size: {:?}", original_size); + println!(" Pixel values shape: {:?}", pixel_values.shape()); + + // Run vision encoder + let mut vision_inputs: HashMap<&str, Value> = HashMap::new(); + vision_inputs.insert("pixel_values", Value::from_array(pixel_values)?.into()); + + let vision_outputs = self.vision_session.run(vision_inputs)?; + + // Extract multi-scale image embeddings + let embeddings_0 = vision_outputs["image_embeddings.0"] + .try_extract_array::()? + .to_owned(); + let embeddings_1 = vision_outputs["image_embeddings.1"] + .try_extract_array::()? + .to_owned(); + let embeddings_2 = vision_outputs["image_embeddings.2"] + .try_extract_array::()? + .to_owned(); + + println!(" Embeddings shapes: {:?}, {:?}, {:?}", + embeddings_0.shape(), + embeddings_1.shape(), + embeddings_2.shape() + ); + + // Cache the embeddings + self.cached_embeddings = Some(CachedEmbeddings { + image_hash, + embeddings_0, + embeddings_1, + embeddings_2, + original_size, + }); + + println!("Image encoding complete"); + Ok(()) + } + + /// Run segmentation with the given prompts + /// Returns 3 masks with different granularities + pub fn segment( + &mut self, + points: &[(f32, f32)], + labels: &[i64], + bbox: Option<(f32, f32, f32, f32)>, + ) -> Result { + // Ensure we have cached embeddings + let cached = self.cached_embeddings.as_ref() + .ok_or_else(|| anyhow::anyhow!("No image encoded. Call encode_image first."))?; + + let original_size = cached.original_size; + + // Transform points to resized image coordinates + let transformed_points = self.image_processor.transform_points(points, original_size); + + // Prepare prompt inputs + // input_points: [batch, 1, num_points, 2] + let num_points = transformed_points.len().max(1); + let mut input_points = Array4::::zeros((1, 1, num_points, 2)); + let mut input_labels = Array3::::zeros((1, 1, num_points)); + + for (i, (x, y)) in transformed_points.iter().enumerate() { + input_points[[0, 0, i, 0]] = *x; + input_points[[0, 0, i, 1]] = *y; + } + for (i, &label) in labels.iter().enumerate() { + input_labels[[0, 0, i]] = label; + } + + // Prepare box input: [batch, num_boxes, 4] + let input_boxes = if let Some((x1, y1, x2, y2)) = bbox { + let transformed = self.image_processor.transform_boxes(&[(x1, y1, x2, y2)], original_size); + let (tx1, ty1, tx2, ty2) = transformed[0]; + let mut boxes = Array3::::zeros((1, 1, 4)); + boxes[[0, 0, 0]] = tx1; + boxes[[0, 0, 1]] = ty1; + boxes[[0, 0, 2]] = tx2; + boxes[[0, 0, 3]] = ty2; + boxes + } else { + // No box - use empty array + Array3::::zeros((1, 0, 4)) + }; + + // Run decoder + let mut decoder_inputs: HashMap<&str, Value> = HashMap::new(); + decoder_inputs.insert("input_points", Value::from_array(input_points)?.into()); + decoder_inputs.insert("input_labels", Value::from_array(input_labels)?.into()); + decoder_inputs.insert("input_boxes", Value::from_array(input_boxes)?.into()); + decoder_inputs.insert("image_embeddings.0", Value::from_array(cached.embeddings_0.clone())?.into()); + decoder_inputs.insert("image_embeddings.1", Value::from_array(cached.embeddings_1.clone())?.into()); + decoder_inputs.insert("image_embeddings.2", Value::from_array(cached.embeddings_2.clone())?.into()); + + let decoder_outputs = self.decoder_session.run(decoder_inputs)?; + + // Extract outputs + let iou_scores_raw = decoder_outputs["iou_scores"] + .try_extract_array::()? + .to_owned(); + let pred_masks = decoder_outputs["pred_masks"] + .try_extract_array::()? + .to_owned(); + let object_score_logits = decoder_outputs["object_score_logits"] + .try_extract_array::()? + .to_owned(); + + // Extract IoU scores for the 3 masks + let iou_scores = vec![ + iou_scores_raw[[0, 0, 0]], + iou_scores_raw[[0, 0, 1]], + iou_scores_raw[[0, 0, 2]], + ]; + + // Compute object presence probability + let object_logit = object_score_logits[[0, 0, 0]]; + let object_score = 1.0 / (1.0 + (-object_logit).exp()); + + // Extract and encode each mask + let shape = pred_masks.shape(); + let h = shape[3]; + let w = shape[4]; + + let mut masks = Vec::with_capacity(3); + for mask_idx in 0..3 { + // Extract mask + let mut mask = Array2::::zeros((h, w)); + for y in 0..h { + for x in 0..w { + mask[[y, x]] = pred_masks[[0, 0, mask_idx, y, x]]; + } + } + + // Postprocess (sigmoid, threshold, resize to original) + let mask_image = self.image_processor.postprocess_mask(&mask.view(), original_size); + + // Encode as base64 PNG + let base64_png = self.image_processor.mask_to_base64_png(&mask_image)?; + masks.push(base64_png); + } + + Ok(Sam3SegmentResult { + masks, + iou_scores, + object_score, + width: original_size.0, + height: original_size.1, + }) + } + + /// Clear cached embeddings to free memory + pub fn clear_cache(&mut self) { + self.cached_embeddings = None; + } + + /// Check if an image is currently cached + pub fn has_cached_embeddings(&self) -> bool { + self.cached_embeddings.is_some() + } + + /// Get a specific mask as a grayscale image for export + pub fn get_mask_for_export(&self, mask_index: usize) -> Result { + let cached = self.cached_embeddings.as_ref() + .ok_or_else(|| anyhow::anyhow!("No image encoded"))?; + + // We need to re-run segmentation to get the mask + // This is a limitation - for now, we'll return an error + // In the future, we could cache the last segmentation result + Err(anyhow::anyhow!("Export requires re-running segmentation. Use the masks from the segment() result.")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_image_hash() { + let data1 = vec![1u8, 2, 3, 4, 5]; + let data2 = vec![1u8, 2, 3, 4, 5]; + let data3 = vec![1u8, 2, 3, 4, 6]; + + let hash1 = Sam3Engine::compute_image_hash(&data1); + let hash2 = Sam3Engine::compute_image_hash(&data2); + let hash3 = Sam3Engine::compute_image_hash(&data3); + + assert_eq!(hash1, hash2); + assert_ne!(hash1, hash3); + } +} diff --git a/src-tauri/src/sam3_image_processor.rs b/src-tauri/src/sam3_image_processor.rs new file mode 100644 index 0000000..f19bd58 --- /dev/null +++ b/src-tauri/src/sam3_image_processor.rs @@ -0,0 +1,198 @@ +// SAM3 Image Processor for Baseweight Canvas +// Handles image preprocessing, coordinate transformation, and mask postprocessing + +use anyhow::Result; +use image::{DynamicImage, GrayImage, RgbaImage, imageops::FilterType}; +use ndarray::Array4; +use std::io::Cursor; +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; + +// ImageNet normalization values (used by SAM3) +const SAM3_MEAN: [f32; 3] = [123.675, 116.28, 103.53]; +const SAM3_STD: [f32; 3] = [58.395, 57.12, 57.375]; + +#[derive(Debug, Clone)] +pub struct Sam3ImageProcessor { + pub image_size: u32, + pub mean: [f32; 3], + pub std: [f32; 3], +} + +impl Default for Sam3ImageProcessor { + fn default() -> Self { + Self { + image_size: 1008, + mean: SAM3_MEAN, + std: SAM3_STD, + } + } +} + +impl Sam3ImageProcessor { + pub fn new() -> Self { + Self::default() + } + + /// Preprocess an image for SAM3 vision encoder + /// Returns: (pixel_values, original_size, resized_size) + /// pixel_values shape: [1, 3, image_size, image_size] + pub fn preprocess(&self, image: DynamicImage) -> Result<(Array4, (u32, u32), (u32, u32))> { + let original_size = (image.width(), image.height()); + + // Convert to RGB + let image = image.to_rgb8(); + + // Resize to target size (1008x1008) + let resized = image::imageops::resize( + &image, + self.image_size, + self.image_size, + FilterType::Lanczos3, + ); + let resized_size = (self.image_size, self.image_size); + + // Create array with shape [1, 3, H, W] and normalize + let mut pixel_values = Array4::::zeros((1, 3, self.image_size as usize, self.image_size as usize)); + + for y in 0..self.image_size as usize { + for x in 0..self.image_size as usize { + let pixel = resized.get_pixel(x as u32, y as u32); + for c in 0..3 { + // Normalize: (pixel - mean) / std + pixel_values[[0, c, y, x]] = (pixel[c] as f32 - self.mean[c]) / self.std[c]; + } + } + } + + Ok((pixel_values, original_size, resized_size)) + } + + /// Transform point coordinates from original image space to resized image space + /// Points are in (x, y) format + pub fn transform_points(&self, points: &[(f32, f32)], original_size: (u32, u32)) -> Vec<(f32, f32)> { + let (orig_w, orig_h) = original_size; + let scale_x = self.image_size as f32 / orig_w as f32; + let scale_y = self.image_size as f32 / orig_h as f32; + + points.iter() + .map(|(x, y)| (x * scale_x, y * scale_y)) + .collect() + } + + /// Transform box coordinates from original image space to resized image space + /// Boxes are in (x1, y1, x2, y2) format + pub fn transform_boxes(&self, boxes: &[(f32, f32, f32, f32)], original_size: (u32, u32)) -> Vec<(f32, f32, f32, f32)> { + let (orig_w, orig_h) = original_size; + let scale_x = self.image_size as f32 / orig_w as f32; + let scale_y = self.image_size as f32 / orig_h as f32; + + boxes.iter() + .map(|(x1, y1, x2, y2)| { + (x1 * scale_x, y1 * scale_y, x2 * scale_x, y2 * scale_y) + }) + .collect() + } + + /// Resize mask from model output size to original image size + /// Returns a grayscale image with binary mask (0 or 255) + pub fn postprocess_mask( + &self, + mask: &ndarray::ArrayView2, + original_size: (u32, u32), + ) -> GrayImage { + let (orig_w, orig_h) = original_size; + let (mask_h, mask_w) = (mask.shape()[0], mask.shape()[1]); + + // Create a grayscale image from the mask + let mut mask_image = GrayImage::new(mask_w as u32, mask_h as u32); + for y in 0..mask_h { + for x in 0..mask_w { + // Apply sigmoid and threshold at 0.5 + let val = mask[[y, x]]; + let sigmoid = 1.0 / (1.0 + (-val).exp()); + let pixel_val = if sigmoid > 0.5 { 255u8 } else { 0u8 }; + mask_image.put_pixel(x as u32, y as u32, image::Luma([pixel_val])); + } + } + + // Resize to original image size + image::imageops::resize( + &mask_image, + orig_w, + orig_h, + FilterType::Lanczos3, + ) + } + + /// Convert a grayscale mask to RGBA with a specified color + /// The mask alpha is based on the mask intensity + pub fn mask_to_rgba(&self, mask: &GrayImage, color: [u8; 3], opacity: u8) -> RgbaImage { + let (w, h) = mask.dimensions(); + let mut rgba = RgbaImage::new(w, h); + + for y in 0..h { + for x in 0..w { + let mask_val = mask.get_pixel(x, y).0[0]; + if mask_val > 127 { + // Foreground pixel - use the specified color with opacity + rgba.put_pixel(x, y, image::Rgba([color[0], color[1], color[2], opacity])); + } else { + // Background pixel - transparent + rgba.put_pixel(x, y, image::Rgba([0, 0, 0, 0])); + } + } + } + + rgba + } + + /// Encode a grayscale mask as base64 PNG + pub fn mask_to_base64_png(&self, mask: &GrayImage) -> Result { + let mut buffer = Cursor::new(Vec::new()); + mask.write_to(&mut buffer, image::ImageFormat::Png)?; + Ok(BASE64.encode(buffer.into_inner())) + } + + /// Encode an RGBA image as base64 PNG + pub fn rgba_to_base64_png(&self, image: &RgbaImage) -> Result { + let mut buffer = Cursor::new(Vec::new()); + image.write_to(&mut buffer, image::ImageFormat::Png)?; + Ok(BASE64.encode(buffer.into_inner())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_image_processor_creation() { + let processor = Sam3ImageProcessor::new(); + assert_eq!(processor.image_size, 1008); + } + + #[test] + fn test_point_transformation() { + let processor = Sam3ImageProcessor::new(); + let original_size = (2016, 2016); // 2x the target size + let points = vec![(100.0, 200.0)]; + let transformed = processor.transform_points(&points, original_size); + + // Points should be scaled by 0.5 + assert!((transformed[0].0 - 50.0).abs() < 0.01); + assert!((transformed[0].1 - 100.0).abs() < 0.01); + } + + #[test] + fn test_box_transformation() { + let processor = Sam3ImageProcessor::new(); + let original_size = (2016, 2016); + let boxes = vec![(100.0, 100.0, 200.0, 200.0)]; + let transformed = processor.transform_boxes(&boxes, original_size); + + assert!((transformed[0].0 - 50.0).abs() < 0.01); + assert!((transformed[0].1 - 50.0).abs() < 0.01); + assert!((transformed[0].2 - 100.0).abs() < 0.01); + assert!((transformed[0].3 - 100.0).abs() < 0.01); + } +} diff --git a/src-tauri/tests/sam3_test.rs b/src-tauri/tests/sam3_test.rs new file mode 100644 index 0000000..0a1a171 --- /dev/null +++ b/src-tauri/tests/sam3_test.rs @@ -0,0 +1,206 @@ +// SAM3 Integration Tests +// Tests SAM3 engine loading, encoding, and segmentation + +use baseweightcanvas_lib::sam3_engine::Sam3Engine; +use std::path::PathBuf; + +fn get_sam3_model_paths() -> (PathBuf, PathBuf) { + let home = std::env::var("HOME").expect("HOME not set"); + let models_dir = PathBuf::from(home).join(".local/share/ai/baseweight/Canvas/models/sam3-tracker"); + + let vision_path = models_dir.join("vision_encoder_q4.onnx"); + let decoder_path = models_dir.join("prompt_encoder_mask_decoder_q4.onnx"); + + (vision_path, decoder_path) +} + +fn get_test_image_path() -> PathBuf { + let home = std::env::var("HOME").expect("HOME not set"); + PathBuf::from(home).join("packraft.jpg") +} + +#[test] +fn test_sam3_model_loading() { + let (vision_path, decoder_path) = get_sam3_model_paths(); + + if !vision_path.exists() || !decoder_path.exists() { + eprintln!("Skipping test - SAM3 models not found at {:?}", vision_path.parent()); + return; + } + + let engine = Sam3Engine::load(&vision_path, &decoder_path); + assert!(engine.is_ok(), "Failed to load SAM3 engine: {:?}", engine.err()); + + println!("SAM3 engine loaded successfully"); +} + +#[test] +fn test_sam3_image_encoding() { + let (vision_path, decoder_path) = get_sam3_model_paths(); + let image_path = get_test_image_path(); + + if !vision_path.exists() || !decoder_path.exists() { + eprintln!("Skipping test - SAM3 models not found"); + return; + } + + if !image_path.exists() { + eprintln!("Skipping test - test image not found at {:?}", image_path); + return; + } + + let mut engine = Sam3Engine::load(&vision_path, &decoder_path) + .expect("Failed to load SAM3 engine"); + + // Load test image + let image_data = std::fs::read(&image_path).expect("Failed to read test image"); + + // Encode image + let result = engine.encode_image(&image_data); + assert!(result.is_ok(), "Failed to encode image: {:?}", result.err()); + + // Check that embeddings are cached + assert!(engine.has_cached_embeddings(), "Embeddings should be cached after encoding"); + + println!("Image encoded successfully"); +} + +#[test] +fn test_sam3_segmentation_with_point() { + let (vision_path, decoder_path) = get_sam3_model_paths(); + let image_path = get_test_image_path(); + + if !vision_path.exists() || !decoder_path.exists() { + eprintln!("Skipping test - SAM3 models not found"); + return; + } + + if !image_path.exists() { + eprintln!("Skipping test - test image not found"); + return; + } + + let mut engine = Sam3Engine::load(&vision_path, &decoder_path) + .expect("Failed to load SAM3 engine"); + + // Load and encode image + let image_data = std::fs::read(&image_path).expect("Failed to read test image"); + engine.encode_image(&image_data).expect("Failed to encode image"); + + // Run segmentation with a single point + let points = vec![(400.0, 300.0)]; + let labels = vec![1i64]; // foreground + let result = engine.segment(&points, &labels, None); + + assert!(result.is_ok(), "Failed to segment: {:?}", result.err()); + + let seg_result = result.unwrap(); + + // Verify we got 3 masks + assert_eq!(seg_result.masks.len(), 3, "Should have 3 masks"); + + // Verify we got 3 IoU scores + assert_eq!(seg_result.iou_scores.len(), 3, "Should have 3 IoU scores"); + + // Verify object score is reasonable + assert!(seg_result.object_score >= 0.0 && seg_result.object_score <= 1.0, + "Object score should be between 0 and 1"); + + // Verify masks are base64 encoded PNGs + for (i, mask) in seg_result.masks.iter().enumerate() { + assert!(!mask.is_empty(), "Mask {} should not be empty", i); + // Base64 PNG starts with specific bytes + let decoded = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, mask); + assert!(decoded.is_ok(), "Mask {} should be valid base64", i); + + let bytes = decoded.unwrap(); + // PNG magic bytes: 89 50 4E 47 + assert!(bytes.len() > 4, "Decoded mask should have PNG data"); + assert_eq!(bytes[0], 0x89, "Mask {} should start with PNG magic", i); + assert_eq!(bytes[1], 0x50, "Mask {} should be PNG", i); + } + + println!("Segmentation completed successfully"); + println!(" IoU scores: {:?}", seg_result.iou_scores); + println!(" Object score: {}", seg_result.object_score); + println!(" Dimensions: {}x{}", seg_result.width, seg_result.height); +} + +#[test] +fn test_sam3_segmentation_with_box() { + let (vision_path, decoder_path) = get_sam3_model_paths(); + let image_path = get_test_image_path(); + + if !vision_path.exists() || !decoder_path.exists() { + eprintln!("Skipping test - SAM3 models not found"); + return; + } + + if !image_path.exists() { + eprintln!("Skipping test - test image not found"); + return; + } + + let mut engine = Sam3Engine::load(&vision_path, &decoder_path) + .expect("Failed to load SAM3 engine"); + + // Load and encode image + let image_data = std::fs::read(&image_path).expect("Failed to read test image"); + engine.encode_image(&image_data).expect("Failed to encode image"); + + // Run segmentation with a bounding box + let points: Vec<(f32, f32)> = vec![]; + let labels: Vec = vec![]; + let bbox = Some((100.0, 100.0, 800.0, 800.0)); + + let result = engine.segment(&points, &labels, bbox); + + assert!(result.is_ok(), "Failed to segment with box: {:?}", result.err()); + + let seg_result = result.unwrap(); + assert_eq!(seg_result.masks.len(), 3, "Should have 3 masks"); + + println!("Box segmentation completed"); + println!(" IoU scores: {:?}", seg_result.iou_scores); +} + +#[test] +fn test_sam3_embedding_caching() { + let (vision_path, decoder_path) = get_sam3_model_paths(); + let image_path = get_test_image_path(); + + if !vision_path.exists() || !decoder_path.exists() { + eprintln!("Skipping test - SAM3 models not found"); + return; + } + + if !image_path.exists() { + eprintln!("Skipping test - test image not found"); + return; + } + + let mut engine = Sam3Engine::load(&vision_path, &decoder_path) + .expect("Failed to load SAM3 engine"); + + let image_data = std::fs::read(&image_path).expect("Failed to read test image"); + + // First encoding + let start1 = std::time::Instant::now(); + engine.encode_image(&image_data).expect("First encode failed"); + let time1 = start1.elapsed(); + + // Second encoding (should use cache) + let start2 = std::time::Instant::now(); + engine.encode_image(&image_data).expect("Second encode failed"); + let time2 = start2.elapsed(); + + // Cache hit should be much faster + assert!(time2 < time1 / 2, "Cached encoding should be faster: {:?} vs {:?}", time2, time1); + + println!("First encoding: {:?}", time1); + println!("Cached encoding: {:?}", time2); + + // Clear cache + engine.clear_cache(); + assert!(!engine.has_cached_embeddings(), "Cache should be cleared"); +} diff --git a/src/App.tsx b/src/App.tsx index cf3cfef..a131eeb 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -7,7 +7,8 @@ import { ImageViewer } from "./components/ImageViewer"; import { ChatPanel } from "./components/ChatPanel"; import { ModelSelectionModal } from "./components/ModelSelectionModal"; import { DownloadModelDialog } from "./components/DownloadModelDialog"; -import type { MediaItem, Model, AvailableModel, OnnxModel, ChatMessage } from "./types"; +import type { MediaItem, Model, AvailableModel, OnnxModel, ChatMessage, Sam3State } from "./types"; +import { initialSam3State } from "./types"; import "./App.css"; // Bundled model that ships with Baseweight Canvas @@ -137,6 +138,12 @@ function App() { const [downloadedModels, setDownloadedModels] = useState([]); const downloadCancelledRef = useRef(false); + // SAM3 state + const [sam3State, setSam3State] = useState(initialSam3State); + const [isSam3DownloadDialogOpen, setIsSam3DownloadDialogOpen] = useState(false); + const [isSam3Downloading, setIsSam3Downloading] = useState(false); + const [sam3FileProgress, setSam3FileProgress] = useState({}); + // Check if bundled model is downloaded on startup useEffect(() => { const checkBundledModel = async () => { @@ -165,6 +172,11 @@ function App() { ...prev, [progress.file]: progress })); + // Also update SAM3 progress if SAM3 download dialog is open + setSam3FileProgress(prev => ({ + ...prev, + [progress.file]: progress + })); }); return () => { @@ -251,6 +263,410 @@ function App() { loadDownloadedModels(); }, []); + // Check if SAM3 model is downloaded on startup + useEffect(() => { + const checkSam3Model = async () => { + try { + const isDownloaded = await invoke('sam3_is_model_downloaded'); + setSam3State(prev => ({ + ...prev, + isModelDownloaded: isDownloaded, + })); + } catch (error) { + console.error('Failed to check SAM3 model:', error); + } + }; + + checkSam3Model(); + }, []); + + // SAM3 handlers + const handleSam3Toggle = async () => { + if (sam3State.isLoading) return; + + if (sam3State.isEnabled) { + // Disable SAM3 and unload model + try { + await invoke('sam3_unload'); + setSam3State(prev => ({ + ...prev, + isEnabled: false, + isModelLoaded: false, + prompts: { points: [], box: null }, + segmentResult: null, + })); + } catch (error) { + console.error('Failed to unload SAM3:', error); + } + } else { + // Enable SAM3 and load model + setSam3State(prev => ({ ...prev, isLoading: true })); + try { + await invoke('sam3_load'); + setSam3State(prev => ({ + ...prev, + isEnabled: true, + isModelLoaded: true, + isLoading: false, + })); + } catch (error) { + console.error('Failed to load SAM3:', error); + setSam3State(prev => ({ ...prev, isLoading: false })); + alert(`Failed to load SAM3: ${error}`); + } + } + }; + + const handleSam3Download = () => { + // Show the download dialog instead of downloading immediately + setIsSam3DownloadDialogOpen(true); + }; + + const handleSam3DownloadConfirm = async () => { + if (sam3State.isLoading || isSam3Downloading) return; + + setIsSam3Downloading(true); + setSam3FileProgress({}); + + try { + await invoke('sam3_download_model'); + setSam3State(prev => ({ + ...prev, + isModelDownloaded: true, + })); + + setIsSam3DownloadDialogOpen(false); + setIsSam3Downloading(false); + setSam3FileProgress({}); + + // After download, automatically load the model + setSam3State(prev => ({ ...prev, isLoading: true })); + try { + await invoke('sam3_load'); + setSam3State(prev => ({ + ...prev, + isEnabled: true, + isModelLoaded: true, + isLoading: false, + })); + } catch (loadError) { + console.error('Failed to load SAM3 after download:', loadError); + setSam3State(prev => ({ ...prev, isLoading: false })); + } + } catch (error) { + console.error('Failed to download SAM3:', error); + setIsSam3Downloading(false); + setIsSam3DownloadDialogOpen(false); + alert(`Failed to download SAM3: ${error}`); + } + }; + + const handleSam3DownloadCancel = () => { + setIsSam3DownloadDialogOpen(false); + setIsSam3Downloading(false); + setSam3FileProgress({}); + }; + + // SAM3 interaction handlers + const handleSam3AddPoint = (point: { x: number; y: number; label: 0 | 1 }) => { + setSam3State(prev => ({ + ...prev, + prompts: { + ...prev.prompts, + points: [...prev.prompts.points, point], + }, + })); + }; + + const handleSam3SetBox = (box: { x1: number; y1: number; x2: number; y2: number } | null) => { + setSam3State(prev => ({ + ...prev, + prompts: { + ...prev.prompts, + box, + }, + })); + }; + + const handleSam3ClearPrompts = () => { + setSam3State(prev => ({ + ...prev, + prompts: { points: [], box: null }, + segmentResult: null, + })); + }; + + const handleSam3MaskIndexChange = (index: number) => { + setSam3State(prev => ({ + ...prev, + selectedMaskIndex: index, + })); + }; + + const handleSam3OpacityChange = (opacity: number) => { + setSam3State(prev => ({ + ...prev, + overlayOpacity: opacity, + })); + }; + + const handleSam3ExportMask = async () => { + if (!sam3State.segmentResult) return; + + try { + // Get the current mask as base64 + const maskBase64 = sam3State.segmentResult.masks[sam3State.selectedMaskIndex]; + + // Convert base64 to blob + const byteCharacters = atob(maskBase64); + const byteNumbers = new Array(byteCharacters.length); + for (let i = 0; i < byteCharacters.length; i++) { + byteNumbers[i] = byteCharacters.charCodeAt(i); + } + const byteArray = new Uint8Array(byteNumbers); + const blob = new Blob([byteArray], { type: 'image/png' }); + + // Create download link + const url = URL.createObjectURL(blob); + const a = document.createElement('a'); + a.href = url; + a.download = `mask_${sam3State.selectedMaskIndex}.png`; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + URL.revokeObjectURL(url); + } catch (error) { + console.error('Failed to export mask:', error); + alert(`Failed to export mask: ${error}`); + } + }; + + const handleSam3CropToMask = async () => { + if (!sam3State.segmentResult || !currentMedia || currentMedia.type !== 'image') return; + + try { + // Get the current mask as base64 + const maskBase64 = sam3State.segmentResult.masks[sam3State.selectedMaskIndex]; + + // Load the mask image + const maskImg = new Image(); + await new Promise((resolve, reject) => { + maskImg.onload = resolve; + maskImg.onerror = reject; + maskImg.src = `data:image/png;base64,${maskBase64}`; + }); + + // Load the original image + const originalImg = new Image(); + originalImg.crossOrigin = 'anonymous'; + await new Promise((resolve, reject) => { + originalImg.onload = resolve; + originalImg.onerror = reject; + originalImg.src = currentMedia.url; + }); + + // Create canvas and apply mask + const canvas = document.createElement('canvas'); + canvas.width = originalImg.width; + canvas.height = originalImg.height; + const ctx = canvas.getContext('2d'); + + if (!ctx) { + throw new Error('Failed to get canvas context'); + } + + // Draw original image + ctx.drawImage(originalImg, 0, 0); + + // Draw mask to another canvas to get data + const maskCanvas = document.createElement('canvas'); + maskCanvas.width = maskImg.width; + maskCanvas.height = maskImg.height; + const maskCtx = maskCanvas.getContext('2d'); + if (!maskCtx) throw new Error('Failed to get mask canvas context'); + maskCtx.drawImage(maskImg, 0, 0); + + // Scale mask to original image size if needed + const scaledMaskCanvas = document.createElement('canvas'); + scaledMaskCanvas.width = originalImg.width; + scaledMaskCanvas.height = originalImg.height; + const scaledMaskCtx = scaledMaskCanvas.getContext('2d'); + if (!scaledMaskCtx) throw new Error('Failed to get scaled mask context'); + scaledMaskCtx.drawImage(maskCanvas, 0, 0, originalImg.width, originalImg.height); + + const maskData = scaledMaskCtx.getImageData(0, 0, originalImg.width, originalImg.height); + + // Apply mask as alpha channel + const imageData = ctx.getImageData(0, 0, originalImg.width, originalImg.height); + for (let i = 0; i < imageData.data.length; i += 4) { + // Use mask luminance as alpha (white = keep, black = transparent) + const maskIndex = i; + const maskValue = maskData.data[maskIndex]; // R channel of mask + imageData.data[i + 3] = maskValue; // Set alpha + } + ctx.putImageData(imageData, 0, 0); + + // Find bounding box of non-transparent pixels + let minX = originalImg.width, minY = originalImg.height, maxX = 0, maxY = 0; + for (let y = 0; y < originalImg.height; y++) { + for (let x = 0; x < originalImg.width; x++) { + const i = (y * originalImg.width + x) * 4; + if (imageData.data[i + 3] > 0) { + minX = Math.min(minX, x); + minY = Math.min(minY, y); + maxX = Math.max(maxX, x); + maxY = Math.max(maxY, y); + } + } + } + + // Crop to bounding box + const cropWidth = maxX - minX + 1; + const cropHeight = maxY - minY + 1; + const croppedCanvas = document.createElement('canvas'); + croppedCanvas.width = cropWidth; + croppedCanvas.height = cropHeight; + const croppedCtx = croppedCanvas.getContext('2d'); + if (!croppedCtx) throw new Error('Failed to get cropped canvas context'); + croppedCtx.drawImage(canvas, minX, minY, cropWidth, cropHeight, 0, 0, cropWidth, cropHeight); + + // Convert to blob URL + const croppedUrl = croppedCanvas.toDataURL('image/png'); + + // Update media item with cropped image + const newMediaItem: MediaItem = { + id: crypto.randomUUID(), + type: 'image', + url: croppedUrl, + filename: `${currentMedia.filename.replace(/\.[^/.]+$/, '')}_cropped.png`, + size: 0, + dimensions: { + width: cropWidth, + height: cropHeight, + }, + createdAt: new Date(), + }; + + setCurrentMedia(newMediaItem); + + // Clear SAM3 prompts after crop + setSam3State(prev => ({ + ...prev, + prompts: { points: [], box: null }, + segmentResult: null, + })); + + } catch (error) { + console.error('Failed to crop to mask:', error); + alert(`Failed to crop to mask: ${error}`); + } + }; + + // Encode image when SAM3 is enabled and image changes + useEffect(() => { + const encodeImage = async () => { + if (!sam3State.isEnabled || !sam3State.isModelLoaded || !currentMedia || currentMedia.type !== 'image') { + return; + } + + setSam3State(prev => ({ ...prev, isEncoding: true })); + + try { + // Get image data from the current media + const img = new Image(); + img.crossOrigin = 'anonymous'; + + await new Promise((resolve, reject) => { + img.onload = resolve; + img.onerror = reject; + img.src = currentMedia.url; + }); + + // Create canvas and get JPEG data + const canvas = document.createElement('canvas'); + canvas.width = img.width; + canvas.height = img.height; + const ctx = canvas.getContext('2d'); + + if (!ctx) { + throw new Error('Failed to get canvas context'); + } + + ctx.drawImage(img, 0, 0); + + // Get image as JPEG bytes + const blob = await new Promise((resolve) => { + canvas.toBlob((b) => resolve(b!), 'image/jpeg', 0.95); + }); + const arrayBuffer = await blob.arrayBuffer(); + const imageBytes = Array.from(new Uint8Array(arrayBuffer)); + + console.log('Encoding image for SAM3:', img.width, 'x', img.height); + await invoke('sam3_encode_image', { imageData: imageBytes }); + + setSam3State(prev => ({ ...prev, isEncoding: false })); + console.log('SAM3 image encoding complete'); + } catch (error) { + console.error('Failed to encode image for SAM3:', error); + setSam3State(prev => ({ ...prev, isEncoding: false })); + } + }; + + encodeImage(); + }, [sam3State.isEnabled, sam3State.isModelLoaded, currentMedia?.url]); + + // Run segmentation when prompts change + useEffect(() => { + const segment = async () => { + if (!sam3State.isEnabled || !sam3State.isModelLoaded || sam3State.isEncoding) { + return; + } + + // Only segment if we have prompts + if (sam3State.prompts.points.length === 0 && !sam3State.prompts.box) { + return; + } + + setSam3State(prev => ({ ...prev, isSegmenting: true })); + + try { + // Prepare points and labels + const points = sam3State.prompts.points.map(p => [p.x, p.y] as [number, number]); + const labels = sam3State.prompts.points.map(p => p.label as number); + + // Prepare box if exists + const box = sam3State.prompts.box + ? [sam3State.prompts.box.x1, sam3State.prompts.box.y1, sam3State.prompts.box.x2, sam3State.prompts.box.y2] as [number, number, number, number] + : null; + + console.log('Running SAM3 segmentation with', points.length, 'points, box:', !!box); + + const result = await invoke<{ + masks: string[]; + iou_scores: number[]; + object_score: number; + width: number; + height: number; + }>('sam3_segment', { points, labels, box }); + + setSam3State(prev => ({ + ...prev, + isSegmenting: false, + segmentResult: result, + })); + + console.log('SAM3 segmentation complete, IoU scores:', result.iou_scores); + } catch (error) { + console.error('Failed to run SAM3 segmentation:', error); + setSam3State(prev => ({ ...prev, isSegmenting: false })); + } + }; + + // Debounce segmentation to avoid rapid calls + const timeoutId = setTimeout(segment, 100); + return () => clearTimeout(timeoutId); + }, [sam3State.isEnabled, sam3State.isModelLoaded, sam3State.isEncoding, sam3State.prompts]); + // Listen for file-opened events from the File menu useEffect(() => { console.log('Setting up file-opened event listener'); @@ -943,6 +1359,25 @@ function App() { onMediaDrop={handleMediaDrop} onLoadImage={handleLoadImage} onImageCrop={handleImageCrop} + sam3Enabled={sam3State.isEnabled} + sam3Loading={sam3State.isLoading} + sam3ModelDownloaded={sam3State.isModelDownloaded} + onSam3Toggle={handleSam3Toggle} + onSam3Download={handleSam3Download} + sam3Encoding={sam3State.isEncoding} + sam3Segmenting={sam3State.isSegmenting} + sam3SegmentResult={sam3State.segmentResult} + sam3SelectedMaskIndex={sam3State.selectedMaskIndex} + sam3OverlayOpacity={sam3State.overlayOpacity} + sam3OverlayColor={sam3State.overlayColor} + sam3Prompts={sam3State.prompts} + onSam3AddPoint={handleSam3AddPoint} + onSam3SetBox={handleSam3SetBox} + onSam3ClearPrompts={handleSam3ClearPrompts} + onSam3MaskIndexChange={handleSam3MaskIndexChange} + onSam3OpacityChange={handleSam3OpacityChange} + onSam3ExportMask={handleSam3ExportMask} + onSam3CropToMask={handleSam3CropToMask} /> + ); } diff --git a/src/components/DownloadModelDialog.tsx b/src/components/DownloadModelDialog.tsx index 118b34c..7853663 100644 --- a/src/components/DownloadModelDialog.tsx +++ b/src/components/DownloadModelDialog.tsx @@ -18,6 +18,12 @@ interface DownloadModelDialogProps { onCancel: () => void; fileProgress: FileProgress; isDownloading: boolean; + // Optional customization for different model types + modelName?: string; + modelDescription?: string; + downloadSize?: string; + source?: string; + sourceRepo?: string; } export const DownloadModelDialog: React.FC = ({ @@ -26,6 +32,11 @@ export const DownloadModelDialog: React.FC = ({ onCancel, fileProgress, isDownloading, + modelName = 'SmolVLM2 2.2B Instruct', + modelDescription = 'Baseweight Canvas ships with SmolVLM2 2.2B Instruct, a compact and efficient vision-language model.', + downloadSize = '~1.3 GB', + source = 'HuggingFace', + sourceRepo = 'ggml-org', }) => { if (!isOpen) return null; @@ -46,7 +57,7 @@ export const DownloadModelDialog: React.FC = ({

- {isDownloading ? 'Downloading Model...' : 'Download SmolVLM2 2.2B Instruct'} + {isDownloading ? 'Downloading Model...' : `Download ${modelName}`}

@@ -54,25 +65,21 @@ export const DownloadModelDialog: React.FC = ({ {!isDownloading ? (

- Baseweight Canvas ships with SmolVLM2 2.2B Instruct, a compact and efficient vision-language model. + {modelDescription}

Model: - SmolVLM2 2.2B Instruct -
-
- Quantization: - Q4_K_M + {modelName}
Download Size: - ~1.3 GB + {downloadSize}
Source: - HuggingFace (ggml-org) + {source} ({sourceRepo})
diff --git a/src/components/ImageViewer.tsx b/src/components/ImageViewer.tsx index 34b2e10..e6c1276 100644 --- a/src/components/ImageViewer.tsx +++ b/src/components/ImageViewer.tsx @@ -1,9 +1,11 @@ import React from 'react'; -import type { MediaItem } from '../types'; +import type { MediaItem, Sam3Point, Sam3Box, Sam3SegmentResult, Sam3Prompts } from '../types'; import type { NormalizedSelection } from '../types/selection'; import { DropZone } from './DropZone'; import { Toolbar, type ToolType } from './Toolbar'; import { SelectionOverlay } from './SelectionOverlay'; +import { Sam3Overlay } from './Sam3Overlay'; +import { Sam3Controls } from './Sam3Controls'; import AudioViewer from './AudioViewer'; import './ImageViewer.css'; @@ -12,9 +14,54 @@ interface ImageViewerProps { onMediaDrop: (files: File[]) => void; onLoadImage?: () => void; onImageCrop?: (croppedImageUrl: string) => void; + // SAM3 toolbar props + sam3Enabled?: boolean; + sam3Loading?: boolean; + sam3ModelDownloaded?: boolean; + onSam3Toggle?: () => void; + onSam3Download?: () => void; + // SAM3 overlay props + sam3Encoding?: boolean; + sam3Segmenting?: boolean; + sam3SegmentResult?: Sam3SegmentResult | null; + sam3SelectedMaskIndex?: number; + sam3OverlayOpacity?: number; + sam3OverlayColor?: string; + sam3Prompts?: Sam3Prompts; + onSam3AddPoint?: (point: Sam3Point) => void; + onSam3SetBox?: (box: Sam3Box | null) => void; + onSam3ClearPrompts?: () => void; + onSam3MaskIndexChange?: (index: number) => void; + onSam3OpacityChange?: (opacity: number) => void; + onSam3ExportMask?: () => void; + onSam3CropToMask?: () => void; } -export const ImageViewer: React.FC = ({ mediaItem, onMediaDrop, onLoadImage, onImageCrop }) => { +export const ImageViewer: React.FC = ({ + mediaItem, + onMediaDrop, + onLoadImage, + onImageCrop, + sam3Enabled, + sam3Loading, + sam3ModelDownloaded, + onSam3Toggle, + onSam3Download, + sam3Encoding, + sam3Segmenting, + sam3SegmentResult, + sam3SelectedMaskIndex = 0, + sam3OverlayOpacity = 50, + sam3OverlayColor = '#00ff00', + sam3Prompts, + onSam3AddPoint, + onSam3SetBox, + onSam3ClearPrompts, + onSam3MaskIndexChange, + onSam3OpacityChange, + onSam3ExportMask, + onSam3CropToMask, +}) => { const [activeTool, setActiveTool] = React.useState(undefined); const [selection, setSelection] = React.useState(null); const imageRef = React.useRef(null); @@ -100,7 +147,16 @@ export const ImageViewer: React.FC = ({ mediaItem, onMediaDrop return (
{mediaItem.type === 'image' && ( - + )}
@@ -118,10 +174,45 @@ export const ImageViewer: React.FC = ({ mediaItem, onMediaDrop isActive={activeTool === 'select'} onSelectionChange={handleSelectionChange} /> + {sam3Enabled && sam3Prompts && onSam3AddPoint && onSam3SetBox && onSam3ClearPrompts && ( + + )}
) : mediaItem.type === 'audio' ? ( ) : null} + {/* SAM3 Controls Panel */} + {sam3Enabled && onSam3MaskIndexChange && onSam3OpacityChange && onSam3ClearPrompts && onSam3ExportMask && onSam3CropToMask && ( + + )}
{mediaItem.type === 'image' && (
diff --git a/src/components/Sam3Controls.css b/src/components/Sam3Controls.css new file mode 100644 index 0000000..6b39db1 --- /dev/null +++ b/src/components/Sam3Controls.css @@ -0,0 +1,217 @@ +/* SAM3 Controls Panel Styles */ + +.bw-sam3-controls { + position: absolute; + top: 60px; + right: 12px; + width: 200px; + background-color: rgba(30, 30, 30, 0.95); + border: 1px solid #4a4a4a; + border-radius: 8px; + padding: 12px; + z-index: 20; + color: #e0e0e0; + font-size: 12px; +} + +.bw-sam3-controls-header { + display: flex; + justify-content: space-between; + align-items: center; + margin-bottom: 12px; + padding-bottom: 8px; + border-bottom: 1px solid #4a4a4a; +} + +.bw-sam3-controls-title { + font-weight: 600; + font-size: 13px; +} + +.bw-sam3-status { + font-size: 11px; + color: #888; + animation: pulse 1.5s ease-in-out infinite; +} + +@keyframes pulse { + 0%, 100% { opacity: 1; } + 50% { opacity: 0.5; } +} + +.bw-sam3-section { + margin-bottom: 12px; +} + +.bw-sam3-section-label { + display: block; + font-size: 11px; + color: #888; + margin-bottom: 6px; + text-transform: uppercase; + letter-spacing: 0.5px; +} + +/* Mask selector buttons */ +.bw-sam3-mask-buttons { + display: flex; + gap: 4px; +} + +.bw-sam3-mask-btn { + flex: 1; + display: flex; + flex-direction: column; + align-items: center; + padding: 6px 4px; + background-color: #3a3a3a; + border: 1px solid #4a4a4a; + border-radius: 4px; + cursor: pointer; + transition: all 0.15s ease; + color: #e0e0e0; +} + +.bw-sam3-mask-btn:hover { + background-color: #454545; + border-color: #5a5a5a; +} + +.bw-sam3-mask-btn--active { + background-color: #0078d4; + border-color: #1084d8; +} + +.bw-sam3-mask-label { + font-size: 10px; + font-weight: 500; +} + +.bw-sam3-mask-iou { + font-size: 10px; + color: #888; + margin-top: 2px; +} + +.bw-sam3-mask-btn--active .bw-sam3-mask-iou { + color: rgba(255, 255, 255, 0.8); +} + +/* Object score bar */ +.bw-sam3-object-score { + position: relative; + height: 20px; + background-color: #3a3a3a; + border-radius: 4px; + overflow: hidden; +} + +.bw-sam3-object-score-bar { + height: 100%; + background: linear-gradient(90deg, #22c55e, #16a34a); + transition: width 0.3s ease; +} + +.bw-sam3-object-score-text { + position: absolute; + top: 50%; + left: 50%; + transform: translate(-50%, -50%); + font-size: 11px; + font-weight: 500; +} + +/* Opacity slider */ +.bw-sam3-slider-row { + display: flex; + align-items: center; + gap: 8px; +} + +.bw-sam3-slider { + flex: 1; + height: 4px; + -webkit-appearance: none; + appearance: none; + background: #3a3a3a; + border-radius: 2px; + cursor: pointer; +} + +.bw-sam3-slider::-webkit-slider-thumb { + -webkit-appearance: none; + appearance: none; + width: 14px; + height: 14px; + background: #0078d4; + border-radius: 50%; + cursor: pointer; +} + +.bw-sam3-slider::-moz-range-thumb { + width: 14px; + height: 14px; + background: #0078d4; + border-radius: 50%; + cursor: pointer; + border: none; +} + +.bw-sam3-slider-value { + font-size: 11px; + color: #888; + min-width: 32px; + text-align: right; +} + +/* Action buttons */ +.bw-sam3-actions { + display: flex; + gap: 6px; + margin-top: 12px; + padding-top: 12px; + border-top: 1px solid #4a4a4a; +} + +.bw-sam3-action-btn { + flex: 1; + padding: 6px 8px; + background-color: #3a3a3a; + border: 1px solid #4a4a4a; + border-radius: 4px; + cursor: pointer; + transition: all 0.15s ease; + color: #e0e0e0; + font-size: 11px; + font-weight: 500; +} + +.bw-sam3-action-btn:hover:not(:disabled) { + background-color: #454545; + border-color: #5a5a5a; +} + +.bw-sam3-action-btn:disabled { + opacity: 0.4; + cursor: not-allowed; +} + +.bw-sam3-action-btn--primary { + background-color: #0078d4; + border-color: #1084d8; +} + +.bw-sam3-action-btn--primary:hover:not(:disabled) { + background-color: #1084d8; + border-color: #2090dc; +} + +/* Instructions */ +.bw-sam3-instructions { + display: flex; + flex-direction: column; + gap: 2px; + margin-top: 12px; + font-size: 10px; + color: #666; +} diff --git a/src/components/Sam3Controls.tsx b/src/components/Sam3Controls.tsx new file mode 100644 index 0000000..2ebfccc --- /dev/null +++ b/src/components/Sam3Controls.tsx @@ -0,0 +1,140 @@ +import React from 'react'; +import type { Sam3SegmentResult } from '../types'; +import './Sam3Controls.css'; + +interface Sam3ControlsProps { + isEnabled: boolean; + isLoading: boolean; + isEncoding: boolean; + isSegmenting: boolean; + segmentResult: Sam3SegmentResult | null; + selectedMaskIndex: number; + overlayOpacity: number; + onMaskIndexChange: (index: number) => void; + onOpacityChange: (opacity: number) => void; + onClearPrompts: () => void; + onExportMask: () => void; + onCropToMask: () => void; +} + +export const Sam3Controls: React.FC = ({ + isEnabled, + isLoading, + isEncoding, + isSegmenting, + segmentResult, + selectedMaskIndex, + overlayOpacity, + onMaskIndexChange, + onOpacityChange, + onClearPrompts, + onExportMask, + onCropToMask, +}) => { + if (!isEnabled) { + return null; + } + + const isProcessing = isLoading || isEncoding || isSegmenting; + + return ( +
+
+ Segment Anything + {isProcessing && ( + + {isLoading ? 'Loading...' : isEncoding ? 'Encoding...' : 'Segmenting...'} + + )} +
+ + {segmentResult && ( + <> + {/* Mask selector */} +
+ Mask Quality +
+ {[0, 1, 2].map((index) => ( + + ))} +
+
+ + {/* Object score */} +
+ Object Confidence +
+
+ + {(segmentResult.object_score * 100).toFixed(0)}% + +
+
+ + )} + + {/* Opacity slider */} +
+ Overlay Opacity +
+ onOpacityChange(parseInt(e.target.value))} + className="bw-sam3-slider" + /> + {overlayOpacity}% +
+
+ + {/* Action buttons */} +
+ + + +
+ + {/* Instructions */} +
+ Left-click: add point + Right-click: exclude point + Drag: draw box +
+
+ ); +}; diff --git a/src/components/Sam3Overlay.css b/src/components/Sam3Overlay.css new file mode 100644 index 0000000..5a401fc --- /dev/null +++ b/src/components/Sam3Overlay.css @@ -0,0 +1,64 @@ +/* SAM3 Overlay Styles */ + +.bw-sam3-overlay { + position: absolute; + top: 0; + left: 0; + width: 100%; + height: 100%; + cursor: crosshair; + z-index: 10; +} + +/* Mask overlay image */ +.bw-sam3-mask { + position: absolute; + top: 0; + left: 0; + width: 100%; + height: 100%; + object-fit: contain; + pointer-events: none; +} + +/* Point markers */ +.bw-sam3-point { + position: absolute; + width: 12px; + height: 12px; + border-radius: 50%; + transform: translate(-50%, -50%); + pointer-events: none; + border: 2px solid white; + box-shadow: 0 0 4px rgba(0, 0, 0, 0.5); +} + +.bw-sam3-point--foreground { + background-color: #22c55e; /* green */ +} + +.bw-sam3-point--background { + background-color: #ef4444; /* red */ +} + +/* Bounding box */ +.bw-sam3-box { + position: absolute; + border: 2px dashed #3b82f6; + background-color: rgba(59, 130, 246, 0.1); + pointer-events: none; +} + +/* Loading indicator */ +.bw-sam3-loading { + position: absolute; + top: 50%; + left: 50%; + transform: translate(-50%, -50%); + background-color: rgba(0, 0, 0, 0.75); + color: white; + padding: 12px 24px; + border-radius: 8px; + font-size: 14px; + pointer-events: none; +} diff --git a/src/components/Sam3Overlay.tsx b/src/components/Sam3Overlay.tsx new file mode 100644 index 0000000..34889d1 --- /dev/null +++ b/src/components/Sam3Overlay.tsx @@ -0,0 +1,227 @@ +import React, { useCallback, useRef, useState } from 'react'; +import type { Sam3Point, Sam3Box, Sam3SegmentResult, Sam3Prompts } from '../types'; +import './Sam3Overlay.css'; + +interface Sam3OverlayProps { + imageElement: HTMLImageElement | null; + imageWidth: number; + imageHeight: number; + isEnabled: boolean; + isEncoding: boolean; + isSegmenting: boolean; + segmentResult: Sam3SegmentResult | null; + selectedMaskIndex: number; + overlayOpacity: number; + overlayColor: string; + prompts: Sam3Prompts; + onAddPoint: (point: Sam3Point) => void; + onSetBox: (box: Sam3Box | null) => void; + onClearPrompts: () => void; +} + +export const Sam3Overlay: React.FC = ({ + imageElement, + imageWidth, + imageHeight, + isEnabled, + isEncoding, + isSegmenting, + segmentResult, + selectedMaskIndex, + overlayOpacity, + overlayColor, + prompts, + onAddPoint, + onSetBox, +}) => { + const overlayRef = useRef(null); + const [isDragging, setIsDragging] = useState(false); + const [dragStart, setDragStart] = useState<{ x: number; y: number } | null>(null); + const [currentBox, setCurrentBox] = useState(null); + + // Convert screen coordinates to image coordinates + const screenToImageCoords = useCallback((clientX: number, clientY: number): { x: number; y: number } | null => { + if (!imageElement || !overlayRef.current) return null; + + const rect = overlayRef.current.getBoundingClientRect(); + const scaleX = imageWidth / rect.width; + const scaleY = imageHeight / rect.height; + + const x = (clientX - rect.left) * scaleX; + const y = (clientY - rect.top) * scaleY; + + // Clamp to image bounds + return { + x: Math.max(0, Math.min(imageWidth, x)), + y: Math.max(0, Math.min(imageHeight, y)), + }; + }, [imageElement, imageWidth, imageHeight]); + + // Handle mouse down - start potential drag + const handleMouseDown = useCallback((e: React.MouseEvent) => { + if (!isEnabled || isEncoding || isSegmenting) return; + + const coords = screenToImageCoords(e.clientX, e.clientY); + if (!coords) return; + + setIsDragging(true); + setDragStart(coords); + setCurrentBox(null); + }, [isEnabled, isEncoding, isSegmenting, screenToImageCoords]); + + // Handle mouse move - update drag box + const handleMouseMove = useCallback((e: React.MouseEvent) => { + if (!isDragging || !dragStart) return; + + const coords = screenToImageCoords(e.clientX, e.clientY); + if (!coords) return; + + // If we've moved more than 10px, start drawing a box + const dx = Math.abs(coords.x - dragStart.x); + const dy = Math.abs(coords.y - dragStart.y); + + if (dx > 10 || dy > 10) { + setCurrentBox({ + x1: Math.min(dragStart.x, coords.x), + y1: Math.min(dragStart.y, coords.y), + x2: Math.max(dragStart.x, coords.x), + y2: Math.max(dragStart.y, coords.y), + }); + } + }, [isDragging, dragStart, screenToImageCoords]); + + // Handle mouse up - finalize action + const handleMouseUp = useCallback((e: React.MouseEvent) => { + if (!isEnabled || !isDragging || !dragStart) { + setIsDragging(false); + setDragStart(null); + return; + } + + const coords = screenToImageCoords(e.clientX, e.clientY); + if (!coords) { + setIsDragging(false); + setDragStart(null); + setCurrentBox(null); + return; + } + + // Check if this was a click (no significant movement) or a drag + const dx = Math.abs(coords.x - dragStart.x); + const dy = Math.abs(coords.y - dragStart.y); + + if (dx < 10 && dy < 10) { + // Click - add a point + const isRightClick = e.button === 2; + onAddPoint({ + x: coords.x, + y: coords.y, + label: isRightClick ? 0 : 1, // 0 = background, 1 = foreground + }); + } else if (currentBox) { + // Drag - set bounding box + onSetBox(currentBox); + } + + setIsDragging(false); + setDragStart(null); + setCurrentBox(null); + }, [isEnabled, isDragging, dragStart, currentBox, screenToImageCoords, onAddPoint, onSetBox]); + + // Handle right-click for background points + const handleContextMenu = useCallback((e: React.MouseEvent) => { + e.preventDefault(); + if (!isEnabled || isEncoding || isSegmenting) return; + + const coords = screenToImageCoords(e.clientX, e.clientY); + if (!coords) return; + + onAddPoint({ + x: coords.x, + y: coords.y, + label: 0, // background point + }); + }, [isEnabled, isEncoding, isSegmenting, screenToImageCoords, onAddPoint]); + + // Convert image coordinates to display coordinates + const imageToDisplayCoords = useCallback((x: number, y: number): { left: string; top: string } => { + if (!overlayRef.current) return { left: '0%', top: '0%' }; + + const leftPercent = (x / imageWidth) * 100; + const topPercent = (y / imageHeight) * 100; + + return { + left: `${leftPercent}%`, + top: `${topPercent}%`, + }; + }, [imageWidth, imageHeight]); + + // Get the selected mask as a data URL + const maskDataUrl = segmentResult?.masks[selectedMaskIndex]; + + if (!isEnabled || !imageElement) { + return null; + } + + return ( +
+ {/* Mask overlay */} + {maskDataUrl && ( + Segmentation mask + )} + + {/* Point markers */} + {prompts.points.map((point, index) => { + const pos = imageToDisplayCoords(point.x, point.y); + return ( +
+ ); + })} + + {/* Bounding box */} + {(prompts.box || currentBox) && (() => { + const box = currentBox || prompts.box!; + const topLeft = imageToDisplayCoords(box.x1, box.y1); + return ( +
+ ); + })()} + + {/* Loading indicator */} + {(isEncoding || isSegmenting) && ( +
+ {isEncoding ? 'Encoding image...' : 'Segmenting...'} +
+ )} +
+ ); +}; diff --git a/src/components/Toolbar.css b/src/components/Toolbar.css index 999d872..57b7b34 100644 --- a/src/components/Toolbar.css +++ b/src/components/Toolbar.css @@ -61,3 +61,44 @@ line-height: 1; user-select: none; } + +/* Toolbar divider */ +.bw-toolbar-divider { + height: 1px; + background-color: #4a4a4a; + margin: 4px 0; +} + +/* SAM3 toggle button */ +.bw-sam3-toggle { + position: relative; +} + +.bw-sam3-toggle.bw-tool-button--loading { + cursor: wait; + animation: pulse 1.5s ease-in-out infinite; +} + +@keyframes pulse { + 0%, 100% { + opacity: 1; + } + 50% { + opacity: 0.6; + } +} + +.bw-sam3-download-badge { + position: absolute; + bottom: 2px; + right: 2px; + font-size: 10px; + line-height: 1; + background-color: #0078d4; + border-radius: 50%; + width: 14px; + height: 14px; + display: flex; + align-items: center; + justify-content: center; +} diff --git a/src/components/Toolbar.tsx b/src/components/Toolbar.tsx index 67624fb..78050fa 100644 --- a/src/components/Toolbar.tsx +++ b/src/components/Toolbar.tsx @@ -7,6 +7,12 @@ interface ToolbarProps { activeTool?: ToolType; hasSelection?: boolean; onToolSelect: (tool: ToolType) => void; + // SAM3 props + sam3Enabled?: boolean; + sam3Loading?: boolean; + sam3ModelDownloaded?: boolean; + onSam3Toggle?: () => void; + onSam3Download?: () => void; } interface ToolButtonProps { @@ -32,13 +38,30 @@ const ToolButton: React.FC = ({ label, icon, active, disabled, ); }; -export const Toolbar: React.FC = ({ activeTool, hasSelection, onToolSelect }) => { +export const Toolbar: React.FC = ({ + activeTool, + hasSelection, + onToolSelect, + sam3Enabled, + sam3Loading, + sam3ModelDownloaded, + onSam3Toggle, + onSam3Download, +}) => { const tools: Array<{ tool: ToolType; label: string; icon: string; disabled?: boolean; disabledReason?: string }> = [ { tool: 'load', label: 'Load', icon: '📁' }, { tool: 'select', label: 'Select', icon: '⬚' }, { tool: 'crop', label: 'Crop', icon: '✂', disabled: !hasSelection, disabledReason: !hasSelection ? 'Select an area first' : undefined }, ]; + const handleSam3Click = () => { + if (!sam3ModelDownloaded && onSam3Download) { + onSam3Download(); + } else if (onSam3Toggle) { + onSam3Toggle(); + } + }; + return (
@@ -54,6 +77,26 @@ export const Toolbar: React.FC = ({ activeTool, hasSelection, onTo onClick={() => onToolSelect(tool)} /> ))} + {/* SAM3 Toggle Button */} +
+
); diff --git a/src/types/index.ts b/src/types/index.ts index baff8b6..7370965 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -106,3 +106,6 @@ export interface OnnxModel { estimatedSizes: { [key: string]: number }; // Size per quantization description?: string; } + +// Re-export SAM3 types +export * from './sam3'; diff --git a/src/types/sam3.ts b/src/types/sam3.ts new file mode 100644 index 0000000..a528d85 --- /dev/null +++ b/src/types/sam3.ts @@ -0,0 +1,94 @@ +// SAM3 (Segment Anything 3) type definitions + +// Point prompt for SAM3 segmentation +export interface Sam3Point { + x: number; + y: number; + label: 0 | 1; // 0 = background (exclude), 1 = foreground (include) +} + +// Bounding box prompt for SAM3 segmentation +export interface Sam3Box { + x1: number; + y1: number; + x2: number; + y2: number; +} + +// Collection of prompts for SAM3 +export interface Sam3Prompts { + points: Sam3Point[]; + box: Sam3Box | null; +} + +// Result from SAM3 segmentation (matches Rust Sam3SegmentResult) +export interface Sam3SegmentResult { + masks: string[]; // 3 base64-encoded PNG masks + iou_scores: number[]; // IoU confidence scores for each mask + object_score: number; // Object presence probability (0-1) + width: number; + height: number; +} + +// SAM3 UI state +export interface Sam3State { + // Loading states + isEnabled: boolean; + isLoading: boolean; // Model loading + isEncoding: boolean; // Vision encoder running + isSegmenting: boolean; // Decoder running + + // Model state + isModelDownloaded: boolean; + isModelLoaded: boolean; + + // Prompts and results + prompts: Sam3Prompts; + segmentResult: Sam3SegmentResult | null; + + // UI configuration + selectedMaskIndex: number; // 0, 1, or 2 + overlayOpacity: number; // 0-100 + overlayColor: string; // Hex color for mask overlay +} + +// Initial SAM3 state +export const initialSam3State: Sam3State = { + isEnabled: false, + isLoading: false, + isEncoding: false, + isSegmenting: false, + isModelDownloaded: false, + isModelLoaded: false, + prompts: { + points: [], + box: null, + }, + segmentResult: null, + selectedMaskIndex: 0, + overlayOpacity: 50, + overlayColor: '#00ff00', // Green +}; + +// SAM3 overlay props +export interface Sam3OverlayProps { + imageElement: HTMLImageElement | null; + imageWidth: number; + imageHeight: number; + state: Sam3State; + onAddPoint: (point: Sam3Point) => void; + onSetBox: (box: Sam3Box | null) => void; + onClearPrompts: () => void; +} + +// SAM3 controls props +export interface Sam3ControlsProps { + state: Sam3State; + onToggle: () => void; + onMaskIndexChange: (index: number) => void; + onOpacityChange: (opacity: number) => void; + onClearPrompts: () => void; + onExportMask: () => void; + onCropToMask: () => void; + onDownloadModel: () => void; +}