test_client.py

#!/usr/bin/env python3
"""
Test client for SPaCial AI OCR Service
Tests the Hugging Face Spaces deployment
"""

import requests
import json
import os
from pathlib import Path

# Configuration
HF_SPACE_URL = "https://your-username-spacial-ai-ocr.hf.space"  # Update with your actual HF Space URL
LOCAL_URL = "http://localhost:7860"  # For local testing

def test_health_check(base_url):
    """Test the health check endpoint"""
    try:
        response = requests.get(f"{base_url}/")
        print(f"Health Check ({base_url}):")
        print(f"Status: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        return response.status_code == 200
    except Exception as e:
        print(f"Health Check Error: {e}")
        return False

def test_ocr_upload(base_url, image_path):
    """Test OCR processing with file upload"""
    if not os.path.exists(image_path):
        print(f"Image file not found: {image_path}")
        return False
    
    try:
        with open(image_path, 'rb') as f:
            files = {'file': f}
            data = {
                'mode': 'accurate',
                'rotation': 0
            }
            
            print(f"\nTesting OCR Upload ({base_url}):")
            print(f"Image: {image_path}")
            print(f"Mode: {data['mode']}")
            print(f"Rotation: {data['rotation']}°")
            
            response = requests.post(f"{base_url}/ocr/process", files=files, data=data)
            
            print(f"Status: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                print(f"Total zones detected: {result['metadata']['total_zones']}")
                
                for i, zone in enumerate(result['zones']):
                    print(f"\nZone {i+1}:")
                    print(f"  Text: {zone['text']}")
                    print(f"  Confidence: {zone['confidence']:.3f}")
                    print(f"  Is Dimension: {zone['is_dimension']}")
                    print(f"  Bounding Box: {zone['bbox']}")
                    
                    if zone['tolerance_info']['plus'] > 0 or zone['tolerance_info']['minus'] > 0:
                        print(f"  Tolerance: +{zone['tolerance_info']['plus']}, -{zone['tolerance_info']['minus']}")
                
                return True
            else:
                print(f"Error: {response.text}")
                return False
                
    except Exception as e:
        print(f"OCR Upload Error: {e}")
        return False

def test_ocr_rotation(base_url, image_path):
    """Test OCR processing with rotation"""
    if not os.path.exists(image_path):
        print(f"Image file not found: {image_path}")
        return False
    
    try:
        with open(image_path, 'rb') as f:
            files = {'file': f}
            data = {
                'mode': 'accurate',
                'rotation': 90  # Test 90° rotation
            }
            
            print(f"\nTesting OCR with Rotation ({base_url}):")
            print(f"Image: {image_path}")
            print(f"Rotation: {data['rotation']}°")
            
            response = requests.post(f"{base_url}/ocr/process", files=files, data=data)
            
            print(f"Status: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                print(f"Total zones detected: {result['metadata']['total_zones']}")
                
                for i, zone in enumerate(result['zones']):
                    print(f"\nZone {i+1} (90° rotated):")
                    print(f"  Text: {zone['text']}")
                    print(f"  Confidence: {zone['confidence']:.3f}")
                    print(f"  Is Dimension: {zone['is_dimension']}")
                
                return True
            else:
                print(f"Error: {response.text}")
                return False
                
    except Exception as e:
        print(f"OCR Rotation Error: {e}")
        return False

def main():
    """Main test function"""
    print("SPaCial AI OCR Service - Test Client")
    print("=" * 50)
    
    # Test image path (update with your test image)
    test_image = "test_image.jpg"
    
    # Check if test image exists
    if not os.path.exists(test_image):
        print(f"Test image not found: {test_image}")
        print("Please place a test image in the same directory as this script.")
        return
    
    # Test local server first
    print("\n1. Testing Local Server")
    print("-" * 30)
    local_health = test_health_check(LOCAL_URL)
    
    if local_health:
        test_ocr_upload(LOCAL_URL, test_image)
        test_ocr_rotation(LOCAL_URL, test_image)
    else:
        print("Local server not available. Skipping local tests.")
    
    # Test Hugging Face Space
    print("\n2. Testing Hugging Face Space")
    print("-" * 30)
    hf_health = test_health_check(HF_SPACE_URL)
    
    if hf_health:
        test_ocr_upload(HF_SPACE_URL, test_image)
        test_ocr_rotation(HF_SPACE_URL, test_image)
    else:
        print("Hugging Face Space not available. Please update HF_SPACE_URL in the script.")
    
    print("\n" + "=" * 50)
    print("Test completed!")

if __name__ == "__main__":
    main()