diff --git a/.github/workflows/build-ocr-wasm.yml b/.github/workflows/build-ocr-wasm.yml new file mode 100644 index 0000000..8f77815 --- /dev/null +++ b/.github/workflows/build-ocr-wasm.yml @@ -0,0 +1,38 @@ +name: Build OCR WASM + +on: + workflow_dispatch: + +permissions: + contents: write + +jobs: + build-ocr: + runs-on: ubuntu-22.04 + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Install Rust + uses: dtolnay/rust-toolchain@stable + with: + targets: wasm32-wasip1 + + - name: Build ocr.wasm + working-directory: packaging/ocr + run: | + cargo build --release --target wasm32-wasip1 + cp target/wasm32-wasip1/release/ocr.wasm ./ocr.wasm + + - name: Commit ocr.wasm + run: | + git config user.name "github-actions" + git config user.email "github-actions@users.noreply.github.com" + git add packaging/ocr/ocr.wasm + git add packaging/ocr/README.md packaging/ocr/Cargo.toml packaging/ocr/src/main.rs + if git diff --cached --quiet; then + echo "No changes to commit" + exit 0 + fi + git commit -m "build: update ocr.wasm" + git push diff --git a/internal/mcp/filesystem.go b/internal/mcp/filesystem.go index 18e7a27..31448ab 100644 --- a/internal/mcp/filesystem.go +++ b/internal/mcp/filesystem.go @@ -12,6 +12,7 @@ import ( "strings" "time" + "github.com/cheikh2shift/godex/internal/ocr" "github.com/cheikh2shift/godex/internal/vision" ) @@ -48,9 +49,14 @@ func NewFileSystemServer(allowedPaths []string, autoConfirm bool) *FileSystemSer }, { Name: "read_image", - Description: "Analyze an image locally and return a summary for the LLM. Optional prompt adds relevance.", + Description: "Analyze an image to classify what's in it (object, scene, etc.). Returns classification labels with confidence scores.", InputSchema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string","description":"Absolute path to the image file"},"prompt":{"type":"string","description":"User prompt for relevance (optional)"}},"required":["path"]}`), }, + { + Name: "read_text", + Description: "Extract text from an image using OCR. Use this when you need to read text, numbers, or written content from images.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string","description":"Absolute path to the image file"}},"required":["path"]}`), + }, { Name: "write_file", Description: "Write content to a file", @@ -136,6 +142,8 @@ func (s *FileSystemServer) CallTool(ctx context.Context, name string, arguments return s.readFile(arguments) case "read_image": return s.readImage(arguments) + case "read_text": + return s.readText(arguments) case "write_file": return s.writeFile(arguments) case "list_directory": @@ -236,6 +244,26 @@ func (s *FileSystemServer) readImage(args map[string]interface{}) (string, error return summary, nil } +func (s *FileSystemServer) readText(args map[string]interface{}) (string, error) { + path, ok := args["path"].(string) + if !ok { + return "", fmt.Errorf("path is required") + } + + if err := s.checkAndMaybeAddPath(path); err != nil { + return "", err + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + text, err := ocr.ExtractText(ctx, path) + if err != nil { + return "", err + } + return text, nil +} + func (s *FileSystemServer) writeFile(args map[string]interface{}) (string, error) { path, ok := args["path"].(string) if !ok { diff --git a/internal/ocr/ocr.go b/internal/ocr/ocr.go new file mode 100644 index 0000000..019f345 --- /dev/null +++ b/internal/ocr/ocr.go @@ -0,0 +1,200 @@ +package ocr + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" + + "github.com/cheikh2shift/godex/internal/providers" + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" +) + +const ( + wasmURL = "https://raw.githubusercontent.com/cheikh2shift/godex/main/packaging/ocr/ocr.wasm" +) + +type OcrResult struct { + Text string `json:"text"` + Confidence float32 `json:"confidence"` +} + +func EnsureRuntime(ctx context.Context) error { + path, err := requiredPath() + if err != nil { + return err + } + if err := ensureFile(ctx, path, wasmURL); err != nil { + return err + } + return nil +} + +func ExtractText(ctx context.Context, imagePath string) (string, error) { + path, err := requiredPath() + if err != nil { + return "", err + } + if !fileExists(path) { + return "", fmt.Errorf("ocr runtime not available") + } + + result, err := runWasm(ctx, path, imagePath) + if err != nil { + return "", err + } + + if providers.DebugMode { + fmt.Printf("[ocr] Extracted text: %s (confidence: %.2f)\n", result.Text, result.Confidence) + } + + if result.Text == "" { + return "No text detected in image", nil + } + + return result.Text, nil +} + +func requiredPath() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("failed to resolve home dir: %w", err) + } + return filepath.Join(home, ".godex", "models", "ocr", "ocr.wasm"), nil +} + +func ensureFile(ctx context.Context, path, url string) error { + if fileExists(path) { + return nil + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return fmt.Errorf("failed to create directory: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return err + } + client := &http.Client{Timeout: 10 * time.Minute} + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("download failed (%s): HTTP %d", url, resp.StatusCode) + } + + tmp := path + ".tmp" + out, err := os.Create(tmp) + if err != nil { + return err + } + + size := resp.ContentLength + name := filepath.Base(path) + if size > 0 { + fmt.Printf("Downloading %s (%.1f MB)...\n", name, float64(size)/1e6) + } + + _, err = io.Copy(out, &progressReader{body: resp.Body, size: size, name: name}) + if err != nil { + out.Close() + _ = os.Remove(tmp) + return err + } + fmt.Printf("Downloaded %s\n", name) + out.Close() + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + return err + } + return nil +} + +type progressReader struct { + body io.Reader + size int64 + name string + last int64 +} + +func (r *progressReader) Read(p []byte) (int, error) { + n, err := r.body.Read(p) + r.last += int64(n) + if r.size > 0 && r.last > 0 { + pct := float64(r.last) / float64(r.size) * 100 + if pct >= 100 { + pct = 100 + } + fmt.Printf("\r %s: %.1f%%", r.name, pct) + } + return n, err +} + +func fileExists(path string) bool { + st, err := os.Stat(path) + return err == nil && !st.IsDir() +} + +func runWasm(ctx context.Context, wasmPath, imagePath string) (OcrResult, error) { + fmt.Printf("Loading ocr.wasm...\n") + wasmBytes, err := os.ReadFile(wasmPath) + if err != nil { + return OcrResult{}, err + } + fmt.Printf("Initializing Wasm runtime...\n") + + r := wazero.NewRuntime(ctx) + defer r.Close(ctx) + _, err = wasi_snapshot_preview1.Instantiate(ctx, r) + if err != nil { + return OcrResult{}, err + } + + imageDir := filepath.Dir(imagePath) + imageName := filepath.Base(imagePath) + fsConfig := wazero.NewFSConfig() + fsConfig = fsConfig.WithDirMount(imageDir, "/image") + + args := []string{ + "ocr", + "/image/" + imageName, + "/model/dummy.onnx", + } + + fmt.Printf("Running OCR on %s...\n", imageName) + + var stdout bytes.Buffer + var stderr bytes.Buffer + modConfig := wazero.NewModuleConfig(). + WithStdout(&stdout). + WithStderr(&stderr). + WithArgs(args...). + WithFSConfig(fsConfig) + + _, err = r.InstantiateWithConfig(ctx, wasmBytes, modConfig) + if err != nil { + return OcrResult{}, fmt.Errorf("ocr wasm failed: %w: %s", err, strings.TrimSpace(stderr.String())) + } + + out := strings.TrimSpace(stdout.String()) + if out == "" { + return OcrResult{}, fmt.Errorf("ocr wasm returned empty output: %s", strings.TrimSpace(stderr.String())) + } + + var result OcrResult + if err := json.Unmarshal([]byte(out), &result); err != nil { + return OcrResult{}, fmt.Errorf("ocr wasm invalid JSON: %w (out=%s)", err, out) + } + + fmt.Printf("OCR complete.\n") + return result, nil +} diff --git a/packaging/ocr/Cargo.toml b/packaging/ocr/Cargo.toml new file mode 100644 index 0000000..45677c3 --- /dev/null +++ b/packaging/ocr/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "ocr" +version = "0.1.0" +edition = "2021" + +[dependencies] +anyhow = "1" +image = { version = "0.25", default-features = false, features = ["png", "jpeg", "gif"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tract-onnx = { version = "0.21", default-features = false } diff --git a/packaging/ocr/README.md b/packaging/ocr/README.md new file mode 100644 index 0000000..5d43d70 --- /dev/null +++ b/packaging/ocr/README.md @@ -0,0 +1,15 @@ +# OCR WASM Runtime + +This directory contains the Rust/WASI OCR module used by GoDex for text extraction from images. + +## Build (local) + +```bash +rustup target add wasm32-wasip1 +cargo build --release --target wasm32-wasip1 +cp target/wasm32-wasip1/release/ocr.wasm ./ocr.wasm +``` + +## Notes +- The workflow `.github/workflows/build-ocr-wasm.yml` builds `ocr.wasm` and commits it into this directory. +- The Go app downloads `ocr.wasm` into `~/.godex/models/ocr/` on startup. diff --git a/packaging/ocr/src/main.rs b/packaging/ocr/src/main.rs new file mode 100644 index 0000000..1a7e230 --- /dev/null +++ b/packaging/ocr/src/main.rs @@ -0,0 +1,257 @@ +use anyhow::{bail, Context, Result}; +use serde::Serialize; +use std::cmp; +use std::env; +use std::fs; +use std::path::Path; +use tract_onnx::prelude::*; + +const REC_WIDTH: u32 = 128; +const REC_HEIGHT: u32 = 32; +const STRIDE_X: u32 = 64; +const STRIDE_Y: u32 = 16; +const MIN_TEXT_WIDTH: usize = 20; + +#[derive(Serialize, Debug)] +struct OcrResult { + text: String, + confidence: f32, +} + +#[derive(Debug, Clone)] +struct TextRegion { + x: usize, + y: usize, + width: usize, + height: usize, + text: String, + confidence: f32, +} + +fn main() -> Result<()> { + let args: Vec = env::args().collect(); + if args.len() < 3 { + bail!("usage: ocr [dict_path]"); + } + + let image_path = &args[1]; + let model_path = &args[2]; + let dict_path = args.get(3).map(|s| s.as_str()); + + let results = run_ocr(image_path, model_path, dict_path)?; + let combined = combine_text(results); + + let result = OcrResult { + text: combined, + confidence: 0.0, + }; + let out = serde_json::to_string(&result)?; + println!("{}", out); + Ok(()) +} + +fn run_ocr(image_path: &str, model_path: &str, dict_path: Option<&str>) -> Result> { + let model = tract_onnx::onnx() + .model_for_path(model_path) + .with_context(|| format!("failed to load model: {}", model_path))? + .into_optimized()? + .into_runnable()?; + + let img = image::open(image_path).with_context(|| format!("failed to open image: {:?}", image_path))?; + let img = img.to_rgb8(); + + let (img_width, img_height) = img.dimensions(); + + let dict = load_dict(dict_path)?; + let mut regions = Vec::new(); + + let mut y = 0u32; + while y < img_height { + let mut x = 0u32; + while x < img_width { + let window = extract_window(&img, x, y, REC_WIDTH, REC_HEIGHT, img_width, img_height); + + if window.width >= MIN_TEXT_WIDTH { + let tensor = match Tensor::from_shape(&[1, 3, REC_HEIGHT as usize, REC_WIDTH as usize], &window.data) { + Ok(t) => t, + Err(_) => { + x = x.saturating_add(STRIDE_X); + continue; + } + }; + + let result = model.run(tvec!(tensor.into())); + if let Ok(result) = result { + if let Some(output) = result.get(0) { + if let Ok((text, confidence)) = ctc_decode(output, &dict) { + if !text.trim().is_empty() { + regions.push(TextRegion { + x: window.x, + y: window.y, + width: window.width, + height: window.height, + text, + confidence, + }); + } + } + } + } + } + + x = x.saturating_add(STRIDE_X); + } + y = y.saturating_add(STRIDE_Y); + } + + Ok(regions) +} + +struct Window { + x: usize, + y: usize, + width: usize, + height: usize, + data: Vec, +} + +fn extract_window(img: &image::RgbImage, start_x: u32, start_y: u32, target_width: u32, target_height: u32, img_width: u32, img_height: u32) -> Window { + let x = start_x as usize; + let y = start_y as usize; + let width = target_width as usize; + let height = target_height as usize; + + let mut data = vec![0f32; 3 * width * height]; + + for dy in 0..target_height { + for dx in 0..target_width { + let src_x = cmp::min(start_x + dx, img_width.saturating_sub(1)); + let src_y = cmp::min(start_y + dy, img_height.saturating_sub(1)); + + let pixel = img.get_pixel(src_x, src_y); + let idx = ((dy * target_width + dx) * 3) as usize; + data[idx] = pixel[0] as f32 / 255.0; + data[idx + 1] = pixel[1] as f32 / 255.0; + data[idx + 2] = pixel[2] as f32 / 255.0; + } + } + + Window { x, y, width, height, data } +} + +fn ctc_decode(output: &Tensor, dict: &[char]) -> Result<(String, f32)> { + let shape = output.shape(); + if shape.len() < 3 { + bail!("unexpected output shape: {:?}", shape); + } + + let seq_len = shape[1]; + let num_classes = shape[2]; + + let arr = output.to_array_view::()?; + + let mut result_text = String::new(); + let mut total_confidence = 0.0f32; + let mut count = 0usize; + let mut last_char: isize = -1; + let mut consecutive_blanks = 0; + + for t in 0..seq_len { + let max_idx: isize = (0..num_classes) + .max_by(|&a, &b| arr[[0, t, a]].partial_cmp(&arr[[0, t, b]]).unwrap()) + .unwrap() as isize; + + if max_idx == 0 { + consecutive_blanks += 1; + if consecutive_blanks >= 2 && last_char != -1 { + result_text.push(' '); + last_char = -1; + } + } else if max_idx != last_char && (max_idx as usize) < dict.len() { + result_text.push(dict[max_idx as usize]); + total_confidence += arr[[0, t, max_idx as usize]]; + count += 1; + consecutive_blanks = 0; + } + + if max_idx != 0 { + last_char = max_idx; + } + } + + let confidence = if count > 0 { + total_confidence / count as f32 + } else { + 0.0 + }; + + Ok((result_text.trim().to_string(), confidence)) +} + +fn combine_text(mut regions: Vec) -> String { + if regions.is_empty() { + return String::new(); + } + + regions.sort_by_key(|r| (r.y, r.x)); + + let mut lines: Vec> = Vec::new(); + let mut current_line: Vec = Vec::new(); + let mut last_y: isize = -1; + let line_gap_threshold: isize = REC_HEIGHT as isize / 2; + + for region in regions { + let region_y = region.y as isize; + if current_line.is_empty() { + current_line.push(region); + last_y = region_y; + } else if (region_y - last_y).abs() < line_gap_threshold { + current_line.push(region); + last_y = region_y; + } else { + if !current_line.is_empty() { + lines.push(current_line); + } + current_line = vec![region]; + last_y = region_y; + } + } + + if !current_line.is_empty() { + lines.push(current_line); + } + + let mut result = String::new(); + for (i, line) in lines.iter().enumerate() { + let mut sorted_line = line.clone(); + sorted_line.sort_by_key(|r| r.x); + + let line_text: String = sorted_line.iter() + .filter(|r| !r.text.trim().is_empty()) + .map(|r| r.text.trim().to_string()) + .collect::>() + .join(" "); + + if !line_text.trim().is_empty() { + if i > 0 { + result.push('\n'); + } + result.push_str(&line_text); + } + } + + result.trim().to_string() +} + +fn load_dict(path: Option<&str>) -> Result> { + let default_dict = " 0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ.,!?;:'\"-()[]{}"; + + if let Some(dict_path) = path { + if Path::new(dict_path).exists() { + let content = fs::read_to_string(dict_path)?; + return Ok(content.chars().collect()); + } + } + + Ok(default_dict.chars().collect()) +}