from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
import io
from app import predict  # Import the predict function from your Gradio code (rename file to app.py if needed)

app = FastAPI()

@app.post("/predict/")
async def predict_endpoint(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")

    pred_labels_and_probs, pred_time, symptoms, causes, treatments = predict(image)

    top_condition = max(pred_labels_and_probs, key=pred_labels_and_probs.get)
    top_confidence = pred_labels_and_probs[top_condition]

    return JSONResponse({
        "prediction": top_condition,
        "confidence": top_confidence,
        "all_probabilities": pred_labels_and_probs,
        "prediction_time": pred_time,
        "symptoms": symptoms,
        "causes": causes,
        "treatments": treatments
    })