File size: 2,436 Bytes
7fe9789 |
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 |
import gradio as gr
import torch
from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
import pytesseract
# Set the Tesseract executable path (for Windows users)
import pytesseract
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:
return {"error": str(e)} # Return error details if any issue occurs
# 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__":
interface.launch(share=True)
|