diff --git a/Cargo.lock b/Cargo.lock index 81b8709..39f9c74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -93,9 +93,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.98" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "atomic_refcell" @@ -605,6 +605,16 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "flatbuffers" +version = "24.12.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" +dependencies = [ + "bitflags 1.3.2", + "rustc_version", +] + [[package]] name = "flate2" version = "1.1.2" @@ -899,10 +909,13 @@ dependencies = [ "hound", "image", "objc", + "ocrs", "once_cell", "pango", "pangocairo", "regex", + "rten", + "rten-imageproc", "serde", "serde_json", "tempfile", @@ -931,7 +944,7 @@ dependencies = [ "paste", "pin-project-lite", "smallvec", - "thiserror 2.0.12", + "thiserror 2.0.18", ] [[package]] @@ -1045,7 +1058,7 @@ dependencies = [ "gstreamer-video-sys", "libc", "once_cell", - "thiserror 2.0.12", + "thiserror 2.0.18", ] [[package]] @@ -1084,6 +1097,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hex" version = "0.4.3" @@ -1615,6 +1634,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "objc" version = "0.2.7" @@ -1633,6 +1662,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "ocrs" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5379fdd3f11522b5a2ff53017a189463dabf5d0a9c915cb3eb97fabec4ea11c" +dependencies = [ + "anyhow", + "rayon", + "rten", + "rten-imageproc", + "rten-tensor", + "thiserror 2.0.18", + "wasm-bindgen", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -1816,7 +1860,7 @@ dependencies = [ "rustc-hash 2.1.1", "rustls", "socket2", - "thiserror 2.0.12", + "thiserror 2.0.18", "tokio", "tracing", "web-time", @@ -1837,7 +1881,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.12", + "thiserror 2.0.18", "tinyvec", "tracing", "web-time", @@ -1933,9 +1977,9 @@ dependencies = [ [[package]] name = "rayon" -version = "1.10.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" dependencies = [ "either", "rayon-core", @@ -1943,9 +1987,9 @@ dependencies = [ [[package]] name = "rayon-core" -version = "1.12.1" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" dependencies = [ "crossbeam-deque", "crossbeam-utils", @@ -2034,6 +2078,106 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rten" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43c230fa4ade87c913f61dbd911b7eb0d49460ceff3f1e4fabc837fac191137c" +dependencies = [ + "flatbuffers", + "num_cpus", + "rayon", + "rten-base", + "rten-gemm", + "rten-model-file", + "rten-shape-inference", + "rten-simd", + "rten-tensor", + "rten-vecmath", + "rustc-hash 2.1.1", + "smallvec", + "typeid", + "wasm-bindgen", +] + +[[package]] +name = "rten-base" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2738cf8bb4c27f828ac788d01ccf4e367e8e773cfec6851f81851b5211de6a79" +dependencies = [ + "rayon", +] + +[[package]] +name = "rten-gemm" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a81a0ca209fb5ce21bd17efa0bd287d5881c6cebfbff0b21c4294a1a14a9e" +dependencies = [ + "rayon", + "rten-base", + "rten-simd", + "rten-tensor", +] + +[[package]] +name = "rten-imageproc" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5f148e7e941fb5727b9046a5fa1b45525543d5105f14b384fd9261df0ee49bc" +dependencies = [ + "rten-tensor", +] + +[[package]] +name = "rten-model-file" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2f8d270f07ab1bbfff47250c6039f6caa5da59d6da7d74f66aa48559aa6fea" +dependencies = [ + "flatbuffers", + "rten-base", +] + +[[package]] +name = "rten-shape-inference" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e8a913c7ca40e2bfbb2a0cd447cce56b33ab19435f56693271a2ef37cf58984" +dependencies = [ + "rten-tensor", + "smallvec", +] + +[[package]] +name = "rten-simd" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b19a0032dfcb70dd20960c1c51a37674b237586cbc1ce586f45b46605d108e82" + +[[package]] +name = "rten-tensor" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05dc744a270aa32d154f1a3df8e48740ccc1be9dfbcf23295ada66d83aa98de6" +dependencies = [ + "rayon", + "rten-base", + "smallvec", + "typeid", +] + +[[package]] +name = "rten-vecmath" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9574ddebf5671bc08ceb76e2e1638fadc57fdeff318634eab2c29e9a803cff64" +dependencies = [ + "rten-base", + "rten-simd", +] + [[package]] name = "rustc-demangle" version = "0.1.25" @@ -2052,6 +2196,15 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "0.38.44" @@ -2126,6 +2279,12 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.219" @@ -2316,11 +2475,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.12" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl 2.0.12", + "thiserror-impl 2.0.18", ] [[package]] @@ -2336,9 +2495,9 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.12" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", @@ -2535,6 +2694,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typeid" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" + [[package]] name = "typenum" version = "1.18.0" diff --git a/Cargo.toml b/Cargo.toml index 89ae0cc..6beac84 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,9 +12,10 @@ name = "gstedgeimpulse" crate-type = ["cdylib"] [features] -default = ["ffi"] +default = ["ffi", "ocrs"] eim = ["edge-impulse-runner/eim"] ffi = ["edge-impulse-runner/ffi"] +ocrs = ["dep:ocrs", "dep:rten", "dep:rten-imageproc"] [dependencies] cairo-rs = "0.20.7" @@ -35,6 +36,9 @@ serde_json = "1.0" tokio = { version = "1.0", features = ["rt-multi-thread"] } edge-impulse-runner = { git = "https://github.com/edgeimpulse/edge-impulse-runner-rs.git", rev = "21449c2", default-features = false } tempfile = "3.10" +ocrs = { version = "0.12.2", optional = true } +rten = { version = "0.24.0", optional = true, default-features = false, features = ["rten_format"] } +rten-imageproc = { version = "0.24.0", optional = true } [dev-dependencies] anyhow = "1.0" diff --git a/README.md b/README.md index 02c88fb..b7ecb74 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [![CI](https://github.com/edgeimpulse/gst-plugins-edgeimpulse/actions/workflows/ci.yml/badge.svg)](https://github.com/edgeimpulse/gst-plugins-edgeimpulse/actions/workflows/ci.yml) [![Docs](https://img.shields.io/badge/docs-latest-blue.svg)](https://edgeimpulse.github.io/gst-plugins-edgeimpulse/) -A GStreamer plugin that enables real-time machine learning inference and data ingestion using Edge Impulse models and APIs. The plugin provides six elements for audio and video inference, visualization, ingestion, and pipeline flow control. +A GStreamer plugin that enables real-time machine learning inference and data ingestion using Edge Impulse models and APIs. The plugin provides seven elements for audio and video inference, visualization, ingestion, and pipeline flow control. ## Architecture Overview @@ -11,6 +11,7 @@ graph LR subgraph "Inference" A[edgeimpulseaudioinfer] V[edgeimpulsevideoinfer] + R[edgeimpulseocr] end subgraph "Flow Control" F[edgeimpulsecontinueif] @@ -24,6 +25,7 @@ graph LR V -- "VideoRegionOfInterestMeta\nInferenceResultMeta" --> O V -- "VideoRegionOfInterestMeta\nInferenceResultMeta" --> F V -- "VideoRegionOfInterestMeta" --> C + R -- "VideoRegionOfInterestMeta" --> O F -- "pass / drop" --> C C -- "CropOriginMeta" --> V A -- "InferenceResultMeta" --> F @@ -35,6 +37,7 @@ graph LR |---------|-------------|-------| | [`edgeimpulseaudioinfer`](docs/edgeimpulseaudioinfer.md) | Runs audio inference (classification, keyword spotting) | Audio | | [`edgeimpulsevideoinfer`](docs/edgeimpulsevideoinfer.md) | Runs video inference (classification, detection, anomaly) | Video | +| [`edgeimpulseocr`](docs/edgeimpulseocr.md) | Recognizes text in video frames and attaches it as ROI metadata | Video | | [`edgeimpulseoverlay`](docs/edgeimpulseoverlay.md) | Draws bounding boxes and labels on video frames | Video | | [`edgeimpulsesink`](docs/edgeimpulsesink.md) | Uploads audio/video to Edge Impulse ingestion API | Audio / Video | | [`edgeimpulsecontinueif`](docs/edgeimpulsecontinueif.md) | Conditional gate — passes or drops buffers based on inference metadata | Any | @@ -312,6 +315,19 @@ Attached by `edgeimpulsecrop` to each cropped buffer, recording where the crop c ``` - **Video Metadata:** Scores → `VideoAnomalyMeta` (see [above](#videoanomalymeta)). Grid cells may also be attached as individual `VideoRegionOfInterestMeta` entries for overlay visualization. +#### 4. Text Recognition (OCR) + +Produced by [`edgeimpulseocr`](docs/edgeimpulseocr.md) rather than `edgeimpulsevideoinfer`, so it uses its own `ocr` element message (one per recognized line) instead of `edge-impulse-video-inference-result`: + +- **Bus Message Example:** + ``` + ocr, text=(string)"HELLO 12345", confidence=(double)1.0, + x=(int)94, y=(int)67, width=(int)451, height=(int)61, + timestamp=(gint64)0 + ``` +- **Video Metadata:** Each recognized line → `VideoRegionOfInterestMeta` (see [above](#videoregionofinterestmeta)) with a `detection` param carrying the text as `label`, so `edgeimpulseoverlay` renders it like any other detection. +- **Backends:** `ocrs` (default) performs recognition in-process with embedded rten models; an `edge-impulse` decode backend is planned. See [`edgeimpulseocr`](docs/edgeimpulseocr.md) for details. + ## Dependencies ### System Dependencies diff --git a/build.rs b/build.rs index 4836365..2823992 100644 --- a/build.rs +++ b/build.rs @@ -150,6 +150,10 @@ pub const VIDEO_ANOMALY_META_API_NAME: &str = "VideoAnomalyMetaAPI{}"; pub const VIDEO_CLASSIFICATION_META_NAME: &str = "VideoClassificationMeta{}"; #[allow(dead_code)] pub const VIDEO_ANOMALY_META_NAME: &str = "VideoAnomalyMeta{}"; + +// OCR element type name +#[allow(dead_code)] +pub const OCR_TYPE_NAME: &str = "EdgeImpulseOcr{}"; "#, plugin_variant, type_suffix, // VIDEO_INFER_TYPE_NAME @@ -166,6 +170,7 @@ pub const VIDEO_ANOMALY_META_NAME: &str = "VideoAnomalyMeta{}"; type_suffix, // VIDEO_ANOMALY_META_API_NAME type_suffix, // VIDEO_CLASSIFICATION_META_NAME type_suffix, // VIDEO_ANOMALY_META_NAME + type_suffix, // OCR_TYPE_NAME ); std::fs::write(&type_names_path, type_names_code).expect("Failed to write type names file"); diff --git a/docs/edgeimpulseocr.md b/docs/edgeimpulseocr.md new file mode 100644 index 0000000..b7a2b54 --- /dev/null +++ b/docs/edgeimpulseocr.md @@ -0,0 +1,124 @@ +# edgeimpulseocr + +Reads text from video frames (optical character recognition) and attaches each recognized line to the buffer as inference metadata. For every line of text found, the element attaches a `VideoRegionOfInterestMeta` — so downstream elements such as `edgeimpulseoverlay` render the text exactly like any other detection — and, optionally, posts an `ocr` message on the bus. + +Recognition runs on a background worker thread, fully decoupled from the streaming thread, so a slow model never stalls the pipeline: frames pass through immediately and results attach to a slightly later frame. + +## Element Details + +- Long name: Edge Impulse OCR +- Class: Filter/Analyzer/Video +- Description: Reads text from video frames and attaches it as ROI metadata + +## Pad Templates + +- Sink pad (Always available): + ``` + video/x-raw + format: RGB + width: [ 1, 2147483647 ] + height: [ 1, 2147483647 ] + ``` +- Source pad (Always available): + ``` + video/x-raw + format: RGB + width: [ 1, 2147483647 ] + height: [ 1, 2147483647 ] + ``` + +## Properties + +1. `backend` (string): + - OCR engine to use. `ocrs` performs recognition in-process with embedded [rten](https://github.com/robertknight/rten) models; `edge-impulse` (decode recognition results from an upstream Edge Impulse model) is planned and currently recognizes nothing. + - Default: `ocrs` + - Flags: readable, writable, changeable in the NULL or READY state + +2. `detection-model` (string): + - Path to a text-detection `.rten` model. Empty uses the model embedded in the plugin. + - Default: `""` (embedded model) + - Flags: readable, writable, changeable in the NULL or READY state + +3. `recognition-model` (string): + - Path to a text-recognition `.rten` model. Empty uses the model embedded in the plugin. + - Default: `""` (embedded model) + - Flags: readable, writable, changeable in the NULL or READY state + +4. `min-confidence` (double): + - Drop recognized lines below this confidence. Note: the `ocrs` backend does not expose a per-line confidence (it reports `1.0`), so this acts as a no-op for that backend today. + - Range: 0.0 -- 1.0 + - Default: 0.0 + - Flags: readable, writable, changeable in PLAYING state + +5. `max-text-length` (unsigned integer): + - Truncate each recognized line to at most this many characters. + - Range: 1 -- 4294967295 + - Default: 256 + - Flags: readable, writable, changeable in PLAYING state + +6. `post-message` (boolean): + - Post an `ocr` element message on the bus for each recognized line. + - Default: true + - Flags: readable, writable, changeable in PLAYING state + +7. `interval` (unsigned integer): + - Run recognition on one frame out of every N (higher values reduce load). + - Range: 1 -- 4294967295 + - Default: 1 + - Flags: readable, writable, changeable in PLAYING state + +## How It Works + +1. Every `interval`-th input frame is copied and handed to a background worker thread; all other frames pass straight through. +2. The worker runs the configured backend (text detection followed by recognition) off the streaming thread. +3. The most recent results are attached to passing buffers as `VideoRegionOfInterestMeta` — one per line, each with a `detection` param carrying the recognized text as `label` and the confidence — so `edgeimpulseoverlay` (or any ROI consumer) can render them. +4. When `post-message` is true, each recognized line is also emitted as an `ocr` bus message. + +Because recognition is asynchronous, results attach to a slightly later frame than the one they were computed from. The `ocr` message's `timestamp` field carries the PTS (in milliseconds) of the source frame the text was read from. + +### `ocr` Bus Message + +``` +ocr, text=(string), confidence=(double), + x=(int), y=(int), width=(int), height=(int), timestamp=(gint64) +``` + +`x`, `y`, `width`, and `height` are the line's bounding box in full-frame pixels; `timestamp` is the source-frame PTS in milliseconds. + +## Backends + +- **`ocrs`** (default): Pure-Rust text detection and recognition using [ocrs](https://github.com/robertknight/ocrs). The default detection and recognition models are embedded in the plugin, so no external files are required. Override them with `detection-model` / `recognition-model` to use your own; the standard models can also be fetched separately with [`examples/download-ocr-models.sh`](../examples/download-ocr-models.sh). This backend is compiled only when the plugin is built with the `ocrs` cargo feature (enabled by default). +- **`edge-impulse`** (planned): Will decode text from an upstream Edge Impulse recognition model, letting you train and deploy custom OCR models via Edge Impulse Studio. Not yet implemented — selecting it currently recognizes nothing. + +> **Build in release mode.** The recognition models are large; debug builds run recognition far too slowly to be usable. Always build and run with `--release`. + +## Example Pipelines + +```bash +# Recognize text from a camera stream and overlay it on the video +gst-launch-1.0 v4l2src ! videoconvert ! video/x-raw,format=RGB ! \ + edgeimpulseocr backend=ocrs ! edgeimpulseoverlay ! autovideosink + +# Recognize text from a single still image (run with -m to print the `ocr` +# bus messages; imagefreeze never ends, so stop with Ctrl+C — or use the +# self-terminating examples/ocr_inference.rs below). +gst-launch-1.0 -m filesrc location=text.png ! decodebin ! imagefreeze ! \ + videoconvert ! video/x-raw,format=RGB ! \ + edgeimpulseocr backend=ocrs ! fakesink + +# Use your own rten models instead of the embedded defaults (Ctrl+C to stop) +gst-launch-1.0 -m filesrc location=text.png ! decodebin ! imagefreeze ! \ + videoconvert ! video/x-raw,format=RGB ! \ + edgeimpulseocr backend=ocrs \ + detection-model=text-detection.rten \ + recognition-model=text-recognition.rten ! \ + fakesink +``` + +An end-to-end example that prints recognized text is available at [`examples/ocr_inference.rs`](../examples/ocr_inference.rs): + +```bash +export GST_PLUGIN_PATH="$(pwd)/target/release:$GST_PLUGIN_PATH" +cargo run --release --no-default-features --features "eim ocrs" \ + --example ocr_inference -- --image text.png +``` diff --git a/examples/download-ocr-models.sh b/examples/download-ocr-models.sh new file mode 100755 index 0000000..9c28671 --- /dev/null +++ b/examples/download-ocr-models.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +set -euo pipefail +HERE="$(dirname "$0")" +DEST="$HERE/assets" # used by the models-gated integration test (explicit paths) +VENDOR="$HERE/../models" # embedded into the .so via include_bytes! (Step 3b) +mkdir -p "$DEST" "$VENDOR" +BASE="https://ocrs-models.s3-accelerate.amazonaws.com" + +# Pinned SHA-256 of the exact model bytes vendored under ../models and embedded +# into the plugin binary. The download is verified against these so a corrupted +# or tampered file can never silently become part of the built .so. +SHA_DETECTION="f15cfb56bd02c4bf478a20343986504a1f01e1665c2b3a0ad66340f054b1b5ca" +SHA_RECOGNITION="e484866d4cce403175bd8d00b128feb08ab42e208de30e42cd9889d8f1735a6e" + +sha256_of() { + if command -v sha256sum >/dev/null 2>&1; then + sha256sum "$1" | awk '{print $1}' + else + shasum -a 256 "$1" | awk '{print $1}' + fi +} + +verify() { + local file="$1" expected="$2" actual + actual="$(sha256_of "$file")" + if [ "$actual" != "$expected" ]; then + echo "ERROR: checksum mismatch for $file" >&2 + echo " expected $expected" >&2 + echo " actual $actual" >&2 + exit 1 + fi +} + +curl -fsSL "$BASE/text-detection.rten" -o "$DEST/text-detection.rten" +curl -fsSL "$BASE/text-recognition.rten" -o "$DEST/text-recognition.rten" +verify "$DEST/text-detection.rten" "$SHA_DETECTION" +verify "$DEST/text-recognition.rten" "$SHA_RECOGNITION" +cp "$DEST/text-detection.rten" "$DEST/text-recognition.rten" "$VENDOR/" +echo "Downloaded and verified ocrs models in $DEST and vendored them into $VENDOR" diff --git a/examples/ocr_inference.rs b/examples/ocr_inference.rs new file mode 100644 index 0000000..4b21b43 --- /dev/null +++ b/examples/ocr_inference.rs @@ -0,0 +1,222 @@ +//! OCR text-recognition example using the `edgeimpulseocr` GStreamer element. +//! +//! Reads a single image, runs text recognition with the built-in pure-Rust +//! `ocrs` backend, and prints every `ocr` message the element posts on the bus. +//! +//! Recognition runs on a worker thread, so results for a still image land a few +//! frames after it enters the element. The pipeline therefore uses `imagefreeze` +//! to keep feeding the frame until the first result arrives (mirroring how a +//! live camera would supply a continuous stream). +//! +//! Build with the `ocrs` feature (which embeds the default detection and +//! recognition models) and disable the default `ffi` backend so no Edge Impulse +//! model needs to be compiled in: +//! +//! ```sh +//! export GST_PLUGIN_PATH="$(pwd)/target/release:$GST_PLUGIN_PATH" +//! cargo run --release --no-default-features --features "eim ocrs" \ +//! --example ocr_inference -- --image path/to/text.png +//! ``` +//! +//! To use your own models instead of the embedded ones, pass their paths with +//! `--detection-model` and `--recognition-model`. + +use clap::Parser; +use gstreamer as gst; +use gstreamer::prelude::*; +use std::collections::HashSet; +use std::error::Error; +use std::path::Path; +use std::time::{Duration, Instant}; + +/// Command line parameters for the OCR example. +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct OcrParams { + /// Path to the input image file + #[arg(short, long)] + image: String, + + /// Path to a text-detection model (optional; uses the embedded model if unset) + #[arg(long, default_value = "")] + detection_model: String, + + /// Path to a text-recognition model (optional; uses the embedded model if unset) + #[arg(long, default_value = "")] + recognition_model: String, + + /// Seconds to wait for the first recognition result before giving up + #[arg(long, default_value = "30")] + timeout: u64, +} + +fn build_pipeline(args: &OcrParams) -> Result> { + gst::init()?; + + if !Path::new(&args.image).exists() { + return Err(format!("Input image file not found: {}", args.image).into()); + } + + let pipeline = gst::Pipeline::new(); + + let source = gst::ElementFactory::make("filesrc") + .name("source") + .property("location", &args.image) + .build()?; + let decodebin = gst::ElementFactory::make("decodebin") + .name("decodebin") + .build()?; + // Repeat the single decoded frame so the async worker keeps receiving it + // until it produces a result. + let imagefreeze = gst::ElementFactory::make("imagefreeze") + .name("imagefreeze") + .build()?; + let videoconvert = gst::ElementFactory::make("videoconvert") + .name("videoconvert") + .build()?; + let caps_filter = gst::ElementFactory::make("capsfilter") + .name("caps_filter") + .property( + "caps", + gst::Caps::builder("video/x-raw") + .field("format", "RGB") + .build(), + ) + .build()?; + + let ocr = gst::ElementFactory::make("edgeimpulseocr") + .name("ocr") + .property("backend", "ocrs") + .build()?; + if !args.detection_model.is_empty() { + ocr.set_property("detection-model", &args.detection_model); + } + if !args.recognition_model.is_empty() { + ocr.set_property("recognition-model", &args.recognition_model); + } + + let sink = gst::ElementFactory::make("fakesink") + .name("sink") + .property("sync", false) + .build()?; + + pipeline.add_many([ + &source, + &decodebin, + &imagefreeze, + &videoconvert, + &caps_filter, + &ocr, + &sink, + ])?; + + // decodebin exposes its source pad dynamically once it detects the format. + imagefreeze.link(&videoconvert)?; + videoconvert.link(&caps_filter)?; + caps_filter.link(&ocr)?; + ocr.link(&sink)?; + + let imagefreeze_sink = imagefreeze.clone(); + decodebin.connect_pad_added(move |_, pad| { + let sink_pad = imagefreeze_sink.static_pad("sink").unwrap(); + if !sink_pad.is_linked() { + if let Err(err) = pad.link(&sink_pad) { + eprintln!("Failed to link decodebin to imagefreeze: {err}"); + } + } + }); + source.link(&decodebin)?; + + Ok(pipeline) +} + +fn example_main() -> Result<(), Box> { + let args = OcrParams::parse(); + + println!("🚀 Starting Edge Impulse OCR inference"); + println!("📁 Input image: {}", args.image); + + let pipeline = build_pipeline(&args)?; + pipeline.set_state(gst::State::Playing)?; + + let bus = pipeline.bus().unwrap(); + let start = Instant::now(); + let hard_timeout = Duration::from_secs(args.timeout); + // Once the first result arrives, wait a short grace period to collect the + // remaining lines from the same frame, then stop (imagefreeze never ends). + let grace = Duration::from_secs(1); + let mut first_result_at: Option = None; + let mut seen = HashSet::new(); + + loop { + if let Some(t) = first_result_at { + if t.elapsed() > grace { + break; + } + } else if start.elapsed() > hard_timeout { + eprintln!("⚠️ No text recognized within {}s", args.timeout); + break; + } + + let Some(msg) = bus.timed_pop(gst::ClockTime::from_mseconds(200)) else { + continue; + }; + use gst::MessageView; + match msg.view() { + MessageView::Error(err) => { + eprintln!( + "❌ Error from {:?}: {} ({})", + err.src().map(|s| s.path_string()), + err.error(), + err.debug().unwrap_or_default() + ); + break; + } + MessageView::Element(element) => { + let Some(s) = element.structure() else { + continue; + }; + if s.name() != "ocr" { + continue; + } + let text = s.get::("text").unwrap_or_default(); + // The same frozen frame is recognized repeatedly; only report + // each distinct string once. + if !seen.insert(text.clone()) { + continue; + } + let confidence = s.get::("confidence").unwrap_or(0.0); + let x = s.get::("x").unwrap_or(0); + let y = s.get::("y").unwrap_or(0); + let width = s.get::("width").unwrap_or(0); + let height = s.get::("height").unwrap_or(0); + println!( + "📖 OCR: {text:?} (confidence {confidence:.2}) at [{x}, {y}, {width}x{height}]" + ); + if first_result_at.is_none() { + first_result_at = Some(Instant::now()); + } + } + _ => {} + } + } + + pipeline.set_state(gst::State::Null)?; + + if seen.is_empty() { + println!("✅ OCR inference finished (no text found)"); + } else { + println!( + "✅ OCR inference finished ({} line(s) recognized)", + seen.len() + ); + } + Ok(()) +} + +fn main() { + if let Err(e) = example_main() { + eprintln!("❌ Error: {e}"); + std::process::exit(1); + } +} diff --git a/models/text-detection.rten b/models/text-detection.rten new file mode 100644 index 0000000..1725cf2 Binary files /dev/null and b/models/text-detection.rten differ diff --git a/models/text-recognition.rten b/models/text-recognition.rten new file mode 100644 index 0000000..464815f Binary files /dev/null and b/models/text-recognition.rten differ diff --git a/src/lib.rs b/src/lib.rs index 3129a29..1164d9b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -78,6 +78,7 @@ mod common; mod crop; mod filter; pub mod meta; +mod ocr; mod overlay; pub mod sink; pub mod video; @@ -89,6 +90,7 @@ fn plugin_init(plugin: &gst::Plugin) -> Result<(), glib::BoolError> { sink::register(plugin)?; filter::register(plugin)?; crop::register(plugin)?; + ocr::register(plugin)?; Ok(()) } diff --git a/src/ocr/backend.rs b/src/ocr/backend.rs new file mode 100644 index 0000000..ad275e3 --- /dev/null +++ b/src/ocr/backend.rs @@ -0,0 +1,60 @@ +//! Backend abstraction for the OCR element. + +/// One recognized line of text with its axis-aligned box in full-frame pixels. +#[derive(Debug, Clone, PartialEq)] +pub struct OcrLine { + pub text: String, + pub confidence: f32, + pub x: u32, + pub y: u32, + pub w: u32, + pub h: u32, +} + +/// Recognizes text in a tightly-packed RGB frame (stride == width * 3). +pub trait OcrBackend: Send { + fn recognize(&mut self, rgb: &[u8], width: u32, height: u32) -> Result, String>; +} + +/// Which recognition engine the element uses. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Backend { + Ocrs, + EdgeImpulse, +} + +impl Backend { + pub fn parse(s: &str) -> Option { + match s { + "ocrs" => Some(Backend::Ocrs), + "edge-impulse" => Some(Backend::EdgeImpulse), + _ => None, + } + } +} + +/// A backend that recognizes nothing. Used as a safe default before a real +/// backend is configured, and for tests. +pub struct NoopBackend; +impl OcrBackend for NoopBackend { + fn recognize(&mut self, _rgb: &[u8], _w: u32, _h: u32) -> Result, String> { + Ok(Vec::new()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_known_backends() { + assert_eq!(Backend::parse("ocrs"), Some(Backend::Ocrs)); + assert_eq!(Backend::parse("edge-impulse"), Some(Backend::EdgeImpulse)); + assert_eq!(Backend::parse("nope"), None); + } + + #[test] + fn noop_backend_returns_no_lines() { + assert!(NoopBackend.recognize(&[0u8; 12], 2, 2).unwrap().is_empty()); + } +} diff --git a/src/ocr/imp.rs b/src/ocr/imp.rs new file mode 100644 index 0000000..74a130c --- /dev/null +++ b/src/ocr/imp.rs @@ -0,0 +1,432 @@ +use gstreamer as gst; +use gstreamer::glib; +use gstreamer::prelude::*; +use gstreamer::subclass::prelude::*; +use gstreamer_base as gst_base; +use gstreamer_base::subclass::prelude::*; +use gstreamer_video as gst_video; +use gstreamer_video::VideoFrameExt; +use gstreamer_video::VideoFrameRef; +use once_cell::sync::Lazy; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc::{sync_channel, SyncSender}; +use std::sync::{Arc, Mutex}; +use std::thread::JoinHandle; + +use crate::ocr::backend::{Backend, NoopBackend, OcrBackend, OcrLine}; +use crate::ocr::shaping::{attach_results, build_ocr_message, filter_and_truncate}; + +include!(concat!(env!("OUT_DIR"), "/type_names.rs")); + +static CAT: Lazy = Lazy::new(|| { + gst::DebugCategory::new( + "edgeimpulseocr", + gst::DebugColorFlags::empty(), + Some("Edge Impulse OCR"), + ) +}); + +#[derive(Debug, Clone)] +pub struct Settings { + pub backend: String, + pub detection_model: String, + pub recognition_model: String, + pub min_confidence: f64, + pub max_text_length: u32, + pub post_message: bool, + pub interval: u32, +} + +impl Default for Settings { + fn default() -> Self { + Self { + backend: "ocrs".into(), + detection_model: String::new(), + recognition_model: String::new(), + min_confidence: 0.0, + max_text_length: 256, + post_message: true, + interval: 1, + } + } +} + +/// Recognition results shared from the worker thread to the streaming thread. +/// `generation` bumps on every new result so the streaming thread can post an +/// `ocr` message once per recognition rather than once per output frame. +#[derive(Default)] +struct Latest { + generation: u64, + lines: Vec, + /// PTS (ms) of the frame these lines were recognized from, carried so + /// `ocr` messages timestamp the recognized frame, not the output frame. + pts_ms: i64, +} + +/// A frame handed to the worker for recognition. +struct FrameJob { + rgb: Vec, + width: u32, + height: u32, + pts_ms: i64, +} + +struct State { + frame_tx: SyncSender, + worker: Option>, + latest: Arc>, + frame_count: u64, + last_posted_generation: u64, +} + +#[derive(Default)] +pub struct EdgeImpulseOcr { + pub(crate) settings: Mutex, + state: Mutex>, + info: Mutex>, + /// Set once the worker thread is observed gone, so we warn only once + /// instead of on every subsequently dropped frame. + worker_gone_logged: AtomicBool, +} + +#[glib::object_subclass] +impl ObjectSubclass for EdgeImpulseOcr { + const NAME: &'static str = OCR_TYPE_NAME; + type Type = super::EdgeImpulseOcr; + type ParentType = gst_base::BaseTransform; +} + +impl ObjectImpl for EdgeImpulseOcr { + fn properties() -> &'static [glib::ParamSpec] { + static PROPERTIES: Lazy> = Lazy::new(|| { + vec![ + glib::ParamSpecString::builder("backend") + .nick("Backend") + .blurb("OCR backend implementation to use") + .default_value(Some("ocrs")) + .mutable_ready() + .build(), + glib::ParamSpecString::builder("detection-model") + .nick("Detection Model") + .blurb("Path to the OCR text detection model file") + .default_value(Some("")) + .mutable_ready() + .build(), + glib::ParamSpecString::builder("recognition-model") + .nick("Recognition Model") + .blurb("Path to the OCR text recognition model file") + .default_value(Some("")) + .mutable_ready() + .build(), + glib::ParamSpecDouble::builder("min-confidence") + .nick("Minimum Confidence") + .blurb("Minimum confidence threshold for OCR text results") + .minimum(0.0) + .maximum(1.0) + .default_value(0.0) + .mutable_playing() + .build(), + glib::ParamSpecUInt::builder("max-text-length") + .nick("Maximum Text Length") + .blurb("Maximum number of characters to recognize per text region") + .minimum(1) + .maximum(u32::MAX) + .default_value(256) + .mutable_playing() + .build(), + glib::ParamSpecBoolean::builder("post-message") + .nick("Post Message") + .blurb("Post OCR results on the GStreamer bus") + .default_value(true) + .mutable_playing() + .build(), + glib::ParamSpecUInt::builder("interval") + .nick("Interval") + .blurb("Process one frame every N input frames") + .minimum(1) + .maximum(u32::MAX) + .default_value(1) + .mutable_playing() + .build(), + ] + }); + PROPERTIES.as_ref() + } + + fn set_property(&self, _id: usize, value: &glib::Value, pspec: &glib::ParamSpec) { + let mut settings = self.settings.lock().unwrap(); + match pspec.name() { + "backend" => settings.backend = value.get().unwrap(), + "detection-model" => settings.detection_model = value.get().unwrap(), + "recognition-model" => settings.recognition_model = value.get().unwrap(), + "min-confidence" => settings.min_confidence = value.get().unwrap(), + "max-text-length" => settings.max_text_length = value.get().unwrap(), + "post-message" => settings.post_message = value.get().unwrap(), + "interval" => settings.interval = value.get().unwrap(), + _ => unimplemented!(), + } + } + + fn property(&self, _id: usize, pspec: &glib::ParamSpec) -> glib::Value { + let settings = self.settings.lock().unwrap(); + match pspec.name() { + "backend" => settings.backend.to_value(), + "detection-model" => settings.detection_model.to_value(), + "recognition-model" => settings.recognition_model.to_value(), + "min-confidence" => settings.min_confidence.to_value(), + "max-text-length" => settings.max_text_length.to_value(), + "post-message" => settings.post_message.to_value(), + "interval" => settings.interval.to_value(), + _ => unimplemented!(), + } + } +} +impl GstObjectImpl for EdgeImpulseOcr {} + +impl ElementImpl for EdgeImpulseOcr { + fn metadata() -> Option<&'static gst::subclass::ElementMetadata> { + static M: Lazy = Lazy::new(|| { + gst::subclass::ElementMetadata::new( + "Edge Impulse OCR", + "Filter/Analyzer/Video", + "Reads text from video frames and attaches it as ROI metadata", + "Fernando Jiménez Moreno ", + ) + }); + Some(&*M) + } + fn pad_templates() -> &'static [gst::PadTemplate] { + static T: Lazy> = Lazy::new(|| { + let caps = gst::Caps::builder("video/x-raw") + .field("format", "RGB") + .field("width", gst::IntRange::new(1, i32::MAX)) + .field("height", gst::IntRange::new(1, i32::MAX)) + .build(); + vec![ + gst::PadTemplate::new( + "sink", + gst::PadDirection::Sink, + gst::PadPresence::Always, + &caps, + ) + .unwrap(), + gst::PadTemplate::new( + "src", + gst::PadDirection::Src, + gst::PadPresence::Always, + &caps, + ) + .unwrap(), + ] + }); + T.as_slice() + } +} + +impl EdgeImpulseOcr { + fn build_backend(settings: &Settings) -> Box { + match Backend::parse(&settings.backend) { + Some(Backend::Ocrs) => Self::build_ocrs(settings), + _ => Box::new(NoopBackend), + } + } + + #[cfg(feature = "ocrs")] + fn build_ocrs(settings: &Settings) -> Box { + // Each model loads from its explicit path when set, else from the + // embedded default (see OcrsBackend::new), so partial configuration + // still yields a working engine rather than a silent no-op. + match crate::ocr::ocrs_backend::OcrsBackend::new( + &settings.detection_model, + &settings.recognition_model, + ) { + Ok(b) => Box::new(b), + Err(e) => { + gst::error!( + CAT, + "Failed to build ocrs backend: {e}; falling back to no-op" + ); + Box::new(NoopBackend) + } + } + } + + #[cfg(not(feature = "ocrs"))] + fn build_ocrs(_settings: &Settings) -> Box { + Box::new(NoopBackend) + } +} + +impl BaseTransformImpl for EdgeImpulseOcr { + const MODE: gst_base::subclass::BaseTransformMode = + gst_base::subclass::BaseTransformMode::AlwaysInPlace; + const PASSTHROUGH_ON_SAME_CAPS: bool = false; + const TRANSFORM_IP_ON_PASSTHROUGH: bool = true; + + fn start(&self) -> Result<(), gst::ErrorMessage> { + let settings = self.settings.lock().unwrap().clone(); + self.worker_gone_logged.store(false, Ordering::Relaxed); + let latest = Arc::new(Mutex::new(Latest::default())); + let latest_worker = latest.clone(); + // Bound of 1: at most one frame waits while another is recognized; + // newer frames are dropped (see `try_send` in transform_ip) so a slow + // backend can never build an unbounded backlog. + let (frame_tx, frame_rx) = sync_channel::(1); + // Build the backend inside the worker so model loading never blocks the + // streaming thread, and recognition runs entirely off it. + let worker = std::thread::spawn(move || { + let mut backend = Self::build_backend(&settings); + while let Ok(job) = frame_rx.recv() { + match backend.recognize(&job.rgb, job.width, job.height) { + Ok(lines) => { + let mut latest = latest_worker.lock().unwrap(); + latest.generation = latest.generation.wrapping_add(1); + latest.lines = lines; + latest.pts_ms = job.pts_ms; + } + Err(e) => gst::warning!(CAT, "OCR worker recognition failed: {e}"), + } + } + }); + *self.state.lock().unwrap() = Some(State { + frame_tx, + worker: Some(worker), + latest, + frame_count: 0, + last_posted_generation: 0, + }); + self.parent_start() + } + + fn stop(&self) -> Result<(), gst::ErrorMessage> { + // Take the state out (releasing the lock immediately), then drop the + // sender so the worker's recv() returns Err and the thread exits, and + // join it. The worker never locks `state`, so this cannot deadlock. + let taken = self.state.lock().unwrap().take(); + if let Some(State { + frame_tx, worker, .. + }) = taken + { + drop(frame_tx); + if let Some(handle) = worker { + let _ = handle.join(); + } + } + *self.info.lock().unwrap() = None; + self.parent_stop() + } + + fn set_caps(&self, incaps: &gst::Caps, _outcaps: &gst::Caps) -> Result<(), gst::LoggableError> { + let info = gst_video::VideoInfo::from_caps(incaps) + .map_err(|_| gst::loggable_error!(CAT, "Failed to parse OCR input caps"))?; + *self.info.lock().unwrap() = Some(info); + Ok(()) + } + + fn unit_size(&self, caps: &gst::Caps) -> Option { + gst_video::VideoInfo::from_caps(caps) + .ok() + .map(|info| info.size()) + } + + fn transform_ip(&self, buf: &mut gst::BufferRef) -> Result { + let (interval, min_confidence, max_text_length, post_message) = { + let settings = self.settings.lock().unwrap(); + ( + // max(1) keeps the modulo below divide-by-zero-safe independent + // of the GObject minimum, so a direct Settings construction + // (e.g. in a unit test) can never panic here. + settings.interval.max(1), + settings.min_confidence, + settings.max_text_length, + settings.post_message, + ) + }; + + let info = self.info.lock().unwrap().clone(); + let Some(info) = info else { + return Ok(gst::FlowSuccess::Ok); + }; + + // Under the state lock (kept cheap): bump the frame counter, grab a + // sender if this frame is due for recognition, snapshot the latest + // results + their source PTS, and decide whether they are new enough to + // post. The frame copy and message posting happen after the lock. + let (sender, raw_lines, result_pts_ms, post_new) = { + let mut guard = self.state.lock().unwrap(); + let Some(state) = guard.as_mut() else { + return Ok(gst::FlowSuccess::Ok); + }; + state.frame_count += 1; + let sender = if state.frame_count % interval as u64 == 0 { + Some(state.frame_tx.clone()) + } else { + None + }; + let (lines, generation, pts_ms) = { + let latest = state.latest.lock().unwrap(); + (latest.lines.clone(), latest.generation, latest.pts_ms) + }; + let post_new = post_message && generation != state.last_posted_generation; + if post_new { + state.last_posted_generation = generation; + } + (sender, lines, pts_ms, post_new) + }; + + // Hand the current frame to the worker, dropping it if the worker is + // busy (or gone) so recognition never stalls the streaming thread. + if let Some(sender) = sender { + let pts_ms = buf.pts().map(|t| t.mseconds() as i64).unwrap_or(0); + let job = { + let frame = VideoFrameRef::from_buffer_ref_readable(buf, &info) + .map_err(|_| gst::FlowError::Error)?; + let width = frame.width(); + let height = frame.height(); + let stride = frame.plane_stride()[0] as usize; + let row_bytes = width as usize * 3; + let src = frame.plane_data(0).map_err(|_| gst::FlowError::Error)?; + let mut rgb = Vec::with_capacity(row_bytes * height as usize); + for row in 0..height as usize { + let start = row * stride; + rgb.extend_from_slice(&src[start..start + row_bytes]); + } + FrameJob { + rgb, + width, + height, + pts_ms, + } + }; + if let Err(std::sync::mpsc::TrySendError::Disconnected(_)) = sender.try_send(job) { + // The worker only disconnects if it panicked; warn once so a + // wedged element that silently drops every frame from here on + // is at least diagnosable. + if !self.worker_gone_logged.swap(true, Ordering::Relaxed) { + gst::error!( + CAT, + obj = self.obj(), + "OCR worker thread has exited; recognition stopped" + ); + } + } + } + + // Filter on the streaming thread so `min-confidence` / `max-text-length` + // stay live-settable (the worker stores raw results); this is cheap. + let lines = filter_and_truncate(raw_lines, min_confidence as f32, max_text_length as usize); + + // Timestamp the frame the text was recognized from (carried through the + // worker), not the current output frame which arrives later. + if post_new { + for line in &lines { + let s = build_ocr_message(line, result_pts_ms); + let _ = self + .obj() + .post_message(gst::message::Element::builder(s).src(&*self.obj()).build()); + } + } + attach_results(buf, &lines); + Ok(gst::FlowSuccess::Ok) + } +} diff --git a/src/ocr/mod.rs b/src/ocr/mod.rs new file mode 100644 index 0000000..8c7c41b --- /dev/null +++ b/src/ocr/mod.rs @@ -0,0 +1,36 @@ +//! # EdgeImpulseOcr — OCR text reader element +//! +//! Runs OCR on RGB frames and attaches recognized text as standard +//! `VideoRegionOfInterestMeta` (label = text) so `edgeimpulseoverlay` renders +//! it, and optionally posts an `ocr` element message on the bus. + +mod backend; +mod imp; +mod shaping; + +#[cfg(feature = "ocrs")] +mod ocrs_backend; + +use gstreamer as gst; +use gstreamer::glib; +use gstreamer::prelude::*; + +glib::wrapper! { + pub struct EdgeImpulseOcr(ObjectSubclass) + @extends gstreamer_base::BaseTransform, gst::Element, gst::Object; +} + +pub fn register(plugin: &gst::Plugin) -> Result<(), glib::BoolError> { + let variant = env!("PLUGIN_VARIANT"); + let name = if variant.is_empty() { + "edgeimpulseocr".to_string() + } else { + format!("edgeimpulseocr_{}", variant) + }; + gst::Element::register( + Some(plugin), + &name, + gst::Rank::NONE, + EdgeImpulseOcr::static_type(), + ) +} diff --git a/src/ocr/ocrs_backend.rs b/src/ocr/ocrs_backend.rs new file mode 100644 index 0000000..4d37127 --- /dev/null +++ b/src/ocr/ocrs_backend.rs @@ -0,0 +1,104 @@ +//! Pure-Rust OCR backend built on the `ocrs` engine (`rten` runtime). +use crate::ocr::backend::{OcrBackend, OcrLine}; +use ocrs::{ImageSource, OcrEngine, OcrEngineParams}; +use rten_imageproc::BoundingRect; + +/// Vendored ocrs models, committed under `/models` for offline, +/// reproducible builds. Populated by `examples/download-ocr-models.sh`. +mod embedded { + pub static DETECTION: &[u8] = include_bytes!("../../models/text-detection.rten"); + pub static RECOGNITION: &[u8] = include_bytes!("../../models/text-recognition.rten"); +} + +pub struct OcrsBackend { + engine: OcrEngine, +} + +fn build_engine(detection: rten::Model, recognition: rten::Model) -> Result { + OcrEngine::new(OcrEngineParams { + detection_model: Some(detection), + recognition_model: Some(recognition), + ..Default::default() + }) + .map_err(|e| e.to_string()) +} + +/// Load a model from `path`, or fall back to the `embedded` bytes when `path` +/// is empty (zero-copy over the `&'static` slice compiled into the binary). +fn load_model(path: &str, embedded: &'static [u8]) -> Result { + if path.is_empty() { + rten::Model::load_static_slice(embedded).map_err(|e| e.to_string()) + } else { + rten::Model::load_file(path).map_err(|e| e.to_string()) + } +} + +impl OcrsBackend { + /// Build the engine, loading each model from its explicit path when set and + /// otherwise from the model embedded in the binary. Two empty paths yield + /// the zero-config all-embedded default; a user may also override just one + /// model (e.g. a custom detector) and keep the embedded other. + pub fn new(detection_model_path: &str, recognition_model_path: &str) -> Result { + let detection = load_model(detection_model_path, embedded::DETECTION)?; + let recognition = load_model(recognition_model_path, embedded::RECOGNITION)?; + Ok(Self { + engine: build_engine(detection, recognition)?, + }) + } +} + +impl OcrBackend for OcrsBackend { + fn recognize(&mut self, rgb: &[u8], width: u32, height: u32) -> Result, String> { + let source = ImageSource::from_bytes(rgb, (width, height)).map_err(|e| e.to_string())?; + let input = self + .engine + .prepare_input(source) + .map_err(|e| e.to_string())?; + let words = self + .engine + .detect_words(&input) + .map_err(|e| e.to_string())?; + let line_rects = self.engine.find_text_lines(&input, &words); + let recognized = self + .engine + .recognize_text(&input, &line_rects) + .map_err(|e| e.to_string())?; + + let mut out = Vec::new(); + for (line, word_rects) in recognized.iter().zip(line_rects.iter()) { + let Some(line) = line else { continue }; + let text = line.to_string(); + if text.trim().is_empty() || word_rects.is_empty() { + continue; + } + let (mut min_x, mut min_y, mut max_x, mut max_y) = + (f32::MAX, f32::MAX, f32::MIN, f32::MIN); + for wr in word_rects { + let r = wr.bounding_rect(); + min_x = min_x.min(r.left()); + min_y = min_y.min(r.top()); + max_x = max_x.max(r.right()); + max_y = max_y.max(r.bottom()); + } + // The detector expands each rect by a few pixels, so clamp both + // corners to the frame before deriving the size — otherwise a word + // touching an edge would over-count width/height. + let x0 = min_x.max(0.0); + let y0 = min_y.max(0.0); + let x1 = max_x.min(width as f32); + let y1 = max_y.min(height as f32); + if x1 <= x0 || y1 <= y0 { + continue; + } + out.push(OcrLine { + text, + confidence: 1.0, // ocrs does not expose a per-line confidence + x: x0 as u32, + y: y0 as u32, + w: (x1 - x0) as u32, + h: (y1 - y0) as u32, + }); + } + Ok(out) + } +} diff --git a/src/ocr/shaping.rs b/src/ocr/shaping.rs new file mode 100644 index 0000000..9cf401c --- /dev/null +++ b/src/ocr/shaping.rs @@ -0,0 +1,131 @@ +//! Pure, host-testable helpers: filtering results and turning an [`OcrLine`] +//! into the ROI metadata / bus message the rest of the stack consumes. +use crate::ocr::backend::OcrLine; +use gstreamer as gst; +use gstreamer_video as gst_video; + +/// Drop empty / low-confidence lines and truncate over-long text (by chars). +pub fn filter_and_truncate( + mut lines: Vec, + min_confidence: f32, + max_len: usize, +) -> Vec { + lines.retain(|l| !l.text.trim().is_empty() && l.confidence >= min_confidence); + for l in &mut lines { + if l.text.chars().count() > max_len { + l.text = l.text.chars().take(max_len).collect(); + } + } + lines +} + +/// Attach one `VideoRegionOfInterestMeta` per line with a `detection` param +/// (`label`, `confidence`) — exactly what `edgeimpulseoverlay` renders. +pub fn attach_results(buf: &mut gst::BufferRef, lines: &[OcrLine]) { + for line in lines { + let mut roi = gst_video::VideoRegionOfInterestMeta::add( + buf, + line.text.as_str(), + (line.x, line.y, line.w, line.h), + ); + let s = gst::Structure::builder("detection") + .field("label", line.text.as_str()) + .field("confidence", line.confidence as f64) + .build(); + roi.add_param(s); + } +} + +/// Build an `ocr` element-message structure for one recognized line. +pub fn build_ocr_message(line: &OcrLine, pts_ms: i64) -> gst::Structure { + gst::Structure::builder("ocr") + .field("text", line.text.as_str()) + .field("confidence", line.confidence as f64) + .field("x", line.x as i32) + .field("y", line.y as i32) + .field("width", line.w as i32) + .field("height", line.h as i32) + .field("timestamp", pts_ms) + .build() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn line(text: &str, conf: f32) -> OcrLine { + OcrLine { + text: text.into(), + confidence: conf, + x: 1, + y: 2, + w: 3, + h: 4, + } + } + + #[test] + fn filter_drops_empty_and_low_confidence() { + let out = filter_and_truncate( + vec![line("hi", 0.9), line(" ", 0.9), line("lo", 0.1)], + 0.5, + 10, + ); + assert_eq!(out.len(), 1); + assert_eq!(out[0].text, "hi"); + } + + #[test] + fn filter_truncates_long_text() { + let out = filter_and_truncate(vec![line("abcdef", 1.0)], 0.0, 3); + assert_eq!(out[0].text, "abc"); + } + + #[test] + fn filter_truncates_multibyte_by_chars() { + // Truncation must count characters, not bytes: a byte-slice + // reimplementation would split the 2-byte 'é' and return "á" or panic. + let out = filter_and_truncate(vec![line("áéíóú", 1.0)], 0.0, 2); + assert_eq!(out[0].text, "áé"); + } + + #[test] + fn attaches_one_roi_meta_per_line() { + gst::init().unwrap(); + let mut buf = gst::Buffer::with_size(64).unwrap(); + let b = buf.get_mut().unwrap(); + attach_results(b, &[line("SN-42", 0.8), line("SN-99", 0.6)]); + let metas: Vec<_> = b + .iter_meta::() + .collect(); + assert_eq!(metas.len(), 2); + // Order-independent: each line must yield its own `detection` param with + // the matching label and confidence (catches overwrite / early-return). + let got: Vec<(String, f64)> = metas + .iter() + .map(|m| { + let p = m.params().find(|p| p.name() == "detection").unwrap(); + ( + p.get::("label").unwrap(), + p.get::("confidence").unwrap(), + ) + }) + .collect(); + assert!(got.contains(&("SN-42".to_string(), 0.8f32 as f64))); + assert!(got.contains(&("SN-99".to_string(), 0.6f32 as f64))); + } + + #[test] + fn builds_ocr_message_fields() { + gst::init().unwrap(); + let s = build_ocr_message(&line("hello", 0.7), 1234); + assert_eq!(s.name(), "ocr"); + assert_eq!(s.get::("text").unwrap(), "hello"); + assert_eq!(s.get::("confidence").unwrap(), 0.7f32 as f64); + assert_eq!(s.get::("x").unwrap(), 1); + assert_eq!(s.get::("y").unwrap(), 2); + assert_eq!(s.get::("width").unwrap(), 3); + assert_eq!(s.get::("height").unwrap(), 4); + assert_eq!(s.get::("timestamp").unwrap(), 1234); + } +} diff --git a/tests/ocr.rs b/tests/ocr.rs new file mode 100644 index 0000000..662e823 --- /dev/null +++ b/tests/ocr.rs @@ -0,0 +1,216 @@ +//! Integration tests for the edgeimpulseocr element. +use gstreamer as gst; +use gstreamer::prelude::*; +use std::sync::{Arc, Mutex}; + +fn init() { + static INIT: std::sync::Once = std::sync::Once::new(); + INIT.call_once(|| gst::init().expect("gst init")); +} + +fn ocr_element_name() -> String { + let variant = env!("PLUGIN_VARIANT"); + if variant.is_empty() { + "edgeimpulseocr".into() + } else { + format!("edgeimpulseocr_{variant}") + } +} + +#[test] +fn passes_buffers_through_unchanged() { + init(); + // The plugin is a cdylib GStreamer loads from GST_PLUGIN_PATH (see the Run + // command below), exactly like tests/e2e.rs. Integration tests do NOT link + // the crate (a cdylib has no rlib artifact), so there is nothing to register + // in-process — the element must be discoverable on GST_PLUGIN_PATH. + if gst::ElementFactory::find(&ocr_element_name()).is_none() { + panic!( + "edgeimpulseocr not found — build the plugin and set \ + GST_PLUGIN_PATH=\"$(pwd)/target/release\" (see the Run command)" + ); + } + + let pipeline = gst::parse::launch(&format!( + "videotestsrc num-buffers=3 ! video/x-raw,format=RGB,width=64,height=64 ! \ + videoconvert ! {} ! appsink name=sink", + ocr_element_name() + )) + .unwrap() + .downcast::() + .unwrap(); + + let sink = pipeline + .by_name("sink") + .unwrap() + .downcast::() + .unwrap(); + let count = Arc::new(Mutex::new(0usize)); + let c2 = count.clone(); + sink.set_callbacks( + gstreamer_app::AppSinkCallbacks::builder() + .new_sample(move |s| { + let _ = s.pull_sample().unwrap(); + *c2.lock().unwrap() += 1; + Ok(gst::FlowSuccess::Ok) + }) + .build(), + ); + + pipeline.set_state(gst::State::Playing).unwrap(); + let bus = pipeline.bus().unwrap(); + for msg in bus.iter_timed(gst::ClockTime::from_seconds(10)) { + use gst::MessageView::*; + match msg.view() { + Eos(..) => break, + Error(e) => panic!("{e:?}"), + _ => {} + } + } + pipeline.set_state(gst::State::Null).unwrap(); + assert_eq!(*count.lock().unwrap(), 3); +} + +#[test] +fn exposes_configurable_properties() { + init(); + let e = gst::ElementFactory::make(&ocr_element_name()) + .build() + .unwrap(); + e.set_property("backend", "tesseract"); + e.set_property("min-confidence", 0.5f64); + e.set_property("max-text-length", 32u32); + e.set_property("post-message", false); + e.set_property("interval", 5u32); + assert_eq!(e.property::("backend"), "tesseract"); + assert_eq!(e.property::("min-confidence"), 0.5); + assert_eq!(e.property::("max-text-length"), 32); + assert!(!e.property::("post-message")); + assert_eq!(e.property::("interval"), 5); +} + +#[test] +fn noop_backend_attaches_no_metas() { + init(); + let got_meta = Arc::new(Mutex::new(false)); + let gm = got_meta.clone(); + // `edge-impulse` currently maps to the no-op backend (see `build_backend`), + // giving a deterministic "recognizes nothing" path — this checks the element + // runs in a real pipeline and attaches no ROI metas when there is no result. + let pipeline = gst::parse::launch(&format!( + "videotestsrc num-buffers=2 ! video/x-raw,format=RGB,width=80,height=48 ! \ + videoconvert ! {} backend=edge-impulse interval=1 ! appsink name=sink", + ocr_element_name() + )) + .unwrap() + .downcast::() + .unwrap(); + let sink = pipeline + .by_name("sink") + .unwrap() + .downcast::() + .unwrap(); + sink.set_callbacks( + gstreamer_app::AppSinkCallbacks::builder() + .new_sample(move |s| { + let sample = s.pull_sample().unwrap(); + if let Some(buf) = sample.buffer() { + if buf + .iter_meta::() + .count() + > 0 + { + *gm.lock().unwrap() = true; + } + } + Ok(gst::FlowSuccess::Ok) + }) + .build(), + ); + pipeline.set_state(gst::State::Playing).unwrap(); + for msg in pipeline + .bus() + .unwrap() + .iter_timed(gst::ClockTime::from_seconds(10)) + { + use gst::MessageView::*; + match msg.view() { + Eos(..) => break, + Error(e) => panic!("{e:?}"), + _ => {} + } + } + pipeline.set_state(gst::State::Null).unwrap(); + assert!( + !*got_meta.lock().unwrap(), + "Noop backend must not attach metas" + ); +} + +/// End-to-end check of the ocrs recognition path. Recognition runs on a worker +/// thread, so results land on a *later* output buffer than the frame that +/// triggered them; feeding the same text frame on a loop (`imagefreeze`) lets +/// the worker catch up and attach a `VideoRegionOfInterestMeta` whose `label` +/// carries the recognized text. Skipped unless `OCR_MODELS_DIR` (containing +/// `text-detection.rten` / `text-recognition.rten`) and `OCR_TEST_IMAGE` (an +/// image with legible text) are both set, so CI — which has no models — stays +/// green. +#[test] +fn worker_recognizes_text_and_attaches_results() { + init(); + let (Ok(models), Ok(image)) = ( + std::env::var("OCR_MODELS_DIR"), + std::env::var("OCR_TEST_IMAGE"), + ) else { + eprintln!("skipping: set OCR_MODELS_DIR and OCR_TEST_IMAGE to run"); + return; + }; + let found = Arc::new(Mutex::new(Vec::::new())); + let f = found.clone(); + let pipeline = gst::parse::launch(&format!( + "filesrc location={image} ! decodebin ! imagefreeze ! videoconvert ! \ + video/x-raw,format=RGB ! {elem} backend=ocrs interval=1 \ + detection-model={models}/text-detection.rten \ + recognition-model={models}/text-recognition.rten ! \ + appsink name=sink max-buffers=2 drop=true", + elem = ocr_element_name(), + )) + .unwrap() + .downcast::() + .unwrap(); + let sink = pipeline + .by_name("sink") + .unwrap() + .downcast::() + .unwrap(); + sink.set_callbacks( + gstreamer_app::AppSinkCallbacks::builder() + .new_sample(move |s| { + let sample = s.pull_sample().unwrap(); + if let Some(buf) = sample.buffer() { + for m in buf.iter_meta::() { + if let Some(p) = m.params().find(|p| p.name() == "detection") { + if let Ok(l) = p.get::("label") { + f.lock().unwrap().push(l); + } + } + } + } + Ok(gst::FlowSuccess::Ok) + }) + .build(), + ); + pipeline.set_state(gst::State::Playing).unwrap(); + let start = std::time::Instant::now(); + while found.lock().unwrap().is_empty() && start.elapsed().as_secs() < 20 { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + pipeline.set_state(gst::State::Null).unwrap(); + let labels = found.lock().unwrap(); + eprintln!("recognized labels: {labels:?}"); + assert!( + !labels.is_empty(), + "worker never attached a recognized line within 20s (did the plugin \ + build with --features ocrs, and are the .rten models in OCR_MODELS_DIR?)" + ); +}