feat: implement OCR functionality with WASM integration for text extraction from images

.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

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 {

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
+}

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 }

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.

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<String> = env::args().collect();
+    if args.len() < 3 {
+        bail!("usage: ocr <image_path> <model_path> [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<Vec<TextRegion>> {
+    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<f32>,
+}
+
+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::<f32>()?;
+
+    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<TextRegion>) -> String {
+    if regions.is_empty() {
+        return String::new();
+    }
+
+    regions.sort_by_key(|r| (r.y, r.x));
+    
+    let mut lines: Vec<Vec<TextRegion>> = Vec::new();
+    let mut current_line: Vec<TextRegion> = 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::<Vec<_>>()
+            .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<Vec<char>> {
+    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())
+}