Spaces:
Sleeping
Sleeping
File size: 1,539 Bytes
b6edda1 55d549a b6edda1 95e79b5 b6edda1 5bc8768 b6edda1 |
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 |
import os
import torch
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image
import io
import re
HF_TOKEN = os.environ.get("HF_TOKEN")
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model_id = "google/paligemma-3b-mix-224"
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).to(device)
processor = AutoProcessor.from_pretrained(model_id)
def extract_text_from_image(image_content):
image = Image.open(io.BytesIO(image_content))
prompt = "Extract the following details from this invoice: Invoice Number, Total Amount, Invoice Date."
inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
input_len = inputs["input_ids"].shape[-1]
with torch.inference_mode():
generation = model.generate(**inputs, max_new_tokens=100, do_sample=False)
generation = generation[0][input_len:]
decoded = processor.decode(generation, skip_special_tokens=True)
return decoded
def extract_invoice_details(text):
details = {}
details['Invoice Number'] = re.search(r'Invoice Number: (\S+)', text).group(1) if re.search(r'Invoice Number: (\S+)', text) else 'N/A'
details['Amount'] = re.search(r'Total Amount Due: (\S+)', text).group(1) if re.search(r'Total Amount Due: (\S+)', text) else 'N/A'
details['Invoice Date'] = re.search(r'Invoice Date: (\S+)', text).group(1) if re.search(r'Invoice Date: (\S+)', text) else 'N/A'
return details
|