# -*- coding: utf-8 -*- """ ------------------------------------------------- @File Name: app.py @Author: Luyao.zhang @Date: 2023/5/15 @Description: ------------------------------------------------- """ HuggingFace = True update_model_id = None 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): try: os.makedirs(base_download_path) except: print("folder exists", base_download_path) pass 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")