File size: 3,042 Bytes
bd746ed a9eefb1 b443de4 a9eefb1 b443de4 b874912 b443de4 bd746ed b443de4 bd746ed b443de4 bd746ed b443de4 |
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 |
import gradio as gr
import torch
from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
import pytesseract
import os
# Explicitly set the Tesseract path for Hugging Face Spaces
pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
# Debugging: Print Tesseract version and PATH details
try:
tesseract_version = pytesseract.get_tesseract_version()
print("Tesseract Version:", tesseract_version)
print("Tesseract Path:", pytesseract.pytesseract.tesseract_cmd)
print("Environment PATH:", os.environ["PATH"])
except Exception as e:
print("Tesseract Debugging Error:", e)
# For local development on Windows
# Uncomment the line below if running locally on Windows
# pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
# Load the model and processor
processor = LayoutLMv3Processor.from_pretrained("quadranttechnologies/Table_OCR")
model = LayoutLMv3ForTokenClassification.from_pretrained("quadranttechnologies/Table_OCR")
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def process_image(image):
try:
# Preprocess the image using the processor
encoding = processor(image, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
# Move inputs to the same device as the model
encoding = {key: val.to(device) for key, val in encoding.items()}
# Perform inference
with torch.no_grad():
outputs = model(**encoding)
predictions = torch.argmax(outputs.logits, dim=-1)
# Extract input IDs, bounding boxes, and predicted labels
words = encoding["input_ids"]
bboxes = encoding["bbox"]
labels = predictions.squeeze().tolist()
# Format output as JSON
structured_output = []
for word_id, bbox, label in zip(words.squeeze().tolist(), bboxes.squeeze().tolist(), labels):
# Decode the word ID to text
word = processor.tokenizer.decode([word_id]).strip()
if word: # Avoid adding empty words
structured_output.append({
"word": word,
"bounding_box": bbox,
"label": model.config.id2label[label] # Convert label ID to label name
})
return structured_output
except Exception as e:
# Debugging: Log any errors encountered during processing
print("Error during processing:", str(e))
return {"error": str(e)}
# Define the Gradio interface
interface = gr.Interface(
fn=process_image,
inputs=gr.Image(type="pil"), # Accepts image input
outputs="json", # Outputs JSON structure
title="Table OCR",
description="Upload an image (e.g., receipt or document) to extract structured information in JSON format."
)
# Launch the app
if __name__ == "__main__":
# Debugging: Check if the app is starting correctly
print("Starting Table OCR App...")
interface.launch(share=True)
|