import gradio as gr
from dotenv import load_dotenv
import os 
import torch
import warnings
from PIL import Image
from util import file_helper
from inference.ocr import prepare_batch_for_inference
from inference.inference_handler import handle

os.system('sudo apt install -y -q tesseract-ocr')
os.system('sudo apt install -y -q libtesseract-dev')

load_dotenv()

def get_model():
    model_dir = "tmp"
    model_filename= 'receipt.pth'
    full_path = os.path.join(model_dir, model_filename)
    if os.path.isfile(full_path):
        return full_path
    
    return file_helper.download_gdrive(os.getenv('MODEL_ID'), model_dir, model_filename)

def run_inference(model_path, images_path):
    try:
        inference_batch = prepare_batch_for_inference(images_path)
        context = {"model_dir": model_path}
        print('handle....')
        handle(inference_batch,context)
    except Exception as err:        
        print('err...', err)


def run(img_path):    
    print('img path: ', img_path)
    model_path = get_model()
    run_inference(model_path, [img_path])

    return Image.open(img_path) 


gr.Markdown('Upload Foto Wajah Kamu (Pastikan hanya terdapat SATU wajah pada)')
iface = gr.Interface(fn=run, inputs=gr.Image(type="filepath"), outputs="image")
iface.launch()