first version
This commit is contained in:
7
CLAUDE.md
Normal file
7
CLAUDE.md
Normal file
@@ -0,0 +1,7 @@
|
||||
|
||||
| Principle | Addresses |
|
||||
|-----------|-----------|
|
||||
| **Think Before Coding** | Wrong assumptions, hidden confusion, missing tradeoffs |
|
||||
| **Simplicity First** | Overcomplication, bloated abstractions |
|
||||
| **Surgical Changes** | Orthogonal edits, touching code you shouldn't |
|
||||
| **Goal-Driven Execution** | Leverage through tests-first, verifiable success criteria |
|
||||
213
README.md
Normal file
213
README.md
Normal file
@@ -0,0 +1,213 @@
|
||||
# SAM-Tool
|
||||
|
||||
Desktop video-frame extraction + image annotation tool, with optional
|
||||
SAM 2 / SAM 3 segmentation backend. Built to replace an earlier
|
||||
Python/Tkinter prototype (`sam_tool_v2.py`) for YOLO-style training
|
||||
dataset prep.
|
||||
|
||||
```
|
||||
┌──────────────────┐ HTTP ┌──────────────────┐
|
||||
│ sam-tool-tauri │ ◀───────────────▶ │ sam2-backend │
|
||||
│ (Tauri/Rust + UI)│ /predict, etc │ (FastAPI + GPU) │
|
||||
└──────────────────┘ └──────────────────┘
|
||||
│ │
|
||||
▼ ▼
|
||||
local images + SAM 2.x / SAM 3
|
||||
sibling JSONs weights on disk
|
||||
```
|
||||
|
||||
The two halves can run on the same machine or separate boxes — point the
|
||||
desktop app at any reachable backend URL from **Settings**.
|
||||
|
||||
---
|
||||
|
||||
## Layout
|
||||
|
||||
```
|
||||
SAM-Tool/
|
||||
├── sam-tool-tauri/ # Desktop app (Tauri 2, React 18, TS, Vite)
|
||||
│ ├── src/ # Frontend
|
||||
│ ├── src-tauri/ # Rust commands (video, decoder, IPC)
|
||||
│ └── package.json # `npm run tauri dev` / `tauri build`
|
||||
│
|
||||
├── sam2-backend/ # FastAPI HTTP server (SAM 2 + SAM 3)
|
||||
│ ├── server.py # /, /models, /predict
|
||||
│ ├── registry.py # Lazy-load + single-resident cache
|
||||
│ ├── adapters/
|
||||
│ │ ├── sam2_adapter.py
|
||||
│ │ ├── sam3_adapter.py
|
||||
│ │ └── utils.py # mask → polygon helper
|
||||
│ ├── models.yaml # Model registry (paths + family)
|
||||
│ ├── run.sh # Convenience launcher (env vars)
|
||||
│ └── README.md # Install + run on the server side
|
||||
│
|
||||
├── sam2/sam2/ # Meta SAM 2 source + checkpoints (vendored)
|
||||
├── files/sam3_weights/ # SAM 3 weights (3.4 GB sam3.pt)
|
||||
├── files/sam3_server.py # Reference SAM3 server (used as a guide)
|
||||
│
|
||||
├── sam_tool_v2.py # Original Python/Tkinter tool (legacy)
|
||||
├── sam_annotation_test/ # Sample data — dashcam MP4 + custom SAM JSON
|
||||
│
|
||||
├── docs/ # ← you are here
|
||||
│ ├── POLYGON_SMOOTHING.md How to tune mask→polygon output
|
||||
│ ├── HF_DEPLOY_OFFLINE.md First-online-then-airgapped HF deploy
|
||||
│ └── BUILD_RELEASE.md Building the Tauri app for release
|
||||
│
|
||||
└── CLAUDE.md # Project rules (think → simplify → surgical)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Quick start (single-machine dev)
|
||||
|
||||
### 1. SAM backend
|
||||
|
||||
```bash
|
||||
cd sam2-backend
|
||||
python3 -m venv .venv && source .venv/bin/activate
|
||||
pip install --upgrade pip
|
||||
|
||||
# Pick the right Torch wheel for your GPU (driver 570+ → cu128):
|
||||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
|
||||
# CPU-only fallback (slow):
|
||||
# pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
|
||||
# SAM 2 source + deps
|
||||
pip install -e ../sam2/sam2
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Optional — SAM 3 support (needs HF account + accepted gating)
|
||||
git clone https://github.com/facebookresearch/sam3.git ..
|
||||
pip install -e ../sam3
|
||||
pip install "transformers>=4.57"
|
||||
hf auth login # paste read token from https://huggingface.co/settings/tokens
|
||||
|
||||
./run.sh # binds 0.0.0.0:9090, auto-picks cuda/cpu
|
||||
```
|
||||
|
||||
Smoke-test:
|
||||
```bash
|
||||
curl http://localhost:9090/ # {"status":"ok",...}
|
||||
curl http://localhost:9090/models # registered models
|
||||
```
|
||||
|
||||
Full backend setup details: [`sam2-backend/README.md`](sam2-backend/README.md).
|
||||
|
||||
### 2. Tauri desktop app
|
||||
|
||||
```bash
|
||||
cd sam-tool-tauri
|
||||
npm install
|
||||
npm run tauri dev # hot-reload dev mode
|
||||
|
||||
# Or for a real build / installer:
|
||||
npm run tauri build # see docs/BUILD_RELEASE.md
|
||||
```
|
||||
|
||||
In the running app:
|
||||
1. **Login** — pick or type a username (no password — used for annotation
|
||||
author stamping).
|
||||
2. **Extract mode** — Upload a video → optionally `Import annotations…`
|
||||
(COCO / SAM JSON) → click `choose output folder…` (or just open a
|
||||
video; the default `<video>_extracted/` folder is auto-adopted) →
|
||||
scrub timeline, press **E** to extract, **U** to undo, **[ … ]** for
|
||||
bulk range.
|
||||
3. **Annotate mode** — Open the extracted folder. Draw bboxes (B) /
|
||||
polygons (P), label them via the picker, switch to select (S) to
|
||||
move/resize/reclass. Settings → enable SAM, pick a model, get
|
||||
polygons back instead of bboxes.
|
||||
|
||||
### 3. Deploying to another PC
|
||||
|
||||
Backend → see the install steps in [`sam2-backend/README.md`](sam2-backend/README.md)
|
||||
("install on a different PC" section).
|
||||
|
||||
Desktop app → [`docs/BUILD_RELEASE.md`](docs/BUILD_RELEASE.md) covers
|
||||
producing `.deb` / AppImage installers.
|
||||
|
||||
---
|
||||
|
||||
## Feature highlights
|
||||
|
||||
### Extract mode
|
||||
- Persistent `ffmpeg -f image2pipe` subprocess + LRU cache + prefetch
|
||||
worker → smooth playback up to ~2× on 1080p.
|
||||
- Per-video default output folder (`<video>_extracted/`) auto-restored on
|
||||
reopen — extracted-frame state and "EXTRACTED" tint reflected
|
||||
immediately.
|
||||
- Timeline ticks mark every previously-extracted frame at a glance.
|
||||
- COCO + custom SAM JSON import, smart overlay with per-class colours,
|
||||
Review mode auto-pauses on annotated frames.
|
||||
- `Space` play/pause · `←→` ±1 frame · `Ctrl+←→` ±10 · `↑↓` speed 0.25–4×
|
||||
· `H` overlay · `E` extract · `U` undo · `[ … ]` bulk range.
|
||||
|
||||
### Annotate mode
|
||||
- Sibling JSON per image with `source_video` field — annotations know
|
||||
where they came from.
|
||||
- Bbox + polygon tools, vertex drag (rAF-throttled), bbox handles, undo
|
||||
(40 deep) + redo, Tab-reclass, inline class creation in the picker.
|
||||
- Multi-select image list (Shift-range, Ctrl-toggle) + bulk delete with
|
||||
two-step confirmation.
|
||||
- A/D image navigation, +/- zoom, cursor-anchored Ctrl+wheel zoom.
|
||||
- SAM mode round-trip with live phase indicator showing model + inference
|
||||
time (auto-fades after 3 s).
|
||||
- Auto-switch to **select** mode after committing each new shape.
|
||||
|
||||
### SAM backend
|
||||
- Pluggable model registry — add SAM-3 / SAM-4 / etc. by dropping an
|
||||
adapter into `adapters/` and one entry into `models.yaml`.
|
||||
- Lazy-load with single-resident eviction (set
|
||||
`SAM_BACKEND_KEEP_ALL=1` to keep all loaded if you have VRAM).
|
||||
- Speaks the Label-Studio-ML `/predict` protocol the Tauri client
|
||||
already used — drop-in compatible.
|
||||
- Local SAM 3 image-predictor + HuggingFace SAM 3 tracker both supported.
|
||||
- Privacy-first: telemetry off by default, fully air-gappable after the
|
||||
first weight download — see
|
||||
[`docs/HF_DEPLOY_OFFLINE.md`](docs/HF_DEPLOY_OFFLINE.md).
|
||||
|
||||
---
|
||||
|
||||
## Documentation index
|
||||
|
||||
| Doc | What it covers |
|
||||
|----------------------------------------------------|-------------------------------------------------------------|
|
||||
| [docs/POLYGON_SMOOTHING.md](docs/POLYGON_SMOOTHING.md) | Tune mask→polygon: point count, kernel size, smoothing on/off |
|
||||
| [docs/HF_DEPLOY_OFFLINE.md](docs/HF_DEPLOY_OFFLINE.md) | Download HF SAM 3 weights, then run fully offline |
|
||||
| [docs/BUILD_RELEASE.md](docs/BUILD_RELEASE.md) | Build the Tauri app for production (deb / AppImage / binary) |
|
||||
| [sam2-backend/README.md](sam2-backend/README.md) | Backend install, run, API reference, model-registry details |
|
||||
| [CLAUDE.md](CLAUDE.md) | Project rules — Think Before Coding · Simplicity First · Surgical Changes · Goal-Driven Execution |
|
||||
|
||||
---
|
||||
|
||||
## Stack
|
||||
|
||||
- **Desktop:** Tauri 2 shell, Rust backend, React 18 + TypeScript, Vite,
|
||||
WebKit2GTK 2.50.
|
||||
- **Backend:** Python 3.10+, FastAPI, uvicorn, Pillow, NumPy, OpenCV
|
||||
(headless), PyTorch ≥ 2.4 with matching CUDA, optional `transformers ≥ 4.57`.
|
||||
- **Models:** Meta SAM 2 / SAM 2.1 (4 Hiera variants), Meta SAM 3.
|
||||
- **Target OS:** Linux (Ubuntu 22.04 / Pop!_OS), AMD APU + NVIDIA RTX
|
||||
GPUs tested.
|
||||
|
||||
---
|
||||
|
||||
## Conventions
|
||||
|
||||
From [CLAUDE.md](CLAUDE.md):
|
||||
|
||||
- **Think Before Coding** — wrong assumptions are the most expensive bugs.
|
||||
- **Simplicity First** — fewer moving parts beats clever abstractions.
|
||||
- **Surgical Changes** — touch what the task needs, leave the rest.
|
||||
- **Goal-Driven Execution** — verify the user-facing behaviour, not just
|
||||
the code path.
|
||||
|
||||
A few specific rules learned the hard way:
|
||||
- `React.StrictMode` double-invokes setState updaters; don't put side
|
||||
effects (like `queueMicrotask` calls or save requests) inside them.
|
||||
Use refs + explicit calls instead.
|
||||
- No `crypto.randomUUID()` without a fallback — webview contexts
|
||||
sometimes lack it; use the `uuid()` helper.
|
||||
- Output folders must live **outside** `sam-tool-tauri/` — the dev
|
||||
watcher restarts the app on any file change inside the project tree.
|
||||
- Backend telemetry stays **off** by default; HF cache moves to a known
|
||||
`HF_HOME` for explicit deployment.
|
||||
226
docs/BUILD_RELEASE.md
Normal file
226
docs/BUILD_RELEASE.md
Normal file
@@ -0,0 +1,226 @@
|
||||
# Building the Tauri app for release (instead of `dev` mode)
|
||||
|
||||
Day-to-day you've been running the desktop app via `npm run tauri dev` —
|
||||
that's the dev server: hot-reload, attached browser devtools, slower
|
||||
binary, and it stays bound to your terminal. For real use (or shipping to
|
||||
another machine), you build a standalone executable + installer.
|
||||
|
||||
---
|
||||
|
||||
## TL;DR
|
||||
|
||||
```bash
|
||||
cd sam-tool-tauri
|
||||
npm install # only on first run / after pulling deps
|
||||
npm run tauri build # produces a release binary + installer
|
||||
```
|
||||
|
||||
Output lands under
|
||||
`sam-tool-tauri/src-tauri/target/release/bundle/`. On Linux you'll find:
|
||||
|
||||
- `bundle/deb/sam-tool-tauri_<version>_amd64.deb` — installable on
|
||||
Debian/Ubuntu hosts (`sudo dpkg -i …`).
|
||||
- `bundle/appimage/sam-tool-tauri_<version>_amd64.AppImage` — single
|
||||
portable file, `chmod +x` and run.
|
||||
- `bundle/rpm/…` if rpm tooling is present.
|
||||
|
||||
Plus the raw binary at
|
||||
`src-tauri/target/release/sam-tool-tauri` (no installer, just `./run` it).
|
||||
|
||||
---
|
||||
|
||||
## Prerequisites (Linux)
|
||||
|
||||
If you've only ever run `tauri dev`, you might be missing the system
|
||||
libraries the bundler needs. On Ubuntu / Debian / Pop!_OS:
|
||||
|
||||
```bash
|
||||
sudo apt update
|
||||
sudo apt install -y \
|
||||
libwebkit2gtk-4.1-dev \
|
||||
libgtk-3-dev \
|
||||
libayatana-appindicator3-dev \
|
||||
librsvg2-dev \
|
||||
patchelf \
|
||||
build-essential \
|
||||
curl wget file \
|
||||
libssl-dev pkg-config
|
||||
```
|
||||
|
||||
You also need a recent Rust toolchain (Tauri 2 requires `rustc >= 1.77`):
|
||||
|
||||
```bash
|
||||
rustup update stable
|
||||
rustc --version # confirm >= 1.77
|
||||
```
|
||||
|
||||
For AppImage bundling specifically, the Tauri CLI auto-fetches the
|
||||
AppImage tools on first build. The first run will be slower; subsequent
|
||||
builds reuse the cache.
|
||||
|
||||
---
|
||||
|
||||
## Build profiles
|
||||
|
||||
Two ways to produce a release binary, with different tradeoffs:
|
||||
|
||||
### A. `npm run tauri build` — full installer build
|
||||
|
||||
```bash
|
||||
npm run tauri build
|
||||
```
|
||||
|
||||
What it does:
|
||||
1. Runs `vite build` → produces optimized JS/CSS in `dist/`.
|
||||
2. Compiles `src-tauri/` Rust crate in `--release` mode.
|
||||
3. Bundles everything into platform-native installers under
|
||||
`bundle/`.
|
||||
|
||||
Pros: production-grade, optimised, signable, distributable.
|
||||
Cons: slow first time (Rust release builds + bundler downloads).
|
||||
|
||||
### B. `cargo build --release` — binary only, skip installers
|
||||
|
||||
If you only want the executable (e.g. for quick testing on a sibling
|
||||
machine on the same OS), drop the bundling step:
|
||||
|
||||
```bash
|
||||
cd sam-tool-tauri
|
||||
npm run build # bundles the frontend → dist/
|
||||
cd src-tauri
|
||||
cargo build --release # binary at target/release/sam-tool-tauri
|
||||
```
|
||||
|
||||
The resulting binary expects `dist/` from the same checkout to be
|
||||
present (Tauri 2 embeds frontend assets, so once compiled it's
|
||||
self-contained). Faster than the full bundle, but no `.deb`/`.AppImage`.
|
||||
|
||||
### C. Debug build — fast compile, big binary
|
||||
|
||||
For iterating on Rust changes without a full release re-compile:
|
||||
|
||||
```bash
|
||||
cd src-tauri
|
||||
cargo build # debug profile — fast compile, big slow binary
|
||||
./target/debug/sam-tool-tauri
|
||||
```
|
||||
|
||||
This is essentially what `tauri dev` runs under the hood, minus the dev
|
||||
server.
|
||||
|
||||
---
|
||||
|
||||
## Where the artefacts go
|
||||
|
||||
```
|
||||
sam-tool-tauri/
|
||||
├── dist/ ← bundled frontend (vite build)
|
||||
└── src-tauri/target/
|
||||
├── debug/sam-tool-tauri ← debug binary
|
||||
└── release/
|
||||
├── sam-tool-tauri ← optimised binary
|
||||
└── bundle/
|
||||
├── deb/sam-tool-tauri_<v>_amd64.deb
|
||||
├── appimage/sam-tool-tauri_<v>_amd64.AppImage
|
||||
└── rpm/… (if `rpm` tools are installed)
|
||||
```
|
||||
|
||||
`bundle/deb/` and `bundle/appimage/` are the things you'd ship to other
|
||||
machines.
|
||||
|
||||
---
|
||||
|
||||
## Versioning
|
||||
|
||||
The version string comes from
|
||||
[`src-tauri/tauri.conf.json`](../sam-tool-tauri/src-tauri/tauri.conf.json) →
|
||||
`"version": "..."`. Bump it before each release; the value is baked
|
||||
into both the binary and the installer filenames.
|
||||
|
||||
---
|
||||
|
||||
## Smoke-test the release build locally
|
||||
|
||||
```bash
|
||||
# Run the binary directly:
|
||||
./src-tauri/target/release/sam-tool-tauri
|
||||
|
||||
# Or install the .deb and launch from your app menu:
|
||||
sudo dpkg -i src-tauri/target/release/bundle/deb/sam-tool-tauri_*_amd64.deb
|
||||
sam-tool-tauri # binary on $PATH after dpkg
|
||||
|
||||
# Or run the AppImage portably:
|
||||
chmod +x src-tauri/target/release/bundle/appimage/sam-tool-tauri_*.AppImage
|
||||
./src-tauri/target/release/bundle/appimage/sam-tool-tauri_*.AppImage
|
||||
```
|
||||
|
||||
The release build behaves exactly like `dev` (same UI, same SAM
|
||||
integration), with these differences:
|
||||
|
||||
- No devtools by default. To re-enable in a release build, edit
|
||||
`src-tauri/tauri.conf.json` → `"app.windows[0].devtools": true`,
|
||||
rebuild.
|
||||
- Faster cold start, smaller memory footprint, optimised JS bundle.
|
||||
- React `StrictMode` still doubles state-updaters in dev — but that's a
|
||||
React-side flag, not Tauri's. The stripped-down behaviour comes from
|
||||
Vite's production mode (`process.env.NODE_ENV === 'production'`).
|
||||
|
||||
---
|
||||
|
||||
## Distributing to another PC
|
||||
|
||||
The user only needs **one** of the bundles:
|
||||
|
||||
**Debian/Ubuntu:**
|
||||
```bash
|
||||
scp src-tauri/target/release/bundle/deb/sam-tool-tauri_*_amd64.deb \
|
||||
target-host:/tmp/
|
||||
ssh target-host "sudo dpkg -i /tmp/sam-tool-tauri_*_amd64.deb"
|
||||
```
|
||||
|
||||
**Anywhere with glibc:**
|
||||
```bash
|
||||
scp src-tauri/target/release/bundle/appimage/sam-tool-tauri_*.AppImage \
|
||||
target-host:/home/user/
|
||||
ssh target-host "chmod +x ~/sam-tool-tauri_*.AppImage && ~/sam-tool-tauri_*.AppImage"
|
||||
```
|
||||
|
||||
The target host needs `webkit2gtk` runtime (most desktops already have
|
||||
it) and `ffmpeg` if they want the Extract mode to decode video. The
|
||||
sibling SAM2 backend and weights are *not* bundled into the desktop app —
|
||||
those live separately and are reached over HTTP, configurable from
|
||||
**Settings → SAM backend URL**.
|
||||
|
||||
---
|
||||
|
||||
## Common build errors
|
||||
|
||||
| Error | Fix |
|
||||
|--------------------------------------------------------------|-------------------------------------------------------------------------------------------|
|
||||
| `error: linker 'cc' not found` | `sudo apt install build-essential` |
|
||||
| `failed to find webkit2gtk-4.1` / `gobject-introspection` | `sudo apt install libwebkit2gtk-4.1-dev libgtk-3-dev` (see prerequisites) |
|
||||
| `failed to bundle .deb` | Missing `dpkg-deb`. `sudo apt install dpkg` |
|
||||
| AppImage step hangs on first build | Tauri downloading appimagetool; wait or pre-fetch under `~/.cache/tauri/` |
|
||||
| `Could not resolve "@tauri-apps/api"` | `npm install` again in `sam-tool-tauri/` |
|
||||
| Build OK on dev machine, fails on Ubuntu 22.04 | webkit2gtk-4.0 vs 4.1 mismatch — Tauri 2 needs 4.1; older Ubuntus may need a backport PPA |
|
||||
|
||||
---
|
||||
|
||||
## Cross-compilation (Windows / macOS targets)
|
||||
|
||||
Out of scope for this repo (Linux-only per CLAUDE.md), but for
|
||||
completeness — Tauri supports it via extra rust targets and a Windows or
|
||||
macOS host machine. See the upstream Tauri docs:
|
||||
https://tauri.app/distribute/
|
||||
|
||||
---
|
||||
|
||||
## Recap — when to run what
|
||||
|
||||
| You want to… | Command |
|
||||
|---------------------------------------------|-------------------------------------------------|
|
||||
| Iterate on UI (hot reload) | `npm run tauri dev` |
|
||||
| Iterate on Rust quickly | `cd src-tauri && cargo build && ./target/debug/sam-tool-tauri` |
|
||||
| Make a release executable, no installer | `npm run build && cd src-tauri && cargo build --release` |
|
||||
| Make installers (.deb / AppImage) | `npm run tauri build` |
|
||||
| Inspect production frontend without Tauri | `npm run preview` (browser-only, no native APIs)|
|
||||
226
docs/HF_DEPLOY_OFFLINE.md
Normal file
226
docs/HF_DEPLOY_OFFLINE.md
Normal file
@@ -0,0 +1,226 @@
|
||||
# Deploying SAM3 via HuggingFace, then going fully offline
|
||||
|
||||
The SAM3 backend uses Meta's `facebook/sam3` weights via HuggingFace. The
|
||||
flow is **online once** (to download + cache the weights) → **then
|
||||
air-gapped** (server runs without ever talking to HuggingFace again).
|
||||
This doc walks through both phases and how to verify your box is actually
|
||||
offline at the end.
|
||||
|
||||
---
|
||||
|
||||
## Phase 0 — prerequisites
|
||||
|
||||
On the target PC:
|
||||
|
||||
- The `sam2-backend/` folder, copied or cloned, with its `.venv`
|
||||
containing `torch`, `transformers>=4.57`, `huggingface_hub`, and the
|
||||
rest of `requirements.txt`. See [`sam2-backend/README.md`](../sam2-backend/README.md)
|
||||
for the full install.
|
||||
- A HuggingFace account that has been **granted access to
|
||||
`facebook/sam3`** — gated. Request access at
|
||||
https://huggingface.co/facebook/sam3 (usually approved within minutes).
|
||||
- A read token from https://huggingface.co/settings/tokens.
|
||||
|
||||
---
|
||||
|
||||
## Phase 1 — first-time online download
|
||||
|
||||
The first call to `Sam3TrackerModel.from_pretrained("facebook/sam3")`
|
||||
streams the weights from HF and writes them to
|
||||
`~/.cache/huggingface/hub/`. Everything afterward can be served from that
|
||||
cache with no network.
|
||||
|
||||
### 1.1 Authenticate
|
||||
|
||||
```bash
|
||||
cd /home/<you>/sam2-backend
|
||||
source .venv/bin/activate
|
||||
hf auth login # paste your read token
|
||||
hf auth whoami # confirm the right account
|
||||
```
|
||||
|
||||
Token is stored at `~/.cache/huggingface/token`. It's a plain file with
|
||||
`chmod 600`.
|
||||
|
||||
### 1.2 Set offline flag to OFF for the first start
|
||||
|
||||
In `run.sh`, the relevant lines are:
|
||||
```bash
|
||||
export HF_HUB_DISABLE_TELEMETRY="${HF_HUB_DISABLE_TELEMETRY:-1}"
|
||||
export DO_NOT_TRACK="${DO_NOT_TRACK:-1}"
|
||||
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-0}" # ← MUST be 0 first time
|
||||
export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-0}" # ← MUST be 0 first time
|
||||
```
|
||||
|
||||
For a one-off override, just don't pass the env vars:
|
||||
```bash
|
||||
./run.sh
|
||||
```
|
||||
|
||||
### 1.3 Make ONE successful predict call
|
||||
|
||||
Either from the Tauri app (Settings → pick `SAM 3 · tracker (HuggingFace)` →
|
||||
draw a bbox), or directly:
|
||||
|
||||
```bash
|
||||
curl -s -X POST "http://localhost:9090/predict?model=sam3_tracker" \
|
||||
-H "Content-Type: application/json" \
|
||||
--data '@some_test_payload.json'
|
||||
```
|
||||
|
||||
A successful response means the weights are now fully cached. Confirm:
|
||||
|
||||
```bash
|
||||
ls -la ~/.cache/huggingface/hub/models--facebook--sam3/
|
||||
# expect a snapshots/<hash>/ subdir with model.safetensors and configs
|
||||
du -sh ~/.cache/huggingface/hub/models--facebook--sam3/
|
||||
# ~3-4 GB for SAM3
|
||||
```
|
||||
|
||||
### 1.4 Optional — pre-warm without a real predict
|
||||
|
||||
If you want to download but not yet wire up Tauri, do it from a Python
|
||||
shell:
|
||||
|
||||
```bash
|
||||
python3 - <<'EOF'
|
||||
from transformers import Sam3TrackerModel, Sam3TrackerProcessor
|
||||
Sam3TrackerProcessor.from_pretrained("facebook/sam3")
|
||||
Sam3TrackerModel.from_pretrained("facebook/sam3")
|
||||
print("cached.")
|
||||
EOF
|
||||
```
|
||||
|
||||
Same effect — fills `~/.cache/huggingface/hub/`.
|
||||
|
||||
---
|
||||
|
||||
## Phase 2 — flip to fully offline
|
||||
|
||||
Once the cache is populated, set the offline flags so the server never
|
||||
tries to revalidate metadata or check for updates.
|
||||
|
||||
### 2.1 Edit `run.sh`
|
||||
|
||||
Change the defaults to `1`:
|
||||
|
||||
```bash
|
||||
export HF_HUB_DISABLE_TELEMETRY="${HF_HUB_DISABLE_TELEMETRY:-1}"
|
||||
export DO_NOT_TRACK="${DO_NOT_TRACK:-1}"
|
||||
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}" # ← 1
|
||||
export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}" # ← 1
|
||||
```
|
||||
|
||||
Or, if you're running via systemd, add to your unit file:
|
||||
|
||||
```ini
|
||||
[Service]
|
||||
Environment=HF_HUB_DISABLE_TELEMETRY=1
|
||||
Environment=DO_NOT_TRACK=1
|
||||
Environment=HF_HUB_OFFLINE=1
|
||||
Environment=TRANSFORMERS_OFFLINE=1
|
||||
```
|
||||
|
||||
Then `sudo systemctl daemon-reload && sudo systemctl restart sam-backend`.
|
||||
|
||||
### 2.2 Restart and verify
|
||||
|
||||
```bash
|
||||
./run.sh
|
||||
```
|
||||
|
||||
You should see clean startup logs **without** any `httpx HTTP Request`
|
||||
lines hitting `huggingface.co`. The `Loading weights: 685/685 [00:00<00:00]`
|
||||
line confirms a cache hit (instant load, no download).
|
||||
|
||||
### 2.3 Pin a custom cache location (optional)
|
||||
|
||||
If you want the cache somewhere predictable (for backup, rsync to other
|
||||
boxes, or to keep `$HOME` clean), set `HF_HOME` in `run.sh`:
|
||||
|
||||
```bash
|
||||
export HF_HOME=/var/cache/sam-backend/huggingface
|
||||
```
|
||||
|
||||
Move the existing cache once:
|
||||
```bash
|
||||
mv ~/.cache/huggingface /var/cache/sam-backend/
|
||||
```
|
||||
|
||||
Subsequent reads come from there.
|
||||
|
||||
---
|
||||
|
||||
## Verifying you're really offline
|
||||
|
||||
Two-stage check:
|
||||
|
||||
### Step 1 — disconnect & restart
|
||||
|
||||
Easiest sanity test: pull the network cable (or `sudo ip link set eth0 down`),
|
||||
restart the server, hit `/predict` from the Tauri app on the same box.
|
||||
If it returns a polygon, you're truly offline.
|
||||
|
||||
### Step 2 — packet-level proof
|
||||
|
||||
If you can't disconnect (e.g., other services need the network), watch
|
||||
for any HF traffic during normal use:
|
||||
|
||||
```bash
|
||||
sudo tcpdump -nn -i any \
|
||||
'host huggingface.co or host cdn-lfs.huggingface.co' \
|
||||
2>/dev/null
|
||||
```
|
||||
|
||||
Restart the backend, hit `/predict` ten times from the Tauri app. With
|
||||
the OFFLINE flags set you should see **zero packets**.
|
||||
|
||||
---
|
||||
|
||||
## Moving the cache to a different machine
|
||||
|
||||
Useful when you want to deploy to a new box without the first-run
|
||||
download (e.g., setting up a second annotation station behind the same
|
||||
firewall).
|
||||
|
||||
```bash
|
||||
# On the source machine:
|
||||
tar czf hf-sam3-cache.tar.gz -C ~/.cache huggingface
|
||||
scp hf-sam3-cache.tar.gz target-box:/tmp/
|
||||
|
||||
# On the target machine:
|
||||
mkdir -p ~/.cache
|
||||
tar xzf /tmp/hf-sam3-cache.tar.gz -C ~/.cache
|
||||
```
|
||||
|
||||
Then start the new backend with `HF_HUB_OFFLINE=1` from day one — it'll
|
||||
load from the imported cache without ever calling out.
|
||||
|
||||
---
|
||||
|
||||
## Common gotchas
|
||||
|
||||
| Symptom | Cause / fix |
|
||||
|--------------------------------------------------------------------------|-----------------------------------------------------------------------------|
|
||||
| `Cannot find the requested files in the cached path` while OFFLINE=1 | Cache is incomplete or in a different `HF_HOME`. Re-run online once. |
|
||||
| `gated repo, access denied` | Token doesn't have `facebook/sam3` access. Accept gate on the model page, then `hf auth login` again. |
|
||||
| Many `HEAD https://huggingface.co/...` lines on startup despite OFFLINE=1 | One of the two env flags is missing. Both `HF_HUB_OFFLINE=1` AND `TRANSFORMERS_OFFLINE=1` are required. |
|
||||
| `Loading weights: ...` takes minutes even with cache | The cache file is on a slow disk (NFS, HDD). Move it to local SSD via `HF_HOME`. |
|
||||
| Token leaked into logs | `hf auth login` only echoes your username, never the token. The token at `~/.cache/huggingface/token` is the only place it's stored. |
|
||||
|
||||
---
|
||||
|
||||
## What HF still sees in offline mode
|
||||
|
||||
**Nothing.** With both `HF_HUB_OFFLINE=1` and `TRANSFORMERS_OFFLINE=1`:
|
||||
|
||||
- `from_pretrained` reads only from `~/.cache/huggingface/hub/` — no API
|
||||
calls, no HEAD revalidation pings.
|
||||
- Telemetry is independently off via `HF_HUB_DISABLE_TELEMETRY=1` /
|
||||
`DO_NOT_TRACK=1`.
|
||||
- Inference itself never touches the network — your images, bboxes, and
|
||||
resulting masks stay on the box regardless of online/offline mode.
|
||||
|
||||
The only thing HF "knows" about you, ever, is "this account downloaded
|
||||
`facebook/sam3` once on date X" — recorded server-side at the moment
|
||||
`from_pretrained` first ran in Phase 1.
|
||||
164
docs/POLYGON_SMOOTHING.md
Normal file
164
docs/POLYGON_SMOOTHING.md
Normal file
@@ -0,0 +1,164 @@
|
||||
# Tuning polygon output (point count + edge smoothness)
|
||||
|
||||
When SAM returns a mask, the backend converts it to a polygon before sending
|
||||
it back to the Tauri annotation tool. Two things you can tune:
|
||||
|
||||
1. **Point count** — how many vertices the polygon has. Each vertex is
|
||||
draggable in the UI, so too many makes editing painful.
|
||||
2. **Edge smoothness** — how the polygon traces curved edges (smooth curves
|
||||
vs. pixel staircases).
|
||||
|
||||
Both are controlled in one file:
|
||||
[`sam2-backend/adapters/utils.py`](../sam2-backend/adapters/utils.py).
|
||||
|
||||
---
|
||||
|
||||
## Where it's called
|
||||
|
||||
Both adapters do the same thing in their predict path:
|
||||
|
||||
```python
|
||||
# sam2-backend/adapters/sam2_adapter.py
|
||||
return PredictResult(
|
||||
polygon=_mask_to_polygon(masks_np[best]),
|
||||
score=best_score,
|
||||
)
|
||||
|
||||
# sam2-backend/adapters/sam3_adapter.py — same call site
|
||||
```
|
||||
|
||||
The defaults are `min_points=4`, `max_points=10`, `smooth=True`. You can
|
||||
override per-call:
|
||||
|
||||
```python
|
||||
polygon=_mask_to_polygon(masks_np[best], min_points=3, max_points=15)
|
||||
```
|
||||
|
||||
…or change the defaults in `utils.py::mask_to_polygon` and they'll apply
|
||||
to every model.
|
||||
|
||||
---
|
||||
|
||||
## Knobs
|
||||
|
||||
### 1. Point count band — `min_points` / `max_points`
|
||||
|
||||
The simplifier uses Douglas–Peucker (`cv2.approxPolyDP`) with a binary
|
||||
search over `epsilon` to land in `[min_points, max_points]`. Within that
|
||||
band it prefers **more** detail (i.e. closer to `max_points`).
|
||||
|
||||
| Band | Use it when… |
|
||||
|------------|-----------------------------------------------------------------------------|
|
||||
| `4..10` | Default. Comfortable to drag, captures common shapes. |
|
||||
| `4..15` | Curvy / organic objects (animals, foliage). 15 vertices follow a curve well. |
|
||||
| `3..6` | Very simple shapes (rectangles, traffic signs). Fewer drag handles. |
|
||||
| `6..20` | High-fidelity work, you don't mind dragging more. |
|
||||
|
||||
**Avoid `max_points > 25`** — past that, the binary search needs many
|
||||
iterations and the marginal benefit per vertex drops sharply. If you need
|
||||
that many points, the bottleneck is mask quality, not simplification.
|
||||
|
||||
### 2. Mask smoothing — `smooth=True | False`
|
||||
|
||||
Before extracting the contour, the helper optionally runs:
|
||||
|
||||
1. **Morphological close + open** (ellipse kernel) — fills single-pixel
|
||||
notches and removes pepper noise.
|
||||
2. **Light Gaussian blur + re-threshold** — sub-pixel edge smoothing so
|
||||
curves don't trace as staircases.
|
||||
|
||||
Pass `smooth=False` to skip both steps and trace the raw mask exactly. Use
|
||||
this when:
|
||||
- The mask has fine features that smoothing erases (thin antennae, sharp
|
||||
90° corners that get rounded).
|
||||
- You're A/B testing whether smoothing is helping or hurting.
|
||||
|
||||
```python
|
||||
polygon=_mask_to_polygon(masks_np[best], smooth=False)
|
||||
```
|
||||
|
||||
### 3. Smoothing kernel size — `_smooth_mask`
|
||||
|
||||
Inside `_smooth_mask` (called by `mask_to_polygon` when `smooth=True`),
|
||||
kernel size scales with the **blob's bounding box diagonal**:
|
||||
|
||||
```python
|
||||
k = int(round(blob_diag * 0.015)) # ← morphology kernel multiplier
|
||||
k = max(3, min(k | 1, 11)) # odd, clamped 3..11
|
||||
|
||||
blur_k = max(3, min(k | 1, 9)) # ← Gaussian blur kernel
|
||||
```
|
||||
|
||||
**The `0.015` multiplier is the main dial.**
|
||||
|
||||
| Multiplier | Effect |
|
||||
|------------|-----------------------------------------------------------------|
|
||||
| `0.008` | Very gentle. Preserves thin features. Use if smoothing dilates the polygon. |
|
||||
| `0.015` | Default. Good balance for typical objects. |
|
||||
| `0.025` | Aggressive. Use only on visibly noisy masks; will round corners. |
|
||||
|
||||
The 11-pixel hard cap means giant masks don't get insanely smoothed.
|
||||
|
||||
### 4. Drop the Gaussian, keep morphology
|
||||
|
||||
If hard 90° corners get rounded but you still want morphology to clean up
|
||||
notches, delete just the Gaussian step in `_smooth_mask`:
|
||||
|
||||
```python
|
||||
# remove these two lines
|
||||
m = cv2.GaussianBlur(m, (blur_k, blur_k), 0)
|
||||
_, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
|
||||
```
|
||||
|
||||
This keeps close+open (which is mostly noise removal) and skips the
|
||||
sub-pixel curve smoothing.
|
||||
|
||||
---
|
||||
|
||||
## Diagnostic recipe — find the right knob in 3 minutes
|
||||
|
||||
1. **Draw a representative bbox in the Tauri app.** Screenshot the result.
|
||||
2. **Set `smooth=False`** in the adapter, restart, draw the same bbox.
|
||||
- If polygon is *better* (less bloat, hugs edges) → smoothing is over-aggressive. Lower the kernel multiplier or delete the Gaussian step.
|
||||
- If polygon is *worse* (jagged stairs) → smoothing was helping. Move on to step 3.
|
||||
3. **Bump `max_points` to 15 or 20**, restart.
|
||||
- If polygon is *better* → simplification was too aggressive. Land on `max_points=12..15` as your new default.
|
||||
- If polygon is the same → the mask itself is the bottleneck (see "If the raw mask is bad" below).
|
||||
4. **Restore `smooth=True`** and try a model swap (SAM 2.1 large ↔ SAM 3 tracker) — the underlying mask differs noticeably between models.
|
||||
|
||||
---
|
||||
|
||||
## If the raw mask is bad
|
||||
|
||||
If after all the above the polygon still looks wrong, the input to the
|
||||
simplifier is bad — fix that, not the polygon code. Common moves:
|
||||
|
||||
- **Pad the bbox before sending** to SAM (Meta's reference scripts use
|
||||
`pad=10` pixels). SAM gives crisper masks with a small box margin.
|
||||
- **Try a larger SAM model** (e.g. switch the dropdown from `SAM 2.1
|
||||
Hiera Small` to `SAM 2.1 Hiera Large`).
|
||||
- **Tighten the bbox** around the actual object — SAM tends to fill the
|
||||
box, so a generous box gets a bloated mask.
|
||||
|
||||
---
|
||||
|
||||
## Quick reference — every dial in one block
|
||||
|
||||
```python
|
||||
# sam2-backend/adapters/utils.py
|
||||
|
||||
mask_to_polygon(
|
||||
mask,
|
||||
min_points=4, # never below this
|
||||
max_points=10, # never above this; fewer = easier to drag
|
||||
smooth=True, # turn off to trace raw mask
|
||||
)
|
||||
|
||||
# Inside _smooth_mask:
|
||||
k = int(round(blob_diag * 0.015)) # ← MAIN DIAL
|
||||
# smaller = preserve detail
|
||||
# bigger = more rounded
|
||||
```
|
||||
|
||||
If you change defaults, restart the backend (`./run.sh`). No frontend
|
||||
restart needed — the polygon arrives over the existing `/predict` call.
|
||||
120
files/README.md
Normal file
120
files/README.md
Normal file
@@ -0,0 +1,120 @@
|
||||
# SAM 3 local segmentation server
|
||||
|
||||
A small FastAPI server that wraps Meta's **SAM 3** (released Nov 19, 2025) so your local annotation tool can send a bbox and get back a segmentation mask.
|
||||
|
||||
## What it does
|
||||
|
||||
Two endpoints, both take an image + bbox:
|
||||
|
||||
| Endpoint | Behavior | Use for |
|
||||
|---|---|---|
|
||||
| `POST /segment/box` | Returns the mask of the **one** object inside your bbox. Uses the SAM 3 tracker (SAM2-style PVS). | Classic annotation — user draws a box around one object, you want that one mask. |
|
||||
| `POST /segment/exemplar` | Treats the bbox as an **exemplar** and returns masks for **all similar** objects in the image. | Auto-propagation — label one thing, get all instances. |
|
||||
|
||||
Mask is returned as a base64 PNG, plus optional **COCO RLE** and **polygon contour** for easy integration with annotation-tool formats (CVAT, Label Studio, COCO JSON, etc).
|
||||
|
||||
## Setup
|
||||
|
||||
### 1. Prerequisites
|
||||
- Python 3.12+
|
||||
- PyTorch 2.7+ with CUDA 12.6 (CPU works too, just slow)
|
||||
- A CUDA-capable GPU is strongly recommended — SAM 3 is 848M params
|
||||
|
||||
### 2. Request access to weights
|
||||
SAM 3 weights are gated on Hugging Face. Request access at https://huggingface.co/facebook/sam3, wait for approval, then:
|
||||
```bash
|
||||
hf auth login
|
||||
```
|
||||
|
||||
### 3. Install
|
||||
```bash
|
||||
# SAM 3 itself
|
||||
git clone https://github.com/facebookresearch/sam3.git
|
||||
cd sam3
|
||||
pip install -e .
|
||||
|
||||
# Server deps
|
||||
pip install fastapi uvicorn pillow numpy requests opencv-python-headless pycocotools
|
||||
# Optional but recommended — gives the clean tracker API used by /segment/box
|
||||
pip install "transformers>=4.57"
|
||||
```
|
||||
|
||||
### 4. Run
|
||||
```bash
|
||||
python sam3_server.py --host 0.0.0.0 --port 8000
|
||||
# or on CPU:
|
||||
python sam3_server.py --device cpu
|
||||
```
|
||||
|
||||
Check it's alive:
|
||||
```bash
|
||||
curl http://localhost:8000/health
|
||||
```
|
||||
|
||||
## Calling it from your annotation tool
|
||||
|
||||
### Python
|
||||
See `client_example.py`. Minimal version:
|
||||
```python
|
||||
import base64, requests
|
||||
img_b64 = base64.b64encode(open("img.jpg","rb").read()).decode()
|
||||
r = requests.post("http://localhost:8000/segment/box", json={
|
||||
"image": img_b64,
|
||||
"bbox": [120, 80, 540, 420], # [x1, y1, x2, y2] pixels
|
||||
})
|
||||
data = r.json()
|
||||
# data["mask"] -> base64 PNG mask
|
||||
# data["polygon"] -> [[x,y], ...] contour
|
||||
# data["rle"] -> COCO RLE dict
|
||||
```
|
||||
|
||||
### JavaScript (browser annotation tools)
|
||||
```js
|
||||
const imgB64 = /* base64 of your image */;
|
||||
const res = await fetch("http://localhost:8000/segment/box", {
|
||||
method: "POST",
|
||||
headers: {"Content-Type": "application/json"},
|
||||
body: JSON.stringify({ image: imgB64, bbox: [x1, y1, x2, y2] })
|
||||
});
|
||||
const { mask, polygon, rle, score } = await res.json();
|
||||
// Draw `mask` (PNG) or `polygon` onto your canvas
|
||||
```
|
||||
|
||||
CORS is enabled for `*`, so browser clients work out of the box — lock this down for production.
|
||||
|
||||
## Request/response reference
|
||||
|
||||
**Request body** (JSON):
|
||||
```json
|
||||
{
|
||||
"image": "<base64 PNG/JPG>", // OR image_path OR image_url
|
||||
"bbox": [x1, y1, x2, y2], // absolute pixels, xyxy
|
||||
"multimask_output": false, // optional, /segment/box only — returns 3 candidates
|
||||
"return_polygon": true,
|
||||
"return_rle": true
|
||||
}
|
||||
```
|
||||
|
||||
**Response**:
|
||||
```json
|
||||
{
|
||||
"mask": "<base64 PNG, 0/255 single channel>",
|
||||
"score": 0.97,
|
||||
"bbox_out": [x1, y1, x2, y2],
|
||||
"rle": {"size": [H, W], "counts": "..."},
|
||||
"polygon": [[x, y], ...],
|
||||
"extra_masks": null
|
||||
}
|
||||
```
|
||||
|
||||
## Notes on the two endpoints
|
||||
|
||||
The SAM 3 reference repo's `Sam3Processor.set_box_prompt` is actually a **concept/exemplar** prompt — it finds *all* objects similar to what's in the box. For a bbox-to-single-mask annotation workflow you specifically want the **tracker** (`Sam3Tracker`), which is SAM 3's SAM2-compatible PVS head. That's what `/segment/box` uses when `transformers` is installed.
|
||||
|
||||
If `transformers` isn't available, `/segment/box` falls back to the exemplar predictor and picks the returned instance whose predicted bbox has the highest IoU with your input — this usually gives you the right object but can pick a similar neighbor in cluttered scenes. Install `transformers` for best results.
|
||||
|
||||
## Performance tips
|
||||
- Keep the process alive; model loading is the slow part (~10–20s).
|
||||
- For batch annotation, call `/segment/box` sequentially — `set_image` caches the image embedding in the session state.
|
||||
- On a single RTX 4090, per-request latency for a 1024px image is typically **~150–300ms** once warm.
|
||||
- For large images, consider resizing on the client to ≤2048px on the long side; SAM 3 resizes internally anyway.
|
||||
50
files/client_example.py
Normal file
50
files/client_example.py
Normal file
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
Example client: how your annotation tool calls the SAM 3 server.
|
||||
|
||||
Run the server first:
|
||||
python sam3_server.py --host 0.0.0.0 --port 8000
|
||||
|
||||
Then, from your annotation tool, POST an image + bbox and render the mask.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
SERVER = "http://localhost:8000"
|
||||
|
||||
|
||||
def segment_with_bbox(image_path: str, bbox_xyxy: list[float]):
|
||||
with open(image_path, "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode("ascii")
|
||||
|
||||
r = requests.post(
|
||||
f"{SERVER}/segment/box",
|
||||
json={
|
||||
"image": img_b64,
|
||||
"bbox": bbox_xyxy, # [x1, y1, x2, y2] in pixels
|
||||
"multimask_output": False,
|
||||
"return_polygon": True, # get contour for polygon-based annotation
|
||||
"return_rle": True, # get COCO RLE for COCO-format export
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
# Decode mask (single-channel PNG, 0 or 255)
|
||||
mask_png = base64.b64decode(data["mask"])
|
||||
mask = Image.open(io.BytesIO(mask_png))
|
||||
mask.save("mask_out.png")
|
||||
|
||||
print(f"score: {data['score']:.3f}")
|
||||
print(f"bbox: {data['bbox_out']}")
|
||||
print(f"polygon: {len(data['polygon']) if data['polygon'] else 0} points")
|
||||
print(f"rle: size={data['rle']['size']} (counts length={len(str(data['rle']['counts']))})")
|
||||
return data
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example: segment the object inside this bbox
|
||||
segment_with_bbox("example.jpg", [120, 80, 540, 420])
|
||||
381
files/sam3_server.py
Normal file
381
files/sam3_server.py
Normal file
@@ -0,0 +1,381 @@
|
||||
"""
|
||||
SAM 3 local inference server for annotation tools.
|
||||
|
||||
Exposes a small FastAPI service that takes an image + bbox and returns a
|
||||
segmentation mask. Designed for the classic annotation workflow:
|
||||
user draws a box around ONE object -> server returns the mask for that object.
|
||||
|
||||
Two endpoints are provided:
|
||||
|
||||
POST /segment/box
|
||||
Uses Sam3Tracker (the SAM2-style Promptable Visual Segmentation head).
|
||||
Given ONE bbox, returns ONE mask for the object inside that box.
|
||||
This is what you want for 99% of annotation-tool use cases.
|
||||
|
||||
POST /segment/exemplar
|
||||
Uses the full SAM 3 image predictor with a box as a visual "exemplar".
|
||||
Given one bbox, returns masks for ALL similar objects in the image.
|
||||
Useful when you want to label one instance and auto-propagate.
|
||||
|
||||
Run:
|
||||
python sam3_server.py --host 0.0.0.0 --port 8000
|
||||
|
||||
Request format (JSON):
|
||||
{
|
||||
"image": "<base64-encoded image>" # OR
|
||||
"image_url": "http://..." # OR
|
||||
"image_path": "/absolute/path.jpg",
|
||||
"bbox": [x1, y1, x2, y2], # absolute pixel coords (xyxy)
|
||||
"multimask_output": false # optional, /segment/box only
|
||||
}
|
||||
|
||||
Response:
|
||||
{
|
||||
"mask": "<base64 PNG, single-channel 0/255>",
|
||||
"rle": {"size":[H,W], "counts":"..."}, # COCO RLE, handy for COCO-format tools
|
||||
"polygon": [[x,y], [x,y], ...], # largest contour, optional
|
||||
"score": 0.97,
|
||||
"bbox_out": [x1,y1,x2,y2]
|
||||
}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# ----- SAM 3 imports ---------------------------------------------------------
|
||||
# The SAM 3 image predictor (concept / exemplar segmentation):
|
||||
from sam3.model_builder import build_sam3_image_model
|
||||
from sam3.model.sam3_image_processor import Sam3Processor
|
||||
|
||||
# The SAM 3 tracker (SAM2-style single-instance PVS on images):
|
||||
# This is exposed in the HF `transformers` integration as Sam3TrackerModel,
|
||||
# and in the reference repo via the image predictor for single-box prompts.
|
||||
# We use HuggingFace transformers for the tracker because its API is clean
|
||||
# and documented for bbox -> single mask. Fall back to the repo's predictor
|
||||
# if transformers isn't available.
|
||||
try:
|
||||
from transformers import Sam3TrackerModel, Sam3TrackerProcessor # type: ignore
|
||||
HAS_TRACKER = True
|
||||
except Exception:
|
||||
HAS_TRACKER = False
|
||||
|
||||
log = logging.getLogger("sam3-server")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Request / response schemas
|
||||
# ============================================================================
|
||||
|
||||
class SegmentRequest(BaseModel):
|
||||
# Exactly one of these three must be provided.
|
||||
image: Optional[str] = Field(None, description="Base64-encoded image bytes.")
|
||||
image_url: Optional[str] = None
|
||||
image_path: Optional[str] = None
|
||||
|
||||
bbox: List[float] = Field(..., description="Box in absolute pixels, [x1,y1,x2,y2].")
|
||||
multimask_output: bool = Field(
|
||||
False,
|
||||
description="If true (tracker only), return the 3 candidate masks instead of picking the highest-scoring one.",
|
||||
)
|
||||
return_polygon: bool = Field(True, description="Also return the largest contour as [[x,y], ...].")
|
||||
return_rle: bool = Field(True, description="Also return COCO RLE encoding of the mask.")
|
||||
|
||||
|
||||
class SegmentResponse(BaseModel):
|
||||
mask: str # base64 PNG
|
||||
score: float
|
||||
bbox_out: List[float]
|
||||
rle: Optional[dict] = None
|
||||
polygon: Optional[List[List[float]]] = None
|
||||
extra_masks: Optional[List[str]] = None # only populated when multimask_output=True
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Model holder (loaded once at startup)
|
||||
# ============================================================================
|
||||
|
||||
class SAM3Models:
|
||||
def __init__(self, device: str, checkpoint_path: Optional[str] = None):
|
||||
self.device = device
|
||||
|
||||
log.info("Loading SAM 3 image model (concept/exemplar predictor)...")
|
||||
if checkpoint_path:
|
||||
self.image_model = build_sam3_image_model(checkpoint_path=checkpoint_path)
|
||||
else:
|
||||
self.image_model = build_sam3_image_model()
|
||||
self.image_model.to(device).eval()
|
||||
self.image_processor = Sam3Processor(self.image_model)
|
||||
|
||||
if HAS_TRACKER:
|
||||
log.info("Loading SAM 3 tracker (single-instance PVS) from HuggingFace...")
|
||||
# facebook/sam3 exposes the tracker weights under this name
|
||||
self.tracker_processor = Sam3TrackerProcessor.from_pretrained("facebook/sam3")
|
||||
self.tracker_model = Sam3TrackerModel.from_pretrained("facebook/sam3").to(device).eval()
|
||||
else:
|
||||
log.warning(
|
||||
"transformers not available or missing Sam3Tracker classes. "
|
||||
"The /segment/box endpoint will fall back to the image predictor "
|
||||
"(still works, but treats the box as an exemplar)."
|
||||
)
|
||||
self.tracker_processor = None
|
||||
self.tracker_model = None
|
||||
|
||||
|
||||
MODELS: Optional[SAM3Models] = None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helpers
|
||||
# ============================================================================
|
||||
|
||||
def _load_image(req: SegmentRequest) -> Image.Image:
|
||||
if req.image:
|
||||
raw = base64.b64decode(req.image)
|
||||
return Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
if req.image_path:
|
||||
return Image.open(req.image_path).convert("RGB")
|
||||
if req.image_url:
|
||||
import requests
|
||||
r = requests.get(req.image_url, timeout=20)
|
||||
r.raise_for_status()
|
||||
return Image.open(io.BytesIO(r.content)).convert("RGB")
|
||||
raise HTTPException(400, "Provide one of `image` (base64), `image_path`, or `image_url`.")
|
||||
|
||||
|
||||
def _mask_to_png_b64(mask: np.ndarray) -> str:
|
||||
"""mask: HxW bool/uint8 -> base64-encoded PNG string."""
|
||||
if mask.dtype != np.uint8:
|
||||
mask = (mask.astype(np.uint8) * 255)
|
||||
img = Image.fromarray(mask, mode="L")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG", optimize=True)
|
||||
return base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def _mask_to_rle(mask: np.ndarray) -> dict:
|
||||
"""COCO-style RLE. Uses pycocotools if available, else a pure-python fallback."""
|
||||
mask = (mask > 0).astype(np.uint8)
|
||||
try:
|
||||
from pycocotools import mask as mutil # type: ignore
|
||||
rle = mutil.encode(np.asfortranarray(mask))
|
||||
rle["counts"] = rle["counts"].decode("ascii")
|
||||
return rle
|
||||
except Exception:
|
||||
# Pure python fallback (uncompressed)
|
||||
h, w = mask.shape
|
||||
flat = mask.T.flatten() # column-major, like COCO
|
||||
counts: List[int] = []
|
||||
last = 0
|
||||
run = 0
|
||||
for v in flat:
|
||||
if v == last:
|
||||
run += 1
|
||||
else:
|
||||
counts.append(run)
|
||||
run = 1
|
||||
last = v
|
||||
counts.append(run)
|
||||
return {"size": [h, w], "counts": counts}
|
||||
|
||||
|
||||
def _mask_to_polygon(mask: np.ndarray) -> Optional[List[List[float]]]:
|
||||
"""Largest-contour polygon, approximated. Returns None if mask is empty."""
|
||||
try:
|
||||
import cv2 # type: ignore
|
||||
except Exception:
|
||||
return None
|
||||
m = (mask > 0).astype(np.uint8) * 255
|
||||
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
|
||||
if not contours:
|
||||
return None
|
||||
c = max(contours, key=cv2.contourArea)
|
||||
# Simplify a bit — saves bytes and is friendlier for annotation tools
|
||||
epsilon = 0.0015 * cv2.arcLength(c, True)
|
||||
c = cv2.approxPolyDP(c, epsilon, True)
|
||||
return [[float(pt[0][0]), float(pt[0][1])] for pt in c]
|
||||
|
||||
|
||||
def _mask_bbox(mask: np.ndarray) -> List[float]:
|
||||
ys, xs = np.where(mask > 0)
|
||||
if len(xs) == 0:
|
||||
return [0.0, 0.0, 0.0, 0.0]
|
||||
return [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())]
|
||||
|
||||
|
||||
def _build_response(mask: np.ndarray, score: float, req: SegmentRequest,
|
||||
extra_masks: Optional[List[np.ndarray]] = None) -> SegmentResponse:
|
||||
return SegmentResponse(
|
||||
mask=_mask_to_png_b64(mask),
|
||||
score=float(score),
|
||||
bbox_out=_mask_bbox(mask),
|
||||
rle=_mask_to_rle(mask) if req.return_rle else None,
|
||||
polygon=_mask_to_polygon(mask) if req.return_polygon else None,
|
||||
extra_masks=[_mask_to_png_b64(m) for m in extra_masks] if extra_masks else None,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# FastAPI app
|
||||
# ============================================================================
|
||||
|
||||
app = FastAPI(title="SAM 3 local server", version="1.0")
|
||||
|
||||
# Allow local annotation tool UIs to call us from the browser
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {
|
||||
"ok": True,
|
||||
"device": MODELS.device if MODELS else "unloaded",
|
||||
"tracker_available": HAS_TRACKER,
|
||||
}
|
||||
|
||||
|
||||
@app.post("/segment/box", response_model=SegmentResponse)
|
||||
@torch.inference_mode()
|
||||
def segment_box(req: SegmentRequest):
|
||||
"""
|
||||
Annotation-tool workflow:
|
||||
user draws a bbox around ONE object -> return ONE mask for that object.
|
||||
|
||||
Uses the SAM3 tracker (SAM2-style PVS) for a clean single-instance result.
|
||||
"""
|
||||
if MODELS is None:
|
||||
raise HTTPException(500, "Models not loaded.")
|
||||
image = _load_image(req)
|
||||
x1, y1, x2, y2 = req.bbox
|
||||
|
||||
# ---- Preferred path: HuggingFace Sam3Tracker -------------------------
|
||||
if HAS_TRACKER and MODELS.tracker_model is not None:
|
||||
inputs = MODELS.tracker_processor(
|
||||
images=image,
|
||||
input_boxes=[[[x1, y1, x2, y2]]], # batch[1] x obj[1] x 4
|
||||
return_tensors="pt",
|
||||
).to(MODELS.device)
|
||||
|
||||
outputs = MODELS.tracker_model(**inputs, multimask_output=req.multimask_output)
|
||||
|
||||
# Post-process to image resolution, binary masks
|
||||
masks = MODELS.tracker_processor.post_process_masks(
|
||||
outputs.pred_masks.cpu(),
|
||||
original_sizes=inputs["original_sizes"].cpu(),
|
||||
reshaped_input_sizes=inputs["reshaped_input_sizes"].cpu(),
|
||||
)[0] # shape: (num_masks, H, W) tensor
|
||||
scores = outputs.iou_scores[0].cpu().numpy().ravel()
|
||||
|
||||
masks_np = masks.numpy().astype(bool)
|
||||
if req.multimask_output and masks_np.shape[0] > 1:
|
||||
best = int(np.argmax(scores))
|
||||
extras = [m for i, m in enumerate(masks_np) if i != best]
|
||||
return _build_response(masks_np[best], scores[best], req, extra_masks=extras)
|
||||
|
||||
best = int(np.argmax(scores)) if masks_np.shape[0] > 1 else 0
|
||||
return _build_response(masks_np[best], float(scores[best]), req)
|
||||
|
||||
# ---- Fallback: use image predictor with box as exemplar, pick the
|
||||
# detection whose predicted box has highest IoU with the user bbox.
|
||||
state = MODELS.image_processor.set_image(image)
|
||||
out = MODELS.image_processor.set_box_prompt(
|
||||
state=state,
|
||||
boxes=[[x1, y1, x2 - x1, y2 - y1]], # xywh as documented
|
||||
box_labels=[1],
|
||||
)
|
||||
masks = out["masks"]
|
||||
boxes = out["boxes"]
|
||||
scores = out["scores"]
|
||||
if len(masks) == 0:
|
||||
raise HTTPException(404, "No segmentation produced for the given bbox.")
|
||||
|
||||
# Pick mask whose predicted bbox most overlaps user bbox
|
||||
def _iou(a, b):
|
||||
ax1, ay1, ax2, ay2 = a
|
||||
bx1, by1, bx2, by2 = b
|
||||
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
||||
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
||||
iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
|
||||
inter = iw * ih
|
||||
ua = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - inter
|
||||
return inter / ua if ua > 0 else 0.0
|
||||
|
||||
user_box = [x1, y1, x2, y2]
|
||||
ious = [_iou(user_box, list(map(float, b))) for b in boxes]
|
||||
idx = int(np.argmax(ious))
|
||||
mask_np = np.asarray(masks[idx]).astype(bool)
|
||||
return _build_response(mask_np, float(scores[idx]), req)
|
||||
|
||||
|
||||
@app.post("/segment/exemplar", response_model=SegmentResponse)
|
||||
@torch.inference_mode()
|
||||
def segment_exemplar(req: SegmentRequest):
|
||||
"""
|
||||
Auto-propagation workflow:
|
||||
user draws a box around ONE instance -> return a merged mask of ALL
|
||||
similar instances found in the image. `extra_masks` contains each
|
||||
individual instance separately.
|
||||
"""
|
||||
if MODELS is None:
|
||||
raise HTTPException(500, "Models not loaded.")
|
||||
image = _load_image(req)
|
||||
x1, y1, x2, y2 = req.bbox
|
||||
|
||||
state = MODELS.image_processor.set_image(image)
|
||||
out = MODELS.image_processor.set_box_prompt(
|
||||
state=state,
|
||||
boxes=[[x1, y1, x2 - x1, y2 - y1]],
|
||||
box_labels=[1],
|
||||
)
|
||||
masks = out["masks"]
|
||||
scores = out["scores"]
|
||||
if len(masks) == 0:
|
||||
raise HTTPException(404, "No instances found.")
|
||||
|
||||
masks_np = [np.asarray(m).astype(bool) for m in masks]
|
||||
merged = np.zeros_like(masks_np[0], dtype=bool)
|
||||
for m in masks_np:
|
||||
merged |= m
|
||||
best_score = float(np.max(scores))
|
||||
return _build_response(merged, best_score, req, extra_masks=masks_np)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Entrypoint
|
||||
# ============================================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--checkpoint", default=None,
|
||||
help="Optional path to sam3 checkpoint (otherwise uses HF default).")
|
||||
args = parser.parse_args()
|
||||
|
||||
global MODELS
|
||||
MODELS = SAM3Models(device=args.device, checkpoint_path=args.checkpoint)
|
||||
log.info("SAM 3 server ready on %s:%d (device=%s)", args.host, args.port, args.device)
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
4
requirements.txt
Normal file
4
requirements.txt
Normal file
@@ -0,0 +1,4 @@
|
||||
opencv-python
|
||||
Pillow
|
||||
numpy
|
||||
|
||||
14
sam-tool-tauri/.gitignore
vendored
Normal file
14
sam-tool-tauri/.gitignore
vendored
Normal file
@@ -0,0 +1,14 @@
|
||||
node_modules
|
||||
dist
|
||||
dist-ssr
|
||||
*.local
|
||||
|
||||
.vite
|
||||
|
||||
# Editor
|
||||
.vscode/*
|
||||
!.vscode/extensions.json
|
||||
.idea
|
||||
|
||||
# System
|
||||
.DS_Store
|
||||
12
sam-tool-tauri/index.html
Normal file
12
sam-tool-tauri/index.html
Normal file
@@ -0,0 +1,12 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>SAM Tool</title>
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
<script type="module" src="/src/main.tsx"></script>
|
||||
</body>
|
||||
</html>
|
||||
2149
sam-tool-tauri/package-lock.json
generated
Normal file
2149
sam-tool-tauri/package-lock.json
generated
Normal file
File diff suppressed because it is too large
Load Diff
27
sam-tool-tauri/package.json
Normal file
27
sam-tool-tauri/package.json
Normal file
@@ -0,0 +1,27 @@
|
||||
{
|
||||
"name": "sam-tool-tauri",
|
||||
"private": true,
|
||||
"version": "0.1.0",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "tsc && vite build",
|
||||
"preview": "vite preview",
|
||||
"tauri": "tauri"
|
||||
},
|
||||
"dependencies": {
|
||||
"@tauri-apps/api": "^2.1.1",
|
||||
"@tauri-apps/plugin-dialog": "^2.2.0",
|
||||
"@tauri-apps/plugin-fs": "^2.2.0",
|
||||
"react": "^18.3.1",
|
||||
"react-dom": "^18.3.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@tauri-apps/cli": "^2.1.0",
|
||||
"@types/react": "^18.3.12",
|
||||
"@types/react-dom": "^18.3.1",
|
||||
"@vitejs/plugin-react": "^4.3.4",
|
||||
"typescript": "^5.7.2",
|
||||
"vite": "^6.0.3"
|
||||
}
|
||||
}
|
||||
2
sam-tool-tauri/src-tauri/.gitignore
vendored
Normal file
2
sam-tool-tauri/src-tauri/.gitignore
vendored
Normal file
@@ -0,0 +1,2 @@
|
||||
target
|
||||
gen/schemas
|
||||
5365
sam-tool-tauri/src-tauri/Cargo.lock
generated
Normal file
5365
sam-tool-tauri/src-tauri/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load Diff
35
sam-tool-tauri/src-tauri/Cargo.toml
Normal file
35
sam-tool-tauri/src-tauri/Cargo.toml
Normal file
@@ -0,0 +1,35 @@
|
||||
[package]
|
||||
name = "sam-tool-tauri"
|
||||
version = "0.1.0"
|
||||
description = "SAM Tool - video frame extraction and annotation"
|
||||
edition = "2021"
|
||||
rust-version = "1.77"
|
||||
|
||||
[lib]
|
||||
name = "sam_tool_tauri_lib"
|
||||
crate-type = ["staticlib", "cdylib", "rlib"]
|
||||
|
||||
[build-dependencies]
|
||||
tauri-build = { version = "2", features = [] }
|
||||
|
||||
[dependencies]
|
||||
tauri = { version = "2", features = [] }
|
||||
tauri-plugin-dialog = "2"
|
||||
tauri-plugin-fs = "2"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
chrono = { version = "0.4", default-features = false, features = ["clock", "serde"] }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "multipart"] }
|
||||
base64 = "0.22"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
|
||||
|
||||
[profile.dev]
|
||||
incremental = true
|
||||
|
||||
[profile.release]
|
||||
codegen-units = 1
|
||||
lto = true
|
||||
opt-level = "s"
|
||||
panic = "abort"
|
||||
strip = true
|
||||
3
sam-tool-tauri/src-tauri/build.rs
Normal file
3
sam-tool-tauri/src-tauri/build.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
fn main() {
|
||||
tauri_build::build()
|
||||
}
|
||||
11
sam-tool-tauri/src-tauri/capabilities/default.json
Normal file
11
sam-tool-tauri/src-tauri/capabilities/default.json
Normal file
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"$schema": "../gen/schemas/desktop-schema.json",
|
||||
"identifier": "default",
|
||||
"description": "Default capabilities for the main window",
|
||||
"windows": ["main"],
|
||||
"permissions": [
|
||||
"core:default",
|
||||
"dialog:default",
|
||||
"fs:default"
|
||||
]
|
||||
}
|
||||
BIN
sam-tool-tauri/src-tauri/icons/128x128.png
Normal file
BIN
sam-tool-tauri/src-tauri/icons/128x128.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.9 KiB |
BIN
sam-tool-tauri/src-tauri/icons/128x128@2x.png
Normal file
BIN
sam-tool-tauri/src-tauri/icons/128x128@2x.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 3.3 KiB |
BIN
sam-tool-tauri/src-tauri/icons/32x32.png
Normal file
BIN
sam-tool-tauri/src-tauri/icons/32x32.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 663 B |
BIN
sam-tool-tauri/src-tauri/icons/icon.png
Normal file
BIN
sam-tool-tauri/src-tauri/icons/icon.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
169
sam-tool-tauri/src-tauri/src/annotate.rs
Normal file
169
sam-tool-tauri/src-tauri/src/annotate.rs
Normal file
@@ -0,0 +1,169 @@
|
||||
use crate::annotation::{AnnotationRecord, ImageAnnotations};
|
||||
use serde::Serialize;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
#[derive(Serialize, Clone)]
|
||||
pub struct ImageEntry {
|
||||
pub filename: String,
|
||||
pub has_annotations: bool,
|
||||
pub annotation_count: u32,
|
||||
pub last_updated: Option<String>,
|
||||
}
|
||||
|
||||
fn updated_at_of(a: &AnnotationRecord) -> &str {
|
||||
match a {
|
||||
AnnotationRecord::Bbox { updated_at, .. } => updated_at,
|
||||
AnnotationRecord::Polygon { updated_at, .. } => updated_at,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_image(name: &str) -> bool {
|
||||
let lower = name.to_lowercase();
|
||||
lower.ends_with(".jpg") || lower.ends_with(".jpeg") || lower.ends_with(".png")
|
||||
}
|
||||
|
||||
fn sibling_json(image_path: &Path) -> PathBuf {
|
||||
image_path.with_extension("json")
|
||||
}
|
||||
|
||||
pub fn list_images(folder: &Path) -> Result<Vec<ImageEntry>, String> {
|
||||
if !folder.exists() {
|
||||
return Err(format!("folder not found: {}", folder.display()));
|
||||
}
|
||||
if !folder.is_dir() {
|
||||
return Err(format!("not a directory: {}", folder.display()));
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
let entries = std::fs::read_dir(folder)
|
||||
.map_err(|e| format!("read_dir {}: {e}", folder.display()))?;
|
||||
for entry in entries {
|
||||
let entry = entry.map_err(|e| format!("dirent: {e}"))?;
|
||||
let path = entry.path();
|
||||
if !path.is_file() {
|
||||
continue;
|
||||
}
|
||||
let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
if !is_image(name) {
|
||||
continue;
|
||||
}
|
||||
let json_path = sibling_json(&path);
|
||||
let (annotation_count, last_updated) = if json_path.exists() {
|
||||
match std::fs::read_to_string(&json_path)
|
||||
.ok()
|
||||
.and_then(|t| serde_json::from_str::<ImageAnnotations>(&t).ok())
|
||||
{
|
||||
Some(ia) => {
|
||||
let latest = ia
|
||||
.annotations
|
||||
.iter()
|
||||
.map(|a| updated_at_of(a).to_string())
|
||||
.max();
|
||||
(ia.annotations.len() as u32, latest)
|
||||
}
|
||||
None => (0, None),
|
||||
}
|
||||
} else {
|
||||
(0, None)
|
||||
};
|
||||
out.push(ImageEntry {
|
||||
filename: name.to_string(),
|
||||
has_annotations: annotation_count > 0,
|
||||
annotation_count,
|
||||
last_updated,
|
||||
});
|
||||
}
|
||||
out.sort_by(|a, b| a.filename.cmp(&b.filename));
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn read_image_bytes(path: &Path) -> Result<Vec<u8>, String> {
|
||||
if !path.exists() {
|
||||
return Err(format!("file not found: {}", path.display()));
|
||||
}
|
||||
std::fs::read(path).map_err(|e| format!("read {}: {e}", path.display()))
|
||||
}
|
||||
|
||||
pub fn load_image_annotations(
|
||||
image_path: &Path,
|
||||
) -> Result<Option<ImageAnnotations>, String> {
|
||||
let json_path = sibling_json(image_path);
|
||||
if !json_path.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let text = std::fs::read_to_string(&json_path)
|
||||
.map_err(|e| format!("read {}: {e}", json_path.display()))?;
|
||||
let ia: ImageAnnotations =
|
||||
serde_json::from_str(&text).map_err(|e| format!("parse sibling JSON: {e}"))?;
|
||||
Ok(Some(ia))
|
||||
}
|
||||
|
||||
pub fn load_classes(folder: &Path) -> Result<Vec<String>, String> {
|
||||
let path = folder.join("classes.txt");
|
||||
if !path.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let text = std::fs::read_to_string(&path)
|
||||
.map_err(|e| format!("read {}: {e}", path.display()))?;
|
||||
Ok(text
|
||||
.lines()
|
||||
.map(|l| l.trim().to_string())
|
||||
.filter(|l| !l.is_empty())
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Delete a batch of image files along with their sibling annotation JSONs.
|
||||
/// Returns the count of images actually removed. Missing images are skipped;
|
||||
/// any hard I/O error aborts with what succeeded so far reflected in the log.
|
||||
pub fn delete_images(paths: &[String]) -> Result<u32, String> {
|
||||
let mut removed: u32 = 0;
|
||||
for p in paths {
|
||||
let path = Path::new(p);
|
||||
if !path.is_file() {
|
||||
continue;
|
||||
}
|
||||
// Only allow image files — guard against a stray path that would
|
||||
// wipe a JSON or classes.txt by name collision.
|
||||
let name = path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or_default();
|
||||
if !is_image(name) {
|
||||
return Err(format!("refusing to delete non-image: {}", path.display()));
|
||||
}
|
||||
std::fs::remove_file(path)
|
||||
.map_err(|e| format!("remove {}: {e}", path.display()))?;
|
||||
let json_path = sibling_json(path);
|
||||
if json_path.exists() {
|
||||
let _ = std::fs::remove_file(&json_path);
|
||||
}
|
||||
removed += 1;
|
||||
}
|
||||
Ok(removed)
|
||||
}
|
||||
|
||||
pub fn append_class(folder: &Path, name: &str) -> Result<Vec<String>, String> {
|
||||
let trimmed = name.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err("class name cannot be empty".into());
|
||||
}
|
||||
if trimmed.len() > 96 {
|
||||
return Err("class name too long".into());
|
||||
}
|
||||
std::fs::create_dir_all(folder)
|
||||
.map_err(|e| format!("mkdir {}: {e}", folder.display()))?;
|
||||
let mut classes = load_classes(folder)?;
|
||||
if classes.iter().any(|n| n == trimmed) {
|
||||
return Err(format!("class '{trimmed}' already exists"));
|
||||
}
|
||||
classes.push(trimmed.to_string());
|
||||
let mut text = String::new();
|
||||
for c in &classes {
|
||||
text.push_str(c);
|
||||
text.push('\n');
|
||||
}
|
||||
let path = folder.join("classes.txt");
|
||||
std::fs::write(&path, text).map_err(|e| format!("write {}: {e}", path.display()))?;
|
||||
Ok(classes)
|
||||
}
|
||||
43
sam-tool-tauri/src-tauri/src/annotation.rs
Normal file
43
sam-tool-tauri/src-tauri/src/annotation.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone)]
|
||||
pub struct ImageAnnotations {
|
||||
pub version: u32,
|
||||
pub image: String,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
/// Absolute path to the source video the frame was extracted from.
|
||||
/// Optional so older JSONs without this field still round-trip cleanly.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub source_video: Option<String>,
|
||||
pub annotations: Vec<AnnotationRecord>,
|
||||
}
|
||||
|
||||
/// Internally-tagged enum. Both variants share id/class/author/timestamps and
|
||||
/// differ only in the shape data (bbox | points). Using a single enum with
|
||||
/// `#[serde(tag = "type")]` gives clean round-trip serde (flatten + tagged
|
||||
/// enums together don't deserialize reliably).
|
||||
#[derive(Serialize, Deserialize, Clone)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum AnnotationRecord {
|
||||
#[serde(rename = "bbox")]
|
||||
Bbox {
|
||||
id: String,
|
||||
class_id: i64,
|
||||
class_name: String,
|
||||
bbox: [f64; 4],
|
||||
author: String,
|
||||
created_at: String,
|
||||
updated_at: String,
|
||||
},
|
||||
#[serde(rename = "polygon")]
|
||||
Polygon {
|
||||
id: String,
|
||||
class_id: i64,
|
||||
class_name: String,
|
||||
points: Vec<[f64; 2]>,
|
||||
author: String,
|
||||
created_at: String,
|
||||
updated_at: String,
|
||||
},
|
||||
}
|
||||
359
sam-tool-tauri/src-tauri/src/coco.rs
Normal file
359
sam-tool-tauri/src-tauri/src/coco.rs
Normal file
@@ -0,0 +1,359 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::path::Path;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CocoFile {
|
||||
images: Vec<CocoImage>,
|
||||
categories: Vec<CocoCategory>,
|
||||
annotations: Vec<CocoAnnotation>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CocoImage {
|
||||
id: i64,
|
||||
#[serde(default)]
|
||||
file_name: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CocoCategory {
|
||||
id: i64,
|
||||
name: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CocoAnnotation {
|
||||
image_id: i64,
|
||||
category_id: i64,
|
||||
#[serde(default)]
|
||||
bbox: Option<Vec<f64>>,
|
||||
#[serde(default)]
|
||||
segmentation: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum Shape {
|
||||
#[serde(rename = "bbox")]
|
||||
Bbox {
|
||||
class_id: i64,
|
||||
class_name: String,
|
||||
bbox: [f64; 4],
|
||||
},
|
||||
#[serde(rename = "polygon")]
|
||||
Polygon {
|
||||
class_id: i64,
|
||||
class_name: String,
|
||||
points: Vec<[f64; 2]>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone)]
|
||||
pub struct Category {
|
||||
pub id: i64,
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone, Default)]
|
||||
pub struct ImportStats {
|
||||
pub total_annotations: usize,
|
||||
pub frame_count: usize,
|
||||
pub categories: usize,
|
||||
pub skipped_rle: usize,
|
||||
pub skipped_no_image: usize,
|
||||
/// Filled in by the import — "coco" or "sam".
|
||||
pub source_format: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone)]
|
||||
pub struct ImportedCoco {
|
||||
pub categories: Vec<Category>,
|
||||
/// Map of frame_idx (as string, per JSON) -> shapes on that frame.
|
||||
pub frames: HashMap<String, Vec<Shape>>,
|
||||
pub stats: ImportStats,
|
||||
}
|
||||
|
||||
enum Format {
|
||||
LabelStudioCoco,
|
||||
SamAnno,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
fn detect_format(v: &serde_json::Value) -> Format {
|
||||
let Some(obj) = v.as_object() else {
|
||||
return Format::Unknown;
|
||||
};
|
||||
if obj.contains_key("images") && obj.contains_key("annotations") {
|
||||
return Format::LabelStudioCoco;
|
||||
}
|
||||
if !obj.is_empty()
|
||||
&& obj
|
||||
.keys()
|
||||
.all(|k| !k.is_empty() && k.chars().all(|c| c.is_ascii_digit()))
|
||||
&& obj.values().all(|v| v.is_object())
|
||||
{
|
||||
return Format::SamAnno;
|
||||
}
|
||||
Format::Unknown
|
||||
}
|
||||
|
||||
pub fn import(path: &Path) -> Result<ImportedCoco, String> {
|
||||
let text = std::fs::read_to_string(path)
|
||||
.map_err(|e| format!("read {}: {e}", path.display()))?;
|
||||
let root: serde_json::Value =
|
||||
serde_json::from_str(&text).map_err(|e| format!("parse JSON: {e}"))?;
|
||||
|
||||
match detect_format(&root) {
|
||||
Format::LabelStudioCoco => {
|
||||
let coco: CocoFile = serde_json::from_value(root)
|
||||
.map_err(|e| format!("parse COCO: {e}"))?;
|
||||
Ok(build_from_coco(coco))
|
||||
}
|
||||
Format::SamAnno => build_from_sam_anno(&root),
|
||||
Format::Unknown => Err(
|
||||
"unrecognized annotation format — expected Label Studio COCO or SAM anno JSON"
|
||||
.into(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Label Studio COCO -----------------------------------------------------
|
||||
|
||||
fn build_from_coco(coco: CocoFile) -> ImportedCoco {
|
||||
let mut image_frame: HashMap<i64, u64> = HashMap::with_capacity(coco.images.len());
|
||||
for img in &coco.images {
|
||||
let frame_idx =
|
||||
trailing_int(&img.file_name).unwrap_or_else(|| img.id.max(0) as u64);
|
||||
image_frame.insert(img.id, frame_idx);
|
||||
}
|
||||
|
||||
let cat_name: HashMap<i64, String> = coco
|
||||
.categories
|
||||
.iter()
|
||||
.map(|c| (c.id, c.name.clone()))
|
||||
.collect();
|
||||
|
||||
let mut frames: HashMap<String, Vec<Shape>> = HashMap::new();
|
||||
let mut stats = ImportStats {
|
||||
categories: coco.categories.len(),
|
||||
source_format: "coco".into(),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
for ann in &coco.annotations {
|
||||
stats.total_annotations += 1;
|
||||
let Some(&frame_idx) = image_frame.get(&ann.image_id) else {
|
||||
stats.skipped_no_image += 1;
|
||||
continue;
|
||||
};
|
||||
let class_name = cat_name
|
||||
.get(&ann.category_id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("class_{}", ann.category_id));
|
||||
let key = frame_idx.to_string();
|
||||
let bucket = frames.entry(key).or_default();
|
||||
|
||||
if let Some(seg) = &ann.segmentation {
|
||||
match seg {
|
||||
serde_json::Value::Array(rings) => {
|
||||
for ring in rings {
|
||||
if let Some(flat) = ring.as_array() {
|
||||
let points = parse_flat_polygon(flat);
|
||||
if points.len() >= 3 {
|
||||
bucket.push(Shape::Polygon {
|
||||
class_id: ann.category_id,
|
||||
class_name: class_name.clone(),
|
||||
points,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(_) => {
|
||||
stats.skipped_rle += 1;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(bbox) = &ann.bbox {
|
||||
if bbox.len() == 4 {
|
||||
bucket.push(Shape::Bbox {
|
||||
class_id: ann.category_id,
|
||||
class_name,
|
||||
bbox: [bbox[0], bbox[1], bbox[2], bbox[3]],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
stats.frame_count = frames.len();
|
||||
let categories = coco
|
||||
.categories
|
||||
.into_iter()
|
||||
.map(|c| Category {
|
||||
id: c.id,
|
||||
name: c.name,
|
||||
})
|
||||
.collect();
|
||||
|
||||
ImportedCoco {
|
||||
categories,
|
||||
frames,
|
||||
stats,
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Custom SAM annotation format ------------------------------------------
|
||||
//
|
||||
// { "<frame_idx>": { "<class_name>": [[id, [tl_x,tl_y], [br_x,br_y],
|
||||
// [[[x,y], ...], ...rings]]] } }
|
||||
// Bbox is redundant with the polygon's tight bounding box, so we emit only
|
||||
// polygons. Class IDs are assigned 0-indexed by sorted class name.
|
||||
fn build_from_sam_anno(root: &serde_json::Value) -> Result<ImportedCoco, String> {
|
||||
let obj = root
|
||||
.as_object()
|
||||
.ok_or_else(|| "SAM anno root is not an object".to_string())?;
|
||||
|
||||
let mut class_names = BTreeSet::<String>::new();
|
||||
let mut frames: HashMap<String, Vec<Shape>> = HashMap::new();
|
||||
let mut stats = ImportStats {
|
||||
source_format: "sam".into(),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
for (frame_key, frame_val) in obj {
|
||||
let Some(cls_obj) = frame_val.as_object() else {
|
||||
continue;
|
||||
};
|
||||
let mut shapes: Vec<Shape> = Vec::new();
|
||||
for (class_name, items) in cls_obj {
|
||||
class_names.insert(class_name.clone());
|
||||
let Some(items) = items.as_array() else {
|
||||
continue;
|
||||
};
|
||||
for item in items {
|
||||
let Some(arr) = item.as_array() else {
|
||||
continue;
|
||||
};
|
||||
if arr.len() < 4 {
|
||||
continue;
|
||||
}
|
||||
stats.total_annotations += 1;
|
||||
|
||||
// arr[3] is a list of polygon rings.
|
||||
let Some(rings) = arr[3].as_array() else {
|
||||
continue;
|
||||
};
|
||||
for ring in rings {
|
||||
let Some(pts) = ring.as_array() else {
|
||||
continue;
|
||||
};
|
||||
let mut points = Vec::with_capacity(pts.len());
|
||||
for p in pts {
|
||||
let Some(pa) = p.as_array() else {
|
||||
continue;
|
||||
};
|
||||
let (Some(x), Some(y)) = (
|
||||
pa.first().and_then(|v| v.as_f64()),
|
||||
pa.get(1).and_then(|v| v.as_f64()),
|
||||
) else {
|
||||
continue;
|
||||
};
|
||||
points.push([x, y]);
|
||||
}
|
||||
if points.len() >= 3 {
|
||||
shapes.push(Shape::Polygon {
|
||||
class_id: 0,
|
||||
class_name: class_name.clone(),
|
||||
points,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if !shapes.is_empty() {
|
||||
frames.insert(frame_key.clone(), shapes);
|
||||
}
|
||||
}
|
||||
|
||||
// Assign class_ids from sorted class names.
|
||||
let categories: Vec<Category> = class_names
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, name)| Category {
|
||||
id: i as i64,
|
||||
name: name.clone(),
|
||||
})
|
||||
.collect();
|
||||
let name_to_id: HashMap<String, i64> = categories
|
||||
.iter()
|
||||
.map(|c| (c.name.clone(), c.id))
|
||||
.collect();
|
||||
|
||||
for shapes in frames.values_mut() {
|
||||
for shape in shapes.iter_mut() {
|
||||
let (cid, cname) = match shape {
|
||||
Shape::Bbox {
|
||||
class_id,
|
||||
class_name,
|
||||
..
|
||||
}
|
||||
| Shape::Polygon {
|
||||
class_id,
|
||||
class_name,
|
||||
..
|
||||
} => (class_id, class_name),
|
||||
};
|
||||
if let Some(&id) = name_to_id.get(cname) {
|
||||
*cid = id;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
stats.categories = categories.len();
|
||||
stats.frame_count = frames.len();
|
||||
|
||||
Ok(ImportedCoco {
|
||||
categories,
|
||||
frames,
|
||||
stats,
|
||||
})
|
||||
}
|
||||
|
||||
// ---- helpers ---------------------------------------------------------------
|
||||
|
||||
fn trailing_int(file_name: &str) -> Option<u64> {
|
||||
let stem = Path::new(file_name)
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or(file_name);
|
||||
let digits: String = stem
|
||||
.chars()
|
||||
.rev()
|
||||
.take_while(|c| c.is_ascii_digit())
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect();
|
||||
if digits.is_empty() {
|
||||
None
|
||||
} else {
|
||||
digits.parse().ok()
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_flat_polygon(flat: &[serde_json::Value]) -> Vec<[f64; 2]> {
|
||||
let mut out = Vec::with_capacity(flat.len() / 2);
|
||||
let mut i = 0;
|
||||
while i + 1 < flat.len() {
|
||||
let x = flat[i].as_f64();
|
||||
let y = flat[i + 1].as_f64();
|
||||
if let (Some(x), Some(y)) = (x, y) {
|
||||
out.push([x, y]);
|
||||
}
|
||||
i += 2;
|
||||
}
|
||||
out
|
||||
}
|
||||
351
sam-tool-tauri/src-tauri/src/decoder.rs
Normal file
351
sam-tool-tauri/src-tauri/src/decoder.rs
Normal file
@@ -0,0 +1,351 @@
|
||||
use crate::video::VideoInfo;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Child, ChildStdout, Command, Stdio};
|
||||
use std::sync::{Arc, Condvar, Mutex};
|
||||
use std::thread::{self, JoinHandle};
|
||||
use std::time::Duration;
|
||||
|
||||
pub const DEFAULT_CACHE_CAPACITY: usize = 500;
|
||||
pub const MIN_CACHE_CAPACITY: usize = 200;
|
||||
pub const MAX_CACHE_CAPACITY: usize = 2000;
|
||||
const READ_BUF_BYTES: usize = 256 * 1024;
|
||||
/// If the requested frame is this far ahead of the worker, kill + reseek.
|
||||
const SEEK_AHEAD_THRESHOLD: u64 = 60;
|
||||
|
||||
struct FrameCache {
|
||||
map: HashMap<u64, Vec<u8>>,
|
||||
order: VecDeque<u64>,
|
||||
capacity: usize,
|
||||
}
|
||||
|
||||
impl FrameCache {
|
||||
fn new(capacity: usize) -> Self {
|
||||
Self {
|
||||
map: HashMap::new(),
|
||||
order: VecDeque::new(),
|
||||
capacity,
|
||||
}
|
||||
}
|
||||
|
||||
fn contains(&self, idx: u64) -> bool {
|
||||
self.map.contains_key(&idx)
|
||||
}
|
||||
|
||||
fn get(&mut self, idx: u64) -> Option<Vec<u8>> {
|
||||
let bytes = self.map.get(&idx)?.clone();
|
||||
if let Some(pos) = self.order.iter().position(|&i| i == idx) {
|
||||
self.order.remove(pos);
|
||||
}
|
||||
self.order.push_back(idx);
|
||||
Some(bytes)
|
||||
}
|
||||
|
||||
fn insert(&mut self, idx: u64, bytes: Vec<u8>) {
|
||||
if self.map.contains_key(&idx) {
|
||||
if let Some(pos) = self.order.iter().position(|&i| i == idx) {
|
||||
self.order.remove(pos);
|
||||
}
|
||||
} else if self.map.len() >= self.capacity {
|
||||
if let Some(oldest) = self.order.pop_front() {
|
||||
self.map.remove(&oldest);
|
||||
}
|
||||
}
|
||||
self.map.insert(idx, bytes);
|
||||
self.order.push_back(idx);
|
||||
}
|
||||
|
||||
fn set_capacity(&mut self, new_cap: usize) {
|
||||
self.capacity = new_cap.clamp(MIN_CACHE_CAPACITY, MAX_CACHE_CAPACITY);
|
||||
while self.map.len() > self.capacity {
|
||||
if let Some(oldest) = self.order.pop_front() {
|
||||
self.map.remove(&oldest);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SharedState {
|
||||
/// Owned by shared state (not the worker closure) so stop_worker can
|
||||
/// kill ffmpeg directly and unblock a pipe read that's waiting for the
|
||||
/// first post-seek frame.
|
||||
child: Option<Child>,
|
||||
cache: FrameCache,
|
||||
/// Frame the worker's pipe will emit next.
|
||||
next_idx: u64,
|
||||
/// Most recent frame the frontend requested — used to bound lookahead.
|
||||
reader_idx: u64,
|
||||
/// Worker reads this many frames past reader_idx before parking.
|
||||
lookahead: u64,
|
||||
eof: bool,
|
||||
error: Option<String>,
|
||||
shutdown: bool,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
state: Mutex<SharedState>,
|
||||
cond: Condvar,
|
||||
}
|
||||
|
||||
pub struct DecoderSession {
|
||||
path: PathBuf,
|
||||
info: VideoInfo,
|
||||
shared: Arc<Shared>,
|
||||
worker: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl DecoderSession {
|
||||
pub fn open(path: PathBuf, info: VideoInfo) -> Self {
|
||||
Self {
|
||||
path,
|
||||
info,
|
||||
shared: Arc::new(Shared {
|
||||
state: Mutex::new(SharedState {
|
||||
child: None,
|
||||
cache: FrameCache::new(DEFAULT_CACHE_CAPACITY),
|
||||
next_idx: 0,
|
||||
reader_idx: 0,
|
||||
lookahead: (DEFAULT_CACHE_CAPACITY as u64) / 2,
|
||||
eof: false,
|
||||
error: None,
|
||||
shutdown: false,
|
||||
}),
|
||||
cond: Condvar::new(),
|
||||
}),
|
||||
worker: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn info(&self) -> &VideoInfo {
|
||||
&self.info
|
||||
}
|
||||
|
||||
pub fn set_cache_capacity(&self, capacity: usize) {
|
||||
let clamped = capacity.clamp(MIN_CACHE_CAPACITY, MAX_CACHE_CAPACITY);
|
||||
let mut st = self.shared.state.lock().unwrap();
|
||||
st.cache.set_capacity(clamped);
|
||||
st.lookahead = (clamped as u64) / 2;
|
||||
}
|
||||
|
||||
pub fn decode_at(&mut self, idx: u64) -> Result<Vec<u8>, String> {
|
||||
if idx >= self.info.total_frames {
|
||||
return Err(format!(
|
||||
"frame {} out of range (total {})",
|
||||
idx, self.info.total_frames
|
||||
));
|
||||
}
|
||||
|
||||
// Update reader cursor, check cache, note whether a seek is needed.
|
||||
let need_seek = {
|
||||
let mut st = self.shared.state.lock().unwrap();
|
||||
st.reader_idx = idx;
|
||||
if let Some(bytes) = st.cache.get(idx) {
|
||||
self.shared.cond.notify_all();
|
||||
return Ok(bytes);
|
||||
}
|
||||
self.worker.is_none()
|
||||
|| idx < st.next_idx
|
||||
|| idx > st.next_idx.saturating_add(SEEK_AHEAD_THRESHOLD)
|
||||
};
|
||||
|
||||
if need_seek {
|
||||
self.seek_and_start_worker(idx)?;
|
||||
} else {
|
||||
self.shared.cond.notify_all();
|
||||
}
|
||||
|
||||
// Wait for the worker to produce idx.
|
||||
let mut st = self.shared.state.lock().unwrap();
|
||||
loop {
|
||||
if st.cache.contains(idx) {
|
||||
return st
|
||||
.cache
|
||||
.get(idx)
|
||||
.ok_or_else(|| "frame vanished from cache".to_string());
|
||||
}
|
||||
if let Some(e) = &st.error {
|
||||
return Err(e.clone());
|
||||
}
|
||||
if st.eof {
|
||||
return Err(format!("frame {} not produced (EOF)", idx));
|
||||
}
|
||||
if st.shutdown {
|
||||
return Err("decoder shut down".into());
|
||||
}
|
||||
st = self.shared.cond.wait(st).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn seek_and_start_worker(&mut self, idx: u64) -> Result<(), String> {
|
||||
self.stop_worker();
|
||||
|
||||
let (child, stdout) = spawn_ffmpeg(&self.path, &self.info, idx)?;
|
||||
|
||||
{
|
||||
let mut st = self.shared.state.lock().unwrap();
|
||||
st.child = Some(child);
|
||||
st.next_idx = idx;
|
||||
st.reader_idx = idx;
|
||||
st.eof = false;
|
||||
st.error = None;
|
||||
st.shutdown = false;
|
||||
}
|
||||
self.shared.cond.notify_all();
|
||||
|
||||
let shared = self.shared.clone();
|
||||
let handle = thread::spawn(move || {
|
||||
worker_loop(stdout, shared);
|
||||
});
|
||||
self.worker = Some(handle);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn stop_worker(&mut self) {
|
||||
// Take ownership of the child out of shared state and kill it FIRST so
|
||||
// a pipe read in the worker unblocks with EOF. Then join the worker.
|
||||
let child_opt = {
|
||||
let mut st = self.shared.state.lock().unwrap();
|
||||
st.shutdown = true;
|
||||
st.child.take()
|
||||
};
|
||||
if let Some(mut child) = child_opt {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
}
|
||||
self.shared.cond.notify_all();
|
||||
|
||||
if let Some(handle) = self.worker.take() {
|
||||
let _ = handle.join();
|
||||
}
|
||||
|
||||
let mut st = self.shared.state.lock().unwrap();
|
||||
st.shutdown = false;
|
||||
st.eof = false;
|
||||
st.error = None;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DecoderSession {
|
||||
fn drop(&mut self) {
|
||||
self.stop_worker();
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_ffmpeg(
|
||||
path: &PathBuf,
|
||||
info: &VideoInfo,
|
||||
idx: u64,
|
||||
) -> Result<(Child, BufReader<ChildStdout>), String> {
|
||||
let path_str = path
|
||||
.to_str()
|
||||
.ok_or_else(|| "non-utf8 path".to_string())?;
|
||||
let t = idx as f64 / info.fps;
|
||||
let mut cmd = Command::new("ffmpeg");
|
||||
cmd.args(["-v", "error", "-threads", "0"]);
|
||||
if t > 1.0 {
|
||||
cmd.args(["-ss", &format!("{:.3}", t - 1.0)]);
|
||||
cmd.args(["-i", path_str]);
|
||||
cmd.args(["-ss", "1.0"]);
|
||||
} else {
|
||||
cmd.args(["-i", path_str]);
|
||||
cmd.args(["-ss", &format!("{:.3}", t)]);
|
||||
}
|
||||
cmd.args([
|
||||
"-an",
|
||||
"-q:v",
|
||||
"3",
|
||||
"-f",
|
||||
"image2pipe",
|
||||
"-vcodec",
|
||||
"mjpeg",
|
||||
"pipe:1",
|
||||
]);
|
||||
cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
|
||||
let mut child = cmd.spawn().map_err(|e| format!("ffmpeg spawn: {e}"))?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| "ffmpeg stdout unavailable".to_string())?;
|
||||
Ok((child, BufReader::with_capacity(READ_BUF_BYTES, stdout)))
|
||||
}
|
||||
|
||||
fn worker_loop(mut stdout: BufReader<ChildStdout>, shared: Arc<Shared>) {
|
||||
loop {
|
||||
// Backpressure: wait if we're already far enough ahead.
|
||||
{
|
||||
let st = shared.state.lock().unwrap();
|
||||
if st.shutdown {
|
||||
return;
|
||||
}
|
||||
let ahead = st.next_idx.saturating_sub(st.reader_idx);
|
||||
if ahead >= st.lookahead {
|
||||
let (st2, _) = shared
|
||||
.cond
|
||||
.wait_timeout(st, Duration::from_millis(100))
|
||||
.unwrap();
|
||||
if st2.shutdown {
|
||||
return;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
match read_jpeg_bytes(&mut stdout) {
|
||||
Ok(bytes) => {
|
||||
let mut st = shared.state.lock().unwrap();
|
||||
if st.shutdown {
|
||||
return;
|
||||
}
|
||||
let idx = st.next_idx;
|
||||
st.cache.insert(idx, bytes);
|
||||
st.next_idx = idx + 1;
|
||||
shared.cond.notify_all();
|
||||
}
|
||||
Err(e) => {
|
||||
let mut st = shared.state.lock().unwrap();
|
||||
if e.contains("pipe closed") || e.contains("EOF") {
|
||||
st.eof = true;
|
||||
} else {
|
||||
st.error = Some(e);
|
||||
}
|
||||
shared.cond.notify_all();
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn read_jpeg_bytes(reader: &mut BufReader<ChildStdout>) -> Result<Vec<u8>, String> {
|
||||
let mut out = Vec::with_capacity(64 * 1024);
|
||||
let mut last = 0u8;
|
||||
let mut in_jpeg = false;
|
||||
loop {
|
||||
let buf = reader.fill_buf().map_err(|e| format!("pipe read: {e}"))?;
|
||||
if buf.is_empty() {
|
||||
return Err("ffmpeg pipe closed before frame completed".into());
|
||||
}
|
||||
let mut consumed = 0usize;
|
||||
for &b in buf {
|
||||
consumed += 1;
|
||||
if !in_jpeg {
|
||||
if last == 0xFF && b == 0xD8 {
|
||||
in_jpeg = true;
|
||||
out.push(0xFF);
|
||||
out.push(0xD8);
|
||||
}
|
||||
last = b;
|
||||
} else {
|
||||
out.push(b);
|
||||
if last == 0xFF && b == 0xD9 {
|
||||
reader.consume(consumed);
|
||||
return Ok(out);
|
||||
}
|
||||
last = b;
|
||||
}
|
||||
}
|
||||
reader.consume(consumed);
|
||||
}
|
||||
}
|
||||
188
sam-tool-tauri/src-tauri/src/extract.rs
Normal file
188
sam-tool-tauri/src-tauri/src/extract.rs
Normal file
@@ -0,0 +1,188 @@
|
||||
use crate::annotation::{AnnotationRecord, ImageAnnotations};
|
||||
use crate::coco::{Category, Shape};
|
||||
use crate::decoder::DecoderSession;
|
||||
use chrono::Utc;
|
||||
use std::path::Path;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn image_filename(idx: u64) -> String {
|
||||
format!("f_{:04}.jpg", idx)
|
||||
}
|
||||
|
||||
pub fn json_filename(idx: u64) -> String {
|
||||
format!("f_{:04}.json", idx)
|
||||
}
|
||||
|
||||
/// Parse a frame index from either the new `f_0042.jpg` pattern or the legacy
|
||||
/// `frame_000042.jpg` pattern. Returns None for unrelated files.
|
||||
fn parse_frame_idx(name: &str) -> Option<u64> {
|
||||
let stem = name.strip_suffix(".jpg")?;
|
||||
if let Some(rest) = stem.strip_prefix("f_") {
|
||||
return rest.parse().ok();
|
||||
}
|
||||
if let Some(rest) = stem.strip_prefix("frame_") {
|
||||
return rest.parse().ok();
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub fn extract_one(
|
||||
decoder: &mut DecoderSession,
|
||||
frame_idx: u64,
|
||||
out_dir: &Path,
|
||||
shapes: &[Shape],
|
||||
author: &str,
|
||||
) -> Result<(), String> {
|
||||
std::fs::create_dir_all(out_dir)
|
||||
.map_err(|e| format!("mkdir {}: {e}", out_dir.display()))?;
|
||||
|
||||
// Snapshot info so we don't hold an immutable borrow across the mutable
|
||||
// decode_at call below.
|
||||
let info = decoder.info().clone();
|
||||
write_folder_source(out_dir, &info.path);
|
||||
|
||||
let jpeg = decoder.decode_at(frame_idx)?;
|
||||
let img_path = out_dir.join(image_filename(frame_idx));
|
||||
std::fs::write(&img_path, &jpeg)
|
||||
.map_err(|e| format!("write {}: {e}", img_path.display()))?;
|
||||
|
||||
let json_path = out_dir.join(json_filename(frame_idx));
|
||||
if shapes.is_empty() {
|
||||
// Clean up any stale sibling JSON from a prior extract with annotations.
|
||||
let _ = std::fs::remove_file(&json_path);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let now = Utc::now().to_rfc3339();
|
||||
let annotations = shapes
|
||||
.iter()
|
||||
.map(|s| shape_to_record(s, author, &now))
|
||||
.collect();
|
||||
|
||||
let ia = ImageAnnotations {
|
||||
version: 1,
|
||||
image: image_filename(frame_idx),
|
||||
width: info.width,
|
||||
height: info.height,
|
||||
source_video: Some(info.path),
|
||||
annotations,
|
||||
};
|
||||
let text = serde_json::to_string_pretty(&ia)
|
||||
.map_err(|e| format!("serialize json: {e}"))?;
|
||||
std::fs::write(&json_path, text)
|
||||
.map_err(|e| format!("write {}: {e}", json_path.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub const FOLDER_SOURCE_FILENAME: &str = "source.json";
|
||||
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
pub struct FolderSource {
|
||||
pub source_video: String,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
/// Write (once, if missing) a folder-level metadata file recording which
|
||||
/// video this output folder was extracted from. Silent on failure — this is
|
||||
/// a best-effort breadcrumb, not a critical write.
|
||||
pub fn write_folder_source(out_dir: &Path, video_path: &str) {
|
||||
let meta_path = out_dir.join(FOLDER_SOURCE_FILENAME);
|
||||
if meta_path.exists() {
|
||||
return;
|
||||
}
|
||||
let meta = FolderSource {
|
||||
source_video: video_path.to_string(),
|
||||
created_at: Utc::now().to_rfc3339(),
|
||||
};
|
||||
if let Ok(text) = serde_json::to_string_pretty(&meta) {
|
||||
let _ = std::fs::write(&meta_path, text);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_folder_source(out_dir: &Path) -> Option<FolderSource> {
|
||||
let meta_path = out_dir.join(FOLDER_SOURCE_FILENAME);
|
||||
let text = std::fs::read_to_string(&meta_path).ok()?;
|
||||
serde_json::from_str(&text).ok()
|
||||
}
|
||||
|
||||
pub fn undo_one(frame_idx: u64, out_dir: &Path) -> Result<(), String> {
|
||||
// Try both naming schemes so legacy extracts are also cleanable.
|
||||
for pair in [
|
||||
(image_filename(frame_idx), json_filename(frame_idx)),
|
||||
(format!("frame_{:06}.jpg", frame_idx), format!("frame_{:06}.json", frame_idx)),
|
||||
] {
|
||||
let img = out_dir.join(&pair.0);
|
||||
let json = out_dir.join(&pair.1);
|
||||
if img.exists() {
|
||||
std::fs::remove_file(&img).map_err(|e| format!("remove {}: {e}", img.display()))?;
|
||||
}
|
||||
if json.exists() {
|
||||
std::fs::remove_file(&json).map_err(|e| format!("remove {}: {e}", json.display()))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn list_extracted(out_dir: &Path) -> Result<Vec<u64>, String> {
|
||||
if !out_dir.exists() {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
let entries = std::fs::read_dir(out_dir)
|
||||
.map_err(|e| format!("read_dir {}: {e}", out_dir.display()))?;
|
||||
for entry in entries {
|
||||
let entry = entry.map_err(|e| format!("dirent: {e}"))?;
|
||||
if let Some(name) = entry.file_name().to_str() {
|
||||
if let Some(n) = parse_frame_idx(name) {
|
||||
out.push(n);
|
||||
}
|
||||
}
|
||||
}
|
||||
out.sort();
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn write_classes_txt(out_dir: &Path, categories: &[Category]) -> Result<(), String> {
|
||||
std::fs::create_dir_all(out_dir)
|
||||
.map_err(|e| format!("mkdir {}: {e}", out_dir.display()))?;
|
||||
let mut sorted = categories.to_vec();
|
||||
sorted.sort_by_key(|c| c.id);
|
||||
let mut text = String::new();
|
||||
for c in &sorted {
|
||||
text.push_str(&c.name);
|
||||
text.push('\n');
|
||||
}
|
||||
let path = out_dir.join("classes.txt");
|
||||
std::fs::write(&path, text).map_err(|e| format!("write {}: {e}", path.display()))
|
||||
}
|
||||
|
||||
fn shape_to_record(shape: &Shape, author: &str, ts: &str) -> AnnotationRecord {
|
||||
match shape {
|
||||
Shape::Bbox {
|
||||
class_id,
|
||||
class_name,
|
||||
bbox,
|
||||
} => AnnotationRecord::Bbox {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
class_id: *class_id,
|
||||
class_name: class_name.clone(),
|
||||
bbox: *bbox,
|
||||
author: author.to_string(),
|
||||
created_at: ts.to_string(),
|
||||
updated_at: ts.to_string(),
|
||||
},
|
||||
Shape::Polygon {
|
||||
class_id,
|
||||
class_name,
|
||||
points,
|
||||
} => AnnotationRecord::Polygon {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
class_id: *class_id,
|
||||
class_name: class_name.clone(),
|
||||
points: points.clone(),
|
||||
author: author.to_string(),
|
||||
created_at: ts.to_string(),
|
||||
updated_at: ts.to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
307
sam-tool-tauri/src-tauri/src/lib.rs
Normal file
307
sam-tool-tauri/src-tauri/src/lib.rs
Normal file
@@ -0,0 +1,307 @@
|
||||
mod annotate;
|
||||
mod annotation;
|
||||
mod coco;
|
||||
mod decoder;
|
||||
mod extract;
|
||||
mod sam;
|
||||
mod settings;
|
||||
mod state;
|
||||
mod user;
|
||||
mod video;
|
||||
|
||||
use decoder::DecoderSession;
|
||||
use serde::Serialize;
|
||||
use state::AppState;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use tauri::{ipc::Response, AppHandle, Emitter, State};
|
||||
use video::VideoInfo;
|
||||
|
||||
#[tauri::command]
|
||||
fn ping() -> String {
|
||||
"pong".into()
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn open_video(state: State<'_, AppState>, path: String) -> Result<VideoInfo, String> {
|
||||
let p = PathBuf::from(&path);
|
||||
if !p.exists() {
|
||||
return Err(format!("file not found: {}", p.display()));
|
||||
}
|
||||
let info = video::probe(&p)?;
|
||||
let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
*guard = Some(DecoderSession::open(p, info.clone()));
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn close_video(state: State<'_, AppState>) -> Result<(), String> {
|
||||
let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
*guard = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn current_video(state: State<'_, AppState>) -> Result<Option<VideoInfo>, String> {
|
||||
let guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
Ok(guard.as_ref().map(|d| d.info().clone()))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn decode_frame(state: State<'_, AppState>, frame_idx: u64) -> Result<Response, String> {
|
||||
let bytes = {
|
||||
let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
let decoder = guard.as_mut().ok_or("no video loaded")?;
|
||||
decoder.decode_at(frame_idx)?
|
||||
};
|
||||
Ok(Response::new(bytes))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn import_coco(path: String) -> Result<coco::ImportedCoco, String> {
|
||||
let p = Path::new(&path);
|
||||
if !p.exists() {
|
||||
return Err(format!("file not found: {}", p.display()));
|
||||
}
|
||||
coco::import(p)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn set_cache_capacity(state: State<'_, AppState>, capacity: usize) -> Result<usize, String> {
|
||||
let clamped = capacity.clamp(decoder::MIN_CACHE_CAPACITY, decoder::MAX_CACHE_CAPACITY);
|
||||
let guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
if let Some(d) = guard.as_ref() {
|
||||
d.set_cache_capacity(clamped);
|
||||
}
|
||||
Ok(clamped)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn extract_frame(
|
||||
state: State<'_, AppState>,
|
||||
frame_idx: u64,
|
||||
out_dir: String,
|
||||
shapes: Vec<coco::Shape>,
|
||||
author: String,
|
||||
) -> Result<(), String> {
|
||||
let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
let decoder = guard.as_mut().ok_or("no video loaded")?;
|
||||
extract::extract_one(decoder, frame_idx, Path::new(&out_dir), &shapes, &author)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn undo_extract(frame_idx: u64, out_dir: String) -> Result<(), String> {
|
||||
extract::undo_one(frame_idx, Path::new(&out_dir))
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone)]
|
||||
struct BulkProgress {
|
||||
done: u64,
|
||||
total: u64,
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn bulk_extract(
|
||||
state: State<'_, AppState>,
|
||||
app: AppHandle,
|
||||
start: u64,
|
||||
end: u64,
|
||||
out_dir: String,
|
||||
shapes_by_frame: HashMap<String, Vec<coco::Shape>>,
|
||||
author: String,
|
||||
) -> Result<u64, String> {
|
||||
if end < start {
|
||||
return Err("end must be >= start".into());
|
||||
}
|
||||
let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
|
||||
let decoder = guard.as_mut().ok_or("no video loaded")?;
|
||||
let dir = PathBuf::from(&out_dir);
|
||||
let total = end - start + 1;
|
||||
let empty: Vec<coco::Shape> = Vec::new();
|
||||
let mut done: u64 = 0;
|
||||
for idx in start..=end {
|
||||
let shapes = shapes_by_frame
|
||||
.get(&idx.to_string())
|
||||
.unwrap_or(&empty);
|
||||
extract::extract_one(decoder, idx, &dir, shapes, &author)?;
|
||||
done += 1;
|
||||
if done == 1 || done % 5 == 0 || done == total {
|
||||
let _ = app.emit("bulk-extract-progress", BulkProgress { done, total });
|
||||
}
|
||||
}
|
||||
Ok(done)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn bulk_undo(
|
||||
app: AppHandle,
|
||||
start: u64,
|
||||
end: u64,
|
||||
out_dir: String,
|
||||
) -> Result<u64, String> {
|
||||
if end < start {
|
||||
return Err("end must be >= start".into());
|
||||
}
|
||||
let dir = PathBuf::from(&out_dir);
|
||||
let total = end - start + 1;
|
||||
let mut done: u64 = 0;
|
||||
for idx in start..=end {
|
||||
extract::undo_one(idx, &dir)?;
|
||||
done += 1;
|
||||
if done == 1 || done % 20 == 0 || done == total {
|
||||
let _ = app.emit("bulk-undo-progress", BulkProgress { done, total });
|
||||
}
|
||||
}
|
||||
Ok(done)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn list_extracted_frames(out_dir: String) -> Result<Vec<u64>, String> {
|
||||
extract::list_extracted(Path::new(&out_dir))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn write_classes_txt(out_dir: String, categories: Vec<coco::Category>) -> Result<(), String> {
|
||||
extract::write_classes_txt(Path::new(&out_dir), &categories)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn list_folder_images(folder: String) -> Result<Vec<annotate::ImageEntry>, String> {
|
||||
annotate::list_images(Path::new(&folder))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn read_image_bytes(path: String) -> Result<Response, String> {
|
||||
let bytes = annotate::read_image_bytes(Path::new(&path))?;
|
||||
Ok(Response::new(bytes))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn load_image_annotations(
|
||||
image_path: String,
|
||||
) -> Result<Option<annotation::ImageAnnotations>, String> {
|
||||
annotate::load_image_annotations(Path::new(&image_path))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn load_classes(folder: String) -> Result<Vec<String>, String> {
|
||||
annotate::load_classes(Path::new(&folder))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn append_class(folder: String, name: String) -> Result<Vec<String>, String> {
|
||||
annotate::append_class(Path::new(&folder), &name)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn delete_images(paths: Vec<String>) -> Result<u32, String> {
|
||||
annotate::delete_images(&paths)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
async fn sam_health_check(url: String) -> Result<bool, String> {
|
||||
sam::health_check(&url).await
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
async fn sam_list_models(url: String) -> Result<sam::ModelList, String> {
|
||||
sam::list_models(&url).await
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
async fn sam_segment_bbox(
|
||||
url: String,
|
||||
image_path: String,
|
||||
bbox: [f64; 4],
|
||||
image_width: u32,
|
||||
image_height: u32,
|
||||
model: Option<String>,
|
||||
) -> Result<Vec<[f64; 2]>, String> {
|
||||
sam::segment_bbox(
|
||||
&url,
|
||||
Path::new(&image_path),
|
||||
bbox,
|
||||
image_width,
|
||||
image_height,
|
||||
model,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn save_image_annotations(
|
||||
image_path: String,
|
||||
payload: annotation::ImageAnnotations,
|
||||
) -> Result<(), String> {
|
||||
let p = Path::new(&image_path);
|
||||
let json_path = p.with_extension("json");
|
||||
if payload.annotations.is_empty() {
|
||||
if json_path.exists() {
|
||||
std::fs::remove_file(&json_path)
|
||||
.map_err(|e| format!("remove {}: {e}", json_path.display()))?;
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
// Backfill source_video from folder-level source.json when the payload
|
||||
// doesn't carry it — handles annotations authored via AnnotateMode that
|
||||
// never went through extract_one.
|
||||
let mut payload = payload;
|
||||
if payload.source_video.is_none() {
|
||||
if let Some(parent) = p.parent() {
|
||||
if let Some(meta) = extract::read_folder_source(parent) {
|
||||
payload.source_video = Some(meta.source_video);
|
||||
}
|
||||
}
|
||||
}
|
||||
let text = serde_json::to_string_pretty(&payload)
|
||||
.map_err(|e| format!("serialize: {e}"))?;
|
||||
std::fs::write(&json_path, text)
|
||||
.map_err(|e| format!("write {}: {e}", json_path.display()))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn read_folder_source(folder: String) -> Option<String> {
|
||||
extract::read_folder_source(Path::new(&folder)).map(|s| s.source_video)
|
||||
}
|
||||
|
||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||
pub fn run() {
|
||||
tauri::Builder::default()
|
||||
.plugin(tauri_plugin_dialog::init())
|
||||
.plugin(tauri_plugin_fs::init())
|
||||
.manage(AppState::new())
|
||||
.invoke_handler(tauri::generate_handler![
|
||||
ping,
|
||||
user::current_user,
|
||||
user::list_known_users,
|
||||
user::set_current_user,
|
||||
user::clear_current_user,
|
||||
open_video,
|
||||
close_video,
|
||||
current_video,
|
||||
decode_frame,
|
||||
set_cache_capacity,
|
||||
import_coco,
|
||||
extract_frame,
|
||||
undo_extract,
|
||||
bulk_extract,
|
||||
bulk_undo,
|
||||
list_extracted_frames,
|
||||
write_classes_txt,
|
||||
list_folder_images,
|
||||
read_image_bytes,
|
||||
load_image_annotations,
|
||||
load_classes,
|
||||
save_image_annotations,
|
||||
append_class,
|
||||
delete_images,
|
||||
read_folder_source,
|
||||
settings::get_settings,
|
||||
settings::save_settings,
|
||||
sam_health_check,
|
||||
sam_list_models,
|
||||
sam_segment_bbox,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
}
|
||||
6
sam-tool-tauri/src-tauri/src/main.rs
Normal file
6
sam-tool-tauri/src-tauri/src/main.rs
Normal file
@@ -0,0 +1,6 @@
|
||||
// Prevents a console window on Windows release builds (harmless on Linux).
|
||||
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
|
||||
|
||||
fn main() {
|
||||
sam_tool_tauri_lib::run()
|
||||
}
|
||||
222
sam-tool-tauri/src-tauri/src/sam.rs
Normal file
222
sam-tool-tauri/src-tauri/src/sam.rs
Normal file
@@ -0,0 +1,222 @@
|
||||
use base64::Engine as _;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||
pub struct ModelInfo {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub family: String,
|
||||
pub available: bool,
|
||||
#[serde(default)]
|
||||
pub checkpoint: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||
pub struct ModelList {
|
||||
#[serde(default)]
|
||||
pub default: Option<String>,
|
||||
pub models: Vec<ModelInfo>,
|
||||
}
|
||||
|
||||
/// Fetch the list of models the backend advertises. The Tauri Settings
|
||||
/// modal uses this to populate a dropdown so users can pick which variant
|
||||
/// to hit.
|
||||
pub async fn list_models(url: &str) -> Result<ModelList, String> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(6))
|
||||
.build()
|
||||
.map_err(|e| format!("http client: {e}"))?;
|
||||
let endpoint = format!("{}/models", url.trim_end_matches('/'));
|
||||
let res = client
|
||||
.get(&endpoint)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("connect to {endpoint}: {e}"))?;
|
||||
if !res.status().is_success() {
|
||||
return Err(format!("GET /models → {}", res.status()));
|
||||
}
|
||||
res.json::<ModelList>()
|
||||
.await
|
||||
.map_err(|e| format!("parse /models: {e}"))
|
||||
}
|
||||
|
||||
/// Minimal XML that makes a Label-Studio-ML SAM backend happy. The names
|
||||
/// (`bbox`, `polygon`, `tag`) are referenced from the request / response by
|
||||
/// their `from_name` / `to_name` fields — if your backend expects different
|
||||
/// names we can wire them up.
|
||||
const LABEL_CONFIG: &str = r##"<View>
|
||||
<Image name="image" value="$image" zoom="true"/>
|
||||
<BrushLabels name="tag" toName="image">
|
||||
<Label value="Object" background="#ff0000"/>
|
||||
</BrushLabels>
|
||||
<RectangleLabels name="bbox" toName="image" smart="true">
|
||||
<Label value="Object" background="#ff0000"/>
|
||||
</RectangleLabels>
|
||||
<PolygonLabels name="polygon" toName="image" smart="true">
|
||||
<Label value="Object" background="#ff0000"/>
|
||||
</PolygonLabels>
|
||||
</View>"##;
|
||||
|
||||
/// Ping the SAM backend's root.
|
||||
pub async fn health_check(url: &str) -> Result<bool, String> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(4))
|
||||
.build()
|
||||
.map_err(|e| format!("http client: {e}"))?;
|
||||
let res = client
|
||||
.get(url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("connect: {e}"))?;
|
||||
Ok(res.status().is_success() || res.status().is_client_error())
|
||||
}
|
||||
|
||||
/// POST the image + bbox prompt to a Label-Studio-ML SAM backend, return the
|
||||
/// polygon in **image-pixel** coords.
|
||||
///
|
||||
/// Wire protocol (LS-ML):
|
||||
/// ```json
|
||||
/// {
|
||||
/// "tasks": [{ "id": 1, "data": { "image": "data:image/jpeg;base64,..." } }],
|
||||
/// "label_config": "<View>...</View>",
|
||||
/// "params": {
|
||||
/// "context": {
|
||||
/// "result": [{
|
||||
/// "original_width": W, "original_height": H, "image_rotation": 0,
|
||||
/// "value": { "x": %, "y": %, "width": %, "height": %,
|
||||
/// "rotation": 0, "rectanglelabels": ["Object"] },
|
||||
/// "id": "box_0",
|
||||
/// "from_name": "bbox", "to_name": "image",
|
||||
/// "type": "rectanglelabels"
|
||||
/// }]
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn segment_bbox(
|
||||
url: &str,
|
||||
image_path: &Path,
|
||||
bbox: [f64; 4],
|
||||
image_width: u32,
|
||||
image_height: u32,
|
||||
model: Option<String>,
|
||||
) -> Result<Vec<[f64; 2]>, String> {
|
||||
let image_bytes = std::fs::read(image_path)
|
||||
.map_err(|e| format!("read image {}: {e}", image_path.display()))?;
|
||||
let image_b64 = base64::engine::general_purpose::STANDARD.encode(&image_bytes);
|
||||
let data_url = format!("data:image/jpeg;base64,{}", image_b64);
|
||||
|
||||
let [x, y, w, h] = bbox;
|
||||
let width_f = image_width.max(1) as f64;
|
||||
let height_f = image_height.max(1) as f64;
|
||||
let x_pct = (x / width_f) * 100.0;
|
||||
let y_pct = (y / height_f) * 100.0;
|
||||
let w_pct = (w / width_f) * 100.0;
|
||||
let h_pct = (h / height_f) * 100.0;
|
||||
|
||||
let body = json!({
|
||||
"tasks": [{
|
||||
"id": 1,
|
||||
"data": { "image": data_url }
|
||||
}],
|
||||
"label_config": LABEL_CONFIG,
|
||||
"params": {
|
||||
"context": {
|
||||
"result": [{
|
||||
"original_width": image_width,
|
||||
"original_height": image_height,
|
||||
"image_rotation": 0,
|
||||
"value": {
|
||||
"x": x_pct,
|
||||
"y": y_pct,
|
||||
"width": w_pct,
|
||||
"height": h_pct,
|
||||
"rotation": 0,
|
||||
"rectanglelabels": ["Object"]
|
||||
},
|
||||
"id": "box_0",
|
||||
"from_name": "bbox",
|
||||
"to_name": "image",
|
||||
"type": "rectanglelabels"
|
||||
}]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(60))
|
||||
.build()
|
||||
.map_err(|e| format!("http client: {e}"))?;
|
||||
|
||||
let endpoint = format!("{}/predict", url.trim_end_matches('/'));
|
||||
let mut req = client.post(&endpoint);
|
||||
if let Some(ref m) = model {
|
||||
if !m.is_empty() {
|
||||
req = req.query(&[("model", m.as_str())]);
|
||||
}
|
||||
}
|
||||
let res = req
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("request to {endpoint}: {e}"))?;
|
||||
|
||||
if !res.status().is_success() {
|
||||
let status = res.status();
|
||||
let body_text = res.text().await.unwrap_or_default();
|
||||
return Err(format!("SAM {} — {}", status, body_text));
|
||||
}
|
||||
let raw: Value = res
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("parse SAM response: {e}"))?;
|
||||
|
||||
extract_polygon(&raw, width_f, height_f)
|
||||
.ok_or_else(|| format!("unrecognized SAM response shape: {}", raw))
|
||||
}
|
||||
|
||||
/// Walk the JSON tree looking for a polygon — LS-ML usually returns
|
||||
/// `{ results: [{ result: [{ value: { points: [[x%, y%], ...] } }] }] }`,
|
||||
/// but we also accept plain `{ polygon: [[x, y], ...] }` in pixels.
|
||||
fn extract_polygon(v: &Value, w: f64, h: f64) -> Option<Vec<[f64; 2]>> {
|
||||
// Plain shape: pixel polygon at the top level.
|
||||
if let Some(poly) = v.get("polygon").and_then(|p| p.as_array()) {
|
||||
let out = parse_points(poly, 1.0, 1.0);
|
||||
if out.len() >= 3 {
|
||||
return Some(out);
|
||||
}
|
||||
}
|
||||
|
||||
// LS-ML shape.
|
||||
let results = v.get("results").and_then(|r| r.as_array())?;
|
||||
let first = results.first()?;
|
||||
let inner = first.get("result").and_then(|r| r.as_array())?;
|
||||
for entry in inner {
|
||||
let value = entry.get("value")?;
|
||||
if let Some(points) = value.get("points").and_then(|p| p.as_array()) {
|
||||
// LS polygons are in percentages of the original image size.
|
||||
let out = parse_points(points, w / 100.0, h / 100.0);
|
||||
if out.len() >= 3 {
|
||||
return Some(out);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn parse_points(arr: &[Value], sx: f64, sy: f64) -> Vec<[f64; 2]> {
|
||||
let mut out = Vec::with_capacity(arr.len());
|
||||
for pt in arr {
|
||||
if let Some(p) = pt.as_array() {
|
||||
if p.len() >= 2 {
|
||||
if let (Some(x), Some(y)) = (p[0].as_f64(), p[1].as_f64()) {
|
||||
out.push([x * sx, y * sy]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
71
sam-tool-tauri/src-tauri/src/settings.rs
Normal file
71
sam-tool-tauri/src-tauri/src/settings.rs
Normal file
@@ -0,0 +1,71 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use tauri::{AppHandle, Manager};
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone)]
|
||||
pub struct Settings {
|
||||
#[serde(default = "default_sam_url")]
|
||||
pub sam_url: String,
|
||||
#[serde(default)]
|
||||
pub sam_enabled: bool,
|
||||
/// Model id chosen from the backend's /models list (e.g.
|
||||
/// "sam2.1_hiera_small"). `None` lets the backend pick its default.
|
||||
#[serde(default)]
|
||||
pub sam_model: Option<String>,
|
||||
#[serde(default = "default_cache_capacity")]
|
||||
pub cache_capacity: usize,
|
||||
}
|
||||
|
||||
fn default_sam_url() -> String {
|
||||
"http://localhost:9090".into()
|
||||
}
|
||||
|
||||
fn default_cache_capacity() -> usize {
|
||||
500
|
||||
}
|
||||
|
||||
impl Default for Settings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
sam_url: default_sam_url(),
|
||||
sam_enabled: false,
|
||||
sam_model: None,
|
||||
cache_capacity: default_cache_capacity(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn settings_path(app: &AppHandle) -> Result<PathBuf, String> {
|
||||
let dir = app
|
||||
.path()
|
||||
.app_local_data_dir()
|
||||
.map_err(|e| format!("app data dir: {e}"))?;
|
||||
fs::create_dir_all(&dir).map_err(|e| format!("mkdir {}: {e}", dir.display()))?;
|
||||
Ok(dir.join("settings.json"))
|
||||
}
|
||||
|
||||
pub fn load(app: &AppHandle) -> Result<Settings, String> {
|
||||
let path = settings_path(app)?;
|
||||
if !path.exists() {
|
||||
return Ok(Settings::default());
|
||||
}
|
||||
let text = fs::read_to_string(&path).map_err(|e| format!("read {}: {e}", path.display()))?;
|
||||
serde_json::from_str(&text).map_err(|e| format!("parse settings: {e}"))
|
||||
}
|
||||
|
||||
pub fn save(app: &AppHandle, s: &Settings) -> Result<(), String> {
|
||||
let path = settings_path(app)?;
|
||||
let text = serde_json::to_string_pretty(s).map_err(|e| format!("serialize: {e}"))?;
|
||||
fs::write(&path, text).map_err(|e| format!("write {}: {e}", path.display()))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn get_settings(app: AppHandle) -> Result<Settings, String> {
|
||||
load(&app)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn save_settings(app: AppHandle, settings: Settings) -> Result<(), String> {
|
||||
save(&app, &settings)
|
||||
}
|
||||
14
sam-tool-tauri/src-tauri/src/state.rs
Normal file
14
sam-tool-tauri/src-tauri/src/state.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
use crate::decoder::DecoderSession;
|
||||
use std::sync::Mutex;
|
||||
|
||||
pub struct AppState {
|
||||
pub decoder: Mutex<Option<DecoderSession>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
decoder: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
76
sam-tool-tauri/src-tauri/src/user.rs
Normal file
76
sam-tool-tauri/src-tauri/src/user.rs
Normal file
@@ -0,0 +1,76 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use tauri::{AppHandle, Manager};
|
||||
|
||||
#[derive(Default, Serialize, Deserialize)]
|
||||
struct UserStore {
|
||||
current_user: Option<String>,
|
||||
#[serde(default)]
|
||||
known_users: Vec<String>,
|
||||
}
|
||||
|
||||
fn store_path(app: &AppHandle) -> Result<PathBuf, String> {
|
||||
let dir = app
|
||||
.path()
|
||||
.app_local_data_dir()
|
||||
.map_err(|e| format!("could not resolve app data dir: {e}"))?;
|
||||
fs::create_dir_all(&dir).map_err(|e| format!("mkdir {}: {e}", dir.display()))?;
|
||||
Ok(dir.join("users.json"))
|
||||
}
|
||||
|
||||
fn load(app: &AppHandle) -> Result<UserStore, String> {
|
||||
let path = store_path(app)?;
|
||||
if !path.exists() {
|
||||
return Ok(UserStore::default());
|
||||
}
|
||||
let text = fs::read_to_string(&path).map_err(|e| format!("read {}: {e}", path.display()))?;
|
||||
serde_json::from_str(&text).map_err(|e| format!("parse users.json: {e}"))
|
||||
}
|
||||
|
||||
fn save(app: &AppHandle, store: &UserStore) -> Result<(), String> {
|
||||
let path = store_path(app)?;
|
||||
let text = serde_json::to_string_pretty(store).map_err(|e| format!("serialize: {e}"))?;
|
||||
fs::write(&path, text).map_err(|e| format!("write {}: {e}", path.display()))
|
||||
}
|
||||
|
||||
fn normalize(name: &str) -> Result<String, String> {
|
||||
let trimmed = name.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err("username cannot be empty".into());
|
||||
}
|
||||
if trimmed.len() > 64 {
|
||||
return Err("username too long (max 64 chars)".into());
|
||||
}
|
||||
Ok(trimmed.to_string())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn current_user(app: AppHandle) -> Result<Option<String>, String> {
|
||||
Ok(load(&app)?.current_user)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn list_known_users(app: AppHandle) -> Result<Vec<String>, String> {
|
||||
Ok(load(&app)?.known_users)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn set_current_user(app: AppHandle, name: String) -> Result<String, String> {
|
||||
let name = normalize(&name)?;
|
||||
let mut store = load(&app)?;
|
||||
if !store.known_users.iter().any(|u| u == &name) {
|
||||
store.known_users.push(name.clone());
|
||||
store.known_users.sort();
|
||||
}
|
||||
store.current_user = Some(name.clone());
|
||||
save(&app, &store)?;
|
||||
Ok(name)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn clear_current_user(app: AppHandle) -> Result<(), String> {
|
||||
let mut store = load(&app)?;
|
||||
store.current_user = None;
|
||||
save(&app, &store)
|
||||
}
|
||||
106
sam-tool-tauri/src-tauri/src/video.rs
Normal file
106
sam-tool-tauri/src-tauri/src/video.rs
Normal file
@@ -0,0 +1,106 @@
|
||||
use serde::Serialize;
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
#[derive(Serialize, Clone, Debug)]
|
||||
pub struct VideoInfo {
|
||||
pub path: String,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub fps: f64,
|
||||
pub total_frames: u64,
|
||||
pub duration_s: f64,
|
||||
/// Suggested output folder for this video, derived from the video path so
|
||||
/// the same folder is produced on each re-open. Frontend uses this to
|
||||
/// auto-restore the extracted-frame set without the user re-picking.
|
||||
pub default_output_folder: String,
|
||||
}
|
||||
|
||||
pub fn default_output_folder(path: &Path) -> String {
|
||||
let stem = path
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or("video");
|
||||
let parent = path.parent().unwrap_or_else(|| Path::new("."));
|
||||
parent
|
||||
.join(format!("{stem}_extracted"))
|
||||
.display()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn parse_rational(s: &str) -> Option<f64> {
|
||||
let (n, d) = s.split_once('/')?;
|
||||
let n: f64 = n.parse().ok()?;
|
||||
let d: f64 = d.parse().ok()?;
|
||||
if d == 0.0 {
|
||||
None
|
||||
} else {
|
||||
Some(n / d)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn probe(path: &Path) -> Result<VideoInfo, String> {
|
||||
let path_str = path.to_str().ok_or_else(|| "invalid path".to_string())?;
|
||||
let out = Command::new("ffprobe")
|
||||
.args([
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"v:0",
|
||||
"-show_entries",
|
||||
"stream=width,height,r_frame_rate,avg_frame_rate,nb_frames,duration",
|
||||
"-show_entries",
|
||||
"format=duration",
|
||||
"-of",
|
||||
"json",
|
||||
path_str,
|
||||
])
|
||||
.output()
|
||||
.map_err(|e| format!("ffprobe spawn failed: {e}"))?;
|
||||
if !out.status.success() {
|
||||
return Err(format!(
|
||||
"ffprobe exited {}: {}",
|
||||
out.status,
|
||||
String::from_utf8_lossy(&out.stderr)
|
||||
));
|
||||
}
|
||||
let json: serde_json::Value =
|
||||
serde_json::from_slice(&out.stdout).map_err(|e| format!("ffprobe json: {e}"))?;
|
||||
let stream = &json["streams"][0];
|
||||
if stream.is_null() {
|
||||
return Err("no video stream found".into());
|
||||
}
|
||||
let width = stream["width"]
|
||||
.as_u64()
|
||||
.ok_or("stream.width missing")? as u32;
|
||||
let height = stream["height"]
|
||||
.as_u64()
|
||||
.ok_or("stream.height missing")? as u32;
|
||||
let fps = stream["r_frame_rate"]
|
||||
.as_str()
|
||||
.and_then(parse_rational)
|
||||
.or_else(|| stream["avg_frame_rate"].as_str().and_then(parse_rational))
|
||||
.ok_or("cannot parse fps")?;
|
||||
let duration_s = stream["duration"]
|
||||
.as_str()
|
||||
.and_then(|s| s.parse::<f64>().ok())
|
||||
.or_else(|| {
|
||||
json["format"]["duration"]
|
||||
.as_str()
|
||||
.and_then(|s| s.parse::<f64>().ok())
|
||||
})
|
||||
.ok_or("cannot parse duration")?;
|
||||
let total_frames = stream["nb_frames"]
|
||||
.as_str()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or_else(|| (fps * duration_s).round() as u64);
|
||||
Ok(VideoInfo {
|
||||
path: path.display().to_string(),
|
||||
width,
|
||||
height,
|
||||
fps,
|
||||
total_frames,
|
||||
duration_s,
|
||||
default_output_folder: default_output_folder(path),
|
||||
})
|
||||
}
|
||||
33
sam-tool-tauri/src-tauri/tauri.conf.json
Normal file
33
sam-tool-tauri/src-tauri/tauri.conf.json
Normal file
@@ -0,0 +1,33 @@
|
||||
{
|
||||
"$schema": "https://schema.tauri.app/config/2",
|
||||
"productName": "SAM Tool",
|
||||
"version": "0.1.0",
|
||||
"identifier": "com.seekright.samtool",
|
||||
"build": {
|
||||
"beforeDevCommand": "npm run dev",
|
||||
"devUrl": "http://localhost:1420",
|
||||
"beforeBuildCommand": "npm run build",
|
||||
"frontendDist": "../dist"
|
||||
},
|
||||
"app": {
|
||||
"windows": [
|
||||
{
|
||||
"title": "SAM Tool",
|
||||
"width": 1280,
|
||||
"height": 860,
|
||||
"resizable": true,
|
||||
"fullscreen": false
|
||||
}
|
||||
],
|
||||
"security": {
|
||||
"csp": null
|
||||
}
|
||||
},
|
||||
"bundle": {
|
||||
"active": false,
|
||||
"targets": "all",
|
||||
"icon": [
|
||||
"icons/icon.png"
|
||||
]
|
||||
}
|
||||
}
|
||||
104
sam-tool-tauri/src/App.tsx
Normal file
104
sam-tool-tauri/src/App.tsx
Normal file
@@ -0,0 +1,104 @@
|
||||
import { useEffect, useState } from "react";
|
||||
import type { Mode, Settings } from "./types";
|
||||
import { api } from "./ipc";
|
||||
import { Login } from "./components/auth/Login";
|
||||
import { TopNav } from "./components/shell/TopNav";
|
||||
import { ExtractMode } from "./modes/ExtractMode";
|
||||
import { AnnotateMode } from "./modes/AnnotateMode";
|
||||
import { ErrorBoundary } from "./components/shell/ErrorBoundary";
|
||||
import { SettingsModal } from "./components/shell/SettingsModal";
|
||||
|
||||
export default function App() {
|
||||
const [username, setUsername] = useState<string | null>(null);
|
||||
const [checking, setChecking] = useState(true);
|
||||
const [mode, setMode] = useState<Mode>("extract");
|
||||
const [settingsOpen, setSettingsOpen] = useState(false);
|
||||
const [settings, setSettings] = useState<Settings | null>(null);
|
||||
const [fullview, setFullview] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
api.currentUser().then((u) => {
|
||||
setUsername(u);
|
||||
setChecking(false);
|
||||
});
|
||||
api.getSettings().then(setSettings).catch(() => {});
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
function onError(e: ErrorEvent) {
|
||||
// eslint-disable-next-line no-console
|
||||
console.error("[uncaught error]", e.message, e.error);
|
||||
}
|
||||
function onRejection(e: PromiseRejectionEvent) {
|
||||
// eslint-disable-next-line no-console
|
||||
console.error("[unhandled rejection]", e.reason);
|
||||
}
|
||||
window.addEventListener("error", onError);
|
||||
window.addEventListener("unhandledrejection", onRejection);
|
||||
return () => {
|
||||
window.removeEventListener("error", onError);
|
||||
window.removeEventListener("unhandledrejection", onRejection);
|
||||
};
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
function onKey(e: KeyboardEvent) {
|
||||
if (e.key === "F11") {
|
||||
e.preventDefault();
|
||||
setFullview((v) => !v);
|
||||
}
|
||||
}
|
||||
window.addEventListener("keydown", onKey);
|
||||
return () => window.removeEventListener("keydown", onKey);
|
||||
}, []);
|
||||
|
||||
if (checking) {
|
||||
return <div className="boot">…</div>;
|
||||
}
|
||||
|
||||
if (!username) {
|
||||
return <Login onLogin={setUsername} />;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className={`app${fullview ? " fullview" : ""}`}>
|
||||
{!fullview && (
|
||||
<TopNav
|
||||
mode={mode}
|
||||
onModeChange={setMode}
|
||||
username={username}
|
||||
onLogout={() => setUsername(null)}
|
||||
onOpenSettings={() => setSettingsOpen(true)}
|
||||
onToggleFullview={() => setFullview(true)}
|
||||
/>
|
||||
)}
|
||||
<main>
|
||||
{mode === "extract" && (
|
||||
<ErrorBoundary label="Extract mode">
|
||||
<ExtractMode username={username} />
|
||||
</ErrorBoundary>
|
||||
)}
|
||||
{mode === "annotate" && (
|
||||
<ErrorBoundary label="Annotate mode">
|
||||
<AnnotateMode username={username} settings={settings} />
|
||||
</ErrorBoundary>
|
||||
)}
|
||||
</main>
|
||||
{fullview && (
|
||||
<button
|
||||
className="exit-fullview"
|
||||
onClick={() => setFullview(false)}
|
||||
title="Exit fullview (F11)"
|
||||
aria-label="Exit fullview"
|
||||
>
|
||||
⛶
|
||||
</button>
|
||||
)}
|
||||
<SettingsModal
|
||||
open={settingsOpen}
|
||||
onClose={() => setSettingsOpen(false)}
|
||||
onSaved={setSettings}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
20
sam-tool-tauri/src/assets/seek-logo.svg
Normal file
20
sam-tool-tauri/src/assets/seek-logo.svg
Normal file
@@ -0,0 +1,20 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 128 128" fill="none" aria-label="SeekRight">
|
||||
<defs>
|
||||
<linearGradient id="srBg" x1="0%" y1="0%" x2="100%" y2="100%">
|
||||
<stop offset="0%" stop-color="#FFD657"/>
|
||||
<stop offset="55%" stop-color="#F5C230"/>
|
||||
<stop offset="100%" stop-color="#E5A800"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="srGloss" x1="0%" y1="0%" x2="0%" y2="100%">
|
||||
<stop offset="0%" stop-color="#FFFFFF" stop-opacity="0.35"/>
|
||||
<stop offset="45%" stop-color="#FFFFFF" stop-opacity="0"/>
|
||||
</linearGradient>
|
||||
</defs>
|
||||
<rect x="6" y="6" width="116" height="116" rx="28" fill="url(#srBg)"/>
|
||||
<rect x="6" y="6" width="116" height="62" rx="28" fill="url(#srGloss)"/>
|
||||
<rect x="6.5" y="6.5" width="115" height="115" rx="27.5" fill="none"
|
||||
stroke="#C08A00" stroke-opacity="0.35" stroke-width="1"/>
|
||||
<text x="22" y="92"
|
||||
font-family="Inter, system-ui, -apple-system, 'Segoe UI', sans-serif"
|
||||
font-weight="900" font-size="64" fill="#1F3488" letter-spacing="-3">SR</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 1.0 KiB |
12
sam-tool-tauri/src/assets/seek-mark.svg
Normal file
12
sam-tool-tauri/src/assets/seek-mark.svg
Normal file
@@ -0,0 +1,12 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 64 64" fill="none" aria-label="Seek">
|
||||
<defs>
|
||||
<linearGradient id="mark" x1="0" y1="0" x2="1" y2="1">
|
||||
<stop offset="0%" stop-color="#F5C838"/>
|
||||
<stop offset="100%" stop-color="#ECBA20"/>
|
||||
</linearGradient>
|
||||
</defs>
|
||||
<path d="M 6 6 L 46 6 Q 54 6 56 14 L 60 54 Q 61 60 54 60 L 6 60 Q 2 60 2 54 L 2 12 Q 2 6 6 6 Z"
|
||||
fill="url(#mark)"/>
|
||||
<text x="10" y="48" font-family="Inter, system-ui, -apple-system, 'Segoe UI', sans-serif"
|
||||
font-weight="900" font-size="44" fill="#1F3488" letter-spacing="-2">S</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 597 B |
98
sam-tool-tauri/src/components/annotate/ClassList.tsx
Normal file
98
sam-tool-tauri/src/components/annotate/ClassList.tsx
Normal file
@@ -0,0 +1,98 @@
|
||||
import { useState } from "react";
|
||||
import { colorForClass } from "../extract/overlay";
|
||||
|
||||
interface Props {
|
||||
classes: string[];
|
||||
activeClassId: number | null;
|
||||
onSelect: (id: number) => void;
|
||||
onAdd?: (name: string) => Promise<void> | void;
|
||||
}
|
||||
|
||||
export function ClassList({ classes, activeClassId, onSelect, onAdd }: Props) {
|
||||
const [adding, setAdding] = useState(false);
|
||||
const [name, setName] = useState("");
|
||||
const [busy, setBusy] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
async function submit() {
|
||||
if (!onAdd) return;
|
||||
const trimmed = name.trim();
|
||||
if (!trimmed) return;
|
||||
setBusy(true);
|
||||
setError(null);
|
||||
try {
|
||||
await onAdd(trimmed);
|
||||
setName("");
|
||||
setAdding(false);
|
||||
} catch (e) {
|
||||
setError(String(e));
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="class-list-wrap">
|
||||
{classes.length === 0 ? (
|
||||
<div className="class-list empty">
|
||||
no <code>classes.txt</code> in folder
|
||||
</div>
|
||||
) : (
|
||||
<ul className="class-list">
|
||||
{classes.map((name, id) => (
|
||||
<li
|
||||
key={`${id}-${name}`}
|
||||
className={id === activeClassId ? "active" : ""}
|
||||
onClick={() => onSelect(id)}
|
||||
title={`Activate "${name}"`}
|
||||
>
|
||||
<span
|
||||
className="swatch"
|
||||
style={{ background: colorForClass(id) }}
|
||||
/>
|
||||
<span className="name">{name}</span>
|
||||
<span className="id">{id}</span>
|
||||
</li>
|
||||
))}
|
||||
</ul>
|
||||
)}
|
||||
|
||||
{onAdd && (
|
||||
<div className="class-add">
|
||||
{adding ? (
|
||||
<div className="class-add-row">
|
||||
<input
|
||||
autoFocus
|
||||
type="text"
|
||||
value={name}
|
||||
maxLength={96}
|
||||
placeholder="class name"
|
||||
onChange={(e) => setName(e.target.value)}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === "Enter") { e.preventDefault(); submit(); }
|
||||
if (e.key === "Escape") {
|
||||
e.preventDefault();
|
||||
setAdding(false); setName(""); setError(null);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
<button disabled={busy || !name.trim()} onClick={submit}>add</button>
|
||||
<button
|
||||
className="link"
|
||||
disabled={busy}
|
||||
onClick={() => { setAdding(false); setName(""); setError(null); }}
|
||||
>
|
||||
cancel
|
||||
</button>
|
||||
</div>
|
||||
) : (
|
||||
<button className="link add-class-trigger" onClick={() => setAdding(true)}>
|
||||
+ add class
|
||||
</button>
|
||||
)}
|
||||
{error && <div className="class-add-error">{error}</div>}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
import { useEffect, useState } from "react";
|
||||
|
||||
interface Props {
|
||||
count: number;
|
||||
sampleNames: string[];
|
||||
busy: boolean;
|
||||
error: string | null;
|
||||
onConfirm: () => void;
|
||||
onCancel: () => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Two-step delete confirmation. Step 1 asks plainly, step 2 warns that the
|
||||
* action is irreversible. The user has to click twice, so accidental Enter
|
||||
* presses can't nuke a folder.
|
||||
*/
|
||||
export function DeleteConfirmModal({
|
||||
count,
|
||||
sampleNames,
|
||||
busy,
|
||||
error,
|
||||
onConfirm,
|
||||
onCancel,
|
||||
}: Props) {
|
||||
const [step, setStep] = useState<1 | 2>(1);
|
||||
|
||||
useEffect(() => {
|
||||
function onKey(e: KeyboardEvent) {
|
||||
if (e.key === "Escape") {
|
||||
e.preventDefault();
|
||||
onCancel();
|
||||
}
|
||||
}
|
||||
window.addEventListener("keydown", onKey, true);
|
||||
return () => window.removeEventListener("keydown", onKey, true);
|
||||
}, [onCancel]);
|
||||
|
||||
const preview = sampleNames.slice(0, 6).join(", ");
|
||||
const more = sampleNames.length > 6 ? ` +${sampleNames.length - 6} more` : "";
|
||||
|
||||
return (
|
||||
<div className="label-picker-backdrop" onMouseDown={(e) => {
|
||||
if (e.target === e.currentTarget && !busy) onCancel();
|
||||
}}>
|
||||
<div className="delete-modal" onMouseDown={(e) => e.stopPropagation()}>
|
||||
<div className="delete-modal-header">
|
||||
{step === 1 ? "Delete images?" : "This cannot be undone"}
|
||||
</div>
|
||||
<div className="delete-modal-body">
|
||||
{step === 1 ? (
|
||||
<>
|
||||
<p>
|
||||
You are about to delete <strong>{count}</strong>{" "}
|
||||
image{count === 1 ? "" : "s"}
|
||||
{" "}and any sibling <code>.json</code> annotations.
|
||||
</p>
|
||||
{preview && (
|
||||
<p className="dim delete-preview">{preview}{more}</p>
|
||||
)}
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<p>
|
||||
<strong>{count}</strong> image{count === 1 ? "" : "s"} will be
|
||||
permanently removed from disk. There is no undo.
|
||||
</p>
|
||||
<p className="dim">
|
||||
Click <strong>Delete permanently</strong> to proceed, or{" "}
|
||||
<strong>Cancel</strong> to back out.
|
||||
</p>
|
||||
</>
|
||||
)}
|
||||
{error && <p className="error">{error}</p>}
|
||||
</div>
|
||||
<div className="delete-modal-actions">
|
||||
<button onClick={onCancel} disabled={busy}>Cancel</button>
|
||||
{step === 1 ? (
|
||||
<button
|
||||
className="danger-ghost"
|
||||
onClick={() => setStep(2)}
|
||||
disabled={busy}
|
||||
autoFocus
|
||||
>
|
||||
Continue
|
||||
</button>
|
||||
) : (
|
||||
<button
|
||||
className="danger"
|
||||
onClick={onConfirm}
|
||||
disabled={busy}
|
||||
>
|
||||
{busy ? "Deleting…" : "Delete permanently"}
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
68
sam-tool-tauri/src/components/annotate/ImageList.tsx
Normal file
68
sam-tool-tauri/src/components/annotate/ImageList.tsx
Normal file
@@ -0,0 +1,68 @@
|
||||
import type { ImageEntry } from "../../types";
|
||||
|
||||
interface Props {
|
||||
images: ImageEntry[];
|
||||
selectedIdx: number | null;
|
||||
selectedSet: Set<number>;
|
||||
onSelect: (idx: number, e: React.MouseEvent) => void;
|
||||
}
|
||||
|
||||
function timeAgo(iso: string | null | undefined): string {
|
||||
if (!iso) return "";
|
||||
const t = new Date(iso).getTime();
|
||||
if (!Number.isFinite(t)) return "";
|
||||
const diff = Math.max(0, Date.now() - t);
|
||||
const s = Math.floor(diff / 1000);
|
||||
if (s < 5) return "now";
|
||||
if (s < 60) return `${s}s`;
|
||||
const m = Math.floor(s / 60);
|
||||
if (m < 60) return `${m}m`;
|
||||
const h = Math.floor(m / 60);
|
||||
if (h < 24) return `${h}h`;
|
||||
const d = Math.floor(h / 24);
|
||||
if (d < 30) return `${d}d`;
|
||||
const mo = Math.floor(d / 30);
|
||||
return `${mo}mo`;
|
||||
}
|
||||
|
||||
export function ImageList({ images, selectedIdx, selectedSet, onSelect }: Props) {
|
||||
if (images.length === 0) {
|
||||
return <div className="image-list empty">no images in folder</div>;
|
||||
}
|
||||
return (
|
||||
<ul className="image-list">
|
||||
{images.map((img, i) => {
|
||||
const isPrimary = i === selectedIdx;
|
||||
const isInSet = selectedSet.has(i);
|
||||
const cls = [
|
||||
isPrimary ? "selected" : "",
|
||||
isInSet && !isPrimary ? "multi-selected" : "",
|
||||
img.has_annotations ? "annotated" : "",
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join(" ");
|
||||
return (
|
||||
<li
|
||||
key={img.filename}
|
||||
className={cls}
|
||||
onClick={(e) => onSelect(i, e)}
|
||||
title={img.filename}
|
||||
>
|
||||
<span className="marker">{img.has_annotations ? "✓" : "·"}</span>
|
||||
<span className="name">{img.filename}</span>
|
||||
{img.annotation_count > 0 && (
|
||||
<span className="count" title={`${img.annotation_count} annotation${img.annotation_count === 1 ? "" : "s"}`}>
|
||||
{img.annotation_count}
|
||||
</span>
|
||||
)}
|
||||
{img.last_updated && (
|
||||
<span className="updated dim" title={img.last_updated}>
|
||||
{timeAgo(img.last_updated)}
|
||||
</span>
|
||||
)}
|
||||
</li>
|
||||
);
|
||||
})}
|
||||
</ul>
|
||||
);
|
||||
}
|
||||
277
sam-tool-tauri/src/components/annotate/LabelPicker.tsx
Normal file
277
sam-tool-tauri/src/components/annotate/LabelPicker.tsx
Normal file
@@ -0,0 +1,277 @@
|
||||
import { useEffect, useMemo, useRef, useState } from "react";
|
||||
import { colorForClass } from "../extract/overlay";
|
||||
|
||||
interface Props {
|
||||
classes: string[];
|
||||
/** Pre-selected class id (last-picked) to highlight when the picker opens. */
|
||||
defaultClassId: number | null;
|
||||
/** Tag shown at the top of the panel — "new bbox", "new polygon", "reclass". */
|
||||
mode: "new-bbox" | "new-polygon" | "reclass";
|
||||
onConfirm: (classId: number) => void;
|
||||
onCancel: () => void;
|
||||
/** When present, shows "+ add '<query>'" below the list when the query has
|
||||
* no matches. Must resolve with the id of the newly created class. */
|
||||
onAdd?: (name: string) => Promise<number>;
|
||||
}
|
||||
|
||||
function initials(name: string): string {
|
||||
return name
|
||||
.split(/[_\s-]+/)
|
||||
.filter(Boolean)
|
||||
.map((p) => p[0]!.toUpperCase())
|
||||
.join("");
|
||||
}
|
||||
|
||||
function scoreMatch(query: string, name: string): number {
|
||||
if (!query) return 1;
|
||||
const qu = query.toUpperCase();
|
||||
const ql = query.toLowerCase();
|
||||
const init = initials(name);
|
||||
const lo = name.toLowerCase();
|
||||
if (lo === ql) return 105;
|
||||
if (init === qu) return 100;
|
||||
if (init.startsWith(qu)) return 90;
|
||||
if (init.includes(qu)) return 75;
|
||||
if (lo.startsWith(ql)) return 70;
|
||||
if (lo.includes(ql)) return 50;
|
||||
return 0;
|
||||
}
|
||||
|
||||
export function LabelPicker({
|
||||
classes,
|
||||
defaultClassId,
|
||||
mode,
|
||||
onConfirm,
|
||||
onCancel,
|
||||
onAdd,
|
||||
}: Props) {
|
||||
const [query, setQuery] = useState("");
|
||||
const [busyAdding, setBusyAdding] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [highlightIdx, setHighlightIdx] = useState(0);
|
||||
const inputRef = useRef<HTMLInputElement>(null);
|
||||
const listRef = useRef<HTMLUListElement>(null);
|
||||
|
||||
const filtered = useMemo(() => {
|
||||
const out = classes
|
||||
.map((name, id) => ({ name, id, score: scoreMatch(query, name) }))
|
||||
.filter((x) => x.score > 0)
|
||||
.sort((a, b) => b.score - a.score || a.id - b.id);
|
||||
return out;
|
||||
}, [classes, query]);
|
||||
|
||||
// Reset highlight when list shape changes; prefer the default class if
|
||||
// present in the current filtered set.
|
||||
useEffect(() => {
|
||||
if (query.trim() === "") {
|
||||
const pos =
|
||||
defaultClassId != null
|
||||
? filtered.findIndex((f) => f.id === defaultClassId)
|
||||
: 0;
|
||||
setHighlightIdx(pos >= 0 ? pos : 0);
|
||||
} else {
|
||||
setHighlightIdx(0);
|
||||
}
|
||||
}, [query, filtered, defaultClassId]);
|
||||
|
||||
useEffect(() => {
|
||||
const el = listRef.current?.querySelector<HTMLElement>("li.highlight");
|
||||
el?.scrollIntoView({ block: "nearest" });
|
||||
}, [highlightIdx]);
|
||||
|
||||
// Force focus on mount — `autoFocus` alone has been unreliable here because
|
||||
// the picker mounts on a mouseup event whose focus target (document.body)
|
||||
// sometimes wins the race. Keeping focus on the input is what makes the
|
||||
// "start typing immediately" UX work.
|
||||
useEffect(() => {
|
||||
const el = inputRef.current;
|
||||
if (!el) return;
|
||||
el.focus();
|
||||
el.select();
|
||||
}, []);
|
||||
|
||||
// Capture all keydowns at the window level while the picker is open. This
|
||||
// routes key events to the picker even if focus escapes (e.g. the user
|
||||
// clicked the backdrop). Esc closes, printable keys feed the input,
|
||||
// navigation/confirm keys run through `onKey`.
|
||||
useEffect(() => {
|
||||
function onWin(e: KeyboardEvent) {
|
||||
if (e.key === "Escape") {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
onCancel();
|
||||
return;
|
||||
}
|
||||
const el = inputRef.current;
|
||||
if (!el) return;
|
||||
// Prevent the parent (AnnotateMode) window handler from hijacking A/D,
|
||||
// space, etc. while the picker is open.
|
||||
if (document.activeElement === el) {
|
||||
// Still guard keys the parent would act on — A/D navigation, etc.
|
||||
if (
|
||||
e.key.length === 1 &&
|
||||
!e.ctrlKey && !e.metaKey && !e.altKey
|
||||
) {
|
||||
e.stopPropagation();
|
||||
}
|
||||
return;
|
||||
}
|
||||
// Input not focused: pull focus back. For printable chars, also append
|
||||
// to the query so the first keystroke isn't lost.
|
||||
if (e.key.length === 1 && !e.ctrlKey && !e.metaKey && !e.altKey) {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
setQuery((q) => q + e.key);
|
||||
el.focus();
|
||||
} else if (
|
||||
e.key === "Enter" ||
|
||||
e.key === "ArrowUp" ||
|
||||
e.key === "ArrowDown" ||
|
||||
e.key === "Tab" ||
|
||||
e.key === "Backspace"
|
||||
) {
|
||||
e.stopPropagation();
|
||||
el.focus();
|
||||
}
|
||||
}
|
||||
window.addEventListener("keydown", onWin, true);
|
||||
return () => window.removeEventListener("keydown", onWin, true);
|
||||
}, [onCancel]);
|
||||
|
||||
async function confirmCurrent() {
|
||||
if (filtered.length > 0) {
|
||||
onConfirm(filtered[highlightIdx]?.id ?? filtered[0].id);
|
||||
return;
|
||||
}
|
||||
// No matches — offer to add a new class from the query.
|
||||
if (onAdd && query.trim()) {
|
||||
await addFromQuery();
|
||||
}
|
||||
}
|
||||
|
||||
async function addFromQuery() {
|
||||
if (!onAdd) return;
|
||||
const q = query.trim();
|
||||
if (!q) return;
|
||||
// If query exactly matches an existing class, select it instead.
|
||||
const exact = classes.findIndex(
|
||||
(c) => c.toLowerCase() === q.toLowerCase()
|
||||
);
|
||||
if (exact >= 0) {
|
||||
onConfirm(exact);
|
||||
return;
|
||||
}
|
||||
setBusyAdding(true);
|
||||
setError(null);
|
||||
try {
|
||||
const newId = await onAdd(q);
|
||||
onConfirm(newId);
|
||||
} catch (e) {
|
||||
setError(String(e));
|
||||
setBusyAdding(false);
|
||||
}
|
||||
}
|
||||
|
||||
function onKey(e: React.KeyboardEvent<HTMLInputElement>) {
|
||||
switch (e.key) {
|
||||
case "Escape":
|
||||
e.preventDefault();
|
||||
onCancel();
|
||||
break;
|
||||
case "Enter":
|
||||
e.preventDefault();
|
||||
confirmCurrent();
|
||||
break;
|
||||
case "ArrowDown":
|
||||
e.preventDefault();
|
||||
setHighlightIdx((i) => Math.min(filtered.length - 1, i + 1));
|
||||
break;
|
||||
case "ArrowUp":
|
||||
e.preventDefault();
|
||||
setHighlightIdx((i) => Math.max(0, i - 1));
|
||||
break;
|
||||
case "Tab":
|
||||
// Tab moves highlight forward (like autocomplete)
|
||||
e.preventDefault();
|
||||
setHighlightIdx((i) => (filtered.length > 0 ? (i + 1) % filtered.length : 0));
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const title =
|
||||
mode === "new-bbox"
|
||||
? "label bbox"
|
||||
: mode === "new-polygon"
|
||||
? "label polygon"
|
||||
: "change class";
|
||||
|
||||
return (
|
||||
<div className="label-picker-backdrop" onMouseDown={(e) => {
|
||||
// Click on backdrop (not modal) cancels.
|
||||
if (e.target === e.currentTarget) onCancel();
|
||||
}}>
|
||||
<div className="label-picker" onMouseDown={(e) => e.stopPropagation()}>
|
||||
<div className="label-picker-header">
|
||||
<span className="mode-tag">{title}</span>
|
||||
<span className="dim">{classes.length} classes</span>
|
||||
</div>
|
||||
<input
|
||||
ref={inputRef}
|
||||
autoFocus
|
||||
type="text"
|
||||
placeholder="type letters — e.g. SL for Street_Light or 'st' for substring"
|
||||
value={query}
|
||||
onChange={(e) => setQuery(e.target.value)}
|
||||
onKeyDown={onKey}
|
||||
/>
|
||||
<ul ref={listRef} className="label-picker-list">
|
||||
{filtered.map((x, i) => (
|
||||
<li
|
||||
key={`${x.id}-${x.name}`}
|
||||
className={i === highlightIdx ? "highlight" : ""}
|
||||
onMouseEnter={() => setHighlightIdx(i)}
|
||||
onClick={() => onConfirm(x.id)}
|
||||
>
|
||||
<span className="swatch" style={{ background: colorForClass(x.id) }} />
|
||||
<span className="name">{x.name}</span>
|
||||
<span className="init">{initials(x.name)}</span>
|
||||
<span className="id">{x.id}</span>
|
||||
</li>
|
||||
))}
|
||||
{filtered.length === 0 && !query.trim() && (
|
||||
<li className="empty">
|
||||
<span className="dim">type to filter, or add new below</span>
|
||||
</li>
|
||||
)}
|
||||
{filtered.length === 0 && query.trim() && !onAdd && (
|
||||
<li className="empty"><span className="dim">no matches</span></li>
|
||||
)}
|
||||
</ul>
|
||||
|
||||
{onAdd && (
|
||||
<div className="label-picker-add-row">
|
||||
<button
|
||||
className="add-inline"
|
||||
onClick={addFromQuery}
|
||||
disabled={busyAdding || !query.trim()}
|
||||
title={
|
||||
query.trim()
|
||||
? `Create "${query.trim()}" as a new class`
|
||||
: "Type a name first, then click to add"
|
||||
}
|
||||
>
|
||||
{query.trim()
|
||||
? `+ add "${query.trim()}" as new class`
|
||||
: "+ add new class (type name above)"}
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
{error && <p className="picker-error">{error}</p>}
|
||||
<div className="label-picker-hint">
|
||||
<kbd>↑↓</kbd> nav · <kbd>Enter</kbd> pick · <kbd>Tab</kbd> next match ·{" "}
|
||||
<kbd>Esc</kbd> cancel
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
106
sam-tool-tauri/src/components/auth/Login.tsx
Normal file
106
sam-tool-tauri/src/components/auth/Login.tsx
Normal file
@@ -0,0 +1,106 @@
|
||||
import { useEffect, useState } from "react";
|
||||
import { api } from "../../ipc";
|
||||
import seekLogo from "../../assets/seek-logo.svg";
|
||||
|
||||
interface Props {
|
||||
onLogin: (username: string) => void;
|
||||
}
|
||||
|
||||
export function Login({ onLogin }: Props) {
|
||||
const [known, setKnown] = useState<string[]>([]);
|
||||
const [name, setName] = useState("");
|
||||
const [mode, setMode] = useState<"pick" | "new">("new");
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [busy, setBusy] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
api.listKnownUsers().then((users) => {
|
||||
setKnown(users);
|
||||
if (users.length > 0) setMode("pick");
|
||||
});
|
||||
}, []);
|
||||
|
||||
async function submit(username: string) {
|
||||
const trimmed = username.trim();
|
||||
if (!trimmed) {
|
||||
setError("Please enter a name");
|
||||
return;
|
||||
}
|
||||
setBusy(true);
|
||||
setError(null);
|
||||
try {
|
||||
const confirmed = await api.setCurrentUser(trimmed);
|
||||
onLogin(confirmed);
|
||||
} catch (e) {
|
||||
setError(String(e));
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="login-page">
|
||||
<div className="login-card">
|
||||
<img className="login-logo" src={seekLogo} alt="Seek" />
|
||||
<p className="subtitle">Who's annotating?</p>
|
||||
|
||||
{mode === "pick" && known.length > 0 && (
|
||||
<div className="known-users">
|
||||
{known.map((u) => (
|
||||
<button
|
||||
key={u}
|
||||
className="known-user"
|
||||
onClick={() => submit(u)}
|
||||
disabled={busy}
|
||||
>
|
||||
{u}
|
||||
</button>
|
||||
))}
|
||||
<button
|
||||
className="link"
|
||||
onClick={() => {
|
||||
setMode("new");
|
||||
setName("");
|
||||
}}
|
||||
>
|
||||
+ new user
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{mode === "new" && (
|
||||
<form
|
||||
onSubmit={(e) => {
|
||||
e.preventDefault();
|
||||
submit(name);
|
||||
}}
|
||||
>
|
||||
<input
|
||||
autoFocus
|
||||
type="text"
|
||||
placeholder="your name"
|
||||
value={name}
|
||||
onChange={(e) => setName(e.target.value)}
|
||||
maxLength={64}
|
||||
/>
|
||||
<div className="actions">
|
||||
{known.length > 0 && (
|
||||
<button
|
||||
type="button"
|
||||
className="link"
|
||||
onClick={() => setMode("pick")}
|
||||
>
|
||||
← back
|
||||
</button>
|
||||
)}
|
||||
<button type="submit" disabled={busy || !name.trim()}>
|
||||
continue
|
||||
</button>
|
||||
</div>
|
||||
</form>
|
||||
)}
|
||||
|
||||
{error && <p className="error">{error}</p>}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
26
sam-tool-tauri/src/components/extract/FrameCanvas.tsx
Normal file
26
sam-tool-tauri/src/components/extract/FrameCanvas.tsx
Normal file
@@ -0,0 +1,26 @@
|
||||
import { forwardRef } from "react";
|
||||
import type { VideoInfo } from "../../types";
|
||||
|
||||
interface Props {
|
||||
video: VideoInfo;
|
||||
extracted?: boolean;
|
||||
}
|
||||
|
||||
export const FrameCanvas = forwardRef<HTMLCanvasElement, Props>(function FrameCanvas(
|
||||
{ video, extracted },
|
||||
ref
|
||||
) {
|
||||
return (
|
||||
<div className="canvas-stage">
|
||||
<div
|
||||
className="canvas-frame"
|
||||
style={{ aspectRatio: `${video.width} / ${video.height}` }}
|
||||
>
|
||||
<canvas ref={ref} width={video.width} height={video.height} />
|
||||
{extracted && (
|
||||
<div className="frame-extracted-banner">EXTRACTED</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
});
|
||||
149
sam-tool-tauri/src/components/extract/overlay.ts
Normal file
149
sam-tool-tauri/src/components/extract/overlay.ts
Normal file
@@ -0,0 +1,149 @@
|
||||
import type { CocoShape } from "../../types";
|
||||
|
||||
const CLASS_COLORS = [
|
||||
"#ff6b6b",
|
||||
"#4ecdc4",
|
||||
"#ffd166",
|
||||
"#a78bfa",
|
||||
"#67e8f9",
|
||||
"#fbbf24",
|
||||
"#86efac",
|
||||
"#f9a8d4",
|
||||
"#a3e635",
|
||||
"#fca5a5",
|
||||
];
|
||||
|
||||
export function colorForClass(classId: number): string {
|
||||
const idx =
|
||||
((classId % CLASS_COLORS.length) + CLASS_COLORS.length) % CLASS_COLORS.length;
|
||||
return CLASS_COLORS[idx];
|
||||
}
|
||||
|
||||
export function drawOverlay(
|
||||
ctx: CanvasRenderingContext2D,
|
||||
shapes: CocoShape[],
|
||||
canvasW: number,
|
||||
canvasH: number
|
||||
) {
|
||||
const baseLineWidth = Math.max(2, Math.round(Math.min(canvasW, canvasH) / 400));
|
||||
const fontSize = Math.max(12, Math.round(Math.min(canvasW, canvasH) / 60));
|
||||
ctx.font = `${fontSize}px sans-serif`;
|
||||
ctx.textBaseline = "top";
|
||||
|
||||
for (const shape of shapes) {
|
||||
const color = colorForClass(shape.class_id);
|
||||
ctx.strokeStyle = color;
|
||||
ctx.lineWidth = baseLineWidth;
|
||||
|
||||
if (shape.type === "bbox") {
|
||||
const [x, y, w, h] = shape.bbox;
|
||||
// Soft class-tinted fill for quick visual ID.
|
||||
ctx.fillStyle = hexWithAlpha(color, 0.1);
|
||||
ctx.fillRect(x, y, w, h);
|
||||
// Dark outer stroke for contrast against any background.
|
||||
ctx.strokeStyle = "rgba(0, 0, 0, 0.55)";
|
||||
ctx.lineWidth = baseLineWidth + Math.max(1, Math.round(baseLineWidth * 0.6));
|
||||
ctx.strokeRect(x, y, w, h);
|
||||
// Bright inner stroke in class color.
|
||||
ctx.strokeStyle = color;
|
||||
ctx.lineWidth = baseLineWidth;
|
||||
ctx.strokeRect(x, y, w, h);
|
||||
drawLabel(ctx, shape.class_name, x, y, h, color, fontSize, canvasW, canvasH);
|
||||
} else {
|
||||
ctx.beginPath();
|
||||
shape.points.forEach(([x, y], i) => {
|
||||
if (i === 0) ctx.moveTo(x, y);
|
||||
else ctx.lineTo(x, y);
|
||||
});
|
||||
ctx.closePath();
|
||||
ctx.stroke();
|
||||
ctx.fillStyle = hexWithAlpha(color, 0.18);
|
||||
ctx.fill();
|
||||
|
||||
if (shape.points.length > 0) {
|
||||
// Anchor label at the bounding box of the polygon so placement is
|
||||
// stable regardless of which vertex happens to be first.
|
||||
const [bx, by, , bh] = boundingBox(shape.points);
|
||||
drawLabel(
|
||||
ctx,
|
||||
shape.class_name,
|
||||
bx,
|
||||
by,
|
||||
bh,
|
||||
color,
|
||||
fontSize,
|
||||
canvasW,
|
||||
canvasH
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Place a label near a shape without letting it fall off the canvas.
|
||||
*
|
||||
* Vertical: prefer *above* the shape. If that clips the top, flip to *below*.
|
||||
* If below also clips, clamp inside the canvas.
|
||||
*
|
||||
* Horizontal: left-align with the shape. If the label would exceed the right
|
||||
* edge, shift left so it ends at the edge. Finally clamp to the left edge.
|
||||
*/
|
||||
function drawLabel(
|
||||
ctx: CanvasRenderingContext2D,
|
||||
text: string,
|
||||
shapeX: number,
|
||||
shapeY: number,
|
||||
shapeH: number,
|
||||
color: string,
|
||||
fontSize: number,
|
||||
canvasW: number,
|
||||
canvasH: number
|
||||
) {
|
||||
const pad = Math.round(fontSize * 0.3);
|
||||
const metrics = ctx.measureText(text);
|
||||
const w = Math.min(metrics.width + pad * 2, canvasW);
|
||||
const h = fontSize + pad * 2;
|
||||
|
||||
// Vertical
|
||||
let ly: number;
|
||||
if (shapeY - h >= 0) {
|
||||
ly = shapeY - h;
|
||||
} else if (shapeY + shapeH + h <= canvasH) {
|
||||
ly = shapeY + shapeH;
|
||||
} else {
|
||||
ly = Math.max(0, Math.min(canvasH - h, shapeY));
|
||||
}
|
||||
|
||||
// Horizontal
|
||||
let lx = shapeX;
|
||||
if (lx + w > canvasW) lx = canvasW - w;
|
||||
if (lx < 0) lx = 0;
|
||||
|
||||
ctx.fillStyle = color;
|
||||
ctx.fillRect(lx, ly, w, h);
|
||||
ctx.fillStyle = "#0e0e14";
|
||||
ctx.fillText(text, lx + pad, ly + pad);
|
||||
}
|
||||
|
||||
function boundingBox(points: [number, number][]): [number, number, number, number] {
|
||||
let minX = Infinity;
|
||||
let minY = Infinity;
|
||||
let maxX = -Infinity;
|
||||
let maxY = -Infinity;
|
||||
for (const [x, y] of points) {
|
||||
if (x < minX) minX = x;
|
||||
if (y < minY) minY = y;
|
||||
if (x > maxX) maxX = x;
|
||||
if (y > maxY) maxY = y;
|
||||
}
|
||||
return [minX, minY, maxX - minX, maxY - minY];
|
||||
}
|
||||
|
||||
function hexWithAlpha(hex: string, alpha: number): string {
|
||||
const c = hex.replace("#", "");
|
||||
const r = parseInt(c.slice(0, 2), 16);
|
||||
const g = parseInt(c.slice(2, 4), 16);
|
||||
const b = parseInt(c.slice(4, 6), 16);
|
||||
return `rgba(${r},${g},${b},${alpha})`;
|
||||
}
|
||||
42
sam-tool-tauri/src/components/shell/ErrorBoundary.tsx
Normal file
42
sam-tool-tauri/src/components/shell/ErrorBoundary.tsx
Normal file
@@ -0,0 +1,42 @@
|
||||
import React from "react";
|
||||
|
||||
interface State {
|
||||
error: Error | null;
|
||||
}
|
||||
|
||||
interface Props {
|
||||
label: string;
|
||||
children: React.ReactNode;
|
||||
}
|
||||
|
||||
export class ErrorBoundary extends React.Component<Props, State> {
|
||||
state: State = { error: null };
|
||||
|
||||
static getDerivedStateFromError(error: Error): State {
|
||||
return { error };
|
||||
}
|
||||
|
||||
componentDidCatch(error: Error, info: React.ErrorInfo) {
|
||||
// Surface in DevTools console for diagnosis.
|
||||
// eslint-disable-next-line no-console
|
||||
console.error(`[${this.props.label}] error:`, error, info);
|
||||
}
|
||||
|
||||
reset = () => this.setState({ error: null });
|
||||
|
||||
render() {
|
||||
if (this.state.error) {
|
||||
return (
|
||||
<div className="error-boundary">
|
||||
<h2>{this.props.label} crashed</h2>
|
||||
<pre>{this.state.error.message}</pre>
|
||||
<pre className="dim">{this.state.error.stack}</pre>
|
||||
<button className="primary" onClick={this.reset}>
|
||||
Try again
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
return this.props.children;
|
||||
}
|
||||
}
|
||||
274
sam-tool-tauri/src/components/shell/SettingsModal.tsx
Normal file
274
sam-tool-tauri/src/components/shell/SettingsModal.tsx
Normal file
@@ -0,0 +1,274 @@
|
||||
import { useEffect, useState } from "react";
|
||||
import { api } from "../../ipc";
|
||||
import type { SamModelInfo, Settings } from "../../types";
|
||||
|
||||
interface Props {
|
||||
open: boolean;
|
||||
onClose: () => void;
|
||||
onSaved?: (s: Settings) => void;
|
||||
}
|
||||
|
||||
type TestState =
|
||||
| { kind: "idle" }
|
||||
| { kind: "testing" }
|
||||
| { kind: "ok" }
|
||||
| { kind: "fail"; message: string };
|
||||
|
||||
type ModelsState =
|
||||
| { kind: "idle" }
|
||||
| { kind: "loading" }
|
||||
| { kind: "ok"; models: SamModelInfo[]; defaultId: string | null }
|
||||
| { kind: "fail"; message: string };
|
||||
|
||||
const URL_PRESETS: { label: string; url: string }[] = [
|
||||
{ label: "localhost", url: "http://localhost:9090" },
|
||||
{ label: "127.0.0.1", url: "http://127.0.0.1:9090" },
|
||||
];
|
||||
|
||||
export function SettingsModal({ open, onClose, onSaved }: Props) {
|
||||
const [settings, setSettings] = useState<Settings | null>(null);
|
||||
const [saving, setSaving] = useState(false);
|
||||
const [test, setTest] = useState<TestState>({ kind: "idle" });
|
||||
const [models, setModels] = useState<ModelsState>({ kind: "idle" });
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
if (!open) return;
|
||||
setError(null);
|
||||
setTest({ kind: "idle" });
|
||||
setModels({ kind: "idle" });
|
||||
api
|
||||
.getSettings()
|
||||
.then((s) => {
|
||||
setSettings(s);
|
||||
// Kick off a model list fetch in the background so the dropdown
|
||||
// populates without the user having to click "test".
|
||||
void loadModels(s.sam_url);
|
||||
})
|
||||
.catch((e) => setError(String(e)));
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [open]);
|
||||
|
||||
async function loadModels(url: string) {
|
||||
if (!url) return;
|
||||
setModels({ kind: "loading" });
|
||||
try {
|
||||
const list = await api.samListModels(url);
|
||||
setModels({ kind: "ok", models: list.models, defaultId: list.default });
|
||||
} catch (e) {
|
||||
setModels({ kind: "fail", message: String(e) });
|
||||
}
|
||||
}
|
||||
|
||||
if (!open || !settings) return open ? (
|
||||
<div className="modal-backdrop">
|
||||
<div className="modal settings-modal">
|
||||
<div className="dim">loading…</div>
|
||||
</div>
|
||||
</div>
|
||||
) : null;
|
||||
|
||||
async function save() {
|
||||
if (!settings) return;
|
||||
setSaving(true);
|
||||
setError(null);
|
||||
try {
|
||||
await api.saveSettings(settings);
|
||||
onSaved?.(settings);
|
||||
onClose();
|
||||
} catch (e) {
|
||||
setError(String(e));
|
||||
} finally {
|
||||
setSaving(false);
|
||||
}
|
||||
}
|
||||
|
||||
async function testConnection() {
|
||||
setTest({ kind: "testing" });
|
||||
try {
|
||||
const ok = await api.samHealthCheck(settings!.sam_url);
|
||||
setTest(ok ? { kind: "ok" } : { kind: "fail", message: "server replied but not 2xx" });
|
||||
if (ok) await loadModels(settings!.sam_url);
|
||||
} catch (e) {
|
||||
setTest({ kind: "fail", message: String(e) });
|
||||
}
|
||||
}
|
||||
|
||||
function setUrl(url: string) {
|
||||
setSettings({ ...settings!, sam_url: url });
|
||||
setTest({ kind: "idle" });
|
||||
setModels({ kind: "idle" });
|
||||
}
|
||||
|
||||
const currentModelExistsInList =
|
||||
models.kind === "ok" &&
|
||||
settings.sam_model != null &&
|
||||
models.models.some((m) => m.id === settings.sam_model);
|
||||
|
||||
return (
|
||||
<div
|
||||
className="modal-backdrop"
|
||||
onMouseDown={(e) => {
|
||||
if (e.target === e.currentTarget) onClose();
|
||||
}}
|
||||
>
|
||||
<div className="modal settings-modal" onMouseDown={(e) => e.stopPropagation()}>
|
||||
<div className="modal-header">
|
||||
<h2>Settings</h2>
|
||||
<button className="modal-close" onClick={onClose} aria-label="close">×</button>
|
||||
</div>
|
||||
|
||||
<section className="settings-section">
|
||||
<div className="settings-row">
|
||||
<label>
|
||||
<div className="settings-label">SAM backend URL</div>
|
||||
<input
|
||||
type="text"
|
||||
value={settings.sam_url}
|
||||
onChange={(e) => setUrl(e.target.value)}
|
||||
placeholder="http://localhost:9090"
|
||||
/>
|
||||
</label>
|
||||
<div className="url-presets">
|
||||
{URL_PRESETS.map((p) => (
|
||||
<button
|
||||
key={p.url}
|
||||
type="button"
|
||||
className="link preset"
|
||||
onClick={() => setUrl(p.url)}
|
||||
title={`use ${p.url}`}
|
||||
>
|
||||
{p.label}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="settings-row test-row">
|
||||
<button
|
||||
className="ghost"
|
||||
disabled={test.kind === "testing"}
|
||||
onClick={testConnection}
|
||||
>
|
||||
{test.kind === "testing" ? "testing…" : "test connection"}
|
||||
</button>
|
||||
{test.kind === "ok" && <span className="chip ok">✓ reachable</span>}
|
||||
{test.kind === "fail" && (
|
||||
<span className="chip fail" title={test.message}>
|
||||
✗ {test.message.slice(0, 60)}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="settings-row">
|
||||
<label>
|
||||
<div className="settings-label">
|
||||
Model
|
||||
{models.kind === "loading" && <span className="dim"> · loading…</span>}
|
||||
{models.kind === "fail" && (
|
||||
<span className="dim" title={models.message}> · couldn't load list</span>
|
||||
)}
|
||||
</div>
|
||||
<select
|
||||
value={settings.sam_model ?? ""}
|
||||
onChange={(e) =>
|
||||
setSettings({
|
||||
...settings,
|
||||
sam_model: e.target.value || null,
|
||||
})
|
||||
}
|
||||
disabled={models.kind !== "ok"}
|
||||
>
|
||||
<option value="">
|
||||
{models.kind === "ok" && models.defaultId
|
||||
? `backend default (${models.defaultId})`
|
||||
: "backend default"}
|
||||
</option>
|
||||
{models.kind === "ok" &&
|
||||
models.models.map((m) => (
|
||||
<option
|
||||
key={m.id}
|
||||
value={m.id}
|
||||
disabled={!m.available}
|
||||
>
|
||||
{m.name}
|
||||
{!m.available ? " · checkpoint missing" : ""}
|
||||
</option>
|
||||
))}
|
||||
{!currentModelExistsInList && settings.sam_model && (
|
||||
<option value={settings.sam_model}>
|
||||
{settings.sam_model} (not in list)
|
||||
</option>
|
||||
)}
|
||||
</select>
|
||||
</label>
|
||||
{models.kind === "fail" && (
|
||||
<button
|
||||
className="link"
|
||||
onClick={() => loadModels(settings.sam_url)}
|
||||
>
|
||||
retry
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="settings-row">
|
||||
<label className="switch-row">
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={settings.sam_enabled}
|
||||
onChange={(e) =>
|
||||
setSettings({ ...settings, sam_enabled: e.target.checked })
|
||||
}
|
||||
/>
|
||||
<span>
|
||||
<strong>SAM mode</strong>
|
||||
<span className="dim"> — auto-segment bboxes via the backend</span>
|
||||
</span>
|
||||
</label>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<section className="settings-section">
|
||||
<div className="settings-row">
|
||||
<label>
|
||||
<div className="settings-label">
|
||||
Frame cache capacity
|
||||
<span className="dim"> · {settings.cache_capacity} frames</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={200}
|
||||
max={2000}
|
||||
step={50}
|
||||
value={settings.cache_capacity}
|
||||
onChange={(e) =>
|
||||
setSettings({
|
||||
...settings,
|
||||
cache_capacity: Number(e.target.value),
|
||||
})
|
||||
}
|
||||
/>
|
||||
<div className="range-ticks">
|
||||
<span>200</span>
|
||||
<span>1000</span>
|
||||
<span>2000</span>
|
||||
</div>
|
||||
</label>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
{error && <div className="picker-error">{error}</div>}
|
||||
|
||||
<div className="modal-footer">
|
||||
<button className="ghost" onClick={onClose} disabled={saving}>
|
||||
cancel
|
||||
</button>
|
||||
<button className="primary" onClick={save} disabled={saving}>
|
||||
{saving ? "saving…" : "save"}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
80
sam-tool-tauri/src/components/shell/TopNav.tsx
Normal file
80
sam-tool-tauri/src/components/shell/TopNav.tsx
Normal file
@@ -0,0 +1,80 @@
|
||||
import { useState } from "react";
|
||||
import type { Mode } from "../../types";
|
||||
import { api } from "../../ipc";
|
||||
import seekLogo from "../../assets/seek-logo.svg";
|
||||
|
||||
interface Props {
|
||||
mode: Mode;
|
||||
onModeChange: (m: Mode) => void;
|
||||
username: string;
|
||||
onLogout: () => void;
|
||||
onOpenSettings: () => void;
|
||||
onToggleFullview: () => void;
|
||||
}
|
||||
|
||||
export function TopNav({
|
||||
mode,
|
||||
onModeChange,
|
||||
username,
|
||||
onLogout,
|
||||
onOpenSettings,
|
||||
onToggleFullview,
|
||||
}: Props) {
|
||||
const [menuOpen, setMenuOpen] = useState(false);
|
||||
|
||||
async function logout() {
|
||||
await api.clearCurrentUser();
|
||||
setMenuOpen(false);
|
||||
onLogout();
|
||||
}
|
||||
|
||||
return (
|
||||
<header className="topnav">
|
||||
<div className="brand-mark">
|
||||
<img className="brand-icon" src={seekLogo} alt="SR" />
|
||||
<span className="brand-text">SAM Tool</span>
|
||||
</div>
|
||||
<nav>
|
||||
<button
|
||||
className={mode === "extract" ? "active" : ""}
|
||||
onClick={() => onModeChange("extract")}
|
||||
>
|
||||
Extract
|
||||
</button>
|
||||
<button
|
||||
className={mode === "annotate" ? "active" : ""}
|
||||
onClick={() => onModeChange("annotate")}
|
||||
>
|
||||
Annotate
|
||||
</button>
|
||||
</nav>
|
||||
<div className="spacer" />
|
||||
<button
|
||||
className="icon-btn"
|
||||
onClick={onToggleFullview}
|
||||
title="Full view (F11)"
|
||||
aria-label="Full view"
|
||||
>
|
||||
⛶
|
||||
</button>
|
||||
<button
|
||||
className="icon-btn"
|
||||
onClick={onOpenSettings}
|
||||
title="Settings"
|
||||
aria-label="Settings"
|
||||
>
|
||||
⚙
|
||||
</button>
|
||||
<div className="user-chip">
|
||||
<button className="chip" onClick={() => setMenuOpen((v) => !v)}>
|
||||
{username} ▾
|
||||
</button>
|
||||
{menuOpen && (
|
||||
<div className="menu" onMouseLeave={() => setMenuOpen(false)}>
|
||||
<button onClick={logout}>Switch user</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</header>
|
||||
);
|
||||
}
|
||||
2063
sam-tool-tauri/src/index.css
Normal file
2063
sam-tool-tauri/src/index.css
Normal file
File diff suppressed because it is too large
Load Diff
94
sam-tool-tauri/src/ipc.ts
Normal file
94
sam-tool-tauri/src/ipc.ts
Normal file
@@ -0,0 +1,94 @@
|
||||
import { invoke } from "@tauri-apps/api/core";
|
||||
import type {
|
||||
CocoCategory,
|
||||
CocoShape,
|
||||
ImageAnnotations,
|
||||
ImageEntry,
|
||||
ImportedCoco,
|
||||
SamModelList,
|
||||
Settings,
|
||||
VideoInfo,
|
||||
} from "./types";
|
||||
|
||||
export const api = {
|
||||
ping: () => invoke<string>("ping"),
|
||||
|
||||
currentUser: () => invoke<string | null>("current_user"),
|
||||
listKnownUsers: () => invoke<string[]>("list_known_users"),
|
||||
setCurrentUser: (name: string) => invoke<string>("set_current_user", { name }),
|
||||
clearCurrentUser: () => invoke<void>("clear_current_user"),
|
||||
|
||||
openVideo: (path: string) => invoke<VideoInfo>("open_video", { path }),
|
||||
closeVideo: () => invoke<void>("close_video"),
|
||||
currentVideo: () => invoke<VideoInfo | null>("current_video"),
|
||||
decodeFrame: (frameIdx: number) =>
|
||||
invoke<ArrayBuffer>("decode_frame", { frameIdx }),
|
||||
setCacheCapacity: (capacity: number) =>
|
||||
invoke<number>("set_cache_capacity", { capacity }),
|
||||
|
||||
importCoco: (path: string) => invoke<ImportedCoco>("import_coco", { path }),
|
||||
|
||||
extractFrame: (frameIdx: number, outDir: string, shapes: CocoShape[], author: string) =>
|
||||
invoke<void>("extract_frame", { frameIdx, outDir, shapes, author }),
|
||||
undoExtract: (frameIdx: number, outDir: string) =>
|
||||
invoke<void>("undo_extract", { frameIdx, outDir }),
|
||||
bulkExtract: (
|
||||
start: number,
|
||||
end: number,
|
||||
outDir: string,
|
||||
shapesByFrame: Record<string, CocoShape[]>,
|
||||
author: string
|
||||
) =>
|
||||
invoke<number>("bulk_extract", {
|
||||
start,
|
||||
end,
|
||||
outDir,
|
||||
shapesByFrame,
|
||||
author,
|
||||
}),
|
||||
bulkUndo: (start: number, end: number, outDir: string) =>
|
||||
invoke<number>("bulk_undo", { start, end, outDir }),
|
||||
listExtractedFrames: (outDir: string) =>
|
||||
invoke<number[]>("list_extracted_frames", { outDir }),
|
||||
writeClassesTxt: (outDir: string, categories: CocoCategory[]) =>
|
||||
invoke<void>("write_classes_txt", { outDir, categories }),
|
||||
|
||||
listFolderImages: (folder: string) =>
|
||||
invoke<ImageEntry[]>("list_folder_images", { folder }),
|
||||
readImageBytes: (path: string) =>
|
||||
invoke<ArrayBuffer>("read_image_bytes", { path }),
|
||||
loadImageAnnotations: (imagePath: string) =>
|
||||
invoke<ImageAnnotations | null>("load_image_annotations", { imagePath }),
|
||||
loadClasses: (folder: string) => invoke<string[]>("load_classes", { folder }),
|
||||
saveImageAnnotations: (imagePath: string, payload: ImageAnnotations) =>
|
||||
invoke<void>("save_image_annotations", { imagePath, payload }),
|
||||
appendClass: (folder: string, name: string) =>
|
||||
invoke<string[]>("append_class", { folder, name }),
|
||||
deleteImages: (paths: string[]) => invoke<number>("delete_images", { paths }),
|
||||
readFolderSource: (folder: string) =>
|
||||
invoke<string | null>("read_folder_source", { folder }),
|
||||
|
||||
getSettings: () => invoke<Settings>("get_settings"),
|
||||
saveSettings: (settings: Settings) =>
|
||||
invoke<void>("save_settings", { settings }),
|
||||
samHealthCheck: (url: string) =>
|
||||
invoke<boolean>("sam_health_check", { url }),
|
||||
samListModels: (url: string) =>
|
||||
invoke<SamModelList>("sam_list_models", { url }),
|
||||
samSegmentBbox: (
|
||||
url: string,
|
||||
imagePath: string,
|
||||
bbox: [number, number, number, number],
|
||||
imageWidth: number,
|
||||
imageHeight: number,
|
||||
model?: string | null
|
||||
) =>
|
||||
invoke<[number, number][]>("sam_segment_bbox", {
|
||||
url,
|
||||
imagePath,
|
||||
bbox,
|
||||
imageWidth,
|
||||
imageHeight,
|
||||
model: model ?? null,
|
||||
}),
|
||||
};
|
||||
10
sam-tool-tauri/src/main.tsx
Normal file
10
sam-tool-tauri/src/main.tsx
Normal file
@@ -0,0 +1,10 @@
|
||||
import React from "react";
|
||||
import ReactDOM from "react-dom/client";
|
||||
import App from "./App";
|
||||
import "./index.css";
|
||||
|
||||
ReactDOM.createRoot(document.getElementById("root")!).render(
|
||||
<React.StrictMode>
|
||||
<App />
|
||||
</React.StrictMode>
|
||||
);
|
||||
1825
sam-tool-tauri/src/modes/AnnotateMode.tsx
Normal file
1825
sam-tool-tauri/src/modes/AnnotateMode.tsx
Normal file
File diff suppressed because it is too large
Load Diff
648
sam-tool-tauri/src/modes/ExtractMode.tsx
Normal file
648
sam-tool-tauri/src/modes/ExtractMode.tsx
Normal file
@@ -0,0 +1,648 @@
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { open as openDialog } from "@tauri-apps/plugin-dialog";
|
||||
import { listen } from "@tauri-apps/api/event";
|
||||
import { api } from "../ipc";
|
||||
import type { CocoShape, ImportedCoco, VideoInfo } from "../types";
|
||||
import { FrameCanvas } from "../components/extract/FrameCanvas";
|
||||
import { drawOverlay } from "../components/extract/overlay";
|
||||
|
||||
const SPEEDS = [0.25, 0.5, 1, 1.5, 2, 3, 4];
|
||||
|
||||
interface Props {
|
||||
username: string;
|
||||
}
|
||||
|
||||
interface BulkProgress {
|
||||
done: number;
|
||||
total: number;
|
||||
phase: "extract" | "undo";
|
||||
}
|
||||
|
||||
export function ExtractMode({ username }: Props) {
|
||||
const [video, setVideo] = useState<VideoInfo | null>(null);
|
||||
const [frameIdx, setFrameIdx] = useState(0);
|
||||
const [playing, setPlaying] = useState(false);
|
||||
const [speed, setSpeed] = useState(1);
|
||||
const [reviewMode, setReviewMode] = useState(false);
|
||||
const [overlayVisible, setOverlayVisible] = useState(true);
|
||||
const [loadError, setLoadError] = useState<string | null>(null);
|
||||
|
||||
const [coco, setCoco] = useState<ImportedCoco | null>(null);
|
||||
|
||||
// Extraction state
|
||||
const [outputFolder, setOutputFolder] = useState<string | null>(null);
|
||||
const [extractedFrames, setExtractedFrames] = useState<Set<number>>(new Set());
|
||||
const [bulkStart, setBulkStart] = useState<number | null>(null);
|
||||
const [bulkEnd, setBulkEnd] = useState<number | null>(null);
|
||||
const [bulkProgress, setBulkProgress] = useState<BulkProgress | null>(null);
|
||||
const [bulkError, setBulkError] = useState<string | null>(null);
|
||||
const [infoCollapsed, setInfoCollapsed] = useState(false);
|
||||
|
||||
const canvasRef = useRef<HTMLCanvasElement>(null);
|
||||
|
||||
const frameMap = useMemo(() => {
|
||||
const m = new Map<number, CocoShape[]>();
|
||||
if (!coco) return m;
|
||||
for (const [k, v] of Object.entries(coco.frames)) m.set(Number(k), v);
|
||||
return m;
|
||||
}, [coco]);
|
||||
|
||||
// Refs mirroring state for effects that shouldn't depend on fast-changing values.
|
||||
const frameIdxRef = useRef(0);
|
||||
const frameMapRef = useRef<Map<number, CocoShape[]>>(new Map());
|
||||
const overlayVisibleRef = useRef(overlayVisible);
|
||||
const annotatedFrames = useRef<Set<number>>(new Set());
|
||||
// Playback clock anchor — set when playback starts, rebased on seeks so a
|
||||
// slider drag during play doesn't get overwritten by the tick loop.
|
||||
const playAnchorRef = useRef<{ time: number; idx: number } | null>(null);
|
||||
|
||||
useEffect(() => { frameIdxRef.current = frameIdx; }, [frameIdx]);
|
||||
useEffect(() => {
|
||||
frameMapRef.current = frameMap;
|
||||
annotatedFrames.current = new Set(frameMap.keys());
|
||||
}, [frameMap]);
|
||||
useEffect(() => { overlayVisibleRef.current = overlayVisible; }, [overlayVisible]);
|
||||
|
||||
// --- decode pump (latest-wins coalescing) ----------------------------
|
||||
const requestedIdx = useRef(0);
|
||||
const renderedIdx = useRef(-1);
|
||||
const decodingInFlight = useRef(false);
|
||||
|
||||
const renderFrame = useCallback(
|
||||
async (bytes: ArrayBuffer, atFrameIdx: number) => {
|
||||
const canvas = canvasRef.current;
|
||||
if (!canvas) return;
|
||||
const blob = new Blob([bytes], { type: "image/jpeg" });
|
||||
let bitmap: ImageBitmap;
|
||||
try { bitmap = await createImageBitmap(blob); }
|
||||
catch (e) { console.error("createImageBitmap failed", e); return; }
|
||||
const ctx = canvas.getContext("2d");
|
||||
if (!ctx) { bitmap.close(); return; }
|
||||
ctx.clearRect(0, 0, canvas.width, canvas.height);
|
||||
ctx.drawImage(bitmap, 0, 0, canvas.width, canvas.height);
|
||||
bitmap.close();
|
||||
|
||||
if (overlayVisibleRef.current) {
|
||||
const shapes = frameMapRef.current.get(atFrameIdx);
|
||||
if (shapes && shapes.length > 0) {
|
||||
drawOverlay(ctx, shapes, canvas.width, canvas.height);
|
||||
}
|
||||
}
|
||||
},
|
||||
[]
|
||||
);
|
||||
|
||||
const pump = useCallback(async () => {
|
||||
if (decodingInFlight.current) return;
|
||||
decodingInFlight.current = true;
|
||||
try {
|
||||
while (requestedIdx.current !== renderedIdx.current) {
|
||||
const want = requestedIdx.current;
|
||||
try {
|
||||
const bytes = await api.decodeFrame(want);
|
||||
if (requestedIdx.current !== want) continue;
|
||||
await renderFrame(bytes, want);
|
||||
renderedIdx.current = want;
|
||||
} catch (e) {
|
||||
if (requestedIdx.current !== want) continue;
|
||||
console.error("decode_frame failed", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
decodingInFlight.current = false;
|
||||
}
|
||||
}, [renderFrame]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!video) return;
|
||||
requestedIdx.current = frameIdx;
|
||||
pump();
|
||||
}, [frameIdx, pump, video]);
|
||||
|
||||
useEffect(() => {
|
||||
renderedIdx.current = -1;
|
||||
pump();
|
||||
}, [overlayVisible, frameMap, pump]);
|
||||
|
||||
// --- playback loop ---------------------------------------------------
|
||||
useEffect(() => {
|
||||
if (!playing || !video) return;
|
||||
playAnchorRef.current = {
|
||||
time: performance.now(),
|
||||
idx: frameIdxRef.current,
|
||||
};
|
||||
|
||||
let raf = 0;
|
||||
const tick = () => {
|
||||
const anchor = playAnchorRef.current;
|
||||
if (!anchor) return;
|
||||
const elapsed = (performance.now() - anchor.time) / 1000;
|
||||
const want = Math.min(
|
||||
video.total_frames - 1,
|
||||
anchor.idx + Math.floor(elapsed * video.fps * speed)
|
||||
);
|
||||
if (reviewMode && annotatedFrames.current.has(want)) {
|
||||
setFrameIdx(want);
|
||||
setPlaying(false);
|
||||
return;
|
||||
}
|
||||
setFrameIdx(want);
|
||||
if (want >= video.total_frames - 1) { setPlaying(false); return; }
|
||||
raf = requestAnimationFrame(tick);
|
||||
};
|
||||
raf = requestAnimationFrame(tick);
|
||||
return () => {
|
||||
cancelAnimationFrame(raf);
|
||||
playAnchorRef.current = null;
|
||||
};
|
||||
}, [playing, video, speed, reviewMode]);
|
||||
|
||||
// --- navigation helpers ----------------------------------------------
|
||||
const seekToFrame = useCallback((idx: number) => {
|
||||
if (!video) return;
|
||||
const clamped = Math.max(0, Math.min(video.total_frames - 1, idx));
|
||||
setFrameIdx(clamped);
|
||||
// If a playback tick is active, rebase its anchor so the loop continues
|
||||
// from the new position instead of snapping back.
|
||||
if (playAnchorRef.current) {
|
||||
playAnchorRef.current = { time: performance.now(), idx: clamped };
|
||||
}
|
||||
}, [video]);
|
||||
|
||||
const stepFrames = useCallback((delta: number) => {
|
||||
if (!video) return;
|
||||
setPlaying(false);
|
||||
setFrameIdx((i) => Math.max(0, Math.min(video.total_frames - 1, i + delta)));
|
||||
}, [video]);
|
||||
|
||||
const togglePlay = useCallback(() => {
|
||||
if (!video) return;
|
||||
setPlaying((p) => !p);
|
||||
}, [video]);
|
||||
|
||||
const bumpSpeed = useCallback((dir: 1 | -1) => {
|
||||
setSpeed((s) => {
|
||||
const i = SPEEDS.indexOf(s);
|
||||
const next = Math.max(0, Math.min(SPEEDS.length - 1, (i < 0 ? 2 : i) + dir));
|
||||
return SPEEDS[next];
|
||||
});
|
||||
}, []);
|
||||
|
||||
// --- extraction helpers ----------------------------------------------
|
||||
const applyOutputFolder = useCallback(async (path: string) => {
|
||||
setOutputFolder(path);
|
||||
try {
|
||||
const existing = await api.listExtractedFrames(path);
|
||||
setExtractedFrames(new Set(existing));
|
||||
} catch (e) {
|
||||
// Folder may not exist yet — that's fine, we'll create it on extract.
|
||||
setExtractedFrames(new Set());
|
||||
console.warn("listExtractedFrames failed (non-fatal):", e);
|
||||
}
|
||||
if (coco && coco.categories.length > 0) {
|
||||
try { await api.writeClassesTxt(path, coco.categories); } catch (e) { console.error(e); }
|
||||
}
|
||||
}, [coco]);
|
||||
|
||||
async function pickOutputFolder() {
|
||||
const selected = await openDialog({ directory: true, multiple: false });
|
||||
if (!selected || typeof selected !== "string") return;
|
||||
await applyOutputFolder(selected);
|
||||
}
|
||||
|
||||
const extractCurrentFrame = useCallback(async () => {
|
||||
if (!outputFolder || !video) return;
|
||||
setBulkError(null);
|
||||
try {
|
||||
const shapes = frameMap.get(frameIdx) ?? [];
|
||||
await api.extractFrame(frameIdx, outputFolder, shapes, username);
|
||||
setExtractedFrames((prev) => new Set(prev).add(frameIdx));
|
||||
} catch (e) {
|
||||
setBulkError(String(e));
|
||||
}
|
||||
}, [outputFolder, video, frameIdx, frameMap, username]);
|
||||
|
||||
const undoCurrentFrame = useCallback(async () => {
|
||||
if (!outputFolder) return;
|
||||
setBulkError(null);
|
||||
try {
|
||||
await api.undoExtract(frameIdx, outputFolder);
|
||||
setExtractedFrames((prev) => {
|
||||
const n = new Set(prev);
|
||||
n.delete(frameIdx);
|
||||
return n;
|
||||
});
|
||||
} catch (e) {
|
||||
setBulkError(String(e));
|
||||
}
|
||||
}, [outputFolder, frameIdx]);
|
||||
|
||||
const setRangeStart = useCallback(() => setBulkStart(frameIdx), [frameIdx]);
|
||||
const setRangeEnd = useCallback(() => setBulkEnd(frameIdx), [frameIdx]);
|
||||
const clearRange = useCallback(() => { setBulkStart(null); setBulkEnd(null); }, []);
|
||||
|
||||
const bulkRange = useMemo(() => {
|
||||
if (bulkStart == null || bulkEnd == null) return null;
|
||||
return {
|
||||
start: Math.min(bulkStart, bulkEnd),
|
||||
end: Math.max(bulkStart, bulkEnd),
|
||||
};
|
||||
}, [bulkStart, bulkEnd]);
|
||||
|
||||
async function runBulkExtract() {
|
||||
if (!outputFolder || !bulkRange) return;
|
||||
setBulkError(null);
|
||||
setBulkProgress({ done: 0, total: bulkRange.end - bulkRange.start + 1, phase: "extract" });
|
||||
const unlisten = await listen<{ done: number; total: number }>(
|
||||
"bulk-extract-progress",
|
||||
(e) => setBulkProgress({ ...e.payload, phase: "extract" })
|
||||
);
|
||||
try {
|
||||
await api.bulkExtract(
|
||||
bulkRange.start,
|
||||
bulkRange.end,
|
||||
outputFolder,
|
||||
coco?.frames ?? {},
|
||||
username
|
||||
);
|
||||
const existing = await api.listExtractedFrames(outputFolder);
|
||||
setExtractedFrames(new Set(existing));
|
||||
} catch (e) {
|
||||
setBulkError(String(e));
|
||||
} finally {
|
||||
unlisten();
|
||||
setBulkProgress(null);
|
||||
}
|
||||
}
|
||||
|
||||
async function runBulkUndo() {
|
||||
if (!outputFolder || !bulkRange) return;
|
||||
setBulkError(null);
|
||||
setBulkProgress({ done: 0, total: bulkRange.end - bulkRange.start + 1, phase: "undo" });
|
||||
const unlisten = await listen<{ done: number; total: number }>(
|
||||
"bulk-undo-progress",
|
||||
(e) => setBulkProgress({ ...e.payload, phase: "undo" })
|
||||
);
|
||||
try {
|
||||
await api.bulkUndo(bulkRange.start, bulkRange.end, outputFolder);
|
||||
const existing = await api.listExtractedFrames(outputFolder);
|
||||
setExtractedFrames(new Set(existing));
|
||||
} catch (e) {
|
||||
setBulkError(String(e));
|
||||
} finally {
|
||||
unlisten();
|
||||
setBulkProgress(null);
|
||||
}
|
||||
}
|
||||
|
||||
// --- keyboard --------------------------------------------------------
|
||||
useEffect(() => {
|
||||
function onKey(e: KeyboardEvent) {
|
||||
const tag = (e.target as HTMLElement | null)?.tagName;
|
||||
if (tag === "INPUT" || tag === "TEXTAREA") return;
|
||||
const step = e.ctrlKey || e.metaKey ? 10 : 1;
|
||||
switch (e.key) {
|
||||
case " ":
|
||||
e.preventDefault(); togglePlay(); break;
|
||||
case "ArrowRight":
|
||||
e.preventDefault(); stepFrames(step); break;
|
||||
case "ArrowLeft":
|
||||
e.preventDefault(); stepFrames(-step); break;
|
||||
case "ArrowUp":
|
||||
e.preventDefault(); bumpSpeed(1); break;
|
||||
case "ArrowDown":
|
||||
e.preventDefault(); bumpSpeed(-1); break;
|
||||
case "h": case "H":
|
||||
e.preventDefault(); setOverlayVisible((v) => !v); break;
|
||||
case "e": case "E":
|
||||
e.preventDefault(); extractCurrentFrame(); break;
|
||||
case "u": case "U":
|
||||
e.preventDefault(); undoCurrentFrame(); break;
|
||||
case "[":
|
||||
e.preventDefault(); setRangeStart(); break;
|
||||
case "]":
|
||||
e.preventDefault(); setRangeEnd(); break;
|
||||
}
|
||||
}
|
||||
window.addEventListener("keydown", onKey);
|
||||
return () => window.removeEventListener("keydown", onKey);
|
||||
}, [togglePlay, stepFrames, bumpSpeed, extractCurrentFrame, undoCurrentFrame, setRangeStart, setRangeEnd]);
|
||||
|
||||
// --- video/coco open -------------------------------------------------
|
||||
useEffect(() => {
|
||||
api.currentVideo().then((v) => {
|
||||
if (!v) return;
|
||||
setVideo(v);
|
||||
// Re-adopt the default folder on mount too, so switching modes and
|
||||
// coming back doesn't drop the extracted-frames overlay.
|
||||
applyOutputFolder(v.default_output_folder);
|
||||
});
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, []);
|
||||
|
||||
async function chooseVideo() {
|
||||
setLoadError(null);
|
||||
const selected = await openDialog({
|
||||
multiple: false,
|
||||
filters: [
|
||||
{
|
||||
name: "Video",
|
||||
extensions: [
|
||||
"mp4","MP4","mov","MOV","avi","AVI","mkv","MKV","webm","WEBM","flv","FLV","wmv","WMV",
|
||||
],
|
||||
},
|
||||
{ name: "All files", extensions: ["*"] },
|
||||
],
|
||||
});
|
||||
if (!selected || typeof selected !== "string") return;
|
||||
try {
|
||||
const info = await api.openVideo(selected);
|
||||
setVideo(info);
|
||||
setFrameIdx(0);
|
||||
setPlaying(false);
|
||||
setSpeed(1);
|
||||
renderedIdx.current = -1;
|
||||
// Reset per-video extraction context, then auto-adopt the per-video
|
||||
// default folder. If the user already extracted frames into it in a
|
||||
// prior session, listExtractedFrames rehydrates them immediately.
|
||||
setBulkStart(null);
|
||||
setBulkEnd(null);
|
||||
await applyOutputFolder(info.default_output_folder);
|
||||
} catch (e) {
|
||||
setLoadError(String(e));
|
||||
}
|
||||
}
|
||||
|
||||
async function chooseCoco() {
|
||||
const selected = await openDialog({
|
||||
multiple: false,
|
||||
filters: [{ name: "Annotations JSON", extensions: ["json"] }, { name: "All files", extensions: ["*"] }],
|
||||
});
|
||||
if (!selected || typeof selected !== "string") return;
|
||||
try {
|
||||
const imported = await api.importCoco(selected);
|
||||
setCoco(imported);
|
||||
// If a folder is already picked, write (or refresh) classes.txt.
|
||||
if (outputFolder && imported.categories.length > 0) {
|
||||
try { await api.writeClassesTxt(outputFolder, imported.categories); } catch (e) { console.error(e); }
|
||||
}
|
||||
} catch (e) {
|
||||
setLoadError(String(e));
|
||||
}
|
||||
}
|
||||
|
||||
function clearCoco() { setCoco(null); }
|
||||
|
||||
// --- render ----------------------------------------------------------
|
||||
|
||||
if (!video) {
|
||||
return (
|
||||
<div className="upload-pane">
|
||||
<button className="primary" onClick={chooseVideo}>Upload Video</button>
|
||||
{loadError && <p className="error">{loadError}</p>}
|
||||
<p className="dim">
|
||||
Supported: mp4, mov, avi, mkv, webm, flv, wmv. Decoded via <code>ffmpeg</code>.
|
||||
</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const t = frameIdx / video.fps;
|
||||
const currentShapes = frameMap.get(frameIdx);
|
||||
const isCurrentExtracted = extractedFrames.has(frameIdx);
|
||||
const pct = bulkProgress ? Math.round((bulkProgress.done / bulkProgress.total) * 100) : 0;
|
||||
|
||||
return (
|
||||
<div className="extract-mode">
|
||||
{!infoCollapsed && (
|
||||
<>
|
||||
<div className="video-info">
|
||||
<span className="path" title={video.path}>{video.path.split("/").pop()}</span>
|
||||
<span className="dim">
|
||||
{video.width}×{video.height} · {video.fps.toFixed(2)} fps · {video.total_frames} frames · {video.duration_s.toFixed(1)}s
|
||||
</span>
|
||||
<button className="link" onClick={chooseVideo}>change…</button>
|
||||
</div>
|
||||
|
||||
<div className="coco-info">
|
||||
{coco ? (
|
||||
<>
|
||||
<span>
|
||||
<span className="dim">{coco.stats.source_format === "sam" ? "SAM" : "COCO"}:</span>{" "}
|
||||
<strong>{coco.stats.total_annotations}</strong> annotations ·{" "}
|
||||
<strong>{coco.stats.frame_count}</strong> frames ·{" "}
|
||||
<strong>{coco.stats.categories}</strong> categories
|
||||
{coco.stats.skipped_rle > 0 && <span className="dim"> · {coco.stats.skipped_rle} RLE skipped</span>}
|
||||
{coco.stats.skipped_no_image > 0 && <span className="dim"> · {coco.stats.skipped_no_image} unmappable</span>}
|
||||
</span>
|
||||
<button className="link" onClick={chooseCoco}>replace…</button>
|
||||
<button className="link" onClick={clearCoco}>clear</button>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<span className="dim">no annotations loaded</span>
|
||||
<button className="link" onClick={chooseCoco}>import annotations…</button>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
<div className={`extract-info${infoCollapsed ? " collapsed" : ""}`}>
|
||||
{outputFolder ? (
|
||||
<>
|
||||
<span>Output: <strong title={outputFolder}>{outputFolder.split("/").pop()}</strong></span>
|
||||
<span className="dim"> · {extractedFrames.size} extracted</span>
|
||||
<button className="link" onClick={pickOutputFolder}>change…</button>
|
||||
</>
|
||||
) : (
|
||||
<button className="link" onClick={pickOutputFolder}>choose output folder…</button>
|
||||
)}
|
||||
<div className="extract-actions">
|
||||
<button
|
||||
className={isCurrentExtracted ? "extracted" : "primary-ghost"}
|
||||
onClick={extractCurrentFrame}
|
||||
disabled={!outputFolder}
|
||||
title="E"
|
||||
>
|
||||
{isCurrentExtracted ? "✓ re-extract (E)" : "Extract (E)"}
|
||||
</button>
|
||||
<button
|
||||
onClick={undoCurrentFrame}
|
||||
disabled={!outputFolder || !isCurrentExtracted}
|
||||
title="U"
|
||||
>
|
||||
Undo (U)
|
||||
</button>
|
||||
<div className="range-picker">
|
||||
<button onClick={setRangeStart} title="[">
|
||||
Start [ {bulkStart ?? "–"}
|
||||
</button>
|
||||
<button onClick={setRangeEnd} title="]">
|
||||
End ] {bulkEnd ?? "–"}
|
||||
</button>
|
||||
{(bulkStart != null || bulkEnd != null) && (
|
||||
<button className="link" onClick={clearRange}>clear</button>
|
||||
)}
|
||||
</div>
|
||||
<button
|
||||
onClick={runBulkExtract}
|
||||
disabled={!outputFolder || !bulkRange || !!bulkProgress}
|
||||
>
|
||||
Bulk Extract{bulkRange && ` (${bulkRange.end - bulkRange.start + 1})`}
|
||||
</button>
|
||||
<button
|
||||
onClick={runBulkUndo}
|
||||
disabled={!outputFolder || !bulkRange || !!bulkProgress}
|
||||
>
|
||||
Bulk Undo
|
||||
</button>
|
||||
<button
|
||||
className="collapse-toggle"
|
||||
onClick={() => setInfoCollapsed((v) => !v)}
|
||||
title={infoCollapsed ? "show info" : "hide info"}
|
||||
>
|
||||
{infoCollapsed ? "v" : "^"}
|
||||
</button>
|
||||
</div>
|
||||
{bulkProgress && (
|
||||
<div className="progress">
|
||||
<div className="progress-bar" style={{ width: `${pct}%` }} />
|
||||
<span>
|
||||
{bulkProgress.phase === "extract" ? "extracting" : "undoing"} ·{" "}
|
||||
{bulkProgress.done} / {bulkProgress.total}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{bulkError && <p className="error">{bulkError}</p>}
|
||||
</div>
|
||||
|
||||
<FrameCanvas ref={canvasRef} video={video} extracted={isCurrentExtracted} />
|
||||
|
||||
<div className="timeline">
|
||||
<div className="timeline-track-wrap">
|
||||
<input
|
||||
type="range"
|
||||
min={0}
|
||||
max={video.total_frames - 1}
|
||||
value={frameIdx}
|
||||
onChange={(e) => seekToFrame(Number(e.target.value))}
|
||||
style={
|
||||
{
|
||||
"--val": `${(frameIdx / Math.max(1, video.total_frames - 1)) * 100}%`,
|
||||
} as React.CSSProperties
|
||||
}
|
||||
/>
|
||||
{extractedFrames.size > 0 && (
|
||||
<ExtractedTicks
|
||||
extracted={extractedFrames}
|
||||
totalFrames={video.total_frames}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
<div className="timeline-readout">
|
||||
frame <strong>{frameIdx}</strong> / {video.total_frames - 1}
|
||||
<span className="dim"> · {formatTime(t)}</span>
|
||||
{currentShapes && currentShapes.length > 0 && (
|
||||
<span className="badge">
|
||||
{currentShapes.length} annotation{currentShapes.length === 1 ? "" : "s"}
|
||||
</span>
|
||||
)}
|
||||
{isCurrentExtracted && <span className="badge extracted">✓ extracted</span>}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="controls">
|
||||
<button className={playing ? "active" : ""} onClick={togglePlay}>
|
||||
{playing ? "❚❚ pause" : "▶ play"}
|
||||
</button>
|
||||
<button onClick={() => stepFrames(-1)} title="←">◀ 1</button>
|
||||
<button onClick={() => stepFrames(1)} title="→">1 ▶</button>
|
||||
<button onClick={() => stepFrames(-10)} title="Ctrl+←">◀◀ 10</button>
|
||||
<button onClick={() => stepFrames(10)} title="Ctrl+→">10 ▶▶</button>
|
||||
<div className="speed">
|
||||
<button onClick={() => bumpSpeed(-1)} title="↓">–</button>
|
||||
<span className="speed-readout">{speed}×</span>
|
||||
<button onClick={() => bumpSpeed(1)} title="↑">+</button>
|
||||
</div>
|
||||
<label className={`review-toggle ${reviewMode ? "on" : ""}`}>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={reviewMode}
|
||||
onChange={(e) => setReviewMode(e.target.checked)}
|
||||
/>
|
||||
Review mode
|
||||
</label>
|
||||
<button
|
||||
className={`overlay-toggle ${overlayVisible ? "on" : ""}`}
|
||||
onClick={() => setOverlayVisible((v) => !v)}
|
||||
title="H"
|
||||
>
|
||||
{overlayVisible ? "👁 overlay on" : "◌ overlay off"}
|
||||
</button>
|
||||
<span className="dim shortcuts">space · ← → · ctrl+← → · ↑ ↓ · H · E · U · [ ]</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function formatTime(s: number) {
|
||||
const m = Math.floor(s / 60);
|
||||
const r = s - m * 60;
|
||||
return `${m}:${r.toFixed(2).padStart(5, "0")}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Canvas overlay that draws one green tick per pixel-column of the timeline
|
||||
* track where at least one extracted frame falls. Canvas avoids the DOM
|
||||
* blow-up of 1-div-per-frame on long videos, and collapses dense neighboring
|
||||
* frames into a single visible stripe at narrow timeline widths.
|
||||
*/
|
||||
function ExtractedTicks({
|
||||
extracted,
|
||||
totalFrames,
|
||||
}: {
|
||||
extracted: Set<number>;
|
||||
totalFrames: number;
|
||||
}) {
|
||||
const ref = useRef<HTMLCanvasElement>(null);
|
||||
useEffect(() => {
|
||||
const canvas = ref.current;
|
||||
if (!canvas) return;
|
||||
let raf = 0;
|
||||
|
||||
const draw = () => {
|
||||
raf = 0;
|
||||
const rect = canvas.getBoundingClientRect();
|
||||
const dpr = window.devicePixelRatio || 1;
|
||||
const w = Math.max(1, Math.floor(rect.width * dpr));
|
||||
const h = Math.max(1, Math.floor(rect.height * dpr));
|
||||
if (canvas.width !== w) canvas.width = w;
|
||||
if (canvas.height !== h) canvas.height = h;
|
||||
const ctx = canvas.getContext("2d");
|
||||
if (!ctx) return;
|
||||
ctx.clearRect(0, 0, w, h);
|
||||
if (extracted.size === 0 || totalFrames <= 0) return;
|
||||
ctx.fillStyle = "rgba(92, 207, 117, 0.85)";
|
||||
const max = Math.max(1, totalFrames - 1);
|
||||
for (const f of extracted) {
|
||||
const x = Math.round((f / max) * (w - 1));
|
||||
ctx.fillRect(x, 0, Math.max(1, Math.floor(dpr)), h);
|
||||
}
|
||||
};
|
||||
|
||||
const schedule = () => {
|
||||
if (raf) return;
|
||||
raf = requestAnimationFrame(draw);
|
||||
};
|
||||
schedule();
|
||||
const ro = new ResizeObserver(schedule);
|
||||
ro.observe(canvas);
|
||||
return () => {
|
||||
if (raf) cancelAnimationFrame(raf);
|
||||
ro.disconnect();
|
||||
};
|
||||
}, [extracted, totalFrames]);
|
||||
return <canvas ref={ref} className="timeline-ticks" aria-hidden="true" />;
|
||||
}
|
||||
109
sam-tool-tauri/src/types.ts
Normal file
109
sam-tool-tauri/src/types.ts
Normal file
@@ -0,0 +1,109 @@
|
||||
export type Mode = "extract" | "annotate";
|
||||
|
||||
export interface VideoInfo {
|
||||
path: string;
|
||||
width: number;
|
||||
height: number;
|
||||
fps: number;
|
||||
total_frames: number;
|
||||
duration_s: number;
|
||||
default_output_folder: string;
|
||||
}
|
||||
|
||||
export interface CocoCategory {
|
||||
id: number;
|
||||
name: string;
|
||||
}
|
||||
|
||||
export interface CocoBboxShape {
|
||||
type: "bbox";
|
||||
class_id: number;
|
||||
class_name: string;
|
||||
bbox: [number, number, number, number];
|
||||
}
|
||||
|
||||
export interface CocoPolygonShape {
|
||||
type: "polygon";
|
||||
class_id: number;
|
||||
class_name: string;
|
||||
points: [number, number][];
|
||||
}
|
||||
|
||||
export type CocoShape = CocoBboxShape | CocoPolygonShape;
|
||||
|
||||
export interface CocoImportStats {
|
||||
total_annotations: number;
|
||||
frame_count: number;
|
||||
categories: number;
|
||||
skipped_rle: number;
|
||||
skipped_no_image: number;
|
||||
source_format: "coco" | "sam" | string;
|
||||
}
|
||||
|
||||
export interface ImportedCoco {
|
||||
categories: CocoCategory[];
|
||||
frames: Record<string, CocoShape[]>;
|
||||
stats: CocoImportStats;
|
||||
}
|
||||
|
||||
export type AnnotationType = "bbox" | "polygon";
|
||||
|
||||
export interface BboxAnnotation {
|
||||
id: string;
|
||||
class_id: number;
|
||||
class_name: string;
|
||||
type: "bbox";
|
||||
bbox: [number, number, number, number];
|
||||
author: string;
|
||||
created_at: string;
|
||||
updated_at: string;
|
||||
}
|
||||
|
||||
export interface PolygonAnnotation {
|
||||
id: string;
|
||||
class_id: number;
|
||||
class_name: string;
|
||||
type: "polygon";
|
||||
points: [number, number][];
|
||||
author: string;
|
||||
created_at: string;
|
||||
updated_at: string;
|
||||
}
|
||||
|
||||
export type Annotation = BboxAnnotation | PolygonAnnotation;
|
||||
|
||||
export interface ImageAnnotations {
|
||||
version: 1;
|
||||
image: string;
|
||||
width: number;
|
||||
height: number;
|
||||
source_video?: string | null;
|
||||
annotations: Annotation[];
|
||||
}
|
||||
|
||||
export interface ImageEntry {
|
||||
filename: string;
|
||||
has_annotations: boolean;
|
||||
annotation_count: number;
|
||||
last_updated: string | null;
|
||||
}
|
||||
|
||||
export interface Settings {
|
||||
sam_url: string;
|
||||
sam_enabled: boolean;
|
||||
sam_model?: string | null;
|
||||
cache_capacity: number;
|
||||
}
|
||||
|
||||
export interface SamModelInfo {
|
||||
id: string;
|
||||
name: string;
|
||||
family: string;
|
||||
available: boolean;
|
||||
checkpoint: string;
|
||||
}
|
||||
|
||||
export interface SamModelList {
|
||||
default: string | null;
|
||||
models: SamModelInfo[];
|
||||
}
|
||||
1
sam-tool-tauri/src/vite-env.d.ts
vendored
Normal file
1
sam-tool-tauri/src/vite-env.d.ts
vendored
Normal file
@@ -0,0 +1 @@
|
||||
/// <reference types="vite/client" />
|
||||
21
sam-tool-tauri/tsconfig.json
Normal file
21
sam-tool-tauri/tsconfig.json
Normal file
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2020",
|
||||
"useDefineForClassFields": true,
|
||||
"lib": ["ES2020", "DOM", "DOM.Iterable"],
|
||||
"module": "ESNext",
|
||||
"skipLibCheck": true,
|
||||
"moduleResolution": "bundler",
|
||||
"allowImportingTsExtensions": true,
|
||||
"resolveJsonModule": true,
|
||||
"isolatedModules": true,
|
||||
"noEmit": true,
|
||||
"jsx": "react-jsx",
|
||||
"strict": true,
|
||||
"noUnusedLocals": true,
|
||||
"noUnusedParameters": true,
|
||||
"noFallthroughCasesInSwitch": true
|
||||
},
|
||||
"include": ["src"],
|
||||
"references": [{ "path": "./tsconfig.node.json" }]
|
||||
}
|
||||
11
sam-tool-tauri/tsconfig.node.json
Normal file
11
sam-tool-tauri/tsconfig.node.json
Normal file
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"composite": true,
|
||||
"skipLibCheck": true,
|
||||
"module": "ESNext",
|
||||
"moduleResolution": "bundler",
|
||||
"allowSyntheticDefaultImports": true,
|
||||
"strict": true
|
||||
},
|
||||
"include": ["vite.config.ts"]
|
||||
}
|
||||
18
sam-tool-tauri/vite.config.ts
Normal file
18
sam-tool-tauri/vite.config.ts
Normal file
@@ -0,0 +1,18 @@
|
||||
import { defineConfig } from "vite";
|
||||
import react from "@vitejs/plugin-react";
|
||||
|
||||
const host = process.env.TAURI_DEV_HOST;
|
||||
|
||||
export default defineConfig(async () => ({
|
||||
plugins: [react()],
|
||||
clearScreen: false,
|
||||
server: {
|
||||
port: 1420,
|
||||
strictPort: true,
|
||||
host: host || false,
|
||||
hmr: host
|
||||
? { protocol: "ws", host, port: 1421 }
|
||||
: undefined,
|
||||
watch: { ignored: ["**/src-tauri/**"] },
|
||||
},
|
||||
}));
|
||||
7
sam2-backend/.gitignore
vendored
Normal file
7
sam2-backend/.gitignore
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
.venv/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pt
|
||||
*.pth
|
||||
.vscode/
|
||||
.idea/
|
||||
191
sam2-backend/README.md
Normal file
191
sam2-backend/README.md
Normal file
@@ -0,0 +1,191 @@
|
||||
# SAM2 Backend
|
||||
|
||||
Standalone HTTP server that runs Meta's SAM-family image predictors and
|
||||
speaks the Label-Studio-ML `/predict` protocol the SAM-Tool Tauri app
|
||||
already uses. Can host locally or on any remote PC on your network — point
|
||||
the Tauri app at `http://<host>:9090` from Settings.
|
||||
|
||||
Ships with adapters for both **SAM 2 / 2.1** and **SAM 3**. The registry
|
||||
(`models.yaml`) is pluggable — new families just drop an adapter file into
|
||||
`adapters/` and register in `adapters/__init__.py`.
|
||||
|
||||
---
|
||||
|
||||
## Layout
|
||||
|
||||
```
|
||||
sam2-backend/
|
||||
├── server.py FastAPI app — /, /models, /predict
|
||||
├── registry.py Model spec loader + lazy-load cache
|
||||
├── adapters/
|
||||
│ ├── base.py Adapter contract
|
||||
│ ├── sam2_adapter.py SAM2 implementation
|
||||
│ └── __init__.py Family → adapter class map
|
||||
├── models.yaml Which models exist + where their checkpoints live
|
||||
├── requirements.txt HTTP server deps (Torch + SAM2 installed separately)
|
||||
├── run.sh Convenience launcher
|
||||
└── README.md
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Install (one-time, on the host PC)
|
||||
|
||||
### 1. Python environment
|
||||
|
||||
```bash
|
||||
cd sam2-backend
|
||||
python -m venv .venv
|
||||
source .venv/bin/activate
|
||||
pip install --upgrade pip
|
||||
```
|
||||
|
||||
### 2. PyTorch (choose one)
|
||||
|
||||
Pick the build that matches your GPU/CPU — do NOT rely on the generic
|
||||
`torch` pulled in transitively by SAM2.
|
||||
|
||||
```bash
|
||||
# CUDA 12.1 (typical modern NVIDIA)
|
||||
pip install torch==2.4.* torchvision==0.19.* --index-url https://download.pytorch.org/whl/cu121
|
||||
|
||||
# CPU-only (works but inference is ~50× slower)
|
||||
pip install torch==2.4.* torchvision==0.19.* --index-url https://download.pytorch.org/whl/cpu
|
||||
```
|
||||
|
||||
### 3. SAM 2 package
|
||||
|
||||
The SAM2 source is already in this repo at `../sam2/sam2/`. Install it as
|
||||
an editable package:
|
||||
|
||||
```bash
|
||||
pip install -e ../sam2/sam2
|
||||
```
|
||||
|
||||
> If you see `RuntimeError: You're likely running Python from the parent
|
||||
> directory of the sam2 repository`, you started Python from the *wrong*
|
||||
> directory — always run `server.py` from `sam2-backend/` (not from the
|
||||
> repo root).
|
||||
|
||||
### 4. SAM 3 package (optional — only if you want SAM3 models)
|
||||
|
||||
```bash
|
||||
# Upstream repo (adjust if you're using a fork):
|
||||
pip install "sam3 @ git+https://github.com/facebookresearch/sam3"
|
||||
# For the HuggingFace tracker path (optional, used by `config: tracker`):
|
||||
pip install transformers
|
||||
```
|
||||
|
||||
The SAM3 image-predictor path uses the `.pt` file in `../sam3_weights/sam3.pt`
|
||||
— already referenced by `models.yaml`. No additional downloads needed for
|
||||
the default SAM3 entry.
|
||||
|
||||
### 5. The rest
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 6. Checkpoints
|
||||
|
||||
**SAM 2 / 2.1** — present at `../sam2/sam2/checkpoints/sam2.1_hiera_{tiny,small,base_plus,large}.pt`.
|
||||
If any are missing, download them with:
|
||||
|
||||
```bash
|
||||
cd ../sam2/sam2/checkpoints
|
||||
./download_ckpts.sh
|
||||
```
|
||||
|
||||
**SAM 3** — present at `../files/sam3_weights/sam3.pt`.
|
||||
|
||||
`models.yaml` points at all of these by default. Edit the file to add,
|
||||
remove, or relocate entries.
|
||||
|
||||
---
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
# From sam2-backend/
|
||||
./run.sh # defaults: 0.0.0.0:9090, auto device
|
||||
HOST=127.0.0.1 PORT=9090 ./run.sh # localhost only
|
||||
DEVICE=cuda:1 ./run.sh # pin a specific GPU
|
||||
|
||||
# or directly
|
||||
python server.py --host 0.0.0.0 --port 9090 --device auto
|
||||
```
|
||||
|
||||
First request for a given model triggers a one-time load (few seconds). By
|
||||
default the previously-loaded model is evicted on switch to keep GPU memory
|
||||
bounded. Set `SAM_BACKEND_KEEP_ALL=1` to keep them all resident.
|
||||
|
||||
---
|
||||
|
||||
## Smoke-test
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:9090/ # {"status":"ok",...}
|
||||
curl -s http://localhost:9090/models # {"default":"...","models":[...]}
|
||||
```
|
||||
|
||||
From the Tauri app: open **Settings**, set the SAM URL to
|
||||
`http://<host>:9090`, enable SAM, pick a model, save. Drawing a bbox should
|
||||
now round-trip through SAM2 and come back as a polygon.
|
||||
|
||||
---
|
||||
|
||||
## API
|
||||
|
||||
### `GET /`
|
||||
```json
|
||||
{"status":"ok","device":"cuda","models":4,"default":"sam2.1_hiera_small"}
|
||||
```
|
||||
|
||||
### `GET /models`
|
||||
```json
|
||||
{
|
||||
"default": "sam2.1_hiera_small",
|
||||
"models": [
|
||||
{"id":"sam2.1_hiera_tiny","name":"SAM 2.1 · Hiera Tiny",
|
||||
"family":"sam2","available":true,
|
||||
"checkpoint":"/abs/path/sam2.1_hiera_tiny.pt"},
|
||||
...
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### `POST /predict?model=<id>`
|
||||
Body: Label-Studio-ML task + bbox prompt (see `server.py` docstring).
|
||||
Optional `?model=<id>` query selects which model. Omit to use the default.
|
||||
|
||||
Response: LS-ML `results[].result[].value.points` as percentages.
|
||||
|
||||
---
|
||||
|
||||
## SAM 3 notes
|
||||
|
||||
Two `config` modes are supported per entry in `models.yaml`:
|
||||
|
||||
| `config` | Loader | Checkpoint | Notes |
|
||||
|-------------------|-----------------------------------------------------------------|--------------------------------------------|-----------------------------------------------------------------------|
|
||||
| `image_predictor` | `sam3.model_builder.build_sam3_image_model` | Local `.pt` file | Default. Box = visual exemplar; picks mask with highest IoU to user box. |
|
||||
| `tracker` | `transformers.Sam3TrackerModel.from_pretrained("facebook/sam3")` | HF model id (set `remote: true` in yaml) | Cleaner single-box → single-mask. Requires `transformers` + first-run network access. |
|
||||
|
||||
The shipped default SAM3 entry uses `image_predictor` + the local weights.
|
||||
To enable the tracker path, uncomment the `sam3_tracker` block in
|
||||
`models.yaml`.
|
||||
|
||||
---
|
||||
|
||||
## Adding another model family later
|
||||
|
||||
1. Drop a new file into `adapters/`, e.g. `adapters/sam4_adapter.py`, that
|
||||
implements `BaseAdapter`.
|
||||
2. Register it in `adapters/__init__.py`:
|
||||
```python
|
||||
ADAPTERS = {"sam2": Sam2Adapter, "sam3": Sam3Adapter, "sam4": Sam4Adapter}
|
||||
```
|
||||
3. Add entries to `models.yaml` with `family: sam4`.
|
||||
4. Restart the server.
|
||||
|
||||
No changes to `server.py` or `registry.py` should be needed.
|
||||
10
sam2-backend/adapters/__init__.py
Normal file
10
sam2-backend/adapters/__init__.py
Normal file
@@ -0,0 +1,10 @@
|
||||
from .base import BaseAdapter, PredictResult
|
||||
from .sam2_adapter import Sam2Adapter
|
||||
from .sam3_adapter import Sam3Adapter
|
||||
|
||||
ADAPTERS: dict[str, type[BaseAdapter]] = {
|
||||
"sam2": Sam2Adapter,
|
||||
"sam3": Sam3Adapter,
|
||||
}
|
||||
|
||||
__all__ = ["BaseAdapter", "PredictResult", "Sam2Adapter", "Sam3Adapter", "ADAPTERS"]
|
||||
34
sam2-backend/adapters/base.py
Normal file
34
sam2-backend/adapters/base.py
Normal file
@@ -0,0 +1,34 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class PredictResult:
|
||||
"""Result of one bbox-prompted segmentation.
|
||||
|
||||
`polygon` is a list of (x, y) pixel coordinates describing the outer
|
||||
contour of the largest mask blob. `score` is the model's confidence for
|
||||
the selected mask (in [0, 1]).
|
||||
"""
|
||||
polygon: list[tuple[float, float]]
|
||||
score: float
|
||||
|
||||
|
||||
class BaseAdapter(Protocol):
|
||||
"""Minimal adapter contract. Each model family implements this."""
|
||||
|
||||
def load(self, config: str, checkpoint: str, device: str) -> None:
|
||||
"""Load weights into memory. Called once, lazily on first predict."""
|
||||
...
|
||||
|
||||
def predict_bbox(
|
||||
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
|
||||
) -> PredictResult:
|
||||
"""Run segmentation with a bbox prompt, return the best polygon."""
|
||||
...
|
||||
|
||||
def unload(self) -> None:
|
||||
"""Free GPU memory. Called when the registry evicts this model."""
|
||||
...
|
||||
82
sam2-backend/adapters/sam2_adapter.py
Normal file
82
sam2-backend/adapters/sam2_adapter.py
Normal file
@@ -0,0 +1,82 @@
|
||||
"""SAM2 adapter — wraps `sam2.sam2_image_predictor.SAM2ImagePredictor`.
|
||||
|
||||
The SAM2 repo uses Hydra to resolve config files relative to the `sam2`
|
||||
package. We rely on that resolution: pass `config="configs/sam2.1/..."` and
|
||||
it works as long as the `sam2` package is importable (pip-installed or on
|
||||
PYTHONPATH).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseAdapter, PredictResult
|
||||
from .utils import mask_to_polygon as _mask_to_polygon
|
||||
|
||||
log = logging.getLogger("sam2-backend.sam2")
|
||||
|
||||
|
||||
class Sam2Adapter(BaseAdapter):
|
||||
def __init__(self) -> None:
|
||||
self._predictor: Optional[object] = None
|
||||
self._device: str = "cpu"
|
||||
|
||||
def load(self, config: str, checkpoint: str, device: str) -> None:
|
||||
# Imports are deferred so the server can start even without the SAM2
|
||||
# deps installed (the model stays in `available: false` until loaded).
|
||||
from sam2.build_sam import build_sam2
|
||||
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
||||
|
||||
log.info("loading SAM2: config=%s ckpt=%s device=%s",
|
||||
config, checkpoint, device)
|
||||
model = build_sam2(config, checkpoint, device=device)
|
||||
self._predictor = SAM2ImagePredictor(model)
|
||||
self._device = device
|
||||
log.info("SAM2 loaded")
|
||||
|
||||
def unload(self) -> None:
|
||||
if self._predictor is None:
|
||||
return
|
||||
self._predictor = None
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
log.info("SAM2 unloaded")
|
||||
|
||||
def predict_bbox(
|
||||
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
|
||||
) -> PredictResult:
|
||||
if self._predictor is None:
|
||||
raise RuntimeError("SAM2 adapter not loaded")
|
||||
|
||||
import torch
|
||||
|
||||
# set_image re-embeds on every call; that's fine for an interactive
|
||||
# one-box-per-request workflow. Batch prompts would want caching.
|
||||
inference_ctx = (
|
||||
torch.inference_mode()
|
||||
if self._device.startswith("cuda") or self._device == "cpu"
|
||||
else torch.no_grad()
|
||||
)
|
||||
with inference_ctx:
|
||||
self._predictor.set_image(image_rgb) # type: ignore[attr-defined]
|
||||
box = np.array(list(bbox_xyxy), dtype=np.float32)
|
||||
masks, scores, _ = self._predictor.predict( # type: ignore[attr-defined]
|
||||
box=box,
|
||||
multimask_output=True,
|
||||
)
|
||||
|
||||
# multimask_output returns 3 candidates — take the highest-scoring one.
|
||||
best_idx = int(np.argmax(scores))
|
||||
best_mask = masks[best_idx]
|
||||
best_score = float(scores[best_idx])
|
||||
polygon = _mask_to_polygon(best_mask)
|
||||
return PredictResult(polygon=polygon, score=best_score)
|
||||
|
||||
|
||||
206
sam2-backend/adapters/sam3_adapter.py
Normal file
206
sam2-backend/adapters/sam3_adapter.py
Normal file
@@ -0,0 +1,206 @@
|
||||
"""SAM3 adapter.
|
||||
|
||||
Mirrors the loading pattern from `files/sam3_server.py` in this repo. Two
|
||||
modes are supported, selected via the model spec's `config` field in
|
||||
models.yaml:
|
||||
|
||||
config: image_predictor — uses `sam3.model_builder.build_sam3_image_model`
|
||||
with a local `.pt` checkpoint. Treats the user's
|
||||
bbox as a visual exemplar and returns the mask
|
||||
whose predicted box has highest IoU with the
|
||||
user box (single-object annotation workflow).
|
||||
|
||||
config: tracker — uses HuggingFace `transformers.Sam3TrackerModel`
|
||||
(SAM2-style PVS on images). `checkpoint` here is
|
||||
a HF model ID like "facebook/sam3" rather than
|
||||
a local .pt. Requires `transformers` installed
|
||||
and network access for first download.
|
||||
|
||||
Both paths return a single polygon + confidence for the best mask, matching
|
||||
the adapter contract the server's `/predict` route expects.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseAdapter, PredictResult
|
||||
from .utils import mask_to_polygon as _mask_to_polygon
|
||||
|
||||
log = logging.getLogger("sam2-backend.sam3")
|
||||
|
||||
|
||||
class Sam3Adapter(BaseAdapter):
|
||||
def __init__(self) -> None:
|
||||
self._mode: str = "image_predictor"
|
||||
self._device: str = "cpu"
|
||||
# image_predictor path
|
||||
self._image_model: Optional[object] = None
|
||||
self._image_processor: Optional[object] = None
|
||||
# tracker path
|
||||
self._tracker_model: Optional[object] = None
|
||||
self._tracker_processor: Optional[object] = None
|
||||
|
||||
def load(self, config: str, checkpoint: str, device: str) -> None:
|
||||
self._device = device
|
||||
mode = (config or "image_predictor").strip().lower()
|
||||
self._mode = mode
|
||||
|
||||
if mode == "tracker":
|
||||
self._load_tracker(checkpoint, device)
|
||||
elif mode in ("image_predictor", "image", "exemplar"):
|
||||
self._load_image_predictor(checkpoint, device)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"unknown SAM3 config mode '{mode}'. "
|
||||
f"expected 'image_predictor' or 'tracker'."
|
||||
)
|
||||
|
||||
def _load_image_predictor(self, checkpoint: str, device: str) -> None:
|
||||
# Deferred imports so the server can start even without sam3 installed
|
||||
# (the model stays unavailable until first use).
|
||||
from sam3.model_builder import build_sam3_image_model
|
||||
from sam3.model.sam3_image_processor import Sam3Processor
|
||||
|
||||
log.info("loading SAM3 image predictor: ckpt=%s device=%s",
|
||||
checkpoint, device)
|
||||
model = (
|
||||
build_sam3_image_model(checkpoint_path=checkpoint)
|
||||
if checkpoint
|
||||
else build_sam3_image_model()
|
||||
)
|
||||
self._image_model = model.to(device).eval()
|
||||
self._image_processor = Sam3Processor(self._image_model)
|
||||
log.info("SAM3 image predictor loaded")
|
||||
|
||||
def _load_tracker(self, hf_id: str, device: str) -> None:
|
||||
from transformers import Sam3TrackerModel, Sam3TrackerProcessor # type: ignore
|
||||
|
||||
model_id = hf_id or "facebook/sam3"
|
||||
log.info("loading SAM3 tracker: hf_id=%s device=%s", model_id, device)
|
||||
self._tracker_processor = Sam3TrackerProcessor.from_pretrained(model_id)
|
||||
self._tracker_model = (
|
||||
Sam3TrackerModel.from_pretrained(model_id).to(device).eval()
|
||||
)
|
||||
log.info("SAM3 tracker loaded")
|
||||
|
||||
def unload(self) -> None:
|
||||
self._image_model = None
|
||||
self._image_processor = None
|
||||
self._tracker_model = None
|
||||
self._tracker_processor = None
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
log.info("SAM3 unloaded")
|
||||
|
||||
def predict_bbox(
|
||||
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
|
||||
) -> PredictResult:
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
pil_image = Image.fromarray(image_rgb)
|
||||
# SAM3 weights are loaded in BFloat16 by `build_sam3_image_model`, but
|
||||
# the processor hands the model Float32 image tensors → matmul dtype
|
||||
# mismatch. Autocast does the input cast for us. Tracker path goes
|
||||
# through HF transformers which already handles dtype internally, but
|
||||
# autocasting is harmless there too.
|
||||
device_type = "cuda" if self._device.startswith("cuda") else "cpu"
|
||||
with torch.inference_mode():
|
||||
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
|
||||
if self._mode == "tracker":
|
||||
return self._predict_tracker(pil_image, bbox_xyxy)
|
||||
return self._predict_image_predictor(pil_image, bbox_xyxy)
|
||||
|
||||
# --------------------------------------------------------------------
|
||||
# Tracker path — HF transformers `Sam3Tracker`
|
||||
# --------------------------------------------------------------------
|
||||
def _predict_tracker(self, image, bbox_xyxy) -> PredictResult:
|
||||
if self._tracker_model is None or self._tracker_processor is None:
|
||||
raise RuntimeError("SAM3 tracker not loaded")
|
||||
x1, y1, x2, y2 = [float(v) for v in bbox_xyxy]
|
||||
inputs = self._tracker_processor(
|
||||
images=image,
|
||||
input_boxes=[[[x1, y1, x2, y2]]], # batch[1] × obj[1] × 4
|
||||
return_tensors="pt",
|
||||
).to(self._device)
|
||||
outputs = self._tracker_model(**inputs, multimask_output=True)
|
||||
# `post_process_masks` lost `reshaped_input_sizes` in newer
|
||||
# transformers — only `masks` and `original_sizes` are accepted now.
|
||||
# Cast to float32 before .numpy() — autocast leaves outputs in
|
||||
# bfloat16, which NumPy can't ingest.
|
||||
masks = self._tracker_processor.post_process_masks(
|
||||
outputs.pred_masks.float().cpu(),
|
||||
original_sizes=inputs["original_sizes"].cpu(),
|
||||
)[0] # (num_masks, H, W)
|
||||
scores = outputs.iou_scores[0].float().cpu().numpy().ravel()
|
||||
masks_np = masks.float().numpy().astype(bool) if masks.dtype != bool else masks.numpy()
|
||||
n_masks = masks_np.shape[0]
|
||||
if n_masks == 0:
|
||||
return PredictResult(polygon=[], score=0.0)
|
||||
# Mask and score counts can disagree (`multimask_output=True` returns 3
|
||||
# IoU scores even when the head only emits 1 mask, e.g. when a
|
||||
# `sam3_video` checkpoint is loaded into `Sam3TrackerModel`). Pick the
|
||||
# best within whichever count is smaller.
|
||||
valid = scores[:n_masks] if scores.size >= n_masks else scores
|
||||
best = int(np.argmax(valid)) if valid.size > 0 else 0
|
||||
best_score = float(valid[best]) if valid.size > 0 else 0.0
|
||||
return PredictResult(
|
||||
polygon=_mask_to_polygon(masks_np[best]),
|
||||
score=best_score,
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------
|
||||
# Image-predictor path — box-as-exemplar, pick highest-IoU result
|
||||
# --------------------------------------------------------------------
|
||||
def _predict_image_predictor(self, image, bbox_xyxy) -> PredictResult:
|
||||
if self._image_processor is None:
|
||||
raise RuntimeError("SAM3 image predictor not loaded")
|
||||
x1, y1, x2, y2 = [float(v) for v in bbox_xyxy]
|
||||
state = self._image_processor.set_image(image) # type: ignore[attr-defined]
|
||||
out = self._image_processor.set_box_prompt( # type: ignore[attr-defined]
|
||||
state=state,
|
||||
boxes=[[x1, y1, x2 - x1, y2 - y1]], # xywh per SAM3 API
|
||||
box_labels=[1],
|
||||
)
|
||||
masks = out.get("masks", [])
|
||||
boxes = out.get("boxes", [])
|
||||
scores = out.get("scores", [])
|
||||
if len(masks) == 0:
|
||||
return PredictResult(polygon=[], score=0.0)
|
||||
# Pick the proposal whose predicted bbox overlaps the user bbox most.
|
||||
user = (x1, y1, x2, y2)
|
||||
ious = [_iou(user, tuple(float(v) for v in b)) for b in boxes]
|
||||
idx = int(np.argmax(ious))
|
||||
mask_np = np.asarray(masks[idx]).astype(bool)
|
||||
return PredictResult(
|
||||
polygon=_mask_to_polygon(mask_np),
|
||||
score=float(scores[idx]),
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# --------------------------------------------------------------------------
|
||||
def _iou(
|
||||
a: tuple[float, float, float, float],
|
||||
b: tuple[float, float, float, float],
|
||||
) -> float:
|
||||
ax1, ay1, ax2, ay2 = a
|
||||
bx1, by1, bx2, by2 = b
|
||||
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
||||
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
||||
iw = max(0.0, ix2 - ix1)
|
||||
ih = max(0.0, iy2 - iy1)
|
||||
inter = iw * ih
|
||||
ua = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - inter
|
||||
return inter / ua if ua > 0 else 0.0
|
||||
|
||||
|
||||
141
sam2-backend/adapters/utils.py
Normal file
141
sam2-backend/adapters/utils.py
Normal file
@@ -0,0 +1,141 @@
|
||||
"""Shared helpers for the SAM-family adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def mask_to_polygon(
|
||||
mask: np.ndarray,
|
||||
*,
|
||||
min_points: int = 4,
|
||||
max_points: int = 10,
|
||||
smooth: bool = True,
|
||||
) -> list[tuple[float, float]]:
|
||||
"""Largest external contour, simplified to a vertex count in
|
||||
[min_points, max_points].
|
||||
|
||||
Pipeline:
|
||||
1. Squeeze + binarise the mask.
|
||||
2. Optional morphological close+open + Gaussian smoothing — this is
|
||||
what turns staircase-jagged SAM mask edges into curves the polygon
|
||||
can actually trace cleanly.
|
||||
3. Pull the full contour (`CHAIN_APPROX_NONE`) so the simplifier has
|
||||
every pixel-edge sample to choose from.
|
||||
4. Binary-search Douglas-Peucker epsilon to hit a vertex count in
|
||||
[min_points, max_points], biased toward more points (more detail).
|
||||
|
||||
Why this matters: the user has to drag every vertex by hand in the
|
||||
annotation UI. Raw `findContours` output is often hundreds of points
|
||||
along a single edge — unusable. Without step 2, even a 15-point output
|
||||
looks choppy because the chosen vertices sit on stair-step pixels.
|
||||
"""
|
||||
m = np.asarray(mask)
|
||||
# Tolerate stray leading singleton dims (some processors return
|
||||
# [1, H, W] or [1, 1, H, W]).
|
||||
m = np.squeeze(m)
|
||||
while m.ndim > 2:
|
||||
m = m[0]
|
||||
if m.ndim != 2 or m.size == 0:
|
||||
return []
|
||||
|
||||
m = m.astype(np.uint8)
|
||||
if m.max() == 1:
|
||||
m = m * 255
|
||||
|
||||
if smooth:
|
||||
m = _smooth_mask(m)
|
||||
|
||||
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
||||
if not contours:
|
||||
return []
|
||||
cnt = max(contours, key=cv2.contourArea)
|
||||
if len(cnt) < 3:
|
||||
return []
|
||||
|
||||
simplified = _simplify_to_band(cnt, min_points=min_points, max_points=max_points)
|
||||
if simplified is None or len(simplified) < 3:
|
||||
return []
|
||||
pts = simplified.reshape(-1, 2)
|
||||
return [(float(x), float(y)) for x, y in pts]
|
||||
|
||||
|
||||
def _smooth_mask(mask_u8: np.ndarray) -> np.ndarray:
|
||||
"""Clean up a binary mask before contour extraction.
|
||||
|
||||
Kernel size scales mildly with the mask's bounding box so a tiny object
|
||||
doesn't get erased and a huge object gets enough smoothing to matter.
|
||||
"""
|
||||
h, w = mask_u8.shape[:2]
|
||||
# Find the actual blob size (not the canvas size) for kernel sizing.
|
||||
ys, xs = np.where(mask_u8 > 0)
|
||||
if xs.size == 0:
|
||||
return mask_u8
|
||||
bw = xs.max() - xs.min() + 1
|
||||
bh = ys.max() - ys.min() + 1
|
||||
blob_diag = max(1.0, float(np.hypot(bw, bh)))
|
||||
# Kernel ~1.5% of blob diagonal, clamped to a usable range.
|
||||
k = int(round(blob_diag * 0.015))
|
||||
k = max(3, min(k | 1, 11)) # odd, between 3 and 11
|
||||
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||||
# Close fills tiny holes / single-pixel notches; open removes specks.
|
||||
m = cv2.morphologyEx(mask_u8, cv2.MORPH_CLOSE, kernel)
|
||||
m = cv2.morphologyEx(m, cv2.MORPH_OPEN, kernel)
|
||||
|
||||
# A light Gaussian + re-threshold smooths the edge by a sub-pixel amount,
|
||||
# which makes the contour follow curves cleanly instead of staircasing.
|
||||
blur_k = max(3, min(k | 1, 9))
|
||||
m = cv2.GaussianBlur(m, (blur_k, blur_k), 0)
|
||||
_, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
|
||||
return m
|
||||
|
||||
|
||||
def _simplify_to_band(
|
||||
contour: np.ndarray,
|
||||
*,
|
||||
min_points: int,
|
||||
max_points: int,
|
||||
iterations: int = 24,
|
||||
) -> np.ndarray | None:
|
||||
"""Binary-search epsilon for `cv2.approxPolyDP` so vertex count lands in
|
||||
[min_points, max_points], biased toward the larger end of the band.
|
||||
|
||||
Strategy:
|
||||
- epsilon=0 → original (many) points
|
||||
- epsilon=arc → degenerates to a single point
|
||||
- smaller epsilon ⇒ more points; larger ⇒ fewer.
|
||||
|
||||
We track the in-band result with the highest vertex count seen
|
||||
(most-detailed pick within the cap) and return it.
|
||||
"""
|
||||
arc = cv2.arcLength(contour, True)
|
||||
if arc == 0:
|
||||
return contour
|
||||
if len(contour) <= max_points:
|
||||
return contour if len(contour) >= min_points else None
|
||||
|
||||
lo, hi = 0.0, arc
|
||||
best: np.ndarray | None = None
|
||||
best_n = -1
|
||||
|
||||
for _ in range(iterations):
|
||||
mid = 0.5 * (lo + hi)
|
||||
approx = cv2.approxPolyDP(contour, mid, True)
|
||||
n = len(approx)
|
||||
if n > max_points:
|
||||
lo = mid
|
||||
elif n < min_points:
|
||||
hi = mid
|
||||
else:
|
||||
if n > best_n:
|
||||
best = approx
|
||||
best_n = n
|
||||
hi = mid
|
||||
if hi - lo < 1e-6 * arc:
|
||||
break
|
||||
|
||||
if best is not None:
|
||||
return best
|
||||
return cv2.approxPolyDP(contour, lo, True)
|
||||
69
sam2-backend/models.yaml
Normal file
69
sam2-backend/models.yaml
Normal file
@@ -0,0 +1,69 @@
|
||||
# Model registry for the SAM backend.
|
||||
#
|
||||
# Add/remove entries freely. Paths may be absolute or relative to this file.
|
||||
# Checkpoints marked as missing on disk are still listed by /models but flagged
|
||||
# `available: false` so the frontend can grey them out.
|
||||
#
|
||||
# `family` selects the adapter code-path. Currently supported:
|
||||
# sam2 — Meta SAM 2 / SAM 2.1 (image predictor, bbox prompts)
|
||||
# sam3 — Meta SAM 3 (image predictor with box-as-exemplar, or HF tracker)
|
||||
#
|
||||
# Per-entry fields:
|
||||
# id — stable key used by /predict?model=<id>
|
||||
# name — human-readable, shown in the Tauri model dropdown
|
||||
# family — adapter to use
|
||||
# config — family-specific (SAM2: hydra yaml path; SAM3: "image_predictor"
|
||||
# or "tracker")
|
||||
# checkpoint — file path (relative paths resolved against this yaml's
|
||||
# directory) OR an opaque identifier when `remote: true`
|
||||
# remote — optional; true = treat checkpoint as a HuggingFace model id
|
||||
# (no on-disk existence check)
|
||||
|
||||
default: sam2.1_hiera_small
|
||||
|
||||
models:
|
||||
- id: sam2.1_hiera_tiny
|
||||
name: "SAM 2.1 · Hiera Tiny"
|
||||
family: sam2
|
||||
config: configs/sam2.1/sam2.1_hiera_t.yaml
|
||||
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_tiny.pt
|
||||
|
||||
- id: sam2.1_hiera_small
|
||||
name: "SAM 2.1 · Hiera Small"
|
||||
family: sam2
|
||||
config: configs/sam2.1/sam2.1_hiera_s.yaml
|
||||
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_small.pt
|
||||
|
||||
- id: sam2.1_hiera_base_plus
|
||||
name: "SAM 2.1 · Hiera Base+"
|
||||
family: sam2
|
||||
config: configs/sam2.1/sam2.1_hiera_b+.yaml
|
||||
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_base_plus.pt
|
||||
|
||||
- id: sam2.1_hiera_large
|
||||
name: "SAM 2.1 · Hiera Large"
|
||||
family: sam2
|
||||
config: configs/sam2.1/sam2.1_hiera_l.yaml
|
||||
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_large.pt
|
||||
|
||||
# ---- SAM 3 -----------------------------------------------------------
|
||||
# Local image-predictor path — uses the .pt weights in this repo. Box is
|
||||
# treated as a visual exemplar; the adapter picks the mask whose predicted
|
||||
# bbox best overlaps the user bbox (classic single-instance annotation).
|
||||
- id: sam3
|
||||
name: "SAM 3 · image predictor"
|
||||
family: sam3
|
||||
config: image_predictor
|
||||
checkpoint: ../files/sam3_weights/sam3.pt
|
||||
|
||||
# Optional HuggingFace tracker path. Requires `pip install transformers`
|
||||
# and pulls weights from HF hub on first use. Gives cleaner single-box
|
||||
# → single-mask results than the exemplar path but needs a network round
|
||||
# trip the first time. Disabled by leaving `remote: false` — flip to
|
||||
# `remote: true` once your machine has transformers installed.
|
||||
# - id: sam3_tracker
|
||||
# name: "SAM 3 · tracker (HuggingFace)"
|
||||
# family: sam3
|
||||
# config: tracker
|
||||
# checkpoint: facebook/sam3
|
||||
# remote: true
|
||||
141
sam2-backend/registry.py
Normal file
141
sam2-backend/registry.py
Normal file
@@ -0,0 +1,141 @@
|
||||
"""Model registry. Reads models.yaml, lazy-loads adapters, caches instances.
|
||||
|
||||
Only one adapter is resident at a time by default — switching models unloads
|
||||
the previous one to keep GPU memory bounded. Set SAM_BACKEND_KEEP_ALL=1 to
|
||||
disable eviction (useful if you have VRAM headroom and switch often).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import yaml
|
||||
|
||||
from adapters import ADAPTERS, BaseAdapter, PredictResult
|
||||
|
||||
log = logging.getLogger("sam2-backend.registry")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelSpec:
|
||||
id: str
|
||||
name: str
|
||||
family: str
|
||||
config: str
|
||||
checkpoint: str
|
||||
# If true, `checkpoint` is treated as an opaque identifier (e.g. a
|
||||
# HuggingFace model id like "facebook/sam3") rather than a file on disk.
|
||||
# Availability check is skipped for remote entries.
|
||||
remote: bool = False
|
||||
available: bool = False
|
||||
# Resolved absolute paths for logging / responses.
|
||||
checkpoint_abs: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Registry:
|
||||
specs: dict[str, ModelSpec] = field(default_factory=dict)
|
||||
default_id: Optional[str] = None
|
||||
device: str = "cpu"
|
||||
keep_all: bool = False
|
||||
_loaded: dict[str, BaseAdapter] = field(default_factory=dict)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def list_specs(self) -> list[ModelSpec]:
|
||||
return list(self.specs.values())
|
||||
|
||||
def resolve(self, model_id: Optional[str]) -> ModelSpec:
|
||||
if not model_id:
|
||||
if self.default_id is None:
|
||||
raise KeyError("no default model configured")
|
||||
model_id = self.default_id
|
||||
if model_id not in self.specs:
|
||||
raise KeyError(f"unknown model: {model_id}")
|
||||
return self.specs[model_id]
|
||||
|
||||
def get_adapter(self, model_id: str) -> BaseAdapter:
|
||||
"""Returns a loaded adapter for `model_id`, loading on first use."""
|
||||
with self._lock:
|
||||
if model_id in self._loaded:
|
||||
return self._loaded[model_id]
|
||||
spec = self.resolve(model_id)
|
||||
if not spec.available:
|
||||
raise FileNotFoundError(
|
||||
f"checkpoint missing: {spec.checkpoint_abs}"
|
||||
)
|
||||
cls = ADAPTERS.get(spec.family)
|
||||
if cls is None:
|
||||
raise ValueError(
|
||||
f"no adapter registered for family '{spec.family}'"
|
||||
)
|
||||
if not self.keep_all:
|
||||
# Evict whichever model is currently resident.
|
||||
for old_id, adapter in list(self._loaded.items()):
|
||||
log.info("evicting %s to make room for %s", old_id, model_id)
|
||||
adapter.unload()
|
||||
del self._loaded[old_id]
|
||||
adapter = cls()
|
||||
adapter.load(spec.config, spec.checkpoint_abs, self.device)
|
||||
self._loaded[model_id] = adapter
|
||||
return adapter
|
||||
|
||||
def predict_bbox(self, model_id: Optional[str], image_rgb, bbox_xyxy) -> tuple[ModelSpec, PredictResult]:
|
||||
spec = self.resolve(model_id)
|
||||
adapter = self.get_adapter(spec.id)
|
||||
# Serialize inference: most of these adapters aren't reentrant.
|
||||
with self._lock:
|
||||
result = adapter.predict_bbox(image_rgb, bbox_xyxy)
|
||||
return spec, result
|
||||
|
||||
|
||||
def load_registry(models_yaml: Path, device: str) -> Registry:
|
||||
with open(models_yaml, "r") as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
base_dir = models_yaml.parent
|
||||
keep_all = os.environ.get("SAM_BACKEND_KEEP_ALL", "0") == "1"
|
||||
reg = Registry(device=device, keep_all=keep_all)
|
||||
reg.default_id = data.get("default")
|
||||
|
||||
for entry in data.get("models", []):
|
||||
remote = bool(entry.get("remote", False))
|
||||
ckpt_str = entry["checkpoint"]
|
||||
if remote:
|
||||
# Opaque identifier — leave as-is, mark available so the UI
|
||||
# shows it. The adapter resolves/downloads on first use.
|
||||
ckpt_abs = ckpt_str
|
||||
available = True
|
||||
else:
|
||||
ckpt = Path(ckpt_str)
|
||||
ckpt_path = (ckpt if ckpt.is_absolute() else (base_dir / ckpt)).resolve()
|
||||
ckpt_abs = str(ckpt_path)
|
||||
available = ckpt_path.is_file()
|
||||
spec = ModelSpec(
|
||||
id=entry["id"],
|
||||
name=entry["name"],
|
||||
family=entry["family"],
|
||||
config=entry["config"],
|
||||
checkpoint=ckpt_str,
|
||||
remote=remote,
|
||||
available=available,
|
||||
checkpoint_abs=ckpt_abs,
|
||||
)
|
||||
reg.specs[spec.id] = spec
|
||||
log.info(
|
||||
"registered model %s (available=%s, ckpt=%s)",
|
||||
spec.id, spec.available, spec.checkpoint_abs,
|
||||
)
|
||||
|
||||
# Fall back the default to the first *available* model if the configured
|
||||
# default's checkpoint is missing.
|
||||
if reg.default_id and not reg.specs.get(reg.default_id, ModelSpec("", "", "", "", "")).available:
|
||||
for s in reg.specs.values():
|
||||
if s.available:
|
||||
reg.default_id = s.id
|
||||
break
|
||||
return reg
|
||||
13
sam2-backend/requirements.txt
Normal file
13
sam2-backend/requirements.txt
Normal file
@@ -0,0 +1,13 @@
|
||||
# Runtime deps for the SAM2 backend HTTP server.
|
||||
# Torch is intentionally NOT pinned here — install the right build for your
|
||||
# CUDA/CPU first (see README.md), then `pip install -r requirements.txt`.
|
||||
|
||||
fastapi>=0.110
|
||||
uvicorn[standard]>=0.27
|
||||
pydantic>=2.6
|
||||
pyyaml>=6.0
|
||||
pillow>=10.0
|
||||
numpy>=1.24,<2.0
|
||||
opencv-python-headless>=4.8
|
||||
hydra-core>=1.3
|
||||
iopath>=0.1.10
|
||||
32
sam2-backend/run.sh
Executable file
32
sam2-backend/run.sh
Executable file
@@ -0,0 +1,32 @@
|
||||
#!/usr/bin/env bash
|
||||
# Convenience launcher for the SAM2 backend.
|
||||
# Assumes a venv at ./.venv with requirements.txt + torch + sam2 installed.
|
||||
# See README.md for install steps.
|
||||
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
if [ -d ".venv" ]; then
|
||||
# shellcheck disable=SC1091
|
||||
source .venv/bin/activate
|
||||
fi
|
||||
|
||||
HOST="${HOST:-0.0.0.0}"
|
||||
PORT="${PORT:-9090}"
|
||||
DEVICE="${DEVICE:-auto}"
|
||||
|
||||
# Privacy: disable HF/transformers telemetry, and (after weights are cached)
|
||||
# refuse any outbound calls to huggingface.co. Inference runs fully local
|
||||
# regardless; these just stop the cosmetic pings on startup.
|
||||
export HF_HUB_DISABLE_TELEMETRY="${HF_HUB_DISABLE_TELEMETRY:-1}"
|
||||
export DO_NOT_TRACK="${DO_NOT_TRACK:-1}"
|
||||
# Flip both OFFLINE flags to 1 once you've done the first `from_pretrained`
|
||||
# and the weights are cached under ~/.cache/huggingface/. Leaving them 0
|
||||
# the first time lets the model download. After that, set them 1 to fully
|
||||
# air-gap and silence the HEAD-request revalidation pings on every startup.
|
||||
# HF_HUB_OFFLINE — used by huggingface_hub
|
||||
# TRANSFORMERS_OFFLINE — used by transformers itself (some legacy paths)
|
||||
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-0}"
|
||||
export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-0}"
|
||||
|
||||
exec python server.py --host "$HOST" --port "$PORT" --device "$DEVICE" "$@"
|
||||
345
sam2-backend/server.py
Normal file
345
sam2-backend/server.py
Normal file
@@ -0,0 +1,345 @@
|
||||
"""FastAPI server exposing SAM-family predictors over HTTP.
|
||||
|
||||
Implements the Label-Studio-ML `/predict` contract the SAM-Tool Tauri app
|
||||
already speaks, and an LS-orthogonal `/models` endpoint so the frontend can
|
||||
offer a model picker.
|
||||
|
||||
Run:
|
||||
python server.py --host 0.0.0.0 --port 9090 --device cuda
|
||||
|
||||
Endpoints:
|
||||
GET / → liveness + backend info
|
||||
GET /models → registered models + availability
|
||||
POST /predict → LS-ML protocol, accepts optional `?model=<id>` query
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException, Query
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel
|
||||
|
||||
from registry import Registry, load_registry
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Logging
|
||||
# --------------------------------------------------------------------------
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||||
stream=sys.stdout,
|
||||
)
|
||||
log = logging.getLogger("sam2-backend.server")
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# App + registry. Registry is loaded inside a FastAPI lifespan so it works
|
||||
# regardless of how uvicorn ends up importing this module — `python
|
||||
# server.py` runs us as `__main__` and uvicorn then re-imports as `server`,
|
||||
# so initialising at module top-level (or only in main()) leaves the live
|
||||
# copy uninitialised.
|
||||
# --------------------------------------------------------------------------
|
||||
REGISTRY: Registry | None = None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI):
|
||||
global REGISTRY
|
||||
if REGISTRY is None:
|
||||
models_path = os.environ.get(
|
||||
"SAM_BACKEND_MODELS",
|
||||
str(Path(__file__).parent / "models.yaml"),
|
||||
)
|
||||
device = _pick_device(os.environ.get("SAM_BACKEND_DEVICE", "auto"))
|
||||
log.info("lifespan: loading registry (models=%s device=%s)",
|
||||
models_path, device)
|
||||
REGISTRY = load_registry(Path(models_path), device=device)
|
||||
log.info(
|
||||
"lifespan: %d model specs (%d available), default=%s",
|
||||
len(REGISTRY.specs),
|
||||
sum(1 for s in REGISTRY.specs.values() if s.available),
|
||||
REGISTRY.default_id,
|
||||
)
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(title="SAM Backend", version="0.1.0", lifespan=lifespan)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
def _registry() -> Registry:
|
||||
if REGISTRY is None:
|
||||
raise RuntimeError("registry not initialised")
|
||||
return REGISTRY
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Schemas
|
||||
# --------------------------------------------------------------------------
|
||||
class ModelInfo(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
family: str
|
||||
available: bool
|
||||
checkpoint: str
|
||||
|
||||
|
||||
class ModelListResponse(BaseModel):
|
||||
default: str | None
|
||||
models: list[ModelInfo]
|
||||
|
||||
|
||||
class HealthResponse(BaseModel):
|
||||
status: str
|
||||
device: str
|
||||
models: int
|
||||
default: str | None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Routes
|
||||
# --------------------------------------------------------------------------
|
||||
@app.get("/", response_model=HealthResponse)
|
||||
def health() -> HealthResponse:
|
||||
reg = _registry()
|
||||
return HealthResponse(
|
||||
status="ok",
|
||||
device=reg.device,
|
||||
models=len(reg.specs),
|
||||
default=reg.default_id,
|
||||
)
|
||||
|
||||
|
||||
@app.get("/models", response_model=ModelListResponse)
|
||||
def list_models() -> ModelListResponse:
|
||||
reg = _registry()
|
||||
return ModelListResponse(
|
||||
default=reg.default_id,
|
||||
models=[
|
||||
ModelInfo(
|
||||
id=s.id,
|
||||
name=s.name,
|
||||
family=s.family,
|
||||
available=s.available,
|
||||
checkpoint=s.checkpoint_abs,
|
||||
)
|
||||
for s in reg.list_specs()
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@app.post("/predict")
|
||||
def predict(payload: dict, model: str | None = Query(default=None)) -> dict:
|
||||
"""Label-Studio-ML protocol.
|
||||
|
||||
Body shape (simplified — we ignore fields we don't need):
|
||||
{
|
||||
"tasks": [{ "data": { "image": "data:image/jpeg;base64,..." } }],
|
||||
"params": {
|
||||
"context": {
|
||||
"result": [{
|
||||
"original_width": W, "original_height": H,
|
||||
"value": { "x": %, "y": %, "width": %, "height": %, ... }
|
||||
}]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Response shape (what sam-tool-tauri / LS Studio expects):
|
||||
{
|
||||
"results": [{
|
||||
"model_version": "<id>",
|
||||
"result": [{
|
||||
"original_width": W, "original_height": H,
|
||||
"value": {
|
||||
"points": [[x%, y%], ...],
|
||||
"closed": true,
|
||||
"polygonlabels": ["Object"]
|
||||
},
|
||||
"from_name": "polygon", "to_name": "image",
|
||||
"type": "polygonlabels"
|
||||
}]
|
||||
}]
|
||||
}
|
||||
"""
|
||||
reg = _registry()
|
||||
|
||||
try:
|
||||
tasks = payload.get("tasks") or []
|
||||
if not tasks:
|
||||
raise HTTPException(400, "tasks[] missing or empty")
|
||||
data = tasks[0].get("data") or {}
|
||||
image_field = data.get("image")
|
||||
if not isinstance(image_field, str):
|
||||
raise HTTPException(400, "tasks[0].data.image missing")
|
||||
image_rgb = _decode_image(image_field)
|
||||
|
||||
ctx_result = (
|
||||
payload.get("params", {})
|
||||
.get("context", {})
|
||||
.get("result", [])
|
||||
)
|
||||
if not ctx_result:
|
||||
raise HTTPException(400, "params.context.result[] missing (bbox prompt)")
|
||||
prompt = ctx_result[0]
|
||||
value = prompt.get("value", {})
|
||||
orig_w = int(prompt.get("original_width") or image_rgb.shape[1])
|
||||
orig_h = int(prompt.get("original_height") or image_rgb.shape[0])
|
||||
x_pct = float(value.get("x", 0.0))
|
||||
y_pct = float(value.get("y", 0.0))
|
||||
w_pct = float(value.get("width", 0.0))
|
||||
h_pct = float(value.get("height", 0.0))
|
||||
|
||||
# Percentages → absolute pixel xyxy, clamped to image bounds.
|
||||
x1 = max(0, int(round(x_pct / 100.0 * orig_w)))
|
||||
y1 = max(0, int(round(y_pct / 100.0 * orig_h)))
|
||||
x2 = min(orig_w, int(round((x_pct + w_pct) / 100.0 * orig_w)))
|
||||
y2 = min(orig_h, int(round((y_pct + h_pct) / 100.0 * orig_h)))
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
raise HTTPException(400, "degenerate bbox after clamping")
|
||||
|
||||
# Model selection: query param wins, then payload.params.model, then default.
|
||||
model_id = (
|
||||
model
|
||||
or payload.get("params", {}).get("model")
|
||||
or payload.get("params", {}).get("model_id")
|
||||
)
|
||||
# Resolve up-front so an unknown id is a clean 400; any error from
|
||||
# the actual predict path falls into the generic 500 with traceback
|
||||
# below (and is no longer mis-reported as "unknown model").
|
||||
try:
|
||||
reg.resolve(model_id)
|
||||
except KeyError as e:
|
||||
raise HTTPException(400, f"unknown model: {e}")
|
||||
|
||||
spec, result = reg.predict_bbox(model_id, image_rgb, (x1, y1, x2, y2))
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(503, str(e))
|
||||
except Exception as e:
|
||||
log.exception("predict failed")
|
||||
raise HTTPException(500, f"{type(e).__name__}: {e}")
|
||||
|
||||
if not result.polygon:
|
||||
raise HTTPException(500, "no polygon returned from model")
|
||||
|
||||
# Convert pixel polygon back to percentages for LS-ML round-trip.
|
||||
points_pct = [
|
||||
[px / orig_w * 100.0, py / orig_h * 100.0]
|
||||
for (px, py) in result.polygon
|
||||
]
|
||||
|
||||
return {
|
||||
"results": [{
|
||||
"model_version": spec.id,
|
||||
"score": result.score,
|
||||
"result": [{
|
||||
"original_width": orig_w,
|
||||
"original_height": orig_h,
|
||||
"image_rotation": 0,
|
||||
"value": {
|
||||
"points": points_pct,
|
||||
"closed": True,
|
||||
"polygonlabels": ["Object"],
|
||||
},
|
||||
"id": "sam_poly_0",
|
||||
"from_name": "polygon",
|
||||
"to_name": "image",
|
||||
"type": "polygonlabels",
|
||||
}],
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# --------------------------------------------------------------------------
|
||||
def _decode_image(field: str) -> np.ndarray:
|
||||
"""Accept `data:<mime>;base64,...` or a bare base64 string."""
|
||||
b64 = field
|
||||
if field.startswith("data:"):
|
||||
comma = field.find(",")
|
||||
if comma == -1:
|
||||
raise HTTPException(400, "malformed data URL")
|
||||
b64 = field[comma + 1:]
|
||||
try:
|
||||
raw = base64.b64decode(b64)
|
||||
except Exception as e:
|
||||
raise HTTPException(400, f"base64 decode failed: {e}")
|
||||
try:
|
||||
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(400, f"image decode failed: {e}")
|
||||
return np.array(img)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Entrypoint
|
||||
# --------------------------------------------------------------------------
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=9090)
|
||||
parser.add_argument(
|
||||
"--models",
|
||||
default=str(Path(__file__).parent / "models.yaml"),
|
||||
help="Path to models.yaml registry file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device",
|
||||
default="auto",
|
||||
help="'auto' (prefer cuda), 'cuda', 'cpu', or specific cuda:N",
|
||||
)
|
||||
parser.add_argument("--reload", action="store_true",
|
||||
help="Enable uvicorn auto-reload (dev only)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Hand config to the worker module via env — the lifespan reads these
|
||||
# whether uvicorn imports this file as `__main__` or as `server`.
|
||||
os.environ["SAM_BACKEND_MODELS"] = args.models
|
||||
os.environ["SAM_BACKEND_DEVICE"] = args.device
|
||||
|
||||
log.info("starting uvicorn on %s:%d (device=%s, models=%s)",
|
||||
args.host, args.port, args.device, args.models)
|
||||
uvicorn.run(
|
||||
"server:app",
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
reload=args.reload,
|
||||
log_level="info",
|
||||
)
|
||||
|
||||
|
||||
def _pick_device(preference: str) -> str:
|
||||
if preference != "auto":
|
||||
return preference
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
except ImportError:
|
||||
pass
|
||||
return "cpu"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1
sam2/sam2
Submodule
1
sam2/sam2
Submodule
Submodule sam2/sam2 added at 2b90b9f5ce
500
sam2/sam_extract.py
Normal file
500
sam2/sam_extract.py
Normal file
@@ -0,0 +1,500 @@
|
||||
"""
|
||||
Runs on video and tracks objects across time, generating masks and polygons for multiple frames.
|
||||
Takes one annotation of each asset and fetches the frames before and after it.
|
||||
"""
|
||||
import os
|
||||
import cv2
|
||||
import ast
|
||||
import json
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from collections import deque
|
||||
from scipy.optimize import linear_sum_assignment
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
# =========================================================
|
||||
# CONFIG
|
||||
# =========================================================
|
||||
VIDEO_PATH = "/media/bs5/New Volume1/Brinda/Dataset/GPS_location/videos/2026_0330_094759_F.MP4"
|
||||
JSON_PATH = "/media/bs5/New Volume1/Brinda/Dataset/GPS_location/output/2026_0330_094759_F.json"
|
||||
OUTPUT_DIR = "/media/bs5/New Volume1/Brinda/Dataset/GPS_location/sam/"
|
||||
|
||||
FRAME_GAP = 5
|
||||
NUM_FRAMES = 5
|
||||
|
||||
SAM_CHECKPOINT = "sam_vit_h_4b8939.pth"
|
||||
MODEL_TYPE = "vit_h"
|
||||
|
||||
|
||||
# =========================================================
|
||||
# OUTPUT DIRS
|
||||
# =========================================================
|
||||
IMG_DIR = os.path.join(OUTPUT_DIR, "images")
|
||||
MASK_DIR = os.path.join(OUTPUT_DIR, "masks")
|
||||
EXTRACT_DIR = os.path.join(OUTPUT_DIR, "masked")
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
os.makedirs(IMG_DIR, exist_ok=True)
|
||||
os.makedirs(MASK_DIR, exist_ok=True)
|
||||
os.makedirs(EXTRACT_DIR, exist_ok=True)
|
||||
|
||||
5
|
||||
# ================== KALMAN FILTER ==================
|
||||
|
||||
class KalmanFilter:
|
||||
"""Kalman filter for smoother track predictions"""
|
||||
def __init__(self):
|
||||
self.kf = cv2.KalmanFilter(8, 4)
|
||||
|
||||
# State transition matrix
|
||||
self.kf.transitionMatrix = np.eye(8, dtype=np.float32)
|
||||
self.kf.transitionMatrix[0, 4] = 1
|
||||
self.kf.transitionMatrix[1, 5] = 1
|
||||
self.kf.transitionMatrix[2, 6] = 1
|
||||
self.kf.transitionMatrix[3, 7] = 1
|
||||
|
||||
# Measurement matrix
|
||||
self.kf.measurementMatrix = np.zeros((4, 8), dtype=np.float32)
|
||||
self.kf.measurementMatrix[0, 0] = 1
|
||||
self.kf.measurementMatrix[1, 1] = 1
|
||||
self.kf.measurementMatrix[2, 2] = 1
|
||||
self.kf.measurementMatrix[3, 3] = 1
|
||||
|
||||
self.kf.processNoiseCov = np.eye(8, dtype=np.float32) * 0.01
|
||||
self.kf.measurementNoiseCov = np.eye(4, dtype=np.float32) * 0.5
|
||||
self.kf.errorCovPost = np.eye(8, dtype=np.float32)
|
||||
|
||||
def predict(self):
|
||||
prediction = self.kf.predict()
|
||||
return prediction[:4].flatten()
|
||||
|
||||
def update(self, measurement):
|
||||
self.kf.correct(np.array(measurement, dtype=np.float32).reshape(4, 1))
|
||||
|
||||
|
||||
# ================== TRACK CLASS ==================
|
||||
|
||||
class Track:
|
||||
def __init__(self, tlwh, score, track_id, class_id=0, class_name='unknown'):
|
||||
self.tlwh = np.asarray(tlwh, dtype=np.float32)
|
||||
self.score = score
|
||||
self.track_id = track_id
|
||||
self.class_id = class_id
|
||||
self.class_name = class_name
|
||||
self.time_since_update = 0
|
||||
self.hits = 1
|
||||
self.age = 1
|
||||
|
||||
# Kalman filter for motion prediction
|
||||
self.kf = KalmanFilter()
|
||||
self.kf.kf.statePost = np.array([*tlwh, 0, 0, 0, 0], dtype=np.float32).reshape(8, 1)
|
||||
|
||||
# Track confidence history for stability
|
||||
self.confidence_history = deque([score], maxlen=10)
|
||||
self.avg_confidence = score
|
||||
|
||||
def predict(self):
|
||||
"""Predict next position using Kalman filter"""
|
||||
self.age += 1
|
||||
self.time_since_update += 1
|
||||
predicted = self.kf.predict()
|
||||
|
||||
if self.time_since_update < 5:
|
||||
self.tlwh = predicted
|
||||
|
||||
def update(self, tlwh, score, class_id, class_name):
|
||||
"""Update track with new detection"""
|
||||
self.tlwh = np.asarray(tlwh, dtype=np.float32)
|
||||
self.score = score
|
||||
self.class_id = class_id
|
||||
self.class_name = class_name
|
||||
self.time_since_update = 0
|
||||
self.hits += 1
|
||||
|
||||
self.confidence_history.append(score)
|
||||
self.avg_confidence = np.mean(list(self.confidence_history))
|
||||
|
||||
self.kf.update(tlwh)
|
||||
|
||||
def is_stable(self):
|
||||
"""Check if track is stable enough to display"""
|
||||
return (self.hits >= 3 and
|
||||
self.avg_confidence > 0.3 and
|
||||
self.time_since_update == 0)
|
||||
|
||||
|
||||
# ================== BYTETRACK TRACKER ==================
|
||||
|
||||
class ByteTrackEnhanced:
|
||||
"""
|
||||
Enhanced ByteTrack with per-class track ID sequencing
|
||||
"""
|
||||
def __init__(self, track_thresh=0.4, match_thresh=0.5, buffer_size=60,
|
||||
high_thresh=0.5, low_thresh=0.1, class_names=None):
|
||||
self.track_thresh = track_thresh
|
||||
self.high_thresh = high_thresh
|
||||
self.low_thresh = low_thresh
|
||||
self.match_thresh = match_thresh
|
||||
self.buffer_size = buffer_size
|
||||
|
||||
self.tracks = []
|
||||
self.class_names = class_names if class_names else []
|
||||
|
||||
# Per-class track ID counters
|
||||
self.next_id_per_class = {}
|
||||
|
||||
def get_next_id(self, class_id):
|
||||
"""Get next track ID for specific class"""
|
||||
if class_id not in self.next_id_per_class:
|
||||
self.next_id_per_class[class_id] = 1
|
||||
|
||||
track_id = self.next_id_per_class[class_id]
|
||||
self.next_id_per_class[class_id] += 1
|
||||
return track_id
|
||||
|
||||
def get_class_name(self, class_id):
|
||||
"""Get class name from class ID"""
|
||||
if class_id < len(self.class_names):
|
||||
return self.class_names[class_id]
|
||||
return f'class_{class_id}'
|
||||
|
||||
def iou(self, a, b):
|
||||
"""Calculate IoU between two boxes"""
|
||||
ax, ay, aw, ah = a
|
||||
bx, by, bw, bh = b
|
||||
ax2, ay2 = ax + aw, ay + ah
|
||||
bx2, by2 = bx + bw, by + bh
|
||||
|
||||
inter_x1 = max(ax, bx)
|
||||
inter_y1 = max(ay, by)
|
||||
inter_x2 = min(ax2, bx2)
|
||||
inter_y2 = min(ay2, by2)
|
||||
|
||||
inter = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
|
||||
union = aw * ah + bw * bh - inter
|
||||
|
||||
return inter / union if union > 0 else 0
|
||||
|
||||
def iou_distance(self, tracks, detections):
|
||||
"""Compute IoU distance matrix"""
|
||||
iou_mat = np.zeros((len(tracks), len(detections)), dtype=np.float32)
|
||||
|
||||
for i, t in enumerate(tracks):
|
||||
for j, d in enumerate(detections):
|
||||
iou = self.iou(t.tlwh, d[:4])
|
||||
|
||||
# Penalize class mismatch
|
||||
if len(d) > 5 and t.class_id != int(d[5]):
|
||||
iou *= 0.3
|
||||
|
||||
# Boost for active tracks
|
||||
if t.time_since_update == 0:
|
||||
iou *= 1.1
|
||||
|
||||
iou_mat[i, j] = iou
|
||||
|
||||
return iou_mat
|
||||
|
||||
def linear_assignment(self, cost_matrix, thresh):
|
||||
"""Perform linear assignment with threshold"""
|
||||
if cost_matrix.size == 0:
|
||||
return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
|
||||
|
||||
matches, unmatched_a, unmatched_b = [], [], []
|
||||
row_ind, col_ind = linear_sum_assignment(-cost_matrix)
|
||||
|
||||
for r, c in zip(row_ind, col_ind):
|
||||
if cost_matrix[r, c] < thresh:
|
||||
unmatched_a.append(r)
|
||||
unmatched_b.append(c)
|
||||
else:
|
||||
matches.append([r, c])
|
||||
|
||||
unmatched_a += list(set(range(cost_matrix.shape[0])) - set(row_ind))
|
||||
unmatched_b += list(set(range(cost_matrix.shape[1])) - set(col_ind))
|
||||
|
||||
return np.array(matches), tuple(unmatched_a), tuple(unmatched_b)
|
||||
|
||||
def update(self, detections):
|
||||
"""Two-stage association with stability fixes"""
|
||||
# Predict all tracks
|
||||
for t in self.tracks:
|
||||
t.predict()
|
||||
|
||||
confirmed_tracks = [t for t in self.tracks if t.hits >= 2]
|
||||
unconfirmed_tracks = [t for t in self.tracks if t.hits < 2]
|
||||
|
||||
# Separate high and low confidence detections
|
||||
high_dets = [d for d in detections if d[4] >= self.high_thresh]
|
||||
low_dets = [d for d in detections if self.low_thresh <= d[4] < self.high_thresh]
|
||||
|
||||
# First association
|
||||
unmatched_tracks = []
|
||||
unmatched_dets = []
|
||||
|
||||
if len(confirmed_tracks) > 0 and len(high_dets) > 0:
|
||||
iou_mat = self.iou_distance(confirmed_tracks, high_dets)
|
||||
matches, u_tracks, u_dets = self.linear_assignment(iou_mat, self.match_thresh)
|
||||
|
||||
for m in matches:
|
||||
track_idx, det_idx = m[0], m[1]
|
||||
class_id = int(high_dets[det_idx][5]) if len(high_dets[det_idx]) > 5 else 0
|
||||
class_name = self.get_class_name(class_id)
|
||||
confirmed_tracks[track_idx].update(high_dets[det_idx][:4],
|
||||
high_dets[det_idx][4],
|
||||
class_id,
|
||||
class_name)
|
||||
|
||||
unmatched_tracks = [confirmed_tracks[i] for i in u_tracks]
|
||||
unmatched_dets = [high_dets[i] for i in u_dets]
|
||||
else:
|
||||
unmatched_tracks = confirmed_tracks
|
||||
unmatched_dets = high_dets
|
||||
|
||||
# Second association with low confidence
|
||||
if len(unmatched_tracks) > 0 and len(low_dets) > 0:
|
||||
iou_mat = self.iou_distance(unmatched_tracks, low_dets)
|
||||
matches, u_tracks, _ = self.linear_assignment(iou_mat, self.match_thresh * 0.6)
|
||||
|
||||
for m in matches:
|
||||
track_idx, det_idx = m[0], m[1]
|
||||
class_id = int(low_dets[det_idx][5]) if len(low_dets[det_idx]) > 5 else 0
|
||||
class_name = self.get_class_name(class_id)
|
||||
unmatched_tracks[track_idx].update(low_dets[det_idx][:4],
|
||||
low_dets[det_idx][4],
|
||||
class_id,
|
||||
class_name)
|
||||
|
||||
unmatched_tracks = [unmatched_tracks[i] for i in u_tracks]
|
||||
|
||||
# Deal with unconfirmed tracks
|
||||
if len(unmatched_dets) > 0:
|
||||
for t in unconfirmed_tracks:
|
||||
iou_mat = self.iou_distance([t], unmatched_dets)
|
||||
if iou_mat.size > 0:
|
||||
best_match = np.argmax(iou_mat[0])
|
||||
if iou_mat[0, best_match] > self.match_thresh * 0.5:
|
||||
class_id = int(unmatched_dets[best_match][5]) if len(unmatched_dets[best_match]) > 5 else 0
|
||||
class_name = self.get_class_name(class_id)
|
||||
t.update(unmatched_dets[best_match][:4],
|
||||
unmatched_dets[best_match][4],
|
||||
class_id,
|
||||
class_name)
|
||||
unmatched_dets.pop(best_match)
|
||||
|
||||
# Create new tracks
|
||||
for d in unmatched_dets:
|
||||
if d[4] > self.track_thresh:
|
||||
overlaps = False
|
||||
for t in self.tracks:
|
||||
if self.iou(t.tlwh, d[:4]) > 0.3:
|
||||
overlaps = True
|
||||
break
|
||||
|
||||
if not overlaps:
|
||||
class_id = int(d[5]) if len(d) > 5 else 0
|
||||
class_name = self.get_class_name(class_id)
|
||||
track_id = self.get_next_id(class_id)
|
||||
self.tracks.append(Track(d[:4], d[4], track_id, class_id, class_name))
|
||||
|
||||
# Remove old tracks
|
||||
self.tracks = [t for t in self.tracks if t.time_since_update < self.buffer_size]
|
||||
|
||||
return self.tracks
|
||||
|
||||
|
||||
|
||||
# =========================================================
|
||||
# CLEAN CLASS NAME
|
||||
# =========================================================
|
||||
def clean_name(name):
|
||||
return name.replace("LEFT_", "").replace("RIGHT_", "").replace("_Start", "").replace("_Stop", "")
|
||||
|
||||
# =========================================================
|
||||
# PARSE JSON
|
||||
# =========================================================
|
||||
def parse_json(path):
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
try:
|
||||
data = json.loads(content) # proper JSON
|
||||
except json.JSONDecodeError:
|
||||
data = ast.literal_eval(content) # handles single quotes etc.
|
||||
|
||||
parsed = {}
|
||||
|
||||
for asset in data["Assets"]:
|
||||
cls = clean_name(asset[0])
|
||||
frame = int(asset[2])
|
||||
|
||||
x1, y1 = asset[3]
|
||||
x2, y2 = asset[4]
|
||||
|
||||
if frame not in parsed:
|
||||
parsed[frame] = []
|
||||
|
||||
parsed[frame].append({
|
||||
"class": cls,
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
"id": asset[1]
|
||||
})
|
||||
|
||||
return parsed
|
||||
|
||||
# =========================================================
|
||||
# SAM
|
||||
# =========================================================
|
||||
def load_sam():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
|
||||
sam.to(device)
|
||||
return SamPredictor(sam)
|
||||
|
||||
def get_mask(image, bbox, predictor):
|
||||
masks, scores, _ = predictor.predict(
|
||||
box=np.array(bbox),
|
||||
multimask_output=True
|
||||
)
|
||||
return masks[np.argmax(scores)]
|
||||
|
||||
# =========================================================
|
||||
# TEMPORAL FRAMES"""
|
||||
Reads YOLO label--> convert to pixel format --> Feed the box to Sam ---> Get mask
|
||||
"""
|
||||
# =========================================================
|
||||
def get_frames(center, total):
|
||||
frames = []
|
||||
for i in range(-NUM_FRAMES, NUM_FRAMES + 1):
|
||||
f = center + i * FRAME_GAP
|
||||
if 0 <= f < total:
|
||||
frames.append(f)
|
||||
return frames
|
||||
|
||||
# =========================================================
|
||||
# MASK → POLYGON
|
||||
# =========================================================
|
||||
def mask_to_polygon(mask):
|
||||
mask_uint8 = (mask * 255).astype(np.uint8)
|
||||
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
polygons = []
|
||||
for cnt in contours:
|
||||
if len(cnt) >= 3:
|
||||
poly = cnt.squeeze().tolist()
|
||||
polygons.append(poly)
|
||||
|
||||
return polygons
|
||||
|
||||
# =========================================================
|
||||
# OVERLAY MASK
|
||||
# =========================================================
|
||||
def overlay_mask(image, mask):
|
||||
overlay = image.copy()
|
||||
color = np.array([255, 0, 0]) # RED
|
||||
|
||||
overlay[mask == 1] = overlay[mask == 1] * 0.5 + color * 0.5
|
||||
return overlay.astype(np.uint8)
|
||||
|
||||
# =========================================================
|
||||
# MAIN PIPELINE
|
||||
# =========================================================
|
||||
def process_video():
|
||||
predictor = load_sam()
|
||||
parsed = parse_json(JSON_PATH)
|
||||
|
||||
cap = cv2.VideoCapture(VIDEO_PATH)
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
bbox_json = {}
|
||||
poly_json = {}
|
||||
|
||||
for center_frame in parsed:
|
||||
|
||||
temporal_frames = get_frames(center_frame, total_frames)
|
||||
|
||||
for obj_idx, obj in enumerate(parsed[center_frame]):
|
||||
cls = obj["class"]
|
||||
init_bbox = obj["bbox"]
|
||||
obj_id = obj["id"]
|
||||
|
||||
tracker = ByteTrackEnhanced(class_names=[cls])
|
||||
|
||||
x1, y1, x2, y2 = init_bbox
|
||||
tlwh = [x1, y1, x2-x1, y2-y1]
|
||||
|
||||
init_det = np.array([[*tlwh, 0.9, 0]])
|
||||
tracker.update(init_det)
|
||||
|
||||
for f in temporal_frames:
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, f)
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
continue
|
||||
|
||||
detections = init_det if f == center_frame else np.empty((0, 6))
|
||||
tracks = tracker.update(detections)
|
||||
|
||||
if len(tracks) == 0:
|
||||
continue
|
||||
|
||||
t = tracks[0]
|
||||
x, y, w, h = t.tlwh
|
||||
bbox = [int(x), int(y), int(x+w), int(y+h)]
|
||||
|
||||
image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
predictor.set_image(image_rgb)
|
||||
|
||||
mask = get_mask(image_rgb, bbox, predictor)
|
||||
mask_uint8 = (mask * 255).astype(np.uint8)
|
||||
|
||||
# SAVE MASK
|
||||
cv2.imwrite(os.path.join(MASK_DIR, f"{f}.png"), mask_uint8)
|
||||
|
||||
# SAVE MASKED IMAGE
|
||||
extracted = cv2.bitwise_and(image_rgb, image_rgb, mask=mask_uint8)
|
||||
cv2.imwrite(os.path.join(EXTRACT_DIR, f"{f}.png"),
|
||||
cv2.cvtColor(extracted, cv2.COLOR_RGB2BGR))
|
||||
|
||||
# OVERLAY IMAGE
|
||||
overlay = overlay_mask(image_rgb, mask)
|
||||
cv2.imwrite(os.path.join(IMG_DIR, f"{f}.png"),
|
||||
cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
|
||||
|
||||
# STORE BBOX JSON
|
||||
if str(f) not in bbox_json:
|
||||
bbox_json[str(f)] = {}
|
||||
|
||||
if cls not in bbox_json[str(f)]:
|
||||
bbox_json[str(f)][cls] = []
|
||||
|
||||
bbox_json[str(f)][cls].append([str(obj_id), [bbox[0], bbox[1]], [bbox[2], bbox[3]]])
|
||||
|
||||
# STORE POLYGON JSON
|
||||
polygons = mask_to_polygon(mask)
|
||||
|
||||
if str(f) not in poly_json:
|
||||
poly_json[str(f)] = {}
|
||||
|
||||
if cls not in poly_json[str(f)]:
|
||||
poly_json[str(f)][cls] = []
|
||||
|
||||
poly_json[str(f)][cls].append([str(obj_id), polygons])
|
||||
|
||||
print(f"[DONE] Frame {f}")
|
||||
|
||||
cap.release()
|
||||
|
||||
# SAVE JSON FILES
|
||||
with open(os.path.join(OUTPUT_DIR, "bboxes.json"), "w") as f:
|
||||
json.dump(bbox_json, f, indent=2)
|
||||
|
||||
with open(os.path.join(OUTPUT_DIR, "polygons.json"), "w") as f:
|
||||
json.dump(poly_json, f, indent=2)
|
||||
|
||||
|
||||
# =========================================================
|
||||
# ENTRY
|
||||
# =========================================================
|
||||
if __name__ == "__main__":
|
||||
process_video()
|
||||
152
sam2/sam_image_mask_extract.py
Normal file
152
sam2/sam_image_mask_extract.py
Normal file
@@ -0,0 +1,152 @@
|
||||
import os
|
||||
import cv2
|
||||
import json
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
# ---------------- CONFIG ----------------
|
||||
SAM_CHECKPOINT = "/media/bs5/New Volume/Brinda/Codes/Florida/SAM/sam_vit_h_4b8939.pth"
|
||||
MODEL_TYPE = "vit_h"
|
||||
|
||||
IMAGE_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/New_Dataset/set2/images_seg"
|
||||
LABEL_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/New_Dataset/set2/labels_seg/"
|
||||
OUTPUT_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/classification_data/"
|
||||
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
os.makedirs(os.path.join(OUTPUT_DIR, "masks"), exist_ok=True)
|
||||
os.makedirs(os.path.join(OUTPUT_DIR, "overlay"), exist_ok=True)
|
||||
os.makedirs(os.path.join(OUTPUT_DIR, "polygons"), exist_ok=True)
|
||||
|
||||
# ---------------- LOAD SAM ----------------
|
||||
sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
|
||||
sam.to(DEVICE)
|
||||
predictor = SamPredictor(sam)
|
||||
|
||||
# ---------------- HELPERS ----------------
|
||||
def yolo_to_xyxy(label, img_w, img_h):
|
||||
"""
|
||||
YOLO format: class cx cy w h (normalized)
|
||||
Convert → pixel [x1, y1, x2, y2]
|
||||
"""
|
||||
cls, cx, cy, w, h = label
|
||||
|
||||
x1 = int((cx - w / 2) * img_w)
|
||||
y1 = int((cy - h / 2) * img_h)
|
||||
x2 = int((cx + w / 2) * img_w)
|
||||
y2 = int((cy + h / 2) * img_h)
|
||||
|
||||
return [x1, y1, x2, y2]
|
||||
|
||||
|
||||
def expand_box(x1, y1, x2, y2, pad, w, h):
|
||||
return [
|
||||
max(0, x1 - pad),
|
||||
max(0, y1 - pad),
|
||||
min(w, x2 + pad),
|
||||
min(h, y2 + pad),
|
||||
]
|
||||
|
||||
|
||||
def mask_to_polygon(mask):
|
||||
contours, _ = cv2.findContours(
|
||||
mask.astype(np.uint8),
|
||||
cv2.RETR_EXTERNAL,
|
||||
cv2.CHAIN_APPROX_SIMPLE
|
||||
)
|
||||
|
||||
polygons = []
|
||||
for cnt in contours:
|
||||
cnt = cnt.squeeze()
|
||||
if len(cnt.shape) == 2 and len(cnt) > 2:
|
||||
polygons.append(cnt.tolist())
|
||||
|
||||
return polygons
|
||||
|
||||
|
||||
# ---------------- PROCESS ----------------
|
||||
for img_file in os.listdir(IMAGE_DIR):
|
||||
|
||||
if not img_file.lower().endswith((".jpg", ".png", ".jpeg")):
|
||||
continue
|
||||
|
||||
image_path = os.path.join(IMAGE_DIR, img_file)
|
||||
label_path = os.path.join(LABEL_DIR, img_file.replace(".jpg", ".txt").replace(".png", ".txt"))
|
||||
|
||||
if not os.path.exists(label_path):
|
||||
continue
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
if image is None:
|
||||
continue
|
||||
|
||||
h, w = image.shape[:2]
|
||||
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
|
||||
predictor.set_image(image_rgb)
|
||||
|
||||
overlay = image.copy()
|
||||
output_data = []
|
||||
|
||||
# -------- Read YOLO labels --------
|
||||
with open(label_path, "r") as f:
|
||||
lines = f.readlines()
|
||||
|
||||
for idx, line in enumerate(lines):
|
||||
parts = list(map(float, line.strip().split()))
|
||||
cls_id = int(parts[0])
|
||||
cx, cy, bw, bh = parts[1:]
|
||||
|
||||
# YOLO → XYXY
|
||||
x1, y1, x2, y2 = yolo_to_xyxy([cls_id, cx, cy, bw, bh], w, h)
|
||||
|
||||
# Expand bbox (important for SAM)
|
||||
x1, y1, x2, y2 = expand_box(x1, y1, x2, y2, pad=10, w=w, h=h)
|
||||
|
||||
box = np.array([x1, y1, x2, y2])
|
||||
|
||||
# -------- SAM Prediction --------
|
||||
masks, scores, _ = predictor.predict(
|
||||
box=box,
|
||||
multimask_output=False
|
||||
)
|
||||
|
||||
mask = masks[0]
|
||||
|
||||
# -------- Save mask --------
|
||||
mask_name = f"{img_file.split('.')[0]}_{idx}.png"
|
||||
mask_path = os.path.join(OUTPUT_DIR, "masks", mask_name)
|
||||
cv2.imwrite(mask_path, mask.astype(np.uint8) * 255)
|
||||
|
||||
# -------- Overlay --------
|
||||
color = np.random.randint(0, 255, (3,))
|
||||
overlay[mask] = (overlay[mask] * 0.5 + color * 0.5).astype(np.uint8)
|
||||
|
||||
# -------- Polygon --------
|
||||
polygons = mask_to_polygon(mask)
|
||||
|
||||
output_data.append({
|
||||
"id": idx,
|
||||
"class": cls_id,
|
||||
"bbox": [x1, y1, x2, y2],
|
||||
"polygon": polygons
|
||||
})
|
||||
|
||||
# -------- Save overlay --------
|
||||
cv2.imwrite(
|
||||
os.path.join(OUTPUT_DIR, "overlay", img_file),
|
||||
overlay
|
||||
)
|
||||
|
||||
# -------- Save JSON --------
|
||||
json_path = os.path.join(
|
||||
OUTPUT_DIR, "polygons", img_file.replace(".jpg", ".json").replace(".png", ".json")
|
||||
)
|
||||
|
||||
with open(json_path, "w") as f:
|
||||
json.dump(output_data, f, indent=4)
|
||||
|
||||
print(" Processing Complete")
|
||||
Reference in New Issue
Block a user