Update app.py

app.py

@@ -1,17 +1,11 @@
 import os
-import streamlit as st  #Web App
-from PIL import Image #Image Processing
-import numpy as np #Image Processing 
 import torch
+import urllib
+from PIL import Image
+import streamlit as st
+from pathlib import Path
 
-#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_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''):
     # Check file(s) for acceptable suffix
     if file and suffix:
@@ -31,52 +25,26 @@ def check_file(file, suffix=''):
         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
+        # '%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
-    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'."
+            assert Path(file).exists() and Path(file).stat(
+            ).st_size > 0, f'File download failed: {url}'  # check
         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
 
-# model = load_model() #load model
+st.title("Hololive Waifu Classification")
+
+image = st.text_input('Image URL', '')
 
 if image != '':
     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")
+    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)