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', '')

if image != '':
    image = check_file(image)
    input_image = Image.open(image)
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
    results = model(input_image, size=640)
    for img in results.render():
        st.image(img)