import gradio as gr

from detection_pipeline import DetectionModel

model = DetectionModel()

def predict(image, threshold):
    preds = model(image)
    preds = list(filter(lambda x: x[4] > threshold, preds))

    output = model.visualize(image, preds)
    return output



interface = gr.Interface(
    fn=predict,
    inputs=[gr.Image(label="Input"), gr.Slider(0, 1, .3, step=.05, label="Threshold")],
    outputs=gr.Image(label="Output")
)

interface.launch()