Spaces:
Sleeping
Sleeping
File size: 2,194 Bytes
778a4f2 00a644a 778a4f2 f29dc51 778333f f29dc51 778333f f29dc51 778a4f2 2f4d6b4 00a644a 2a392dd 8bfe677 778a4f2 |
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 |
import cv2
import gradio as gr
from ultralytics import YOLO
def inference(path:str, threshold:float=0.6):
print("trying inference with path", path)
if path is None:
return None,0
model = YOLO('yolo8n_small.pt') # Caution new API since Ultralytics 8.0.43
outputs = model.predict(source=path, show=False, conf=threshold) # new API with Ultralytics 8.0.43. Accepts 'show', 'classes', 'stream', 'conf' (default is 0.25)
image = cv2.imread(path)
counter = 0
for output in outputs: # mono item batch
conf = output.boxes.conf # the tensor of detection confidences
xyxy = output.boxes.xyxy
cls = output.boxes.cls # 0 is 'person' and 5 is 'bus' 16 is dog
nb=cls.size(dim=0)
for i in range(nb):
box = xyxy[i]
if conf[i]<threshold:
break
cv2.rectangle(
image,
(int(box[0]), int(box[1])),
(int(box[2]), int(box[3])),
color=(0, 0, 255),
thickness=2,
)
# font = cv2.FONT_HERSHEY_SIMPLEX
# fontScale = 1
# fontColor = (255,255,255)
# thickness = 1
# lineType = 2
# cv2.putText(image,
# str(cls),
# (int(box[0]), int(box[2])),
# font,
# fontScale,
# fontColor,
# thickness,
# lineType)
counter+=1
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), counter
gr.Interface(
fn = inference,
inputs = [ gr.components.Image(type="filepath", label="Input"), gr.Slider(minimum=0.0, maximum=0.95, step=0.05, value=0.4, label="Confidence threshold") ],
outputs = [ gr.components.Image(type="numpy", label="Output"), gr.Label(label="Number of legos detected for given confidence threshold") ],
title="YOLOv8n LEGO Detection Demo",
description="LEGO Detection DEMO dataset_small",
examples=[ ['sample1.jpg'],['sample2.jpg'], ['hard.jpg']],
allow_flagging="never"
).launch(debug=True, enable_queue=True) |