# -*- coding: utf-8 -*- HuggingFace = True update_model_id = None base_download_path = "models" 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" ) mnames = { 'cat_body_parts':['.pt','Cat_Body_Parts'], 'road_defects':['.pt','Road_Defects'], '13_animals':['.pt','13-Kinds-Animals'], 'digger':['.pt','Digger'], 'fire_smoke':['.pt','Fire_Smoke'], 'gun_knife':['.pt','Gun-Knife'], 'Type_of_car':['.pt','Car_Types'], 'working_security':['.pt', 'Working-Safety'], 'crowded_human':['.pt','Crowded-Human'], 'face_mask_eyeballs':['.pt', 'Face_Mask_Eyeballs'], 'forklift':['.pt','Forklift'], 'Human_Palm':['.pt','Human_Palm'], 'vehicles_plates':['.pt','Vehicles_Plates'], 'like_dislike':['.pt','Like_Dislike'], 'yolo_gender':['.pt','Human_Gender'] } query_params = st.experimental_get_query_params() qmodel = 'Crowded-Human' if 'model' in query_params: qmodel = query_params['model'][0] #------------------------ 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 model_name in mnames: for i in range(0, model_count): if HuggingFace is False: model_name = os.getenv("m{}_name".format(i)) model_extname = os.getenv("m{}_type".format(i)) model_desc = os.getenv("m{}_desc".format(i)) model_url = os.getenv("m{}_url".format(i)) else: model_name = st.secrets["m{}_name".format(i)] model_extname = st.secrets["m{}_type".format(i)] model_desc = st.secrets["m{}_desc".format(i)] model_url = st.secrets["m{}_url".format(i)] #model_extname = mnames[model_name][0] #model_desc = mnames[model_name][1] 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 = model_url #download_link = "https://drive.google.com/file/d/{}/view?usp=sharing".format(gdrive_id) #print('wget -O {} --content-disposition "{}"'.format(path_model, download_link)) #os.system( 'wget -O {} --content-disposition "{}"'.format(path_model, download_link)) print('wget {} '.format(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) ) 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_selectbox = "Image" confidence = float(st.sidebar.slider( "Adjust Model Confidence", 5, 100, 25)) / 100 drawBox = st.sidebar.checkbox("Draw box", value=True) drawLabel = st.sidebar.checkbox("disply label", value=True) drawScore = st.sidebar.checkbox("disply confidence", value=True) source_img = None if source_selectbox == config.SOURCES_LIST[0]: # Image infer_uploaded_image(confidence, drawBox, drawLabel, drawScore, 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")