import gradio as gr
import yolov5
import os
from transformers import pipeline

#ImageClassifier = pipeline(task="image-classification", model="")

model = yolov5.load('./gentle-meadow.pt', device='cpu')

def predict(image):
   results = model([image], size=224)
    #predictions = imageClassifier(image)
    # classMappings = {
    #     'police': "Police / Authorized Personnel",
    #     'public': 'Unauthorized Person'
    # }
    # output = {}
    # for item in predictions:
    #     output[classMappings[item['label']]] = item['score']

    return results.render()[0]

demo = gr.Interface(fn=predict, 
             inputs=gr.inputs.Image(type="pil"),
             outputs=gr.outputs.Image(type="pil"),
             )
             
demo.launch()