Spaces:
Runtime error
Runtime error
File size: 3,301 Bytes
1a1ee1f 52e322a 1a1ee1f |
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 101 102 |
from PIL import Image
import torch
from numpy import random
from utils.general import (non_max_suppression, scale_coords)
from utils.plots import plot_one_box
from models.models import *
from utils.general import *
import gradio as gr
import requests
import gdown
url = 'https://drive.google.com/u/0/uc?id=1Tdn3yqpZ79X7R1Ql0zNlNScB1Dv9Fp76&export=download'
output = 'yolor_p6.pt'
gdown.download(url, output, quiet=False)
url1 = 'https://cdn.pixabay.com/photo/2014/09/07/21/52/city-438393_1280.jpg'
r = requests.get(url1, allow_redirects=True)
open("city1.jpg", 'wb').write(r.content)
url2 = 'https://cdn.pixabay.com/photo/2016/02/19/11/36/canal-1209808_1280.jpg'
r = requests.get(url2, allow_redirects=True)
open("city2.jpg", 'wb').write(r.content)
conf_thres = 0.4
iou_thres = 0.5
def load_classes(path):
# Loads *.names file at 'path'
with open(path, 'r') as f:
names = f.read().split('\n')
return list(filter(None, names)) # filter removes empty strings (such as last line)
def detect(pil_img,names):
img_np = np.array(pil_img)
img = torch.from_numpy(img_np)
img = img.float()
img /= 255.0 # 0 - 255 to 0.0 - 1.0
# Inference
pred = model(img.unsqueeze(0).permute(0,3,1,2), augment=False)[0].detach()
# Apply NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)
# Process detections
for i, det in enumerate(pred): # detections per image
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img_np.shape, det[:, :4], img_np.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
# Write results
for *xyxy, conf, cls in det:
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy, img_np, label=label, color=colors[int(cls)], line_thickness=3)
return Image.fromarray(img_np)
with torch.no_grad():
cfg = 'cfg/yolor_p6.cfg'
imgsz = 1280
names = 'data/coco.names'
weights = 'yolor_p6.pt'
# Load model
model = Darknet(cfg, imgsz)
model.load_state_dict(torch.load(weights)['model'])
model.eval()
# Get names and colors
names = load_classes(names)
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
def inference(image):
image = image.resize(size=(imgsz, imgsz))
return detect(image, names)
title = "YOLOR P6"
description = "demo for YOLOR. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below.\nModel: YOLOR-P6"
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2105.04206'>You Only Learn One Representation: Unified Network for Multiple Tasks</a> | <a href='https://github.com/WongKinYiu/yolor'>Github Repo</a></p>"
gr.Interface(
inference,
[gr.inputs.Image(type="pil", label="Input")],
gr.outputs.Image(type="numpy", label="Output"),
title=title,
description=description,
article=article,
examples=[
["city1.jpg"],
["city2.jpg"]
]).launch()
|