HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add Streamlit Inference Python model arg (#14563)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
rulosant 2024-07-21 15:54:52 -03:00 • committed by GitHub
parent 291883a23f
commit f4af1bccc6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 23 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
ultralytics/solutions/streamlit_inference.py
View file

@ -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