|
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 |
|
|
|
|
|
processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector') |
|
model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector') |
|
|
|
|
|
def predict_image(image): |
|
try: |
|
|
|
inputs = processor(images=image, return_tensors="pt") |
|
outputs = model(**inputs) |
|
logits = outputs.logits |
|
|
|
|
|
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) |
|
|
|
|
|
st.title("NSFW Image Classifier") |
|
|
|
|
|
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\"}'") |
|
|
|
|
|
image_url = st.text_input("Enter Image URL", placeholder="Enter image URL here") |
|
if image_url: |
|
try: |
|
|
|
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...") |
|
|
|
|
|
prediction = predict_image(image) |
|
st.write(f"Predicted Class: {prediction}") |
|
except Exception as e: |
|
st.write(f"Error: {e}") |
|
|
|
|
|
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: |
|
|
|
response = requests.get(image_url) |
|
image = Image.open(BytesIO(response.content)) |
|
|
|
|
|
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) |
|
|