diff --git a/.github/workflows/build-ocr-wasm.yml b/.github/workflows/build-ocr-wasm.yml deleted file mode 100644 index 8f77815..0000000 --- a/.github/workflows/build-ocr-wasm.yml +++ /dev/null @@ -1,38 +0,0 @@ -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/.github/workflows/build-vision-wasm.yml b/.github/workflows/build-vision-wasm.yml deleted file mode 100644 index f5c57eb..0000000 --- a/.github/workflows/build-vision-wasm.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: Build Vision WASM - -on: - workflow_dispatch: - -permissions: - contents: write - -jobs: - build-vision: - 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 vision.wasm - working-directory: packaging/vision - run: | - cargo build --release --target wasm32-wasip1 - cp target/wasm32-wasip1/release/vision.wasm ./vision.wasm - - - name: Commit vision.wasm - run: | - git config user.name "github-actions" - git config user.email "github-actions@users.noreply.github.com" - git add packaging/vision/vision.wasm - git add packaging/vision/README.md packaging/vision/Cargo.toml packaging/vision/src/main.rs - if git diff --cached --quiet; then - echo "No changes to commit" - exit 0 - fi - git commit -m "build: update vision.wasm" - git push diff --git a/call_api.sh b/call_api.sh new file mode 100755 index 0000000..c61f804 --- /dev/null +++ b/call_api.sh @@ -0,0 +1,15 @@ +#!/bin/bash +cat > /tmp/test_vision.json << 'EOF' +{ + "model": "lightonocr", + "prompt": "Extract all text from this image:\n", + "multimodal_data": ["FILEPATH:/home/cheikh-seck/.godex/models/vision/test.png"], + "n_predict": 1024, + "temperature": 0.0, + "top_p": 0.9 +} +EOF + +curl -s -X POST http://127.0.0.1:18888/completion \ + -H "Content-Type: application/json" \ + -d @/tmp/test_vision.json | jq -r '.content' 2>/dev/null diff --git a/cmd/godex/main.go b/cmd/godex/main.go index 6d3aa9e..db53ed7 100644 --- a/cmd/godex/main.go +++ b/cmd/godex/main.go @@ -31,8 +31,8 @@ import ( "github.com/cheikh2shift/godex/internal/history" "github.com/cheikh2shift/godex/internal/hive" "github.com/cheikh2shift/godex/internal/mcp" + "github.com/cheikh2shift/godex/internal/ml" "github.com/cheikh2shift/godex/internal/providers" - "github.com/cheikh2shift/godex/internal/vision" "github.com/cheikh2shift/godex/internal/wizard" "github.com/cheikh2shift/godex/modelquery" ) @@ -195,8 +195,67 @@ func main() { provider.LlamaServerURL = strings.TrimSpace(llamaServerURL) } - if err := vision.EnsureRuntime(context.Background()); err != nil { - log.Printf("[Vision] Setup failed, continuing without vision: %v", err) + if cfg.VisionEnabled == nil { + fmt.Print("Enable vision features (image analysis)? (y/N): ") + var response string + fmt.Scanln(&response) + enabled := strings.ToLower(strings.TrimSpace(response)) == "y" + cfg.VisionEnabled = &enabled + if err := config.Save(configPath, cfg); err != nil { + log.Printf("failed to save vision preference: %v", err) + } + } + + if cfg.VisionEnabled != nil && *cfg.VisionEnabled { + if !cfg.VisionCLIDownloaded { + fmt.Println("Vision/OCR requires Tesseract to be installed.") + fmt.Println("") + + var installCmd string + switch runtime.GOOS { + case "darwin": + installCmd = "brew install tesseract" + case "linux": + installCmd = "sudo apt-get update && sudo apt-get install -y tesseract-ocr" + case "windows": + installCmd = "Download from: https://github.com/UB-Mannheim/tesseract/wiki" + default: + installCmd = "Install tesseract-ocr via your package manager" + } + + fmt.Printf("To install, run:\n %s\n", installCmd) + fmt.Println("") + fmt.Print("Install tesseract now? (y/N): ") + + var response string + fmt.Scanln(&response) + if strings.ToLower(strings.TrimSpace(response)) == "y" { + cfg.VisionCLIDownloaded = true + if err := ml.DownloadVisionServer(context.Background()); err != nil { + log.Printf("[Vision] Setup failed: %v", err) + cfg.VisionCLIDownloaded = false + } + if err := config.Save(configPath, cfg); err != nil { + log.Printf("failed to save vision preference: %v", err) + } + } else { + fmt.Println("Vision disabled. You can enable it later in config.") + *cfg.VisionEnabled = false + if err := config.Save(configPath, cfg); err != nil { + log.Printf("failed to save vision preference: %v", err) + } + } + } + + if cfg.VisionCLIDownloaded { + if err := ml.EnsureVisionModel(context.Background()); err != nil { + log.Printf("[Vision] Model setup failed: %v", err) + } else { + if err := ml.StartVisionServer(context.Background()); err != nil { + log.Printf("[Vision] Server start failed: %v", err) + } + } + } } llmProvider, err := agent.GetProvider(provider) @@ -204,12 +263,20 @@ func main() { log.Fatalf("failed to create provider: %v", err) } providers.DebugMode = debugMode + ml.DebugMode = debugMode + ml.PromptInstall = func(tool string) bool { + fmt.Printf("Tool '%s' is needed. Install now? (y/N): ", tool) + var response string + fmt.Scanln(&response) + return strings.ToLower(strings.TrimSpace(response)) == "y" + } if provider == nil { log.Fatalf("no provider configured; use --wizard to create one") } ctx, cancel := context.WithCancel(context.Background()) + defer ml.StopVisionServer() servers, mcpLogs := initMCPServers(ctx, provider, autoConfirm) statusCh := make(chan string, 8) @@ -993,8 +1060,10 @@ promptLoop: toolResults = append(toolResults, errMsg) hasError = true } else { - if toolName == "read_image" { + if toolName == "read_image" || toolName == "read_text" { toolResults = append(toolResults, result) + } else if toolName == "read_pdf" { + toolResults = append(toolResults, truncate(result, 5000)) } else { toolResults = append(toolResults, truncate(result, 2500)) } @@ -1988,8 +2057,10 @@ func runToolLoop(ctx context.Context, provider *config.Provider, servers []MCPSe toolResults = append(toolResults, errMsg) hasError = true } else { - if toolName == "read_image" { + if toolName == "read_image" || toolName == "read_text" { toolResults = append(toolResults, result) + } else if toolName == "read_pdf" { + toolResults = append(toolResults, truncate(result, 5000)) } else { toolResults = append(toolResults, truncate(result, 2500)) } @@ -2375,6 +2446,8 @@ func executeToolCallsInParallel(ctx context.Context, servers []MCPServer, toolCa } else { if toolCalls[result.index]["name"].(string) == "read_image" { results[result.index] = result.result + } else if toolCalls[result.index]["name"].(string) == "read_pdf" { + results[result.index] = truncate(result.result, 5000) } else { results[result.index] = truncate(result.result, 2500) } diff --git a/go.mod b/go.mod index 235f72f..50b3ed2 100644 --- a/go.mod +++ b/go.mod @@ -10,7 +10,6 @@ require ( github.com/chromedp/chromedp v0.14.2 github.com/mark3labs/mcp-go v0.45.0 github.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06 - github.com/tetratelabs/wazero v1.11.0 golang.org/x/net v0.48.0 golang.org/x/term v0.38.0 google.golang.org/genai v1.49.0 diff --git a/go.sum b/go.sum index 7cac0e1..e1352a9 100644 --- a/go.sum +++ b/go.sum @@ -189,8 +189,6 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= -github.com/tetratelabs/wazero v1.11.0 h1:+gKemEuKCTevU4d7ZTzlsvgd1uaToIDtlQlmNbwqYhA= -github.com/tetratelabs/wazero v1.11.0/go.mod h1:eV28rsN8Q+xwjogd7f4/Pp4xFxO7uOGbLcD/LzB1wiU= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= diff --git a/internal/agent/client.go b/internal/agent/client.go index 01ae68a..64ce9b0 100644 --- a/internal/agent/client.go +++ b/internal/agent/client.go @@ -341,6 +341,23 @@ func LaunchToolsInParallel(ctx context.Context, provider *config.Provider, toolC } resultPreview := resp.Result + toolName := toolCalls[i].Name + + if toolName == "read_image" { + // Full output for read_image + resultPreview = resp.Result + } else if toolName == "read_pdf" { + // Truncate read_pdf but with higher limit + if len(resultPreview) > 5000 { + resultPreview = resultPreview[:5000] + "\n... (truncated)" + } + } else { + // Truncate other tools + if len(resultPreview) > 2500 { + resultPreview = resultPreview[:2500] + "\n... (truncated)" + } + } + if len(resultPreview) > 60 { resultPreview = resultPreview[:60] + "..." } diff --git a/internal/config/config.go b/internal/config/config.go index 8920299..ddbbf83 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -37,8 +37,10 @@ type Provider struct { // Config is a list of providers, with a default provider name for the CLI. type Config struct { - Providers []Provider `yaml:"providers"` - DefaultProvider string `yaml:"default_provider,omitempty"` + Providers []Provider `yaml:"providers"` + DefaultProvider string `yaml:"default_provider,omitempty"` + VisionEnabled *bool `yaml:"vision_enabled,omitempty"` + VisionCLIDownloaded bool `yaml:"vision_cli_downloaded,omitempty"` } // Load reads the YAML configuration from the provided path. diff --git a/internal/mcp/filesystem.go b/internal/mcp/filesystem.go index 31448ab..652be56 100644 --- a/internal/mcp/filesystem.go +++ b/internal/mcp/filesystem.go @@ -12,8 +12,7 @@ import ( "strings" "time" - "github.com/cheikh2shift/godex/internal/ocr" - "github.com/cheikh2shift/godex/internal/vision" + "github.com/cheikh2shift/godex/internal/ml" ) func getIntArg(args map[string]interface{}, key, alias string) (float64, bool) { @@ -49,14 +48,24 @@ func NewFileSystemServer(allowedPaths []string, autoConfirm bool) *FileSystemSer }, { Name: "read_image", - 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"]}`), + Description: "Extract text from an image. Use this tool ONLY to extract and transcribe text from images - for OCR, reading screenshots, documents, or any image containing text.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string","description":"Absolute path to the image file"}},"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: "read_pdf", + Description: "Extract text from a PDF file. Use 'pages' to specify page range (e.g., '1-5' or '1,3,5').", + InputSchema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string","description":"Absolute path to the PDF file"},"pages":{"type":"string","description":"Page range to extract (e.g., '1-5' or '1,3,5'). Leave empty for all pages"}},"required":["path"]}`), + }, + { + Name: "read_text", + Description: "Extract text from an image using OCR.", + 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", @@ -144,6 +153,8 @@ func (s *FileSystemServer) CallTool(ctx context.Context, name string, arguments return s.readImage(arguments) case "read_text": return s.readText(arguments) + case "read_pdf": + return s.readPdf(arguments) case "write_file": return s.writeFile(arguments) case "list_directory": @@ -233,11 +244,10 @@ func (s *FileSystemServer) readImage(args map[string]interface{}) (string, error return "", err } - prompt, _ := args["prompt"].(string) - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) defer cancel() - summary, err := vision.SummarizeImage(ctx, prompt, path) + summary, err := ml.QueryImage(ctx, path, "") if err != nil { return "", err } @@ -254,10 +264,32 @@ func (s *FileSystemServer) readText(args map[string]interface{}) (string, error) return "", err } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) defer cancel() - text, err := ocr.ExtractText(ctx, path) + text, err := ml.QueryImage(ctx, path, "Extract all text visible in this image. List each piece of text you see.") + if err != nil { + return "", err + } + return text, nil +} + +func (s *FileSystemServer) readPdf(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 + } + + pages, _ := args["pages"].(string) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + text, err := ml.ExtractPDF(ctx, path, pages) if err != nil { return "", err } diff --git a/internal/ml/vision.go b/internal/ml/vision.go new file mode 100644 index 0000000..47bbc93 --- /dev/null +++ b/internal/ml/vision.go @@ -0,0 +1,236 @@ +package ml + +import ( + "bytes" + "context" + "fmt" + "log" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "sync" +) + +var DebugMode bool + +type VisionResult struct { + Response string `json:"response"` + Duration float64 `json:"duration_ms"` +} + +var ( + visionServer *exec.Cmd + visionServerOnce sync.Once + visionServerErr error + tesseractPath string +) + +func visionDebug(format string, args ...interface{}) { + if DebugMode { + log.Printf("[vision] "+format, args...) + } +} + +var PromptInstall func(string) bool + +func EnsureOCR(ctx context.Context) error { + tesseractPath, err := exec.LookPath("tesseract") + if err == nil { + visionDebug("Found tesseract at: %s", tesseractPath) + return nil + } + + visionDebug("Tesseract not found, prompting user...") + + if PromptInstall != nil && PromptInstall("tesseract-ocr (for reading text in images)") { + var cmd *exec.Cmd + switch runtime.GOOS { + case "darwin": + cmd = exec.CommandContext(ctx, "brew", "install", "tesseract") + case "linux": + cmd = exec.CommandContext(ctx, "sh", "-c", "sudo apt-get update && sudo apt-get install -y tesseract-ocr") + } + if cmd != nil { + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + cmd.Run() + } + + tesseractPath, err := exec.LookPath("tesseract") + if err == nil { + visionDebug("Found tesseract at: %s", tesseractPath) + return nil + } + } + + return fmt.Errorf("tesseract not installed. Please run: sudo apt-get install -y tesseract-ocr") +} + +func EnsurePDFTools(ctx context.Context) error { + cmd := exec.CommandContext(ctx, "pdftotext", "-v") + if err := cmd.Run(); err != nil { + visionDebug("pdftotext not found, prompting user...") + + if PromptInstall != nil && PromptInstall("poppler-utils (for reading PDFs)") { + var installCmd *exec.Cmd + switch runtime.GOOS { + case "darwin": + installCmd = exec.CommandContext(ctx, "brew", "install", "poppler") + case "linux": + installCmd = exec.CommandContext(ctx, "sh", "-c", "sudo apt-get update && sudo apt-get install -y poppler-utils") + } + if installCmd != nil { + installCmd.Stdout = os.Stdout + installCmd.Stderr = os.Stderr + installCmd.Run() + } + + cmd := exec.CommandContext(ctx, "pdftotext", "-v") + if cmd.Run() == nil { + return nil + } + } + + return fmt.Errorf("pdftotext not installed. Run: sudo apt-get install -y poppler-utils") + } + return nil +} + +func QueryImage(ctx context.Context, imagePath, prompt string) (string, error) { + if err := EnsureOCR(ctx); err != nil { + return "", err + } + + visionDebug("Running OCR on: %s", imagePath) + + cmd := exec.CommandContext(ctx, "tesseract", imagePath, "stdout") + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return "", fmt.Errorf("tesseract failed: %w - %s", err, stderr.String()) + } + + result := strings.TrimSpace(stdout.String()) + if result == "" { + return "No text detected in image", nil + } + + return result, nil +} + +type pageRange struct { + start int + end int +} + +func parsePageRange(pages string) pageRange { + pr := pageRange{} + if pages == "" { + return pr + } + + parts := strings.Split(pages, "-") + if len(parts) == 2 { + fmt.Sscanf(parts[0], "%d", &pr.start) + fmt.Sscanf(parts[1], "%d", &pr.end) + } else if len(parts) == 1 { + fmt.Sscanf(parts[0], "%d", &pr.start) + pr.end = pr.start + } + return pr +} + +func ExtractPDF(ctx context.Context, pdfPath, pages string) (string, error) { + if err := EnsurePDFTools(ctx); err != nil { + return "", err + } + + visionDebug("Extracting PDF: %s (pages: %s)", pdfPath, pages) + + args := []string{"-layout"} + + if pages != "" { + parsed := parsePageRange(pages) + if parsed.start > 0 { + args = append(args, "-f", fmt.Sprintf("%d", parsed.start)) + } + if parsed.end > 0 { + args = append(args, "-l", fmt.Sprintf("%d", parsed.end)) + } + } + + args = append(args, pdfPath, "-") + + cmd := exec.CommandContext(ctx, "pdftotext", args...) + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return "", fmt.Errorf("pdftotext failed: %w - %s", err, stderr.String()) + } + + result := strings.TrimSpace(stdout.String()) + if result == "" { + return "No text detected in PDF", nil + } + + return result, nil +} + +func StopVisionServer() { + if visionServer != nil && visionServer.Process != nil { + visionServer.Process.Signal(os.Kill) + visionServer = nil + visionServerOnce = sync.Once{} + } +} + +func StartVisionServer(ctx context.Context) error { + return EnsureOCR(ctx) +} + +func checkServerHealth(url string) bool { + return false +} + +func getGodexBinDir() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".godex", "bin"), nil +} + +func DownloadVisionServer(ctx context.Context) error { + visionDebug("Attempting to install tesseract...") + + var cmd *exec.Cmd + switch runtime.GOOS { + case "darwin": + cmd = exec.CommandContext(ctx, "brew", "install", "tesseract") + case "linux": + cmd = exec.CommandContext(ctx, "sh", "-c", "sudo apt-get update && sudo apt-get install -y tesseract-ocr") + default: + return fmt.Errorf("unsupported OS: %s", runtime.GOOS) + } + + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + cmd.Stdin = os.Stdin + + if err := cmd.Run(); err != nil { + visionDebug("Auto-install failed: %v", err) + return EnsureOCR(ctx) + } + + return EnsureOCR(ctx) +} + +func EnsureVisionModel(ctx context.Context) error { + return EnsureOCR(ctx) +} diff --git a/internal/ocr/ocr.go b/internal/ocr/ocr.go deleted file mode 100644 index 019f345..0000000 --- a/internal/ocr/ocr.go +++ /dev/null @@ -1,200 +0,0 @@ -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/internal/vision/vision.go b/internal/vision/vision.go deleted file mode 100644 index 016ba4d..0000000 --- a/internal/vision/vision.go +++ /dev/null @@ -1,324 +0,0 @@ -package vision - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "image" - _ "image/gif" - _ "image/jpeg" - _ "image/png" - "io" - "math" - "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/vision/vision.wasm" - modelURL = "https://huggingface.co/onnx-community/mobilenetv3_large_100.miil_in21k_ft_in1k/resolve/main/onnx/model.onnx" - labelsURL = "https://huggingface.co/datasets/huggingface/label-files/raw/main/imagenet-1k-id2label.json" -) - -type Summary struct { - Primary string - Top5 []LabelScore - Stats ImageStats -} - -type LabelScore struct { - Label string `json:"label"` - Score float32 `json:"score"` -} - -type ImageStats struct { - Width int - Height int - Aspect float64 - Brightness float64 -} - -type wasmResult struct { - Label string `json:"label"` - Top5 []LabelScore `json:"top5"` -} - -func EnsureRuntime(ctx context.Context) error { - paths, err := requiredPaths() - if err != nil { - return err - } - if err := ensureFile(ctx, paths.WasmPath, wasmURL); err != nil { - return err - } - if err := ensureFile(ctx, paths.ModelPath, modelURL); err != nil { - return err - } - if err := ensureFile(ctx, paths.LabelsPath, labelsURL); err != nil { - return err - } - return nil -} - -func SummarizeImage(ctx context.Context, prompt, imagePath string) (string, error) { - paths, err := requiredPaths() - if err != nil { - return "", err - } - if !fileExists(paths.WasmPath) || !fileExists(paths.ModelPath) || !fileExists(paths.LabelsPath) { - return "", fmt.Errorf("vision runtime not available") - } - - res, err := runWasm(ctx, paths, imagePath) - if err != nil { - return "", err - } - visionDebug("Inference result: primary=%s top5=%v", res.Primary, res.Top5) - stats, _ := computeStats(imagePath) - lines := buildSummaryLines(prompt, res, stats) - return strings.Join(lines, "\n"), nil -} - -func visionDebug(format string, args ...interface{}) { - if providers.DebugMode { - fmt.Printf("[vision] "+format+"\n", args...) - } -} - -type runtimePaths struct { - BaseDir string - WasmPath string - ModelPath string - LabelsPath string -} - -func requiredPaths() (runtimePaths, error) { - home, err := os.UserHomeDir() - if err != nil { - return runtimePaths{}, fmt.Errorf("failed to resolve home dir: %w", err) - } - base := filepath.Join(home, ".godex", "models", "vision") - return runtimePaths{ - BaseDir: base, - WasmPath: filepath.Join(base, "vision.wasm"), - ModelPath: filepath.Join(base, "mobilenetv3-large.onnx"), - LabelsPath: filepath.Join(base, "imagenet-labels.json"), - }, 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, paths runtimePaths, imagePath string) (Summary, error) { - fmt.Printf("Loading vision.wasm...\n") - wasmBytes, err := os.ReadFile(paths.WasmPath) - if err != nil { - return Summary{}, 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 Summary{}, err - } - - imageDir := filepath.Dir(imagePath) - imageName := filepath.Base(imagePath) - fsConfig := wazero.NewFSConfig() - fsConfig = fsConfig.WithDirMount(imageDir, "/image") - fsConfig = fsConfig.WithDirMount(paths.BaseDir, "/model") - - args := []string{ - "vision", - "/image/" + imageName, - "/model/" + filepath.Base(paths.ModelPath), - "/model/" + filepath.Base(paths.LabelsPath), - } - - fmt.Printf("Running model inference 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 Summary{}, fmt.Errorf("vision wasm failed: %w: %s", err, strings.TrimSpace(stderr.String())) - } - - out := strings.TrimSpace(stdout.String()) - if out == "" { - return Summary{}, fmt.Errorf("vision wasm returned empty output: %s", strings.TrimSpace(stderr.String())) - } - var parsed wasmResult - if err := json.Unmarshal([]byte(out), &parsed); err != nil { - return Summary{}, fmt.Errorf("vision wasm invalid JSON: %w (out=%s)", err, out) - } - - fmt.Printf("Inference complete.\n") - return Summary{Primary: parsed.Label, Top5: parsed.Top5}, nil -} - -func computeStats(path string) (ImageStats, error) { - f, err := os.Open(path) - if err != nil { - return ImageStats{}, err - } - defer f.Close() - - img, _, err := image.Decode(f) - if err != nil { - return ImageStats{}, err - } - b := img.Bounds() - w := b.Dx() - h := b.Dy() - aspect := 0.0 - if h > 0 { - aspect = float64(w) / float64(h) - } - brightness := sampleBrightness(img) - return ImageStats{Width: w, Height: h, Aspect: aspect, Brightness: brightness}, nil -} - -func sampleBrightness(img image.Image) float64 { - b := img.Bounds() - w := b.Dx() - h := b.Dy() - if w == 0 || h == 0 { - return 0 - } - stepX := int(math.Max(1, float64(w)/64)) - stepY := int(math.Max(1, float64(h)/64)) - - var total float64 - var count float64 - for y := b.Min.Y; y < b.Max.Y; y += stepY { - for x := b.Min.X; x < b.Max.X; x += stepX { - r, g, b, _ := img.At(x, y).RGBA() - lum := 0.299*float64(r) + 0.587*float64(g) + 0.114*float64(b) - lum = lum / 65535.0 - total += lum - count++ - } - } - if count == 0 { - return 0 - } - return total / count -} - -func buildSummaryLines(prompt string, summary Summary, stats ImageStats) []string { - lines := make([]string, 0, 6) - lines = append(lines, "Local image analysis:") - if summary.Primary != "" { - lines = append(lines, fmt.Sprintf("Primary label: %s", summary.Primary)) - } - if len(summary.Top5) > 0 { - parts := make([]string, 0, len(summary.Top5)) - for _, item := range summary.Top5 { - parts = append(parts, fmt.Sprintf("%s (%.2f)", item.Label, item.Score)) - } - lines = append(lines, "Top-5 labels: "+strings.Join(parts, ", ")) - } - if stats.Width > 0 && stats.Height > 0 { - lines = append(lines, fmt.Sprintf("Image stats: %dx%d (aspect %.2f), brightness %.2f", stats.Width, stats.Height, stats.Aspect, stats.Brightness)) - } - if summary.Primary != "" { - lines = append(lines, fmt.Sprintf("Summary: likely contains %s", summary.Primary)) - } - if strings.TrimSpace(prompt) != "" { - lines = append(lines, fmt.Sprintf("Relevance to question: %s", strings.TrimSpace(prompt))) - } else { - lines = append(lines, "Relevance to question: (no prompt provided)") - } - for len(lines) < 5 { - lines = append(lines, "Additional context: (unavailable)") - } - return lines -} diff --git a/ocr_tool.sh b/ocr_tool.sh new file mode 100644 index 0000000..f271f14 --- /dev/null +++ b/ocr_tool.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# Install tesseract OCR: +# sudo apt-get update && sudo apt-get install -y tesseract-ocr + +IMAGE_PATH="$1" +if [ -z "$IMAGE_PATH" ]; then + echo "Usage: $0 " + exit 1 +fi + +if ! command -v tesseract &> /dev/null; then + echo "Error: tesseract not installed. Run: sudo apt-get install -y tesseract-ocr" + exit 1 +fi + +tesseract "$IMAGE_PATH" stdout diff --git a/packaging/ocr/Cargo.toml b/packaging/ocr/Cargo.toml deleted file mode 100644 index 45677c3..0000000 --- a/packaging/ocr/Cargo.toml +++ /dev/null @@ -1,11 +0,0 @@ -[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 deleted file mode 100644 index 5d43d70..0000000 --- a/packaging/ocr/README.md +++ /dev/null @@ -1,15 +0,0 @@ -# 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 deleted file mode 100644 index 1a7e230..0000000 --- a/packaging/ocr/src/main.rs +++ /dev/null @@ -1,257 +0,0 @@ -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()) -} diff --git a/packaging/vision/Cargo.toml b/packaging/vision/Cargo.toml deleted file mode 100644 index b3afc8f..0000000 --- a/packaging/vision/Cargo.toml +++ /dev/null @@ -1,11 +0,0 @@ -[package] -name = "vision" -version = "0.1.0" -edition = "2021" - -[dependencies] -anyhow = "1" -serde = { version = "1", features = ["derive"] } -serde_json = "1" -image = { version = "0.25", default-features = false, features = ["png", "jpeg", "gif"] } -tract-onnx = { version = "0.21", default-features = false } diff --git a/packaging/vision/README.md b/packaging/vision/README.md deleted file mode 100644 index 05f925e..0000000 --- a/packaging/vision/README.md +++ /dev/null @@ -1,15 +0,0 @@ -# Vision WASM Runtime - -This directory contains the Rust/WASI image analysis module used by GoDex for local image analysis. - -## Build (local) - -```bash -rustup target add wasm32-wasip1 -cargo build --release --target wasm32-wasip1 -cp target/wasm32-wasip1/release/vision.wasm ./vision.wasm -``` - -## Notes -- The workflow `.github/workflows/build-vision-wasm.yml` builds `vision.wasm` and commits it into this directory. -- The Go app downloads `vision.wasm`, the MobileNetV2 ONNX model, and labels into `~/.godex/models/vision/` on startup. diff --git a/packaging/vision/src/main.rs b/packaging/vision/src/main.rs deleted file mode 100644 index 7874bae..0000000 --- a/packaging/vision/src/main.rs +++ /dev/null @@ -1,114 +0,0 @@ -use anyhow::{bail, Context, Result}; -use image::imageops::FilterType; -use serde::Serialize; -use std::env; -use std::fs; -use std::path::Path; -use tract_onnx::prelude::*; - -const INPUT_SIZE: u32 = 224; - -#[derive(Serialize)] -struct ResultPayload { - label: String, - top5: Vec, -} - -#[derive(Serialize)] -struct LabelScore { - label: String, - score: f32, -} - -fn main() -> Result<()> { - let args: Vec = env::args().collect(); - if args.len() < 4 { - bail!("usage: vision "); - } - - let image_path = &args[1]; - let model_path = &args[2]; - let labels_path = &args[3]; - - let labels = load_labels(labels_path)?; - let input = load_image(image_path)?; - - let model = tract_onnx::onnx() - .model_for_path(model_path) - .with_context(|| format!("failed to load model: {}", model_path))? - .into_optimized()? - .into_runnable()?; - - let result = model.run(tvec!(input.into()))?; - let scores = result[0].to_array_view::()?; - let scores = scores.iter().cloned().collect::>(); - - let mut indexed: Vec<(usize, f32)> = scores.into_iter().enumerate().collect(); - indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - - let top5 = indexed - .iter() - .take(5) - .map(|(idx, score)| LabelScore { - label: labels.get(*idx).cloned().unwrap_or_else(|| "unknown".to_string()), - score: *score, - }) - .collect::>(); - - let label = top5 - .first() - .map(|l| l.label.clone()) - .unwrap_or_else(|| "unknown".to_string()); - - let payload = ResultPayload { label, top5 }; - let out = serde_json::to_string(&payload)?; - println!("{}", out); - Ok(()) -} - -fn load_labels>(path: P) -> Result> { - let content = fs::read_to_string(&path)?; - let labels = content - .lines() - .map(|l| l.trim().to_string()) - .filter(|l| !l.is_empty()) - .collect::>(); - if labels.is_empty() { - bail!("labels file is empty"); - } - Ok(labels) -} - -fn load_image>(path: P) -> Result { - let img = image::open(&path).with_context(|| format!("failed to open image: {:?}", path.as_ref()))?; - let img = img.to_rgb8(); - let img = image::imageops::resize(&img, INPUT_SIZE, INPUT_SIZE, FilterType::Triangle); - - let (w, h) = img.dimensions(); - if w != INPUT_SIZE || h != INPUT_SIZE { - bail!("failed to resize image"); - } - - let mut data = vec![0f32; 3 * INPUT_SIZE as usize * INPUT_SIZE as usize]; - for y in 0..INPUT_SIZE { - for x in 0..INPUT_SIZE { - let pixel = img.get_pixel(x, y); - let idx = (0usize * INPUT_SIZE as usize * INPUT_SIZE as usize) - + (y as usize * INPUT_SIZE as usize) - + (x as usize); - data[idx] = pixel[0] as f32 / 255.0; - let idx = (1usize * INPUT_SIZE as usize * INPUT_SIZE as usize) - + (y as usize * INPUT_SIZE as usize) - + (x as usize); - data[idx] = pixel[1] as f32 / 255.0; - let idx = (2usize * INPUT_SIZE as usize * INPUT_SIZE as usize) - + (y as usize * INPUT_SIZE as usize) - + (x as usize); - data[idx] = pixel[2] as f32 / 255.0; - } - } - - let tensor = Tensor::from_shape(&[1, 3, INPUT_SIZE as usize, INPUT_SIZE as usize], &data) - .context("failed to create tensor")?; - Ok(tensor) -} diff --git a/packaging/vision/vision.wasm b/packaging/vision/vision.wasm deleted file mode 100755 index 3e26321..0000000 Binary files a/packaging/vision/vision.wasm and /dev/null differ diff --git a/start_server.sh b/start_server.sh new file mode 100755 index 0000000..3100022 --- /dev/null +++ b/start_server.sh @@ -0,0 +1,12 @@ +#!/bin/bash +~/.godex/bin/llama-b8665/llama-server \ + -m ~/.godex/models/vision/LightOnOCR-2-1B-Q4_K_M.gguf \ + --mmproj ~/.godex/models/vision/LightOnOCR-2-1B-mmproj-f16.gguf \ + --port 18888 --host 127.0.0.1 \ + -ngl 99 \ + -t 16 -tb 16 \ + -c 16384 \ + -fa on \ + --image-min-tokens 256 \ + --image-max-tokens 1024 \ + -v diff --git a/test-vision.sh b/test-vision.sh new file mode 100755 index 0000000..cfce814 --- /dev/null +++ b/test-vision.sh @@ -0,0 +1,14 @@ +#!/bin/bash +IMAGE_PATH=~/.godex/models/vision/test.png + +cat > /tmp/test_vision.json << 'EOF' +{ + "model": "gemma-4", + "prompt": "user\n\nDescribe this image\nmodel\n", + "multimodal_data": ["FILEPATH:/home/cheikh-seck/.godex/models/vision/test.png"], + "n_predict": 200 +} +EOF + +echo "=== Testing /completion with FILEPATH ===" +curl -s -X POST http://127.0.0.1:18888/completion -H "Content-Type: application/json" -d @/tmp/test_vision.json