update

app.py

@@ -1,4 +1,77 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+-------------------------------------------------
+   @File Name:     app.py
+   @Author:        Luyao.zhang
+   @Date:          2023/5/15
+   @Description:
+-------------------------------------------------
+"""
+from pathlib import Path
+from PIL import Image
 import streamlit as st
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+import config
+from utils import load_model, infer_uploaded_image, infer_uploaded_video, infer_uploaded_webcam
+
+# 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")
+
+# sidebar
+st.sidebar.header("DL Model Config")
+
+# model options
+task_type = st.sidebar.selectbox(
+    "Select Task",
+    ["Detection"]
+)
+
+model_type = None
+if task_type == "Detection":
+    model_type = st.sidebar.selectbox(
+        "Select Model",
+        config.DETECTION_MODEL_LIST
+    )
+else:
+    st.error("Currently only 'Detection' function is implemented")
+
+confidence = float(st.sidebar.slider(
+    "Select Model Confidence", 30, 100, 50)) / 100
+
+model_path = ""
+if model_type:
+    model_path = Path(config.DETECTION_MODEL_DIR, str(model_type))
+else:
+    st.error("Please Select Model in Sidebar")
+
+# load pretrained DL model
+try:
+    model = load_model(model_path)
+except Exception as e:
+    st.error(f"Unable to load model. Please check the specified path: {model_path}")
+
+# image/video options
+st.sidebar.header("Image/Video Config")
+source_selectbox = st.sidebar.selectbox(
+    "Select Source",
+    config.SOURCES_LIST
+)
+
+source_img = None
+if source_selectbox == config.SOURCES_LIST[0]: # Image
+    infer_uploaded_image(confidence, model)
+elif source_selectbox == config.SOURCES_LIST[1]: # Video
+    infer_uploaded_video(confidence, model)
+elif source_selectbox == config.SOURCES_LIST[2]: # Webcam
+    infer_uploaded_webcam(confidence, model)
+else:
+    st.error("Currently only 'Image' and 'Video' source are implemented")

config.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+-------------------------------------------------
+   @File Name:     config.py
+   @Author:        Luyao.zhang
+   @Date:          2023/5/16
+   @Description: configuration file
+-------------------------------------------------
+"""
+from pathlib import Path
+import sys
+
+# Get the absolute path of the current file
+file_path = Path(__file__).resolve()
+
+# Get the parent directory of the current file
+root_path = file_path.parent
+
+# Add the root path to the sys.path list if it is not already there
+if root_path not in sys.path:
+    sys.path.append(str(root_path))
+
+# Get the relative path of the root directory with respect to the current working directory
+ROOT = root_path.relative_to(Path.cwd())
+
+
+# Source
+SOURCES_LIST = ["Image", "Video", "Webcam"]
+
+
+# DL model config
+DETECTION_MODEL_DIR = ROOT / 'models' / 'detection'
+YOLOv8n = DETECTION_MODEL_DIR / "yolov8n.pt"
+YOLOv8s = DETECTION_MODEL_DIR / "yolov8s.pt"
+YOLOv8m = DETECTION_MODEL_DIR / "yolov8m.pt"
+YOLOv8l = DETECTION_MODEL_DIR / "yolov8l.pt"
+YOLOv8x = DETECTION_MODEL_DIR / "yolov8x.pt"
+
+DETECTION_MODEL_LIST = [
+    "yolov8n.pt",
+    "yolov8s.pt",
+    "yolov8m.pt",
+    "yolov8l.pt",
+    "yolov8x.pt"]

utils.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+-------------------------------------------------
+   @File Name:     utils.py
+   @Author:        Luyao.zhang
+   @Date:          2023/5/16
+   @Description:
+-------------------------------------------------
+"""
+from ultralytics import YOLO
+import streamlit as st
+import cv2
+from PIL import Image
+import tempfile
+
+
+def _display_detected_frames(conf, model, st_frame, image):
+    """
+    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
+    """
+    # 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
+                   )
+
+
+@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
+            )
+
+    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]
+
+                with col2:
+                    st.image(res_plotted,
+                             caption="Detected Image",
+                             use_column_width=True)
+                    try:
+                        with st.expander("Detection Results"):
+                            for box in boxes:
+                                st.write(box.xywh)
+                    except Exception as ex:
+                        st.write("No image is uploaded yet!")
+                        st.write(ex)
+
+
+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()
+                    while (vid_cap.isOpened()):
+                        success, image = vid_cap.read()
+                        if success:
+                            _display_detected_frames(conf,
+                                                     model,
+                                                     st_frame,
+                                                     image
+                                                     )
+                        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()
+        while not flag:
+            success, image = vid_cap.read()
+            if success:
+                _display_detected_frames(
+                    conf,
+                    model,
+                    st_frame,
+                    image
+                )
+            else:
+                vid_cap.release()
+                break
+    except Exception as e:
+        st.error(f"Error loading video: {str(e)}")