Spaces:
Runtime error
Runtime error
File size: 2,991 Bytes
988a945 304ad99 2ea737d 988a945 307f250 41f68b4 4c1a80d 41f68b4 19a7d81 304ad99 7988f17 19a7d81 0d3a1e8 76818d8 7147635 efc152c 2ea737d 5b4740e 2ea737d 5b4740e 2ea737d 5b4740e 2ea737d 5b4740e 2ea737d 5b4740e 2ea737d 5b4740e 2ea737d 5b4740e 64fa63b 5b4740e 2ea737d 7988f17 efc152c 7988f17 19a7d81 9fa2277 ad1c17e 1b61e41 ad1c17e ae0f8d4 19a7d81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import gradio as gr
from PIL import Image
import os
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision
from ultralytics.utils.plotting import Annotator, colors
import numpy as np
import yaml
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = YOLO('Models/best.pt')
model = model.to(device)
def load_img (filename):
if isinstance(img,str):
img = get_url_img(img) if img.startswith('http') else Image.open(img).convert('RGB')
return img
def process_img(image):
with torch.no_grad():
result = model(source=image)
lbel=''
if len(result[0].boxes)>0:
ann=Annotator(im=image)
boxes=result[0].boxes.xyxy
for element in boxes:
box=np.array(element.cpu()).flatten()
if result[0].boxes.cls[0].cpu().numpy()==2:
lbel='car'
if result[0].boxes.cls[0].cpu().numpy()==0:
lbel='bicycle'
if result[0].boxes.cls[0].cpu().numpy()==1:
lbel='bus'
if result[0].boxes.cls[0].cpu().numpy()==3:
lbel='motorcycle'
if result[0].boxes.cls[0].cpu().numpy()==4:
lbel='person'
if result[0].boxes.cls[0].cpu().numpy()==5:
lbel='train'
if result[0].boxes.cls[0].cpu().numpy()==6:
lbel='truck'
ann.box_label(box=box, label=lbel, color=(0,128,0))
vis=ann.result()
else:
vis = image
return vis
title = "Efficient Hazy Vehicle Detection ✏️[] 🤗"
description = ''' ## [Efficient Hazy Vehicle Detection](https://github.com/cidautai)
[Paula Garrido Mellado](https://github.com/paugar5)
Fundación Cidaut
> **Disclaimer:** please remember this is not a product, thus, you will notice some limitations.
**This demo expects an image with some degradations.**
Due to the GPU memory limitations, the app might crash if you feed a high-resolution image (2K, 4K).
<br>
'''
examples = [['examples/dusttornado.jpg'],
['examples/foggy.jpg'],
['examples/haze.jpg'],
["examples/mist.jpg"],
["examples/rain_storm.jpg"],
["examples/sand_storm.jpg"],
["examples/snow_storm.jpg"]]
css = """
.image-frame img, .image-container img {
width: auto;
height: auto;
max-width: none;
}
"""
demo = gr.Interface(
fn = process_img,
inputs = [
gr.Image(type = 'pil', label = 'input')
],
outputs = [gr.Image(type='pil', label = 'output')],
title = title,
description = description,
examples = examples,
css = css
)
if __name__ == '__main__':
demo.launch()
|