File size: 3,974 Bytes
78f16cf |
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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
#!/usr/bin/env python3
"""
granite-docling ONNX Usage Example with ONNX Runtime
Demonstrates how to use the converted granite-docling model for document processing
"""
import onnxruntime as ort
import numpy as np
from PIL import Image
import json
def load_granite_docling_onnx(model_path: str):
"""Load granite-docling ONNX model"""
print(f"Loading granite-docling ONNX model from: {model_path}")
session = ort.InferenceSession(model_path)
# Print model information
print("Model Information:")
print(f" Inputs:")
for inp in session.get_inputs():
print(f" {inp.name}: {inp.shape} ({inp.type})")
print(f" Outputs:")
for out in session.get_outputs():
print(f" {out.name}: {out.shape} ({out.type})")
return session
def preprocess_document_image(image_path: str) -> np.ndarray:
"""Preprocess document image for granite-docling inference"""
# Load and resize image to 512x512 (SigLIP2 requirement)
image = Image.open(image_path).convert('RGB')
image = image.resize((512, 512))
# Convert to numpy array and normalize
pixel_values = np.array(image).astype(np.float32) / 255.0
# Normalize using SigLIP2 parameters (from preprocessor_config.json)
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
pixel_values = (pixel_values - mean) / std
# Reshape to [batch_size, channels, height, width]
pixel_values = pixel_values.transpose(2, 0, 1) # HWC -> CHW
pixel_values = pixel_values[np.newaxis, :] # Add batch dimension
return pixel_values
def create_text_inputs(prompt: str = "Convert this document to DocTags:") -> tuple:
"""Create text inputs for granite-docling"""
# Simple tokenization (in practice, use proper tokenizer)
# This is a simplified example - use actual granite-docling tokenizer
tokens = [1] + [i for i in range(2, len(prompt.split()) + 2)] + [2] # Simple token mapping
input_ids = np.array([tokens], dtype=np.int64)
attention_mask = np.ones((1, len(tokens)), dtype=np.int64)
return input_ids, attention_mask
def run_granite_docling_inference(session, image_path: str):
"""Run complete granite-docling inference"""
print(f"Processing document: {image_path}")
# Prepare inputs
pixel_values = preprocess_document_image(image_path)
input_ids, attention_mask = create_text_inputs()
print(f"Input shapes:")
print(f" pixel_values: {pixel_values.shape}")
print(f" input_ids: {input_ids.shape}")
print(f" attention_mask: {attention_mask.shape}")
# Run inference
outputs = session.run(None, {
'pixel_values': pixel_values,
'input_ids': input_ids,
'attention_mask': attention_mask
})
logits = outputs[0]
print(f"Output logits shape: {logits.shape}")
# Decode logits to tokens (simplified)
predicted_tokens = np.argmax(logits, axis=-1)
print(f"Predicted tokens shape: {predicted_tokens.shape}")
# In practice, decode tokens to DocTags markup using proper tokenizer
print("✅ Inference completed successfully")
return predicted_tokens
def main():
"""Main example usage"""
model_path = "model.onnx" # Path to downloaded ONNX model
try:
# Load model
session = load_granite_docling_onnx(model_path)
# Run inference on example document
# (Replace with actual document image path)
image_path = "example_document.png"
if os.path.exists(image_path):
result = run_granite_docling_inference(session, image_path)
print("✅ granite-docling ONNX inference successful!")
else:
print("⚠️ No example document provided")
print(" Create a test document image to run inference")
except Exception as e:
print(f"❌ Example failed: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
import os
main() |