|
from flask import Flask, request, jsonify |
|
from tensorflow.keras.models import load_model |
|
from tensorflow.keras.preprocessing import image |
|
import numpy as np |
|
import requests |
|
from io import BytesIO |
|
from PIL import Image |
|
|
|
|
|
import os |
|
from dotenv import load_dotenv |
|
load_dotenv() |
|
|
|
API_KEY = os.getenv('API_KEY') |
|
|
|
app = Flask(__name__) |
|
|
|
|
|
img_width, img_height = 150, 150 |
|
|
|
|
|
model = load_model('inceptionv3_nsfw_model.h5') |
|
|
|
def classify_image(img): |
|
img = img.resize((img_width, img_height)) |
|
img_array = image.img_to_array(img) |
|
img_array = np.expand_dims(img_array, axis=0) / 255.0 |
|
|
|
predictions = model.predict(img_array) |
|
class_names = ['sexy', 'hentai', 'porn', 'neutral', 'drawings'] |
|
result = dict(zip(class_names, predictions[0].astype(float))) |
|
|
|
return result |
|
|
|
def check_api_key(request): |
|
api_key = request.headers.get('x-api-key') |
|
return api_key == API_KEY |
|
|
|
@app.route('/classify', methods=['POST']) |
|
def classify(): |
|
if not check_api_key(request): |
|
return jsonify({'error': 'Invalid or missing API key'}), 403 |
|
|
|
data = request.json |
|
img_url = data.get('image_url') |
|
|
|
if not img_url: |
|
return jsonify({'error': 'No image URL provided'}), 400 |
|
|
|
try: |
|
response = requests.get(img_url) |
|
img = Image.open(BytesIO(response.content)) |
|
except Exception as e: |
|
return jsonify({'error': str(e)}), 400 |
|
|
|
result = classify_image(img) |
|
return jsonify(result) |
|
|
|
if __name__ == '__main__': |
|
app.run(host='0.0.0.0', port=5000) |
|
|