import os
import gradio as gr

from transformers import pipeline
from helper import render_results_in_image

od_pipe = pipeline("object-detection", "facebook/detr-resnet-50")

def get_pipeline_prediction(pil_image):
    pipeline_output = od_pipe(pil_image)
    processed_image = render_results_in_image(pil_image, pipeline_output)
    return processed_image

demo = gr.Interface(
    fn = get_pipeline_prediction,
    inputs = gr.Image(label="Input image", type="pil"),
    outputs = gr.Image(label="Output image with predicted instances", type="pil"), 
    title="Object Detection with DETR-ResNet-50",
    description="Model Info: [facebook/detr-resnet-50](https://huggingface.co/facebook/detr-resnet-50)",
    allow_flagging="never",
    examples=["catdog.jpg", "bus.jpg"]
)

demo.launch()