minimal app

Makefile

@@ -0,0 +1,15 @@
+install:
+	pip install --upgrade pip
+	pip install -r requirements.txt
+
+test:
+	python -m pytest -vv app.py
+
+format:
+	black *.py
+
+lint:
+	pylint --disable=R, C *.py
+
+all:
+	install lint test format

app.py

@@ -0,0 +1,184 @@
+import glob
+import streamlit as st
+import wget
+from PIL import Image
+import torch
+import cv2
+import os
+import time
+
+from ultralyticsplus import YOLO, render_result
+
+
+st.set_page_config(layout="wide")
+
+cfg_model_path = 'models/yolov5s.pt'
+model = None
+confidence = .25
+
+
+def image_input(data_src):
+    img_file = None
+    if data_src == 'Sample data':
+        # get all sample images
+        img_path = glob.glob('data/sample_images/*')
+        img_slider = st.slider("Select a test image.", min_value=1, max_value=len(img_path), step=1)
+        img_file = img_path[img_slider - 1]
+    else:
+        img_bytes = st.sidebar.file_uploader("Upload an image", type=['png', 'jpeg', 'jpg'])
+        if img_bytes:
+            img_file = "data/uploaded_data/upload." + img_bytes.name.split('.')[-1]
+            Image.open(img_bytes).save(img_file)
+
+    if img_file:
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(img_file, caption="Selected Image")
+        with col2:
+            img = infer_image(img_file)
+            st.image(img, caption="Model prediction")
+
+
+def video_input(data_src):
+    vid_file = None
+    if data_src == 'Sample data':
+        vid_file = "data/sample_videos/sample.mp4"
+    else:
+        vid_bytes = st.sidebar.file_uploader("Upload a video", type=['mp4', 'mpv', 'avi'])
+        if vid_bytes:
+            vid_file = "data/uploaded_data/upload." + vid_bytes.name.split('.')[-1]
+            with open(vid_file, 'wb') as out:
+                out.write(vid_bytes.read())
+
+    if vid_file:
+        cap = cv2.VideoCapture(vid_file)
+        custom_size = st.sidebar.checkbox("Custom frame size")
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        if custom_size:
+            width = st.sidebar.number_input("Width", min_value=120, step=20, value=width)
+            height = st.sidebar.number_input("Height", min_value=120, step=20, value=height)
+
+        fps = 0
+        st1, st2, st3 = st.columns(3)
+        with st1:
+            st.markdown("## Height")
+            st1_text = st.markdown(f"{height}")
+        with st2:
+            st.markdown("## Width")
+            st2_text = st.markdown(f"{width}")
+        with st3:
+            st.markdown("## FPS")
+            st3_text = st.markdown(f"{fps}")
+
+        st.markdown("---")
+        output = st.empty()
+        prev_time = 0
+        curr_time = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                st.write("Can't read frame, stream ended? Exiting ....")
+                break
+            frame = cv2.resize(frame, (width, height))
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            output_img = infer_image(frame)
+            output.image(output_img)
+            curr_time = time.time()
+            fps = 1 / (curr_time - prev_time)
+            prev_time = curr_time
+            st1_text.markdown(f"**{height}**")
+            st2_text.markdown(f"**{width}**")
+            st3_text.markdown(f"**{fps:.2f}**")
+
+        cap.release()
+
+
+def infer_image(img, size=None):
+    model.conf = confidence
+    result = model(img, size=size) if size else model(img)
+    result.render()
+    image = Image.fromarray(result.ims[0])
+    return image
+
+
+@st.experimental_singleton
+def load_model(path, device):
+    model_ = torch.hub.load('ultralytics/yolov5', 'custom', path=path, force_reload=True)
+    model_.to(device)
+    print("model to ", device)
+    return model_
+
+
+@st.experimental_singleton
+def download_model(url):
+    model_file = wget.download(url, out="models")
+    return model_file
+
+
+def get_user_model():
+    model_src = st.sidebar.radio("Model source", ["file upload", "url"])
+    model_file = None
+    if model_src == "file upload":
+        model_bytes = st.sidebar.file_uploader("Upload a model file", type=['pt'])
+        if model_bytes:
+            model_file = "models/uploaded_" + model_bytes.name
+            with open(model_file, 'wb') as out:
+                out.write(model_bytes.read())
+    else:
+        url = st.sidebar.text_input("model url")
+        if url:
+            model_file_ = download_model(url)
+            if model_file_.split(".")[-1] == "pt":
+                model_file = model_file_
+
+    return model_file
+
+def main():
+    # global variables
+    global model, confidence, cfg_model_path
+
+    st.title("Object Recognition Dashboard")
+
+    st.sidebar.title("Settings")
+
+    # device options
+    if torch.cuda.is_available():
+        device_option = st.sidebar.radio("Select Device", ['cpu', 'cuda'], disabled=False, index=0)
+    else:
+        device_option = st.sidebar.radio("Select Device", ['cpu', 'cuda'], disabled=True, index=0)
+
+    # load model
+    model = YOLO('ultralyticsplus/yolov8s')
+
+    # confidence slider
+    confidence = st.sidebar.slider('Confidence', min_value=0.1, max_value=1.0, value=.45)
+
+    # custom classes
+    if st.sidebar.checkbox("Custom Classes"):
+        model_names = list(model.names.values())
+        assigned_class = st.sidebar.multiselect("Select Classes", model_names, default=[model_names[0]])
+        classes = [model_names.index(name) for name in assigned_class]
+        model.classes = classes
+    else:
+        model.classes = list(model.names.keys())
+
+    st.sidebar.markdown("---")
+
+    # input options
+    input_option = st.sidebar.radio("Select input type: ", ['image', 'video'])
+
+    # input src option
+    data_src = st.sidebar.radio("Select input source: ", ['Sample data', 'Upload your own data'])
+
+    if input_option == 'image':
+        image_input(data_src)
+    else:
+        video_input(data_src)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except SystemExit:
+        pass

requirements.txt

@@ -0,0 +1,26 @@
+# Base ----------------------------------------
+ultralyticsplus
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python-headless
+Pillow>=7.1.2
+PyYAML>=5.3.1
+requests>=2.23.0
+scipy>=1.4.1  # Google Colab version
+torch>=1.7.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+protobuf<4.21.5  # https://github.com/ultralytics/yolov5/issues/8012
+
+
+# Plotting ------------------------------------
+pandas>=1.1.4
+seaborn>=0.11.0
+
+
+# Extras --------------------------------------
+ipython  # interactive notebook
+psutil  # system utilization
+thop  # FLOPs computation
+streamlit
+wget