import numpy as np
from PIL import Image
import gradio as gr
from onnx import hub
import onnxruntime as ort
import onnx



onnx_model = onnx.load("M-Raw.onnx")
text = onnx.checker.check_model(onnx_model)
print("The model is checked")

def snap(image):
    image = Image.fromarray(image) # np to pil
    print(image)
    print("-----------")
    image = image.resize((640, 640))
    print(image)
    print("-----------")
    image = np.asarray(image)
    print(image)
    print("-----------")
    ort_sess = ort.InferenceSession("M-Raw.onnx")
    output = ort_sess.run(["output0"], {"images": np.array([image])})
    return [output]


demo = gr.Interface(
    snap,
    [gr.Image(source="webcam", tool=None, streaming=True)],
    ["image"],
)

if __name__ == "__main__":
    demo.launch()