Update app.py

app.py

@@ -1,45 +1,71 @@
-import easyocr as ocr  #OCR
 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")
 
-title = st.text_input('Image URL', '')
-st.write('URL: ', title)
+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'])
+# 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(): 
-    reader = ocr.Reader(['en'],model_storage_directory='.')
-    return reader 
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5x6', path='best.pt')  # local model
 
 reader = load_model() #load model
 
-if image is not None:
-
+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! "):
+    #with st.spinner("🤖 AI is at Work! "):
         
 
-        result = reader.readtext(np.array(input_image))
+        #result = reader.readtext(np.array(input_image))
 
-        result_text = [] #empty list for results
+        #result_text = [] #empty list for results
 
 
-        for text in result:
-            result_text.append(text[1])
+        #for text in result:
+            #result_text.append(text[1])
 
-        st.write(result_text)
+        #st.write(result_text)
     #st.success("Here you go!")
-    st.balloons()
-else:
-    st.write("Upload an Image")
+    #st.balloons()
+#else:
+    #st.write("Upload an Image")
 
-st.caption("neko941")
+st.caption("authors: neko941")