CoolFace
Apppublic

Propelis/QC_Rules

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
test_google_doc_ai.py183 linesDownload Raw Back to root
1#!/usr/bin/env python3
2"""
3Test script for Google Document AI functionality.
4This script demonstrates the text extraction with bounding boxes and height calculation.
5"""
6
7import os
8import sys
9from pathlib import Path
10
11# Add the src directory to the path so we can import our modules
12sys.path.append(str(Path(__file__).parent / "src"))
13
14from extract_text.google_document_api import GoogleDocumentAPI
15
16def test_google_doc_ai():
17    """Test the Google Document AI functionality with a sample PDF."""
18    
19    # Path to the credentials file
20    credentials_path = "src/extract_text/photon-services-f0d3ec1417d0.json"
21    
22    # Path to a test PDF file
23    test_pdf_path = "requirements_library/client-requirements/Kir-Kat/kitkat-f1.pdf"
24    
25    # Check if files exist
26    if not os.path.exists(credentials_path):
27        print(f"❌ Credentials file not found: {credentials_path}")
28        print("Please ensure the Google Cloud credentials file is in the correct location.")
29        return
30    
31    if not os.path.exists(test_pdf_path):
32        print(f"❌ Test PDF file not found: {test_pdf_path}")
33        print("Please ensure the test PDF file exists.")
34        return
35    
36    print("🔍 Testing Google Document AI functionality...")
37    print(f"📄 Using PDF: {test_pdf_path}")
38    print(f"🔑 Using credentials: {credentials_path}")
39    print("-" * 80)
40    
41    try:
42        # Initialize the Google Document API
43        print("1. Initializing Google Document API...")
44        doc_api = GoogleDocumentAPI(credentials_path)
45        print("✅ Google Document API initialized successfully")
46        
47        # Process the document
48        print("\n2. Processing document...")
49        document = doc_api.process_document(test_pdf_path)
50        print("✅ Document processed successfully")
51        
52        # Get basic text
53        print("\n3. Extracting basic text...")
54        basic_text = doc_api.get_document_text(document, page_number=0)
55        print(f"📝 Basic text length: {len(basic_text)} characters")
56        print(f"📝 First 200 characters: {basic_text[:200]}...")
57        
58        # Extract text with bounding boxes and height
59        print("\n4. Extracting text with bounding boxes and height...")
60        text_blocks = doc_api.extract_text_with_bounding_boxes(document)
61        print(f"📊 Found {len(text_blocks)} text blocks")
62        
63        # Display sample text blocks
64        print("\n5. Sample text blocks with height information:")
65        print("-" * 80)
66        for i, block in enumerate(text_blocks[:10]):  # Show first 10 blocks
67            print(f"Block {i+1}:")
68            print(f"  Page: {block['page_number']}")
69            print(f"  Height: {block['height']:.2f} mm")
70            print(f"  Style: {block['style']}")
71            print(f"  Text: {block['text'][:100]}{'...' if len(block['text']) > 100 else ''}")
72            print(f"  Bounding Box: {block['bounding_box']}")
73            print()
74        
75        # Generate markdown table
76        print("\n6. Generating markdown table...")
77        markdown_table = doc_api.extract_text_with_markdown_table(document)
78        print("📋 Markdown table generated successfully")
79        
80        # Test the new extract_text_heights_mm function
81        print("\n7. Testing extract_text_heights_mm function...")
82        heights_mm = doc_api.extract_text_heights_mm(document)
83        print(f"📏 Found {len(heights_mm)} lines with height in mm")
84        
85        # Display sample heights
86        print("\n📏 Sample line heights (mm):")
87        print("-" * 60)
88        for i, (page_num, line_text, height_mm) in enumerate(heights_mm[:10]):
89            print(f"Line {i+1}: Page {page_num}, Height={height_mm}mm | Text: {line_text[:50]}...")
90        
91        # Save results to files
92        print("\n8. Saving results to files...")
93        
94        # Save raw text blocks
95        with open("test_results_text_blocks.txt", "w", encoding="utf-8") as f:
96            f.write("Text Blocks with Height Information:\n")
97            f.write("=" * 50 + "\n\n")
98            for i, block in enumerate(text_blocks):
99                f.write(f"Block {i+1}:\n")
100                f.write(f"  Page: {block['page_number']}\n")
101                f.write(f"  Height: {block['height']:.2f} mm\n")
102                f.write(f"  Style: {block['style']}\n")
103                f.write(f"  Text: {block['text']}\n")
104                f.write(f"  Bounding Box: {block['bounding_box']}\n")
105                f.write("-" * 40 + "\n")
106        
107        # Save markdown table
108        with open("test_results_markdown_table.md", "w", encoding="utf-8") as f:
109            f.write("# Google Document AI Results\n\n")
110            f.write("## Text Blocks with Height Information\n\n")
111            f.write(markdown_table)
112        
113        # Save basic text
114        with open("test_results_basic_text.txt", "w", encoding="utf-8") as f:
115            f.write("Basic Extracted Text:\n")
116            f.write("=" * 30 + "\n\n")
117            f.write(basic_text)
118        
119        print("✅ Results saved to:")
120        print("   - test_results_text_blocks.txt")
121        print("   - test_results_markdown_table.md")
122        print("   - test_results_basic_text.txt")
123        
124        # Save heights data
125        with open("test_results_heights_mm.txt", "w", encoding="utf-8") as f:
126            f.write("Line Heights in Millimeters:\n")
127            f.write("=" * 40 + "\n\n")
128            for i, (page_num, line_text, height_mm) in enumerate(heights_mm):
129                f.write(f"Line {i+1}: Page {page_num}, Height={height_mm}mm\n")
130                f.write(f"Text: {line_text}\n")
131                f.write("-" * 40 + "\n")
132        
133        print("   - test_results_heights_mm.txt")
134        
135        # Display statistics
136        print("\n9. Statistics:")
137        print("-" * 30)
138        heights = [block['height'] for block in text_blocks]
139        if heights:
140            print(f"📏 Height statistics:")
141            print(f"   Min height: {min(heights):.2f} mm")
142            print(f"   Max height: {max(heights):.2f} mm")
143            print(f"   Average height: {sum(heights)/len(heights):.2f} mm")
144        
145        # Count styles
146        styles = {}
147        for block in text_blocks:
148            style = block['style']
149            styles[style] = styles.get(style, 0) + 1
150        
151        print(f"\n🎨 Style distribution:")
152        for style, count in sorted(styles.items(), key=lambda x: x[1], reverse=True):
153            print(f"   {style}: {count} blocks")
154        
155        print("\n🎉 Test completed successfully!")
156        
157    except Exception as e:
158        print(f"❌ Error during testing: {str(e)}")
159        import traceback
160        traceback.print_exc()
161
162def display_markdown_preview():
163    """Display a preview of the generated markdown table."""
164    try:
165        with open("test_results_markdown_table.md", "r", encoding="utf-8") as f:
166            content = f.read()
167        
168        print("\n📋 Markdown Table Preview:")
169        print("=" * 80)
170        print(content)
171        
172    except FileNotFoundError:
173        print("❌ Markdown table file not found. Run the test first.")
174
175if __name__ == "__main__":
176    print("🚀 Google Document AI Test Script")
177    print("=" * 50)
178    
179    # Run the main test
180    test_google_doc_ai()
181    
182    # Display markdown preview
183    display_markdown_preview()