yeftakun commited on
Commit
3155a4d
1 Parent(s): 697c4ae

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -4
app.py CHANGED
@@ -5,8 +5,13 @@ import requests
5
  from io import BytesIO
6
 
7
  # Load the model and processor
8
- processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
9
- model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector')
 
 
 
 
 
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 input
31
- image_url = st.text_input("Enter Image URL", placeholder="Enter image URL here")
 
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.")