import gradio as gr import torch from PIL import Image # Images torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg') torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', 'bus.jpg') # Model model_name ='fire_model.pt' # force_reload=True to update if model_name: model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_name, force_reload=True) else: model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) def yolo(im, size=640): g = (size / max(im.size)) # gain im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS) # resize results = model(im) # inference results.render() # updates results.imgs with boxes and labels return Image.fromarray(results.imgs[0]) inputs = gr.inputs.Image(type='pil', label="Original Image") outputs = gr.outputs.Image(type="pil", label="Output Image") title = "YOLOv5" description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use." article = "

THis Demo is meant to detect specific models of fire extinguishers , trained on an artificially generated dataset from IFC MODEL with Blender" \ "The aim is to simulate real world fire extinguishers as much possible in order for the object detector to recognizeit" \ examples = [['23327.png'], ['download.jpg']] gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch( debug=True) # try again