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