Spaces:
Runtime error
Runtime error
| 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 | |
| st.title("Hololive Waifu Classification") | |
| image = st.text_input('Image URL', '') | |
| imgsz = st.text_input('Image Size', 1280) | |
| conf = st.text_input('Confidence threshold', 0.25) | |
| iou = st.text_input('IoU threshold', 0.45) | |
| 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.text_input('Maximum number of detections per image', 1000) | |
| if st.button('Excute'): | |
| image = check_file(image) | |
| input_image = Image.open(image) | |
| model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt') | |
| 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 | |
| results = model(input_image, size=int(imgsz)) | |
| for img in results.render(): | |
| st.image(img) | |
| st.write(results.pandas().xyxy[0]) |