import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from PIL import Image
import pytesseract  # Install using `pip install pytesseract` and ensure Tesseract is installed

# Load your fine-tuned model and tokenizer
model_name = "quadranttechnologies/Receipt_Image_Analyzer"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define a function to preprocess the image and predict
def analyze_receipt(image):
    # Perform OCR to extract text from the image
    extracted_text = pytesseract.image_to_string(image)
    
    # Tokenize the extracted text
    inputs = tokenizer(extracted_text, return_tensors="pt", truncation=True, padding=True)
    
    # Get model predictions
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = logits.argmax(-1).item()
    
    # Optionally return extracted text and prediction as JSON
    result = {
        "extracted_text": extracted_text,
        "predicted_class": predicted_class
    }
    return result

# Create a Gradio interface
interface = gr.Interface(
    fn=analyze_receipt,
    inputs=gr.inputs.Image(type="pil"),  # Accept image input
    outputs="json",  # Return JSON output
    title="Receipt Image Analyzer",
    description="Upload a receipt image to analyze and classify its contents.",
)

# Launch the Gradio app
interface.launch()