Spaces:
Sleeping
Sleeping
from typing import Tuple | |
import gradio as gr | |
import numpy as np | |
import supervision as sv | |
from inference import get_model | |
MARKDOWN = """ | |
<h1 style='text-align: center'>Detect Something π</h1> | |
Welcome to Detect Something! Just a simple demo to showcase the detection capabilities of various YOLOv8 models. πππ | |
A simple project just for fun for on the go object detection. π | |
Inspired from YOLO-ARENA by SkalskiP. π | |
Powered by Roboflow [Inference](https://github.com/roboflow/inference) and | |
[Supervision](https://github.com/roboflow/supervision). π₯ | |
""" | |
IMAGE_EXAMPLES = [ | |
['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, 0.3, 0.3], | |
['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, 0.3, 0.3], | |
['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.3, 0.3, 0.3], | |
] | |
YOLO_V8N_MODEL = get_model(model_id="yolov8n-640") | |
YOLO_V8S_MODEL = get_model(model_id="yolov8s-640") | |
YOLO_V8M_MODEL = get_model(model_id="yolov8m-640") | |
LABEL_ANNOTATORS = sv.LabelAnnotator(text_color=sv.Color.black()) | |
BOUNDING_BOX_ANNOTATORS = sv.BoundingBoxAnnotator() | |
def detect_and_annotate( | |
model, | |
input_image: np.ndarray, | |
confidence_threshold: float, | |
iou_threshold: float, | |
class_id_mapping: dict = None | |
) -> np.ndarray: | |
result = model.infer( | |
input_image, | |
confidence=confidence_threshold, | |
iou_threshold=iou_threshold | |
)[0] | |
detections = sv.Detections.from_inference(result) | |
if class_id_mapping: | |
detections.class_id = np.array([ | |
class_id_mapping[class_id] | |
for class_id | |
in detections.class_id | |
]) | |
labels = [ | |
f"{class_name} ({confidence:.2f})" | |
for class_name, confidence | |
in zip(detections['class_name'], detections.confidence) | |
] | |
annotated_image = input_image.copy() | |
annotated_image = BOUNDING_BOX_ANNOTATORS.annotate( | |
scene=annotated_image, detections=detections) | |
annotated_image = LABEL_ANNOTATORS.annotate( | |
scene=annotated_image, detections=detections, labels=labels) | |
return annotated_image | |
def process_image( | |
input_image: np.ndarray, | |
yolo_v8_confidence_threshold: float, | |
yolo_v9_confidence_threshold: float, | |
yolo_v10_confidence_threshold: float, | |
iou_threshold: float | |
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: | |
# Validate iou_threshold before using it | |
if iou_threshold is None or not isinstance(iou_threshold, float): | |
iou_threshold = 0.3 # Default value, adjust as necessary | |
yolo_v8n_annotated_image = detect_and_annotate( | |
YOLO_V8N_MODEL, input_image, yolo_v8_confidence_threshold, iou_threshold) | |
yolo_v8s_annotated_image = detect_and_annotate( | |
YOLO_V8S_MODEL, input_image, yolo_v9_confidence_threshold, iou_threshold) | |
yolo_8m_annotated_image = detect_and_annotate( | |
YOLO_V8M_MODEL, input_image, yolo_v10_confidence_threshold, iou_threshold) | |
return ( | |
yolo_v8n_annotated_image, | |
yolo_v8s_annotated_image, | |
yolo_8m_annotated_image | |
) | |
yolo_v8N_confidence_threshold_component = gr.Slider( | |
minimum=0, | |
maximum=1.0, | |
value=0.3, | |
step=0.01, | |
label="YOLOv8N Confidence Threshold", | |
info=( | |
"The confidence threshold for the YOLO model. Lower the threshold to " | |
"reduce false negatives, enhancing the model's sensitivity to detect " | |
"sought-after objects. Conversely, increase the threshold to minimize false " | |
"positives, preventing the model from identifying objects it shouldn't." | |
)) | |
yolo_v8S_confidence_threshold_component = gr.Slider( | |
minimum=0, | |
maximum=1.0, | |
value=0.3, | |
step=0.01, | |
label="YOLOv8S Confidence Threshold", | |
info=( | |
"The confidence threshold for the YOLO model. Lower the threshold to " | |
"reduce false negatives, enhancing the model's sensitivity to detect " | |
"sought-after objects. Conversely, increase the threshold to minimize false " | |
"positives, preventing the model from identifying objects it shouldn't." | |
)) | |
yolo_v8M_confidence_threshold_component = gr.Slider( | |
minimum=0, | |
maximum=1.0, | |
value=0.3, | |
step=0.01, | |
label="YOLOv8M Confidence Threshold", | |
info=( | |
"The confidence threshold for the YOLO model. Lower the threshold to " | |
"reduce false negatives, enhancing the model's sensitivity to detect " | |
"sought-after objects. Conversely, increase the threshold to minimize false " | |
"positives, preventing the model from identifying objects it shouldn't." | |
)) | |
iou_threshold_component = gr.Slider( | |
minimum=0, | |
maximum=1.0, | |
value=0.5, | |
step=0.01, | |
label="IoU Threshold", | |
info=( | |
"The Intersection over Union (IoU) threshold for non-maximum suppression. " | |
"Decrease the value to lessen the occurrence of overlapping bounding boxes, " | |
"making the detection process stricter. On the other hand, increase the value " | |
"to allow more overlapping bounding boxes, accommodating a broader range of " | |
"detections." | |
)) | |
with gr.Blocks() as demo: | |
gr.Markdown(MARKDOWN) | |
with gr.Accordion("Configuration", open=False): | |
with gr.Row(): | |
yolo_v8N_confidence_threshold_component.render() | |
yolo_v8S_confidence_threshold_component.render() | |
yolo_v8M_confidence_threshold_component.render() | |
iou_threshold_component.render() | |
with gr.Row(): | |
input_image_component = gr.Image( | |
type='pil', | |
label='Input' | |
) | |
yolo_v8n_output_image_component = gr.Image( | |
type='pil', | |
label='YOLOv8N' | |
) | |
with gr.Row(): | |
yolo_v8s_output_image_component = gr.Image( | |
type='pil', | |
label='YOLOv8S' | |
) | |
yolo_v8m_output_image_component = gr.Image( | |
type='pil', | |
label='YOLOv8M' | |
) | |
submit_button_component = gr.Button( | |
value='Submit', | |
scale=1, | |
variant='primary' | |
) | |
gr.Examples( | |
fn=process_image, | |
examples=IMAGE_EXAMPLES, | |
inputs=[ | |
input_image_component, | |
yolo_v8N_confidence_threshold_component, | |
yolo_v8S_confidence_threshold_component, | |
yolo_v8M_confidence_threshold_component, | |
iou_threshold_component | |
], | |
outputs=[ | |
yolo_v8n_output_image_component, | |
yolo_v8s_output_image_component, | |
yolo_v8m_output_image_component | |
] | |
) | |
submit_button_component.click( | |
fn=process_image, | |
inputs=[ | |
input_image_component, | |
yolo_v8N_confidence_threshold_component, | |
yolo_v8S_confidence_threshold_component, | |
yolo_v8M_confidence_threshold_component, | |
iou_threshold_component | |
], | |
outputs=[ | |
yolo_v8n_output_image_component, | |
yolo_v8s_output_image_component, | |
yolo_v8m_output_image_component | |
] | |
) | |
demo.launch(debug=False, show_error=True, max_threads=1) |