Update app.py

app.py

@@ -4,7 +4,6 @@ from PIL import Image
 import requests
 from io import BytesIO
 import json
-from flask import Flask, request, jsonify
 
 # Load the model and processor
 processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
@@ -26,52 +25,57 @@ def predict_image(image):
  except Exception as e:
  return str(e)
 
-# Streamlit app
+# Streamlit app for UI and API endpoint
 st.title("NSFW Image Classifier")
 
-# Display API usage instructions
-st.write("You can use this app with the API endpoint below. Send a POST request with the image URL to get classification.")
-st.write("Example URL to use with curl:")
-st.code("curl -X POST https://huggingface.co/spaces/yeftakun/nsfw_api2/api/classify -H 'Content-Type: application/json' -d '{\"image_url\": \"https://example.jpg\"}'")
-
 # URL input for UI
-image_url = st.text_input("Enter Image URL", placeholder="Enter image URL here")
-if image_url:
+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 URL
- response = requests.get(image_url)
+ # 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
+ # 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}")
 
-# API Endpoint using Flask
-app = Flask(__name__)
-
-@app.route('/api/classify', methods=['POST'])
-def classify():
- data = request.json
- image_url = data.get('image_url')
- 
- if not image_url:
- return jsonify({"error": "Image URL is required"}), 400
- 
- try:
- # Load image from URL
- response = requests.get(image_url)
- image = Image.open(BytesIO(response.content))
- 
- # Predict image
- prediction = predict_image(image)
- return jsonify({"predicted_class": prediction})
- except Exception as e:
- return jsonify({"error": str(e)}), 500
-
-if __name__ == '__main__':
- app.run(port=5000)
+# 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.")