neko941's picture
Update app.py
c1115fd
raw
history blame
2.54 kB
import streamlit as st #Web App
from PIL import Image #Image Processing
import numpy as np #Image Processing
import torch
#title
st.title("Hololive Waifu Classification")
image = st.text_input('Image URL', '')
# st.write('URL: ', title)
#image uploader
# image = st.file_uploader(label = "Upload your image here",type=['png','jpg','jpeg'])
def check_file(file, suffix=''):
# Search/download file (if necessary) and return path
check_suffix(file, suffix) # optional
file = str(file) # convert to str()
if os.path.isfile(file) or not file: # exists
return file
elif file.startswith(('http:/', 'https:/')): # download
url = file # warning: Pathlib turns :// -> :/
file = Path(urllib.parse.unquote(file).split('?')[0]).name # '%2F' to '/', split https://url.com/file.txt?auth
if os.path.isfile(file):
print(f'Found {url} locally at {file}') # file already exists
else:
print(f'Downloading {url} to {file}...')
torch.hub.download_url_to_file(url, file)
assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
return file
elif file.startswith('clearml://'): # ClearML Dataset ID
assert 'clearml' in sys.modules, "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
return file
else: # search
files = []
for d in 'data', 'models', 'utils': # search directories
files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True)) # find file
assert len(files), f'File not found: {file}' # assert file was found
assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
return files[0]
@st.cache
def load_model():
model = torch.hub.load('ultralytics/yolov5', 'yolov5x6', path='best.pt') # local model
reader = load_model() #load model
if title is not None:
image = check_file(image)
st.write('File: ', image)
input_image = Image.open(image) #read image
st.image(input_image) #display image
#with st.spinner("πŸ€– AI is at Work! "):
#result = reader.readtext(np.array(input_image))
#result_text = [] #empty list for results
#for text in result:
#result_text.append(text[1])
#st.write(result_text)
#st.success("Here you go!")
#st.balloons()
#else:
#st.write("Upload an Image")
st.caption("authors: neko941")