Spaces:
Running
Running
| import streamlit as st | |
| 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') | |
| model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector') | |
| # Define prediction function | |
| def predict_image(image): | |
| try: | |
| # Process the image and make prediction | |
| inputs = processor(images=image, return_tensors="pt") | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| # Get predicted class | |
| predicted_class_idx = logits.argmax(-1).item() | |
| predicted_label = model.config.id2label[predicted_class_idx] | |
| return predicted_label | |
| except Exception as e: | |
| return str(e) | |
| # Streamlit app for UI and API endpoint | |
| 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) | |
| # 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.") |