test_speed.py

import time
import os
import sys
import cv2

# Ensure we're running from the project root
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from app.extractor_new import extract_document

def test_speed(image_path):
    print(f"Testing speed on: {image_path}")
    start_time = time.time()
    result = extract_document(image_path)
    end_time = time.time()
    
    duration = end_time - start_time
    print(f"\nExtraction took: {duration:.2f} seconds")
    
    # Count blocks
    pages = result.pages
    block_counts = {}
    for page in pages:
        for block in page.blocks:
            block_counts[block.type] = block_counts.get(block.type, 0) + 1
            
    print("Blocks detected:")
    for b_type, count in block_counts.items():
        print(f"  - {b_type}: {count}")

if __name__ == "__main__":
    image_path = "temp-uploads/proc_e1056d38.jpg"
    if not os.path.exists(image_path):
        # Find any other jpg if this one is gone
        jpgs = [f for f in os.listdir("temp-uploads") if f.endswith(".jpg")]
        if not jpgs:
            print("No images found to test.")
            sys.exit(1)
        image_path = os.path.join("temp-uploads", jpgs[0])
        
    test_speed(image_path)