app.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
-------------------------------------------------
   @File Name:     app.py
   @Author:        Luyao.zhang
   @Date:          2023/5/15
   @Description:
-------------------------------------------------
"""
from pathlib import Path
from PIL import Image
import streamlit as st

import config
from utils import load_model, infer_uploaded_image, infer_uploaded_video, infer_uploaded_webcam
from download_gd import *
import os

if not os.path.exists('./crowded_human_yolov8s.pt'):
    download_file_from_google_drive('1qCXBDy3YuxS9bqfRIc--cRLo1L3DAYq2', './crowded_human_yolov8s.pt')

if os.path.exists('./crowded_human_yolov8s.pt'):
    file_exist = "./crowded_human_yolov8s.pt"
else:
    file_exist = "No file"

# setting page layout
st.set_page_config(
    page_title="YOLO.dog",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
    )

# main page heading
st.title(file_exist)

# sidebar
st.sidebar.header("DL Model Config")

# model options
task_type = "Detection"
model_type = "crowded_human"

confidence = float(st.sidebar.slider(
    "Select Model Confidence", 30, 100, 50)) / 100

if model_type:
    #model_path = Path(config.DETECTION_MODEL_DIR, str(model_type))
    model_path = "./crowded_human_yolov8s.pt"
else:
    st.error("Please Select Model in Sidebar")

# load pretrained DL model
try:
    print('model_path', model_path)
    model = load_model(model_path)
except Exception as e:
    st.error(f"Unable to load model. Please check the specified path: {model_path}")

# image/video options
st.sidebar.header("Image/Video Config")
source_selectbox = st.sidebar.selectbox(
    "Select Source",
    config.SOURCES_LIST
)

source_img = None
if source_selectbox == config.SOURCES_LIST[0]: # Image
    infer_uploaded_image(confidence, model)
elif source_selectbox == config.SOURCES_LIST[1]: # Video
    infer_uploaded_video(confidence, model)
elif source_selectbox == config.SOURCES_LIST[2]: # Webcam
    infer_uploaded_webcam(confidence, model)
else:
    st.error("Currently only 'Image' and 'Video' source are implemented")