Update app.py

app.py

@@ -1,3 +1,4 @@
+import cv2
 from ultralytics import YOLO
 import streamlit as st
 
@@ -24,6 +25,114 @@ def yolo_tracking(frame):
 
     return rgb_frame
 
+def _display_detected_frames(conf, model, st_count, 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)
+    
+    inText = 'Vehicle In'
+    outText = 'Vehicle Out'
+    if config.OBJECT_COUNTER1 != None:
+        for _, (key, value) in enumerate(config.OBJECT_COUNTER1.items()):
+            inText += ' - ' + str(key) + ": " +str(value)
+    if config.OBJECT_COUNTER != None:
+        for _, (key, value) in enumerate(config.OBJECT_COUNTER.items()):
+            outText += ' - ' + str(key) + ": " +str(value)
+    
+    # Plot the detected objects on the video frame
+    st_count.write(inText + '\n\n' + outText)
+    res_plotted = res[0].plot()
+    st_frame.image(res_plotted,
+                   caption='Detected Video',
+                   channels="BGR",
+                   use_column_width=True
+                   )
+
+
+
+
+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:
+                    config.OBJECT_COUNTER1 = None
+                    config.OBJECT_COUNTER = None
+                    tfile = tempfile.NamedTemporaryFile()
+                    tfile.write(source_video.read())
+                    vid_cap = cv2.VideoCapture(
+                        tfile.name)
+                    st_count = st.empty()
+                    st_frame = st.empty()
+                    while (vid_cap.isOpened()):
+                        success, image = vid_cap.read()
+                        if success:
+                            _display_detected_frames(conf,
+                                                     model,
+                                                     st_count,
+                                                     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_count = st.empty()
+        st_frame = st.empty()
+        while not flag:
+            success, image = vid_cap.read()
+            if success:
+                _display_detected_frames(
+                    conf,
+                    model,
+                    st_count,
+                    st_frame,
+                    image
+                )
+            else:
+                vid_cap.release()
+                break
+    except Exception as e:
+        st.error(f"Error loading video: {str(e)}")
+
 # Setting page layout
 st.set_page_config(
     page_title="Object Detection using YOLOv8",
@@ -35,3 +144,34 @@ st.set_page_config(
 # Main page heading
 st.title("Object Detection And Tracking using YOLOv8")
 
+# sidebar
+st.sidebar.header("DL Model Config")
+
+# model options
+task_type = st.sidebar.selectbox(
+    "Select Task",
+    ["Detection"]
+)
+
+confidence = float(st.sidebar.slider(
+    "Select Model Confidence", 30, 100, 50)) / 100
+
+# image/video options
+st.sidebar.header("Image/Video Config")
+source_selectbox = st.sidebar.selectbox(
+    "Select Source",
+    config.SOURCES_LIST
+)
+
+SOURCES_LIST = ["Video", "Webcam"]
+
+source_img = None
+#if source_selectbox == config.SOURCES_LIST[0]: # Image
+#    infer_uploaded_image(confidence, model)
+#elif source_selectbox == config.SOURCES_LIST[1]: # Video
+if source_selectbox == SOURCES_LIST[0]: # Video
+    infer_uploaded_video(confidence, model)
+elif source_selectbox == SOURCES_LIST[1]: # Webcam
+    infer_uploaded_webcam(confidence, model)
+else:
+    st.error("Currently only 'Image' and 'Video' source are implemented")