soroban-abacus-flashcards/apps/web/scripts/test_boundary_model.py

#!/usr/bin/env python3
"""
Test the boundary detector model locally with a sample image.
Compares Python predictions with what the browser should produce.

Usage:
    python scripts/test_boundary_model.py <image_path>
    python scripts/test_boundary_model.py data/vision-training/boundary-frames/device123/frame001.png
"""

import sys
import os
from pathlib import Path
import json

# Add the training scripts to path
sys.path.insert(0, str(Path(__file__).parent))

import numpy as np
from PIL import Image

# Try to import TensorFlow
try:
    import tensorflow as tf
    print(f"TensorFlow version: {tf.__version__}")
except ImportError:
    print("TensorFlow not installed. Run: pip install tensorflow")
    sys.exit(1)


def load_model():
    """Load the boundary detector model."""
    model_dir = Path(__file__).parent.parent / "public/models/abacus-boundary-detector"
    model_path = model_dir / "model.json"

    if not model_path.exists():
        print(f"Model not found at: {model_path}")
        sys.exit(1)

    print(f"Loading model from: {model_dir}")

    # Try loading as SavedModel format first (graph model)
    saved_model_dir = model_dir

    # For tfjs_graph_model, we need to convert back or use a different approach
    # Let's try loading the original Keras model if it exists
    keras_path = model_dir / "column-classifier.keras"  # Wrong name, let me check

    # Actually, let's load from the saved_model that was used for conversion
    saved_model_path = Path(__file__).parent / "train-boundary-detector/saved_model"

    if saved_model_path.exists():
        print(f"Loading from SavedModel: {saved_model_path}")
        model = tf.saved_model.load(str(saved_model_path))
        return model, "saved_model"

    # Try loading from keras file
    keras_path = model_dir.parent.parent / "data/vision-training/boundary-detector-model.keras"
    if keras_path.exists():
        print(f"Loading from Keras: {keras_path}")
        model = tf.keras.models.load_model(str(keras_path))
        return model, "keras"

    # Look for any .keras or .h5 file
    for pattern in ["*.keras", "*.h5"]:
        for p in model_dir.glob(pattern):
            print(f"Loading from: {p}")
            model = tf.keras.models.load_model(str(p))
            return model, "keras"

    print("Could not find a loadable model format.")
    print("Available files in model dir:", list(model_dir.iterdir()))
    sys.exit(1)


def preprocess_image(image_path: str) -> np.ndarray:
    """Preprocess an image the same way the browser does."""
    img = Image.open(image_path).convert("RGB")
    print(f"Original image size: {img.size}")

    # Resize to model input size (224x224)
    img = img.resize((224, 224), Image.BILINEAR)

    # Convert to numpy array and normalize to [0, 1]
    img_array = np.array(img, dtype=np.float32) / 255.0

    # Add batch dimension
    img_array = np.expand_dims(img_array, axis=0)

    print(f"Preprocessed shape: {img_array.shape}")
    print(f"Value range: [{img_array.min():.3f}, {img_array.max():.3f}]")

    return img_array


def load_ground_truth(image_path: str) -> dict | None:
    """Load ground truth corners from accompanying JSON file."""
    json_path = Path(image_path).with_suffix(".json")
    if json_path.exists():
        with open(json_path) as f:
            return json.load(f)
    return None


def main():
    if len(sys.argv) < 2:
        print(__doc__)

        # Try to find a sample image
        boundary_frames_dir = Path(__file__).parent.parent / "data/vision-training/boundary-frames"
        if boundary_frames_dir.exists():
            for device_dir in boundary_frames_dir.iterdir():
                if device_dir.is_dir():
                    for png in device_dir.glob("*.png"):
                        print(f"\nExample: python {sys.argv[0]} {png}")
                        break
                    break
        sys.exit(1)

    image_path = sys.argv[1]
    if not os.path.exists(image_path):
        print(f"Image not found: {image_path}")
        sys.exit(1)

    # Load model
    model, model_type = load_model()
    print(f"Model loaded ({model_type})")

    # Preprocess image
    input_data = preprocess_image(image_path)

    # Run inference
    print("\nRunning inference...")
    if model_type == "saved_model":
        # For SavedModel, use the serving signature
        infer = model.signatures["serving_default"]
        # Get the input tensor name
        input_name = list(infer.structured_input_signature[1].keys())[0]
        result = infer(**{input_name: tf.constant(input_data)})
        # Get output
        output_name = list(result.keys())[0]
        predictions = result[output_name].numpy()[0]
    else:
        # For Keras model
        predictions = model.predict(input_data, verbose=0)[0]

    print(f"\nRaw predictions: {predictions}")
    print(f"Predictions length: {len(predictions)}")

    # Parse corners (order: topLeft, topRight, bottomLeft, bottomRight)
    corners = {
        "topLeft": {"x": float(predictions[0]), "y": float(predictions[1])},
        "topRight": {"x": float(predictions[2]), "y": float(predictions[3])},
        "bottomLeft": {"x": float(predictions[4]), "y": float(predictions[5])},
        "bottomRight": {"x": float(predictions[6]), "y": float(predictions[7])},
    }

    print("\nParsed corners:")
    for name, coord in corners.items():
        print(f"  {name}: ({coord['x']:.4f}, {coord['y']:.4f})")

    # Load and compare with ground truth
    ground_truth = load_ground_truth(image_path)
    if ground_truth:
        print("\nGround truth corners:")
        for name in ["topLeft", "topRight", "bottomLeft", "bottomRight"]:
            gt = ground_truth["corners"][name]
            pred = corners[name]
            error_x = abs(pred["x"] - gt["x"])
            error_y = abs(pred["y"] - gt["y"])
            print(f"  {name}: ({gt['x']:.4f}, {gt['y']:.4f}) | error: ({error_x:.4f}, {error_y:.4f})")

        # Calculate overall error
        total_error = 0
        for name in ["topLeft", "topRight", "bottomLeft", "bottomRight"]:
            gt = ground_truth["corners"][name]
            pred = corners[name]
            total_error += abs(pred["x"] - gt["x"]) + abs(pred["y"] - gt["y"])
        avg_error = total_error / 8
        print(f"\nAverage coordinate error: {avg_error:.4f}")
    else:
        print("\nNo ground truth JSON file found.")

    # Basic sanity checks
    print("\nSanity checks:")
    checks_passed = True

    # Check 1: Top left should be to the left of top right
    if corners["topLeft"]["x"] >= corners["topRight"]["x"]:
        print("  ❌ topLeft.x >= topRight.x (left should be less than right)")
        checks_passed = False
    else:
        print("  ✓ topLeft.x < topRight.x")

    # Check 2: Bottom left should be to the left of bottom right
    if corners["bottomLeft"]["x"] >= corners["bottomRight"]["x"]:
        print("  ❌ bottomLeft.x >= bottomRight.x")
        checks_passed = False
    else:
        print("  ✓ bottomLeft.x < bottomRight.x")

    # Check 3: Top should be above bottom (y increases downward)
    if corners["topLeft"]["y"] >= corners["bottomLeft"]["y"]:
        print("  ❌ topLeft.y >= bottomLeft.y (top should be above bottom)")
        checks_passed = False
    else:
        print("  ✓ topLeft.y < bottomLeft.y")

    if corners["topRight"]["y"] >= corners["bottomRight"]["y"]:
        print("  ❌ topRight.y >= bottomRight.y")
        checks_passed = False
    else:
        print("  ✓ topRight.y < bottomRight.y")

    # Check 4: Reasonable size
    width = (corners["topRight"]["x"] - corners["topLeft"]["x"] +
             corners["bottomRight"]["x"] - corners["bottomLeft"]["x"]) / 2
    height = (corners["bottomLeft"]["y"] - corners["topLeft"]["y"] +
              corners["bottomRight"]["y"] - corners["topRight"]["y"]) / 2

    print(f"\n  Average width: {width:.4f}")
    print(f"  Average height: {height:.4f}")

    if width < 0.1:
        print("  ❌ Width too small (<0.1)")
        checks_passed = False
    if height < 0.1:
        print("  ❌ Height too small (<0.1)")
        checks_passed = False

    print(f"\nOverall: {'✓ PASS' if checks_passed else '❌ FAIL'}")


if __name__ == "__main__":
    main()