Spaces:
Runtime error
Runtime error
File size: 3,953 Bytes
e9e5323 c1115fd c9ce858 97b86b5 2ccd650 94239c6 97b86b5 c1115fd c9ce858 c1115fd c9ce858 c1115fd 97b86b5 d5cf266 c9ce858 0492776 360336c ddab68e e280fa2 28026d6 d5cf266 0492776 5e3b678 e46d69b 0492776 97b86b5 e46d69b 0492776 6b95018 d5cf266 0492776 c9ce858 6b95018 |
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 |
import os
import torch
import urllib
from PIL import Image
import streamlit as st
from pathlib import Path
def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''):
# Check file(s) for acceptable suffix
if file and suffix:
if isinstance(suffix, str):
suffix = [suffix]
for f in file if isinstance(file, (list, tuple)) else [file]:
s = Path(f).suffix.lower() # file suffix
if len(s):
assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}"
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 :// -> :/
# '%2F' to '/', split https://url.com/file.txt?auth
file = Path(urllib.parse.unquote(file).split('?')[0]).name
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
def read_pretrain(path):
return torch.hub.load('ultralytics/yolov5', 'custom', path=path
default_pretrained = '2022.11.04-YOLOv5x6_1280-Hololive_Waifu_Classification.pt'
with st.spinner('Loading the model...'):
read_pretrain(default_pretrained)
st.title("Hololive Waifu Classification")
image = st.text_input('Image URL', '')
st.info(
'Images for quick tesing:\n \n \n'
' - https://i.imgur.com/tFZwWYw.jpg'
'\n \n \n'
' - https://static.wikia.nocookie.net/omniversal-battlefield/images/b/bd/Council.jpg'
'\n \n \n'
' - https://rare-gallery.com/uploads/posts/951368-anime-anime-girls-digital-art-artwork-2D-portrait.jpg'
'\n \n \n'
' - https://megapx-assets.dcard.tw/images/65993ab1-fe08-43be-87cd-2ecd201cacbd/1280.jpeg'
'\n \n \n'
' - https://img.esportsku.com/wp-content/uploads//2021/07/hololive-en.png')
pretrained = st.selectbox('Select pre-trained', ('2022.11.04-YOLOv5x6_1280-Hololive_Waifu_Classification.pt', '2022.11.01-YOLOv5x6_1280-Hololive_Waifu_Classification.pt'))
imgsz = st.number_input(label='Image Size', min_value=None, max_value=None, value=1280, step=1)
conf = st.slider(label='Confidence threshold', min_value=0.0, max_value=1.0, value=0.25, step=0.01)
iou = st.slider(label='IoU threshold', min_value=0.0, max_value=1.0, value=0.45, step=0.01)
multi_label = st.selectbox('Multiple labels per box', (False, True))
agnostic = st.selectbox('Class-agnostic', (False, True))
amp = st.selectbox('Automatic Mixed Precision inference', (False, True))
max_det = st.number_input(label='Maximum number of detections per image', min_value=None, max_value=None, value=1000, step=1)
if st.button('Excute'):
with st.spinner('Loading the image...'):
image_path = check_file(image)
input_image = Image.open(image_path)
if default_pretrained != pretrained:
with st.spinner('Loading the model...'):
# model = torch.hub.load('ultralytics/yolov5', 'custom', path=os.path.join('pretrained', pretrained))
model = torch.hub.load('ultralytics/yolov5', 'custom', path=pretrained)
with st.spinner('Updating configuration...'):
model.conf = float(conf)
model.max_det = int(max_det)
model.iou = float(iou)
model.agnostic = agnostic
model.multi_label = multi_label
model.amp = amp
with st.spinner('Predicting...'):
results = model(input_image, size=int(imgsz))
for img in results.render():
st.image(img)
st.write(results.pandas().xyxy[0])
os.remove(image_path)
|