参考:
https://github.com/ultralytics/ultralytics/blob/2330caa50a8a8e0bb61408df8dca0721fb350dbe/ultralytics/solutions/streamlit_inference.py
版本:
ultralytics 8.2.27
# Ultralytics YOLO 🚀, AGPL-3.0 licenseimport cv2
import streamlit as st
import torchfrom ultralytics import YOLO# Hide main menu style
menu_style_cfg = """<style>MainMenu {visibility: hidden;}</style>"""# Main title of streamlit application
main_title_cfg = """<div><h1 style="color:#FF64DA; text-align:center; font-size:40px; font-family: 'Archivo', sans-serif; margin-top:-50px;">Ultralytics YOLOv8 Streamlit Application</h1></div>"""# Subtitle of streamlit application
sub_title_cfg = """<div><h4 style="color:#042AFF; text-align:center; font-family: 'Archivo', sans-serif; margin-top:-15px;">Experience real-time object detection on your webcam with the power of Ultralytics YOLOv8! 🚀</h4></div>"""def inference():# Set html page configurationst.set_page_config(page_title="Ultralytics Streamlit App", layout="wide", initial_sidebar_state="auto")# Append the custom HTMLst.markdown(menu_style_cfg, unsafe_allow_html=True)st.markdown(main_title_cfg, unsafe_allow_html=True)st.markdown(sub_title_cfg, unsafe_allow_html=True)# Add elements to vertical setting menust.sidebar.title("USER Configuration")yolov8_model = st.sidebar.radio("Model", ("yolov8n", "yolov8s", "yolov8m", "yolov8l"))conf_thres = st.sidebar.text_input("Confidence Threshold", "0.25")nms_thres = st.sidebar.text_input("NMS Threshold", "0.45")col1, col2 = st.columns(2)org_frame = col1.empty()ann_frame = col2.empty()if st.sidebar.button("Start"):model = YOLO(yolov8_model + ".pt") # Load the yolov8 modelvideocapture = cv2.VideoCapture(0) # Capture the webcamif not videocapture.isOpened():st.error("Could not open webcam.")stop_button = st.button("Stop") # Button to stop the inference# execute code until webcam closedwhile videocapture.isOpened():success, frame = videocapture.read()if not success:st.warning("Failed to read frame from webcam. Please make sure the webcam is connected properly.")break# Store model predictionsresults = model(frame, conf=float(conf_thres), iou=float(nms_thres))annotated_frame = results[0].plot() # Add annotations on frame# display frameorg_frame.image(frame, channels="BGR")ann_frame.image(annotated_frame, channels="BGR")if stop_button:videocapture.release() # Release the capturetorch.cuda.empty_cache() # Clear CUDA memoryst.stop() # Stop streamlit app# Release the capturevideocapture.release()# Clear CUDA memorytorch.cuda.empty_cache()# Destroy windowcv2.destroyAllWindows()# Main function call
if __name__ == "__main__":try:inference()except SystemExit:pass
运行:
streamlit run .\YOLO\streamlit_yolo.py