Spaces:
Runtime error
Runtime error
| import torch | |
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image | |
| model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) | |
| model.conf = 0.25 | |
| model.iou = 0.45 | |
| model.agnostic = False | |
| model.multi_label = False | |
| model.max_det = 1000 | |
| def detect(img): | |
| results = model(img, size=640) | |
| predictions = results.pred[0] | |
| boxes = predictions[:, :4] # x1, y1, x2, y2 | |
| scores = predictions[:, 4] | |
| categories = predictions[:, 5] | |
| new_image = np.squeeze(results.render()) | |
| return new_image | |
| css = ".output-image, .input-image, .image-preview {height: 600px !important}" | |
| iface = gr.Interface(fn=detect, | |
| inputs=gr.inputs.Image(type="numpy",), | |
| outputs=gr.outputs.Image(type="numpy",), | |
| css=css, | |
| enable_queue=True) | |
| iface.launch(debug=True, inline=True) |