aikolipda's picture
Update app.py
3d9b03d
raw
history blame
993 Bytes
from ultralytics import YOLO
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
# Load a model
model = YOLO("best.pt")
def inference(gr_input):
"""
Inference function for gradio.
"""
pred = model(gr_input)
draw_prediction = ImageDraw.Draw(gr_input)
boxes_predict = pred[0].boxes
boxes = boxes_predict.xyxy.tolist()
scores = boxes_predict.conf.tolist()
for score, box in zip(scores, boxes):
x, y, x2, y2 = tuple(box)
draw_prediction.rectangle((x, y, x2, y2), outline="red", width=2)
return gr_input
def main():
imagein = gr.inputs.Image(label="Input Image", type="pil")
imageout = gr.outputs.Image(label="Predicted Image", type="pil")
interface = gr.Interface(
fn=inference,
inputs=imagein,
outputs=imageout,
title="Potholes detection",
interpretation="default",
)
# launch demo
interface.launch(enable_queue=True)
if __name__ == "__main__":
main()