main.py

import os
import numpy as np
import tensorflow as tf
from PIL import Image
import io
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse
from contextlib import asynccontextmanager

# --- CONFIGURACIÓN Y MAPEO DE CLASES ---

MODEL_PATHS = {
    "basico": "modelos/modelo_basico.keras",
    "cnn": "modelos/modelo_cnn_2.keras",
    "data_aug": "modelos/modelo_cnn_color_data.keras",
    "dropout": "modelos/modelo_cnn_color_data_dropout2.keras",
    "cnn_gris": "modelos/modelo_cnn_gris.keras"
}

# Diccionario completo de las 38 clases (PlantVillage estándar)
PLANT_CLASSES = [
    "Manzana: Escara", "Manzana: Podredumbre negra", "Manzana: Roya del cedro", "Manzana: Sana",
    "Arándano: Sano", "Cereza: Oídio", "Cereza: Sana", "Maíz: Cercospora (Mancha gris)",
    "Maíz: Roya común", "Maíz: Tizón del norte", "Maíz: Sano", "Uva: Podredumbre negra",
    "Uva: Escariosis", "Uva: Mildiu", "Uva: Sana", "Naranja: Huanglongbing (Greening)",
    "Melocotón: Mancha bacteriana", "Melocotón: Sano", "Pimiento: Mancha bacteriana", "Pimiento: Sano",
    "Patata: Tizón temprano", "Patata: Tizón tardío", "Patata: Sana", "Frambuesa: Sana",
    "Soja: Sana", "Calabaza: Oídio", "Fresa: Mancha foliar", "Fresa: Sana",
    "Tomate: Mancha bacteriana", "Tomate: Tizón temprano", "Tomate: Tizón tardío", "Tomate: Moho foliar",
    "Tomate: Mancha Septoria", "Tomate: Araña roja (Ácaros)", "Tomate: Mancha diana",
    "Tomate: Virus del rizado amarillo", "Tomate: Virus del mosaico", "Tomate: Sano"
]

models = {}

# --- ARQUITECTURAS MANUALES ---

def crear_arquitectura_color(input_shape=(200, 200, 3)):
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=input_shape),
        tf.keras.layers.Resizing(200, 200),
        tf.keras.layers.Rescaling(1./255),
        tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(128, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(256, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dense(38, activation='softmax')
    ])
    return model

def crear_arquitectura_dropout(input_shape=(200, 200, 3)):
    """Arquitectura corregida para coincidir con los pesos (Shape 256, 256)."""
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=input_shape),
        tf.keras.layers.Resizing(200, 200),
        tf.keras.layers.Rescaling(1./255),
        
        tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        
        tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        
        tf.keras.layers.Conv2D(128, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        
        # Añadimos una capa extra de 256 para que el Global sea de 256
        tf.keras.layers.Conv2D(256, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        
        tf.keras.layers.GlobalAveragePooling2D(), 
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(38, activation='softmax')
    ])
    return model

def crear_arquitectura_gris(input_shape=(200, 200, 1)):
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=input_shape),
        tf.keras.layers.Resizing(200, 200),
        tf.keras.layers.Rescaling(1./255),
        tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(128, (3, 3), padding='same', activation='relu'),
        tf.keras.layers.BatchNormalization(axis=3),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(38, activation='softmax')
    ])
    return model

# --- LÓGICA DE CARGA Y APP ---

def load_all_models():
    for mid, path in MODEL_PATHS.items():
        if not os.path.exists(path):
            print(f"⚠️ Archivo no encontrado: {path}")
            continue
        try:
            models[mid] = tf.keras.models.load_model(path, compile=False)
            print(f"✅ {mid} cargado directamente.")
        except Exception:
            try:
                if mid == "cnn_gris": model = crear_arquitectura_gris()
                elif mid == "dropout": model = crear_arquitectura_dropout()
                else: model = crear_arquitectura_color()
                model.load_weights(path)
                models[mid] = model
                print(f"✅ {mid} reconstruido manualmente.")
            except Exception as e:
                print(f"❌ Error crítico en {mid}: {e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    load_all_models()
    yield
    models.clear()

app = FastAPI(lifespan=lifespan)

@app.post("/predict_all")
async def predict_all(file: UploadFile = File(...)):
    if not models:
        raise HTTPException(status_code=500, detail="No hay modelos cargados.")
    
    try:
        contents = await file.read()
        results = []
        names = {
            "basico": "Modelo Lineal Básico",
            "cnn": "CNN Color V2",
            "dropout": "CNN Dropout (Mejorada)",
            "data_aug": "CNN Data Augmentation",
            "cnn_gris": "CNN Escala de Grises"
        }

        for mid, model in models.items():
            try:
                image_bytes = io.BytesIO(contents)
                img = Image.open(image_bytes).resize((200, 200))
                
                # Preprocesamiento según el modelo
                if "gris" in mid: 
                    img = img.convert("L")
                    img_array = np.array(img).astype('float32') / 255.0
                    img_array = np.expand_dims(img_array, axis=(0, -1))
                else:
                    img = img.convert("RGB")
                    img_array = np.array(img).astype('float32') / 255.0
                    img_array = np.expand_dims(img_array, axis=0)
                
                prediction = model.predict(img_array, verbose=0)
                
                # Lógica para obtener el nombre de la clase
                idx = np.argmax(prediction[0])
                conf = float(prediction[0][idx])
                
                # Intentar obtener el nombre de PLANT_CLASSES, si no, usar el índice
                label = PLANT_CLASSES[idx] if idx < len(PLANT_CLASSES) else f"Clase {idx}"
                
                # Clasificación binaria simple para la UI (Sana/Enferma)
                estado = "Sana" if "Sana" in label or "Sano" in label else "Enferma"

                results.append({
                    "id": mid,
                    "name": names.get(mid, mid),
                    "prediction": label,  # Ahora muestra el nombre específico
                    "status": estado,     # Para filtros rápidos en el frontend
                    "confidence": round(conf * 100, 2)
                })
                
            except Exception as e:
                print(f"Error prediciendo con {mid}: {e}")
                continue

        return results

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error procesando imagen: {str(e)}")

@app.get("/")
def home():
    return FileResponse("index.html")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)