Spaces:
Running
Running
| import streamlit as st | |
| 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 | |
| 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 | |
| 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: | |
| try: | |
| # Load image from URL | |
| response = requests.get(image_url) | |
| 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 | |
| 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__) | |
| 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) | |