#!/usr/bin/env python
# coding: utf-8

# In[ ]:


# importing required libraries
from transformers import pipeline
import gradio as gr
from PIL import Image, ImageDraw

# main function for object detection
def detector(raw):
    # Resize the image
    WIDTH = 800
    width, height = raw.size
    ratio = float(WIDTH) / float(width)
    new_h = height * ratio
    ip_img = raw.resize((int(WIDTH), int(new_h)), Image.Resampling.LANCZOS)

    # load the model pipeline and predict
    outs = pipeline(model="hustvl/yolos-tiny")(ip_img)

    # draw the image on the canvas
    draw = ImageDraw.Draw(ip_img)

    # draw the boxes with labels
    for object in outs:
        score = f"{object['score']*100:.2f}%"
        label = object['label']
        xmin, ymin, xmax, ymax = object['box'].values()
        draw.rectangle((xmin, ymin, xmax, ymax), outline='red', width=1)
        draw.text((xmin, ymin), f"{label}: {score}", fill="black")
        
    return ip_img

demo = gr.Interface(fn=detector, 
             inputs=gr.Image(type='pil'),
             outputs=gr.Image(type='pil'), allow_flagging=False)
demo.queue(True)
demo.launch(debug=True, inline=False, show_api=False, share=False)