Spaces:
Runtime error
Runtime error
| 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() | |