receipt_extractor.py

# receipt_extractor.py
import torch
from transformers import Pix2StructProcessor, Pix2StructForConditionalGeneration
from PIL import Image
import time

# Load model sekali saja (global)
processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-base")
model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-base")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def extract_receipt_data(image_path):
    """Ekstrak data dari gambar nota menggunakan Pix2Struct."""
    image = Image.open(image_path).convert("RGB")
    
    # Daftar pertanyaan yang akan diajukan
    questions = [
        "What is the total amount?",
        "What is the date?",
        "List all items and their prices one by one?"
    ]
    
    extracted = {}
    total_time = 0
    
    for q in questions:
        inputs = processor(images=image, text=q, return_tensors="pt").to(device)
        start = time.time()
        outputs = model.generate(**inputs, max_new_tokens=100)
        duration = time.time() - start
        total_time += duration
        answer = processor.decode(outputs[0], skip_special_tokens=True)
        extracted[q] = answer
    
    extracted["inference_time"] = total_time
    return extracted