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() |