Update app.py

app.py

@@ -1,39 +1,54 @@
-from fastapi import FastAPI, UploadFile, File
-from fastapi.middleware.cors import CORSMiddleware
-import requests
-import base64
-import logging
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+import io
+import uvicorn
 
-# Configuración básica
-app = FastAPI()
+# Cargar etiquetas y modelo
+with open("labels_mobilenet_quant_v1_224.txt", "r") as f:
+    labels = f.read().splitlines()
+
+interpreter = tf.lite.Interpreter(model_path="mobilenet_v1_1.0_224_quant.tflite")
+interpreter.allocate_tensors()
+
+input_details = interpreter.get_input_details()
+output_details = interpreter.get_output_details()
+
+# Clasificación
+def classify_image(image: Image.Image):
+    image = image.convert("RGB").resize((224, 224))
+    input_data = np.expand_dims(np.array(image, dtype=np.uint8), axis=0)
+
+    interpreter.set_tensor(input_details[0]["index"], input_data)
+    interpreter.invoke()
+    output_data = interpreter.get_tensor(output_details[0]["index"])
 
-# Agregar soporte de CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Cambia esto a tus dominios específicos si es necesario
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+    # Ajuste de cuantización
+    output_scale, output_zero_point = output_details[0]["quantization"]
+    output = output_scale * (output_data.astype(np.float32) - output_zero_point)
 
-# URL del modelo en Hugging Face Space
-HUGGINGFACE_API = "https://mobilenetv1-tflite-demo.myloiose.hf.space/predict"
+    pred_idx = np.argmax(output[0])
+    pred_label = labels[pred_idx] if pred_idx < len(labels) else "Etiqueta no encontrada"
+
+    # Softmax
+    exp_scores = np.exp(output[0] - np.max(output[0]))
+    probabilities = exp_scores / np.sum(exp_scores)
+    confidence = probabilities[pred_idx]
+
+    return pred_label, float(confidence)
+
+# Crear API
+app = FastAPI()
 
-# Endpoint para hacer la predicción
-@app.post("/predict/")
+@app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # Leer el archivo y convertirlo a base64
-    img_bytes = await file.read()
-    img_b64 = "data:image/jpeg;base64," + base64.b64encode(img_bytes).decode()
-
-    # Crear el payload
-    payload = {"data": [img_b64]}
-
-    try:
-        # Realizar la petición al Space de Hugging Face
-        response = requests.post(HUGGINGFACE_API, json=payload)
-        response.raise_for_status()  # Asegura que no haya errores de red
-        return response.json()  # Retorna la respuesta del Space
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Error al conectar con Hugging Face: {str(e)}")
-        return {"error": "Error en la predicción", "details": str(e)}
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents))
+    label, conf = classify_image(image)
+    return JSONResponse({"label": label, "confidence": f"{conf*100:.2f}%"})
 
+# Para correr localmente
+# if __name__ == "__main__":
+#     uvicorn.run(app, host="0.0.0.0", port=7860)