File size: 1,893 Bytes
92992b8 |
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 |
import gradio as gr
import torch
from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
# 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):
# 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
# Define the Gradio interface
interface = gr.Interface(
fn=process_image,
inputs=gr.Image(type="pil"), # Accepts image input
outputs="json" # Outputs JSON structure
)
# Launch the app
if __name__ == "__main__":
interface.launch(share=False)
|