Spaces:
Build error
Build error
File size: 2,688 Bytes
d8e1649 |
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 |
import gradio as gr
from gradio.outputs import Label
import cv2
import requests
import os
import numpy as np
from ultralytics import YOLO
import yolov5
file_urls = [
'https://c8.alamy.com/zooms/9/382c1e254fe14207998df9ec56354291/wce5tj.jpg',
'https://c8.alamy.com/comp/W20YYR/freight-container-on-the-back-of-a-truck-stuck-in-traffic-on-the-interstate-in-georgia-usa-W20YYR.jpg',
'https://www.shutterstock.com/shutterstock/photos/318604739/display_1500/stock-photo-highway-and-container-truck-at-china-318604739.jpg'
]
def download_file(url, save_name):
url = url
if not os.path.exists(save_name):
file = requests.get(url)
open(save_name, 'wb').write(file.content)
for i, url in enumerate(file_urls):
download_file(
file_urls[i],
f"image_{i}.jpg"
)
# model_path = 'Container_YOLOV5'
def yolov5_inference(
image: gr.inputs.Image = None,
model_path: gr.inputs.Dropdown = None,
image_size: gr.inputs.Slider = 640,
conf_threshold: gr.inputs.Slider = 0.25,
iou_threshold: gr.inputs.Slider = 0.45 ):
model = yolov5.load(model_path, device="cpu")
model.conf = conf_threshold
model.iou = iou_threshold
results = model([image], size=image_size)
crops = results.crop(save=False)
img_crops = []
for i in range(len(crops)):
img_crops.append(crops[i]["im"][..., ::-1])
return results.render()[0], img_crops
inputs = [
gr.inputs.Image(type="pil", label="Input Image"),
gr.inputs.Dropdown(["Crime_Y5.pt","yolov5s.pt", "yolov5m.pt", "yolov5l.pt", "yolov5x.pt"], label="Model", default = 'Crime_Y5.pt'),
gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
]
outputs = gr.outputs.Image(type="filepath", label="Output Image")
outputs_crops = gr.Gallery(label="Object crop")
title = "Container code detection - YOLO V5"
description = "YOLOv5 is a family of object detection models pretrained on COCO dataset. This model is a pip implementation of the original YOLOv5 model."
examples = [['1.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]
,['2.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]
,['4.jpg', 'Crime_Y5.pt', 640, 0.35, 0.45]]
demo_app = gr.Interface(
fn=yolov5_inference,
inputs=inputs,
outputs=[outputs,outputs_crops],
title=title,
examples=examples,
cache_examples=True,
live=True,
theme='huggingface',
)
demo_app.launch(debug=True, enable_queue=True, width=50, height=50) |