yolov8 目标检测快速streamlit可视化界面

发布于:2024-07-04 ⋅ 阅读:(170) ⋅ 点赞:(0)

参考:
https://github.com/ultralytics/ultralytics/blob/2330caa50a8a8e0bb61408df8dca0721fb350dbe/ultralytics/solutions/streamlit_inference.py

版本:

ultralytics                    8.2.27 
# Ultralytics YOLO 🚀, AGPL-3.0 license

import cv2
import streamlit as st
import torch

from 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 configuration
    st.set_page_config(page_title="Ultralytics Streamlit App", layout="wide", initial_sidebar_state="auto")

    # Append the custom HTML
    st.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 menu
    st.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 model
        videocapture = cv2.VideoCapture(0)  # Capture the webcam

        if not videocapture.isOpened():
            st.error("Could not open webcam.")

        stop_button = st.button("Stop")  # Button to stop the inference
        # execute code until webcam closed
        while 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 predictions
            results = model(frame, conf=float(conf_thres), iou=float(nms_thres))
            annotated_frame = results[0].plot()  # Add annotations on frame

            # display frame
            org_frame.image(frame, channels="BGR")
            ann_frame.image(annotated_frame, channels="BGR")

            if stop_button:
                videocapture.release()  # Release the capture
                torch.cuda.empty_cache()  # Clear CUDA memory
                st.stop()  # Stop streamlit app

        # Release the capture
        videocapture.release()

    # Clear CUDA memory
    torch.cuda.empty_cache()

    # Destroy window
    cv2.destroyAllWindows()


# Main function call
if __name__ == "__main__":
    try:
        inference()
    except SystemExit:
        pass

运行:
streamlit run .\YOLO\streamlit_yolo.py

在这里插入图片描述


网站公告

今日签到

点亮在社区的每一天
去签到