Spaces:
No application file
No application file
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 | |
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)) | |