import gradio as gr
import inference_flask as util
import cv2
from PIL import Image

model, transform, device = util.load_model()

def inference_img(input_img):
    global model, transform, device
    input_img = Image.fromarray(input_img)
    output = util.image_inference(model, transform, device, input_img, raw_img_input=True, return_img=True)
    output = cv2.putText(output[0],str(output[1]), (0,70), cv2.FONT_HERSHEY_SIMPLEX, 
                   1, (255, 0, 0), 2, cv2.LINE_AA)
    return output

demo = gr.Interface(inference_img, gr.Image(), "image")
demo.launch(server_port=7860)