import spaces
import gradio as gr
from extractors.idefics2json.receipt_extractor import ReceiptExtractor
from indexify_extractor_sdk import Content

receipt_extractor = ReceiptExtractor()

@spaces.GPU
def img2json(image_filepath):
	if image_filepath is None:
		raise gr.Error("Please provide some input image: either upload an image file or use the camera")
	
	with open(image_filepath, "rb") as f:
		image_data = f.read()
	
	content = Content(content_type="image/jpg", data=image_data)
	
	result = receipt_extractor.extract(content)
	text_content = next(content.data.decode('utf-8') for content in result)
	
	return text_content

with gr.Blocks(
	title="Finetuned Idefics2 for Image to JSON with Indexify"
) as receipt_demo:

	gr.HTML("<h1 style='text-align: center'>Finetuned Idefics2 for Image to JSON with Indexify</h1>")
	gr.HTML("<p style='text-align: center'>Indexify is a scalable realtime and continuous indexing and structured extraction engine for unstructured data to build generative AI applications</p>")
	gr.HTML("<h3 style='text-align: center'>If you like this demo, please ⭐ Star us on <a href='https://github.com/tensorlakeai/indexify' target='_blank'>GitHub</a>!</h3>")

	with gr.Row():
		with gr.Column():
			gr.HTML(
				"<p><b>Step 1:</b> Upload an image file or capture with your camera.</p>"

				"<p style='color: #A0A0A0;'>Use this demo for single image file only. "
				"You can extract from image files continuously and try various other extractors locally with "
				"<a href='https://getindexify.io/'>Indexify</a>.</p>"
			)

			image_file = gr.Image(sources=["webcam", "upload"], type="filepath")

		with gr.Column():

			gr.HTML("<p><b>Step 2:</b> Run the extractor.</p>")

			go_button = gr.Button(
				value="Run extractor",
				variant="primary", # make "primary" so it stands out (default is "secondary")
			)

			model_output_text_box = gr.Textbox(
				label="Extractor Output",
				elem_id="model_output_text_box",
			)

	with gr.Row():

		gr.HTML(
			"<p style='text-align: center'>"
				"Developed with 🫶 by <a href='https://getindexify.io/' target='_blank'>Indexify</a> | "
				"a <a href='https://www.tensorlake.ai/' target='_blank'>Tensorlake</a> product"
			"</p>"
		)

	go_button.click(
		fn=img2json, 
		inputs = [image_file],
		outputs = [model_output_text_box]
	)

demo = gr.TabbedInterface([receipt_demo], ["Receipt Extraction"], theme=gr.themes.Soft())

demo.queue()
demo.launch()