import uvicorn
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import tensorflow as tf
import numpy as np
from PIL import Image
from io import BytesIO

# Load model
model = tf.keras.models.load_model("Image_model.keras", compile=False) # Ignore custom metric
IMG_SIZE = (299, 299)

app = FastAPI(title="Hate Speech Image Classifier")

def preprocess_image(img: Image.Image):
    img = img.resize(IMG_SIZE)
    img_array = np.array(img) / 255.0
    return np.expand_dims(img_array, axis=0)


@app.post("/predict/")
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        # Read image
        contents = await file.read()
        img = Image.open(BytesIO(contents)).convert("RGB")
        
        # Preprocess
        img_array = preprocess_image(img)
        
        # Predict
        prediction = model.predict(img_array)[0][0]
        label = "Hate Speech" if prediction >= 0.5 else "Non-Hate Speech"
        confidence = round(float(prediction), 4)
        
        return JSONResponse({
            "prediction": label,
            "confidence": confidence
        })
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

@app.get("/")
async def root():
    return {"message": "Welcome to Hate Speech Image Classifier API"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)