File size: 4,470 Bytes
6a3a328 6bbd9db 2589871 6bbd9db 4f70fa6 6a3a328 14a2f01 6a3a328 5da8fea 2a77066 2589871 46779cd 2589871 2a77066 4659419 4f70fa6 b61deaf 6bbd9db 4f70fa6 6bbd9db 4f70fa6 2a77066 6a3a328 9e66fa1 6a3a328 2a77066 6a3a328 d0dd251 4f70fa6 2a77066 6a3a328 a3d2a2b 6a3a328 5489b52 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
# -*- coding: utf-8 -*-
"""
-------------------------------------------------
@File Name: app.py
@Author: Luyao.zhang
@Date: 2023/5/15
@Description:
-------------------------------------------------
"""
HuggingFace = True
update_model_id = 0
base_download_path = "downloaded"
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
query_params = st.experimental_get_query_params()
#------------
if update_model_id is not None:
if HuggingFace is False:
model_name = os.getenv("m{}_name".format(update_model_id))
model_extname = os.getenv("m{}_type".format(update_model_id))
else:
model_name = st.secrets["m{}_name".format(update_model_id)]
model_extname = st.secrets["m{}_type".format(update_model_id)]
path_model = os.path.join(base_download_path, model_name + model_extname)
if os.path.exists(path_model):
try:
os.remove(path_model)
except:
print("Cannot remove", path_model)
pass
#-------------
qmodel = 'crowded_human'
if 'model' in query_params:
qmodel = query_params['model'][0]
# setting page layout
st.set_page_config(
page_title="YOLO.dog",
page_icon="🤖",
layout="wide",
initial_sidebar_state="expanded"
)
#------------------------
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)
if qmodel not in models_list:
qmodel = models_list[0]
# main page heading
#st.title( model_info[qmodel] )
# sidebar
st.sidebar.header("Model Config")
# model options
task_type = "Detection"
model_type = st.sidebar.selectbox(
"Models list",
models_list,
index=models_list.index(qmodel) )
confidence = float(st.sidebar.slider(
"Select Model Confidence", 10, 100, 5)) / 100
if model_type:
st.header('{} Model Trial'.format(model_type),divider='rainbow')
st.subheader('Use your :blue[photo/video] :sunglasses:')
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")
|