From f4af1bccc66bfe5e80643469a6dc82fefea62015 Mon Sep 17 00:00:00 2001 From: rulosant Date: Sun, 21 Jul 2024 15:54:52 -0300 Subject: [PATCH] Add Streamlit Inference Python `model` arg (#14563) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- docs/en/guides/streamlit-live-inference.md | 20 ++++++++++++++++++-- ultralytics/solutions/streamlit_inference.py | 7 +++++-- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/docs/en/guides/streamlit-live-inference.md b/docs/en/guides/streamlit-live-inference.md index b88ae9f5..7deb8c70 100644 --- a/docs/en/guides/streamlit-live-inference.md +++ b/docs/en/guides/streamlit-live-inference.md @@ -31,7 +31,7 @@ Streamlit makes it simple to build and deploy interactive web applications. Comb === "Python" - ```Python + ```python from ultralytics import solutions solutions.inference() @@ -47,6 +47,21 @@ Streamlit makes it simple to build and deploy interactive web applications. Comb This will launch the Streamlit application in your default web browser. You will see the main title, subtitle, and the sidebar with configuration options. Select your desired YOLOv8 model, set the confidence and NMS thresholds, and click the "Start" button to begin the real-time object detection. +You can optionally supply a specific model in Python: + +!!! Example "Streamlit Application with a custom model" + + === "Python" + + ```python + from ultralytics import solutions + + # Pass a model as an argument + solutions.inference(model="path/to/model.pt") + + ### Make sure to run the file using command `streamlit run ` + ``` + ## Conclusion By following this guide, you have successfully created a real-time object detection application using Streamlit and Ultralytics YOLOv8. This application allows you to experience the power of YOLOv8 in detecting objects through your webcam, with a user-friendly interface and the ability to stop the video stream at any time. @@ -82,8 +97,9 @@ Then, you can create a basic Streamlit application to run live inference: === "Python" - ```Python + ```python from ultralytics import solutions + solutions.inference() ### Make sure to run the file using command `streamlit run ` diff --git a/ultralytics/solutions/streamlit_inference.py b/ultralytics/solutions/streamlit_inference.py index fbb764d1..99916552 100644 --- a/ultralytics/solutions/streamlit_inference.py +++ b/ultralytics/solutions/streamlit_inference.py @@ -10,7 +10,7 @@ from ultralytics.utils.checks import check_requirements from ultralytics.utils.downloads import GITHUB_ASSETS_STEMS -def inference(): +def inference(model=None): """Runs real-time object detection on video input using Ultralytics YOLOv8 in a Streamlit application.""" check_requirements("streamlit>=1.29.0") # scope imports for faster ultralytics package load speeds import streamlit as st @@ -67,7 +67,10 @@ def inference(): vid_file_name = 0 # Add dropdown menu for model selection - available_models = (x.replace("yolo", "YOLO") for x in GITHUB_ASSETS_STEMS if x.startswith("yolov8")) + available_models = [x.replace("yolo", "YOLO") for x in GITHUB_ASSETS_STEMS if x.startswith("yolov8")] + if model: + available_models.insert(0, model) + selected_model = st.sidebar.selectbox("Model", available_models) with st.spinner("Model is downloading..."): model = YOLO(f"{selected_model.lower()}.pt") # Load the YOLO model