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 = "
You Only Learn One Representation: Unified Network for Multiple Tasks | Github Repo
" 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()