add "vision" github action

.github/workflows/build-vision-wasm.yml

@@ -0,0 +1,38 @@
+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-wasi
+
+      - name: Build vision.wasm
+        working-directory: packaging/vision
+        run: |
+          cargo build --release --target wasm32-wasi
+          cp target/wasm32-wasi/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

cmd/godex/main.go

@@ -32,6 +32,7 @@ import (
 	"github.com/cheikh2shift/godex/internal/hive"
 	"github.com/cheikh2shift/godex/internal/mcp"
 	"github.com/cheikh2shift/godex/internal/providers"
+	"github.com/cheikh2shift/godex/internal/vision"
 	"github.com/cheikh2shift/godex/internal/wizard"
 	"github.com/cheikh2shift/godex/modelquery"
 )
@@ -193,6 +194,11 @@ func main() {
 	if strings.TrimSpace(llamaServerURL) != "" {
 		provider.LlamaServerURL = strings.TrimSpace(llamaServerURL)
 	}
+
+	if err := vision.EnsureRuntime(context.Background()); err != nil {
+		log.Printf("[Vision] Setup failed, continuing without vision: %v", err)
+	}
+
 	llmProvider, err := agent.GetProvider(provider)
 	if err != nil {
 		log.Fatalf("failed to create provider: %v", err)
@@ -924,7 +930,7 @@ promptLoop:
 			var toolExecError bool
 
 			if len(toolCalls) > 1 && llmProvider != nil {
-				toolResults, toolExecError = executeToolCallsInParallel(roundCtx, servers, toolCalls, toolTimeout, llmProvider.SupportsNativeToolCalls(), true, false)
+				toolResults, toolExecError = executeToolCallsInParallel(roundCtx, servers, toolCalls, toolTimeout, llmProvider.SupportsNativeToolCalls(), true, false, input)
 
 			} else {
 				hasError := false
@@ -932,6 +938,12 @@ promptLoop:
 					toolName := tc["name"].(string)
 					args := tc["arguments"].(map[string]interface{})
 
+					if toolName == "read_image" && input != "" {
+						if _, ok := args["prompt"]; !ok {
+							args["prompt"] = input
+						}
+					}
+
 					toolDesc := getToolDescription(servers, toolName)
 					if toolDesc != "" {
 						fmt.Print("\033[90m")
@@ -981,7 +993,11 @@ promptLoop:
 						toolResults = append(toolResults, errMsg)
 						hasError = true
 					} else {
-						toolResults = append(toolResults, truncate(result, 2500))
+						if toolName == "read_image" {
+							toolResults = append(toolResults, result)
+						} else {
+							toolResults = append(toolResults, truncate(result, 2500))
+						}
 					}
 				}
 				toolExecError = hasError
@@ -1921,13 +1937,19 @@ func runToolLoop(ctx context.Context, provider *config.Provider, servers []MCPSe
 		var toolResults []string
 		var toolExecError bool
 		if len(toolCalls) > 1 && llmProvider != nil {
-			toolResults, toolExecError = executeToolCallsInParallel(ctx, servers, toolCalls, toolTimeout, llmProvider.SupportsNativeToolCalls(), verbose, autoDenyRestrictedPaths)
+			toolResults, toolExecError = executeToolCallsInParallel(ctx, servers, toolCalls, toolTimeout, llmProvider.SupportsNativeToolCalls(), verbose, autoDenyRestrictedPaths, input)
 		} else {
 			hasError := false
 			for _, tc := range toolCalls {
 				toolName := tc["name"].(string)
 				args := tc["arguments"].(map[string]interface{})
 
+				if toolName == "read_image" && input != "" {
+					if _, ok := args["prompt"]; !ok {
+						args["prompt"] = input
+					}
+				}
+
 				if onToolCall != nil {
 					onToolCall(toolName)
 				}
@@ -1966,7 +1988,11 @@ func runToolLoop(ctx context.Context, provider *config.Provider, servers []MCPSe
 					toolResults = append(toolResults, errMsg)
 					hasError = true
 				} else {
-					toolResults = append(toolResults, truncate(result, 2500))
+					if toolName == "read_image" {
+						toolResults = append(toolResults, result)
+					} else {
+						toolResults = append(toolResults, truncate(result, 2500))
+					}
 				}
 			}
 			toolExecError = hasError
@@ -2252,7 +2278,7 @@ func callTool(ctx context.Context, servers []MCPServer, name string, args map[st
 	return "", fmt.Errorf("tool %s not found", name)
 }
 
-func executeToolCallsInParallel(ctx context.Context, servers []MCPServer, toolCalls []map[string]interface{}, timeoutSecs int, supportsNativeToolCalls bool, verbose bool, autoDenyRestrictedPaths bool) ([]string, bool) {
+func executeToolCallsInParallel(ctx context.Context, servers []MCPServer, toolCalls []map[string]interface{}, timeoutSecs int, supportsNativeToolCalls bool, verbose bool, autoDenyRestrictedPaths bool, promptContext string) ([]string, bool) {
 	if len(toolCalls) == 0 {
 		return nil, false
 	}
@@ -2295,6 +2321,12 @@ func executeToolCallsInParallel(ctx context.Context, servers []MCPServer, toolCa
 			toolName := call["name"].(string)
 			args := call["arguments"].(map[string]interface{})
 
+			if toolName == "read_image" && promptContext != "" {
+				if _, ok := args["prompt"]; !ok {
+					args["prompt"] = promptContext
+				}
+			}
+
 			if supportsNativeToolCalls && toolName == "run_command" {
 				if cmd, ok := args["command"].(string); ok {
 					cmd = fixDriveLetterPath(cmd)
@@ -2341,7 +2373,11 @@ func executeToolCallsInParallel(ctx context.Context, servers []MCPServer, toolCa
 			results[result.index] = fmt.Sprintf("ERROR: %v", result.err)
 			hasError = true
 		} else {
-			results[result.index] = truncate(result.result, 2500)
+			if toolCalls[result.index]["name"].(string) == "read_image" {
+				results[result.index] = result.result
+			} else {
+				results[result.index] = truncate(result.result, 2500)
+			}
 		}
 	}
 
@@ -2435,7 +2471,7 @@ func normalizeToolPathArgs(toolName string, args map[string]interface{}) {
 
 func isFilesystemTool(toolName string) bool {
 	switch toolName {
-	case "read_file", "write_file", "list_directory", "create_directory", "delete_file", "search_files", "get_file_info":
+	case "read_file", "read_image", "write_file", "list_directory", "create_directory", "delete_file", "search_files", "get_file_info":
 		return true
 	default:
 		return false

go.mod

@@ -10,6 +10,7 @@ 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

go.sum

@@ -189,6 +189,8 @@ 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=

internal/mcp/filesystem.go

@@ -10,6 +10,9 @@ import (
 	"path/filepath"
 	"regexp"
 	"strings"
+	"time"
+
+	"github.com/cheikh2shift/godex/internal/vision"
 )
 
 func getIntArg(args map[string]interface{}, key, alias string) (float64, bool) {
@@ -43,6 +46,11 @@ func NewFileSystemServer(allowedPaths []string, autoConfirm bool) *FileSystemSer
 				Description: "Read contents of a file",
 				InputSchema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string","description":"Absolute path to the file"}},"required":["path"]}`),
 			},
+			{
+				Name:        "read_image",
+				Description: "Analyze an image locally and return a summary for the LLM. Optional prompt adds relevance.",
+				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:        "write_file",
 				Description: "Write content to a file",
@@ -126,6 +134,8 @@ func (s *FileSystemServer) CallTool(ctx context.Context, name string, arguments
 	switch name {
 	case "read_file":
 		return s.readFile(arguments)
+	case "read_image":
+		return s.readImage(arguments)
 	case "write_file":
 		return s.writeFile(arguments)
 	case "list_directory":
@@ -205,6 +215,27 @@ func (s *FileSystemServer) readFile(args map[string]interface{}) (string, error)
 	return string(content), nil
 }
 
+func (s *FileSystemServer) readImage(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
+	}
+
+	prompt, _ := args["prompt"].(string)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+	defer cancel()
+
+	summary, err := vision.SummarizeImage(ctx, prompt, path)
+	if err != nil {
+		return "", err
+	}
+	return summary, nil
+}
+
 func (s *FileSystemServer) writeFile(args map[string]interface{}) (string, error) {
 	path, ok := args["path"].(string)
 	if !ok {

internal/vision/vision.go

@@ -0,0 +1,282 @@
+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/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/webnn/mobilenet-v2/resolve/main/onnx/mobilenetv2-10.onnx"
+	labelsURL = "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.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
+	}
+	stats, _ := computeStats(imagePath)
+	lines := buildSummaryLines(prompt, res, stats)
+	return strings.Join(lines, "\n"), nil
+}
+
+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, "mobilenetv2-10.onnx"),
+		LabelsPath: filepath.Join(base, "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
+	}
+	if _, err := io.Copy(out, resp.Body); err != nil {
+		out.Close()
+		_ = os.Remove(tmp)
+		return err
+	}
+	out.Close()
+	if err := os.Rename(tmp, path); err != nil {
+		_ = os.Remove(tmp)
+		return err
+	}
+	return nil
+}
+
+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) {
+	wasmBytes, err := os.ReadFile(paths.WasmPath)
+	if err != nil {
+		return Summary{}, err
+	}
+
+	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),
+	}
+
+	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)
+	}
+
+	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
+}

packaging/vision/Cargo.toml

@@ -0,0 +1,12 @@
+[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, features = ["onnx"] }
+ndarray = "0.15"

packaging/vision/README.md

@@ -0,0 +1,15 @@
+# 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-wasi
+cargo build --release --target wasm32-wasi
+cp target/wasm32-wasi/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.

packaging/vision/src/main.rs

@@ -0,0 +1,107 @@
+use anyhow::{bail, Context, Result};
+use image::{imageops::FilterType, GenericImageView};
+use ndarray::Array4;
+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<LabelScore>,
+}
+
+#[derive(Serialize)]
+struct LabelScore {
+    label: String,
+    score: f32,
+}
+
+fn main() -> Result<()> {
+    let args: Vec<String> = env::args().collect();
+    if args.len() < 4 {
+        bail!("usage: vision <image_path> <model_path> <labels_path>");
+    }
+
+    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))?
+        .with_input_fact(0, InferenceFact::dt_shape(f32::datum_type(), tvec!(1, 3, INPUT_SIZE as usize, INPUT_SIZE as usize)))?
+        .into_optimized()?
+        .into_runnable()?;
+
+    let result = model.run(tvec!(input.into()))?;
+    let scores = result[0].to_array_view::<f32>()?;
+    let scores = scores.iter().cloned().collect::<Vec<f32>>();
+
+    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::<Vec<_>>();
+
+    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<P: AsRef<Path>>(path: P) -> Result<Vec<String>> {
+    let content = fs::read_to_string(&path)?;
+    let labels = content
+        .lines()
+        .map(|l| l.trim().to_string())
+        .filter(|l| !l.is_empty())
+        .collect::<Vec<_>>();
+    if labels.is_empty() {
+        bail!("labels file is empty");
+    }
+    Ok(labels)
+}
+
+fn load_image<P: AsRef<Path>>(path: P) -> Result<Array4<f32>> {
+    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 array = Array4::<f32>::zeros((1, 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 r = pixel[0] as f32 / 255.0;
+            let g = pixel[1] as f32 / 255.0;
+            let b = pixel[2] as f32 / 255.0;
+            array[[0, 0, y as usize, x as usize]] = r;
+            array[[0, 1, y as usize, x as usize]] = g;
+            array[[0, 2, y as usize, x as usize]] = b;
+        }
+    }
+    Ok(array)
+}