Spaces:
Sleeping
Sleeping
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 | |