Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -5,8 +5,13 @@ import requests
|
|
5 |
from io import BytesIO
|
6 |
|
7 |
# Load the model and processor
|
8 |
-
|
9 |
-
|
|
|
|
|
|
|
|
|
|
|
10 |
|
11 |
# Define prediction function
|
12 |
def predict_image(image):
|
@@ -27,8 +32,9 @@ def predict_image(image):
|
|
27 |
# Streamlit app
|
28 |
st.title("NSFW Image Classifier")
|
29 |
|
30 |
-
# URL
|
31 |
-
image_url = st.
|
|
|
32 |
if image_url:
|
33 |
try:
|
34 |
# Load image from URL
|
@@ -43,3 +49,5 @@ if image_url:
|
|
43 |
st.write(f"Predicted Class: {prediction}")
|
44 |
except Exception as e:
|
45 |
st.write(f"Error: {e}")
|
|
|
|
|
|
5 |
from io import BytesIO
|
6 |
|
7 |
# Load the model and processor
|
8 |
+
@st.experimental_singleton
|
9 |
+
def load_model():
|
10 |
+
processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
|
11 |
+
model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector')
|
12 |
+
return processor, model
|
13 |
+
|
14 |
+
processor, model = load_model()
|
15 |
|
16 |
# Define prediction function
|
17 |
def predict_image(image):
|
|
|
32 |
# Streamlit app
|
33 |
st.title("NSFW Image Classifier")
|
34 |
|
35 |
+
# Get image URL from query parameters
|
36 |
+
image_url = st.experimental_get_query_params().get('image_url', [None])[0]
|
37 |
+
|
38 |
if image_url:
|
39 |
try:
|
40 |
# Load image from URL
|
|
|
49 |
st.write(f"Predicted Class: {prediction}")
|
50 |
except Exception as e:
|
51 |
st.write(f"Error: {e}")
|
52 |
+
else:
|
53 |
+
st.write("Please provide an image URL using the 'image_url' query parameter.")
|