import gradio as gr
import cv2
from transformers import pipeline

model = pipeline('object-detection')

def draw_box(image):
  img = cv2.imread(image)
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  test = model(image)

  for objects in test:
    if objects['score'] < .5:
      continue

    coord = objects['box']
    label = objects['label']
    color = (0,0,255)
    img = cv2.rectangle(img, (coord['xmin'],coord['ymin']) , (coord['xmax'],coord['ymax']), color,1 )
    img  = cv2.putText(img,label,(coord['xmin'], coord['ymin']-10), cv2.FONT_HERSHEY_PLAIN, 1, color , 2)

  return img


with gr.Blocks() as demo:


  gr.Markdown("""# Object Detection using the Transformers library <br> 
    Enter an Image on the left and view the localized objects on the right.
  """)

  with gr.Row():
    inp = gr.Image( type='filepath')
    out = gr.Image()
  btn = gr.Button('Detect Objects')


  btn.click(fn = draw_box, inputs = inp, outputs = out)

demo.launch()