Update app.py

app.py

@@ -3,15 +3,12 @@ from transformers import ViTImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import requests
 from io import BytesIO
+import json
+from flask import Flask, request, jsonify
 
 # Load the model and processor
-@st.cache_data
-def load_model():
-    processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
-    model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector')
-    return processor, model
-
-processor, model = load_model()
+processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
+model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector')
 
 # Define prediction function
 def predict_image(image):
@@ -32,10 +29,13 @@ def predict_image(image):
 # Streamlit app
 st.title("NSFW Image Classifier")
 
-# Get image URL from query parameters
-query_params = st.query_params()
-image_url = query_params.get('image_url', [None])[0]
+# 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:
     try:
         # Load image from URL
@@ -50,5 +50,28 @@ if image_url:
         st.write(f"Predicted Class: {prediction}")
     except Exception as e:
         st.write(f"Error: {e}")
-else:
-    st.write("Please provide an image URL using the 'image_url' query parameter.")
+
+# 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)