Spaces:
Sleeping
Sleeping
import os | |
import easyocr | |
import gradio as gr | |
from PIL import Image | |
from llama_index.core import Settings | |
from llama_index.llms.gemini import Gemini | |
from llama_index.core import Document, VectorStoreIndex | |
from llama_index.embeddings.gemini import GeminiEmbedding | |
from llama_index.core import load_index_from_storage, StorageContext | |
reader = easyocr.Reader(['en']) | |
llm = Gemini(api_key=os.getenv('GEMINI_API_KEY'), model_name="models/gemini-2.0-flash") | |
gemini_embedding_model = GeminiEmbedding(api_key=os.getenv('GEMINI_API_KEY'), model_name="models/embedding-001") | |
# Set Global settings | |
Settings.llm = llm | |
Settings.embed_model = gemini_embedding_model | |
def ocr_inference(img_path, width_ths): | |
output = reader.readtext(img_path, detail=0, slope_ths=0.7, ycenter_ths=0.9, | |
height_ths=0.8, width_ths=width_ths, add_margin=0.2) | |
output = "\n".join(output) | |
doc = Document(text = output) | |
index = VectorStoreIndex.from_documents([doc]) | |
index.storage_context.persist(persist_dir = "./receiptsembeddings") | |
return output | |
def inference(question): | |
persist_dir = "./receiptsembeddings" | |
storage_context = StorageContext.from_defaults(persist_dir = persist_dir) | |
index = load_index_from_storage(storage_context) | |
query_engine = index.as_query_engine() | |
response = query_engine.query(question) | |
return response | |
title = "Receipt RAG" | |
description = "A simple Gradio interface to query receipts using RAG" | |
examples = [["data/receipt_00000.JPG", 7.7], | |
["data/receipt_00001.jpg", 7.7]] | |
with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
gr.Markdown(f"# {title}\n{description}") | |
with gr.Row(): | |
with gr.Column(): | |
image = gr.Image(width=320, height=320, label="Input Receipt") | |
width_ths = gr.Slider(0, 10, 7.7, 0.1, label="Width Threshold to Merge Bounding Boxes") | |
with gr.Row(): | |
clear_btn = gr.ClearButton(components=[image, width_ths]) | |
submit_btn = gr.Button("Submit", variant='primary') | |
with gr.Column(): | |
ocr_out = gr.Textbox(label="OCR Output", type="text") | |
submit_btn.click(ocr_inference, inputs=[image, width_ths], outputs=ocr_out) | |
with gr.Row(): | |
with gr.Column(): | |
text = gr.Textbox(label="Question", type="text") | |
with gr.Row(): | |
chat_clear_btn = gr.ClearButton(components=[text]) | |
chat_submit_btn = gr.Button("Submit", variant='primary') | |
with gr.Column(): | |
chat_out = gr.Textbox(label="Response", type="text") | |
chat_submit_btn.click(inference, inputs=[text], outputs=[chat_out]) | |
examples_obj = gr.Examples(examples=examples, inputs=[image, width_ths]) | |
demo.launch() | |