from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
import tensorflow as tf
import numpy as np
import cv2

app = FastAPI()

# Cargar el modelo TFLite
interpreter = tf.lite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()

# Obtener detalles de las entradas y salidas del modelo
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Función para preprocesar la imagen
def preprocess_image(image):
    image = cv2.resize(image, (224, 224))
    image = image / 255.0
    image = np.expand_dims(image, axis=0).astype(np.float32)
    return image

# Ruta de predicción
@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    try:
        # Leer la imagen
        image = await file.read()
        image = cv2.imdecode(np.frombuffer(image, np.uint8), cv2.IMREAD_COLOR)
        
        # Preprocesar la imagen
        processed_image = preprocess_image(image)
        
        # Realizar la predicción
        interpreter.set_tensor(input_details[0]['index'], processed_image)
        interpreter.invoke()
        output_data = interpreter.get_tensor(output_details[0]['index'])

        # Determinar la clase y la confianza
        class_idx = np.argmax(output_data[0])
        labels = ['Benign', 'Malignant']
        result = labels[class_idx]
        confidence = float(output_data[0][class_idx])

        return {"class": result, "confidence": confidence}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))