first commit

.gitignore

@@ -0,0 +1,148 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Local folder
+local_folder
+project_demo
+project_demo/
+logs/
+local_folder/
+/demo.py
+demo.py
+runs/
+
+# Large folders
+weights/
+videos/
+images/
+
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+.vscode
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+venv_/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

app.py

@@ -0,0 +1,134 @@
+# Python In-built packages
+from pathlib import Path
+import PIL
+import time
+
+# External packages
+import streamlit as st
+import cv2
+
+# Local Modules
+import settings
+import helper
+
+# Setting page layout
+st.set_page_config(
+    page_title="Object Detection using YOLOv9",
+    page_icon="🤖",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Main page heading
+st.title("Object Detection using YOLOv9")
+
+# Sidebar
+st.sidebar.header("ML Model Config")
+
+# Model Options
+model_type = st.sidebar.radio(
+    "Select Task", ['Detection', 'Segmentation'])
+confidence = float(st.sidebar.slider(
+    "Select Model Confidence", 25, 100, 40)) / 100
+
+# Selecting Detection Or Segmentation
+if model_type == 'Detection':
+    model_path = Path(settings.DETECTION_MODEL)
+elif model_type == 'Segmentation':
+    model_path = Path(settings.SEGMENTATION_MODEL)
+
+# Load Pre-trained ML Model
+try:
+    model = helper.load_model(model_path)
+except Exception as ex:
+    st.error(f"Unable to load model. Check the specified path: {model_path}")
+    st.error(ex)
+
+st.sidebar.header("Image/Video Config")
+source_radio = st.sidebar.radio(
+    "Select Source", settings.SOURCES_LIST)
+source_img = None
+
+# If image is selected
+if source_radio == settings.IMAGE:
+    source_img = st.sidebar.file_uploader(
+        "Choose an image...", type=("jpg", "jpeg", "png", 'bmp', 'webp'))
+    col1, col2 = st.columns(2)
+    with col1:
+        try:
+            if source_img is None:
+                default_image_path = str(settings.DEFAULT_IMAGE)
+                default_image = PIL.Image.open(default_image_path)
+                st.image(default_image_path, caption="Default Image",
+                         use_column_width=True)
+            else:
+                uploaded_image = PIL.Image.open(source_img)
+                st.image(source_img, caption="Uploaded Image",
+                         use_column_width=True)
+        except Exception as ex:
+            st.error("Error occurred while opening the image.")
+            st.error(ex)
+    with col2:
+        if source_img is None:
+            default_detected_image_path = str(settings.DEFAULT_DETECT_IMAGE)
+            default_detected_image = PIL.Image.open(
+                default_detected_image_path)
+            st.image(default_detected_image_path, caption='Detected Image',
+                     use_column_width=True)
+        else:
+            if st.sidebar.button('Detect Objects'):
+                res = model.predict(uploaded_image,
+                                    conf=confidence
+                                    )
+                boxes = res[0].boxes
+                res_plotted = res[0].plot()[:, :, ::-1]
+                st.image(res_plotted, caption='Detected Image',
+                         use_column_width=True)
+                try:
+                    with st.expander("Detection Results"):
+                        for box in boxes:
+                            st.write(box.data)
+                except Exception as ex:
+                    # st.write(ex)
+                    st.write("No image is uploaded yet!")
+
+elif source_radio == settings.VIDEO:
+    helper.play_stored_video(confidence, model)
+
+elif source_radio == settings.YOUTUBE:
+    helper.play_youtube_video(confidence, model)
+
+elif source_radio == settings.WEBCAM:  # Adding webcam option
+    st.sidebar.markdown("**Webcam**")
+    webcam = cv2.VideoCapture(1)
+    if not webcam.isOpened():
+        st.error("Unable to open webcam. Please make sure it's connected.")
+    else:
+        webcam_status = st.sidebar.empty()
+        start_stop_button = st.sidebar.button("Start Webcam")
+
+        if start_stop_button:
+            webcam_started = True
+            webcam_status.text("Webcam Started!")
+
+            while webcam_started:
+                ret, frame = webcam.read()
+                if not ret:
+                    st.error("Failed to capture frame from webcam")
+                    break
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+                # Perform object detection on the frame
+                res = model.predict(frame, conf=confidence)
+                boxes = res[0].boxes
+                frame = res[0].plot()[:, :, ::-1]
+
+                # Display the frame
+                st.image(frame, channels="RGB", use_column_width=True)
+
+                stop_button = st.sidebar.button("Stop Webcam" + str(hash(time.time())))  # Unique key
+                if stop_button:
+                    webcam_started = False
+                    webcam_status.text("Webcam Stopped!")
+
+            webcam.release()

helper.py

@@ -0,0 +1,74 @@
+from ultralytics import YOLO
+import streamlit as st
+import cv2
+import pafy
+
+import settings
+
+def load_model(model_path):
+    model = YOLO(model_path)
+    return model
+
+def _display_detected_frames(conf, model, st_frame, image):
+    # Resize the image to a standard size
+    image = cv2.resize(image, (720, int(720*(9/16))))
+    
+    # Predict the objects in the image using the YOLOv9 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
+                   )
+
+def play_youtube_video(conf, model):
+    source_youtube = st.sidebar.text_input("YouTube Video url")
+    if st.sidebar.button('Detect Objects'):
+        try:
+            video = pafy.new(source_youtube)
+            best = video.getbest(preftype="mp4")
+            vid_cap = cv2.VideoCapture(best.url)
+            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.sidebar.error("Error loading video: " + str(e))
+
+def play_stored_video(conf, model):
+    source_vid = st.sidebar.selectbox(
+        "Choose a video...", settings.VIDEOS_DICT.keys())
+    with open(settings.VIDEOS_DICT.get(source_vid), 'rb') as video_file:
+        video_bytes = video_file.read()
+    
+    if video_bytes:
+        st.video(video_bytes)
+    if st.sidebar.button('Detect Video Objects'):
+        try:
+            vid_cap = cv2.VideoCapture(
+                str(settings.VIDEOS_DICT.get(source_vid)))
+            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.sidebar.error("Error loading video: " + str(e))

packages.txt

@@ -0,0 +1,3 @@
+freeglut3-dev
+libgtk2.0-dev
+libgl1-mesa-glx

requirements.txt

@@ -0,0 +1,78 @@
+altair==5.2.0
+attrs==23.2.0
+blinker==1.7.0
+cachetools==5.3.3
+certifi==2024.2.2
+charset-normalizer==3.3.2
+click==8.1.7
+contourpy==1.2.0
+cycler==0.12.1
+filelock==3.13.1
+fonttools==4.50.0
+fsspec==2024.3.1
+gitdb==4.0.11
+GitPython==3.1.42
+idna==3.6
+Jinja2==3.1.3
+jsonschema==4.21.1
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.8.3
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.19.3
+nvidia-nvjitlink-cu12==12.4.99
+nvidia-nvtx-cu12==12.1.105
+opencv-python==4.9.0.80
+packaging==23.2
+pandas==2.2.1
+pillow==10.2.0
+protobuf==4.25.3
+psutil==5.9.8
+py-cpuinfo==9.0.0
+pyarrow==15.0.2
+pydeck==0.8.1b0
+Pygments==2.17.2
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.34.0
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.0
+scipy==1.12.0
+seaborn==0.13.2
+six==1.16.0
+smmap==5.0.1
+streamlit==1.32.2
+sympy==1.12
+tenacity==8.2.3
+thop==0.1.1.post2209072238
+toml==0.10.2
+toolz==0.12.1
+torch==2.2.1
+torchvision==0.17.1
+tornado==6.4
+tqdm==4.66.2
+triton==2.2.0
+typing_extensions==4.10.0
+tzdata==2024.1
+ultralytics==8.1.30
+urllib3==2.2.1
+watchdog==4.0.0
+pafy
+youtube-dl

settings.py

@@ -0,0 +1,47 @@
+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())
+
+# Sources
+IMAGE = 'Image'
+VIDEO = 'Video'
+YOUTUBE = 'YouTube'
+WEBCAM = 'Webcam'
+SOURCES_LIST = [IMAGE, VIDEO, YOUTUBE, WEBCAM]
+
+# Images config
+IMAGES_DIR = ROOT / 'images'
+DEFAULT_IMAGE = IMAGES_DIR / 'office_4.jpg'
+DEFAULT_DETECT_IMAGE = IMAGES_DIR / 'office_4_detected.jpg'
+
+# Videos config
+VIDEO_DIR = ROOT / 'videos'
+VIDEO_1_PATH = VIDEO_DIR / 'video_1.mp4'
+VIDEO_2_PATH = VIDEO_DIR / 'video_2.mp4'
+VIDEO_3_PATH = VIDEO_DIR / 'video_3.mp4'
+VIDEO_4_PATH = VIDEO_DIR / 'video_4.mp4'
+VIDEO_5_PATH = VIDEO_DIR / 'video_5.mp4'
+VIDEOS_DICT = {
+    'video_1': VIDEO_1_PATH,
+    'video_2': VIDEO_2_PATH,
+    'video_3': VIDEO_3_PATH,
+    'video_4': VIDEO_4_PATH,
+    'video_5': VIDEO_5_PATH,
+}
+
+# ML Model config
+MODEL_DIR = ROOT / 'weights'
+DETECTION_MODEL = MODEL_DIR / 'yolov9c.pt'
+SEGMENTATION_MODEL = MODEL_DIR / 'yolov9c-seg.pt'