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.")