File size: 3,531 Bytes
6a3a328 6bbd9db 4f70fa6 6a3a328 14a2f01 6a3a328 5da8fea 4659419 4f70fa6 b61deaf 6bbd9db 4f70fa6 6bbd9db 4f70fa6 6a3a328 4f70fa6 6a3a328 d0dd251 4f70fa6 6a3a328 0f897c1 4f70fa6 6a3a328 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
# -*- coding: utf-8 -*-
"""
-------------------------------------------------
@File Name: app.py
@Author: Luyao.zhang
@Date: 2023/5/15
@Description:
-------------------------------------------------
"""
HuggingFace = True
if HuggingFace is False:
from dotenv import load_dotenv
load_dotenv()
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
import os
# setting page layout
st.set_page_config(
page_title="YOLO.dog",
page_icon="🤖",
layout="wide",
initial_sidebar_state="expanded"
)
base_download_path = "downloaded"
#------------------------
if not os.path.exists(base_download_path):
os.makedirs(base_download_path)
if HuggingFace is False:
model_count = int(os.getenv("model_count"))
else:
model_count = int(st.secrets["model_count"])
model_info = {}
models_list = []
for i in range(0, model_count):
if HuggingFace is False:
model_name = os.getenv("m{}_name".format(i))
gdrive_id = os.getenv("m{}_griv".format(i))
model_extname = os.getenv("m{}_type".format(i))
model_desc = os.getenv("m{}_desc".format(i))
else:
model_name = st.secrets["m{}_name".format(i)]
gdrive_id = st.secrets["m{}_griv".format(i)]
model_extname = st.secrets["m{}_type".format(i)]
model_desc = st.secrets["m{}_desc".format(i)]
print(i, model_name, gdrive_id, model_extname, model_desc)
path_model = os.path.join(base_download_path, model_name + model_extname)
print('path_model', path_model)
model_info.update( {model_desc:path_model} )
models_list.append(model_desc)
if not os.path.exists(path_model):
download_link = "https://drive.google.com/uc?export=download&confirm=t&id={}".format(gdrive_id)
#subprocess.Popen( 'gdown {}'.format(download_link)
#if gdrive_id[:4] == "http":
print('wget -O {} --content-disposition "{}"'.format(path_model, download_link))
os.system( 'wget -O {} --content-disposition "{}"'.format(path_model, download_link))
#else:
# download_file_from_google_drive(gdrive_id, path_model)
print('models_list', models_list)
# main page heading
st.title("Models Demo")
# sidebar
st.sidebar.header("DL Model Config")
# model options
task_type = "Detection"
model_type = st.sidebar.selectbox(
"Model types",
tuple(models_list))
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 = model_info[model_type]
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}")
else:
st.error("Please Select Model in Sidebar")
# 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")
|