Update app.py

app.py

@@ -3,7 +3,6 @@ from transformers import ViTImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import requests
 from io import BytesIO
-import json
 
 # Load the model and processor
 processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
@@ -25,57 +24,17 @@ def predict_image(image):
  except Exception as e:
  return str(e)
 
-# Streamlit app for UI and API endpoint
+# Streamlit app
 st.title("NSFW Image Classifier")
 
-# URL input for UI
-image_url_ui = st.text_input("Enter Image URL", placeholder="Enter image URL here")
-
-# API endpoint for classification (POST request)
-@st.experimental_singleton # Ensure a single instance for performance
-def api_endpoint():
- if request.method == 'POST':
- data = request.json
- if 'image_url' in data:
- try:
- image_url = data['image_url']
- # Load image from URL
- response = requests.get(image_url)
- image = Image.open(BytesIO(response.content))
-
- # Predict and return result as JSON
- prediction = predict_image(image)
- return json.dumps({'predicted_class': prediction})
- except Exception as e:
- return json.dumps({'error': str(e)}), 500 # Internal Server Error
- else:
- return json.dumps({'error': 'Missing "image_url" in request body'}), 400 # Bad Request
- else:
- return json.dumps({'error': 'Only POST requests are allowed'}), 405 # Method Not Allowed
-
-st.experimental_next_router(api_endpoint) # Register the API endpoint
-
-if image_url_ui:
- try:
- # Load image from UI input (if URL is provided)
- response = requests.get(image_url_ui)
- image = Image.open(BytesIO(response.content))
- st.image(image, caption='Image from URL', use_column_width=True)
- st.write("")
- st.write("Classifying...")
-
- # Predict and display result (for UI)
- prediction = predict_image(image)
- st.write(f"Predicted Class: {prediction}")
- except Exception as e:
- st.write(f"Error: {e}")
-
-# Display API endpoint information
-space_url = st.session_state.get('huggingface_space_url') # Assuming it's available
-if space_url:
- api_endpoint_url = f"{space_url}/api/classify" # Construct the URL based on Space URL
- st.write(f"You can also use this API endpoint to classify images:")
- st.write(f"```curl")
- st.write(f"curl -X POST -H 'Content-Type: application/json' -d '{{ \"image_url\": \"https://example.jpg\" }}' {api_endpoint_url}")
- st.write(f"```")
- st.write(f"This will return the predicted class in JSON format.")
+# Upload image file
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+ image = Image.open(uploaded_file)
+ st.image(image, caption='Uploaded Image.', use_column_width=True)
+ st.write("")
+ st.write("Classifying...")
+
+ # Predict and display result
+ prediction = predict_image(image)
+ st.write(f"Predicted Class: {prediction}")