File size: 2,997 Bytes
2a729e6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
"""

Test script for OmniParser API

"""
import requests
import json
import base64
from pathlib import Path

BASE_URL = "http://localhost:8000"


def health_check():
    """Check API health"""
    print("πŸ₯ Health Check...")
    response = requests.get(f"{BASE_URL}/health")
    print(f"Status: {response.status_code}")
    print(json.dumps(response.json(), indent=2))
    print()


def parse_file(image_path: str):
    """Parse image file"""
    print(f"πŸ“Έ Parsing file: {image_path}")
    
    if not Path(image_path).exists():
        print(f"❌ File not found: {image_path}")
        return
    
    with open(image_path, "rb") as f:
        files = {"file": f}
        response = requests.post(f"{BASE_URL}/parse", files=files)
    
    if response.status_code == 200:
        result = response.json()
        print(f"βœ… Found {len(result['elements'])} UI elements")
        print(f"   Image size: {result['image_width']}x{result['image_height']}")
        print(f"   Processing time: {result['processing_time']:.2f}s")
        print("\n   Elements:")
        for elem in result['elements']:
            print(f"   - {elem['label']}: bbox={elem['bbox']}, confidence={elem['confidence']}")
    else:
        print(f"❌ Error: {response.status_code}")
        print(response.text)
    print()


def parse_base64(image_path: str):
    """Parse base64-encoded image"""
    print(f"πŸ“· Parsing base64 image: {image_path}")
    
    if not Path(image_path).exists():
        print(f"❌ File not found: {image_path}")
        return
    
    # Read and encode image
    with open(image_path, "rb") as f:
        image_data = base64.b64encode(f.read()).decode('utf-8')
    
    payload = {
        "image_base64": image_data,
        "extract_text": True,
        "extract_icons": True
    }
    
    response = requests.post(f"{BASE_URL}/parse-base64", json=payload)
    
    if response.status_code == 200:
        result = response.json()
        print(f"βœ… Found {len(result['elements'])} UI elements")
        print(f"   Processing time: {result['processing_time']:.2f}s")
    else:
        print(f"❌ Error: {response.status_code}")
        print(response.text)
    print()


if __name__ == "__main__":
    print("=" * 60)
    print("OmniParser API Test Suite")
    print("=" * 60)
    print()
    
    # Health check
    health_check()
    
    # Test with a sample image (if available)
    sample_images = [
        "screenshot.png",
        "test_image.png",
        "../screenshots/example.png"
    ]
    
    for img in sample_images:
        if Path(img).exists():
            parse_file(img)
            parse_base64(img)
            break
    else:
        print("⚠️  No test images found. Upload an image and try again.")
        print("   Expected: screenshot.png or test_image.png")
    
    print("=" * 60)
    print("βœ… Test suite completed")
    print("=" * 60)