Initial commit

apd_utils.py

@@ -0,0 +1,34 @@
+import os
+import gdown
+import subprocess
+
+def write_video(source_vid):
+    os.makedirs('temp', exist_ok=True)
+    temp_uploaded_path = f'temp/{source_vid.name}'
+
+    with open(temp_uploaded_path, mode='wb') as temp:
+        temp.write(source_vid.read())
+
+    return temp_uploaded_path
+
+def convert_video(in_path, out_path):
+    command = [
+        'ffmpeg',
+        '-i', in_path,
+        '-vcodec', 'libx264',
+        '-y',
+        out_path
+    ]
+    subprocess.run(command)
+
+def download_model(url):
+    os.makedirs('models', exist_ok=True)
+    gdown.download(url, 'models/ckpt_best_1.pth', fuzzy=True)
+
+def delete_temp():
+    files = os.listdir('./temp')
+    
+    for file_name in files:
+        file_path = os.path.join(folder_path, file_name)
+        if os.path.isfile(file_path):
+            os.remove(file_path)

app.py

@@ -0,0 +1,110 @@
+from super_gradients.training import models
+from apd_utils import write_video, convert_video
+import torch, PIL, os
+import streamlit as st
+
+CLASSES = ['Dust Mask', 'Eye Wear', 'Glove', 'Protective Boots', 'Protective Helmet', 'Safety Vest', 'Shield']
+SOURCES = ['Images', 'Videos']
+
+# Setting page layout
+st.set_page_config(
+    page_title="PPE Object Detection using YOLO-NAS",
+    page_icon="👷",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Main page heading
+st.title("PPE Object Detection using YOLO-NAS")
+
+# Sidebar
+st.sidebar.header("YOLO-NAS Model Config")
+
+# Model Options
+confidence = float(st.sidebar.slider(
+    "Select Model Confidence", 0, 100, 40)) / 100
+
+st.sidebar.header("Image/Video Config")
+source_radio = st.sidebar.radio("Select Source", SOURCES)
+
+source_img = None
+source_vid = None
+
+#with st.spinner('Downloading model..'):
+    #model_url = 'https://drive.google.com/file/d/1XOq3OkpQ3OgibjHmYOCMsQPBtqjdf2i3/view?usp=sharing'
+    #download_model(model_url)
+
+model = models.get('yolo_nas_m',
+                        num_classes=len(CLASSES),
+                        checkpoint_path="./ckpt_best_yolonas.pth")
+
+device = 'cuda' if torch.cuda.is_available() else "cpu"
+device = 'cpu'
+
+if source_radio == 'Images':
+    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:
+                st.image('default_img.png', 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:
+            st.image('default_img_res.png', caption="Detected Objects",
+                    use_column_width=True)
+        else:
+            if st.sidebar.button('Detect Objects'):
+                res = model.to(device).predict(uploaded_image,
+                                    conf=confidence)
+                st.image(res.draw(), caption='Detected Image',
+                        use_column_width=True)
+
+elif source_radio == 'Videos':
+    source_vid = st.sidebar.file_uploader(
+        "Choose a video ...", type=("mp4", "mov", "webM"))
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        if source_vid is None:
+            st.image('default_img.png', caption="Default Image",
+                    use_column_width=True)
+        else:
+            try:
+                uploaded_video = source_vid.getvalue()
+                st.video(uploaded_video)
+            except Exception as ex:
+                st.error("Error occurred while opening the video.")
+                st.error(ex)
+    with col2:
+        if source_vid is None:
+            st.image('default_img_res.png', caption="Detected Objects",
+                    use_column_width=True)
+        else:
+            if st.sidebar.button('Detect Objects'):
+                temp_uploaded_path = write_video(source_vid)
+                res = model.to(device).predict(temp_uploaded_path, conf=confidence)
+
+                with st.spinner('Processing video ...'):
+                    in_temp_res_path = "./temp/result.mp4"
+                    out_temp_res_path = "./temp/result2.mp4"
+
+                    res.save(in_temp_res_path)
+                    convert_video(in_temp_res_path, out_temp_res_path)
+                st.video(out_temp_res_path)
+
+                os.remove(temp_uploaded_path)
+                os.remove(in_temp_res_path)
+                os.remove(out_temp_res_path)
+else:
+    st.error("Please select a valid source type!")

ckpt_best_yolonas.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a3a5dc7442e873aa894e672388910e87456292e9e4ff941630757fab2754e83
+size 681102850

packages.txt

@@ -0,0 +1,7 @@
+freeglut3-dev
+libgtk2.0-dev
+libgl1-mesa-glx
+tesseract-ocr
+libtesseract-dev
+libtesseract4
+tesseract-ocr-all

requirements.txt

@@ -0,0 +1,168 @@
+absl-py==2.1.0
+alabaster==0.7.16
+albumentations==1.3.1
+altair==5.2.0
+antlr4-python3-runtime==4.9.3
+arabic-reshaper==3.0.0
+asn1crypto==1.5.1
+attrs==23.2.0
+Babel==2.14.0
+beautifulsoup4==4.12.3
+blinker==1.7.0
+boto3==1.34.54
+botocore==1.34.54
+build==1.1.1
+cachetools==5.3.3
+certifi==2024.2.2
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+coloredlogs==15.0.1
+contourpy==1.2.0
+coverage==5.3.1
+cryptography==42.0.5
+cssselect2==0.7.0
+cycler==0.12.1
+data-gradients==0.3.2
+Deprecated==1.2.14
+docutils==0.17.1
+einops==0.3.2
+filelock==3.13.1
+flatbuffers==23.5.26
+fonttools==4.49.0
+fsspec==2024.2.0
+future==1.0.0
+gdown==5.1.0
+gitdb==4.0.11
+GitPython==3.1.42
+grpcio==1.62.0
+html5lib==1.1
+humanfriendly==10.0
+hydra-core==1.3.2
+idna==3.6
+imagededup==0.3.2
+imageio==2.34.0
+imagesize==1.4.1
+importlib-metadata==7.0.1
+importlib_resources==6.1.2
+Jinja2==3.1.3
+jmespath==1.0.1
+joblib==1.3.2
+json-tricks==3.16.1
+jsonschema==4.21.1
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+lazy_loader==0.3
+lxml==5.1.0
+Markdown==3.5.2
+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.23.0
+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.3.101
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+onnx==1.13.0
+onnxruntime==1.13.1
+onnxsim==0.4.35
+opencv-python==4.9.0.80
+opencv-python-headless==4.9.0.80
+oscrypto==1.3.0
+packaging==23.2
+pandas==2.2.1
+pillow==10.2.0
+pip-tools==7.4.0
+platformdirs==4.2.0
+protobuf==3.20.3
+psutil==5.9.8
+pyarrow==15.0.0
+pycocotools==2.0.6
+pycparser==2.21
+pydeck==0.8.1b0
+pyDeprecate==0.3.2
+Pygments==2.17.2
+pyHanko==0.21.0
+pyhanko-certvalidator==0.26.3
+pyparsing==2.4.5
+pypdf==4.0.2
+pypng==0.20220715.0
+pyproject_hooks==1.0.0
+PySocks==1.7.1
+python-bidi==0.4.2
+python-dateutil==2.9.0.post0
+pytz==2024.1
+PyWavelets==1.5.0
+PyYAML==6.0.1
+qrcode==7.4.2
+qudida==0.0.4
+rapidfuzz==3.6.1
+referencing==0.33.0
+reportlab==3.6.13
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.0
+s3transfer==0.10.0
+scikit-image==0.22.0
+scikit-learn==1.4.1.post1
+scipy==1.12.0
+seaborn==0.13.2
+six==1.16.0
+smmap==5.0.1
+snowballstemmer==2.2.0
+soupsieve==2.5
+Sphinx==4.0.3
+sphinx-rtd-theme==1.3.0
+sphinxcontrib-applehelp==1.0.8
+sphinxcontrib-devhelp==1.0.6
+sphinxcontrib-htmlhelp==2.0.5
+sphinxcontrib-jquery==4.1
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==1.0.7
+sphinxcontrib-serializinghtml==1.1.10
+streamlit==1.31.1
+stringcase==1.2.0
+super-gradients==3.6.0
+svglib==1.5.1
+sympy==1.12
+tenacity==8.2.3
+tensorboard==2.16.2
+tensorboard-data-server==0.7.2
+termcolor==1.1.0
+threadpoolctl==3.3.0
+tifffile==2024.2.12
+tinycss2==1.2.1
+toml==0.10.2
+tomli==2.0.1
+toolz==0.12.1
+torch==2.2.1
+torchmetrics==0.8.0
+torchvision==0.17.1
+tornado==6.4
+tqdm==4.66.2
+treelib==1.6.1
+triton==2.2.0
+typing_extensions==4.10.0
+tzdata==2024.1
+tzlocal==5.2
+uritools==4.0.2
+urllib3==1.26.18
+validators==0.22.0
+watchdog==4.0.0
+webencodings==0.5.1
+Werkzeug==3.0.1
+wrapt==1.16.0
+xhtml2pdf==0.2.11
+zipp==3.17.0