115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test client for the OCR API
|
|
This script demonstrates how to use the OCR API endpoints
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# API base URL
|
|
BASE_URL = "http://localhost:8000"
|
|
API_KEY = os.getenv("API_KEY")
|
|
|
|
def test_health():
|
|
"""Test the health endpoint"""
|
|
try:
|
|
response = requests.get(f"{BASE_URL}/health")
|
|
print(f"Health check: {response.status_code} - {response.json()}")
|
|
return response.status_code == 200
|
|
except requests.exceptions.ConnectionError:
|
|
print("❌ Could not connect to the API. Make sure it's running on localhost:8000")
|
|
return False
|
|
|
|
def test_pdf_ocr(file_path):
|
|
"""Test PDF OCR endpoint"""
|
|
if not Path(file_path).exists():
|
|
print(f"❌ File not found: {file_path}")
|
|
return False
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
files = {'file': f}
|
|
print(f"📄 Processing PDF: {file_path}")
|
|
response = requests.post(f"{BASE_URL}/ocr/pdf", files=files)
|
|
|
|
if response.status_code == 200:
|
|
result = response.json()
|
|
print(f"✅ PDF OCR successful!")
|
|
print(f" Pages: {result['pages']}")
|
|
print(f" Processing time: {result['processing_time']}s")
|
|
print(f" Confidence: {result['confidence']}%")
|
|
print(f" Text preview: {result['text'][:200]}...")
|
|
return True
|
|
else:
|
|
print(f"❌ PDF OCR failed: {response.status_code} - {response.text}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error processing PDF: {str(e)}")
|
|
return False
|
|
|
|
def test_image_ocr(file_path):
|
|
"""Test image OCR endpoint"""
|
|
if not Path(file_path).exists():
|
|
print(f"❌ File not found: {file_path}")
|
|
return False
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
files = {'file': f}
|
|
print(f"🖼️ Processing image: {file_path}")
|
|
response = requests.post(f"{BASE_URL}/ocr/image", files=files)
|
|
|
|
if response.status_code == 200:
|
|
result = response.json()
|
|
print(f"✅ Image OCR successful!")
|
|
print(f" Processing time: {result['processing_time']}s")
|
|
print(f" Confidence: {result['confidence']}%")
|
|
print(f" Text: {result['text']}")
|
|
return True
|
|
else:
|
|
print(f"❌ Image OCR failed: {response.status_code} - {response.text}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error processing image: {str(e)}")
|
|
return False
|
|
|
|
def main():
|
|
"""Main test function"""
|
|
print("🧪 OCR API Test Client")
|
|
print("=" * 50)
|
|
|
|
# Test health endpoint
|
|
if not test_health():
|
|
sys.exit(1)
|
|
|
|
print("\n" + "=" * 50)
|
|
|
|
# Test PDF OCR if file provided
|
|
if len(sys.argv) > 1:
|
|
file_path = sys.argv[1]
|
|
file_ext = Path(file_path).suffix.lower()
|
|
|
|
if file_ext == '.pdf':
|
|
test_pdf_ocr(file_path)
|
|
elif file_ext in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
|
|
test_image_ocr(file_path)
|
|
else:
|
|
print(f"❌ Unsupported file type: {file_ext}")
|
|
print("Supported types: .pdf, .png, .jpg, .jpeg, .bmp, .tiff, .tif")
|
|
else:
|
|
print("📝 Usage: python test_client.py <file_path>")
|
|
print(" Example: python test_client.py document.pdf")
|
|
print(" Example: python test_client.py image.png")
|
|
|
|
if __name__ == "__main__":
|
|
main() |