Update main.py

main.py

@@ -1,13 +1,78 @@
 import streamlit as st
+from yolo_predictions import Yolo_Pred
+from PIL import Image
+import numpy as np
 
-st.set_page_config(page_title="YOLO Object Detection Application", layout="centered")
+# Set page config at the start
+st.set_page_config(
+    page_title='YOLO V5 OBJECT DETECTION APPLICATION',
+    layout='centered',
+    page_icon='https://cdn.analyticsvidhya.com/wp-content/uploads/2018/12/yologo_2-850x451.png'
+)
 
-st.sidebar.title("Navigation")
-page = st.sidebar.radio("Go to", ["Home", "Image Object Detection"])
+def home():
+    st.title('YOLO V5 OBJECT DETECTION APPLICATION')
+    st.caption("Our project demonstrates ***object detection*** application using the ***YOLO***(You Only Look Once) model version 5 and the ***Streamlit*** framework, combined with computer vision techniques using ***OpenCV***. Object detection plays a crucial role in many computer vision applications, allowing us to identify and locate multiple objects within an image or a video stream. The YOLO model version 5 is a state-of-the-art deep learning algorithm known for its high accuracy and real-time performance. With Streamlit, we create an interactive and user-friendly web application that enables users to upload images or provide a video stream for real-time object detection. The Streamlit framework simplifies the development process by providing intuitive widgets and UI components. To process the uploaded images or video stream, we leverage OpenCV, a powerful library for computer vision tasks. OpenCV allows us to perform various image processing operations, such as resizing, color conversions, and pre-processing required for object detection. We integrate the YOLO model version 5 into our application to detect objects accurately. Our application provides a seamless user experience, allowing users to visualize the detected objects, their bounding boxes, and class labels. Additionally, we can leverage Streamlit's capabilities to incorporate additional functionalities, such as filtering detected objects based on confidence scores or displaying statistical insights about the detected objects. By combining the YOLO model version 5, Streamlit, and OpenCV, our project empowers users to perform efficient and accurate object detection tasks in real-time. This application finds applications in various domains, including surveillance, autonomous vehicles, and object recognition systems.")
+    st.markdown("""
+    ### This app detects objects from images and videos 
+    - Our model detects 20 objects from images and videos:
+        - person        
+        - car
+        - chair
+        - bottle
+        - pottedplant
+        - bird
+        - dog
+        - sofa
+        - bicycle
+        - horse
+        - boat
+        - motorbike
+        - cat
+        - tvmonitor
+        - cow
+        - sheep
+        - aeroplane
+        - train
+        - diningtable
+        - bus
 
-if page == "Home":
-    import home
-    home.show()
-elif page == "Image Object Detection":
-    import pages.1_yolo_image as yolo_image
-    yolo_image.show()
+    - [Click here for image object detection](?page=yolo_image)
+    """)
+
+def yolo_image():
+    st.title('YOLO OBJECT DETECTION FROM IMAGES')
+    st.write('Please upload images to get detections')
+
+    with st.spinner('Model loading ....'):
+        yolo = Yolo_Pred(onnx_model='./models/best.onnx', data_yaml='./models/data.yaml')
+        st.balloons()
+
+    # Upload image
+    image_file = st.file_uploader("Choose an image", type=["jpg", "png"])
+    if image_file is not None:
+        prediction = False
+        # Read and display the image
+        st.info('The image uploaded successfully')
+        image = Image.open(image_file)
+        st.image(image)
+        button = st.button('GET PREDICTIONS')
+        if button:
+            image_array = np.array(image)
+            img_pred = yolo.predictions(image_array)
+            pred_img_obj = Image.fromarray(img_pred)
+            prediction = True
+        if prediction:
+            st.image(pred_img_obj)
+
+def main():
+    query_params = st.experimental_get_query_params()
+    page = query_params.get("page", ["home"])[0]
+
+    if page == "home":
+        home()
+    elif page == "yolo_image":
+        yolo_image()
+
+if __name__ == "__main__":
+    main()