Spaces:
Sleeping
Sleeping
| 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) |