from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from PIL import Image
import numpy as np
import tensorflow as tf
import pickle
import io

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/", StaticFiles(directory=".", html=True), name="static")

model = tf.keras.models.load_model("model_2.h5")
with open("label_encoder.pkl", "rb") as f:
    label_encoder = pickle.load(f)

class_names = label_encoder.inverse_transform(np.arange(len(label_encoder.classes_)))

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    image = image.resize((224, 224))
    img_array = np.array(image) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    prediction = model.predict(img_array)[0]
    predicted_index = np.argmax(prediction)
    predicted_label = class_names[predicted_index]
    confidence = float(prediction[predicted_index]) * 100

    all_probs = {
        class_names[i]: float(prob)
        for i, prob in enumerate(prediction)
    }

    return JSONResponse(content={
        "predicted_label": predicted_label,
        "confidence": round(confidence, 2),
        "all_probabilities": all_probs
    })

@app.get("/")
def redirect_to_ui():
    return RedirectResponse("/index.html")