test_local.py

#!/usr/bin/env python3
"""
Test FoundationPose Space locally before deploying to Hugging Face.

This script tests both placeholder and real modes (if weights available).
"""

import os
import sys
import time
from pathlib import Path

import cv2
import numpy as np

# Set to test placeholder mode
os.environ["USE_REAL_MODEL"] = "false"

print("=" * 60)
print("FoundationPose Local Test")
print("=" * 60)
print()

# Import after setting environment variable
try:
    from app import pose_estimator
    print("✓ Successfully imported app.py")
except Exception as e:
    print(f"✗ Failed to import app.py: {e}")
    sys.exit(1)

print(f"Mode: {'Real' if pose_estimator.use_real_model else 'Placeholder'}")
print()


def test_placeholder_mode():
    """Test the Space in placeholder mode."""
    print("Test 1: Placeholder Mode")
    print("-" * 40)

    # Create dummy reference images
    ref_images = []
    for i in range(5):
        img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
        ref_images.append(img)

    # Test registration
    print("Registering object with 5 reference images...")
    start = time.time()
    success = pose_estimator.register_object(
        object_id="test_object",
        reference_images=ref_images,
        camera_intrinsics={"fx": 500, "fy": 500, "cx": 320, "cy": 240}
    )
    elapsed = time.time() - start

    if success:
        print(f"✓ Registration successful ({elapsed:.2f}s)")
    else:
        print(f"✗ Registration failed")
        return False

    # Test pose estimation
    print("Estimating pose from query image...")
    query_img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)

    start = time.time()
    result = pose_estimator.estimate_pose(
        object_id="test_object",
        query_image=query_img,
        camera_intrinsics={"fx": 500, "fy": 500, "cx": 320, "cy": 240}
    )
    elapsed = time.time() - start

    if result["success"]:
        num_poses = len(result["poses"])
        print(f"✓ Pose estimation successful ({elapsed:.2f}s)")
        print(f"  Detected poses: {num_poses}")
        if num_poses == 0 and "note" in result:
            print(f"  Note: {result['note']}")
        return True
    else:
        print(f"✗ Pose estimation failed: {result.get('error', 'Unknown')}")
        return False


def test_with_reference_images():
    """Test with actual reference images if available."""
    print()
    print("Test 2: Real Reference Images")
    print("-" * 40)

    # Check for reference images
    ref_dir = Path("../training/perception/reference/target_cube")
    if not ref_dir.exists():
        print("⊘ Reference images not found, skipping")
        print(f"  Expected at: {ref_dir}")
        return True

    # Load reference images
    ref_images = []
    for img_path in sorted(ref_dir.glob("*.jpg")):
        img = cv2.imread(str(img_path))
        if img is not None:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            ref_images.append(img)

    if not ref_images:
        print("⊘ No .jpg files found in reference directory")
        return True

    print(f"Found {len(ref_images)} reference images")

    # Test registration
    print("Registering target_cube...")
    start = time.time()
    success = pose_estimator.register_object(
        object_id="target_cube",
        reference_images=ref_images
    )
    elapsed = time.time() - start

    if success:
        print(f"✓ Registration successful ({elapsed:.2f}s)")
    else:
        print(f"✗ Registration failed")
        return False

    # Test pose estimation with first reference image as query
    print("Estimating pose (using first reference image as query)...")
    start = time.time()
    result = pose_estimator.estimate_pose(
        object_id="target_cube",
        query_image=ref_images[0]
    )
    elapsed = time.time() - start

    if result["success"]:
        num_poses = len(result["poses"])
        print(f"✓ Pose estimation successful ({elapsed:.2f}s)")
        print(f"  Detected poses: {num_poses}")

        if num_poses > 0:
            pose = result["poses"][0]
            print(f"  Position: ({pose['position']['x']:.3f}, {pose['position']['y']:.3f}, {pose['position']['z']:.3f})")
            print(f"  Confidence: {pose['confidence']:.3f}")
        else:
            print(f"  Note: {result.get('note', 'No poses detected')}")

        return True
    else:
        print(f"✗ Pose estimation failed: {result.get('error', 'Unknown')}")
        return False


def test_api_format():
    """Test that API format matches expected structure."""
    print()
    print("Test 3: API Format Validation")
    print("-" * 40)

    # Create test object
    ref_img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
    pose_estimator.register_object("api_test", [ref_img])

    # Get result
    result = pose_estimator.estimate_pose("api_test", ref_img)

    # Check format
    required_keys = ["success", "poses"]
    optional_keys = ["error", "note"]

    print("Checking response format...")

    for key in required_keys:
        if key in result:
            print(f"  ✓ Has '{key}' field")
        else:
            print(f"  ✗ Missing '{key}' field")
            return False

    if result["success"]:
        if len(result["poses"]) > 0:
            pose = result["poses"][0]
            pose_required = ["object_id", "position", "orientation", "confidence", "dimensions"]

            for key in pose_required:
                if key in pose:
                    print(f"  ✓ Pose has '{key}' field")
                else:
                    print(f"  ✗ Pose missing '{key}' field")
                    return False

            # Check nested structure
            if isinstance(pose["position"], dict) and "x" in pose["position"]:
                print(f"  ✓ Position format correct")
            else:
                print(f"  ✗ Position format incorrect")
                return False

            if isinstance(pose["orientation"], dict) and "w" in pose["orientation"]:
                print(f"  ✓ Orientation format correct")
            else:
                print(f"  ✗ Orientation format incorrect")
                return False
        else:
            print(f"  ℹ No poses detected (OK for placeholder mode)")

    print("✓ API format valid")
    return True


def main():
    """Run all tests."""
    print("Starting tests...")
    print()

    tests = [
        ("Placeholder Mode", test_placeholder_mode),
        ("Reference Images", test_with_reference_images),
        ("API Format", test_api_format),
    ]

    results = []
    for name, test_func in tests:
        try:
            success = test_func()
            results.append((name, success))
        except Exception as e:
            print(f"✗ Exception in {name}: {e}")
            results.append((name, False))

    # Summary
    print()
    print("=" * 60)
    print("Test Summary")
    print("=" * 60)

    passed = sum(1 for _, success in results if success)
    total = len(results)

    for name, success in results:
        status = "✓ PASS" if success else "✗ FAIL"
        print(f"{status}: {name}")

    print()
    print(f"Results: {passed}/{total} tests passed")

    if passed == total:
        print()
        print("🎉 All tests passed! Ready to deploy.")
        print()
        print("Next steps:")
        print("  1. Run './deploy.sh' to deploy to Hugging Face")
        print("  2. Or start locally: python app.py")
        return 0
    else:
        print()
        print("⚠ Some tests failed. Fix issues before deploying.")
        return 1


if __name__ == "__main__":
    sys.exit(main())