vending_app.py

from ultralytics import YOLO
import streamlit as st
import cv2
from PIL import Image
import tempfile
import time

def _display_detected_frames(conf, model, st_frame, image, box_b, box_r):
    """
    Display the detected objects on a video frame using the YOLOv8 model.
    :param conf (float): Confidence threshold for object detection.
    :param model (YOLOv8): An instance of the `YOLOv8` class containing the YOLOv8 model.
    :param st_frame (Streamlit object): A Streamlit object to display the detected video.
    :param image (numpy array): A numpy array representing the video frame.
    :return: None
    """
    t1 = time.time()
    # Resize the image to a standard size
    image = cv2.resize(image, (720, int(720 * (9 / 16))))

    # Predict the objects in the image using YOLOv8 model
    res = model.predict(image, conf=conf)

    # Plot the detected objects on the video frame
    res_plotted = res[0].plot()
    st_frame.image(res_plotted,
                   caption='Detected Video',
                   channels="BGR",
                   use_column_width=True
                   )
    num_brown = 0
    num_red = 0
    for pred in res[0].boxes.cls:
        if pred == 0:
            num_brown += 1
        elif pred == 1:
            num_red += 1
    t2 = time.time()
    
    box_b.subheader(f"Number of brown: {num_brown}")
    box_r.subheader(f"Number of red: {num_red}")
    box_r.text(f"Time: {t2-t1:.4f}s") 
    


@st.cache_resource
def load_model(model_path):
    """
    Loads a YOLO object detection model from the specified model_path.

    Parameters:
        model_path (str): The path to the YOLO model file.

    Returns:
        A YOLO object detection model.
    """
    model = YOLO(model_path)
    return model


def infer_uploaded_image(conf, model):
    """
    Execute inference for uploaded image
    :param conf: Confidence of YOLOv8 model
    :param model: An instance of the `YOLOv8` class containing the YOLOv8 model.
    :return: None
    """
    source_img = st.sidebar.file_uploader(
        label="Choose an image...",
        type=("jpg", "jpeg", "png", 'bmp', 'webp')
    )

    col1, col2 = st.columns(2)

    with col1:
        if source_img:
            uploaded_image = Image.open(source_img)
            # adding the uploaded image to the page with caption
            st.image(
                image=source_img,
                caption="Uploaded Image",
                use_column_width=True
            )
    t1 = time.time()
    if source_img:
        if st.button("Execution"):
            with st.spinner("Running..."):
                res = model.predict(uploaded_image,
                                    conf=conf)
                boxes = res[0].boxes
                res_plotted = res[0].plot()[:, :, ::-1]

                num_brown = 0
                num_red = 0
                for pred in res[0].boxes.cls:
                    if pred == 0:
                        num_brown += 1
                    elif pred == 1:
                        num_red += 1

                with col2:
                    t2 = time.time()
                    st.image(res_plotted,
                             caption="Detected Image",
                             use_column_width=True)
                    st.subheader(f"Number of brown: {num_brown}")
                    st.subheader(f"Number of red: {num_red}")
                    st.text(f"Time: {t2-t1:.4f}s")


def infer_uploaded_video(conf, model):
    """
    Execute inference for uploaded video
    :param conf: Confidence of YOLOv8 model
    :param model: An instance of the `YOLOv8` class containing the YOLOv8 model.
    :return: None
    """
    source_video = st.sidebar.file_uploader(
        label="Choose a video..."
    )

    if source_video:
        st.video(source_video)

    if source_video:
        if st.button("Execution"):
            with st.spinner("Running..."):
                try:
                    tfile = tempfile.NamedTemporaryFile()
                    tfile.write(source_video.read())
                    vid_cap = cv2.VideoCapture(
                        tfile.name)
                    st_frame = st.empty()
                    b1 = st.empty()
                    b2 = st.empty()
                    while (vid_cap.isOpened()):
                        success, image = vid_cap.read()
                        if success:
                            _display_detected_frames(conf,
                                                     model,
                                                     st_frame,
                                                     image,
                                                     b1,
                                                     b2
                                                     )
                        else:
                            vid_cap.release()
                            break
                except Exception as e:
                    st.error(f"Error loading video: {e}")


def infer_uploaded_webcam(conf, model):
    """
    Execute inference for webcam.
    :param conf: Confidence of YOLOv8 model
    :param model: An instance of the `YOLOv8` class containing the YOLOv8 model.
    :return: None
    """
    try:
        flag = st.button(
            label="Stop running"
        )
        vid_cap = cv2.VideoCapture(0)  # local camera
        st_frame = st.empty()
        b1 = st.empty()
        b2 = st.empty()
        while not flag:
            success, image = vid_cap.read()
            if success:
                _display_detected_frames(
                    conf,
                    model,
                    st_frame,
                    image,
                    b1,
                    b2
                )
            else:
                vid_cap.release()
                break
    except Exception as e:
        st.error(f"Error loading video: {str(e)}")

def main():
    #!/usr/bin/env python
    # -*- coding: utf-8 -*-
    from pathlib import Path
    from PIL import Image
    import streamlit as st


    # setting page layout
    st.set_page_config(
        page_title="Interactive Interface for YOLOv8",
        page_icon="🤖",
        layout="wide",
        initial_sidebar_state="expanded"
        )

    # main page heading
    st.title("Interactive Interface for YOLOv8")

    # load pretrained DL model
    model = load_model('/code/Vending-YOLOv8n.onnx')

    # image/video options
    st.sidebar.header("Image/Video Config")
    source_selectbox = st.sidebar.selectbox(
        "Select Source",
        ["Image", "Video", "Webcam"]
    )

    source_img = None
    confidence = 0.5
    if source_selectbox == "Image": # Image
        infer_uploaded_image(confidence, model)
    elif source_selectbox == "Video": # Video
        infer_uploaded_video(confidence, model)
    elif source_selectbox == "Webcam": # Webcam
        infer_uploaded_webcam(confidence, model)
    else:
        st.error("Currently only 'Image' and 'Video' source are implemented")

if __name__ == "__main__":
    main()