yaswanth-b/Visual_Pollution_Detector

1import gradio as gr2import torch3import yolov54 5def yolov5_inference(6    image: gr.inputs.Image = None,7 8):9    """10    YOLOv5 inference function11    Args:12        image: Input image13        model_path: Path to the model14        image_size: Image size15        conf_threshold: Confidence threshold16        iou_threshold: IOU threshold17    Returns:18        Rendered image19    """20    model = yolov5.load('wts.pt', device="cpu")21    model.conf = 0.2522    model.iou = 0.4523    results = model([image], size=640)24    return results.render()[0]25        26 27inputs = [28    gr.inputs.Image(type="pil", label="Input Image"),29]30 31outputs = gr.outputs.Image(type="filepath", label="Output Image")32title = "Visual Pollution Detection"33description = "<p style='text-align: center'>Built using YOLOv5 and Pytorch."34 35examples = ['ex1.jpg', 'ex2.jpg', 'ex3.jpg', 'ex4.jpg', 'ex5.jpg','ex6.jpg', 'ex7.jpg']36demo_app = gr.Interface(37    fn=yolov5_inference,38    inputs=inputs,39    outputs=outputs,40    title=title,41    examples=examples,42    cache_examples=False,43    live=True,44    theme='huggingface',45)46demo_app.launch(debug=True, enable_queue=True)