hmrhmr/IA-proyect-uem

1import gradio as gr2import numpy as np3import tensorflow as tf4import cv25from PIL import Image6 7# Cargar modelo TFLite8interpreter = tf.lite.Interpreter(model_path="traffic_sign_lite.tflite")9interpreter.allocate_tensors()10 11input_details = interpreter.get_input_details()12output_details = interpreter.get_output_details()13 14# Diccionario de clases15class_names = {16    0: 'Limite 20km/h',17    1: 'Limite 30km/h',18    2: 'Limite 50km/h',19    3: 'Limite 60km/h',20    4: 'Limite 70km/h',21    5: 'Limite 80km/h',22    6: 'Fin limite 80km/h',23    7: 'Limite 100km/h',24    8: 'Limite 120km/h',25    9: 'Prohibido adelantar',26    10: 'Prohibido adelantar (>3.5t)',27    11: 'Cruce con prioridad',28    12: 'Via prioritaria',29    13: 'Ceda el paso',30    14: 'Stop',31    15: 'Prohibido el paso',32    16: 'Prohibido (>3.5t)',33    17: 'Prohibida la entrada',34    18: 'Peligro general',35    19: 'Curva peligrosa izq.',36    20: 'Curva peligrosa der.',37    21: 'Doble curva',38    22: 'Baches',39    23: 'Pavimento deslizante',40    24: 'Estrecho por derecha',41    25: 'Obras',42    26: 'Semaforo',43    27: 'Peatones',44    28: 'Ninos cruzando',45    29: 'Ciclistas',46    30: 'Nieve/hielo',47    31: 'Animales salvajes',48    32: 'Fin prohibiciones',49    33: 'Giro obligatorio derecha',50    34: 'Giro obligatorio izq.',51    35: 'Solo adelante',52    36: 'Recto o derecha',53    37: 'Recto o izquierda',54    38: 'Mantener derecha',55    39: 'Mantener izquierda',56    40: 'Rotonda obligatoria',57    41: 'Fin prohib. adelantar',58    42: 'Fin prohib. adelantar (>3.5t)'59}60 61def predict_traffic_sign(image):62    try:63        if isinstance(image, Image.Image):64            image = np.array(image)65        66        img = cv2.resize(image, (32, 32))67        img = img.astype('float32') / 255.068        img = np.expand_dims(img, axis=0)69        70        interpreter.set_tensor(input_details[0]['index'], img)71        interpreter.invoke()72        predictions = interpreter.get_tensor(output_details[0]['index'])[0]73        74        top_5_indices = np.argsort(predictions)[-5:][::-1]75        76        results = {}77        for idx in top_5_indices:78            class_name = class_names.get(idx, f"Clase {idx}")79            confidence = float(predictions[idx])80            results[class_name] = confidence81        82        return results83        84    except Exception as e:85        return {"Error": str(e)}86 87demo = gr.Interface(88    fn=predict_traffic_sign,89    inputs=gr.Image(type="pil", label="Sube una imagen de señal de trafico"),90    outputs=gr.Label(num_top_classes=5, label="Prediccion"),91    title="Clasificador de Señales de Trafico",92    description="""93    ## Sube una imagen de una señal de trafico94    95    Identifica 43 tipos diferentes de señales96    Muestra las 5 predicciones mas probables97    Precision del modelo: ~90%98    """,99    theme="default"100)101 102if __name__ == "__main__":103    demo.launch()