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

# Ambil api key dari .env
import os
from dotenv import load_dotenv
load_dotenv()

API_KEY = os.getenv('API_KEY')

app = Flask(__name__)

# Definisikan ukuran gambar
img_width, img_height = 150, 150

# Memuat model
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)